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

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

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

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
        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).to(self.device)
        y_tensor = torch.LongTensor(y).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:
            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).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(),
            "seq_mean": self.seq_mean,
            "seq_std": self.seq_std,
            "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:
            checkpoint = torch.load(path, map_location=self.device, weights_only=False)
            self.model.load_state_dict(checkpoint["model_state"])
            self.seq_mean = checkpoint["seq_mean"]
            self.seq_std = checkpoint["seq_std"]
            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()
        X_scaled = self.scaler.fit_transform(X_tabular)

        # 학습/검증 분리 (시계열이므로 뒤쪽 20%를 검증으로)
        split_idx = int(len(X_scaled) * 0.8)
        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 = []

        # LightGBM
        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

        # LSTM
        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

        # RF
        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]:
        """단일 종목 예측.

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

        Returns:
            {
                "ensemble_score": int (0~100),
                "confidence": float (0~1),
                "signal": str ("BUY"/"HOLD"/"SELL"),
                "model_scores": {"lgbm": ..., "lstm": ..., "rf": ...},
                "probabilities": {"up": ..., "hold": ..., "down": ...}
            }
        """
        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]

        # 앙상블 스코어: UP 확률 기반 (0~100)
        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),
            },
        }

    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")

        # 메트릭 저장
        import json
        with open(self.model_dir / "metrics.json", "w") as f:
            json.dump(self.training_metrics, f, indent=2, ensure_ascii=False, default=str)

        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