auto-stock/backend/app/services/ml_retrain_pipeline.py

"""
ML 자가 재학습 파이프라인

주기적(매주 일요일 02:00 KST) 또는 성능 저하 감지 시 LightGBM 모델을 자동 재학습합니다.

기능:
  1. 재학습 트리거: 주기적(매주 일요일 02:00) + 성능 저하 감지(최근 20건 승률 vs 전체 평균)
  2. 학습 데이터: stock_daily_prices + trade_history 기반 피처 엔지니어링
  3. 모델: LightGBM (rolling 6개월 윈도우, Train/Val 80/20 split)
  4. 버전 관리: ~/.cache/auto-stock/models/ 에 model_YYYYMMDD.pkl 저장, 최근 5개 유지
  5. 성능 모니터링: validation accuracy, precision, recall 로깅
  6. 스케줄러 통합: main.py 에서 ml_retrain_scheduler() 호출

ARCHITECTURE:
  - DB 읽기는 AsyncSessionLocal 직접 사용 (SELECT only)
  - DB 쓰기가 필요하면 db_task_queue 경유 (현재는 쓰기 없음)
  - 무거운 학습은 asyncio.to_thread 로 워커 스레드 실행
  - from app.core.config import settings 싱글톤 사용
"""

import asyncio
import logging
import os
import shutil
from datetime import datetime, date, timedelta
from pathlib import Path
from typing import Dict, Any, Optional, Tuple, List

import numpy as np
import pandas as pd
import joblib
import pytz

from app.core.config import settings
from app.core.database import AsyncSessionLocal
from app.services.feature_engineering import (
    extract_features,
    db_data_to_dataframe,
    FEATURE_COLUMNS,
)

logger = logging.getLogger(__name__)

# ─── 상수 ─────────────────────────────────────────────────────────────────
KST = pytz.timezone("Asia/Seoul")

# 모델 저장 디렉토리
MODEL_BASE_DIR = Path(os.path.expanduser("~")) / ".cache" / "auto-stock" / "models"

# 재학습 파이프라인 피처 (요구사항 기반 7개 — 기존 feature_engineering 20개와 호환)
RETRAIN_FEATURE_COLUMNS = [
    "rsi_14",       # RSI(14)
    "bb_position",  # BB %B(20)
    "atr_ratio",    # ATR / close
    "volume_ratio", # vol / avg20
    "ma5_ma20_gap", # MA5/MA20 이격
    "trend_strength",  # ADX 근사치 (feature_engineering 의 trend_strength)
    "macd",         # MACD
]

# 라벨링 기준
LABEL_FORECAST_DAYS = 5
LABEL_BUY_THRESHOLD = 0.01   # +1%
LABEL_SELL_THRESHOLD = -0.01  # -1%

# 성능 저하 감지 기준
PERF_DEGRADE_LOOKBACK = 20  # 최근 N건
PERF_DEGRADE_THRESHOLD = 0.15  # 15%p 하락

# Rolling window
ROLLING_WINDOW_MONTHS = 6

# 모델 버전 유지 개수
MAX_MODEL_VERSIONS = 5


# ─── 피처 엔지니어링 (재학습 전용) ─────────────────────────────────────────

def _compute_retrain_features(df: pd.DataFrame) -> pd.DataFrame:
    """stock_daily_prices DataFrame 에서 재학습용 피처를 계산합니다.

    기존 feature_engineering.extract_features() 를 기반으로 하되,
    추가 피처(atr_ratio, ma5_ma20_gap)를 포함합니다.

    Args:
        df: columns=[date, open, high, low, close, volume], 시간순 정렬

    Returns:
        동일 인덱스의 피처 DataFrame (RETRAIN_FEATURE_COLUMNS 컬럼)
    """
    # 기존 20개 피처 계산
    base_features = extract_features(df)

    close = df["close"].values.astype(float)
    high = df["high"].values.astype(float)
    low = df["low"].values.astype(float)
    n = len(close)

    # ATR(14) / close 비율
    atr_ratio = np.full(n, np.nan, dtype=float)
    tr = np.zeros(n, dtype=float)
    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]),
        )
    for i in range(14, n):
        atr_14 = np.mean(tr[i - 13 : i + 1])
        if close[i] > 0:
            atr_ratio[i] = atr_14 / close[i]

    # MA5 / MA20 이격
    ma5_ma20_gap = np.full(n, np.nan, dtype=float)
    for i in range(19, n):
        ma5 = np.mean(close[i - 4 : i + 1])
        ma20 = np.mean(close[i - 19 : i + 1])
        if ma20 > 0:
            ma5_ma20_gap[i] = (ma5 - ma20) / ma20

    result = pd.DataFrame(index=df.index)
    result["rsi_14"] = base_features["rsi_14"]
    result["bb_position"] = base_features["bb_position"]
    result["atr_ratio"] = atr_ratio
    result["volume_ratio"] = base_features["volume_ratio"]
    result["ma5_ma20_gap"] = ma5_ma20_gap
    result["trend_strength"] = base_features["trend_strength"]
    result["macd"] = base_features["macd"]

    return result


def _prepare_retrain_labels(close: np.ndarray) -> np.ndarray:
    """5일 후 수익률 기반 라벨 생성.

    Returns:
        labels: (N,) array, 0=SELL, 1=HOLD, 2=BUY, NaN=미확정
    """
    n = len(close)
    labels = np.full(n, np.nan, dtype=float)
    for i in range(n - LABEL_FORECAST_DAYS):
        if close[i] > 0:
            ret = (close[i + LABEL_FORECAST_DAYS] - close[i]) / close[i]
            if ret > LABEL_BUY_THRESHOLD:
                labels[i] = 2  # BUY
            elif ret < LABEL_SELL_THRESHOLD:
                labels[i] = 0  # SELL
            else:
                labels[i] = 1  # HOLD
    return labels


def _build_training_dataset(
    all_features: pd.DataFrame,
    labels: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray]:
    """NaN 이 아닌 유효 행만 추출하여 (X, y) 반환.

    Returns:
        X: (N_valid, len(RETRAIN_FEATURE_COLUMNS))
        y: (N_valid,)
    """
    feat_arr = all_features[RETRAIN_FEATURE_COLUMNS].values
    n = len(feat_arr)

    valid_mask = np.ones(n, dtype=bool)
    valid_mask[np.isnan(labels)] = False
    for i in range(n):
        if valid_mask[i] and np.any(np.isnan(feat_arr[i])):
            valid_mask[i] = False

    valid_idx = np.where(valid_mask)[0]
    if len(valid_idx) == 0:
        return np.array([]), np.array([])

    X = feat_arr[valid_idx]
    y = labels[valid_idx].astype(int)
    return X, y


# ─── 모델 학습 ─────────────────────────────────────────────────────────────

def _train_lightgbm(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: np.ndarray,
    y_val: np.ndarray,
) -> Tuple[Any, Dict[str, Any]]:
    """LightGBM 모델 학습 및 검증 메트릭 반환.

    Returns:
        (model, metrics_dict)
    """
    import lightgbm as lgb
    from sklearn.metrics import accuracy_score, precision_score, recall_score

    model = lgb.LGBMClassifier(
        n_estimators=200,
        max_depth=6,
        learning_rate=0.05,
        num_leaves=31,
        min_child_samples=10,
        subsample=0.8,
        colsample_bytree=0.8,
        reg_alpha=0.1,
        reg_lambda=0.1,
        random_state=42,
        verbose=-1,
        num_class=3,
        objective="multiclass",
    )

    model.fit(
        X_train, y_train,
        eval_set=[(X_val, y_val)],
    )

    y_pred = model.predict(X_val)

    accuracy = float(accuracy_score(y_val, y_pred))
    precision = float(precision_score(y_val, y_pred, average="macro", zero_division=0))
    recall = float(recall_score(y_val, y_pred, average="macro", zero_division=0))

    # 피처 중요도
    importance = model.feature_importances_
    top_features = sorted(
        zip(RETRAIN_FEATURE_COLUMNS, importance),
        key=lambda x: x[1],
        reverse=True,
    )

    class_distribution = {
        "train": {int(k): int(v) for k, v in zip(*np.unique(y_train, return_counts=True))},
        "val": {int(k): int(v) for k, v in zip(*np.unique(y_val, return_counts=True))},
    }

    metrics = {
        "accuracy": round(accuracy, 4),
        "precision": round(precision, 4),
        "recall": round(recall, 4),
        "val_samples": len(y_val),
        "train_samples": len(y_train),
        "class_distribution": class_distribution,
        "feature_importance": [
            {"name": name, "importance": int(imp)} for name, imp in top_features
        ],
    }

    return model, metrics


# ─── 모델 버전 관리 ────────────────────────────────────────────────────────

def _save_model_version(
    model: Any,
    scaler: Any,
    metrics: Dict[str, Any],
    model_dir: Path,
) -> Path:
    """모델을 날짜 기반 버전으로 저장하고 최근 MAX_MODEL_VERSIONS 개만 유지.

    Returns:
        저장된 모델 파일 경로
    """
    model_dir.mkdir(parents=True, exist_ok=True)

    today_str = datetime.now(KST).strftime("%Y%m%d")
    version_dir = model_dir / f"v_{today_str}"
    version_dir.mkdir(parents=True, exist_ok=True)

    model_path = version_dir / f"model_{today_str}.pkl"
    scaler_path = version_dir / "scaler.pkl"
    metrics_path = version_dir / "metrics.json"

    joblib.dump(model, model_path)
    if scaler is not None:
        joblib.dump(scaler, scaler_path)

    from app.core.json_helpers import _dumps
    with open(metrics_path, "w") as f:
        f.write(_dumps(metrics))

    # 최신 심볼릭 (복사) — latest 디렉토리
    latest_dir = model_dir / "latest"
    if latest_dir.exists():
        shutil.rmtree(latest_dir)
    shutil.copytree(version_dir, latest_dir)

    # 오래된 버전 정리
    _cleanup_old_versions(model_dir)

    logger.info(f"모델 저장 완료: {version_dir}")
    return model_path


def _cleanup_old_versions(model_dir: Path):
    """MAX_MODEL_VERSIONS 개를 초과하는 오래된 버전 디렉토리를 삭제합니다."""
    version_dirs = sorted(
        [d for d in model_dir.iterdir() if d.is_dir() and d.name.startswith("v_")],
        key=lambda d: d.name,
        reverse=True,
    )
    for old_dir in version_dirs[MAX_MODEL_VERSIONS:]:
        shutil.rmtree(old_dir, ignore_errors=True)
        logger.info(f"오래된 모델 버전 삭제: {old_dir.name}")


def load_latest_model(
    model_dir: Optional[Path] = None,
) -> Tuple[Optional[Any], Optional[Any], Optional[Dict[str, Any]]]:
    """최신 모델, scaler, metrics 를 로드합니다.

    Returns:
        (model, scaler, metrics) 또는 (None, None, None)
    """
    if model_dir is None:
        model_dir = MODEL_BASE_DIR
    latest_dir = model_dir / "latest"
    if not latest_dir.exists():
        return None, None, None

    try:
        # 모델 파일 찾기
        model_files = list(latest_dir.glob("model_*.pkl"))
        if not model_files:
            return None, None, None

        model = joblib.load(model_files[0])

        scaler = None
        scaler_path = latest_dir / "scaler.pkl"
        if scaler_path.exists():
            scaler = joblib.load(scaler_path)

        metrics = None
        metrics_path = latest_dir / "metrics.json"
        if metrics_path.exists():
            import json
            with open(metrics_path) as f:
                metrics = json.load(f)

        return model, scaler, metrics
    except Exception as e:
        logger.error(f"모델 로드 실패: {e}")
        return None, None, None


# ─── DB 데이터 로드 (읽기 전용) ────────────────────────────────────────────

async def _load_daily_prices(
    months: int = ROLLING_WINDOW_MONTHS,
) -> pd.DataFrame:
    """stock_daily_prices 에서 최근 N개월 데이터를 로드합니다.

    Returns:
        DataFrame[stock_code, date, open, high, low, close, volume]
    """
    from sqlalchemy import select, and_
    from app.models.models import StockDailyPrice

    cutoff = date.today() - timedelta(days=months * 30)

    async with AsyncSessionLocal() as session:
        stmt = (
            select(StockDailyPrice)
            .where(StockDailyPrice.trade_date >= cutoff)
            .order_by(StockDailyPrice.stock_code, StockDailyPrice.trade_date)
        )
        result = await session.execute(stmt)
        rows = result.scalars().all()

    if not rows:
        return pd.DataFrame()

    records = []
    for row in rows:
        records.append({
            "stock_code": row.stock_code,
            "date": str(row.trade_date),
            "open": float(row.open_price) if row.open_price else 0,
            "high": float(row.high_price) if row.high_price else 0,
            "low": float(row.low_price) if row.low_price else 0,
            "close": float(row.close_price) if row.close_price else 0,
            "volume": float(row.volume) if row.volume else 0,
        })
    return pd.DataFrame(records)


async def _load_trade_history(
    months: int = ROLLING_WINDOW_MONTHS,
) -> pd.DataFrame:
    """trade_history 에서 최근 N개월 체결 매매 이력을 로드합니다.

    Returns:
        DataFrame[stock_code, order_type, executed_price, profit_loss_rate, created_at]
    """
    from sqlalchemy import select, and_
    from app.models.models import TradeHistory

    cutoff = datetime.now(KST) - timedelta(days=months * 30)

    async with AsyncSessionLocal() as session:
        stmt = (
            select(TradeHistory)
            .where(
                and_(
                    TradeHistory.status == "EXECUTED",
                    TradeHistory.created_at >= cutoff,
                )
            )
            .order_by(TradeHistory.created_at)
        )
        result = await session.execute(stmt)
        rows = result.scalars().all()

    if not rows:
        return pd.DataFrame()

    records = []
    for row in rows:
        records.append({
            "stock_code": row.stock_code,
            "order_type": row.order_type,
            "executed_price": float(row.executed_price) if row.executed_price else 0,
            "profit_loss_rate": float(row.profit_loss_rate) if row.profit_loss_rate else 0,
            "created_at": row.created_at,
        })
    return pd.DataFrame(records)


# ─── 성능 저하 감지 ────────────────────────────────────────────────────────

async def check_performance_degradation() -> Tuple[bool, Dict[str, Any]]:
    """최근 20건 매매의 승률이 전체 평균보다 15%p 이상 하락했는지 확인합니다.

    Returns:
        (degraded: bool, details: dict)
    """
    trade_df = await _load_trade_history(months=12)
    if trade_df.empty or len(trade_df) < PERF_DEGRADE_LOOKBACK:
        return False, {"reason": "데이터 부족", "total_trades": len(trade_df) if not trade_df.empty else 0}

    # 매도 건만 (수익률이 있는 건)
    sell_df = trade_df[
        (trade_df["order_type"].str.upper() == "SELL")
        & (trade_df["profit_loss_rate"] != 0)
    ].copy()

    if len(sell_df) < PERF_DEGRADE_LOOKBACK:
        return False, {"reason": f"매도 건수 부족: {len(sell_df)}", "total_sells": len(sell_df)}

    # 전체 승률
    total_wins = (sell_df["profit_loss_rate"] > 0).sum()
    total_win_rate = total_wins / len(sell_df)

    # 최근 N건 승률
    recent_sells = sell_df.tail(PERF_DEGRADE_LOOKBACK)
    recent_wins = (recent_sells["profit_loss_rate"] > 0).sum()
    recent_win_rate = recent_wins / len(recent_sells)

    degraded = (total_win_rate - recent_win_rate) >= PERF_DEGRADE_THRESHOLD

    details = {
        "total_win_rate": round(float(total_win_rate), 4),
        "recent_win_rate": round(float(recent_win_rate), 4),
        "gap": round(float(total_win_rate - recent_win_rate), 4),
        "threshold": PERF_DEGRADE_THRESHOLD,
        "degraded": degraded,
        "total_sells": len(sell_df),
        "recent_lookback": PERF_DEGRADE_LOOKBACK,
    }

    if degraded:
        logger.warning(
            f"ML 성능 저하 감지: 전체 승률 {total_win_rate:.1%} vs "
            f"최근 {PERF_DEGRADE_LOOKBACK}건 승률 {recent_win_rate:.1%} "
            f"(차이 {total_win_rate - recent_win_rate:.1%}p >= {PERF_DEGRADE_THRESHOLD:.0%}p)"
        )

    return degraded, details


# ─── 메인 재학습 실행 ──────────────────────────────────────────────────────

async def run_retrain(
    model_dir: Optional[Path] = None,
    force: bool = False,
) -> Dict[str, Any]:
    """ML 재학습 파이프라인을 실행합니다.

    Args:
        model_dir: 모델 저장 디렉토리 (기본: ~/.cache/auto-stock/models/)
        force: True 시 성능 저하 여부와 무관하게 강제 재학습

    Returns:
        재학습 결과 메트릭
    """
    if model_dir is None:
        model_dir = MODEL_BASE_DIR

    start_time = datetime.now(KST)
    logger.info("ML 재학습 파이프라인 시작...")

    # 1. 일봉 데이터 로드
    logger.info(f"최근 {ROLLING_WINDOW_MONTHS}개월 일봉 데이터 로딩...")
    price_df = await _load_daily_prices(months=ROLLING_WINDOW_MONTHS)
    if price_df.empty:
        msg = "재학습 중단: stock_daily_prices 데이터 없음"
        logger.warning(msg)
        return {"status": "skipped", "reason": msg}

    stock_codes = price_df["stock_code"].unique()
    logger.info(f"로드된 종목 수: {len(stock_codes)}, 총 행: {len(price_df)}")

    # 2. 종목별 피처 + 라벨 생성
    all_X_list: List[np.ndarray] = []
    all_y_list: List[np.ndarray] = []

    for code in stock_codes:
        code_df = price_df[price_df["stock_code"] == code].copy()
        if len(code_df) < 80:  # 최소 80일 필요 (60일 피처 + 20일 여유)
            continue

        code_df = code_df.reset_index(drop=True)

        try:
            features = _compute_retrain_features(code_df)
            labels = _prepare_retrain_labels(code_df["close"].values.astype(float))
            X, y = _build_training_dataset(features, labels)

            if len(X) > 0:
                all_X_list.append(X)
                all_y_list.append(y)
        except Exception as e:
            logger.debug(f"종목 {code} 피처 추출 실패: {e}")
            continue

    if not all_X_list:
        msg = "재학습 중단: 유효 학습 데이터 없음"
        logger.warning(msg)
        return {"status": "skipped", "reason": msg}

    X_all = np.vstack(all_X_list)
    y_all = np.concatenate(all_y_list)

    logger.info(
        f"학습 데이터 준비 완료: {len(X_all)} 샘플, "
        f"클래스 분포: {dict(zip(*np.unique(y_all, return_counts=True)))}"
    )

    if len(X_all) < 100:
        msg = f"재학습 중단: 학습 데이터 부족 ({len(X_all)}개 < 100개)"
        logger.warning(msg)
        return {"status": "skipped", "reason": msg}

    # 3. Train/Val split (80/20, 시계열이므로 뒤쪽 20%를 검증용)
    split_idx = int(len(X_all) * 0.8)
    X_train, X_val = X_all[:split_idx], X_all[split_idx:]
    y_train, y_val = y_all[:split_idx], y_all[split_idx:]

    # 4. 스케일링
    from sklearn.preprocessing import StandardScaler
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_val_scaled = scaler.transform(X_val)

    # 5. 학습 (CPU 바운드이므로 워커 스레드에서 실행)
    logger.info("LightGBM 학습 시작 (워커 스레드)...")
    model, metrics = await asyncio.to_thread(
        _train_lightgbm, X_train_scaled, y_train, X_val_scaled, y_val
    )

    logger.info(
        f"학습 완료 - Accuracy: {metrics['accuracy']:.4f}, "
        f"Precision: {metrics['precision']:.4f}, "
        f"Recall: {metrics['recall']:.4f}"
    )

    # 6. 모델 저장 (버전 관리)
    saved_path = await asyncio.to_thread(
        _save_model_version, model, scaler, metrics, model_dir
    )

    elapsed = (datetime.now(KST) - start_time).total_seconds()
    result = {
        "status": "completed",
        "model_path": str(saved_path),
        "elapsed_seconds": round(elapsed, 1),
        "total_samples": len(X_all),
        "stocks_used": len(all_X_list),
        "metrics": metrics,
        "retrained_at": start_time.isoformat(),
    }

    logger.info(f"ML 재학습 파이프라인 완료 ({elapsed:.1f}초)")
    return result


# ─── 스케줄러 ──────────────────────────────────────────────────────────────

async def ml_retrain_scheduler():
    """ML 재학습 스케줄러.

    - 매주 일요일 02:00 KST 주기적 재학습
    - 매일 09:30 KST 성능 저하 감지 -> 자동 재학습 트리거
    - 1시간 주기로 루프 (sleep 3600초)
    """
    logger.info("ML 재학습 스케줄러 시작")
    await asyncio.sleep(30)  # 앱 초기화 대기

    last_scheduled_retrain: Optional[date] = None
    last_degradation_check: Optional[date] = None

    while True:
        try:
            now = datetime.now(KST)
            today = now.date()

            # ── 주기적 재학습: 매주 일요일 02:00 KST ──
            if (
                now.weekday() == 6  # Sunday
                and now.hour >= 2
                and last_scheduled_retrain != today
            ):
                logger.info("주기적 ML 재학습 트리거 (매주 일요일 02:00)")
                last_scheduled_retrain = today
                try:
                    result = await run_retrain(force=True)
                    logger.info(f"주기적 재학습 결과: {result.get('status')}")
                except Exception as e:
                    logger.error(f"주기적 재학습 실패: {e}", exc_info=True)

            # ── 성능 저하 감지: 매일 09:30 KST (장 시작 전) ──
            if (
                now.hour >= 9
                and now.minute >= 30
                and last_degradation_check != today
            ):
                last_degradation_check = today
                try:
                    degraded, details = await check_performance_degradation()
                    if degraded:
                        logger.warning(f"성능 저하 감지 -> 자동 재학습 트리거: {details}")
                        result = await run_retrain(force=True)
                        logger.info(f"성능 저하 재학습 결과: {result.get('status')}")
                    else:
                        logger.debug(f"성능 저하 미감지: {details}")
                except Exception as e:
                    logger.error(f"성능 저하 감지/재학습 실패: {e}", exc_info=True)

        except asyncio.CancelledError:
            logger.info("ML 재학습 스케줄러 종료")
            return
        except Exception as e:
            logger.error(f"ML 재학습 스케줄러 오류: {e}", exc_info=True)

        await asyncio.sleep(3600)  # 1시간 주기


# ─── 외부 호출용 함수 ──────────────────────────────────────────────────────

async def get_retrain_status() -> Dict[str, Any]:
    """현재 모델 상태와 최근 재학습 정보를 반환합니다."""
    model, scaler, metrics = await asyncio.to_thread(load_latest_model)

    model_dir = MODEL_BASE_DIR
    version_dirs = []
    if model_dir.exists():
        version_dirs = sorted(
            [d.name for d in model_dir.iterdir() if d.is_dir() and d.name.startswith("v_")],
            reverse=True,
        )

    return {
        "model_loaded": model is not None,
        "scaler_loaded": scaler is not None,
        "latest_metrics": metrics,
        "model_dir": str(model_dir),
        "available_versions": version_dirs[:MAX_MODEL_VERSIONS],
    }