## 加载包
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
sql = f"""

SELECT 
    date, 
    instrument,

FROM cn_stock_factors

WHERE date BETWEEN DATE '{sd}' - INTERVAL 0 DAY AND '{ed}'
AND st_status = 0
AND trading_days > 270
AND list_sector NOT in (2, 4)

QUALIFY COLUMNS(*) IS NOT NULL
AND price_limit_status_1 = 3
AND price_limit_status_2 = 3
AND close > m_avg(close, 5)
AND close < 30
AND (5 < pe_ttm AND pe_ttm < 30)
AND net_profit_yoy_mrq > 0.1

ORDER BY date, volume DESC

"""

# 数据提取
data = dai.query(sql).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 = 3
    context.target_hold_count = 5
    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

    if len(target_hold_instruments) < context.target_hold_count and len(target_hold_instruments) != 0:
        weight = 1 / len(target_hold_instruments)
    else:
        weight = 1 / context.target_hold_count

    for instrument in sell_set:
        context.order_target_percent(instrument, 0)
    for instrument in target_hold_instruments:
        pl = dai.query(f"SELECT date, instrument, price_limit_status FROM cn_stock_factors WHERE date = '{current_date}' AND instrument = '{instrument}'").df()["price_limit_status"].iloc[0]
        if pl != 3:
            context.order_target_percent(instrument, weight)

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=1000000,
    frequency='daily',
    product_type='股票',
    before_start_days=0,
    volume_limit=1,
    order_price_field_buy='open',
    order_price_field_sell='close',
    benchmark='000300.SH',
    plot_charts=True,
    disable_cache=False,
    debug=False,
    backtest_only=False,
    m_cached=False
)