## 加载包
import dai
import pandas as pd
import numpy as np
import math
import warnings
from datetime import datetime, timedelta

from bigmodule import M
from bigtrader.finance.commission import PerOrder

## 设置开始和结束时间
sd = '2020-01-01'
ed = datetime.now().date().strftime("%Y-%m-%d")

## 提因子用的SQL
alpha_sql = f"""
            -- 以下代码为SQL代码，模块化运行时需要在“输入特征（DAI SQL）”中运行，代码数据提取需要以“dai.query(" SQL代码 ")”的形式运行
            SELECT 
                date,
                instrument, 
                -1 *   m_sum(turn, 250) AS factor 
            FROM cn_stock_bar1d  
            ORDER BY date, instrument 
"""

sql = f"""
    WITH
    data_alpha AS (
        {alpha_sql}
    ),
    data_alpha_origin AS (
        SELECT *
        FROM data_alpha
        QUALIFY COLUMNS(*) IS NOT NULL AND factor != 'Infinity' AND factor != '-Infinity'
    ),
    data_alpha_process AS (
        SELECT 
            date,
            instrument,
            factor,
            clip(factor, c_avg(factor) - 3 * c_std(factor), c_avg(factor) + 3 * c_std(factor)) AS clipped_factor,
            c_normalize(clipped_factor) AS normalized_factor,
            c_neutralize(normalized_factor, sw2021_level1, LOG(total_market_cap)) AS neutralized_factor,
        FROM data_alpha_origin JOIN cn_stock_factors_base USING (date, instrument)
        WHERE 1=1
        AND amount > 0
        AND st_status = 0
        AND trading_days > 252
        AND (instrument LIKE '%SH' OR instrument LIKE '%SZ')
        QUALIFY COLUMNS(*) IS NOT NULL
        ORDER BY date, instrument
    )
    SELECT 
        date, 
        instrument, 
        neutralized_factor AS factor 
    FROM data_alpha_process 
    ORDER BY date, factor DESC

"""

# 数据提取
data = dai.query(sql, filters={'date':[sd, ed]}).df()
data

def m_initialize_bigquant_run(context):
    context.set_commission(PerOrder(buy_cost=0.0003, sell_cost=0.0013, min_cost=5))
    context.holding_days = 1
    context.target_hold_count = 10
    context.target_percent_per_instrument = 1.0 / context.target_hold_count

def m_before_trading_start_bigquant_run(context, data):
    pass

def m_handle_tick_bigquant_run(context, tick):
    pass

def m_handle_data_bigquant_run(context, data):
    if context.trading_day_index % context.holding_days != 0:
        return
    current_date = data.current_dt.strftime("%Y-%m-%d")
    current_day_data = context.data[context.data["date"] == current_date]
    current_day_data = current_day_data.head(context.target_hold_count)
    target_hold_instruments = set(current_day_data["instrument"])
    current_hold_instruments = set(context.get_account_positions().keys())
    sell_set = current_hold_instruments - target_hold_instruments
    buy_set  = target_hold_instruments  - current_hold_instruments
    for instrument in sell_set:
        context.order_target_percent(instrument, 0)
    for instrument in buy_set:
        context.order_target_percent(instrument, context.target_percent_per_instrument)

def m_handle_trade_bigquant_run(context, trade):
    pass

def m_handle_order_bigquant_run(context, order):
    pass

def m_after_trading_bigquant_run(context, data):
    pass

m = M.bigtrader.v14(
    data=data,
    start_date='',
    end_date='',
    initialize=m_initialize_bigquant_run,
    before_trading_start=m_before_trading_start_bigquant_run,
    handle_tick=m_handle_tick_bigquant_run,
    handle_data=m_handle_data_bigquant_run,
    handle_trade=m_handle_trade_bigquant_run,
    handle_order=m_handle_order_bigquant_run,
    after_trading=m_after_trading_bigquant_run,
    capital_base=500000,
    frequency='daily',
    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,
    debug=False,
    backtest_only=False,
    m_cached=False
)