"""
Technical Analysis Indicator Functions

Standalone TA computation functions (rma, sma, ema, atr, rsi, dmi, adx, macd).
PineScript-compatible implementations using Wilder's RMA.

Usage:
    from sixtysix.ta import rma, sma, ema, atr, rsi, dmi, adx, macd
"""

from __future__ import annotations
import numpy as np
import pandas as pd


def rma(series: pd.Series, period: int) -> pd.Series:
    """
    Wilder's RMA — exact match for PineScript ta.rma().

    Seeds with SMA of first `period` values, then applies
    exponential smoothing with alpha=1/period. All other TA
    helpers (atr, rsi, adx) use this internally.
    """
    arr = series.values.astype(float)
    out = np.full_like(arr, np.nan)

    # Find first non-NaN
    first = 0
    while first < len(arr) and np.isnan(arr[first]):
        first += 1

    if first + period > len(arr):
        return pd.Series(out, index=series.index)

    # Seed with SMA of first `period` valid values
    rma_val = np.nanmean(arr[first:first + period])
    out[first + period - 1] = rma_val

    alpha = 1.0 / period
    for i in range(first + period, len(arr)):
        rma_val = (1 - alpha) * rma_val + alpha * arr[i]
        out[i] = rma_val

    return pd.Series(out, index=series.index)


def sma(series: pd.Series, period: int) -> pd.Series:
    """Simple Moving Average."""
    return series.rolling(period).mean()


def ema(series: pd.Series, period: int) -> pd.Series:
    """Exponential Moving Average (span-based, matches pandas default)."""
    return series.ewm(span=period, adjust=False).mean()


def atr(df: pd.DataFrame, period: int = 14) -> pd.Series:
    """
    Average True Range — exact match for PineScript ta.atr().

    Uses Wilder's RMA (SMA-seeded, alpha=1/period).
    """
    high = df['high']
    low = df['low']
    close = df['close'].shift(1)

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

    tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
    return rma(tr, period)


def rsi(df: pd.DataFrame, period: int = 14) -> pd.Series:
    """
    Relative Strength Index — exact match for PineScript ta.rsi().

    Uses Wilder's RMA (SMA-seeded) for gain/loss smoothing.
    """
    delta = df['close'].diff()
    gain = delta.where(delta > 0, 0.0)
    loss = -delta.where(delta < 0, 0.0)

    avg_gain = rma(gain, period)
    avg_loss = rma(loss, period)

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


def dmi(df: pd.DataFrame, di_period: int = 14, adx_period: int = 14):
    """
    Directional Movement Index — exact match for PineScript ta.dmi().

    Uses Wilder's RMA (SMA-seeded) for all smoothing.
    Returns (plus_di, minus_di, adx) like Pine's [diplus, diminus, adx] = ta.dmi(14, 14).
    """
    high = df['high']
    low = df['low']
    close = df['close']

    # True Range
    tr1 = high - low
    tr2 = abs(high - close.shift(1))
    tr3 = abs(low - close.shift(1))
    tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

    # Directional Movement
    up_move = high - high.shift(1)
    down_move = low.shift(1) - low

    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)

    # Wilder's smoothing (RMA)
    atr_smooth = rma(tr, di_period)
    smooth_plus_dm = rma(plus_dm, di_period)
    smooth_minus_dm = rma(minus_dm, di_period)

    # +DI and -DI
    plus_di = 100 * smooth_plus_dm / atr_smooth
    minus_di = 100 * smooth_minus_dm / atr_smooth

    # DX -> ADX
    dx = 100 * abs(plus_di - minus_di) / (plus_di + minus_di)
    adx_val = rma(dx, adx_period)

    return plus_di, minus_di, adx_val


def adx(df: pd.DataFrame, period: int = 14) -> pd.Series:
    """
    Average Directional Index — convenience wrapper around dmi().

    Returns ADX line only.
    """
    _, _, adx_val = dmi(df, di_period=period, adx_period=period)
    return adx_val


def macd(df: pd.DataFrame, fast: int = 12, slow: int = 26, signal: int = 9) -> pd.DataFrame:
    """
    MACD indicator.

    Returns DataFrame with columns: 'macd', 'signal', 'histogram'.
    """
    ema_fast = df['close'].ewm(span=fast, adjust=False).mean()
    ema_slow = df['close'].ewm(span=slow, adjust=False).mean()
    macd_line = ema_fast - ema_slow
    signal_line = macd_line.ewm(span=signal, adjust=False).mean()
    histogram = macd_line - signal_line
    return pd.DataFrame({
        'macd': macd_line,
        'signal': signal_line,
        'histogram': histogram,
    })
