"""
볼린저 밴드·일목균형표 계산 (모든 봉 간격 공용).
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from config import BB_PERIOD, BB_STD


def add_bollinger(
    df: pd.DataFrame,
    period: int = BB_PERIOD,
    std_mult: float = BB_STD,
) -> pd.DataFrame:
    """MA, Upper, Lower, BB_Width, bb_pos 컬럼 추가."""
    out = df.copy()
    if "MA" not in out.columns:
        out["MA"] = out["Close"].rolling(period).mean()
    if "Upper" not in out.columns or "Lower" not in out.columns:
        std = out["Close"].rolling(period).std()
        out["STD"] = std
        out["Upper"] = out["MA"] + std_mult * std
        out["Lower"] = out["MA"] - std_mult * std
    ma = out["MA"].replace(0, np.nan)
    band = (out["Upper"] - out["Lower"]).replace(0, np.nan)
    out["bb_pos"] = ((out["Close"] - out["Lower"]) / band).clip(0, 1)
    out["BB_Width"] = band / ma * 100
    return out


def add_ichimoku(
    df: pd.DataFrame,
    tenkan: int = 9,
    kijun: int = 26,
    senkou_b_period: int = 52,
) -> pd.DataFrame:
    """
    일목균형표 라인·구름 위치 컬럼 추가 (해당 봉 시점, 미래 데이터 미사용).

    Returns:
        ichi_tenkan, ichi_kijun, ichi_span_a, ichi_span_b,
        ichi_cloud_top, ichi_cloud_bottom
    """
    out = df.copy()
    h = out["High"].astype(float)
    l = out["Low"].astype(float)
    c = out["Close"].astype(float)

    out["ichi_tenkan"] = (h.rolling(tenkan).max() + l.rolling(tenkan).min()) / 2
    out["ichi_kijun"] = (h.rolling(kijun).max() + l.rolling(kijun).min()) / 2
    out["ichi_span_a"] = (out["ichi_tenkan"] + out["ichi_kijun"]) / 2
    out["ichi_span_b"] = (h.rolling(senkou_b_period).max() + l.rolling(senkou_b_period).min()) / 2
    out["ichi_cloud_top"] = np.maximum(out["ichi_span_a"], out["ichi_span_b"])
    out["ichi_cloud_bottom"] = np.minimum(out["ichi_span_a"], out["ichi_span_b"])
    return out