"""
波动率指标 — 布林带、ATR

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

import math


def bollinger(
    data: list[float],
    period: int = 20,
    std: float = 2.0,
):
    """布林带 (Bollinger Bands)

    使用流式计算方差，O(n) 复杂度。

    Args:
        data: 价格序列（通常为收盘价）
        period: 中轨 SMA 周期
        std: 标准差倍数

    Returns:
        (upper, mid, lower) 三个等长序列
    """
    n = len(data)
    upper = [0.0] * n
    mid = [0.0] * n
    lower = [0.0] * n

    if n < period:
        return upper, mid, lower

    # 初始窗口的 sum 和 sum_sq
    window_sum = 0.0
    window_sum_sq = 0.0
    for i in range(period):
        v = data[i]
        window_sum += v
        window_sum_sq += v * v

    # 第一个点
    mean = window_sum / period
    mid[period - 1] = mean
    variance = (window_sum_sq / period) - (mean * mean)
    stdev = math.sqrt(max(variance, 0.0))
    upper[period - 1] = mean + std * stdev
    lower[period - 1] = mean - std * stdev

    # 滑动窗口计算后续点
    for i in range(period, n):
        old_val = data[i - period]
        new_val = data[i]
        window_sum += new_val - old_val
        window_sum_sq += new_val * new_val - old_val * old_val

        mean = window_sum / period
        mid[i] = mean
        variance = (window_sum_sq / period) - (mean * mean)
        stdev = math.sqrt(max(variance, 0.0))
        upper[i] = mean + std * stdev
        lower[i] = mean - std * stdev

    return upper, mid, lower


def bollinger_upper(data: list[float], period: int = 20, std: float = 2.0) -> list[float]:
    """布林带上轨"""
    upper, _, _ = bollinger(data, period, std)
    return upper


def bollinger_mid(data: list[float], period: int = 20) -> list[float]:
    """布林带中轨"""
    from .trend import sma as _sma
    return _sma(data, period)


def bollinger_lower(data: list[float], period: int = 20, std: float = 2.0) -> list[float]:
    """布林带下轨"""
    _, _, lower = bollinger(data, period, std)
    return lower


def atr(
    high: list[float],
    low: list[float],
    close: list[float],
    period: int = 14,
) -> list[float]:
    """平均真实波幅 (ATR)

    使用 Wilder 平滑算法。

    Args:
        high: 最高价序列
        low: 最低价序列
        close: 收盘价序列
        period: 周期

    Returns:
        与输入等长的 ATR 序列，前 period 位置为 0
    """
    n = len(close)
    result = [0.0] * n
    if n < period + 1:
        return result

    # 计算 True Range
    tr = [0.0] * n
    tr[0] = high[0] - low[0]
    for i in range(1, n):
        tr[i] = max(
            high[i] - low[i],
            abs(high[i] - close[i - 1]),
            abs(low[i] - close[i - 1]),
        )

    # 初始 ATR 为前 period 个 TR 的均值
    result[period] = sum(tr[1:period + 1]) / period

    # Wilder 平滑
    for i in range(period + 1, n):
        result[i] = (result[i - 1] * (period - 1) + tr[i]) / period

    return result