"""
Technical Analysis Indicator Functions

Standalone TA computation functions for use in strategies and indicators.
PineScript-compatible implementations using Wilder's RMA where applicable.

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

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) -> tuple[pd.Series, pd.Series, pd.Series]:
    """
    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,
    })


# ── Bollinger Bands ──


def bbands(
    df: pd.DataFrame, period: int = 20, std_dev: float = 2.0,
) -> pd.DataFrame:
    """
    Bollinger Bands.

    Returns DataFrame with columns: 'upper', 'middle', 'lower', 'bandwidth', 'pctb'.

    Args:
        df: DataFrame with 'close' column
        period: SMA lookback period
        std_dev: Standard deviation multiplier for bands

    Example:
        bb = ta.bbands(df, 20, 2.0)
        bb['upper']      # Upper band
        bb['lower']      # Lower band
        bb['pctb']       # %B: position within bands (0-1)
        bb['bandwidth']  # Band width as % of middle
    """
    close = df['close']
    middle = close.rolling(period).mean()
    std = close.rolling(period).std()
    upper = middle + std_dev * std
    lower = middle - std_dev * std
    bandwidth = ((upper - lower) / middle) * 100
    pctb = (close - lower) / (upper - lower)

    return pd.DataFrame({
        'upper': upper,
        'middle': middle,
        'lower': lower,
        'bandwidth': bandwidth,
        'pctb': pctb,
    })


# ── Supertrend ──


def supertrend(
    df: pd.DataFrame, period: int = 10, multiplier: float = 3.0,
) -> pd.DataFrame:
    """
    Supertrend indicator.

    Returns DataFrame with columns: 'supertrend', 'direction'.
    Direction: 1 = uptrend (price above supertrend), -1 = downtrend.

    Args:
        df: DataFrame with 'high', 'low', 'close' columns
        period: ATR period
        multiplier: ATR multiplier for band width

    Example:
        st = ta.supertrend(df, 10, 3.0)
        st['direction']   # 1 = bullish, -1 = bearish
        st['supertrend']  # Support/resistance line
    """
    high = df['high'].values
    low = df['low'].values
    close = df['close'].values
    n = len(close)

    atr_series = atr(df, period)
    atr_vals = atr_series.values

    mid = (high + low) / 2.0
    upper = mid + multiplier * atr_vals
    lower = mid - multiplier * atr_vals

    direction = np.ones(n)
    st_line = np.full(n, np.nan)

    # Find first valid ATR bar
    first = 0
    while first < n and np.isnan(atr_vals[first]):
        first += 1
    if first >= n:
        return pd.DataFrame({
            'supertrend': pd.Series(st_line, index=df.index),
            'direction': pd.Series(direction, index=df.index),
        })

    direction[first] = 1
    st_line[first] = lower[first]

    for i in range(first + 1, n):
        # Tighten bands: only move in the trend direction
        if close[i - 1] > upper[i - 1]:
            upper[i] = upper[i]  # reset
        else:
            upper[i] = min(upper[i], upper[i - 1])

        if close[i - 1] < lower[i - 1]:
            lower[i] = lower[i]  # reset
        else:
            lower[i] = max(lower[i], lower[i - 1])

        # Flip direction on close crossing band
        if direction[i - 1] == 1:
            if close[i] < lower[i]:
                direction[i] = -1
            else:
                direction[i] = 1
        else:
            if close[i] > upper[i]:
                direction[i] = 1
            else:
                direction[i] = -1

        st_line[i] = lower[i] if direction[i] == 1 else upper[i]

    return pd.DataFrame({
        'supertrend': pd.Series(st_line, index=df.index),
        'direction': pd.Series(direction, index=df.index),
    })


# ── Renko ──


def renko(
    df: pd.DataFrame,
    brick_size: float | None = None,
    atr_period: int = 14,
    method: str = 'atr',
    pct: float | None = None,
) -> pd.DataFrame:
    """
    Renko overlay — maps renko bricks back to the bar chart.

    Returns DataFrame with columns: 'trend', 'top', 'bottom'.
    Trend: 1 = up brick, -1 = down brick, 0 = no brick yet.
    Top/bottom: current brick boundaries.

    Args:
        df: DataFrame with 'high', 'low', 'close' columns
        brick_size: Fixed brick size in price units (used when method='fixed')
        atr_period: ATR period for dynamic brick sizing (used when method='atr')
        method: 'atr' | 'fixed' | 'pct'
                'atr': ATR-based dynamic bricks (updates every 20 bars)
                'fixed': constant size in price units (brick_size required)
                'pct': percentage of current price (pct required)
        pct: Brick size as percentage of close (used when method='pct').
             E.g., 0.5 = 0.5% of close price per brick.

    Example:
        # ATR-based (adapts to volatility)
        r = ta.renko(df, atr_period=14, method='atr')

        # Fixed price amount
        r = ta.renko(df, brick_size=2.0, method='fixed')

        # Percentage of price (0.5% per brick)
        r = ta.renko(df, pct=0.5, method='pct')

        r['trend']   # 1 = up, -1 = down, 0 = no brick
        r['top']     # Top of current brick
        r['bottom']  # Bottom of current brick
    """
    close = df['close'].values
    n = len(close)

    trend = np.zeros(n)
    top = np.full(n, np.nan)
    bottom = np.full(n, np.nan)

    _empty = pd.DataFrame({
        'trend': pd.Series(trend, index=df.index),
        'top': pd.Series(top, index=df.index),
        'bottom': pd.Series(bottom, index=df.index),
    })

    if n < 20:
        return _empty

    if method == 'pct':
        if pct is None or pct <= 0:
            raise ValueError("pct must be positive when method='pct'")
        atr_vals = None
        start = 0
        bs = close[0] * pct / 100.0
    elif method == 'fixed':
        if brick_size is None or brick_size <= 0:
            raise ValueError("brick_size must be positive when method='fixed'")
        atr_vals = None
        start = 0
        bs = brick_size
    else:  # 'atr'
        atr_vals = atr(df, atr_period).values
        # Find first valid ATR for seeding
        seed = atr_period
        while seed < n and (np.isnan(atr_vals[seed]) or atr_vals[seed] <= 0):
            seed += 1
        if seed >= n:
            return _empty
        start = seed
        bs = atr_vals[seed]

    # Seed first brick
    b_bottom = close[start]
    b_top = b_bottom + bs
    cur_trend = 0

    for i in range(start + 1, n):
        # Update brick size dynamically
        if method == 'atr' and i % 20 == 0 and atr_vals[i] > 0 and not np.isnan(atr_vals[i]):
            bs = atr_vals[i]
        elif method == 'pct':
            bs = close[i] * pct / 100.0

        if bs < 1e-10:
            bs = close[i] * 0.01

        if close[i] >= b_top + bs:
            num = int((close[i] - b_top) / bs)
            b_bottom = b_top + (num - 1) * bs
            b_top = b_bottom + bs
            cur_trend = 1
        elif close[i] <= b_bottom - bs:
            num = int((b_bottom - close[i]) / bs)
            b_top = b_bottom - (num - 1) * bs
            b_bottom = b_top - bs
            cur_trend = -1

        trend[i] = cur_trend
        top[i] = b_top
        bottom[i] = b_bottom

    return pd.DataFrame({
        'trend': pd.Series(trend, index=df.index),
        'top': pd.Series(top, index=df.index),
        'bottom': pd.Series(bottom, index=df.index),
    })
