"""
SMA 双均线交叉策略

策略逻辑:
- 当短期均线上穿长期均线时 (金叉)，买入
- 当短期均线下穿长期均线时 (死叉)，卖出

指标计算 (使用 ta-lib):
- SMA10: 10 日简单移动平均线
- SMA30: 30 日简单移动平均线
- SMA60: 60 日简单移动平均线
- SMA120: 120 日简单移动平均线
"""

from backtesting import Strategy
from backtesting.lib import crossover


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

    使用 ta-lib 库计算 SMA 指标

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

    返回:
        DataFrame, 添加了指标列:
        - sma10: 10 日简单移动平均线
        - sma30: 30 日简单移动平均线
        - sma60: 60 日简单移动平均线
        - sma120: 120 日简单移动平均线
    """
    data = data.copy()

    # 计算不同周期的移动平均线
    data["sma10"] = talib.SMA(data["Close"], timeperiod=10)
    data["sma30"] = talib.SMA(data["Close"], timeperiod=30)
    data["sma60"] = talib.SMA(data["Close"], timeperiod=60)
    data["sma120"] = talib.SMA(data["Close"], timeperiod=120)

    return data


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

    返回:
        SmaCross 类
    """
    return SmaCross


class SmaCross(Strategy):
    """
    SMA 双均线交叉策略

    参数:
        short_period: 短期均线周期 (默认: 10)
        long_period: 长期均线周期 (默认: 30)
    """

    # 可配置参数
    short_period = 10
    long_period = 30

    def init(self):
        """
        初始化策略
        注册指标到 backtesting 框架
        """
        self.sma_short = self.I(lambda x: x, self.data.sma10)
        self.sma_long = self.I(lambda x: x, self.data.sma30)

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

        金叉: 短期均线上穿长期均线 → 买入
        死叉: 短期均线下穿长期均线 → 卖出
        """
        # 金叉：短期均线上穿长期均线
        if crossover(self.data.sma10, self.data.sma30):
            self.buy()  # 开多仓

        # 死叉：短期均线下穿长期均线
        elif self.position.size > 0 and crossover(self.data.sma30, self.data.sma10):
            self.position.close()  # 开空仓


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