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

# 本代码由可视化策略环境自动生成 2023年5月6日 09:27
# 本代码单元只能在可视化模式下编辑。您也可以拷贝代码，粘贴到新建的代码单元或者策略，然后修改。
 
# 显式导入 BigQuant 相关 SDK 模块
from bigdatasource.api import DataSource
from biglearning.api import M
from biglearning.api import tools as T
from biglearning.module2.common.data import Outputs

# Python 代码入口函数，input_1/2/3 对应三个输入端，data_1/2/3 对应三个输出端
def m3_run_bigquant_run(input_1, input_2, input_3):
    # 示例代码如下。在这里编写您的代码
    df = input_1.read()
    
    from joblib import Parallel, delayed
    import copy
    lst = df.instrument.unique().tolist()
    
    def calc_data(df,instrument):
        
        import warnings
        warnings.filterwarnings('ignore')
        df_ = df.copy()

        df_ = df_[df_['instrument'] == instrument]

        df_['year'] = df_.date.dt.year
        df_['month'] = df_.date.dt.month
        df_['day'] = df_.date.dt.day
        df_['hour'] = df_.date.dt.hour

        df_['pct_change_h'] = df_['close']/df_['open']

        def calc_open(df):
            idx_list = df.index.tolist()
            idx = idx_list[0]
            idx_4 = idx_list[3]
            open_ = df.loc[idx,'open']
            close_ = df.loc[idx_4,'close']
            df['today_open'] = open_
            df['today_close'] = close_
            return df
        df_ = df_.groupby(['year','month','day']).apply(calc_open)
        
        def calc_pre_close(df):
            df.sort_values(by='date',inplace=True)
            df['pre_close'] = df['today_close'].shift(4)
            return df
        
        df_ = df_.groupby('instrument').apply(calc_pre_close)

        df_['up'] = df_['close']/df_['pre_close']

        return df_

    
    
    
    results = Parallel(n_jobs=-1)(delayed(calc_data)(df,ins) for ins in lst)
    
    df = pd.concat(results)

    data_1 = DataSource.write_df(df)

    return Outputs(data_1=data_1, data_2=None, data_3=None)

# 后处理函数，可选。输入是主函数的输出，可以在这里对数据做处理，或者返回更友好的outputs数据格式。此函数输出不会被缓存。
def m3_post_run_bigquant_run(outputs):
    return outputs

# 交易引擎：初始化函数，只执行一次
def m4_initialize_bigquant_run(context):
    #=========================== 加载预测数据  ===================================
    
    context.ranker_prediction = context.options['data'].read_df()
    context.ranker_prediction.set_index('date',inplace=True)
    
    #===========================  设置交易参数  ====================================
    # 系统已经设置了默认的交易手续费和滑点，要修改手续费可使用如下函数
    context.set_commission(PerOrder(buy_cost=0.00015, sell_cost=0.0013, min_cost=5))
    #===========================  设置持仓参数  ====================================
    #最大持仓量
    context.stock_count = 5
    # 设置每只股票占用的最大资金比例
    context.max_cash_per_instrument = 0.5
    
    #持仓天数
    context.hold_days = 5
    
    #股票止损列表
    context.stop_list = {}
    
    #股票止损卖出后最小买入间隔
    context.min_stop_days = 20
    
    #=======================================功能开关
    #是否止损
    context.stop_win = True

    context.set_stock_t1(1)

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

# 交易引擎：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):
    import datetime
    #初始化
    buy_list = []  #买入列表
    sell_list = [] #卖出列表
    target_list = []  #目标列表
    #==================== 数据准备
    today = data.current_dt.strftime('%Y-%m-%d')
    time = data.current_dt
    
    holding_list = list(context.get_account_positions().keys())  #目前持仓列表
    #能卖则立即平仓
    for ins in holding_list:
        context.order_target(ins,0)  
    
    try:
        today_data  = context.ranker_prediction.loc[time,:]
        today_data.reset_index(inplace=True)
    except:
        return

    if len(today_data)>0:
        #=========================================
        # 事件筛选        
        today_data = today_data[today_data['up'] >= 1.085]
        today_data.sort_values(by='pct_change_h',ascending=False,inplace=True)
        #==========================================
 
        #开仓
        target_list = today_data.instrument.tolist()
        len_ = len(target_list)
        if len_ > 1:
            buy_num = 1
        else:
            buy_num = 0
        if len_>0:
            target_list = target_list[:buy_num]
            for ins in target_list:
                context.order_percent(ins,0.25)

    

# 交易引擎：成交回报处理函数，每个成交发生时执行一次
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

#读取数据模块
m1 = M.use_datasource.v2(
    datasource_id='bar60m_CN_STOCK_A',
    start_date='2023-01-01',
    end_date='2023-05-05',
    before_start_days=30
)

#自定义python模块
m3 = M.cached.v3(
    input_1=m1.data,
    run=m3_run_bigquant_run,
    post_run=m3_post_run_bigquant_run,
    input_ports='',
    params='{}',
    output_ports=''
)

m2 = M.instruments.v2(
    start_date='2023-01-01',
    end_date='2023-05-05',
    market='CN_STOCK_A',
    instrument_list='',
    max_count=0
)

m4 = M.hftrade.v2(
    instruments=m2.data,
    options_data=m3.data_1,
    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=2000000.001,
    frequency='minute',
    price_type='真实价格',
    product_type='股票',
    before_start_days='0',
    volume_limit=0,
    order_price_field_buy='open',
    order_price_field_sell='close',
    benchmark='000300.HIX',
    plot_charts=True,
    disable_cache=False,
    replay_bdb=False,
    show_debug_info=False,
    backtest_only=False
)