import pandas as pd

from signalab.data.market_data import MarketData
from signalab.indicator._utils import _mid_price, _smma
from signalab.indicator.base import Indicator


class RSI(Indicator):
    """
    Relative Strength Index.

    Generates:
        signal = 1  (BUY)  when RSI < buy_level  (oversold)
        signal = -1 (SELL) when RSI > sell_level (overbought)
        signal = 0  (HOLD) otherwise
    """

    def __init__(
        self,
        period: int = 9,
        buy_level: float = 40.0,
        sell_level: float = 60.0,
        candles_type: str = "5m",
        label: str = "rsi",
    ):
        super().__init__(name="RSI", label=label)
        self.period = period
        self.buy_level = buy_level
        self.sell_level = sell_level
        self.candles_type = candles_type

    def compute(self, data: MarketData) -> pd.DataFrame:
        candles = data.query(self.candles_type)
        price = _mid_price(candles)

        delta = price.diff()
        gain = delta.clip(lower=0)
        loss = -delta.clip(upper=0)

        # Use SMMA (Wilder's smoothing) — standard for RSI
        avg_gain = _smma(gain, self.period)
        avg_loss = _smma(loss, self.period)

        rs = avg_gain / avg_loss.replace(0, float("inf"))
        rsi = 100 - (100 / (1 + rs))

        df = pd.DataFrame(index=candles.index)
        df["rsi"] = rsi
        prev_rsi = rsi.shift(1)

        # Buy: RSI was below buy_level, now crosses back above it (oversold → recovering)
        buy_signal = (prev_rsi < self.buy_level) & (rsi >= self.buy_level)

        # Sell: RSI was above sell_level, now crosses back below it (overbought → falling)
        sell_signal = (prev_rsi > self.sell_level) & (rsi <= self.sell_level)

        df["signal"] = 0
        df.loc[buy_signal, "signal"] = 1
        df.loc[sell_signal, "signal"] = -1
        df.dropna(inplace=True)
        return df

    def __str__(self) -> str:
        return (
            f"RSI(period={self.period}, "
            f"buy={self.buy_level}, sell={self.sell_level}, "
            f"candles={self.candles_type})"
        )
