"""
Day 8：技术因子单因子选股策略
100天经典策略学习计划

测试因子：MACD、RSI、KDJ、布林带
方法：根据技术指标信号选股，等权持有
股票池：沪深300
调仓频率：每月

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

注意：技术指标通常适合高频交易，月度调仓可能不匹配
"""


# ========== 切换因子 ==========
# 可选: 'macd', 'rsi', 'kdj', 'bollinger'
FACTOR_NAME = 'macd'
# ==============================


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

    # 技术因子：根据不同指标选股
    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

    # 获取足够长的历史数据
    prices = get_price(
        stocks,
        end_date=context.current_dt,
        count=60,  # 足够计算各种指标
        frequency='daily',
        fields=['close', 'high', 'low'],
        panel=False
    )

    if prices is None or prices.empty:
        return None

    results = []

    for stock in stocks:
        stock_data = prices[prices['code'] == stock].sort_values('time')
        if len(stock_data) < 30:
            continue

        close = stock_data['close'].values
        high = stock_data['high'].values
        low = stock_data['low'].values

        try:
            if g.factor == 'macd':
                score = calculate_macd(close)
            elif g.factor == 'rsi':
                score = calculate_rsi(close)
            elif g.factor == 'kdj':
                score = calculate_kdj(close, high, low)
            elif g.factor == 'bollinger':
                score = calculate_bollinger(close)
            else:
                continue

            if score is not None and not np.isnan(score):
                results.append({'code': stock, g.factor: score})
        except:
            continue

    if not results:
        return None

    df = pd.DataFrame(results)
    return df


def calculate_macd(close):
    """
    MACD指标
    选股逻辑：DIF > 0 且 DIF > DEA（金叉且在零轴上方）
    返回：DIF值（越大越好）
    """
    import pandas as pd

    prices = pd.Series(close)

    # 计算EMA
    ema_fast = prices.ewm(span=12, adjust=False).mean()
    ema_slow = prices.ewm(span=26, adjust=False).mean()

    # DIF = 快线 - 慢线
    dif = ema_fast - ema_slow

    # DEA = DIF的9日EMA
    dea = dif.ewm(span=9, adjust=False).mean()

    # 只选DIF > 0 且 DIF > DEA的股票
    if dif.iloc[-1] > 0 and dif.iloc[-1] > dea.iloc[-1]:
        return float(dif.iloc[-1])
    else:
        return -999  # 不符合条件的给负分


def calculate_rsi(close, period=14):
    """
    RSI指标
    选股逻辑：RSI在40-60区间（避免超买超卖）
    返回：距离50的负距离（越接近50越好）
    """
    import pandas as pd

    prices = pd.Series(close)
    delta = prices.diff()

    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)

    avg_gain = gain.rolling(window=period).mean()
    avg_loss = loss.rolling(window=period).mean()

    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))

    rsi_value = rsi.iloc[-1]

    # RSI在40-60区间最好，距离50越近越好
    if 40 <= rsi_value <= 60:
        return 100 - abs(rsi_value - 50)  # 距离50越近分数越高
    else:
        return -999  # 超买超卖区域不选


def calculate_kdj(close, high, low, period=9):
    """
    KDJ指标
    选股逻辑：K > D 且 J在20-80区间
    返回：K值（越大越好）
    """
    import pandas as pd

    df = pd.DataFrame({'close': close, 'high': high, 'low': low})

    # 计算RSV
    low_min = df['low'].rolling(window=period).min()
    high_max = df['high'].rolling(window=period).max()
    rsv = (df['close'] - low_min) / (high_max - low_min) * 100

    # 计算K、D、J
    k = rsv.ewm(com=2, adjust=False).mean()
    d = k.ewm(com=2, adjust=False).mean()
    j = 3 * k - 2 * d

    k_value = k.iloc[-1]
    d_value = d.iloc[-1]
    j_value = j.iloc[-1]

    # K > D 且 J在20-80区间
    if k_value > d_value and 20 <= j_value <= 80:
        return float(k_value)
    else:
        return -999


def calculate_bollinger(close, period=20, std_dev=2):
    """
    布林带指标
    选股逻辑：价格接近下轨（均值回归机会）
    返回：(下轨 - 价格) / 标准差（越大越好，表示越接近下轨）
    """
    import pandas as pd

    prices = pd.Series(close)

    # 计算布林带
    ma = prices.rolling(window=period).mean()
    std = prices.rolling(window=period).std()

    upper = ma + std_dev * std
    lower = ma - std_dev * std

    current_price = prices.iloc[-1]
    current_lower = lower.iloc[-1]
    current_std = std.iloc[-1]

    # 价格越接近下轨，分数越高
    # 但不选已经跌破下轨太多的（可能是基本面恶化）
    distance = (current_lower - current_price) / current_std

    if -0.5 < distance < 0.5:  # 在下轨附近
        return float(distance + 1)  # 转换为正分
    else:
        return -999


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)