# coding:utf-8
from datetime import datetime as dt
import os
import pandas as pd
from xtquant.xttrader import XtQuantTrader, XtQuantTraderCallback
from xtquant.xttype import StockAccount
from xtquant import xtdata, xtconstant


df = pd.DataFrame(columns=['time', 'open', 'close', 'high', 'low', 'volume', 'amount'])
def avg_price(df):
    print('avg_price')
    daily_avg_price = []
    for i in range(1, len(df) + 1):
        daily_avg_price.append(round(sum(df['amount'][:i]) / sum(df['volume'][:i]) / 10, 2))
    df['daily_avg_price'] = daily_avg_price
    return df


def over_avg(datas):
    global df
    for stock_code in datas:
        df_temp = pd.DataFrame(datas[stock_code])
        df_temp['time'] = df_temp['time'].apply(lambda x: dt.fromtimestamp(x / 1000.0))
        df = pd.concat([df, df_temp], axis=0, ignore_index=True)
        df = avg_price(df)
        df = df[['time', 'open', 'close', 'high', 'low', 'volume', 'amount', 'daily_avg_price']]
        print(stock_code, '\n', df)


if __name__ == '__main__':
    print('start')
    kzz = xtdata.get_stock_list_in_sector('沪深转债')
    print(kzz, len(kzz))
    # exit()
    t = dt.now().strftime('%Y%m%d')
    # stocks = xtdata.get_stock_list_in_sector('沪深A股')
    # xtdata.subscribe_whole_quote(kzz, callback=over_avg)
    xtdata.subscribe_quote('128130.SZ', period='1m', start_time=t, count=-1, callback=over_avg)
    xtdata.run()