"""
앙상블 모델 학습 스크립트

데이터 소스를 선택하여 LightGBM + PyTorch LSTM + Random Forest 앙상블 모델을 학습합니다.
- DB 모드(기본): stock_daily_prices 테이블의 KRX 히스토리컬 데이터 사용
- API 모드: KIS API에서 실시간 차트 데이터 수집

사용법:
    python train_ensemble.py                        # DB 기반 전체 학습
    python train_ensemble.py --source db --limit 20 # DB에서 20종목
    python train_ensemble.py --source api --test    # KIS API로 10종목 (기존 방식)
    python train_ensemble.py --min-days 500         # 최소 500일 데이터 있는 종목만
"""

import asyncio
import sys
import argparse
import logging
from pathlib import Path
from datetime import datetime

import numpy as np
import pandas as pd

# 프로젝트 루트를 Python 경로에 추가
sys.path.insert(0, str(Path(__file__).parent))

from app.services.feature_engineering import (
    chart_data_to_dataframe,
    db_data_to_dataframe,
    extract_features,
    prepare_training_data,
)
from app.services.ml_ensemble import MLEnsemble

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


# ─── DB 기반 데이터 수집 ───

async def collect_training_data_from_db(limit: int = 200, min_days: int = 200) -> tuple:
    """DB의 stock_daily_prices에서 학습 데이터 수집

    Args:
        limit: 최대 종목 수
        min_days: 최소 데이터 일수 (이보다 적으면 제외)

    Returns:
        (X_tabular, X_seq, y)
    """
    from sqlalchemy import text
    from app.core.database import AsyncSessionLocal

    all_X_tabular = []
    all_X_seq = []
    all_y = []

    async with AsyncSessionLocal() as session:
        # 충분한 데이터가 있는 종목 조회
        result = await session.execute(
            text("""
                SELECT stock_code, stock_name, COUNT(*) as cnt
                FROM stock_daily_prices
                GROUP BY stock_code, stock_name
                HAVING COUNT(*) >= :min_days
                ORDER BY cnt DESC
                LIMIT :limit
            """),
            {"min_days": min_days, "limit": limit},
        )
        stock_list = [(row[0], row[1], row[2]) for row in result.fetchall()]

        if not stock_list:
            logger.error(f"DB에 {min_days}일 이상 데이터가 있는 종목이 없습니다")
            logger.error("먼저 python collect_krx_data.py 를 실행하세요")
            return np.array([]), np.array([]), np.array([])

        logger.info(f"DB에서 {len(stock_list)}개 종목 조회 (최소 {min_days}일)")

        for i, (stock_code, stock_name, day_count) in enumerate(stock_list):
            logger.info(f"[{i+1}/{len(stock_list)}] {stock_code} ({stock_name}): {day_count}일 데이터")

            try:
                # DB에서 OHLCV 조회
                rows = await session.execute(
                    text("""
                        SELECT trade_date, open_price, high_price, low_price,
                               close_price, volume
                        FROM stock_daily_prices
                        WHERE stock_code = :code
                        ORDER BY trade_date ASC
                    """),
                    {"code": stock_code},
                )
                data = [
                    {
                        "trade_date": r[0],
                        "open_price": r[1],
                        "high_price": r[2],
                        "low_price": r[3],
                        "close_price": r[4],
                        "volume": r[5],
                    }
                    for r in rows.fetchall()
                ]

                if len(data) < min_days:
                    continue

                # DataFrame 변환 및 피처 추출
                df = db_data_to_dataframe(data)
                features_df = extract_features(df)

                # 학습 데이터 준비
                X_tab, X_seq, y, indices = prepare_training_data(
                    df, features_df,
                    forecast_days=3,
                    up_threshold=0.015,
                    down_threshold=-0.015,
                )

                if len(X_tab) == 0:
                    logger.warning(f"  {stock_code}: 유효 샘플 없음")
                    continue

                all_X_tabular.append(X_tab)
                all_X_seq.append(X_seq)
                all_y.append(y)

                logger.info(
                    f"  {stock_code}: {len(X_tab)}개 샘플 "
                    f"(UP:{np.sum(y==2)} HOLD:{np.sum(y==1)} DOWN:{np.sum(y==0)})"
                )

            except Exception as e:
                logger.error(f"  {stock_code} 처리 실패: {e}")

    if not all_X_tabular:
        return np.array([]), np.array([]), np.array([])

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

    return X_tabular, X_seq, y


# ─── KIS API 기반 데이터 수집 (기존 방식) ───

async def collect_training_data_from_api(stock_list: list, kis_api) -> tuple:
    """여러 종목의 차트 데이터를 KIS API에서 수집하여 학습 데이터로 변환"""
    all_X_tabular = []
    all_X_seq = []
    all_y = []

    for i, (stock_code, stock_name) in enumerate(stock_list):
        logger.info(f"[{i+1}/{len(stock_list)}] 데이터 수집: {stock_code} ({stock_name})")

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

            if not chart_data or len(chart_data) < 80:
                logger.warning(f"  {stock_code}: 데이터 부족 ({len(chart_data) if chart_data else 0}일)")
                continue

            df = chart_data_to_dataframe(chart_data)
            features_df = extract_features(df)

            X_tab, X_seq, y, indices = prepare_training_data(
                df, features_df,
                forecast_days=3,
                up_threshold=0.015,
                down_threshold=-0.015,
            )

            if len(X_tab) == 0:
                logger.warning(f"  {stock_code}: 유효 샘플 없음")
                continue

            all_X_tabular.append(X_tab)
            all_X_seq.append(X_seq)
            all_y.append(y)

            logger.info(f"  {stock_code}: {len(X_tab)}개 샘플 (UP:{np.sum(y==2)} HOLD:{np.sum(y==1)} DOWN:{np.sum(y==0)})")

        except Exception as e:
            logger.error(f"  {stock_code} 수집 실패: {e}")

        await asyncio.sleep(3)

    if not all_X_tabular:
        return np.array([]), np.array([]), np.array([])

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

    return X_tabular, X_seq, y


async def get_stock_list_from_api(kis_api, limit: int = 50) -> list:
    """거래량 상위 종목 조회 (KIS API)"""
    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_keywords = ["ETF", "레버리지", "인버스", "KODEX", "선물", "TIGER", "ARIRANG"]
                if any(k in name for k in skip_keywords):
                    continue
                stocks.append((item["mksc_shrn_iscd"], name))
    except Exception as e:
        logger.error(f"KOSPI 종목 조회 실패: {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_keywords = ["ETF", "레버리지", "인버스", "KODEX", "선물", "TIGER", "ARIRANG"]
                if any(k in name for k in skip_keywords):
                    continue
                stocks.append((item["mksc_shrn_iscd"], name))
    except Exception as e:
        logger.error(f"KOSDAQ 종목 조회 실패: {e}")

    logger.info(f"총 {len(stocks)}개 종목 조회 완료")
    return stocks[:limit]


# ─── 메인 ───

async def main():
    parser = argparse.ArgumentParser(description="앙상블 모델 학습")
    parser.add_argument("--source", choices=["db", "api"], default="db",
                        help="데이터 소스 (기본: db)")
    parser.add_argument("--limit", type=int, default=200,
                        help="학습 종목 수 (기본: 200)")
    parser.add_argument("--test", action="store_true",
                        help="테스트 모드 (10종목)")
    parser.add_argument("--min-days", type=int, default=200,
                        help="DB 모드: 최소 데이터 일수 (기본: 200)")
    args = parser.parse_args()

    limit = 10 if args.test else args.limit
    logger.info(f"데이터 소스: {args.source.upper()}")
    logger.info(f"{'테스트 모드' if args.test else '전체 학습'}: {limit}개 종목")

    start_time = datetime.now()

    # 데이터 수집
    if args.source == "db":
        logger.info("=" * 60)
        logger.info("DB(stock_daily_prices)에서 학습 데이터 수집 중...")
        X_tabular, X_seq, y = await collect_training_data_from_db(
            limit=limit, min_days=args.min_days
        )
    else:
        # KIS API 모드
        from app.services.korea_investment_api import KoreaInvestmentAPI
        kis_api = KoreaInvestmentAPI()

        stock_list = await get_stock_list_from_api(kis_api, limit=limit)
        if not stock_list:
            logger.error("종목 리스트를 가져올 수 없습니다")
            return

        logger.info("=" * 60)
        logger.info("KIS API에서 학습 데이터 수집 중...")
        X_tabular, X_seq, y = await collect_training_data_from_api(stock_list, kis_api)

    if len(X_tabular) == 0:
        logger.error("학습 데이터가 없습니다")
        return

    logger.info(f"총 학습 데이터: {len(X_tabular):,}개 샘플")
    logger.info(f"  클래스 분포: UP={np.sum(y==2):,}, HOLD={np.sum(y==1):,}, DOWN={np.sum(y==0):,}")

    # 앙상블 학습
    logger.info("=" * 60)
    logger.info("앙상블 모델 학습 중...")
    ensemble = MLEnsemble()
    metrics = ensemble.train(X_tabular, X_seq, y)

    # 결과 출력
    duration = (datetime.now() - start_time).total_seconds()
    logger.info("=" * 60)
    logger.info("학습 완료!")
    logger.info(f"  소요 시간: {duration:.1f}초 ({duration/60:.1f}분)")
    logger.info(f"  학습 결과:")
    for key, value in metrics.items():
        logger.info(f"    {key}: {value}")
    logger.info("=" * 60)


if __name__ == "__main__":
    asyncio.run(main())