"""
MACD 趋势跟踪策略

策略逻辑:
- 当 MACD 线上穿信号线时 (金叉)，且价格 > EMA 时，买入
- 当 MACD 线下穿信号线时 (死叉)，或价格 < EMA 时，卖出

指标计算:
- MACD(10, 20, 9): 快线 10 日，慢线 20 日，信号线 9 日
- EMA: 200 日指数移动平均线（趋势确认）

参数选择理由:
- 快线 10: 比标准 12 更敏感，适应 A 股较高波动性
- 慢线 20: 比标准 26 更快响应，同时保持趋势跟踪稳定性
- 信号线 9: 保持标准，避免信号过于频繁
- EMA: 被广泛认可为牛熊分界线，避免逆势交易

趋势过滤:
- EMA 上方: 确认为上升趋势，允许开多仓
- EMA 下方: 确认为下降趋势，不开多仓，强制平仓

Author: Sisyphus
Date: 2025-01-27
"""

from backtesting import Strategy
from backtesting.lib import crossover


def calculate_indicators(data):
    """
    计算策略所需的技术指标

    使用 ta-lib 库计算 MACD 和 EMA 指标

    参数:
        data: DataFrame, 包含 [Open, High, Low, Close, Volume, factor]

    返回:
        DataFrame, 添加了指标列:
        - macd: MACD 线 (macd)
        - signal: MACD 信号线 (DEA)
        - hist: MACD 柱状图 (Histogram)
        - ema: 日指数移动平均线
    """
    data = data.copy()

    # 计算 MACD 指标 (10, 20, 9)
    # talib.MACD 返回三个值: (macd, macdsignal, macdhist)
    macd, macdsignal, macdhist = talib.MACD(data["Close"], fastperiod=10, slowperiod=20, signalperiod=9)

    data["macd"] = macd
    data["signal"] = macdsignal
    data["hist"] = macdhist

    # 计算 EMA 趋势线
    data["ema"] = talib.SMA(data["Close"], timeperiod=120)

    return data


def get_strategy():
    """
    返回策略类

    返回:
        MacdTrendStrategy 类
    """
    return MacdTrendStrategy


class MacdTrendStrategy(Strategy):
    """
    MACD 趋势跟踪策略

    结合 MACD 金叉/死叉信号和 EMA 趋势过滤

    参数:
        fast_period: MACD 快线周期 (默认: 10)
        slow_period: MACD 慢线周期 (默认: 20)
        signal_period: MACD 信号线周期 (默认: 9)
    """

    # 可配置参数
    fast_period = 10
    slow_period = 20
    signal_period = 9

    def init(self):
        """
        初始化策略
        注册指标到 backtesting 框架
        """
        # 注册 MACD 线
        self.macd = self.I(lambda x: x, self.data.macd)

        # 注册 MACD 信号线
        self.signal = self.I(lambda x: x, self.data.signal)

        # 注册 EMA 趋势线
        self.ema = self.I(lambda x: x, self.data.ema)

    def next(self):
        """
        每个时间步的决策逻辑

        买入条件:
        - MACD 金叉 (MACD 线上穿信号线)
        - 价格 > EMA (确认上升趋势)

        卖出条件:
        - MACD 死叉 (MACD 线下穿信号线)
        - 或价格 < EMA (趋势转向，强制平仓)
        """
        # 买入条件: MACD 金叉 AND 价格 > EMA
        if crossover(self.macd, self.signal) and self.data.Close[-1] > self.ema[-1]:
            self.buy()  # 开多仓

        # 卖出条件: MACD 死叉 OR 价格 < EMA
        elif self.position.size > 0 and (crossover(self.signal, self.macd) or self.data.Close[-1] < self.ema[-1]):
            self.position.close()  # 平掉多仓


# 导入 talib (必须在文件末尾，因为 calculate_indicators 函数中使用了 talib)
import talib