## 加载包
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 = f"""
SELECT
    date,
    instrument,
    m_lag(IF(m_avg(close,5) > m_avg(close,10) AND m_avg(close,10) > m_avg(close,20) AND m_avg(close,20) > m_avg(close,40) AND m_avg(close,40)>m_avg(close,120) AND low < m_avg(close,10),1,0),1) AS buy_signal,
    m_lag(IF(m_avg(close, 5) < m_avg(close, 40), 1,0),1) AS sell_signal
FROM cn_stock_factors
QUALIFY COLUMNS(*) IS NOT NULL
AND date BETWEEN DATE '{sd}' - INTERVAL 0 DAY AND '{ed}'
ORDER BY date,instrument
"""
data = dai.query(sql,full_db_scan=True).df()
data

# 交易引擎：初始化函数，只执行一次
def m_initialize_bigquant_run(context):
    # 系统已经设置了默认的交易手续费和滑点，要修改手续费可使用如下函数
    context.set_commission(PerOrder(buy_cost=0.0003, sell_cost=0.0013, min_cost=5))
    # 持有1天
    context.holding_days = 1
    # 设置买入股票数量
    context.target_hold_count = 20
    # 每只股票的目标权重
    context.target_percent_per_instrument = 1.0 / context.target_hold_count

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

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

# 交易引擎：K先处理函数，每个K线执行一次
def m_handle_data_bigquant_run(context, data):

    # 每5个交易日调仓一次
    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]
    
    # 获取当日符合买入/卖出条件的股票列表
    try:
        buy_stock = set(current_day_data[current_day_data['buy_signal'] == 1].iloc[:context.target_hold_count]['instrument']) # 当日符合买入条件的股票(20只)
    except:
        buy_stock=[]
    try:
        sell_stock = set(current_day_data[current_day_data['sell_signal']==1]['instrument'])# 当日符合卖出条件的股票
    except:
        sell_stock = []
    
    # 获取当前已持有股票
    current_hold_instruments = set(context.get_account_positions().keys())
    hold_num=len(current_hold_instruments)

    # 需要卖出的股票:已有持仓中符合卖出条件的股票
    sell_set = [i for i in current_hold_instruments if i in sell_stock]
    # 需要买入的股票:没有持仓且符合买入条件的股票
    buy_set = [i for i in buy_stock if i not in current_hold_instruments]

    # 卖出不在目标持有列表中的股票
    for instrument in sell_set:
        context.order_target_percent(instrument, 0)
        hold_num-=1

    # 当日还允许买入建仓的股票数目
    stock_can_buy_num = context.target_hold_count - hold_num
    stock_to_buy_num = min(stock_can_buy_num,len(buy_set))
    
    # 如果当天没有买入的股票，就返回
    if stock_to_buy_num == 0:
        return
    # 记录已经买入的股票数量
    buy_num = 0
    # 买入目标持有列表中的股票
    for instrument in buy_set:
        if buy_num < stock_to_buy_num:
            context.order_target_percent(instrument, context.target_percent_per_instrument)
            buy_num += 1

# 交易引擎：成交回报处理函数，每个成交发生时执行一次
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='open',
    benchmark='000300.SH',
    plot_charts=True,
    disable_cache=False,
    debug=False,
    backtest_only=False,
    m_cached=False
)