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 ADX(Indicator):
    """
    Average Directional Index — measures trend strength, not direction.

    Generates:
        signal = 1  (BUY)  when ADX > threshold AND +DI > -DI (strong uptrend)
        signal = -1 (SELL) when ADX > threshold AND -DI > +DI (strong downtrend)
        signal = 0  (HOLD) when ADX < threshold (no clear trend)
    """

    def __init__(
        self,
        period: int = 10,
        threshold: float = 25.0,
        candles_type: str = "5m",
        label: str = "adx",
    ):
        super().__init__(name="ADX", label=label)
        self.period = period
        self.threshold = threshold
        self.candles_type = candles_type

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

        # Use mid high/low if available, otherwise approximate from 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 = _mid_price(candles)
        else:
            close = _mid_price(candles)
            high = close
            low = close

        # Directional movement
        up_move = high.diff()
        down_move = -low.diff()

        plus_dm = up_move.where((up_move > down_move) & (up_move > 0), 0.0)
        minus_dm = down_move.where((down_move > up_move) & (down_move > 0), 0.0)

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

        # Wilder's smoothing
        atr = _smma(tr, self.period)
        plus_di = 100 * _smma(plus_dm, self.period) / atr
        minus_di = 100 * _smma(minus_dm, self.period) / atr

        dx = (
            100
            * (plus_di - minus_di).abs()
            / (plus_di + minus_di).replace(0, float("inf"))
        )
        adx = _smma(dx, self.period)

        df = pd.DataFrame(index=candles.index)
        df["adx"] = adx
        df["plus_di"] = plus_di
        df["minus_di"] = minus_di
        df["signal"] = 0
        df.loc[(adx > self.threshold) & (plus_di > minus_di), "signal"] = 1
        df.loc[(adx > self.threshold) & (minus_di > plus_di), "signal"] = -1
        df.dropna(inplace=True)
        return df

    def __str__(self) -> str:
        return (
            f"ADX(period={self.period}, threshold={self.threshold}, "
            f"candles={self.candles_type})"
        )
