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

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

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

## 设置开始和结束时间
sd = T.live_run_param('trading_date', '2023-01-01')
ed = T.live_run_param('trading_date', '2024-02-23')

## 提因子用的SQL
sql = f"""
WITH 
data_origin AS (
    SELECT 
        date, 
        instrument,
        float_market_cap
    FROM cn_stock_factors_base

    WHERE date BETWEEN DATE '{sd}' - INTERVAL 0 DAY AND '{ed}'
    AND st_status = 0
    AND list_days > 365
    AND suspended = 0
    AND list_sector = 1

    QUALIFY COLUMNS(*) IS NOT NULL

    ORDER BY date, float_market_cap
),
data_date AS (
    SELECT
        date,
        LAG(date, 1) OVER (ORDER BY date) AS _last_trading_date,
        date - INTERVAL 1 DAY AS _last_date,
        LAG(IF(_last_date = _last_trading_date, 0, 1), -1) OVER (ORDER BY date) AS trade_day
    FROM all_trading_days
    WHERE date BETWEEN DATE '{sd}' - INTERVAL 0 DAY AND '{ed}'
    AND market_code = 'CN'
    QUALIFY COLUMNS(*) IS NOT NULL
)
SELECT * 
FROM data_date JOIN data_origin USING (date)
ORDER BY date, float_market_cap

"""

data = dai.query(sql).df()
data

def m_initialize_bigquant_run(context):
    context.set_commission(PerOrder(buy_cost=0.0001, sell_cost=0.0013, min_cost=5))
    context.holding_days = 5
    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):
    
    current_date = data.current_dt.strftime("%Y-%m-%d")
    current_day_data = context.data[context.data["date"] == current_date]

    if current_day_data['trade_day'].iloc[0] != 1:
        return
    
    current_day_data = current_day_data.iloc[: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.v9(
    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='open',
    benchmark='000300.SH',
    plot_charts=True,
    disable_cache=True,
    debug=False,
    backtest_only=False
)