"""Moving averages: SMA, EMA, WMA, DEMA, TEMA, Hull MA."""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

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


class SMA(Indicator[NDArray[np.float64]]):
    """Simple Moving Average. Wilder (1978)."""

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

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

    @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)
        result = np.full(n, np.nan, dtype=np.float64)
        for i in range(self._period - 1, n):
            result[i] = np.mean(prices[i - self._period + 1 : i + 1])
        return result


class EMA(Indicator[NDArray[np.float64]]):
    """Exponential Moving Average. alpha = 2 / (period + 1)."""

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

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

    @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)
        result = np.full(n, np.nan, dtype=np.float64)
        if n < self._period:
            return result
        result[self._period - 1] = np.mean(prices[: self._period])
        for i in range(self._period, n):
            result[i] = self._alpha * prices[i] + (1 - self._alpha) * result[i - 1]
        return result


class WMA(Indicator[NDArray[np.float64]]):
    """Weighted Moving Average: linear weights 1..period."""

    def __init__(self, period: int) -> None:
        if period < 1:
            raise ValueError("period must be >= 1")
        self._period = period
        self._weights = np.arange(1, period + 1, dtype=np.float64) / np.arange(
            1, period + 1, dtype=np.float64
        ).sum()

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

    @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)
        result = np.full(n, np.nan, dtype=np.float64)
        for i in range(self._period - 1, n):
            result[i] = np.dot(self._weights, prices[i - self._period + 1 : i + 1])
        return result


class DEMA(Indicator[NDArray[np.float64]]):
    """Double Exponential Moving Average. 2*EMA - EMA(EMA)."""

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

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

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

    def calculate(self, *inputs: NDArray[np.float64]) -> NDArray[np.float64]:
        prices = ensure_float64_1d(inputs[0])
        ema1 = self._ema.calculate(prices)
        ema2 = self._ema.calculate(ema1)
        return np.where(np.isnan(ema2), np.nan, 2 * ema1 - ema2)


class TEMA(Indicator[NDArray[np.float64]]):
    """Triple Exponential Moving Average. 3*EMA - 3*EMA(EMA) + EMA(EMA(EMA))."""

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

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

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

    def calculate(self, *inputs: NDArray[np.float64]) -> NDArray[np.float64]:
        prices = ensure_float64_1d(inputs[0])
        ema1 = self._ema.calculate(prices)
        ema2 = self._ema.calculate(ema1)
        ema3 = self._ema.calculate(ema2)
        return np.where(np.isnan(ema3), np.nan, 3 * ema1 - 3 * ema2 + ema3)


class HullMA(Indicator[NDArray[np.float64]]):
    """Hull Moving Average: WMA(2*WMA(n/2) - WMA(n), sqrt(n))."""

    def __init__(self, period: int) -> None:
        if period < 2:
            raise ValueError("period must be >= 2")
        self._period = period
        half = period // 2
        self._sqrt_n = int(np.sqrt(period))
        if self._sqrt_n < 1:
            self._sqrt_n = 1
        self._wma_half = WMA(half)
        self._wma_full = WMA(period)
        self._wma_sqrt = WMA(self._sqrt_n)

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

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

    def calculate(self, *inputs: NDArray[np.float64]) -> NDArray[np.float64]:
        prices = ensure_float64_1d(inputs[0])
        wma_h = self._wma_half.calculate(prices)
        wma_f = self._wma_full.calculate(prices)
        raw = 2 * wma_h - wma_f
        return self._wma_sqrt.calculate(raw)
