import tushare as ts
import MySQLdb
import math
import time

from config import CONFIG_DB 

db = MySQLdb.connect(CONFIG_DB['host'], CONFIG_DB['user'], CONFIG_DB['password'], CONFIG_DB['database'], charset='utf8')
pro = ts.pro_api('f4b683b73b5208105564cb2a086f9c687268b8c4a50d7534853889e3')

def str2float( str ):
    try:
        val = float(str)
        if math.isinf(val) or math.isnan(val) :
            val = 0
        return val
    except Exception as e:
        return 0.0

def queryts( code ) :
    df = pro.income(**{
        "ts_code": code,
        "ann_date": "",
        "f_ann_date": "",
        "start_date": 20220630,
        "end_date": 20230604,
        "period": "",
        "report_type": "",
        "comp_type": "",
        "end_type": "",
        "is_calc": "",
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",
        "ann_date",
        "f_ann_date",
        "end_date",
        "report_type",
        "comp_type",
        "end_type",
        "basic_eps",
        "diluted_eps",
        "total_revenue",
        "revenue",
        "int_income",
        "prem_earned",
        "comm_income",
        "n_commis_income",
        "n_oth_income",
        "n_oth_b_income",
        "prem_income",
        "out_prem",
        "une_prem_reser",
        "reins_income",
        "n_sec_tb_income",
        "n_sec_uw_income",
        "n_asset_mg_income",
        "oth_b_income",
        "fv_value_chg_gain",
        "invest_income",
        "ass_invest_income",
        "forex_gain",
        "total_cogs",
        "oper_cost",
        "int_exp",
        "comm_exp",
        "biz_tax_surchg",
        "sell_exp",
        "admin_exp",
        "fin_exp",
        "assets_impair_loss",
        "prem_refund",
        "compens_payout",
        "reser_insur_liab",
        "div_payt",
        "reins_exp",
        "oper_exp",
        "compens_payout_refu",
        "insur_reser_refu",
        "reins_cost_refund",
        "other_bus_cost",
        "operate_profit",
        "non_oper_income",
        "non_oper_exp",
        "nca_disploss",
        "total_profit",
        "income_tax",
        "n_income",
        "n_income_attr_p",
        "minority_gain",
        "oth_compr_income",
        "t_compr_income",
        "compr_inc_attr_p",
        "compr_inc_attr_m_s",
        "ebit",
        "ebitda",
        "insurance_exp",
        "undist_profit",
        "distable_profit",
        "rd_exp",
        "fin_exp_int_exp",
        "fin_exp_int_inc",
        "transfer_surplus_rese",
        "transfer_housing_imprest",
        "transfer_oth",
        "adj_lossgain",
        "withdra_legal_surplus",
        "withdra_legal_pubfund",
        "withdra_biz_devfund",
        "withdra_rese_fund",
        "withdra_oth_ersu",
        "workers_welfare",
        "distr_profit_shrhder",
        "prfshare_payable_dvd",
        "comshare_payable_dvd",
        "capit_comstock_div",
        "net_after_nr_lp_correct",
        "oth_income",
        "asset_disp_income",
        "continued_net_profit",
        "end_net_profit",
        "credit_impa_loss",
        "net_expo_hedging_benefits",
        "oth_impair_loss_assets",
        "total_opcost",
        "amodcost_fin_assets",
        "update_flag"
    ])
    return df

def getsqlstr(row):
    sqlstr = "insert into t_ods_hs_finance_income(ts_code, ann_date, f_ann_date, end_date, report_type, comp_type, end_type, \
            basic_eps, diluted_eps, total_revenue, revenue, int_income, prem_earned, comm_income, n_commis_income, n_oth_income, \
            n_oth_b_income, prem_income, out_prem, une_prem_reser, reins_income, n_sec_tb_income, n_sec_uw_income, n_asset_mg_income, \
            oth_b_income, fv_value_chg_gain, invest_income, ass_invest_income, forex_gain, total_cogs, oper_cost, int_exp, comm_exp, \
            biz_tax_surchg, sell_exp, admin_exp, fin_exp, assets_impair_loss, prem_refund, compens_payout, reser_insur_liab, div_payt, \
            reins_exp, oper_exp, compens_payout_refu, insur_reser_refu, reins_cost_refund, other_bus_cost, operate_profit, non_oper_income, \
            non_oper_exp, nca_disploss, total_profit, income_tax, n_income, n_income_attr_p, minority_gain, oth_compr_income, t_compr_income, \
            compr_inc_attr_p, compr_inc_attr_m_s, ebit, ebitda, insurance_exp, undist_profit, distable_profit, rd_exp, fin_exp_int_exp, \
            fin_exp_int_inc, transfer_surplus_rese, transfer_housing_imprest, transfer_oth, adj_lossgain, withdra_legal_surplus, \
            withdra_legal_pubfund, withdra_biz_devfund, withdra_rese_fund, withdra_oth_ersu, workers_welfare, distr_profit_shrhder, \
            prfshare_payable_dvd, comshare_payable_dvd, capit_comstock_div, net_after_nr_lp_correct, oth_income, asset_disp_income, \
            continued_net_profit, end_net_profit, credit_impa_loss, net_expo_hedging_benefits, oth_impair_loss_assets, total_opcost, \
            amodcost_fin_assets, update_flag) \
            values('%s', '%s', '%s', '%s', '%s', '%s', '%s',  \
            %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, \
            %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, \
            %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, \
            %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, \
            %f, %f, %f, %f, %f, %f, '%s') \
            on duplicate key update  report_type='%s', comp_type='%s', end_type='%s', \
            basic_eps=%f, diluted_eps=%f, total_revenue=%f, revenue=%f, int_income=%f, prem_earned=%f, \
            comm_income=%f, n_commis_income=%f, n_oth_income=%f, n_oth_b_income=%f, prem_income=%f, \
            out_prem=%f, une_prem_reser=%f, reins_income=%f, n_sec_tb_income=%f, n_sec_uw_income=%f, \
            n_asset_mg_income=%f, oth_b_income=%f, fv_value_chg_gain=%f, invest_income=%f, ass_invest_income=%f, \
            forex_gain=%f, total_cogs=%f, oper_cost=%f, int_exp=%f, comm_exp=%f, biz_tax_surchg=%f, \
            sell_exp=%f, admin_exp=%f, fin_exp=%f, assets_impair_loss=%f, prem_refund=%f, compens_payout=%f, \
            reser_insur_liab=%f, div_payt=%f, reins_exp=%f, oper_exp=%f, compens_payout_refu=%f, \
            insur_reser_refu=%f, reins_cost_refund=%f, other_bus_cost=%f, operate_profit=%f, non_oper_income=%f, \
            non_oper_exp=%f, nca_disploss=%f, total_profit=%f, income_tax=%f, n_income=%f, n_income_attr_p=%f, \
            minority_gain=%f, oth_compr_income=%f, t_compr_income=%f, compr_inc_attr_p=%f, compr_inc_attr_m_s=%f, \
            ebit=%f, ebitda=%f, insurance_exp=%f, undist_profit=%f, distable_profit=%f, rd_exp=%f, fin_exp_int_exp=%f, \
            fin_exp_int_inc=%f, transfer_surplus_rese=%f, transfer_housing_imprest=%f, transfer_oth=%f, \
            adj_lossgain=%f, withdra_legal_surplus=%f, withdra_legal_pubfund=%f, withdra_biz_devfund=%f, \
            withdra_rese_fund=%f, withdra_oth_ersu=%f, workers_welfare=%f, distr_profit_shrhder=%f,  \
            prfshare_payable_dvd=%f, comshare_payable_dvd=%f, capit_comstock_div=%f, net_after_nr_lp_correct=%f, \
            oth_income=%f, asset_disp_income=%f, continued_net_profit=%f, end_net_profit=%f, credit_impa_loss=%f, \
            net_expo_hedging_benefits=%f, oth_impair_loss_assets=%f, total_opcost=%f, amodcost_fin_assets=%f " % (
                getattr(row, 'ts_code'), getattr(row, 'ann_date'), getattr(row, 'f_ann_date'), getattr(row, 'end_date'), 
                getattr(row, 'report_type'), getattr(row, 'comp_type'), getattr(row, 'end_type'), 
                str2float(getattr(row, 'basic_eps')), str2float(getattr(row, 'diluted_eps')), str2float(getattr(row, 'total_revenue')), 
                str2float(getattr(row, 'revenue')), str2float(getattr(row, 'int_income')), str2float(getattr(row, 'prem_earned')), 
                str2float(getattr(row, 'comm_income')), str2float(getattr(row, 'n_commis_income')), str2float(getattr(row, 'n_oth_income')), 
                str2float(getattr(row, 'n_oth_b_income')), str2float(getattr(row, 'prem_income')), str2float(getattr(row, 'out_prem')), 
                str2float(getattr(row, 'une_prem_reser')), str2float(getattr(row, 'reins_income')), str2float(getattr(row, 'n_sec_tb_income')), 
                str2float(getattr(row, 'n_sec_uw_income')), str2float(getattr(row, 'n_asset_mg_income')), str2float(getattr(row, 'oth_b_income')), 
                str2float(getattr(row, 'fv_value_chg_gain')), str2float(getattr(row, 'invest_income')), str2float(getattr(row, 'ass_invest_income')), 
                str2float(getattr(row, 'forex_gain')), str2float(getattr(row, 'total_cogs')), str2float(getattr(row, 'oper_cost')), 
                str2float(getattr(row, 'int_exp')), str2float(getattr(row, 'comm_exp')), str2float(getattr(row, 'biz_tax_surchg')), 
                str2float(getattr(row, 'sell_exp')), str2float(getattr(row, 'admin_exp')), str2float(getattr(row, 'fin_exp')), 
                str2float(getattr(row, 'assets_impair_loss')), str2float(getattr(row, 'prem_refund')), str2float(getattr(row, 'compens_payout')), 
                str2float(getattr(row, 'reser_insur_liab')), str2float(getattr(row, 'div_payt')), str2float(getattr(row, 'reins_exp')), 
                str2float(getattr(row, 'oper_exp')), str2float(getattr(row, 'compens_payout_refu')), str2float(getattr(row, 'insur_reser_refu')), 
                str2float(getattr(row, 'reins_cost_refund')), str2float(getattr(row, 'other_bus_cost')), str2float(getattr(row, 'operate_profit')), 
                str2float(getattr(row, 'non_oper_income')), str2float(getattr(row, 'non_oper_exp')), str2float(getattr(row, 'nca_disploss')), 
                str2float(getattr(row, 'total_profit')), str2float(getattr(row, 'income_tax')), str2float(getattr(row, 'n_income')), 
                str2float(getattr(row, 'n_income_attr_p')), str2float(getattr(row, 'minority_gain')), str2float(getattr(row, 'oth_compr_income')), 
                str2float(getattr(row, 't_compr_income')), str2float(getattr(row, 'compr_inc_attr_p')), str2float(getattr(row, 'compr_inc_attr_m_s')), 
                str2float(getattr(row, 'ebit')), str2float(getattr(row, 'ebitda')), str2float(getattr(row, 'insurance_exp')), 
                str2float(getattr(row, 'undist_profit')), str2float(getattr(row, 'distable_profit')), str2float(getattr(row, 'rd_exp')), 
                str2float(getattr(row, 'fin_exp_int_exp')), str2float(getattr(row, 'fin_exp_int_inc')), str2float(getattr(row, 'transfer_surplus_rese')), 
                str2float(getattr(row, 'transfer_housing_imprest')), str2float(getattr(row, 'transfer_oth')), str2float(getattr(row, 'adj_lossgain')), 
                str2float(getattr(row, 'withdra_legal_surplus')), str2float(getattr(row, 'withdra_legal_pubfund')), str2float(getattr(row, 'withdra_biz_devfund')), 
                str2float(getattr(row, 'withdra_rese_fund')), str2float(getattr(row, 'withdra_oth_ersu')), str2float(getattr(row, 'workers_welfare')), 
                str2float(getattr(row, 'distr_profit_shrhder')), str2float(getattr(row, 'prfshare_payable_dvd')), str2float(getattr(row, 'comshare_payable_dvd')), 
                str2float(getattr(row, 'capit_comstock_div')), str2float(getattr(row, 'net_after_nr_lp_correct')), str2float(getattr(row, 'oth_income')), 
                str2float(getattr(row, 'asset_disp_income')), str2float(getattr(row, 'continued_net_profit')), str2float(getattr(row, 'end_net_profit')), 
                str2float(getattr(row, 'credit_impa_loss')), str2float(getattr(row, 'net_expo_hedging_benefits')), str2float(getattr(row, 'oth_impair_loss_assets')),
                str2float(getattr(row, 'total_opcost')), str2float(getattr(row, 'amodcost_fin_assets')), getattr(row, 'update_flag'), 
                 
                getattr(row, 'report_type'), getattr(row, 'comp_type'), getattr(row, 'end_type'), 
                str2float(getattr(row, 'basic_eps')), str2float(getattr(row, 'diluted_eps')), str2float(getattr(row, 'total_revenue')), 
                str2float(getattr(row, 'revenue')), str2float(getattr(row, 'int_income')), str2float(getattr(row, 'prem_earned')), 
                str2float(getattr(row, 'comm_income')), str2float(getattr(row, 'n_commis_income')), str2float(getattr(row, 'n_oth_income')), 
                str2float(getattr(row, 'n_oth_b_income')), str2float(getattr(row, 'prem_income')), str2float(getattr(row, 'out_prem')), 
                str2float(getattr(row, 'une_prem_reser')), str2float(getattr(row, 'reins_income')), str2float(getattr(row, 'n_sec_tb_income')), 
                str2float(getattr(row, 'n_sec_uw_income')), str2float(getattr(row, 'n_asset_mg_income')), str2float(getattr(row, 'oth_b_income')), 
                str2float(getattr(row, 'fv_value_chg_gain')), str2float(getattr(row, 'invest_income')), str2float(getattr(row, 'ass_invest_income')), 
                str2float(getattr(row, 'forex_gain')), str2float(getattr(row, 'total_cogs')), str2float(getattr(row, 'oper_cost')), 
                str2float(getattr(row, 'int_exp')), str2float(getattr(row, 'comm_exp')), str2float(getattr(row, 'biz_tax_surchg')), 
                str2float(getattr(row, 'sell_exp')), str2float(getattr(row, 'admin_exp')), str2float(getattr(row, 'fin_exp')), 
                str2float(getattr(row, 'assets_impair_loss')), str2float(getattr(row, 'prem_refund')), str2float(getattr(row, 'compens_payout')), 
                str2float(getattr(row, 'reser_insur_liab')), str2float(getattr(row, 'div_payt')), str2float(getattr(row, 'reins_exp')), 
                str2float(getattr(row, 'oper_exp')), str2float(getattr(row, 'compens_payout_refu')), str2float(getattr(row, 'insur_reser_refu')), 
                str2float(getattr(row, 'reins_cost_refund')), str2float(getattr(row, 'other_bus_cost')), str2float(getattr(row, 'operate_profit')), 
                str2float(getattr(row, 'non_oper_income')), str2float(getattr(row, 'non_oper_exp')), str2float(getattr(row, 'nca_disploss')), 
                str2float(getattr(row, 'total_profit')), str2float(getattr(row, 'income_tax')), str2float(getattr(row, 'n_income')), 
                str2float(getattr(row, 'n_income_attr_p')), str2float(getattr(row, 'minority_gain')), str2float(getattr(row, 'oth_compr_income')), 
                str2float(getattr(row, 't_compr_income')), str2float(getattr(row, 'compr_inc_attr_p')), str2float(getattr(row, 'compr_inc_attr_m_s')), 
                str2float(getattr(row, 'ebit')), str2float(getattr(row, 'ebitda')), str2float(getattr(row, 'insurance_exp')), 
                str2float(getattr(row, 'undist_profit')), str2float(getattr(row, 'distable_profit')), str2float(getattr(row, 'rd_exp')), 
                str2float(getattr(row, 'fin_exp_int_exp')), str2float(getattr(row, 'fin_exp_int_inc')), str2float(getattr(row, 'transfer_surplus_rese')), 
                str2float(getattr(row, 'transfer_housing_imprest')), str2float(getattr(row, 'transfer_oth')), str2float(getattr(row, 'adj_lossgain')), 
                str2float(getattr(row, 'withdra_legal_surplus')), str2float(getattr(row, 'withdra_legal_pubfund')), str2float(getattr(row, 'withdra_biz_devfund')), 
                str2float(getattr(row, 'withdra_rese_fund')), str2float(getattr(row, 'withdra_oth_ersu')), str2float(getattr(row, 'workers_welfare')), 
                str2float(getattr(row, 'distr_profit_shrhder')), str2float(getattr(row, 'prfshare_payable_dvd')), str2float(getattr(row, 'comshare_payable_dvd')), 
                str2float(getattr(row, 'capit_comstock_div')), str2float(getattr(row, 'net_after_nr_lp_correct')), str2float(getattr(row, 'oth_income')), 
                str2float(getattr(row, 'asset_disp_income')), str2float(getattr(row, 'continued_net_profit')), str2float(getattr(row, 'end_net_profit')), 
                str2float(getattr(row, 'credit_impa_loss')), str2float(getattr(row, 'net_expo_hedging_benefits')), str2float(getattr(row, 'oth_impair_loss_assets')),
                str2float(getattr(row, 'total_opcost')), str2float(getattr(row, 'amodcost_fin_assets')) 
            )
    return sqlstr

def process():
    stock_list = []
    cursor = db.cursor()
    total_stocks = 0
    total_rows = 0
    fd = open("./income.log", 'w+')
    # get stock list
    sqlstr = "select ts_code from t_ods_hs_base_list where list_status = 'L' "
    try:
        cursor.execute(sqlstr)
        results = cursor.fetchall()
        for row in results :
            stock_list.append(row[0])
            total_stocks = total_stocks+1
    except Exception as ex:
        print (ex)
        print ("get stock list error!")
        db.close()
        return
    # get income data
    for st_code in stock_list :
        df = queryts(st_code)
        loop = 0
        for row in df.itertuples():
            sql = getsqlstr(row)
            try:
                cursor.execute(sql)
                db.commit()
                loop = loop+1
            except Exception as ex:
                print (ex)
                print ("get stock %s error!" % st_code )
                db.rollback()
                return
        print("get %s income data succ! total rows: %d" % (st_code, loop))
        fd.write("get %s income data succ! total rows: %d \n" % (st_code, loop))
        fd.flush()
        total_rows = total_rows+loop
        time.sleep(0.5)
    # get all data  
    print("All stock income data process succ! total stocks: %d, total rows: %d " % (total_stocks, total_rows))  
    fd.write("All stock income data process succ! total stocks: %d, total rows: %d \n" % (total_stocks, total_rows))   

    # close db and file
    fd.close()
    db.close()


process()