scripts/model_training.py

"""
Model Training Pipeline
GPU-accelerated stacked ensemble: XGBoost + LightGBM + PyTorch NN -> Logistic Regression meta-learner.
Includes hyperparameter tuning, 5-fold cross-validation, SHAP interpretation, and model saving.
"""

import os
import warnings
import itertools
from pathlib import Path

import numpy as np
import pandas as pd
import joblib

from sklearn.model_selection import StratifiedKFold
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score, log_loss, roc_auc_score, roc_curve,
    brier_score_loss, confusion_matrix,
)
from sklearn.preprocessing import StandardScaler

import xgboost as xgb
import lightgbm as lgb

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns

import shap

warnings.filterwarnings("ignore")

DATA_DIR = Path(__file__).parent.parent / "data"
MODELS_DIR = Path(__file__).parent.parent / "models"
PLOTS_DIR = Path(__file__).parent.parent / "plots"
TRAINING_CSV = DATA_DIR / "training_data.csv"
SCALER_PATH = MODELS_DIR / "scaler.pkl"
FEATURE_NAMES_PATH = MODELS_DIR / "feature_names.pkl"

MODEL_PATHS = {
    "xgb": MODELS_DIR / "xgb_model.json",
    "lgb": MODELS_DIR / "lgb_model.txt",
    "nn": MODELS_DIR / "nn_model.pt",
    "meta": MODELS_DIR / "meta_learner.pkl",
    "nn_temp": MODELS_DIR / "nn_temperature.pkl",
}

RANDOM_STATE = 42
EARLY_STOPPING_ROUNDS = 50
N_FOLDS = 5
N_TUNE_TRIALS = 30

NN_HIDDEN_LAYERS = [128, 64]
NN_DROPOUT = 0.25
NN_BATCH_SIZE = 64
NN_EPOCHS = 500
NN_LR = 0.0008
NN_PATIENCE = 50
NN_LABEL_SMOOTHING = 0.05
TEST_SPLIT_DATE = 0.80

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

XGB_PARAM_GRID = {
    "n_estimators": [300, 500, 800],
    "max_depth": [4, 6, 8],
    "learning_rate": [0.01, 0.03, 0.05, 0.1],
    "subsample": [0.7, 0.8, 1.0],
    "colsample_bytree": [0.6, 0.8, 1.0],
    "gamma": [0, 0.1, 0.2],
    "reg_alpha": [0, 0.1, 1.0],
    "reg_lambda": [0.1, 1.0, 2.0],
    "min_child_weight": [1, 3, 5],
}

LGB_PARAM_GRID = {
    "n_estimators": [300, 500, 800],
    "max_depth": [-1, 5, 7, 9],
    "learning_rate": [0.01, 0.03, 0.05, 0.1],
    "num_leaves": [31, 63, 127],
    "subsample": [0.7, 0.8, 1.0],
    "colsample_bytree": [0.6, 0.8, 1.0],
    "reg_alpha": [0, 0.1, 1.0],
    "reg_lambda": [0.1, 1.0, 2.0],
    "min_child_samples": [5, 10, 20],
}


class UFCFightNet(nn.Module):
    """PyTorch neural network base learner for UFC fight prediction."""

    def __init__(self, input_dim, hidden_layers=None, dropout=0.25):
        super().__init__()
        if hidden_layers is None:
            hidden_layers = [128, 64]

        layers = []
        prev_dim = input_dim
        for h_dim in hidden_layers:
            layers.extend([
                nn.Linear(prev_dim, h_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
            ])
            prev_dim = h_dim

        layers.append(nn.Linear(prev_dim, 1))

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x).squeeze()


class EarlyStopping:
    def __init__(self, patience=30, min_delta=0.0, verbose=True):
        self.patience = patience
        self.min_delta = min_delta
        self.verbose = verbose
        self.counter = 0
        self.best_loss = None
        self.early_stop = False

    def __call__(self, val_loss, model, path):
        if self.best_loss is None:
            self.best_loss = val_loss
            self.save_checkpoint(model, path)
        elif val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            if self.verbose:
                print(f"  EarlyStopping counter: {self.counter}/{self.patience}")
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.save_checkpoint(model, path)
            self.counter = 0

    def save_checkpoint(self, model, path):
        if self.verbose:
            print(f"  Validation loss improved. Saving model to {path}")
        torch.save(model.state_dict(), path)


def load_data():
    """Load the preprocessed training data."""
    print("\n" + "=" * 60)
    print("  Loading Training Data")
    print("=" * 60)

    if not TRAINING_CSV.exists():
        print(f"  ERROR: {TRAINING_CSV} not found. Run feature_engineering.py first.")
        raise FileNotFoundError(f"Training data not found at {TRAINING_CSV}")

    df = pd.read_csv(TRAINING_CSV)
    print(f"  Loaded {len(df)} samples with {len(df.columns) - 1} features")

    target_col = "target"
    if target_col not in df.columns:
        raise ValueError(f"Target column '{target_col}' not found in data")

    X = df.drop(columns=[target_col])
    y = df[target_col].values

    # Shuffle since synthetic dates create distribution artifacts
    rng = np.random.RandomState(RANDOM_STATE)
    shuffle_idx = rng.permutation(len(X))
    X_arr = X.values.astype(np.float32)[shuffle_idx]
    y_arr = y[shuffle_idx]

    split_idx = int(len(X_arr) * TEST_SPLIT_DATE)
    X_train, X_test = X_arr[:split_idx], X_arr[split_idx:]
    y_train, y_test = y_arr[:split_idx], y_arr[split_idx:]

    print(f"  Train: {len(X_train)} samples | Test: {len(X_test)} samples")
    print(f"  Train target distribution: A={y_train.sum():.0f} ({y_train.mean():.2%})")
    print(f"  Test target distribution:  A={y_test.sum():.0f} ({y_test.mean():.2%})")

    return X_train, X_test, y_train.astype(np.float32), y_test.astype(np.float32), list(X.columns)


def _random_search_from_grid(rng, param_grid):
    """Sample one random combination from a parameter grid."""
    params = {}
    for key, values in param_grid.items():
        params[key] = values[rng.randint(len(values))]
    return params


def tune_xgboost_cv(X_train, y_train, n_trials=30):
    """Random search hyperparameter tuning for XGBoost with CV."""
    print("\n" + "=" * 60)
    print("  Tuning XGBoost Hyperparameters (GPU)")
    print("=" * 60)

    rng = np.random.RandomState(RANDOM_STATE)
    skf = StratifiedKFold(n_splits=3, shuffle=True, random_state=RANDOM_STATE)
    best_score = -1
    best_params = None

    for i in range(n_trials):
        params = _random_search_from_grid(rng, XGB_PARAM_GRID)
        trial_model = xgb.XGBClassifier(
            **params,
            tree_method="hist",
            device="cuda",
            random_state=RANDOM_STATE,
            eval_metric="logloss",
            verbosity=0,
        )
        cv_scores = []
        for train_idx, val_idx in skf.split(X_train, y_train):
            X_tr_fold, X_val_fold = X_train[train_idx], X_train[val_idx]
            y_tr_fold, y_val_fold = y_train[train_idx], y_train[val_idx]
            trial_model.fit(X_tr_fold, y_tr_fold, eval_set=[(X_val_fold, y_val_fold)],
                           verbose=False)
            val_pred = trial_model.predict_proba(X_val_fold)[:, 1]
            cv_scores.append(roc_auc_score(y_val_fold, val_pred))

        mean_score = np.mean(cv_scores)
        if (i + 1) % 10 == 0:
            print(f"  Trial {i+1:2d}/{n_trials} | CV AUC: {mean_score:.4f} | lr={params['learning_rate']} md={params['max_depth']} est={params['n_estimators']}")

        if mean_score > best_score:
            best_score = mean_score
            best_params = params

    print(f"  Best XGBoost CV AUC: {best_score:.4f}")
    print(f"  Best params: {best_params}")
    return best_params


def train_xgboost_tuned(X_train, y_train, X_val, y_val, params):
    """Train XGBoost with best found parameters."""
    print("\n" + "=" * 60)
    print("  Training XGBoost (GPU) with Best Params")
    print("=" * 60)
    model = xgb.XGBClassifier(
        **params,
        tree_method="hist",
        device="cuda",
        random_state=RANDOM_STATE,
        early_stopping_rounds=EARLY_STOPPING_ROUNDS,
        eval_metric="logloss",
        verbosity=0,
    )
    model.fit(X_train, y_train, eval_set=[(X_val, y_val)], verbose=False)
    print(f"  Train logloss: {log_loss(y_train, model.predict_proba(X_train)[:,1]):.4f}")
    print(f"  Val logloss:   {log_loss(y_val, model.predict_proba(X_val)[:,1]):.4f}")
    return model


def tune_lightgbm_cv(X_train, y_train, n_trials=30):
    """Random search hyperparameter tuning for LightGBM with CV."""
    print("\n" + "=" * 60)
    print("  Tuning LightGBM Hyperparameters (GPU)")
    print("=" * 60)

    rng = np.random.RandomState(RANDOM_STATE)
    skf = StratifiedKFold(n_splits=3, shuffle=True, random_state=RANDOM_STATE)
    best_score = -1
    best_params = None

    for i in range(n_trials):
        params = _random_search_from_grid(rng, LGB_PARAM_GRID)
        trial_model = lgb.LGBMClassifier(
            **params,
            device_type="gpu",
            random_state=RANDOM_STATE,
            verbose=-1,
        )
        cv_scores = []
        for train_idx, val_idx in skf.split(X_train, y_train):
            X_tr_fold, X_val_fold = X_train[train_idx], X_train[val_idx]
            y_tr_fold, y_val_fold = y_train[train_idx], y_train[val_idx]
            trial_model.fit(X_tr_fold, y_tr_fold,
                           eval_set=[(X_val_fold, y_val_fold)],
                           eval_metric="logloss",
                           callbacks=[lgb.early_stopping(15, verbose=False),
                                     lgb.log_evaluation(0)])
            val_pred = trial_model.predict_proba(X_val_fold)[:, 1]
            cv_scores.append(roc_auc_score(y_val_fold, val_pred))

        mean_score = np.mean(cv_scores)
        if (i + 1) % 10 == 0:
            print(f"  Trial {i+1:2d}/{n_trials} | CV AUC: {mean_score:.4f} | lr={params['learning_rate']} nl={params['num_leaves']} md={params['max_depth']}")

        if mean_score > best_score:
            best_score = mean_score
            best_params = params

    print(f"  Best LightGBM CV AUC: {best_score:.4f}")
    print(f"  Best params: {best_params}")
    return best_params


def train_lightgbm_tuned(X_train, y_train, X_val, y_val, params):
    """Train LightGBM with best found parameters."""
    print("\n" + "=" * 60)
    print("  Training LightGBM (GPU) with Best Params")
    print("=" * 60)
    model = lgb.LGBMClassifier(
        **params,
        device_type="gpu",
        random_state=RANDOM_STATE,
        verbose=-1,
    )
    model.fit(X_train, y_train,
              eval_set=[(X_val, y_val)],
              eval_metric="logloss",
              callbacks=[lgb.early_stopping(EARLY_STOPPING_ROUNDS, verbose=False),
                        lgb.log_evaluation(0)])
    print(f"  Train logloss: {log_loss(y_train, model.predict_proba(X_train)[:,1]):.4f}")
    print(f"  Val logloss:   {log_loss(y_val, model.predict_proba(X_val)[:,1]):.4f}")
    return model


def train_neural_network(X_train, y_train, X_val, y_val, input_dim):
    """Train PyTorch neural network with GPU acceleration."""
    print("\n" + "=" * 60)
    print(f"  Training Neural Network Base Learner ({DEVICE.upper()})")
    print("=" * 60)

    train_dataset = TensorDataset(
        torch.tensor(X_train, dtype=torch.float32),
        torch.tensor(y_train, dtype=torch.float32),
    )
    val_dataset = TensorDataset(
        torch.tensor(X_val, dtype=torch.float32),
        torch.tensor(y_val, dtype=torch.float32),
    )

    train_loader = DataLoader(train_dataset, batch_size=NN_BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=NN_BATCH_SIZE * 2, shuffle=False)

    model = UFCFightNet(input_dim, NN_HIDDEN_LAYERS, NN_DROPOUT).to(DEVICE)

    class_counts = [len(y_train) - y_train.sum(), y_train.sum()]
    pos_weight = torch.tensor([class_counts[0] / max(class_counts[1], 1)]).to(DEVICE)
    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
    optimizer = optim.AdamW(model.parameters(), lr=NN_LR, weight_decay=5e-4)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=10,
                                                     min_lr=1e-6)

    checkpoint_path = MODEL_PATHS["nn"]
    early_stopping = EarlyStopping(patience=NN_PATIENCE)

    for epoch in range(NN_EPOCHS):
        model.train()
        train_loss = 0.0
        for batch_X, batch_y in train_loader:
            batch_X, batch_y = batch_X.to(DEVICE), batch_y.to(DEVICE)
            optimizer.zero_grad()
            outputs = model(batch_X)
            targets_smooth = batch_y * (1 - NN_LABEL_SMOOTHING) + 0.5 * NN_LABEL_SMOOTHING
            loss = criterion(outputs, targets_smooth)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            train_loss += loss.item() * len(batch_X)
        train_loss /= len(train_dataset)

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for batch_X, batch_y in val_loader:
                batch_X, batch_y = batch_X.to(DEVICE), batch_y.to(DEVICE)
                outputs = model(batch_X)
                val_loss += criterion(outputs, batch_y).item() * len(batch_X)
        val_loss /= len(val_dataset)

        scheduler.step(val_loss)

        if (epoch + 1) % 50 == 0 or epoch == 0:
            print(f"  Epoch {epoch+1:3d}/{NN_EPOCHS} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")

        early_stopping(val_loss, model, checkpoint_path)
        if early_stopping.early_stop:
            print(f"  Early stopping triggered at epoch {epoch+1}")
            break

    model.load_state_dict(torch.load(checkpoint_path, weights_only=True))
    model.eval()

    with torch.no_grad():
        train_preds = model(torch.tensor(X_train, dtype=torch.float32).to(DEVICE)).sigmoid().cpu().numpy()
        val_preds = model(torch.tensor(X_val, dtype=torch.float32).to(DEVICE)).sigmoid().cpu().numpy()

    # Compute optimal temperature scaling
    val_logits = model(torch.tensor(X_val, dtype=torch.float32).to(DEVICE)).detach().cpu().numpy()
    best_t = 1.0
    best_loss = float("inf")
    for t in np.arange(0.5, 5.1, 0.25):
        p = 1.0 / (1.0 + np.exp(-val_logits / t))
        p = p.clip(1e-15, 1 - 1e-15)
        nll = -np.mean(y_val * np.log(p) + (1 - y_val) * np.log(1 - p))
        if nll < best_loss:
            best_loss = nll
            best_t = t
    joblib.dump(best_t, MODEL_PATHS["nn_temp"])
    print(f"  Optimal temperature: {best_t:.2f} (NLL: {best_loss:.4f})")

    print(f"  NN train logloss: {log_loss(y_train, train_preds):.4f}")
    print(f"  NN val logloss:   {log_loss(y_val, val_preds):.4f}")
    print(f"  NN trained on {DEVICE.upper()}")

    return model


def build_meta_features(xgb_model, lgb_model, nn_model, X):
    """Generate out-of-fold meta-features from all base learners."""
    xgb_proba = xgb_model.predict_proba(X)[:, 1].reshape(-1, 1)
    lgb_proba = lgb_model.predict_proba(X)[:, 1].reshape(-1, 1)

    if isinstance(nn_model, UFCFightNet):
        nn_model.eval()
        with torch.no_grad():
            tensor_X = torch.tensor(X, dtype=torch.float32).to(DEVICE)
            nn_proba = nn_model(tensor_X).sigmoid().cpu().numpy().reshape(-1, 1)
    else:
        nn_proba = np.zeros((len(X), 1))

    meta_X = np.hstack([xgb_proba, lgb_proba, nn_proba])

    if len(meta_X) >= 3:
        meta_X = np.hstack([
            meta_X,
            (meta_X[:, 0] + meta_X[:, 1] + meta_X[:, 2]).reshape(-1, 1) / 3,  # avg proba
            np.max(meta_X, axis=1).reshape(-1, 1),  # max proba
            np.min(meta_X, axis=1).reshape(-1, 1),  # min proba
            (np.max(meta_X, axis=1) - np.min(meta_X, axis=1)).reshape(-1, 1),  # disagreement
        ])

    return meta_X


def train_meta_learner(meta_X_train, y_train):
    """Train logistic regression meta-learner on base learner outputs."""
    print("\n" + "=" * 60)
    print("  Training Meta-Learner (Logistic Regression)")
    print("=" * 60)

    meta_model = LogisticRegression(
        C=1.0,
        penalty="l2",
        solver="lbfgs",
        max_iter=1000,
        random_state=RANDOM_STATE,
    )
    meta_model.fit(meta_X_train, y_train)

    train_preds = meta_model.predict_proba(meta_X_train)[:, 1]
    print(f"  Meta-learner train accuracy: {accuracy_score(y_train, (train_preds > 0.5).astype(int)):.4f}")
    print(f"  Meta-learner train logloss:  {log_loss(y_train, train_preds):.4f}")

    return meta_model


def evaluate_ensemble(models, X_test, y_test, feature_names):
    """Evaluate the full stacked ensemble and generate plots."""
    print("\n" + "=" * 60)
    print("  Model Evaluation")
    print("=" * 60)

    xgb_model = models["xgb"]
    lgb_model = models["lgb"]
    nn_model = models["nn"]
    meta_model = models["meta"]

    xgb_proba = xgb_model.predict_proba(X_test)[:, 1]
    lgb_proba = lgb_model.predict_proba(X_test)[:, 1]

    nn_model.eval()
    with torch.no_grad():
        nn_proba = nn_model(torch.tensor(X_test, dtype=torch.float32).to(DEVICE)).sigmoid().cpu().numpy()
    nn_proba = nn_proba.clip(0, 1)

    meta_X_test = build_meta_features(xgb_model, lgb_model, nn_model, X_test)
    ensemble_proba = meta_model.predict_proba(meta_X_test)[:, 1]

    results = {}
    for name, proba in [("XGBoost", xgb_proba), ("LightGBM", lgb_proba),
                         ("NeuralNet", nn_proba), ("Ensemble", ensemble_proba)]:
        pred_binary = (proba > 0.5).astype(int)
        results[name] = {
            "Accuracy": accuracy_score(y_test, pred_binary),
            "LogLoss": log_loss(y_test, proba),
            "ROC-AUC": roc_auc_score(y_test, proba),
            "Brier": brier_score_loss(y_test, proba),
        }

    results_df = pd.DataFrame(results).T
    print("\n" + results_df.round(4).to_string())

    print(f"\n  Best Model: {results_df['ROC-AUC'].idxmax()} (ROC-AUC: {results_df['ROC-AUC'].max():.4f})")

    os.makedirs(PLOTS_DIR, exist_ok=True)

    plt.figure(figsize=(10, 8))
    for name, proba in [("XGBoost", xgb_proba), ("LightGBM", lgb_proba),
                         ("NeuralNet", nn_proba), ("Ensemble", ensemble_proba)]:
        fpr, tpr, _ = roc_curve(y_test, proba)
        roc_auc = roc_auc_score(y_test, proba)
        plt.plot(fpr, tpr, lw=2, label=f"{name} (AUC = {roc_auc:.3f})")
    plt.plot([0, 1], [0, 1], "k--", lw=1, label="Random")
    plt.xlabel("False Positive Rate", fontsize=12)
    plt.ylabel("True Positive Rate", fontsize=12)
    plt.title("ROC Curves - UFC Fight Prediction Ensemble", fontsize=14)
    plt.legend(loc="lower right", fontsize=10)
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(PLOTS_DIR / "roc_curve.png", dpi=150)
    plt.close()
    print(f"  Saved ROC curve to {PLOTS_DIR / 'roc_curve.png'}")

    plt.figure(figsize=(8, 6))
    cm = confusion_matrix(y_test, (ensemble_proba > 0.5).astype(int))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                xticklabels=["Fighter B Wins", "Fighter A Wins"],
                yticklabels=["Fighter B Wins", "Fighter A Wins"])
    plt.title("Confusion Matrix - Ensemble Model", fontsize=14)
    plt.xlabel("Predicted", fontsize=12)
    plt.ylabel("Actual", fontsize=12)
    plt.tight_layout()
    plt.savefig(PLOTS_DIR / "confusion_matrix.png", dpi=150)
    plt.close()

    return results_df


def compute_shap_values(models, X_sample, feature_names):
    """Compute SHAP values using TreeExplainer for XGBoost and DeepExplainer for NN."""
    print("\n" + "=" * 60)
    print("  Computing SHAP Feature Importance")
    print("=" * 60)

    os.makedirs(PLOTS_DIR, exist_ok=True)

    sample_size = min(200, len(X_sample))
    X_sample = X_sample[:sample_size]

    try:
        print("  Computing TreeExplainer for XGBoost...")
        explainer_xgb = shap.TreeExplainer(models["xgb"])
        shap_values_xgb = explainer_xgb.shap_values(X_sample)

        if isinstance(shap_values_xgb, list):
            shap_values_xgb = shap_values_xgb[1] if len(shap_values_xgb) > 1 else shap_values_xgb[0]

        plt.figure(figsize=(12, 10))
        shap.summary_plot(
            shap_values_xgb, X_sample,
            feature_names=feature_names,
            show=False,
            max_display=20,
        )
        plt.title("XGBoost SHAP Feature Importance", fontsize=14)
        plt.tight_layout()
        plt.savefig(PLOTS_DIR / "shap_xgb_summary.png", dpi=150, bbox_inches="tight")
        plt.close()
        print(f"  Saved XGBoost SHAP plot to {PLOTS_DIR / 'shap_xgb_summary.png'}")

        plt.figure(figsize=(10, 8))
        mean_shap = np.abs(shap_values_xgb).mean(axis=0)
        top_idx = np.argsort(mean_shap)[-20:]
        plt.barh(range(len(top_idx)), mean_shap[top_idx])
        plt.yticks(range(len(top_idx)), [feature_names[i] for i in top_idx])
        plt.xlabel("Mean |SHAP Value|", fontsize=12)
        plt.title("XGBoost Top 20 Features by SHAP", fontsize=14)
        plt.tight_layout()
        plt.savefig(PLOTS_DIR / "shap_xgb_bar.png", dpi=150)
        plt.close()
        print(f"  Saved XGBoost SHAP bar plot to {PLOTS_DIR / 'shap_xgb_bar.png'}")

    except Exception as e:
        print(f"  XGBoost SHAP failed: {e}")

    try:
        print("  Computing TreeExplainer for LightGBM...")
        explainer_lgb = shap.TreeExplainer(models["lgb"])
        shap_values_lgb = explainer_lgb.shap_values(X_sample)

        if isinstance(shap_values_lgb, list):
            shap_values_lgb = shap_values_lgb[1] if len(shap_values_lgb) > 1 else shap_values_lgb[0]

        plt.figure(figsize=(12, 10))
        shap.summary_plot(
            shap_values_lgb, X_sample,
            feature_names=feature_names,
            show=False,
            max_display=20,
        )
        plt.title("LightGBM SHAP Feature Importance", fontsize=14)
        plt.tight_layout()
        plt.savefig(PLOTS_DIR / "shap_lgb_summary.png", dpi=150, bbox_inches="tight")
        plt.close()
        print(f"  Saved LightGBM SHAP plot to {PLOTS_DIR / 'shap_lgb_summary.png'}")

    except Exception as e:
        print(f"  LightGBM SHAP failed: {e}")

    try:
        print("  Computing DeepExplainer for Neural Network...")
        nn_model = models["nn"]
        nn_model.eval()
        background = torch.tensor(X_sample[:100], dtype=torch.float32).to(DEVICE)
        explainer_nn = shap.DeepExplainer(nn_model, background)
        shap_values_nn = explainer_nn.shap_values(
            torch.tensor(X_sample[:100], dtype=torch.float32).to(DEVICE)
        )
        if isinstance(shap_values_nn, list):
            shap_values_nn = shap_values_nn[0]

        shap_values_nn_np = np.array(shap_values_nn)
        if len(shap_values_nn_np.shape) > 2:
            shap_values_nn_np = shap_values_nn_np.reshape(shap_values_nn_np.shape[0], shap_values_nn_np.shape[2])

        plt.figure(figsize=(12, 10))
        shap.summary_plot(
            shap_values_nn_np,
            X_sample[:shap_values_nn_np.shape[0]],
            feature_names=feature_names,
            show=False,
            max_display=20,
        )
        plt.title("Neural Network SHAP Feature Importance", fontsize=14)
        plt.tight_layout()
        plt.savefig(PLOTS_DIR / "shap_nn_summary.png", dpi=150, bbox_inches="tight")
        plt.close()
        print(f"  Saved Neural Network SHAP plot to {PLOTS_DIR / 'shap_nn_summary.png'}")

    except Exception as e:
        print(f"  Neural Network SHAP failed: {e}")

    print("  SHAP analysis complete.")


def save_models(models):
    """Save all trained models to disk."""
    print("\n" + "=" * 60)
    print("  Saving Models")
    print("=" * 60)

    os.makedirs(MODELS_DIR, exist_ok=True)

    models["xgb"].save_model(str(MODEL_PATHS["xgb"]))
    print(f"  XGBoost saved to {MODEL_PATHS['xgb']}")

    models["lgb"].booster_.save_model(str(MODEL_PATHS["lgb"]))
    print(f"  LightGBM saved to {MODEL_PATHS['lgb']}")

    torch.save(models["nn"].state_dict(), MODEL_PATHS["nn"])
    print(f"  Neural Network saved to {MODEL_PATHS['nn']}")

    joblib.dump(models["meta"], MODEL_PATHS["meta"])
    print(f"  Meta-learner saved to {MODEL_PATHS['meta']}")

    print("  All models saved successfully!")


def main():
    os.makedirs(MODELS_DIR, exist_ok=True)
    os.makedirs(PLOTS_DIR, exist_ok=True)

    X_train, X_test, y_train, y_test, feature_names = load_data()

    # Step 1: Hyperparameter tuning on training set via cross-validation
    xgb_best_params = tune_xgboost_cv(X_train, y_train, n_trials=N_TUNE_TRIALS)
    lgb_best_params = tune_lightgbm_cv(X_train, y_train, n_trials=N_TUNE_TRIALS)

    # Step 2: Train final models using train/val split
    val_size = int(len(X_train) * 0.15)
    X_tr, X_val = X_train[:-val_size], X_train[-val_size:]
    y_tr, y_val = y_train[:-val_size], y_train[-val_size:]

    print(f"\n  Training split: {len(X_tr)} train | {len(X_val)} validation | {len(X_test)} test")

    xgb_model = train_xgboost_tuned(X_tr, y_tr, X_val, y_val, xgb_best_params)
    lgb_model = train_lightgbm_tuned(X_tr, y_tr, X_val, y_val, lgb_best_params)
    nn_model = train_neural_network(X_tr, y_tr, X_val, y_val, input_dim=X_train.shape[1])

    print("\n" + "=" * 60)
    print("  Building Meta-Features")
    print("=" * 60)

    meta_X_train = build_meta_features(xgb_model, lgb_model, nn_model, X_train)
    meta_model = train_meta_learner(meta_X_train, y_train)

    models = {"xgb": xgb_model, "lgb": lgb_model, "nn": nn_model, "meta": meta_model}

    results = evaluate_ensemble(models, X_test, y_test, feature_names)

    compute_shap_values(models, X_test, feature_names)

    save_models(models)

    print("\n" + "=" * 60)
    print("  Training Complete! Summary:")
    print("=" * 60)
    print(f"  Ensemble ROC-AUC:  {results['ROC-AUC']['Ensemble']:.4f}")
    print(f"  Ensemble Accuracy: {results['Accuracy']['Ensemble']:.4f}")
    print(f"  Ensemble LogLoss:  {results['LogLoss']['Ensemble']:.4f}")
    print(f"\n  Best single model: {results['ROC-AUC'].idxmax()} ({results['ROC-AUC'].max():.4f})")
    print(f"\n  Next step: python scripts/predict_fight.py -a 'Fighter A' -b 'Fighter B'")

    return models, results


if __name__ == "__main__":
    main()