"""Drop-in replacement for the `talib.abstract` calls this project makes.

Same call signatures, same column names, same NaN warm-up lengths, so call
sites only change their import line.

Only what the codebase actually uses is implemented: SMA, MA, EMA, RSI, ATR,
MACD, BBANDS. Numerical agreement with TA-Lib is enforced by
`chanlun/tests/test_ta_compat.py`, which skips when talib is absent.

The warm-up conventions below are TA-Lib's, not the textbook ones, and they
differ between functions — getting them wrong shifts every downstream Chan
structure by a bar:

  SMA/BBANDS  first value at index period-1
  EMA         seeded with the SMA of the first `period` values, at index period-1
  RSI/ATR     Wilder smoothing (alpha = 1/period), first value at index period
"""

from __future__ import annotations

import numpy as np
import pandas as pd

__all__ = ["SMA", "MA", "EMA", "RSI", "ATR", "MACD", "BBANDS"]


def _series(data, price: str = "close") -> pd.Series:
    """Accept the abstract-API shapes: DataFrame, Series, or ndarray."""
    if isinstance(data, pd.DataFrame):
        return data[price].astype(float)
    if isinstance(data, pd.Series):
        return data.astype(float)
    return pd.Series(np.asarray(data, dtype=float))


def _recursive(values: np.ndarray, seed: float, start: int, alpha: float, n: int) -> np.ndarray:
    """out[start] = seed; out[i] = alpha*values[i] + (1-alpha)*out[i-1].

    Delegates the recursion to pandas' C implementation rather than a Python
    loop — `research/` runs this over long histories.
    """
    out = np.full(n, np.nan)
    if start >= n:
        return out
    tail = values[start:].astype(float).copy()
    tail[0] = seed
    out[start:] = pd.Series(tail).ewm(alpha=alpha, adjust=False).mean().to_numpy()
    return out


def SMA(data, timeperiod: int = 30, price: str = "close") -> pd.Series:
    s = _series(data, price)
    return s.rolling(window=timeperiod, min_periods=timeperiod).mean()


def MA(data, timeperiod: int = 30, matype: int = 0, price: str = "close") -> pd.Series:
    if matype != 0:
        raise NotImplementedError(f"MA matype={matype} is not used by this codebase")
    return SMA(data, timeperiod, price=price)


def _ema(x: np.ndarray, period: int, start: int) -> np.ndarray:
    """EMA whose first output lands on `start`, seeded by the SMA of the
    `period` values ending there.

    `start` is a parameter because MACD needs the fast EMA to begin later than
    it naturally would; see the note in MACD().
    """
    n = x.size
    if n <= start or start < period - 1:
        return np.full(n, np.nan)
    seed = x[start - period + 1: start + 1].mean()
    return _recursive(x, seed, start, 2.0 / (period + 1.0), n)


def EMA(data, timeperiod: int = 30, price: str = "close") -> pd.Series:
    s = _series(data, price)
    x = s.to_numpy(dtype=float)
    return pd.Series(_ema(x, timeperiod, timeperiod - 1), index=s.index)


def RSI(data, timeperiod: int = 14, price: str = "close") -> pd.Series:
    s = _series(data, price)
    x = s.to_numpy(dtype=float)
    n = x.size
    out = np.full(n, np.nan)
    if n <= timeperiod:
        return pd.Series(out, index=s.index)

    delta = np.diff(x)
    gain = np.where(delta > 0.0, delta, 0.0)
    loss = np.where(delta < 0.0, -delta, 0.0)

    # delta[k] corresponds to bar k+1, so the first `timeperiod` deltas seed bar `timeperiod`.
    alpha = 1.0 / timeperiod
    avg_gain = _recursive(gain, gain[:timeperiod].mean(), timeperiod - 1, alpha, n - 1)
    avg_loss = _recursive(loss, loss[:timeperiod].mean(), timeperiod - 1, alpha, n - 1)

    ag = avg_gain[timeperiod - 1:]
    al = avg_loss[timeperiod - 1:]
    with np.errstate(divide="ignore", invalid="ignore"):
        rsi = np.where(al == 0.0, 100.0, 100.0 - 100.0 / (1.0 + ag / al))
    out[timeperiod:] = rsi
    return pd.Series(out, index=s.index)


def ATR(data, timeperiod: int = 14) -> pd.Series:
    if not isinstance(data, pd.DataFrame):
        raise TypeError("ATR needs a DataFrame with high/low/close")
    high = data["high"].to_numpy(dtype=float)
    low = data["low"].to_numpy(dtype=float)
    close = data["close"].to_numpy(dtype=float)
    n = high.size
    out = np.full(n, np.nan)
    if n <= timeperiod:
        return pd.Series(out, index=data.index)

    prev_close = close[:-1]
    tr = np.maximum.reduce([
        high[1:] - low[1:],
        np.abs(high[1:] - prev_close),
        np.abs(low[1:] - prev_close),
    ])

    # tr[k] is bar k+1; the first `timeperiod` true ranges seed bar `timeperiod`.
    smoothed = _recursive(tr, tr[:timeperiod].mean(), timeperiod - 1, 1.0 / timeperiod, n - 1)
    out[timeperiod:] = smoothed[timeperiod - 1:]
    return pd.Series(out, index=data.index)


def MACD(
    data,
    fastperiod: int = 12,
    slowperiod: int = 26,
    signalperiod: int = 9,
    price: str = "close",
) -> pd.DataFrame:
    if slowperiod < fastperiod:
        fastperiod, slowperiod = slowperiod, fastperiod

    s = _series(data, price)
    x = s.to_numpy(dtype=float)
    n = x.size

    macd = np.full(n, np.nan)
    signal = np.full(n, np.nan)
    hist = np.full(n, np.nan)
    empty = pd.DataFrame({"macd": macd, "macdsignal": signal, "macdhist": hist}, index=s.index)

    # Both EMAs emit their first value on the same bar. That makes the slow one
    # ordinary, but re-seeds the fast one from the SMA of the `fastperiod`
    # values ending there instead of carrying the recursion forward from bar
    # fastperiod-1 — the two disagree by ~0.2 on a 100-priced series.
    macd_start = slowperiod - 1
    if n <= macd_start:
        return empty
    line = _ema(x, fastperiod, macd_start) - _ema(x, slowperiod, macd_start)

    # The signal EMA runs over the MACD line, so everything shifts by another
    # signalperiod-1 bars, and TA-Lib trims the MACD line to match.
    valid = line[macd_start:]
    if valid.size < signalperiod:
        return empty
    sig = _recursive(
        valid, valid[:signalperiod].mean(), signalperiod - 1, 2.0 / (signalperiod + 1.0), valid.size
    )
    start = macd_start + signalperiod - 1
    macd[start:] = line[start:]
    signal[macd_start:] = sig
    hist = macd - signal

    return pd.DataFrame({"macd": macd, "macdsignal": signal, "macdhist": hist}, index=s.index)


def BBANDS(
    data,
    timeperiod: int = 5,
    nbdevup: float = 2.0,
    nbdevdn: float = 2.0,
    matype: int = 0,
    price: str = "close",
) -> pd.DataFrame:
    if matype != 0:
        raise NotImplementedError(f"BBANDS matype={matype} is not used by this codebase")
    s = _series(data, price)
    middle = s.rolling(window=timeperiod, min_periods=timeperiod).mean()
    # TA-Lib uses the population standard deviation.
    std = s.rolling(window=timeperiod, min_periods=timeperiod).std(ddof=0)
    return pd.DataFrame(
        {
            "upperband": middle + nbdevup * std,
            "middleband": middle,
            "lowerband": middle - nbdevdn * std,
        },
        index=s.index,
    )