"""
ML 앙상블 모델 서비스

3개 모델을 결합하여 주식 매수/매도 신호를 예측합니다:
- LightGBM (40%): 탭뷸러 피처 기반 분류
- PyTorch LSTM (35%): 시계열 패턴 학습
- Random Forest (25%): 피처 조합 기반 분류

앙상블 스코어: 0~100 (높을수록 매수 신호 강함)
"""

import asyncio
import logging
import os
from pathlib import Path
from typing import Dict, Any, Optional, Tuple

import numpy as np
import joblib

from app.services.feature_engineering import (
    FEATURE_COLUMNS,
    LSTM_SEQUENCE_FEATURES,
    extract_single_features,
    chart_data_to_dataframe,
    extract_features,
    prepare_training_data,
)

logger = logging.getLogger(__name__)

# 모델 저장 디렉토리
MODEL_DIR = Path(__file__).parent.parent.parent / "models" / "ensemble"


# ─── PyTorch LSTM 모델 정의 ───

def _get_torch():
    """PyTorch lazy import"""
    try:
        import torch
        import torch.nn as nn
        return torch, nn
    except ImportError:
        return None, None


class StockLSTM:
    """PyTorch LSTM 모델 래퍼"""

    def __init__(self, input_size: int = 5, hidden_size: int = 64, num_layers: int = 2,
                 num_classes: int = 3, dropout: float = 0.3):
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.num_classes = num_classes
        self.dropout = dropout
        self.model = None
        self.device = None
        # 정규화 파라미터 — train_model() 또는 load() 전까지 None
        self.seq_mean: Optional[np.ndarray] = None
        self.seq_std: Optional[np.ndarray] = None
        self._build()

    def _build(self):
        torch, nn = _get_torch()
        if torch is None:
            logger.warning("PyTorch not available, LSTM model disabled")
            return

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        class LSTMNet(nn.Module):
            def __init__(self_, input_size, hidden_size, num_layers, num_classes, dropout):
                super().__init__()
                self_.lstm = nn.LSTM(
                    input_size=input_size,
                    hidden_size=hidden_size,
                    num_layers=num_layers,
                    batch_first=True,
                    dropout=dropout if num_layers > 1 else 0,
                )
                self_.fc1 = nn.Linear(hidden_size, 32)
                self_.relu = nn.ReLU()
                self_.drop = nn.Dropout(dropout)
                self_.fc2 = nn.Linear(32, num_classes)

            def forward(self_, x):
                out, _ = self_.lstm(x)
                out = out[:, -1, :]  # 마지막 타임스텝
                out = self_.drop(self_.relu(self_.fc1(out)))
                out = self_.fc2(out)
                return out

        self.model = LSTMNet(
            self.input_size, self.hidden_size, self.num_layers,
            self.num_classes, self.dropout,
        ).to(self.device)

    def train_model(self, X_seq: np.ndarray, y: np.ndarray,
                    epochs: int = 50, lr: float = 0.001, batch_size: int = 32) -> Dict[str, float]:
        """LSTM 모델 학습"""
        torch, nn = _get_torch()
        if torch is None or self.model is None:
            return {"error": "PyTorch not available"}

        from torch.utils.data import TensorDataset, DataLoader

        # 시퀀스 피처 정규화 (채널별)
        self.seq_mean = X_seq.mean(axis=(0, 1))
        self.seq_std = X_seq.std(axis=(0, 1)) + 1e-8
        X_norm = (X_seq - self.seq_mean) / self.seq_std

        X_tensor = torch.FloatTensor(X_norm.copy()).to(self.device)
        y_tensor = torch.LongTensor(y.copy()).to(self.device)

        dataset = TensorDataset(X_tensor, y_tensor)
        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

        criterion = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)

        self.model.train()
        best_loss = float("inf")

        for epoch in range(epochs):
            total_loss = 0.0
            correct = 0
            total = 0
            for X_batch, y_batch in loader:
                optimizer.zero_grad()
                outputs = self.model(X_batch)
                loss = criterion(outputs, y_batch)
                loss.backward()
                optimizer.step()

                total_loss += loss.item() * len(y_batch)
                _, predicted = torch.max(outputs, 1)
                correct += (predicted == y_batch).sum().item()
                total += len(y_batch)

            avg_loss = total_loss / total
            accuracy = correct / total
            if avg_loss < best_loss:
                best_loss = avg_loss

            if (epoch + 1) % 10 == 0:
                logger.info(f"LSTM Epoch [{epoch+1}/{epochs}] Loss: {avg_loss:.4f} Acc: {accuracy:.4f}")

        return {"final_loss": best_loss, "final_accuracy": accuracy}

    def predict_proba(self, X_seq: np.ndarray) -> np.ndarray:
        """예측 확률 반환 (N, 3)"""
        torch, nn = _get_torch()
        if torch is None or self.model is None or self.seq_mean is None or self.seq_std is None:
            return np.full((len(X_seq), 3), 1.0 / 3)

        self.model.eval()
        X_norm = (X_seq - self.seq_mean) / self.seq_std
        X_tensor = torch.FloatTensor(X_norm.copy()).to(self.device)

        with torch.no_grad():
            outputs = self.model(X_tensor)
            probs = torch.softmax(outputs, dim=1).cpu().numpy()
        return probs

    def save(self, path: Path):
        torch, _ = _get_torch()
        if torch is None or self.model is None:
            return
        torch.save({
            "model_state": self.model.state_dict(),
            # numpy 배열을 그대로 저장하면 weights_only=True 로드가 거부된다
            # (2026-09-23: LSTM이 매번 로드 실패해 앙상블에서 조용히 빠져 있었음) → list 로 저장
            "seq_mean": self.seq_mean.tolist() if self.seq_mean is not None else None,
            "seq_std": self.seq_std.tolist() if self.seq_std is not None else None,
            "config": {
                "input_size": self.input_size,
                "hidden_size": self.hidden_size,
                "num_layers": self.num_layers,
                "num_classes": self.num_classes,
                "dropout": self.dropout,
            },
        }, path)

    def load(self, path: Path) -> bool:
        torch, _ = _get_torch()
        if torch is None:
            return False
        try:
            # [이슈5 수정] weights_only=True로 변경하여 임의 코드 실행(pickle RCE) 방지.
            # 단, 구버전 PyTorch(<2.0)에서는 weights_only 미지원 — except로 폴백.
            try:
                checkpoint = torch.load(path, map_location=self.device, weights_only=True)
            except TypeError:
                # PyTorch < 2.0: weights_only 파라미터 미지원
                checkpoint = torch.load(path, map_location=self.device)  # noqa: S614
            except Exception as _wo_err:
                # 구 체크포인트(numpy 배열 포함)는 weights_only=True 로 못 읽는다.
                # 우리 학습 산출물(MODEL_DIR 내부)일 때만 완화해서 재시도한다.
                try:
                    _inside = Path(path).resolve().is_relative_to(MODEL_DIR.resolve())
                except Exception:
                    _inside = False
                if not _inside:
                    raise
                logger.warning(
                    f"LSTM 체크포인트 weights_only 로드 실패 → 자체 학습 산출물로 재시도: {_wo_err}"
                )
                checkpoint = torch.load(path, map_location=self.device, weights_only=False)  # noqa: S614
            self.model.load_state_dict(checkpoint["model_state"])
            _mean, _std = checkpoint["seq_mean"], checkpoint["seq_std"]
            self.seq_mean = None if _mean is None else np.asarray(_mean, dtype=np.float32)
            self.seq_std = None if _std is None else np.asarray(_std, dtype=np.float32)
            self.model.eval()
            return True
        except Exception as e:
            logger.error(f"Failed to load LSTM model: {e}")
            return False


# ─── 앙상블 모델 클래스 ───

class MLEnsemble:
    """3개 모델(LightGBM + LSTM + RF)을 앙상블하여 주식 매수/매도 예측"""

    def __init__(self, model_dir: Path = MODEL_DIR,
                 weights: Tuple[float, float, float] = (0.40, 0.35, 0.25)):
        self.model_dir = model_dir
        self.weights = weights  # (lgbm, lstm, rf)

        self.lgbm_model = None
        self.lstm_model = None
        self.rf_model = None
        self.scaler = None  # tabular feature scaler

        self.is_trained = False
        self.training_metrics: Dict[str, Any] = {}

    def _ensure_dir(self):
        self.model_dir.mkdir(parents=True, exist_ok=True)

    def train(self, X_tabular: np.ndarray, X_seq: np.ndarray, y: np.ndarray) -> Dict[str, Any]:
        """전체 앙상블 학습

        Args:
            X_tabular: (N, 20) 탭뷸러 피처
            X_seq: (N, 60, 5) LSTM 시퀀스
            y: (N,) 레이블 (0=DOWN, 1=HOLD, 2=UP)

        Returns:
            학습 메트릭
        """
        self._ensure_dir()
        metrics = {}

        if len(X_tabular) < 50:
            return {"error": f"학습 데이터 부족: {len(X_tabular)}개 (최소 50개 필요)"}

        logger.info(f"앙상블 학습 시작: {len(X_tabular)}개 샘플, 클래스 분포: {np.bincount(y)}")

        # 탭뷸러 피처 스케일링
        from sklearn.preprocessing import StandardScaler
        self.scaler = StandardScaler()

        # 학습/검증 분리 (시계열이므로 뒤쪽 20%를 검증으로)
        split_idx = int(len(X_tabular) * 0.8)

        # S2 fix: scaler는 학습 데이터에만 fit하여 data leakage 방지
        self.scaler.fit(X_tabular[:split_idx])
        X_scaled = self.scaler.transform(X_tabular)
        X_train, X_val = X_scaled[:split_idx], X_scaled[split_idx:]
        y_train, y_val = y[:split_idx], y[split_idx:]
        X_seq_train, X_seq_val = X_seq[:split_idx], X_seq[split_idx:]

        # 1. LightGBM
        logger.info(">>> LightGBM 학습 중...")
        try:
            import lightgbm as lgb
            self.lgbm_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",
            )
            self.lgbm_model.fit(
                X_train, y_train,
                eval_set=[(X_val, y_val)],
            )
            lgbm_val_pred = self.lgbm_model.predict(X_val)
            lgbm_acc = np.mean(lgbm_val_pred == y_val)
            metrics["lgbm_accuracy"] = float(lgbm_acc)
            logger.info(f"LightGBM 검증 정확도: {lgbm_acc:.4f}")

            importance = self.lgbm_model.feature_importances_
            top_features = sorted(
                zip(FEATURE_COLUMNS, importance),
                key=lambda x: x[1], reverse=True
            )[:5]
            metrics["lgbm_top_features"] = [
                {"name": name, "importance": int(imp)} for name, imp in top_features
            ]
        except Exception as e:
            logger.error(f"LightGBM 학습 실패: {e}")
            metrics["lgbm_error"] = str(e)

        # 2. PyTorch LSTM
        logger.info(">>> PyTorch LSTM 학습 중...")
        try:
            self.lstm_model = StockLSTM(
                input_size=len(LSTM_SEQUENCE_FEATURES),
                hidden_size=64,
                num_layers=2,
                num_classes=3,
                dropout=0.3,
            )
            lstm_metrics = self.lstm_model.train_model(
                X_seq_train, y_train,
                epochs=50, lr=0.001, batch_size=32,
            )
            lstm_val_probs = self.lstm_model.predict_proba(X_seq_val)
            lstm_val_pred = np.argmax(lstm_val_probs, axis=1)
            lstm_acc = np.mean(lstm_val_pred == y_val)
            metrics["lstm_accuracy"] = float(lstm_acc)
            metrics["lstm_training"] = lstm_metrics
            logger.info(f"LSTM 검증 정확도: {lstm_acc:.4f}")
        except Exception as e:
            logger.error(f"LSTM 학습 실패: {e}")
            metrics["lstm_error"] = str(e)

        # 3. Random Forest
        logger.info(">>> Random Forest 학습 중...")
        try:
            from sklearn.ensemble import RandomForestClassifier
            self.rf_model = RandomForestClassifier(
                n_estimators=100,
                max_depth=8,
                min_samples_split=10,
                min_samples_leaf=5,
                random_state=42,
                n_jobs=-1,
            )
            self.rf_model.fit(X_train, y_train)
            rf_val_pred = self.rf_model.predict(X_val)
            rf_acc = np.mean(rf_val_pred == y_val)
            metrics["rf_accuracy"] = float(rf_acc)
            logger.info(f"Random Forest 검증 정확도: {rf_acc:.4f}")
        except Exception as e:
            logger.error(f"Random Forest 학습 실패: {e}")
            metrics["rf_error"] = str(e)

        # 앙상블 검증
        ensemble_probs = self._ensemble_predict_proba(X_val, X_seq_val)
        if ensemble_probs is not None:
            ensemble_pred = np.argmax(ensemble_probs, axis=1)
            ensemble_acc = np.mean(ensemble_pred == y_val)
            metrics["ensemble_accuracy"] = float(ensemble_acc)
            logger.info(f"앙상블 검증 정확도: {ensemble_acc:.4f}")

        self._save_models()
        self.is_trained = True
        self.training_metrics = metrics

        logger.info(f"앙상블 학습 완료: {metrics}")
        return metrics

    def _ensemble_predict_proba(self, X_tabular: np.ndarray, X_seq: np.ndarray) -> Optional[np.ndarray]:
        """3개 모델의 가중 평균 확률"""
        probs_list = []
        weights_list = []

        if self.lgbm_model is not None:
            try:
                lgbm_probs = self.lgbm_model.predict_proba(X_tabular)
                probs_list.append(lgbm_probs)
                weights_list.append(self.weights[0])
            except Exception:
                pass

        if self.lstm_model is not None and self.lstm_model.model is not None:
            try:
                lstm_probs = self.lstm_model.predict_proba(X_seq)
                probs_list.append(lstm_probs)
                weights_list.append(self.weights[1])
            except Exception:
                pass

        if self.rf_model is not None:
            try:
                rf_probs = self.rf_model.predict_proba(X_tabular)
                probs_list.append(rf_probs)
                weights_list.append(self.weights[2])
            except Exception:
                pass

        if not probs_list:
            return None

        total_weight = sum(weights_list)
        weighted_probs = sum(
            p * (w / total_weight) for p, w in zip(probs_list, weights_list)
        )
        return weighted_probs

    def predict(self, chart_data: list) -> Dict[str, Any]:
        """단일 종목 예측 (동기).

        이벤트 루프 내 async 핸들러에서는 predict_async() 를 사용하세요.

        Args:
            chart_data: KIS API 차트 데이터 (최신순)

        Returns:
            {ensemble_score, confidence, signal, model_scores, probabilities}
        """
        if not self.is_trained:
            return {"error": "모델이 학습되지 않았습니다", "ensemble_score": 50, "confidence": 0.0, "signal": "HOLD"}

        tabular, sequence = extract_single_features(chart_data)
        if tabular is None or sequence is None:
            return {"error": "피처 추출 실패 (데이터 부족)", "ensemble_score": 50, "confidence": 0.0, "signal": "HOLD"}

        tabular_scaled = self.scaler.transform(tabular.reshape(1, -1))
        sequence_batch = sequence.reshape(1, 60, len(LSTM_SEQUENCE_FEATURES))

        probs = self._ensemble_predict_proba(tabular_scaled, sequence_batch)
        if probs is None:
            return {"error": "모든 모델 예측 실패", "ensemble_score": 50, "confidence": 0.0, "signal": "HOLD"}

        prob_down, prob_hold, prob_up = probs[0]
        ensemble_score = int(prob_up * 100)
        max_prob = max(prob_down, prob_hold, prob_up)
        confidence = float(max_prob)

        if prob_up > prob_down and prob_up > prob_hold and prob_up > 0.4:
            signal = "BUY"
        elif prob_down > prob_up and prob_down > prob_hold and prob_down > 0.4:
            signal = "SELL"
        else:
            signal = "HOLD"

        model_scores = {}
        if self.lgbm_model is not None:
            try:
                lp = self.lgbm_model.predict_proba(tabular_scaled)[0]
                model_scores["lgbm"] = {"up": float(lp[2]), "hold": float(lp[1]), "down": float(lp[0])}
            except Exception:
                pass
        if self.lstm_model is not None and self.lstm_model.model is not None:
            try:
                lp = self.lstm_model.predict_proba(sequence_batch)[0]
                model_scores["lstm"] = {"up": float(lp[2]), "hold": float(lp[1]), "down": float(lp[0])}
            except Exception:
                pass
        if self.rf_model is not None:
            try:
                lp = self.rf_model.predict_proba(tabular_scaled)[0]
                model_scores["rf"] = {"up": float(lp[2]), "hold": float(lp[1]), "down": float(lp[0])}
            except Exception:
                pass

        return {
            "ensemble_score": ensemble_score,
            "confidence": round(confidence, 4),
            "signal": signal,
            "model_scores": model_scores,
            "probabilities": {
                "up": round(float(prob_up), 4),
                "hold": round(float(prob_hold), 4),
                "down": round(float(prob_down), 4),
            },
        }

    # ─── 비동기 래퍼 (이슈 7: 이벤트 루프 블로킹 방지) ──────────────────
    async def predict_async(self, chart_data: list) -> Dict[str, Any]:
        """asyncio.to_thread로 CPU 추론을 워커 스레드에서 실행.

        async 핸들러/스케줄러에서는 이 메서드를 사용하세요.
        무거운 RandomForest 배치 학습에는 ProcessPoolExecutor도 고려 가능.
        """
        return await asyncio.to_thread(self.predict, chart_data)

    async def train_async(
        self, X_tabular: np.ndarray, X_seq: np.ndarray, y: np.ndarray
    ) -> Dict[str, Any]:
        """asyncio.to_thread로 CPU 학습을 워커 스레드에서 실행."""
        return await asyncio.to_thread(self.train, X_tabular, X_seq, y)
    # ─────────────────────────────────────────────────────────────────────

    def _save_models(self):
        """모델 파일 저장"""
        self._ensure_dir()

        if self.lgbm_model is not None:
            joblib.dump(self.lgbm_model, self.model_dir / "lgbm_model.pkl")
        if self.rf_model is not None:
            joblib.dump(self.rf_model, self.model_dir / "rf_model.pkl")
        if self.lstm_model is not None:
            self.lstm_model.save(self.model_dir / "lstm_model.pt")
        if self.scaler is not None:
            joblib.dump(self.scaler, self.model_dir / "scaler.pkl")

        from app.core.json_helpers import _dumps
        with open(self.model_dir / "metrics.json", "w") as f:
            f.write(_dumps(self.training_metrics))

        logger.info(f"모델 저장 완료: {self.model_dir}")

    def load_models(self) -> bool:
        """저장된 모델 로드"""
        try:
            scaler_path = self.model_dir / "scaler.pkl"
            if not scaler_path.exists():
                logger.warning("저장된 앙상블 모델이 없습니다")
                return False

            self.scaler = joblib.load(scaler_path)

            lgbm_path = self.model_dir / "lgbm_model.pkl"
            if lgbm_path.exists():
                self.lgbm_model = joblib.load(lgbm_path)

            rf_path = self.model_dir / "rf_model.pkl"
            if rf_path.exists():
                self.rf_model = joblib.load(rf_path)

            lstm_path = self.model_dir / "lstm_model.pt"
            if lstm_path.exists():
                self.lstm_model = StockLSTM(
                    input_size=len(LSTM_SEQUENCE_FEATURES),
                    hidden_size=64, num_layers=2, num_classes=3, dropout=0.3,
                )
                self.lstm_model.load(lstm_path)

            metrics_path = self.model_dir / "metrics.json"
            if metrics_path.exists():
                import json
                with open(metrics_path) as f:
                    self.training_metrics = json.load(f)

            self.is_trained = True
            logger.info("앙상블 모델 로드 완료")
            return True

        except Exception as e:
            logger.error(f"앙상블 모델 로드 실패: {e}")
            return False

    def get_status(self) -> Dict[str, Any]:
        """모델 상태 조회"""
        return {
            "is_trained": self.is_trained,
            "models": {
                "lgbm": self.lgbm_model is not None,
                "lstm": self.lstm_model is not None and self.lstm_model.model is not None,
                "rf": self.rf_model is not None,
            },
            "weights": {"lgbm": self.weights[0], "lstm": self.weights[1], "rf": self.weights[2]},
            "metrics": self.training_metrics,
            "model_dir": str(self.model_dir),
        }


# ─── 싱글톤 인스턴스 ───

_ensemble_instance: Optional[MLEnsemble] = None


def get_ensemble() -> MLEnsemble:
    """앙상블 모델 싱글톤 인스턴스 반환"""
    global _ensemble_instance
    if _ensemble_instance is None:
        _ensemble_instance = MLEnsemble()
        _ensemble_instance.load_models()
    return _ensemble_instance