"""Momentum indicators: RSI, Stochastic."""

import pandas as pd
import numpy as np


def rsi(closes: pd.Series, period: int = 14) -> pd.Series:
    """RSI using Wilder's smoothing (EMA-based)."""
    delta = closes.diff()
    gain = delta.clip(lower=0)
    loss = -delta.clip(upper=0)
    avg_gain = gain.ewm(com=period - 1, min_periods=period).mean()
    avg_loss = loss.ewm(com=period - 1, min_periods=period).mean()
    rs = avg_gain / avg_loss
    rsi_vals = pd.Series(np.where(avg_loss == 0, 100.0, 100 - (100 / (1 + rs))), index=closes.index)
    # Keep NaN where we don't have enough data
    rsi_vals[:period] = np.nan
    return rsi_vals


def stochastic(
    highs: pd.Series,
    lows: pd.Series,
    closes: pd.Series,
    k_period: int = 14,
    d_period: int = 3,
) -> tuple[pd.Series, pd.Series]:
    """Stochastic Oscillator.

    Returns: (%K, %D)
    """
    lowest_low = lows.rolling(window=k_period).min()
    highest_high = highs.rolling(window=k_period).max()
    denom = highest_high - lowest_low
    k = 100 * (closes - lowest_low) / denom.replace(0, float("nan"))
    d = k.rolling(window=d_period).mean()
    return k, d