import pandas as pd
import numpy as np
from biglearning.module2.common.data import Outputs
from zipline.finance.commission import PerOrder

from bigdatasource.api import DataSource
from biglearning.api import M
from biglearning.api import tools as T
from biglearning.module2.common.data import Outputs

import warnings
warnings.filterwarnings('ignore')
import dai

sd = '2015-01-01'
ed = '2023-09-01'

datasource_id = "hf_alpha_fzzq"


df = dai.query(f"SELECT * FROM {datasource_id}", filters={"date": [sd,ed]}).df()

df = df[['date','instrument','alpha_91007']]

df

df['ins'] = df['instrument']
df

test = dai.query(f"SELECT * FROM cn_stock_index_info ", filters={"date": [sd,ed]}).df()
test

test[test.name.str.contains('中证500')]

df_w = dai.query("""
    SELECT *
    FROM cn_stock_index_weight
    WHERE date > '2015-01-01'
    AND instrument = '000905.SH';
""").df()

df_w = df_w.rename(columns={
    'member_code': 'ins',
    'instrument':'index_code'
    })
df_w

df

df = pd.merge(df_w,df,left_on=['date','ins'],right_on=['date','ins'],how='left')
df.dropna(subset=['instrument'],inplace=True)
df = df[df.date >= '2016-01-01']
df

df['instrument'] = df['instrument'].apply(lambda x: x+'A')
df

df.sort_values(by='date',inplace=True)
sd_ = np.str(df.date.unique()[0])
ed_ = np.str(df.date.unique()[-1])


m1 = M.instruments.v2(
    start_date=sd_,
    end_date=ed_,
    market='CN_STOCK_A',
    instrument_list='',
    max_count=0
)


df = DataSource.write_df(df)

# 交易引擎：初始化函数，只执行一次
def m4_initialize_bigquant_run(context):

    context.ranker_prediction = context.options['data'].read_df()
    context.ranker_prediction.set_index('date',inplace=True)
    


# 交易引擎：每个单位时间开盘前调用一次。
def m4_before_trading_start_bigquant_run(context, data):
    # 盘前处理，订阅行情等
    #context.subscribe_bar(context.instruments,'1m')
    pass

# 交易引擎：tick数据处理函数，每个tick执行一次
def m4_handle_tick_bigquant_run(context, tick):
    pass

def handle_bar(context, bar):
    pass


# 交易引擎：bar数据处理函数，每个时间单位执行一次
def m4_handle_data_bigquant_run(context,data):


    remainder = context.trading_day_index % 2
    #如果没到调仓期直接结束运行
    if remainder !=0:
        return



    import datetime
    #初始化
    buy_list = []  #买入列表
    sell_list = [] #卖出列表
    
    #==================== 数据准备
    today = data.current_dt.strftime('%Y-%m-%d')
    time = data.current_dt

    account_pos = context.get_account_positions()
    holding_list = list({key: value for key, value in account_pos.items() if value.avail_qty > 0}.keys())
    holding_num = len(holding_list)

    #读取当日数据
    try:
        today_data  = context.ranker_prediction.loc[today,:]
        today_data.reset_index(inplace=True)
    except:
        return

    today_data.dropna(inplace=True)
    today_data.sort_values(by='alpha_91007',ascending = True , inplace=True)



    
    #构建买入列表
    target_list = today_data.instrument.to_list()[:50] #选股
    weight_list = today_data.weight.to_list()[:50] #权重列表
    weight_sum = today_data.iloc[:50].weight.sum()
    mul = 1/weight_sum  #权重倍数


    #构建卖出列表
    for ins in holding_list:
        if ins not in target_list:
            sell_list.append(ins)

    #构建买入列表
    for ins in target_list:
        if ins not in holding_list:
            buy_list.append(ins)
            
    for ins in sell_list:
        context.order_target(ins,0)


    for ins in buy_list:
        idx = target_list.index(ins)
        weight_ = weight_list[idx] #按权重 * 倍数 买入
        weight_ = weight_ * mul
        if weight_ >= 0.04:
            weight_ = 0.04
        context.order_target_percent(ins, weight_)


    

# 交易引擎：成交回报处理函数，每个成交发生时执行一次
def m4_handle_trade_bigquant_run(context, trade):
    pass

# 交易引擎：委托回报处理函数，每个委托变化时执行一次
def m4_handle_order_bigquant_run(context, order):
    pass

# 交易引擎：盘后处理函数，每日盘后执行一次
def m4_after_trading_bigquant_run(context, data):
    pass

trade = M.hftrade.v2(
    instruments=m1.data,
    options_data=df,
    start_date='',
    end_date='',
    initialize=m4_initialize_bigquant_run,
    before_trading_start=m4_before_trading_start_bigquant_run,
    handle_tick=m4_handle_tick_bigquant_run,
    handle_data=m4_handle_data_bigquant_run,
    handle_trade=m4_handle_trade_bigquant_run,
    handle_order=m4_handle_order_bigquant_run,
    after_trading=m4_after_trading_bigquant_run,
    capital_base=100000000,
    frequency='daily',
    price_type='真实价格',
    product_type='股票',
    before_start_days='0',
    volume_limit=0,
    order_price_field_buy='close',
    order_price_field_sell='open',
    benchmark='000852.HIX',
    plot_charts=True,
    disable_cache=True,
    replay_bdb=False,
    show_debug_info=False,
    backtest_only=False
)