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

import pandas as pd
import numpy as np
import math
import dai
import warnings
import datetime
from datetime import timedelta

from zipline.finance.commission import PerOrder
from zipline.api import get_open_orders
from zipline.api import symbol

from bigtrader.sdk import *
from bigtrader.utils.my_collections import NumPyDeque
from bigtrader.constant import OrderType
from bigtrader.constant import Direction

sql = """
SELECT * FROM minute_alpha
INNER JOIN weituo_alpha USING (date, instrument)
INNER JOIN (
    SELECT date, instrument, m_lead(close, 1) / close - 1 AS label
    FROM cn_stock_bar1d
)
USING (date, instrument)
QUALIFY COLUMNS(*) IS NOT NULL
ORDER BY date
"""

dai.query(sql, filters={'date': ['2023-01-01', '2023-10-10']}).df()

def get_data(sd, ed, training=True):
    if training:
        # 训练集是需要标签的
        sql = """
        SELECT * FROM minute_alpha
        INNER JOIN weituo_alpha USING (date, instrument)
        INNER JOIN (
            SELECT date, instrument, m_lead(close, 1) / close - 1 AS label
            FROM cn_stock_bar1d
        )
        USING (date, instrument)

        INNER JOIN (
            -- 必须要保证是在中证1000当中的股票池
            SELECT date, instrument AS _index_code, member_code AS instrument
            FROM cn_stock_index_component
            WHERE _index_code = '000852.SH'
        )
        USING (date, instrument)

        QUALIFY COLUMNS(*) IS NOT NULL
        ORDER BY date
        """
        df = dai.query(sql, filters={'date': [sd, ed]}).df()
    
    else:
        # 测试集是不需要标签的
        sql = """
        SELECT * FROM minute_alpha
        INNER JOIN weituo_alpha USING (date, instrument)
        INNER JOIN (
            -- 必须要保证是在中证1000当中的股票池
            SELECT date, instrument AS _index_code, member_code AS instrument
            FROM cn_stock_index_component
            WHERE _index_code = '000852.SH'
        )
        USING (date, instrument)
        QUALIFY COLUMNS(*) IS NOT NULL
        ORDER BY date
        """
        df = dai.query(sql, filters={'date': [sd, ed]}).df()
    return df

def scroling(train_sd, train_ed, test_sd, test_ed):
    # 加载训练集
    train_df = get_data(train_sd, train_ed)
    label = np.array(train_df['label']).reshape(-1, 1)
    feature = np.array(train_df.drop(['date', 'instrument', 'label'], axis=1))
    beta = np.linalg.inv(feature.T@feature)@feature.T@label
    
    # 加载回测集
    back_df = get_data(test_sd, test_ed, training=False)
    feature = np.array(back_df.drop(['date', 'instrument'], axis=1))
    ypre = (feature@beta).reshape(1, -1)[0]
    pred = back_df[['date', 'instrument']]
    pred['pred_label'] = ypre
    return pred

import pandas as pd

# 设置开始日期和结束日期
start_date = '2020-01-01'
end_date = '2020-12-31'

# 生成每个月第一天的日期范围
dates = pd.date_range(start=start_date, end=end_date, freq='MS')  # 'MS' 表示月初（Month Start）
date = []
for d in dates:
    date.append(d.strftime('%Y-%m-%d'))
date = date[:-1]

date

table = []
for i in range(len(date)-2):
    train_sd = date[i]
    train_ed = date[i+1]
    test_sd = date[i+1]
    test_ed = date[i+2]
    print(train_sd, train_ed, test_sd, test_ed)
    try:
        df = scroling(train_sd, train_ed, test_sd, test_ed)
    except:
        continue
    table.append(df)

pre_table = pd.concat(table, axis=0)
pre_table

# 开始回测(回测准备工作)
instruments = {'market': 'CN_STOCK_A', 'instruments': list(pre_table.instrument.unique()), 'start_date': '2020-02-03', 'end_date': '2020-10-30'}
instruments = DataSource.write_pickle(instruments)

df = DataSource.write_df(pre_table)

# 交易引擎：初始化函数，只执行一次
def m4_initialize_bigquant_run(context):
    # 加载预测数据
    context.df = context.options['data'].read_df()

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

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

# 交易引擎：bar数据处理函数，每个时间单位执行一次
def m4_handle_data_bigquant_run(context, data):
    dt = data.current_dt.strftime('%Y-%m-%d')

    # 获取数据
    df = context.df[context.df['date']==dt].sort_values('pred_label', ascending=False)
    instruments = list(df[df['pred_label']>0].instrument)[:10]
    
    # 获取持仓信息
    holding = context.get_account_positions()
    holding_list = list(holding.keys())

    # 卖出不在买入池中的股票
    for ins in holding_list:
        if ins not in instruments and data.can_trade(ins):
            context.order_target(ins, 0)
            holding_list.remove(ins)
    
    # 买入持仓中没有的票
    for ins in instruments:
        if ins not in holding_list and data.can_trade(ins) and len(holding_list)<10:
            context.order_target_percent(ins, 1/10)
            holding_list.append(ins)


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


m4 = M.hftrade.v2(
    instruments=instruments,
    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=1000000,
    frequency='daily',
    price_type='真实价格',
    product_type='股票',
    before_start_days='0',
    volume_limit=1,
    order_price_field_buy='open',
    order_price_field_sell='open',
    benchmark='000300.SH',
    plot_charts=True,
    disable_cache=False,
    replay_bdb=False,
    show_debug_info=False,
    backtest_only=False
)