from xtquant import xtdata
from datetime import datetime as dt
import pandas as pd
import math
from sqlalchemy import create_engine
import multiprocessing as mp
from apscheduler.schedulers.blocking import BlockingScheduler

# pd.set_option('display.max_rows', None) # 设置显示最大行


path = 'C:\\qmt\\userdata_mini'

field = ['time', 'open', 'close', 'high', 'low', 'volume', 'amount']
cpu_count = mp.cpu_count()


def to_sql(stock_list, eng_back, eng_front):
    print(dt.now(), '开始循环入库！')
    for stock in stock_list:
        print(stock)
        data = xtdata.get_market_data(field, [stock], '1d', end_time='', count=-1, dividend_type='back')
        df = pd.concat([data[i].loc[stock].T for i in ['time', 'open', 'high', 'low', 'close', 'volume', 'amount']],
                       axis=1)
        df.columns = ['time', 'open', 'high', 'low', 'close', 'volume', 'amount']
        df['time'] = df['time'].apply(lambda x: dt.fromtimestamp(x / 1000.0))
        df.reset_index(drop=True, inplace=True)
        print(df)
        df.to_sql('%s_1d' % stock, con=eng_back, index=True, if_exists='append')

    for stock in stock_list:
        print(stock)
        data = xtdata.get_market_data(field, [stock], '1d', end_time='', count=-1, dividend_type='front')
        df = pd.concat([data[i].loc[stock].T for i in ['time', 'open', 'high', 'low', 'close', 'volume', 'amount']],
                       axis=1)
        df.columns = ['time', 'open', 'high', 'low', 'close', 'volume', 'amount']
        df['time'] = df['time'].apply(lambda x: dt.fromtimestamp(x / 1000.0))
        df.reset_index(drop=True, inplace=True)
        print(df)
        df.to_sql('%s_1d' % stock, con=eng_front, index=True, if_exists='append')


def download_data(stock_list, eng_back, eng_front):
    print(dt.now(), '开始下载！')
    xtdata.download_history_data2(stock_list=stock_list, period='1d', start_time='', end_time='')
    print(dt.now(), '下载完成，准备入库！')
    to_sql(stock_list, eng_back, eng_front)


# def to_df(key, values, engine):
#     print('to_df')
#     pass


if __name__ == '__main__':
    stocks = xtdata.get_stock_list_in_sector('沪深A股')
    field = ['time', 'open', 'close', 'high', 'low', 'volume', 'amount']
    cpu_count = mp.cpu_count()
    stocks.sort()
    step = math.ceil(len(stocks) / cpu_count)
    eng_b = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/qmt_stocks?charset=utf8')
    eng_f = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/qmt_stocks_front?charset=utf8')

    download_data(stocks, eng_b, eng_f)

    # scheduler = BlockingScheduler()
    # scheduler.add_job(func=download_data, trigger='cron', hour='15', minute='45', args=[stocks, eng_b, eng_f],
    #                   timezone="Asia/Shanghai")
    # try:
    #     scheduler.start()
    # except (KeyboardInterrupt, SystemExit):
    #     pass