import pandas as pd

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


class BollingerBands(Indicator):
    def __init__(
        self,
        window: int = 20,
        std_multiplier: float = 2.0,
        candles_type: str = "15m",
        label: str = "bb",
    ):
        super().__init__(name="Bollinger Bands", label=label)
        self.window = window
        self.std_multiplier = std_multiplier
        self.candles_type = candles_type

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

        df = pd.DataFrame(index=candle_data.index)

        df["price"] = (candle_data["ask_close"] + candle_data["bid_close"]) / 2
        df["avg"] = df["price"].rolling(window=self.window).mean()
        df["std"] = df["price"].rolling(window=self.window).std() * self.std_multiplier
        df["upper"] = df["avg"] + df["std"]
        df["lower"] = df["avg"] - df["std"]

        price = df["price"]
        lower = df["lower"]
        upper = df["upper"]

        # Was the previous candle outside the band?
        prev_below_lower = price.shift(1) < lower.shift(1)
        prev_above_upper = price.shift(1) > upper.shift(1)

        # Did the current candle close back inside?
        now_above_lower = price >= lower  # re-entered from below  → buy
        now_below_upper = price <= upper  # re-entered from above  → sell

        df["signal"] = 0
        df.loc[prev_below_lower & now_above_lower, "signal"] = 1  # buy crossover
        df.loc[prev_above_upper & now_below_upper, "signal"] = -1  # sell crossover

        df.dropna(inplace=True)
        df.drop(columns=["price", "std"], inplace=True)

        return df

    def __str__(self) -> str:
        return f"BB(window={self.window}, std_multiplier={self.std_multiplier}, candles={self.candles_type})"
