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

测试因子：20日价格动量、60日价格动量、120日价格动量
方法：选动量最强（涨幅最大）的N只股票，等权持有
股票池：沪深300
调仓频率：每月

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


# ========== 切换因子 ==========
# 可选: 'price_mom_20d', 'price_mom_60d', 'price_mom_120d'
FACTOR_NAME = 'price_mom_20d'
# ==============================


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)

    df = get_factor_data(stocks, context)

    if df is None or df.empty:
        return

    # 动量因子：越高越好（买赢家）
    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]}')
    if not df.empty:
        log.info(f'动量范围: [{df[g.factor].min():.2%}, {df[g.factor].max():.2%}]')

    adjust_portfolio(context, target)


def get_factor_data(stocks, context):
    """获取动量因子数据"""
    import pandas as pd

    # 确定回望周期
    if g.factor == 'price_mom_20d':
        lookback = 20
    elif g.factor == 'price_mom_60d':
        lookback = 60
    elif g.factor == 'price_mom_120d':
        lookback = 120
    else:
        return None

    # 获取历史收盘价（多取一些天数确保足够交易日）
    prices = get_price(
        stocks,
        end_date=context.current_dt,
        count=lookback + 1,
        frequency='daily',
        fields=['close'],
        panel=False
    )

    if prices is None or prices.empty:
        return None

    # pivot: 行=日期, 列=股票代码, 值=收盘价
    pivot = prices.pivot(index='time', columns='code', values='close')

    # 计算区间收益率 = (最新价 / 起始价) - 1
    mom = (pivot.iloc[-1] / pivot.iloc[0] - 1)

    df = pd.DataFrame({
        'code': mom.index,
        g.factor: mom.values
    })
    df = df.dropna()

    # 过滤极端值（排除涨幅超200%或跌幅超50%的异常股票）
    df = df[df[g.factor].between(-0.5, 2.0)]

    return df


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)