"""Technical indicators module for jbqlab.

This module provides common technical analysis indicators that can be
used to build custom strategies.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

__all__ = [
    # Trend Indicators
    "sma",
    "ema",
    "wma",
    "dema",
    "tema",
    # Momentum Indicators
    "rsi",
    "stochastic",
    "macd_indicator",
    "roc",
    "momentum",
    # Volatility Indicators
    "bollinger_bands",
    "atr",
    "keltner_channel",
    "donchian_channel",
    # Volume Indicators
    "obv",
    "vwap",
    # Other
    "adx",
    "cci",
    "williams_r",
]


# =============================================================================
# Trend Indicators
# =============================================================================


def sma(series: pd.Series, window: int) -> pd.Series:
    """Simple Moving Average.

    Args:
        series: Price series.
        window: Lookback window.

    Returns:
        SMA series.
    """
    return series.rolling(window=window, min_periods=window).mean()


def ema(series: pd.Series, span: int) -> pd.Series:
    """Exponential Moving Average.

    Args:
        series: Price series.
        span: EMA span.

    Returns:
        EMA series.
    """
    return series.ewm(span=span, adjust=False).mean()


def wma(series: pd.Series, window: int) -> pd.Series:
    """Weighted Moving Average.

    Args:
        series: Price series.
        window: Lookback window.

    Returns:
        WMA series with linearly increasing weights.
    """
    weights = np.arange(1, window + 1)
    return series.rolling(window=window).apply(
        lambda x: np.sum(weights * x) / weights.sum(), raw=True
    )


def dema(series: pd.Series, span: int) -> pd.Series:
    """Double Exponential Moving Average.

    DEMA = 2 * EMA - EMA(EMA)

    Args:
        series: Price series.
        span: EMA span.

    Returns:
        DEMA series.
    """
    ema1 = ema(series, span)
    ema2 = ema(ema1, span)
    return 2 * ema1 - ema2


def tema(series: pd.Series, span: int) -> pd.Series:
    """Triple Exponential Moving Average.

    TEMA = 3 * EMA - 3 * EMA(EMA) + EMA(EMA(EMA))

    Args:
        series: Price series.
        span: EMA span.

    Returns:
        TEMA series.
    """
    ema1 = ema(series, span)
    ema2 = ema(ema1, span)
    ema3 = ema(ema2, span)
    return 3 * ema1 - 3 * ema2 + ema3


# =============================================================================
# Momentum Indicators
# =============================================================================


def rsi(series: pd.Series, window: int = 14) -> pd.Series:
    """Relative Strength Index.

    Args:
        series: Price series.
        window: Lookback window (default 14).

    Returns:
        RSI series (0-100).
    """
    delta = series.diff()
    gain = delta.where(delta > 0, 0)
    loss = (-delta).where(delta < 0, 0)

    avg_gain = gain.ewm(alpha=1 / window, min_periods=window, adjust=False).mean()
    avg_loss = loss.ewm(alpha=1 / window, min_periods=window, adjust=False).mean()

    rs = avg_gain / avg_loss
    return 100 - (100 / (1 + rs))


def stochastic(
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    k_window: int = 14,
    d_window: int = 3,
) -> tuple[pd.Series, pd.Series]:
    """Stochastic Oscillator.

    Args:
        high: High price series.
        low: Low price series.
        close: Close price series.
        k_window: %K window (default 14).
        d_window: %D smoothing window (default 3).

    Returns:
        Tuple of (%K, %D) series.
    """
    lowest_low = low.rolling(window=k_window).min()
    highest_high = high.rolling(window=k_window).max()

    k = 100 * (close - lowest_low) / (highest_high - lowest_low)
    d = k.rolling(window=d_window).mean()

    return k, d


def macd_indicator(
    series: pd.Series,
    fast: int = 12,
    slow: int = 26,
    signal: int = 9,
) -> tuple[pd.Series, pd.Series, pd.Series]:
    """MACD (Moving Average Convergence Divergence).

    Args:
        series: Price series.
        fast: Fast EMA period (default 12).
        slow: Slow EMA period (default 26).
        signal: Signal line period (default 9).

    Returns:
        Tuple of (MACD line, Signal line, Histogram).
    """
    fast_ema = ema(series, fast)
    slow_ema = ema(series, slow)

    macd_line = fast_ema - slow_ema
    signal_line = ema(macd_line, signal)
    histogram = macd_line - signal_line

    return macd_line, signal_line, histogram


def roc(series: pd.Series, window: int = 10) -> pd.Series:
    """Rate of Change.

    Args:
        series: Price series.
        window: Lookback window.

    Returns:
        ROC series (percentage change).
    """
    return (series - series.shift(window)) / series.shift(window) * 100


def momentum(series: pd.Series, window: int = 10) -> pd.Series:
    """Momentum (difference from n periods ago).

    Args:
        series: Price series.
        window: Lookback window.

    Returns:
        Momentum series.
    """
    return series - series.shift(window)


# =============================================================================
# Volatility Indicators
# =============================================================================


def bollinger_bands(
    series: pd.Series,
    window: int = 20,
    num_std: float = 2.0,
) -> tuple[pd.Series, pd.Series, pd.Series]:
    """Bollinger Bands.

    Args:
        series: Price series.
        window: SMA window (default 20).
        num_std: Number of standard deviations (default 2.0).

    Returns:
        Tuple of (Upper band, Middle band, Lower band).
    """
    middle = sma(series, window)
    std = series.rolling(window=window).std()

    upper = middle + (num_std * std)
    lower = middle - (num_std * std)

    return upper, middle, lower


def atr(
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    window: int = 14,
) -> pd.Series:
    """Average True Range.

    Args:
        high: High price series.
        low: Low price series.
        close: Close price series.
        window: Lookback window (default 14).

    Returns:
        ATR series.
    """
    prev_close = close.shift(1)

    tr1 = high - low
    tr2 = (high - prev_close).abs()
    tr3 = (low - prev_close).abs()

    true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

    return true_range.rolling(window=window).mean()


def keltner_channel(
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    ema_window: int = 20,
    atr_window: int = 10,
    atr_mult: float = 2.0,
) -> tuple[pd.Series, pd.Series, pd.Series]:
    """Keltner Channel.

    Args:
        high: High price series.
        low: Low price series.
        close: Close price series.
        ema_window: EMA window for middle band.
        atr_window: ATR window.
        atr_mult: ATR multiplier.

    Returns:
        Tuple of (Upper channel, Middle line, Lower channel).
    """
    middle = ema(close, ema_window)
    atr_val = atr(high, low, close, atr_window)

    upper = middle + (atr_mult * atr_val)
    lower = middle - (atr_mult * atr_val)

    return upper, middle, lower


def donchian_channel(
    high: pd.Series,
    low: pd.Series,
    window: int = 20,
) -> tuple[pd.Series, pd.Series, pd.Series]:
    """Donchian Channel.

    Args:
        high: High price series.
        low: Low price series.
        window: Lookback window.

    Returns:
        Tuple of (Upper channel, Middle, Lower channel).
    """
    upper = high.rolling(window=window).max()
    lower = low.rolling(window=window).min()
    middle = (upper + lower) / 2

    return upper, middle, lower


# =============================================================================
# Volume Indicators
# =============================================================================


def obv(close: pd.Series, volume: pd.Series) -> pd.Series:
    """On-Balance Volume.

    Args:
        close: Close price series.
        volume: Volume series.

    Returns:
        OBV series.
    """
    direction = close.diff().apply(lambda x: 1 if x > 0 else (-1 if x < 0 else 0))
    return (volume * direction).cumsum()


def vwap(
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    volume: pd.Series,
) -> pd.Series:
    """Volume Weighted Average Price (cumulative).

    Args:
        high: High price series.
        low: Low price series.
        close: Close price series.
        volume: Volume series.

    Returns:
        VWAP series.
    """
    typical_price = (high + low + close) / 3
    return (typical_price * volume).cumsum() / volume.cumsum()


# =============================================================================
# Other Indicators
# =============================================================================


def adx(
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    window: int = 14,
) -> tuple[pd.Series, pd.Series, pd.Series]:
    """Average Directional Index.

    Args:
        high: High price series.
        low: Low price series.
        close: Close price series.
        window: Lookback window (default 14).

    Returns:
        Tuple of (ADX, +DI, -DI).
    """
    # True Range
    prev_close = close.shift(1)
    tr1 = high - low
    tr2 = (high - prev_close).abs()
    tr3 = (low - prev_close).abs()
    true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

    # Directional Movement
    plus_dm = high.diff()
    minus_dm = -low.diff()

    plus_dm = plus_dm.where((plus_dm > minus_dm) & (plus_dm > 0), 0)
    minus_dm = minus_dm.where((minus_dm > plus_dm) & (minus_dm > 0), 0)

    # Smoothed
    tr_smooth = true_range.ewm(alpha=1 / window, adjust=False).mean()
    plus_dm_smooth = plus_dm.ewm(alpha=1 / window, adjust=False).mean()
    minus_dm_smooth = minus_dm.ewm(alpha=1 / window, adjust=False).mean()

    # Directional Indicators
    plus_di = 100 * plus_dm_smooth / tr_smooth
    minus_di = 100 * minus_dm_smooth / tr_smooth

    # DX and ADX
    dx = 100 * (plus_di - minus_di).abs() / (plus_di + minus_di)
    adx_val = dx.ewm(alpha=1 / window, adjust=False).mean()

    return adx_val, plus_di, minus_di


def cci(
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    window: int = 20,
) -> pd.Series:
    """Commodity Channel Index.

    Args:
        high: High price series.
        low: Low price series.
        close: Close price series.
        window: Lookback window.

    Returns:
        CCI series.
    """
    typical_price = (high + low + close) / 3
    sma_tp = sma(typical_price, window)

    # Mean Absolute Deviation
    mad = typical_price.rolling(window=window).apply(
        lambda x: np.mean(np.abs(x - np.mean(x))), raw=True
    )

    return (typical_price - sma_tp) / (0.015 * mad)


def williams_r(
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    window: int = 14,
) -> pd.Series:
    """Williams %R.

    Args:
        high: High price series.
        low: Low price series.
        close: Close price series.
        window: Lookback window.

    Returns:
        Williams %R series (-100 to 0).
    """
    highest_high = high.rolling(window=window).max()
    lowest_low = low.rolling(window=window).min()

    return -100 * (highest_high - close) / (highest_high - lowest_low)
