"""Trend: ADX (+DI/-DI), Aroon, Supertrend."""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

from pyalbedo.indicators.base import Indicator
from pyalbedo.indicators.volatility import ATR
from pyalbedo.utils.validation import ensure_float64_1d


class ADX(Indicator[NDArray[np.float64]]):
    """
    Average Directional Index with +DI and -DI (Wilder smoothing).
    Returns (adx, plus_di, minus_di) as 2D array shape (n, 3).
    """

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

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

    @property
    def lookback(self) -> int:
        return 2 * self._period  # conservative

    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)
        n = len(close)
        prev_high = np.roll(high, 1)
        prev_high[0] = high[0]
        prev_low = np.roll(low, 1)
        prev_low[0] = low[0]
        prev_close = np.roll(close, 1)
        prev_close[0] = close[0]
        plus_dm = np.where((high - prev_high) > (prev_low - low), np.maximum(high - prev_high, 0), 0)
        minus_dm = np.where((prev_low - low) > (high - prev_high), np.maximum(prev_low - low, 0), 0)
        tr = np.maximum(
            high - low,
            np.maximum(np.abs(high - prev_close), np.abs(low - prev_close)),
        )
        alpha = 1.0 / self._period
        plus_di = np.full(n, np.nan, dtype=np.float64)
        minus_di = np.full(n, np.nan, dtype=np.float64)
        dx = np.full(n, np.nan, dtype=np.float64)
        adx = np.full(n, np.nan, dtype=np.float64)
        smooth_plus = np.nan
        smooth_minus = np.nan
        smooth_tr = np.nan
        for i in range(1, n):
            if np.isnan(smooth_plus):
                smooth_plus = np.mean(plus_dm[1 : self._period + 1])
                smooth_minus = np.mean(minus_dm[1 : self._period + 1])
                smooth_tr = np.mean(tr[1 : self._period + 1])
            else:
                smooth_plus = (1 - alpha) * smooth_plus + alpha * plus_dm[i]
                smooth_minus = (1 - alpha) * smooth_minus + alpha * minus_dm[i]
                smooth_tr = (1 - alpha) * smooth_tr + alpha * tr[i]
            if i >= self._period and smooth_tr > 0:
                plus_di[i] = 100.0 * smooth_plus / smooth_tr
                minus_di[i] = 100.0 * smooth_minus / smooth_tr
                di_sum = plus_di[i] + minus_di[i]
                if di_sum > 0:
                    dx[i] = 100.0 * np.abs(plus_di[i] - minus_di[i]) / di_sum
        for i in range(2 * self._period - 1, n):
            if not np.isnan(dx[i - self._period + 1 : i + 1]).any():
                adx[i] = np.mean(dx[i - self._period + 1 : i + 1])
        return np.column_stack((adx, plus_di, minus_di))


class Aroon(Indicator[NDArray[np.float64]]):
    """
    Aroon Up and Down: 100 * (period - periods_since_high) / period, same for low.
    Returns (aroon_up, aroon_down) as 2D array shape (n, 2).
    """

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

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

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

    def calculate(
        self,
        high: NDArray[np.float64],
        low: NDArray[np.float64],
    ) -> NDArray[np.float64]:
        high = ensure_float64_1d(high)
        low = ensure_float64_1d(low)
        n = len(high)
        aroon_up = np.full(n, np.nan, dtype=np.float64)
        aroon_down = np.full(n, np.nan, dtype=np.float64)
        for i in range(self._period, n):
            window_h = high[i - self._period + 1 : i + 1]
            window_l = low[i - self._period + 1 : i + 1]
            idx_high = np.argmax(window_h)
            idx_low = np.argmin(window_l)
            aroon_up[i] = 100.0 * (self._period - (self._period - 1 - idx_high)) / self._period
            aroon_down[i] = 100.0 * (self._period - (self._period - 1 - idx_low)) / self._period
        return np.column_stack((aroon_up, aroon_down))


class Supertrend(Indicator[NDArray[np.float64]]):
    """
    Supertrend (ATR-based): upper_band = (high+low)/2 + mult*ATR, lower_band = (high+low)/2 - mult*ATR;
    trend line and direction. Returns (supertrend_line, direction) as 2D array shape (n, 2).
    direction: 1 = uptrend, -1 = downtrend.
    """

    def __init__(self, period: int = 10, multiplier: float = 3.0) -> None:
        if period < 1 or multiplier <= 0:
            raise ValueError("period >= 1, multiplier > 0")
        self._period = period
        self._mult = multiplier
        self._atr = ATR(period)

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

    @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)
        n = len(close)
        atr = self._atr.calculate(high, low, close)
        hl2 = (high + low) / 2.0
        upper = hl2 + self._mult * atr
        lower = hl2 - self._mult * atr
        supertrend = np.full(n, np.nan, dtype=np.float64)
        direction = np.full(n, np.nan, dtype=np.float64)
        cur_upper = upper[self._period]
        cur_lower = lower[self._period]
        cur_dir = 1.0
        for i in range(self._period, n):
            if np.isnan(atr[i]):
                continue
            if close[i] > cur_upper:
                cur_dir = 1.0
                cur_lower = lower[i]
                cur_upper = upper[i]
            elif close[i] < cur_lower:
                cur_dir = -1.0
                cur_upper = upper[i]
                cur_lower = lower[i]
            else:
                if cur_dir == 1.0:
                    cur_lower = max(lower[i], cur_lower)
                    if lower[i] > cur_lower:
                        cur_lower = lower[i]
                else:
                    cur_upper = min(upper[i], cur_upper)
                    if upper[i] < cur_upper:
                        cur_upper = upper[i]
            supertrend[i] = cur_lower if cur_dir == 1.0 else cur_upper
            direction[i] = cur_dir
        return np.column_stack((supertrend, direction))
