auto-stock/backend/app/api/endpoints/ml_ensemble.py

"""
ML 앙상블 API 엔드포인트

- POST /train: 모델 학습
- GET /predict/{stock_code}: 개별 종목 예측
- GET /status: 모델 상태 조회
- GET /backtest: 백테스트 결과
"""

from fastapi import APIRouter, Depends, HTTPException, Query, BackgroundTasks
from app.core.security import require_api_key
import logging
import asyncio

from app.services.ml_ensemble import get_ensemble
from app.services.korea_investment_api import korea_investment_api
from app.core.config import settings

router = APIRouter()
logger = logging.getLogger(__name__)


@router.get("/status")
async def get_model_status():
    """앙상블 모델 상태 조회"""
    ensemble = get_ensemble()
    status = ensemble.get_status()
    status["config"] = {
        "enabled": settings.ML_ENSEMBLE_ENABLED,
        "rule_weight": settings.ML_ENSEMBLE_RULE_WEIGHT,
        "ml_weight": settings.ML_ENSEMBLE_ML_WEIGHT,
        "buy_threshold": settings.ML_ENSEMBLE_BUY_THRESHOLD,
        "confidence_min": settings.ML_ENSEMBLE_CONFIDENCE_MIN,
    }
    return status


@router.get("/predict/{stock_code}")
async def predict_stock(stock_code: str):
    """개별 종목 ML 앙상블 예측"""
    ensemble = get_ensemble()

    if not ensemble.is_trained:
        raise HTTPException(
            status_code=503,
            detail="모델이 학습되지 않았습니다. POST /ml-ensemble/train 으로 학습을 먼저 실행하세요."
        )

    try:
        chart_resp = await korea_investment_api.get_daily_chart(stock_code, "D")
        chart_data = chart_resp.get("output", [])

        if not chart_data or len(chart_data) < 61:
            raise HTTPException(status_code=400, detail=f"차트 데이터 부족: {len(chart_data) if chart_data else 0}일 (최소 61일 필요)")

        result = ensemble.predict(chart_data)
        result["stock_code"] = stock_code

        return result

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"예측 실패 ({stock_code}): {e}")
        raise HTTPException(status_code=500, detail=str(e))


async def _train_task(limit: int):
    """백그라운드 학습 작업"""
    import numpy as np
    from app.services.feature_engineering import (
        chart_data_to_dataframe, extract_features, prepare_training_data
    )
    from app.services.korea_investment_api import KoreaInvestmentAPI

    kis_api = KoreaInvestmentAPI()

    # 종목 조회
    stocks = []
    try:
        kospi = await kis_api.get_volume_rank("J")
        if kospi and kospi.get("output"):
            for item in kospi["output"]:
                name = item.get("hts_kor_isnm", "")
                skip = ["ETF", "레버리지", "인버스", "KODEX", "선물", "TIGER", "ARIRANG"]
                if not any(k in name for k in skip):
                    stocks.append((item["mksc_shrn_iscd"], name))
    except Exception as e:
        logger.error(f"종목 조회 실패: {e}")

    await asyncio.sleep(3)

    try:
        kosdaq = await kis_api.get_volume_rank("Q")
        if kosdaq and kosdaq.get("output"):
            for item in kosdaq["output"]:
                name = item.get("hts_kor_isnm", "")
                skip = ["ETF", "레버리지", "인버스", "KODEX", "선물", "TIGER", "ARIRANG"]
                if not any(k in name for k in skip):
                    stocks.append((item["mksc_shrn_iscd"], name))
    except Exception as e:
        logger.error(f"종목 조회 실패: {e}")

    stocks = stocks[:limit]
    logger.info(f"학습 대상: {len(stocks)}개 종목")

    # 데이터 수집
    all_X_tab, all_X_seq, all_y = [], [], []
    for i, (code, name) in enumerate(stocks):
        try:
            chart_resp = await kis_api.get_daily_chart(code, "D")
            chart_data = chart_resp.get("output", [])
            if not chart_data or len(chart_data) < 80:
                continue

            df = chart_data_to_dataframe(chart_data)
            feat = extract_features(df)
            X_tab, X_seq, y, _ = prepare_training_data(df, feat)
            if len(X_tab) > 0:
                all_X_tab.append(X_tab)
                all_X_seq.append(X_seq)
                all_y.append(y)
                logger.info(f"[{i+1}/{len(stocks)}] {code}: {len(X_tab)}개 샘플")
        except Exception as e:
            logger.warning(f"[{i+1}] {code} 실패: {e}")
        await asyncio.sleep(3)

    if not all_X_tab:
        logger.error("학습 데이터 없음")
        return

    X_tabular = np.vstack(all_X_tab)
    X_seq = np.vstack(all_X_seq)
    y = np.concatenate(all_y)

    logger.info(f"총 {len(X_tabular)}개 샘플로 학습 시작")

    ensemble = get_ensemble()
    metrics = ensemble.train(X_tabular, X_seq, y)
    logger.info(f"학습 완료: {metrics}")


@router.post("/train", dependencies=[Depends(require_api_key)])
async def train_model(
    background_tasks: BackgroundTasks,
    limit: int = Query(50, description="학습 종목 수"),
):
    """앙상블 모델 학습 (백그라운드 실행)"""
    background_tasks.add_task(_train_task, limit)
    return {
        "message": f"학습이 백그라운드에서 시작되었습니다 ({limit}개 종목)",
        "status": "training_started",
    }


@router.get("/backtest")
async def get_backtest():
    """모델 학습 메트릭 및 백테스트 결과 조회"""
    ensemble = get_ensemble()
    if not ensemble.is_trained:
        raise HTTPException(status_code=503, detail="모델이 학습되지 않았습니다")

    return {
        "is_trained": True,
        "metrics": ensemble.training_metrics,
        "weights": {
            "lgbm": ensemble.weights[0],
            "lstm": ensemble.weights[1],
            "rf": ensemble.weights[2],
        },
    }