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 ATR(Indicator):
    """Average True Range volatility indicator.

    Notes
    -----
    ATR is reported in price units and does not encode direction.
    """

    def __init__(
        self,
        period: int = 14,
        candles_type: str = "5m",
        label: str = "atr",
    ):
        """Initialize ATR configuration.

        Parameters
        ----------
        period : int, default=14
            Wilder smoothing period.
        candles_type : str, default="5m"
            Candle timeframe queried from :class:`MarketData`.
        label : str, default="atr"
            Output label namespace.
        """
        super().__init__(name="ATR", label=label)
        self.period = period
        self.candles_type = candles_type

    def compute(self, data: MarketData) -> pd.DataFrame:
        """Compute Average True Range values.

        Parameters
        ----------
        data : MarketData
            Source market data.

        Returns
        -------
        pd.DataFrame
            Dataframe with one ``atr`` column.
        """
        candles = data.query(self.candles_type)

        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 = _mid_price(candles)
        else:
            close = _mid_price(candles)
            high = close
            low = close

        prev_close = close.shift(1)
        tr = pd.concat(
            [
                high - low,
                (high - prev_close).abs(),
                (low - prev_close).abs(),
            ],
            axis=1,
        ).max(axis=1)

        df = pd.DataFrame(index=candles.index)
        df["atr"] = _smma(tr, self.period)
        df.dropna(inplace=True)
        return df

    def __str__(self) -> str:
        """Return a compact string representation.

        Returns
        -------
        str
            Indicator summary.
        """
        return f"ATR(period={self.period}, candles={self.candles_type})"
