"""
SMA 双均线交叉策略

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

指标计算:
- SMA10: 10 日简单移动平均线
- SMA30: 30 日简单移动平均线
- SMA60: 60 日简单移动平均线
- SMA120: 120 日简单移动平均线
"""

import pandas as pd
from backtesting import Strategy
from backtesting.lib import crossover


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

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

    返回:
        DataFrame, 添加了指标列
    """
    data = data.copy()

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

    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.position.close()  # 先平掉现有仓位
            self.buy()  # 开多仓

        # 死叉：短期均线下穿长期均线
        elif crossover(self.data.sma30, self.data.sma10):
            self.position.close()  # 先平掉现有仓位
            self.sell()  # 开空仓