"""Volatility: ATR, Bollinger Bands, Keltner Channels."""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

from pyalbedo.indicators.base import Indicator
from pyalbedo.indicators.moving_averages import EMA, SMA
from pyalbedo.utils.validation import ensure_float64_1d


class ATR(Indicator[NDArray[np.float64]]):
    """
    Average True Range (Wilder smoothing).
    TR = max(high-low, |high-prev_close|, |low-prev_close|); ATR = EMA(TR, period).
    """

    def __init__(self, period: int = 14) -> None:
        if period < 1:
            raise ValueError("period must be >= 1")
        self._period = period
        self._ema = EMA(period)

    @property
    def name(self) -> str:
        return f"ATR({self._period})"

    @property
    def lookback(self) -> int:
        return self._period + 1

    def calculate(
        self,
        high: NDArray[np.float64],
        low: NDArray[np.float64],
        close: NDArray[np.float64],
    ) -> NDArray[np.float64]:
        high = ensure_float64_1d(high)
        low = ensure_float64_1d(low)
        close = ensure_float64_1d(close)
        prev_close = np.roll(close, 1)
        prev_close[0] = close[0]
        tr = np.maximum(
            high - low,
            np.maximum(
                np.abs(high - prev_close),
                np.abs(low - prev_close),
            ),
        )
        return self._ema.calculate(tr)


class BollingerBands(Indicator[NDArray[np.float64]]):
    """
    Bollinger Bands: middle = SMA(close), upper = middle + k*std, lower = middle - k*std.
    Returns (upper, middle, lower) as 2D array shape (n, 3).
    """

    def __init__(self, period: int = 20, num_std: float = 2.0) -> None:
        if period < 1 or num_std < 0:
            raise ValueError("period >= 1 and num_std >= 0")
        self._period = period
        self._num_std = num_std
        self._sma = SMA(period)

    @property
    def name(self) -> str:
        return f"BollingerBands({self._period},{self._num_std})"

    @property
    def lookback(self) -> int:
        return self._period

    def calculate(self, *inputs: NDArray[np.float64]) -> NDArray[np.float64]:
        prices = ensure_float64_1d(inputs[0])
        n = len(prices)
        middle = self._sma.calculate(prices)
        upper = np.full(n, np.nan, dtype=np.float64)
        lower = np.full(n, np.nan, dtype=np.float64)
        for i in range(self._period - 1, n):
            std = np.std(prices[i - self._period + 1 : i + 1])
            upper[i] = middle[i] + self._num_std * std
            lower[i] = middle[i] - self._num_std * std
        return np.column_stack((upper, middle, lower))


class KeltnerChannels(Indicator[NDArray[np.float64]]):
    """
    Keltner Channels: middle = EMA(close), upper = middle + k*ATR, lower = middle - k*ATR.
    Returns (upper, middle, lower) as 2D array shape (n, 3).
    """

    def __init__(self, ema_period: int = 20, atr_period: int = 10, atr_mult: float = 2.0) -> None:
        if ema_period < 1 or atr_period < 1 or atr_mult < 0:
            raise ValueError("invalid Keltner params")
        self._ema_period = ema_period
        self._atr_period = atr_period
        self._atr_mult = atr_mult
        self._ema = EMA(ema_period)
        self._atr = ATR(atr_period)

    @property
    def name(self) -> str:
        return f"KeltnerChannels({self._ema_period},{self._atr_period},{self._atr_mult})"

    @property
    def lookback(self) -> int:
        return max(self._ema_period, self._atr_period + 1)

    def calculate(
        self,
        high: NDArray[np.float64],
        low: NDArray[np.float64],
        close: NDArray[np.float64],
    ) -> NDArray[np.float64]:
        high = ensure_float64_1d(high)
        low = ensure_float64_1d(low)
        close = ensure_float64_1d(close)
        middle = self._ema.calculate(close)
        atr = self._atr.calculate(high, low, close)
        upper = middle + self._atr_mult * atr
        lower = middle - self._atr_mult * atr
        return np.column_stack((upper, middle, lower))
