"""
动量指标 — RSI、Stochastic

所有函数返回与输入等长的 list[float]，不足周期位置填 0.0。
"""


def rsi(data: list[float], period: int = 14) -> list[float]:
    """相对强弱指数 (RSI)

    使用 Wilder 平滑算法，Wilder's RSI = 100 - [100 / (1 + avg_gain / avg_loss)]

    Args:
        data: 价格序列
        period: 周期（默认 14）

    Returns:
        与 data 等长的 RSI 序列 [0, 100]，前 period 位置为 0
    """
    n = len(data)
    result = [0.0] * n
    if n < period + 1:
        return result

    # 计算价格变化
    changes = [data[i] - data[i - 1] for i in range(1, n)]

    # 初始平均涨幅和跌幅（Simple average of first `period` changes）
    gains = [max(c, 0) for c in changes[:period]]
    losses = [abs(min(c, 0)) for c in changes[:period]]
    avg_gain = sum(gains) / period
    avg_loss = sum(losses) / period

    # 计算第一个 RSI
    if avg_loss == 0:
        result[period] = 100.0
    else:
        rs = avg_gain / avg_loss
        result[period] = 100.0 - (100.0 / (1.0 + rs))

    # Wilder 平滑后续值
    for i in range(period, n - 1):
        change = changes[i]
        gain = max(change, 0.0)
        loss = abs(min(change, 0.0))

        avg_gain = (avg_gain * (period - 1) + gain) / period
        avg_loss = (avg_loss * (period - 1) + loss) / period

        if avg_loss == 0:
            result[i + 1] = 100.0
        else:
            rs = avg_gain / avg_loss
            result[i + 1] = 100.0 - (100.0 / (1.0 + rs))

    return result


def stoch(
    high: list[float],
    low: list[float],
    close: list[float],
    k_period: int = 14,
    k_smooth: int = 3,
    d_smooth: int = 3,
):
    """Stochastic 指标 (KDJ 中的 K/D)

    %K = 100 * (close - lowest_low) / (highest_high - lowest_low)
    %K_smoothed = SMA(%K, k_smooth)
    %D = SMA(%K_smoothed, d_smooth)

    Args:
        high: 最高价序列
        low: 最低价序列
        close: 收盘价序列
        k_period: %K 窗口
        k_smooth: %K 平滑周期
        d_smooth: %D 平滑周期

    Returns:
        (k_values, d_values) 两个等长序列，范围 [0, 100]
    """
    n = len(close)
    k_raw = [0.0] * n
    k_values = [0.0] * n
    d_values = [0.0] * n

    if n < k_period:
        return k_values, d_values

    # 计算原始 %K
    for i in range(k_period - 1, n):
        highest = max(high[i - k_period + 1 : i + 1])
        lowest = min(low[i - k_period + 1 : i + 1])
        if highest != lowest:
            k_raw[i] = 100.0 * (close[i] - lowest) / (highest - lowest)
        else:
            k_raw[i] = 50.0

    # 平滑 %K
    from .trend import sma as _sma
    k_smoothed = _sma(k_raw, k_smooth)
    d_smoothed = _sma(k_smoothed, d_smooth)

    return k_smoothed, d_smoothed


def stoch_k(
    high: list[float],
    low: list[float],
    close: list[float],
    k_period: int = 14,
    k_smooth: int = 3,
) -> list[float]:
    """Stochastic %K"""
    k, _ = stoch(high, low, close, k_period, k_smooth)
    return k


def stoch_d(
    high: list[float],
    low: list[float],
    close: list[float],
    k_period: int = 14,
    k_smooth: int = 3,
    d_smooth: int = 3,
) -> list[float]:
    """Stochastic %D"""
    _, d = stoch(high, low, close, k_period, k_smooth, d_smooth)
    return d