data_processor.py

"""데이터 정리 및 기술적 지표 계산 (ta 라이브러리)"""

import pandas as pd
import numpy as np
from ta.trend import SMAIndicator, EMAIndicator, MACD
from ta.momentum import RSIIndicator, StochasticOscillator, ROCIndicator, WilliamsRIndicator
from ta.volatility import BollingerBands, AverageTrueRange
from ta.volume import OnBalanceVolumeIndicator

import config


def clean_data(df: pd.DataFrame) -> pd.DataFrame:
    """결측치 처리, 정렬, 타입 검증"""
    df = df.copy()
    df = df.sort_index()
    # 전체 NaN 행 제거
    df = df.dropna(how="all")
    # 소규모 갭 forward-fill (최대 3일)
    df = df.ffill(limit=3)
    # 숫자 타입 보장
    for col in ["Open", "High", "Low", "Close", "Volume"]:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")
    return df


def add_moving_averages(df: pd.DataFrame) -> pd.DataFrame:
    """SMA, EMA 이동평균선 추가"""
    close = df["Close"]
    df["SMA_50"] = SMAIndicator(close, window=config.SMA_SHORT).sma_indicator()
    df["SMA_200"] = SMAIndicator(close, window=config.SMA_LONG).sma_indicator()
    df["EMA_12"] = EMAIndicator(close, window=config.EMA_SHORT).ema_indicator()
    df["EMA_26"] = EMAIndicator(close, window=config.EMA_LONG).ema_indicator()
    return df


def add_rsi(df: pd.DataFrame) -> pd.DataFrame:
    """RSI 상대강도지수 추가"""
    df["RSI"] = RSIIndicator(df["Close"], window=config.RSI_WINDOW).rsi()
    return df


def add_macd(df: pd.DataFrame) -> pd.DataFrame:
    """MACD, Signal, Histogram 추가"""
    macd = MACD(
        df["Close"],
        window_fast=config.MACD_FAST,
        window_slow=config.MACD_SLOW,
        window_sign=config.MACD_SIGNAL,
    )
    df["MACD"] = macd.macd()
    df["MACD_Signal"] = macd.macd_signal()
    df["MACD_Hist"] = macd.macd_diff()
    return df


def add_bollinger_bands(df: pd.DataFrame) -> pd.DataFrame:
    """볼린저밴드 상/중/하단, 밴드폭, %B 추가"""
    bb = BollingerBands(df["Close"], window=config.BB_WINDOW, window_dev=config.BB_STD)
    df["BB_Upper"] = bb.bollinger_hband()
    df["BB_Middle"] = bb.bollinger_mavg()
    df["BB_Lower"] = bb.bollinger_lband()
    df["BB_Width"] = bb.bollinger_wband()
    df["BB_PctB"] = bb.bollinger_pband()
    return df


def add_atr(df: pd.DataFrame) -> pd.DataFrame:
    """ATR 평균진폭 추가"""
    df["ATR"] = AverageTrueRange(
        df["High"], df["Low"], df["Close"], window=config.ATR_WINDOW
    ).average_true_range()
    return df


def add_momentum_indicators(df: pd.DataFrame) -> pd.DataFrame:
    """스토캐스틱, ROC, 윌리엄스 %R 추가"""
    stoch = StochasticOscillator(
        df["High"], df["Low"], df["Close"], window=config.STOCH_WINDOW
    )
    df["Stoch_K"] = stoch.stoch()
    df["Stoch_D"] = stoch.stoch_signal()

    df["ROC"] = ROCIndicator(df["Close"], window=config.ROC_WINDOW).roc()

    df["Williams_R"] = WilliamsRIndicator(
        df["High"], df["Low"], df["Close"], lbp=config.STOCH_WINDOW
    ).williams_r()
    return df


def add_volume_indicators(df: pd.DataFrame) -> pd.DataFrame:
    """OBV, 거래량 비율 추가"""
    df["OBV"] = OnBalanceVolumeIndicator(df["Close"], df["Volume"]).on_balance_volume()
    df["Volume_SMA"] = df["Volume"].rolling(window=config.VOLUME_SMA_WINDOW).mean()
    df["Volume_Ratio"] = df["Volume"] / df["Volume_SMA"]
    return df


def add_all_features(df: pd.DataFrame) -> pd.DataFrame:
    """모든 기술적 지표를 계산하고 warm-up NaN 제거"""
    print("[데이터 가공] 기술적 지표 계산 중...")
    df = clean_data(df)
    df = add_moving_averages(df)
    df = add_rsi(df)
    df = add_macd(df)
    df = add_bollinger_bands(df)
    df = add_atr(df)
    df = add_momentum_indicators(df)
    df = add_volume_indicators(df)

    # warm-up 기간 NaN 제거 (SMA_200이 가장 긴 윈도우)
    rows_before = len(df)
    df = df.dropna()
    rows_dropped = rows_before - len(df)
    print(f"[데이터 가공] 완료: {len(df)}일 데이터, {df.shape[1]}개 컬럼 (warm-up {rows_dropped}일 제거)")
    return df