"""
动量突破策略。

价格突破 N 日新高 → 买入
价格跌破 N 日均线 → 平仓
"""

import pandas as pd

from backtest.base import BaseStrategy
from backtest.signal import factor_to_threshold_signal


class MomentumBreakoutStrategy(BaseStrategy):
    """动量突破策略。"""

    category = "trend"

    def __init__(self, lookback: int = 20, exit_period: int = 10):
        self.lookback = lookback
        self.exit_period = exit_period
        self.name = f"mom_breakout_{lookback}"

    def generate_signals(self, factor_df: pd.DataFrame) -> pd.Series:
        if "close" not in factor_df.columns:
            raise ValueError("factor_df 缺少 'close' 列")

        close = factor_df["close"]
        # 买入信号：突破 N 日新高
        rolling_high = close.rolling(window=self.lookback, min_periods=self.lookback).max()
        breakout = close >= rolling_high.shift(1)

        # 平仓信号：跌破 exit 日均线
        exit_ma = close.rolling(window=self.exit_period, min_periods=self.exit_period).mean()

        signals = pd.Series(0, index=close.index)
        signals[breakout] = 1
        signals[close < exit_ma] = 0

        return self._dedup(signals)

    @staticmethod
    def _dedup(signals: pd.Series) -> pd.Series:
        """只保留第一个买入和第一个卖出信号。"""
        result = signals.copy()
        prev = -1
        for i in range(len(result)):
            if result.iloc[i] == prev:
                result.iloc[i] = -1
            else:
                prev = result.iloc[i]
        return result[result != -1].reindex(signals.index).fillna(-1)