import pandas as pd

from signalab.data.market_data import MarketData
from signalab.indicator.base import Indicator


class Stochastic(Indicator):
    """
    Stochastic Oscillator.

    %K = (close - lowest_low) / (highest_high - lowest_low) * 100
    %K_slowed = SMA(%K, slowing)        — "slow %K"
    %D = SMA(%K_slowed, d_period)       — signal line

    Generates:
        signal = 1  (BUY)  when %K_slowed < lower (oversold)
        signal = -1 (SELL) when %K_slowed > upper (overbought)
        signal = 0  (HOLD) otherwise
    """

    def __init__(
        self,
        k_period: int = 5,
        d_period: int = 3,
        slowing: int = 3,
        upper: float = 80.0,
        lower: float = 20.0,
        candles_type: str = "5m",
        label: str = "stoch",
        crossover_only: bool = True,
    ):
        super().__init__(name="Stochastic", label=label)
        self.k_period = k_period
        self.d_period = d_period
        self.slowing = slowing
        self.upper = upper
        self.lower = lower
        self.candles_type = candles_type
        self.crossover_only = crossover_only

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

        # Use mid high/low/close
        if "ask_high" in candles.columns and "bid_high" in candles.columns:
            high = (candles["ask_high"] + candles["bid_high"]) / 2
            low = (candles["ask_low"] + candles["bid_low"]) / 2
            close = (candles["ask_close"] + candles["bid_close"]) / 2
        else:
            close = (candles["ask_close"] + candles["bid_close"]) / 2
            high = close
            low = close

        lowest_low = low.rolling(window=self.k_period).min()
        highest_high = high.rolling(window=self.k_period).max()

        denom = highest_high - lowest_low
        # Avoid division by zero when high == low (flat candles)
        denom = denom.replace(0, float("nan"))

        raw_k = (close - lowest_low) / denom * 100
        slow_k = raw_k.rolling(window=self.slowing).mean()  # slowing applied to %K
        slow_d = slow_k.rolling(window=self.d_period).mean()  # %D = SMA of slow %K

        df = pd.DataFrame(index=candles.index)
        df["k"] = slow_k
        df["d"] = slow_d
        df["signal"] = 0

        if self.crossover_only:
            prev_k = slow_k.shift(1)
            df.loc[(prev_k < self.lower) & (slow_k >= self.lower), "signal"] = 1
            df.loc[(prev_k > self.upper) & (slow_k <= self.upper), "signal"] = -1
        else:
            df.loc[slow_k < self.lower, "signal"] = 1
            df.loc[slow_k > self.upper, "signal"] = -1

        df.dropna(inplace=True)
        return df

    def __str__(self) -> str:
        return (
            f"Stochastic(k={self.k_period}, d={self.d_period}, "
            f"slowing={self.slowing}, upper={self.upper}, lower={self.lower}, "
            f"candles={self.candles_type})"
        )
