"""均线类策略：双均线交叉、单均线（价格上穿/下穿）。"""
from __future__ import annotations

import pandas as pd

from ...indicators import ma
from ..broker import PaperBroker
from .base import Strategy


class MACrossStrategy(Strategy):
    """双均线交叉：快线上穿慢线买入，下穿卖出（金叉/死叉）。"""

    def __init__(self, fast: float = 5, slow: float = 20):
        self.fast = int(fast)
        self.slow = int(slow)
        self._ind: pd.DataFrame | None = None
        self._prev_fast: float | None = None
        self._prev_slow: float | None = None

    def compute(self, close: pd.Series, high: pd.Series, low: pd.Series) -> pd.DataFrame:
        self._ind = pd.DataFrame({"fast": ma(close, self.fast), "slow": ma(close, self.slow)})
        return self._ind

    def on_bar(self, i: int, row: pd.Series, broker: PaperBroker) -> None:
        f = float(self._ind["fast"].iloc[i])
        s = float(self._ind["slow"].iloc[i])
        if self._prev_fast is None:
            self._prev_fast, self._prev_slow = f, s
            return

        golden = self._prev_fast <= self._prev_slow and f > s
        death = self._prev_fast >= self._prev_slow and f < s

        price = float(row["close"])
        ts = row["ts"]
        if golden and pd.notna(f):
            broker.buy_max(ts, price)
        elif death and broker.position > 0 and pd.notna(f):
            broker.sell_all(ts, price)

        self._prev_fast, self._prev_slow = f, s


class SingleMAStrategy(Strategy):
    """单均线：收盘价上穿均线买入，下穿均线卖出。"""

    def __init__(self, period: float = 20):
        self.period = int(period)
        self._ind: pd.DataFrame | None = None
        self._prev_above: bool | None = None

    def compute(self, close: pd.Series, high: pd.Series, low: pd.Series) -> pd.DataFrame:
        self._ind = pd.DataFrame({"ma": ma(close, self.period)})
        return self._ind

    def on_bar(self, i: int, row: pd.Series, broker: PaperBroker) -> None:
        m = float(self._ind["ma"].iloc[i])
        price = float(row["close"])
        ts = row["ts"]
        above = price > m
        if self._prev_above is None or pd.isna(m):
            self._prev_above = above
            return

        cross_up = above and not self._prev_above        # 上穿
        cross_down = (not above) and self._prev_above    # 下穿

        if cross_up:
            broker.buy_max(ts, price)
        elif cross_down and broker.position > 0:
            broker.sell_all(ts, price)

        self._prev_above = above