"""
ml/ensemble.py — Ensemble modeling: XGBoost + LightGBM + LSTM/Transformer.

Combines multiple model families for robust predictions via:
  1. Weighted averaging (simple, robust)
  2. Stacking meta-learner (learns optimal blend)

GPU-accelerated LSTM for temporal pattern capture.
"""
from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression

logger = logging.getLogger("trading_system.ml.ensemble")


@dataclass
class EnsemblePrediction:
    """Output from ensemble model."""
    probs: np.ndarray           # (n_samples, n_classes)
    member_probs: dict[str, np.ndarray]  # model_name -> probs
    member_weights: dict[str, float]     # model_name -> weight


# ═══════════════════════════════════════════════════════════════════════════════
#  INDIVIDUAL MODEL FACTORIES
# ═══════════════════════════════════════════════════════════════════════════════

def _create_xgboost(n_classes: int = 3, **kwargs):
    import xgboost as xgb
    defaults = dict(
        n_estimators=100, max_depth=4, learning_rate=0.04,
        subsample=0.75, colsample_bytree=0.65,
        reg_alpha=2.0, reg_lambda=4.0,
        objective="multi:softprob", num_class=n_classes,
        eval_metric="mlogloss", verbosity=0,
        use_label_encoder=False, min_child_weight=15,
    )
    defaults.update(kwargs)
    return xgb.XGBClassifier(**defaults)


def _create_lightgbm(n_classes: int = 3, **kwargs):
    import lightgbm as lgb
    defaults = dict(
        n_estimators=100, max_depth=4, learning_rate=0.04,
        subsample=0.75, colsample_bytree=0.65,
        reg_alpha=2.0, reg_lambda=4.0,
        objective="multiclass", num_class=n_classes,
        min_child_samples=25, verbose=-1,
    )
    defaults.update(kwargs)
    return lgb.LGBMClassifier(**defaults)


class LSTMSignalModel:
    """Lightweight LSTM for temporal pattern recognition.

    Architecture: Input → LSTM(hidden) → Dropout → FC → Softmax
    Uses GPU if available. Falls back to CPU gracefully.
    """

    def __init__(
        self,
        input_dim: int = 30,
        hidden_dim: int = 64,
        num_layers: int = 2,
        dropout: float = 0.3,
        seq_len: int = 20,
        n_classes: int = 3,
        lr: float = 0.001,
        epochs: int = 30,
        batch_size: int = 128,
    ):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.dropout = dropout
        self.seq_len = seq_len
        self.n_classes = n_classes
        self.lr = lr
        self.epochs = epochs
        self.batch_size = batch_size
        self._model = None
        self._device = None
        self._fitted = False
        self.classes_ = np.arange(n_classes)

    def _build_model(self):
        import torch
        import torch.nn as nn

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

        class _LSTMNet(nn.Module):
            def __init__(self_, input_dim, hidden_dim, num_layers, dropout, n_classes):
                super().__init__()
                self_.lstm = nn.LSTM(
                    input_dim, hidden_dim, num_layers,
                    batch_first=True, dropout=dropout if num_layers > 1 else 0,
                )
                self_.dropout = nn.Dropout(dropout)
                self_.fc = nn.Linear(hidden_dim, n_classes)

            def forward(self_, x):
                # x: (batch, seq_len, features)
                lstm_out, _ = self_.lstm(x)
                last_hidden = lstm_out[:, -1, :]  # use last timestep
                out = self_.dropout(last_hidden)
                return self_.fc(out)

        self._model = _LSTMNet(
            self.input_dim, self.hidden_dim, self.num_layers,
            self.dropout, self.n_classes,
        ).to(self._device)

    def _make_sequences(self, X: np.ndarray) -> np.ndarray:
        """Convert flat feature matrix into overlapping sequences."""
        n = len(X)
        if n <= self.seq_len:
            # Pad with zeros
            padded = np.zeros((self.seq_len, X.shape[1]))
            padded[-n:] = X
            return padded[np.newaxis, ...]

        sequences = []
        for i in range(self.seq_len, n):
            sequences.append(X[i - self.seq_len:i])
        return np.array(sequences)

    def fit(self, X: np.ndarray, y: np.ndarray):
        """Train LSTM on feature sequences."""
        import torch
        import torch.nn as nn

        self.input_dim = X.shape[1]
        self._build_model()

        # Create sequences
        seqs = self._make_sequences(X)
        labels = y[self.seq_len:]  # align with sequences

        if len(seqs) != len(labels):
            min_len = min(len(seqs), len(labels))
            seqs = seqs[:min_len]
            labels = labels[:min_len]

        if len(seqs) < 50:
            logger.warning("Too few sequences for LSTM training (%d)", len(seqs))
            self._fitted = False
            return

        # Tensors
        X_t = torch.FloatTensor(seqs).to(self._device)
        y_t = torch.LongTensor(labels).to(self._device)

        # Class weights for imbalanced data
        classes, counts = np.unique(labels, return_counts=True)
        weight = torch.FloatTensor(
            [len(labels) / (len(classes) * c) for c in counts]
        ).to(self._device)

        optimizer = torch.optim.Adam(self._model.parameters(), lr=self.lr,
                                     weight_decay=1e-4)
        criterion = nn.CrossEntropyLoss(weight=weight)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, patience=5, factor=0.5,
        )

        self._model.train()
        dataset = torch.utils.data.TensorDataset(X_t, y_t)
        loader = torch.utils.data.DataLoader(
            dataset, batch_size=self.batch_size, shuffle=True,
        )

        best_loss = float('inf')
        patience_counter = 0

        for epoch in range(self.epochs):
            epoch_loss = 0.0
            for batch_X, batch_y in loader:
                optimizer.zero_grad()
                output = self._model(batch_X)
                loss = criterion(output, batch_y)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self._model.parameters(), 1.0)
                optimizer.step()
                epoch_loss += loss.item()

            avg_loss = epoch_loss / len(loader)
            scheduler.step(avg_loss)

            if avg_loss < best_loss - 0.001:
                best_loss = avg_loss
                patience_counter = 0
            else:
                patience_counter += 1
                if patience_counter >= 8:
                    break

        self._fitted = True

    def predict_proba(self, X) -> np.ndarray:
        """Predict class probabilities."""
        import torch

        if not self._fitted or self._model is None:
            n = len(X) if hasattr(X, '__len__') else 1
            return np.full((n, self.n_classes), 1.0 / self.n_classes)

        if isinstance(X, pd.DataFrame):
            X = X.values

        seqs = self._make_sequences(X)

        self._model.eval()
        with torch.no_grad():
            X_t = torch.FloatTensor(seqs).to(self._device)
            outputs = self._model(X_t)
            probs = torch.softmax(outputs, dim=1).cpu().numpy()

        # Pad beginning (no sequence available) with uniform
        n_orig = len(X)
        n_pred = len(probs)
        if n_pred < n_orig:
            pad = np.full((n_orig - n_pred, self.n_classes), 1.0 / self.n_classes)
            probs = np.vstack([pad, probs])

        return probs


# ═══════════════════════════════════════════════════════════════════════════════
#  ENSEMBLE COMBINER
# ═══════════════════════════════════════════════════════════════════════════════

class EnsembleModel:
    """Combine XGBoost + LightGBM + LSTM for robust predictions.

    Combination strategies:
      - "average": Weighted average of member probabilities.
      - "stacking": Train logistic meta-learner on out-of-fold predictions.

    Walk-forward compatible: can train/predict incrementally.
    """

    def __init__(
        self,
        combination: str = "average",
        use_lstm: bool = True,
        n_classes: int = 3,
        xgb_weight: float = 0.35,
        lgb_weight: float = 0.35,
        lstm_weight: float = 0.30,
        lstm_seq_len: int = 20,
        lstm_hidden: int = 64,
        lstm_epochs: int = 25,
    ):
        self.combination = combination
        self.use_lstm = use_lstm
        self.n_classes = n_classes
        self.xgb_weight = xgb_weight
        self.lgb_weight = lgb_weight
        self.lstm_weight = lstm_weight
        self.lstm_seq_len = lstm_seq_len
        self.lstm_hidden = lstm_hidden
        self.lstm_epochs = lstm_epochs

        self._xgb = None
        self._lgb = None
        self._lstm = None
        self._stacker = None
        self._fitted_members: set[str] = set()

    def fit(self, X: np.ndarray, y: np.ndarray):
        """Train all ensemble members.

        Args:
            X: Feature matrix (n_samples, n_features).
            y: Labels (n_samples,), integer class indices.
        """
        if isinstance(X, pd.DataFrame):
            feature_names = list(X.columns)
            X_arr = X.values
        else:
            feature_names = None
            X_arr = X

        # Mask NaN
        mask = ~np.isnan(X_arr).any(axis=1) & ~np.isnan(y)
        X_clean = X_arr[mask]
        y_clean = y[mask]

        if len(X_clean) < 100 or len(np.unique(y_clean)) < 2:
            logger.warning("Insufficient data for ensemble (%d samples)", len(X_clean))
            return

        # Train XGBoost
        try:
            self._xgb = _create_xgboost(self.n_classes)
            if feature_names:
                self._xgb.fit(pd.DataFrame(X_clean, columns=feature_names), y_clean)
            else:
                self._xgb.fit(X_clean, y_clean)
            self._fitted_members.add("xgboost")
        except Exception as e:
            logger.warning("XGBoost training failed: %s", e)

        # Train LightGBM
        try:
            self._lgb = _create_lightgbm(self.n_classes)
            if feature_names:
                self._lgb.fit(pd.DataFrame(X_clean, columns=feature_names), y_clean)
            else:
                self._lgb.fit(X_clean, y_clean)
            self._fitted_members.add("lightgbm")
        except Exception as e:
            logger.warning("LightGBM training failed: %s", e)

        # Train LSTM
        if self.use_lstm:
            try:
                self._lstm = LSTMSignalModel(
                    input_dim=X_clean.shape[1],
                    hidden_dim=self.lstm_hidden,
                    seq_len=self.lstm_seq_len,
                    n_classes=self.n_classes,
                    epochs=self.lstm_epochs,
                )
                self._lstm.fit(X_clean, y_clean)
                if self._lstm._fitted:
                    self._fitted_members.add("lstm")
            except Exception as e:
                logger.warning("LSTM training failed: %s", e)

        # Stacking meta-learner (if sufficient data)
        if self.combination == "stacking" and len(self._fitted_members) >= 2:
            self._fit_stacker(X_clean, y_clean, feature_names)

    def _fit_stacker(self, X: np.ndarray, y: np.ndarray, feature_names=None):
        """Train stacking meta-learner on in-sample predictions."""
        meta_features = self._get_meta_features(X, feature_names)
        if meta_features is not None and len(meta_features) > 50:
            self._stacker = LogisticRegression(
                C=0.5, max_iter=300, multi_class="multinomial",
            )
            self._stacker.fit(meta_features, y)

    def _get_meta_features(self, X: np.ndarray, feature_names=None) -> np.ndarray | None:
        """Get stacked features from member predictions."""
        parts = []
        if "xgboost" in self._fitted_members:
            if feature_names:
                p = self._xgb.predict_proba(pd.DataFrame(X, columns=feature_names))
            else:
                p = self._xgb.predict_proba(X)
            parts.append(p)
        if "lightgbm" in self._fitted_members:
            if feature_names:
                p = self._lgb.predict_proba(pd.DataFrame(X, columns=feature_names))
            else:
                p = self._lgb.predict_proba(X)
            parts.append(p)
        if "lstm" in self._fitted_members:
            p = self._lstm.predict_proba(X)
            parts.append(p)

        if not parts:
            return None
        return np.hstack(parts)

    def predict_proba(self, X) -> EnsemblePrediction:
        """Get ensemble probability predictions.

        Returns EnsemblePrediction with combined and per-member probs.
        """
        if isinstance(X, pd.DataFrame):
            feature_names = list(X.columns)
            X_arr = X.values
        else:
            feature_names = None
            X_arr = X

        X_clean = np.nan_to_num(X_arr, nan=0.0)
        n = len(X_clean)
        member_probs = {}

        # Get member predictions
        if "xgboost" in self._fitted_members:
            try:
                if feature_names:
                    p = self._xgb.predict_proba(pd.DataFrame(X_clean, columns=feature_names))
                else:
                    p = self._xgb.predict_proba(X_clean)
                p = self._pad_probs(p, self.n_classes)
                member_probs["xgboost"] = p
            except Exception:
                pass

        if "lightgbm" in self._fitted_members:
            try:
                if feature_names:
                    p = self._lgb.predict_proba(pd.DataFrame(X_clean, columns=feature_names))
                else:
                    p = self._lgb.predict_proba(X_clean)
                p = self._pad_probs(p, self.n_classes)
                member_probs["lightgbm"] = p
            except Exception:
                pass

        if "lstm" in self._fitted_members:
            try:
                p = self._lstm.predict_proba(X_clean)
                p = self._pad_probs(p, self.n_classes)
                member_probs["lstm"] = p
            except Exception:
                pass

        if not member_probs:
            uniform = np.full((n, self.n_classes), 1.0 / self.n_classes)
            return EnsemblePrediction(
                probs=uniform, member_probs={}, member_weights={},
            )

        # Combine
        if self.combination == "stacking" and self._stacker is not None:
            meta = np.hstack(list(member_probs.values()))
            combined = self._stacker.predict_proba(meta)
            weights = {k: 1.0 / len(member_probs) for k in member_probs}
        else:
            # Weighted average
            weight_map = {
                "xgboost": self.xgb_weight,
                "lightgbm": self.lgb_weight,
                "lstm": self.lstm_weight,
            }
            active = {k: weight_map.get(k, 0.33) for k in member_probs}
            total_w = sum(active.values())
            weights = {k: v / total_w for k, v in active.items()}

            combined = np.zeros((n, self.n_classes))
            for name, w in weights.items():
                combined += w * member_probs[name]

        return EnsemblePrediction(
            probs=combined,
            member_probs=member_probs,
            member_weights=weights,
        )

    def _pad_probs(self, probs: np.ndarray, n_classes: int) -> np.ndarray:
        """Ensure probability matrix has correct number of columns."""
        if probs.shape[1] >= n_classes:
            return probs[:, :n_classes]
        padded = np.full((len(probs), n_classes), 1.0 / n_classes)
        padded[:, :probs.shape[1]] = probs
        return padded