"""
Day 7：流动性因子单因子选股策略
100天经典策略学习计划

测试因子：换手率、Amihud非流动性、近似价差
方法：选低换手率/高Amihud/高价差的N只（流动性差=溢价）
股票池：沪深300
调仓频率：每月

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


# ========== 切换因子 ==========
# 可选: 'turnover', 'amihud', 'spread'
FACTOR_NAME = 'turnover'
# ==============================


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


def set_params():
    g.stock_pool = '000300.XSHG'
    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

    # 换手率：越低越好（流动性差=溢价）
    # Amihud/价差：越高越好（流动性差=溢价）
    if g.factor == 'turnover':
        df = df.sort_values(g.factor, ascending=True)
    else:
        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 get_factor_data(stocks, context):
    """获取流动性因子数据"""
    import pandas as pd
    import numpy as np

    if g.factor == 'turnover':
        # 换手率：过去20天平均换手率
        prices = get_price(
            stocks,
            end_date=context.current_dt,
            count=20,
            frequency='daily',
            fields=['volume', 'money'],
            panel=False
        )

        if prices is None or prices.empty:
            return None

        # 计算日均换手率（成交额/流通市值的近似）
        # 聚宽的volume是成交量，money是成交额
        turnover = prices.groupby('code')['money'].mean()

        # 获取流通市值
        q = query(
            valuation.code,
            valuation.circulating_market_cap
        ).filter(
            valuation.code.in_(stocks)
        )
        cap_df = get_fundamentals(q, date=context.current_dt)

        if cap_df.empty:
            return None

        cap_df = cap_df.set_index('code')

        # 计算换手率 = 日均成交额 / 流通市值（亿元）
        df = pd.DataFrame({
            'code': turnover.index,
            'turnover': turnover.values / (cap_df.loc[turnover.index, 'circulating_market_cap'].values * 1e8)
        })
        df = df.dropna()
        df = df[df['turnover'] > 0]

        return df

    elif g.factor == 'amihud':
        # Amihud非流动性 = avg(|return| / volume)
        prices = get_price(
            stocks,
            end_date=context.current_dt,
            count=21,  # 需要21天计算20天收益率
            frequency='daily',
            fields=['close', 'money'],
            panel=False
        )

        if prices is None or prices.empty:
            return None

        pivot_close = prices.pivot(index='time', columns='code', values='close')
        pivot_money = prices.pivot(index='time', columns='code', values='money')

        # 计算日收益率
        returns = pivot_close.pct_change().iloc[1:]
        money = pivot_money.iloc[1:]

        # Amihud = mean(|return| / money)
        # money单位是元，转成亿元
        amihud = (returns.abs() / (money / 1e8)).mean()

        df = pd.DataFrame({
            'code': amihud.index,
            'amihud': amihud.values
        })
        df = df.dropna()
        df = df[df['amihud'] > 0]

        # 过滤极端值
        df = df[df['amihud'] < df['amihud'].quantile(0.95)]

        return df

    elif g.factor == 'spread':
        # 近似价差 = (High - Low) / Close
        prices = get_price(
            stocks,
            end_date=context.current_dt,
            count=20,
            frequency='daily',
            fields=['high', 'low', 'close'],
            panel=False
        )

        if prices is None or prices.empty:
            return None

        # 计算日内价差
        prices['daily_spread'] = (prices['high'] - prices['low']) / prices['close']

        # 取20天平均
        spread = prices.groupby('code')['daily_spread'].mean()

        df = pd.DataFrame({
            'code': spread.index,
            'spread': spread.values
        })
        df = df.dropna()
        df = df[df['spread'] > 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)