"""
Day 2：质量因子单因子选股策略
100天经典策略学习计划

测试因子：ROE、ROA、毛利率、资产周转率
方法：选因子值最高的N只股票（高质量），等权持有
股票池：沪深300
调仓频率：每月

使用方法：修改 FACTOR_NAME 切换测试不同因子
回测建议：2020-01-01 至 2026-02-01，初始资金100万
"""


# ========== 切换因子 ==========
# 可选: 'roe', 'roa', 'gross_profit_margin', 'asset_turnover'
FACTOR_NAME = 'roe'
# ==============================


def initialize(context):
    set_params()
    set_backtest()
    run_monthly(rebalance, monthday=1, time='09:31')


def set_params():
    g.stock_pool = '000300.XSHG'  # 沪深300
    g.stock_num = 20              # 持仓数量
    g.factor = FACTOR_NAME


def set_backtest():
    set_benchmark('000300.XSHG')
    set_option('use_real_price', True)
    set_slippage(FixedSlippage(0.02))
    set_commission(PerTrade(buy_cost=0.0003, sell_cost=0.0013, min_cost=5))
    log.set_level('order', 'error')


def rebalance(context):
    # 获取股票池
    stocks = get_index_stocks(g.stock_pool)
    stocks = filter_stocks(stocks)

    # 查询因子数据
    q = build_query(stocks)
    df = get_fundamentals(q, date=context.current_dt)

    if df.empty:
        return

    # 资产周转率需要手动计算
    if g.factor == 'asset_turnover':
        df = df.dropna(subset=['total_operating_revenue', 'total_assets'])
        df = df[df['total_assets'] > 0]
        df['asset_turnover'] = df['total_operating_revenue'] / df['total_assets']
        df = df[df['asset_turnover'] > 0]

    # 按因子排序，选最高的N只（质量因子越高越好）
    df = df.sort_values(g.factor, ascending=False)
    target = df['code'].head(g.stock_num).tolist()

    log.info(f'[{g.factor}] 选出{len(target)}只，前3: {target[:3]}')

    # 调仓
    adjust_portfolio(context, target)


def build_query(stocks):
    """根据因子构建查询"""
    if g.factor == 'roe':
        return query(
            valuation.code,
            indicator.roe
        ).filter(
            valuation.code.in_(stocks),
            indicator.roe > 0,     # 排除亏损
            indicator.roe < 50     # 排除极端值
        )
    elif g.factor == 'roa':
        return query(
            valuation.code,
            indicator.roa
        ).filter(
            valuation.code.in_(stocks),
            indicator.roa > 0
        )
    elif g.factor == 'gross_profit_margin':
        return query(
            valuation.code,
            indicator.gross_profit_margin
        ).filter(
            valuation.code.in_(stocks),
            indicator.gross_profit_margin > 0
        )
    elif g.factor == 'asset_turnover':
        # 聚宽无直接字段，从营收和总资产手动计算
        return query(
            valuation.code,
            income.total_operating_revenue,
            balance.total_assets
        ).filter(
            valuation.code.in_(stocks),
            income.total_operating_revenue > 0,
            balance.total_assets > 0
        )


def filter_stocks(stocks):
    """过滤ST和停牌"""
    current_data = get_current_data()
    return [s for s in stocks
            if not current_data[s].is_st
            and not current_data[s].paused]


def adjust_portfolio(context, target):
    """调仓：先卖后买"""
    # 卖出不在目标中的
    for stock in list(context.portfolio.positions):
        if stock not in target:
            order_target(stock, 0)

    # 等权买入
    if target:
        per_value = context.portfolio.total_value * 0.95 / len(target)
        for stock in target:
            order_target_value(stock, per_value)