models/neural_models.py

#!/usr/bin/env python3
"""
NBA Quant AI — Neural Network Models (2025-2026 SOTA)
======================================================
Real, production-grade neural architectures for NBA game prediction.

Models implemented:
  1. LSTMSequenceModel      — Bidirectional LSTM over last N games
  2. TransformerAttentionModel — Self-attention over game history
  3. TabNetModel             — Attention-based tabular learning (Arik & Pfister 2021)
  4. FTTransformerModel      — Feature Tokenizer + Transformer (Gorishniy et al. 2021)
  5. DeepEnsemble            — N independent nets, averaged predictions
  6. ConformalPredictionWrapper — Calibrated prediction intervals (any base model)
  7. AutoGluonEnsemble       — Auto-stacking over hundreds of configs

All models:
  - Handle NaN gracefully (median imputation)
  - Work with 6000+ features
  - Use early stopping
  - CPU-only PyTorch (no CUDA needed)
  - Fit in 16 GB RAM (HF Spaces free tier)

THIS RUNS ON HF SPACES ONLY — NOT ON VM.
"""

from __future__ import annotations

import copy
import json
import math
import os
import pickle
import warnings
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

warnings.filterwarnings("ignore", category=UserWarning)

# ---------------------------------------------------------------------------
# Lazy imports — heavy libraries loaded only when a model is instantiated
# ---------------------------------------------------------------------------

def _import_torch():
    """Import torch lazily to avoid startup cost."""
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader, TensorDataset
    return torch, nn, optim, DataLoader, TensorDataset


# ---------------------------------------------------------------------------
# Base class — common interface for all models
# ---------------------------------------------------------------------------

class BaseNBAModel(ABC):
    """Abstract base for all NBA prediction models."""

    def __init__(self, **params):
        self.params = params
        self._scaler: Optional[StandardScaler] = None
        self._feature_medians: Optional[np.ndarray] = None
        self._is_fitted = False

    # --- public interface ---------------------------------------------------

    @abstractmethod
    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
    ) -> "BaseNBAModel":
        """Train the model. Returns self."""
        ...

    @abstractmethod
    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """Return P(home_win) for each row — shape (n,)."""
        ...

    def get_params(self) -> Dict[str, Any]:
        """Return hyperparameter dict (JSON-serialisable)."""
        return {k: v for k, v in self.params.items() if _is_jsonable(v)}

    def save(self, path: Union[str, Path]) -> None:
        """Persist to disk."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "BaseNBAModel":
        """Load from disk."""
        with open(path, "rb") as f:
            obj = pickle.load(f)
        return obj

    # --- NaN handling & scaling --------------------------------------------

    def _impute(self, X: np.ndarray, fit: bool = False) -> np.ndarray:
        """Replace NaN/Inf with column medians. If *fit*, compute medians first."""
        X = np.array(X, dtype=np.float32)
        X = np.where(np.isfinite(X), X, np.nan)
        if fit:
            self._feature_medians = np.nanmedian(X, axis=0)
            self._feature_medians = np.where(
                np.isfinite(self._feature_medians), self._feature_medians, 0.0
            )
        medians = self._feature_medians if self._feature_medians is not None else np.zeros(X.shape[1])
        inds = np.where(np.isnan(X))
        X[inds] = np.take(medians, inds[1])
        return X

    def _scale(self, X: np.ndarray, fit: bool = False) -> np.ndarray:
        """Standard-scale features."""
        if fit:
            self._scaler = StandardScaler()
            return self._scaler.fit_transform(X).astype(np.float32)
        if self._scaler is not None:
            return self._scaler.transform(X).astype(np.float32)
        return X.astype(np.float32)

    def _prepare(self, X: np.ndarray, fit: bool = False) -> np.ndarray:
        """Impute + scale."""
        X = self._impute(X, fit=fit)
        X = self._scale(X, fit=fit)
        return X

    def _auto_val_split(
        self,
        X: np.ndarray,
        y: np.ndarray,
        X_val: Optional[np.ndarray],
        y_val: Optional[np.ndarray],
        val_frac: float = 0.15,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """If no validation set provided, carve one from the tail (time-ordered)."""
        if X_val is not None and y_val is not None:
            return X, y, X_val, y_val
        split = int(len(X) * (1 - val_frac))
        return X[:split], y[:split], X[split:], y[split:]


# ===========================================================================
# 1. LSTM Game Sequence Model
# ===========================================================================

class LSTMSequenceModel(BaseNBAModel):
    """
    Bidirectional LSTM over the last *seq_len* games of features per team.

    Input shape: (batch, seq_len, n_features)
    Architecture: BiLSTM(128) -> BiLSTM(64) -> Dense(32) -> Sigmoid

    For flat input (n_samples, n_features), the model internally reshapes
    using a sliding window of *seq_len* rows, treating consecutive games as
    the sequence dimension.  For true per-team sequences, pass 3-D arrays
    directly.
    """

    def __init__(
        self,
        seq_len: int = 10,
        hidden1: int = 128,
        hidden2: int = 64,
        dense_dim: int = 32,
        dropout: float = 0.3,
        lr: float = 1e-3,
        weight_decay: float = 1e-5,
        batch_size: int = 256,
        epochs: int = 120,
        patience: int = 15,
        **kw,
    ):
        super().__init__(
            seq_len=seq_len, hidden1=hidden1, hidden2=hidden2,
            dense_dim=dense_dim, dropout=dropout, lr=lr,
            weight_decay=weight_decay, batch_size=batch_size,
            epochs=epochs, patience=patience, **kw,
        )
        self.seq_len = seq_len
        self.hidden1 = hidden1
        self.hidden2 = hidden2
        self.dense_dim = dense_dim
        self.dropout = dropout
        self.lr = lr
        self.weight_decay = weight_decay
        self.batch_size = batch_size
        self.epochs = epochs
        self.patience = patience
        self._net = None

    # --- PyTorch module (defined inside method to keep torch lazy) ----------

    @staticmethod
    def _build_net(n_features: int, cfg: dict):
        torch, nn, _, _, _ = _import_torch()

        class BiLSTMNet(nn.Module):
            def __init__(self):
                super().__init__()
                self.lstm1 = nn.LSTM(
                    input_size=n_features,
                    hidden_size=cfg["hidden1"],
                    batch_first=True,
                    bidirectional=True,
                    dropout=cfg["dropout"] if cfg["hidden2"] else 0,
                )
                self.lstm2 = nn.LSTM(
                    input_size=cfg["hidden1"] * 2,  # bidirectional doubles
                    hidden_size=cfg["hidden2"],
                    batch_first=True,
                    bidirectional=True,
                )
                self.dropout = nn.Dropout(cfg["dropout"])
                self.fc1 = nn.Linear(cfg["hidden2"] * 2, cfg["dense_dim"])
                self.relu = nn.ReLU()
                self.fc2 = nn.Linear(cfg["dense_dim"], 1)

            def forward(self, x):
                # x: (batch, seq_len, features)
                out, _ = self.lstm1(x)
                out = self.dropout(out)
                out, _ = self.lstm2(out)
                # Take last hidden state
                out = out[:, -1, :]
                out = self.dropout(out)
                out = self.relu(self.fc1(out))
                out = self.dropout(out)
                return torch.sigmoid(self.fc2(out)).squeeze(-1)

        return BiLSTMNet()

    # --- Sequence construction from flat arrays ----------------------------

    def _make_sequences(
        self, X: np.ndarray, y: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Convert flat (n_games, n_features) into (n_sequences, seq_len, n_features).
        Uses a sliding window — game i maps to window [i-seq_len+1 .. i].
        The first seq_len-1 games are dropped (not enough history).
        """
        if X.ndim == 3:
            return X, y  # already sequential
        seqs, labels = [], []
        for i in range(self.seq_len - 1, len(X)):
            seqs.append(X[i - self.seq_len + 1 : i + 1])
            labels.append(y[i])
        return np.array(seqs, dtype=np.float32), np.array(labels, dtype=np.float32)

    # --- fit / predict -----------------------------------------------------

    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
    ) -> "LSTMSequenceModel":
        torch, nn, optim, DataLoader, TensorDataset = _import_torch()

        # Prepare
        X_train = self._prepare(X_train, fit=True)
        X_train, y_train, X_val, y_val = self._auto_val_split(X_train, y_train, X_val, y_val)
        if X_val is not None:
            X_val = self._prepare(X_val)

        # Build sequences
        X_tr_seq, y_tr_seq = self._make_sequences(X_train, y_train)
        X_va_seq, y_va_seq = self._make_sequences(X_val, y_val)

        n_features = X_tr_seq.shape[2]
        self._net = self._build_net(n_features, {
            "hidden1": self.hidden1, "hidden2": self.hidden2,
            "dense_dim": self.dense_dim, "dropout": self.dropout,
        })

        optimizer = optim.AdamW(
            self._net.parameters(), lr=self.lr, weight_decay=self.weight_decay
        )
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode="min", factor=0.5, patience=5, min_lr=1e-6
        )
        criterion = nn.BCELoss()

        train_ds = TensorDataset(
            torch.from_numpy(X_tr_seq), torch.from_numpy(y_tr_seq)
        )
        train_dl = DataLoader(train_ds, batch_size=self.batch_size, shuffle=True)

        val_X_t = torch.from_numpy(X_va_seq)
        val_y_t = torch.from_numpy(y_va_seq)

        best_val_loss = float("inf")
        best_state = None
        wait = 0

        self._net.train()
        for epoch in range(self.epochs):
            epoch_loss = 0.0
            for xb, yb in train_dl:
                optimizer.zero_grad()
                preds = self._net(xb)
                loss = criterion(preds, yb)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self._net.parameters(), 1.0)
                optimizer.step()
                epoch_loss += loss.item() * len(xb)
            epoch_loss /= len(train_ds)

            # Validation
            self._net.eval()
            with torch.no_grad():
                val_preds = self._net(val_X_t)
                val_loss = criterion(val_preds, val_y_t).item()
            self._net.train()

            scheduler.step(val_loss)

            if val_loss < best_val_loss - 1e-6:
                best_val_loss = val_loss
                best_state = copy.deepcopy(self._net.state_dict())
                wait = 0
            else:
                wait += 1
                if wait >= self.patience:
                    break

        if best_state is not None:
            self._net.load_state_dict(best_state)
        self._net.eval()
        self._is_fitted = True
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        torch, _, _, _, _ = _import_torch()
        assert self._is_fitted, "Model not fitted yet"

        X = self._prepare(X)
        # If flat, create sequences with padding for early games
        if X.ndim == 2:
            seqs = []
            for i in range(len(X)):
                start = max(0, i - self.seq_len + 1)
                seq = X[start : i + 1]
                if len(seq) < self.seq_len:
                    pad = np.zeros((self.seq_len - len(seq), X.shape[1]), dtype=np.float32)
                    seq = np.concatenate([pad, seq], axis=0)
                seqs.append(seq)
            X_seq = np.array(seqs, dtype=np.float32)
        else:
            X_seq = X.astype(np.float32)

        self._net.eval()
        with torch.no_grad():
            preds = self._net(torch.from_numpy(X_seq))
        return preds.numpy()


# ===========================================================================
# 2. Transformer Attention Model
# ===========================================================================

class TransformerAttentionModel(BaseNBAModel):
    """
    Self-attention over team performance history.

    Architecture:
        Linear projection -> Positional encoding ->
        TransformerEncoder (2 layers, 4 heads) ->
        Global average pool -> Dense -> Sigmoid

    For flat input the model treats each game as one token in a
    sequence of *seq_len* tokens (same sliding-window as LSTM model).
    """

    def __init__(
        self,
        seq_len: int = 10,
        d_model: int = 128,
        n_heads: int = 4,
        n_layers: int = 2,
        dim_ff: int = 256,
        dropout: float = 0.2,
        lr: float = 5e-4,
        weight_decay: float = 1e-4,
        batch_size: int = 256,
        epochs: int = 120,
        patience: int = 15,
        **kw,
    ):
        super().__init__(
            seq_len=seq_len, d_model=d_model, n_heads=n_heads,
            n_layers=n_layers, dim_ff=dim_ff, dropout=dropout,
            lr=lr, weight_decay=weight_decay, batch_size=batch_size,
            epochs=epochs, patience=patience, **kw,
        )
        self.seq_len = seq_len
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.dim_ff = dim_ff
        self.dropout = dropout
        self.lr = lr
        self.weight_decay = weight_decay
        self.batch_size = batch_size
        self.epochs = epochs
        self.patience = patience
        self._net = None

    @staticmethod
    def _build_net(n_features: int, cfg: dict):
        torch, nn, _, _, _ = _import_torch()

        class PositionalEncoding(nn.Module):
            """Sinusoidal positional encoding for game order."""
            def __init__(self, d_model: int, max_len: int = 200):
                super().__init__()
                pe = torch.zeros(max_len, d_model)
                position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
                div_term = torch.exp(
                    torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
                )
                pe[:, 0::2] = torch.sin(position * div_term)
                pe[:, 1::2] = torch.cos(position * div_term[: d_model // 2])  # handle odd d_model
                pe = pe.unsqueeze(0)  # (1, max_len, d_model)
                self.register_buffer("pe", pe)

            def forward(self, x):
                return x + self.pe[:, : x.size(1), :]

        class TransformerNet(nn.Module):
            def __init__(self):
                super().__init__()
                self.input_proj = nn.Linear(n_features, cfg["d_model"])
                self.pos_enc = PositionalEncoding(cfg["d_model"], max_len=cfg["seq_len"] + 10)
                self.layer_norm_in = nn.LayerNorm(cfg["d_model"])
                encoder_layer = nn.TransformerEncoderLayer(
                    d_model=cfg["d_model"],
                    nhead=cfg["n_heads"],
                    dim_feedforward=cfg["dim_ff"],
                    dropout=cfg["dropout"],
                    batch_first=True,
                    activation="gelu",
                )
                self.encoder = nn.TransformerEncoder(
                    encoder_layer, num_layers=cfg["n_layers"]
                )
                self.dropout = nn.Dropout(cfg["dropout"])
                self.fc1 = nn.Linear(cfg["d_model"], cfg["d_model"] // 2)
                self.gelu = nn.GELU()
                self.fc2 = nn.Linear(cfg["d_model"] // 2, 1)

            def forward(self, x):
                # x: (batch, seq_len, n_features)
                x = self.input_proj(x)
                x = self.pos_enc(x)
                x = self.layer_norm_in(x)
                x = self.encoder(x)
                # Global average pooling across sequence dim
                x = x.mean(dim=1)
                x = self.dropout(x)
                x = self.gelu(self.fc1(x))
                x = self.dropout(x)
                return torch.sigmoid(self.fc2(x)).squeeze(-1)

        return TransformerNet()

    def _make_sequences(self, X: np.ndarray, y: np.ndarray):
        if X.ndim == 3:
            return X, y
        seqs, labels = [], []
        for i in range(self.seq_len - 1, len(X)):
            seqs.append(X[i - self.seq_len + 1 : i + 1])
            labels.append(y[i])
        return np.array(seqs, dtype=np.float32), np.array(labels, dtype=np.float32)

    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
    ) -> "TransformerAttentionModel":
        torch, nn, optim, DataLoader, TensorDataset = _import_torch()

        X_train = self._prepare(X_train, fit=True)
        X_train, y_train, X_val, y_val = self._auto_val_split(X_train, y_train, X_val, y_val)
        if X_val is not None:
            X_val = self._prepare(X_val)

        X_tr_seq, y_tr_seq = self._make_sequences(X_train, y_train)
        X_va_seq, y_va_seq = self._make_sequences(X_val, y_val)

        n_features = X_tr_seq.shape[2]
        self._net = self._build_net(n_features, {
            "d_model": self.d_model, "n_heads": self.n_heads,
            "n_layers": self.n_layers, "dim_ff": self.dim_ff,
            "dropout": self.dropout, "seq_len": self.seq_len,
        })

        optimizer = optim.AdamW(
            self._net.parameters(), lr=self.lr, weight_decay=self.weight_decay
        )
        scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
            optimizer, T_0=10, T_mult=2, eta_min=1e-6
        )
        criterion = nn.BCELoss()

        train_ds = TensorDataset(
            torch.from_numpy(X_tr_seq), torch.from_numpy(y_tr_seq)
        )
        train_dl = DataLoader(train_ds, batch_size=self.batch_size, shuffle=True)

        val_X_t = torch.from_numpy(X_va_seq)
        val_y_t = torch.from_numpy(y_va_seq)

        best_val_loss = float("inf")
        best_state = None
        wait = 0

        self._net.train()
        for epoch in range(self.epochs):
            epoch_loss = 0.0
            for xb, yb in train_dl:
                optimizer.zero_grad()
                preds = self._net(xb)
                loss = criterion(preds, yb)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self._net.parameters(), 1.0)
                optimizer.step()
                epoch_loss += loss.item() * len(xb)
            epoch_loss /= len(train_ds)
            scheduler.step(epoch + epoch_loss)  # warm restart input

            self._net.eval()
            with torch.no_grad():
                val_preds = self._net(val_X_t)
                val_loss = criterion(val_preds, val_y_t).item()
            self._net.train()

            if val_loss < best_val_loss - 1e-6:
                best_val_loss = val_loss
                best_state = copy.deepcopy(self._net.state_dict())
                wait = 0
            else:
                wait += 1
                if wait >= self.patience:
                    break

        if best_state is not None:
            self._net.load_state_dict(best_state)
        self._net.eval()
        self._is_fitted = True
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        torch, _, _, _, _ = _import_torch()
        assert self._is_fitted, "Model not fitted yet"

        X = self._prepare(X)
        if X.ndim == 2:
            seqs = []
            for i in range(len(X)):
                start = max(0, i - self.seq_len + 1)
                seq = X[start : i + 1]
                if len(seq) < self.seq_len:
                    pad = np.zeros((self.seq_len - len(seq), X.shape[1]), dtype=np.float32)
                    seq = np.concatenate([pad, seq], axis=0)
                seqs.append(seq)
            X_seq = np.array(seqs, dtype=np.float32)
        else:
            X_seq = X.astype(np.float32)

        self._net.eval()
        with torch.no_grad():
            preds = self._net(torch.from_numpy(X_seq))
        return preds.numpy()


# ===========================================================================
# 3. TabNet — Attention-based Tabular Model
# ===========================================================================

class TabNetModel(BaseNBAModel):
    """
    TabNet (Arik & Pfister 2021) — SOTA attention-based tabular learning.

    Uses sequential attention to select features at each decision step,
    providing built-in interpretability via attention masks.

    Wraps pytorch_tabnet.TabNetClassifier with NaN handling and
    early stopping.
    """

    def __init__(
        self,
        n_d: int = 32,
        n_a: int = 32,
        n_steps: int = 5,
        gamma: float = 1.5,
        lambda_sparse: float = 1e-4,
        n_independent: int = 2,
        n_shared: int = 2,
        lr: float = 2e-2,
        batch_size: int = 1024,
        virtual_batch_size: int = 256,
        epochs: int = 200,
        patience: int = 20,
        mask_type: str = "entmax",
        **kw,
    ):
        super().__init__(
            n_d=n_d, n_a=n_a, n_steps=n_steps, gamma=gamma,
            lambda_sparse=lambda_sparse, n_independent=n_independent,
            n_shared=n_shared, lr=lr, batch_size=batch_size,
            virtual_batch_size=virtual_batch_size, epochs=epochs,
            patience=patience, mask_type=mask_type, **kw,
        )
        self.n_d = n_d
        self.n_a = n_a
        self.n_steps = n_steps
        self.gamma = gamma
        self.lambda_sparse = lambda_sparse
        self.n_independent = n_independent
        self.n_shared = n_shared
        self.lr = lr
        self.batch_size = batch_size
        self.virtual_batch_size = virtual_batch_size
        self.epochs = epochs
        self.patience = patience
        self.mask_type = mask_type
        self._clf = None
        self._feature_importances: Optional[np.ndarray] = None

    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
    ) -> "TabNetModel":
        from pytorch_tabnet.tab_model import TabNetClassifier

        X_train = self._impute(X_train, fit=True)
        X_train, y_train, X_val, y_val = self._auto_val_split(X_train, y_train, X_val, y_val)
        if X_val is not None:
            X_val = self._impute(X_val)

        y_train = y_train.astype(np.int64)
        y_val = y_val.astype(np.int64)

        self._clf = TabNetClassifier(
            n_d=self.n_d,
            n_a=self.n_a,
            n_steps=self.n_steps,
            gamma=self.gamma,
            lambda_sparse=self.lambda_sparse,
            n_independent=self.n_independent,
            n_shared=self.n_shared,
            optimizer_fn=None,  # default Adam
            optimizer_params={"lr": self.lr},
            mask_type=self.mask_type,
            scheduler_fn=None,
            scheduler_params=None,
            verbose=0,
            device_name="cpu",
        )

        self._clf.fit(
            X_train=X_train,
            y_train=y_train,
            eval_set=[(X_val, y_val)],
            eval_name=["val"],
            eval_metric=["logloss"],
            max_epochs=self.epochs,
            patience=self.patience,
            batch_size=self.batch_size,
            virtual_batch_size=min(self.virtual_batch_size, self.batch_size),
            drop_last=False,
        )

        self._feature_importances = self._clf.feature_importances_
        self._is_fitted = True
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        assert self._is_fitted, "Model not fitted yet"
        X = self._impute(X)
        proba = self._clf.predict_proba(X)  # shape (n, 2)
        return proba[:, 1]

    def get_feature_importances(self) -> Optional[np.ndarray]:
        """Return TabNet attention-based feature importances."""
        return self._feature_importances

    def explain(self, X: np.ndarray) -> np.ndarray:
        """Return per-sample feature attention masks."""
        assert self._is_fitted, "Model not fitted yet"
        X = self._impute(X)
        masks, _ = self._clf.explain(X)
        return masks


# ===========================================================================
# 4. FT-Transformer (Feature Tokenizer + Transformer)
# ===========================================================================

class FTTransformerModel(BaseNBAModel):
    """
    FT-Transformer (Gorishniy et al. 2021) — confirmed SOTA for tabular
    data in 2025-2026 benchmarks.

    Each numerical feature is projected into a *d_token*-dimensional embedding.
    A [CLS] token is prepended. Self-attention across all feature tokens
    captures cross-feature interactions.  The [CLS] representation feeds a
    classification head.

    Because the full 6000+ features would create 6000+ tokens (too large for
    self-attention on CPU), we first apply a learned linear bottleneck to
    reduce to *n_tokens* feature groups.
    """

    def __init__(
        self,
        n_tokens: int = 128,
        d_token: int = 64,
        n_heads: int = 4,
        n_layers: int = 3,
        dim_ff: int = 256,
        dropout: float = 0.2,
        attention_dropout: float = 0.1,
        lr: float = 1e-4,
        weight_decay: float = 1e-5,
        batch_size: int = 512,
        epochs: int = 120,
        patience: int = 15,
        **kw,
    ):
        super().__init__(
            n_tokens=n_tokens, d_token=d_token, n_heads=n_heads,
            n_layers=n_layers, dim_ff=dim_ff, dropout=dropout,
            attention_dropout=attention_dropout, lr=lr,
            weight_decay=weight_decay, batch_size=batch_size,
            epochs=epochs, patience=patience, **kw,
        )
        self.n_tokens = n_tokens
        self.d_token = d_token
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.dim_ff = dim_ff
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.lr = lr
        self.weight_decay = weight_decay
        self.batch_size = batch_size
        self.epochs = epochs
        self.patience = patience
        self._net = None

    @staticmethod
    def _build_net(n_features: int, cfg: dict):
        torch, nn, _, _, _ = _import_torch()

        class FTTransformerNet(nn.Module):
            """
            Feature Tokenizer + Transformer.

            1) Bottleneck: Linear(n_features -> n_tokens)  — group features
            2) Token embed: each of *n_tokens* scalars -> d_token vector
            3) Prepend [CLS] token
            4) TransformerEncoder
            5) [CLS] output -> classification head
            """

            def __init__(self):
                super().__init__()
                n_tok = cfg["n_tokens"]
                d_tok = cfg["d_token"]

                # Bottleneck projection: reduce 6000 features to n_tokens groups
                self.bottleneck = nn.Linear(n_features, n_tok)
                self.bn_norm = nn.LayerNorm(n_tok)

                # Per-token embedding: each scalar -> d_token vector
                # Implemented as a shared Linear(1 -> d_token) + per-token bias
                self.token_weight = nn.Parameter(torch.randn(n_tok, d_tok) * 0.02)
                self.token_bias = nn.Parameter(torch.zeros(n_tok, d_tok))

                # [CLS] token
                self.cls_token = nn.Parameter(torch.randn(1, 1, d_tok) * 0.02)

                # Transformer
                self.layer_norm = nn.LayerNorm(d_tok)
                encoder_layer = nn.TransformerEncoderLayer(
                    d_model=d_tok,
                    nhead=cfg["n_heads"],
                    dim_feedforward=cfg["dim_ff"],
                    dropout=cfg["dropout"],
                    batch_first=True,
                    activation="gelu",
                )
                self.encoder = nn.TransformerEncoder(
                    encoder_layer, num_layers=cfg["n_layers"]
                )

                # Head
                self.head = nn.Sequential(
                    nn.LayerNorm(d_tok),
                    nn.Linear(d_tok, d_tok // 2),
                    nn.GELU(),
                    nn.Dropout(cfg["dropout"]),
                    nn.Linear(d_tok // 2, 1),
                )

            def forward(self, x):
                # x: (batch, n_features)
                batch_size = x.size(0)

                # Bottleneck: (batch, n_features) -> (batch, n_tokens)
                x = self.bn_norm(self.bottleneck(x))

                # Token embedding: (batch, n_tokens) -> (batch, n_tokens, d_token)
                # x_i * weight_i + bias_i for each token
                x = x.unsqueeze(-1) * self.token_weight.unsqueeze(0) + self.token_bias.unsqueeze(0)

                # Prepend [CLS]
                cls = self.cls_token.expand(batch_size, -1, -1)
                x = torch.cat([cls, x], dim=1)  # (batch, 1 + n_tokens, d_token)

                x = self.layer_norm(x)
                x = self.encoder(x)

                # Extract [CLS] output
                cls_out = x[:, 0, :]
                return torch.sigmoid(self.head(cls_out)).squeeze(-1)

        return FTTransformerNet()

    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
    ) -> "FTTransformerModel":
        torch, nn, optim, DataLoader, TensorDataset = _import_torch()

        X_train = self._prepare(X_train, fit=True)
        X_train, y_train, X_val, y_val = self._auto_val_split(X_train, y_train, X_val, y_val)
        if X_val is not None:
            X_val = self._prepare(X_val)

        y_train = y_train.astype(np.float32)
        y_val = y_val.astype(np.float32)

        n_features = X_train.shape[1]
        self._net = self._build_net(n_features, {
            "n_tokens": min(self.n_tokens, n_features),
            "d_token": self.d_token,
            "n_heads": self.n_heads,
            "n_layers": self.n_layers,
            "dim_ff": self.dim_ff,
            "dropout": self.dropout,
        })

        optimizer = optim.AdamW(
            self._net.parameters(), lr=self.lr, weight_decay=self.weight_decay
        )
        scheduler = optim.lr_scheduler.OneCycleLR(
            optimizer, max_lr=self.lr * 10, total_steps=self.epochs,
            pct_start=0.1, anneal_strategy="cos",
        )
        criterion = nn.BCELoss()

        train_ds = TensorDataset(
            torch.from_numpy(X_train), torch.from_numpy(y_train)
        )
        train_dl = DataLoader(train_ds, batch_size=self.batch_size, shuffle=True)

        val_X_t = torch.from_numpy(X_val)
        val_y_t = torch.from_numpy(y_val)

        best_val_loss = float("inf")
        best_state = None
        wait = 0

        self._net.train()
        for epoch in range(self.epochs):
            epoch_loss = 0.0
            for xb, yb in train_dl:
                optimizer.zero_grad()
                preds = self._net(xb)
                loss = criterion(preds, yb)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self._net.parameters(), 1.0)
                optimizer.step()
                epoch_loss += loss.item() * len(xb)
            epoch_loss /= len(train_ds)
            scheduler.step()

            self._net.eval()
            with torch.no_grad():
                val_preds = self._net(val_X_t)
                val_loss = criterion(val_preds, val_y_t).item()
            self._net.train()

            if val_loss < best_val_loss - 1e-6:
                best_val_loss = val_loss
                best_state = copy.deepcopy(self._net.state_dict())
                wait = 0
            else:
                wait += 1
                if wait >= self.patience:
                    break

        if best_state is not None:
            self._net.load_state_dict(best_state)
        self._net.eval()
        self._is_fitted = True
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        torch, _, _, _, _ = _import_torch()
        assert self._is_fitted, "Model not fitted yet"

        X = self._prepare(X)
        X_t = torch.from_numpy(X)

        self._net.eval()
        # Batch to avoid OOM on large inputs
        preds_list = []
        bs = self.batch_size
        for i in range(0, len(X_t), bs):
            with torch.no_grad():
                p = self._net(X_t[i : i + bs])
            preds_list.append(p.numpy())
        return np.concatenate(preds_list)


# ===========================================================================
# 5. Deep Ensemble
# ===========================================================================

class DeepEnsemble(BaseNBAModel):
    """
    Train N independent neural networks with different random seeds.

    Average their predictions for:
      - Better calibration (ensemble smoothing)
      - Uncertainty estimation (prediction variance)

    Each member is a simple but effective MLP with skip connections (ResNet-style),
    which is the 2025 consensus best architecture for tabular deep learning
    when ensembled (Kadra et al. 2021 "Well-Tuned Simple Nets").
    """

    def __init__(
        self,
        n_members: int = 10,
        hidden_dims: Tuple[int, ...] = (512, 256, 128),
        dropout: float = 0.3,
        lr: float = 1e-3,
        weight_decay: float = 1e-4,
        batch_size: int = 512,
        epochs: int = 100,
        patience: int = 12,
        **kw,
    ):
        super().__init__(
            n_members=n_members, hidden_dims=list(hidden_dims),
            dropout=dropout, lr=lr, weight_decay=weight_decay,
            batch_size=batch_size, epochs=epochs, patience=patience, **kw,
        )
        self.n_members = n_members
        self.hidden_dims = hidden_dims
        self.dropout = dropout
        self.lr = lr
        self.weight_decay = weight_decay
        self.batch_size = batch_size
        self.epochs = epochs
        self.patience = patience
        self._members: List = []

    @staticmethod
    def _build_mlp(n_features: int, hidden_dims: Tuple[int, ...], dropout: float, seed: int):
        """Build one ResNet-style MLP member."""
        torch, nn, _, _, _ = _import_torch()
        torch.manual_seed(seed)

        class ResBlock(nn.Module):
            """Pre-activation residual block."""
            def __init__(self, dim: int, drop: float):
                super().__init__()
                self.net = nn.Sequential(
                    nn.LayerNorm(dim),
                    nn.GELU(),
                    nn.Linear(dim, dim),
                    nn.Dropout(drop),
                    nn.LayerNorm(dim),
                    nn.GELU(),
                    nn.Linear(dim, dim),
                    nn.Dropout(drop),
                )

            def forward(self, x):
                return x + self.net(x)

        layers = []
        in_dim = n_features
        for h_dim in hidden_dims:
            layers.append(nn.Linear(in_dim, h_dim))
            layers.append(nn.GELU())
            layers.append(nn.Dropout(dropout))
            # Add residual block at each hidden layer
            layers.append(ResBlock(h_dim, dropout))
            in_dim = h_dim
        layers.append(nn.Linear(in_dim, 1))

        class EnsembleMLP(nn.Module):
            def __init__(self, layer_list):
                super().__init__()
                self.net = nn.Sequential(*layer_list)

            def forward(self, x):
                return torch.sigmoid(self.net(x)).squeeze(-1)

        return EnsembleMLP(layers)

    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
    ) -> "DeepEnsemble":
        torch, nn, optim, DataLoader, TensorDataset = _import_torch()

        X_train = self._prepare(X_train, fit=True)
        X_train, y_train, X_val, y_val = self._auto_val_split(X_train, y_train, X_val, y_val)
        if X_val is not None:
            X_val = self._prepare(X_val)

        y_train = y_train.astype(np.float32)
        y_val = y_val.astype(np.float32)
        n_features = X_train.shape[1]

        val_X_t = torch.from_numpy(X_val)
        val_y_t = torch.from_numpy(y_val)
        criterion = nn.BCELoss()

        self._members = []
        for member_idx in range(self.n_members):
            seed = 42 + member_idx * 1337
            net = self._build_mlp(n_features, self.hidden_dims, self.dropout, seed)

            # Each member gets a different random seed for data shuffling too
            torch.manual_seed(seed)
            np.random.seed(seed)

            optimizer = optim.AdamW(
                net.parameters(), lr=self.lr, weight_decay=self.weight_decay
            )
            scheduler = optim.lr_scheduler.ReduceLROnPlateau(
                optimizer, mode="min", factor=0.5, patience=5, min_lr=1e-6
            )

            train_ds = TensorDataset(
                torch.from_numpy(X_train), torch.from_numpy(y_train)
            )
            train_dl = DataLoader(train_ds, batch_size=self.batch_size, shuffle=True)

            best_val_loss = float("inf")
            best_state = None
            wait = 0

            net.train()
            for epoch in range(self.epochs):
                for xb, yb in train_dl:
                    optimizer.zero_grad()
                    preds = net(xb)
                    loss = criterion(preds, yb)
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
                    optimizer.step()

                net.eval()
                with torch.no_grad():
                    vp = net(val_X_t)
                    vl = criterion(vp, val_y_t).item()
                net.train()
                scheduler.step(vl)

                if vl < best_val_loss - 1e-6:
                    best_val_loss = vl
                    best_state = copy.deepcopy(net.state_dict())
                    wait = 0
                else:
                    wait += 1
                    if wait >= self.patience:
                        break

            if best_state is not None:
                net.load_state_dict(best_state)
            net.eval()
            self._members.append(net)

        self._is_fitted = True
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """Return mean prediction across ensemble members."""
        torch, _, _, _, _ = _import_torch()
        assert self._is_fitted and self._members, "Model not fitted yet"

        X = self._prepare(X)
        X_t = torch.from_numpy(X)

        all_preds = []
        for net in self._members:
            net.eval()
            with torch.no_grad():
                p = net(X_t).numpy()
            all_preds.append(p)

        return np.mean(all_preds, axis=0)

    def predict_uncertainty(self, X: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Return (mean_prediction, std_prediction) across ensemble members.
        High std = high model uncertainty = less confident prediction.
        """
        torch, _, _, _, _ = _import_torch()
        assert self._is_fitted and self._members, "Model not fitted yet"

        X = self._prepare(X)
        X_t = torch.from_numpy(X)

        all_preds = []
        for net in self._members:
            net.eval()
            with torch.no_grad():
                p = net(X_t).numpy()
            all_preds.append(p)

        stacked = np.array(all_preds)  # (n_members, n_samples)
        return stacked.mean(axis=0), stacked.std(axis=0)


# ===========================================================================
# 6. Conformal Prediction Wrapper
# ===========================================================================

class ConformalPredictionWrapper(BaseNBAModel):
    """
    Wraps ANY model to provide calibrated prediction intervals with
    guaranteed coverage.

    Uses split conformal prediction:
      1. Train base model on training set
      2. Compute non-conformity scores on calibration holdout
      3. At inference, use quantile of scores to produce prediction sets

    For binary classification:
      - Returns P(home_win) from base model (point prediction)
      - Also provides prediction_set() that returns {0}, {1}, or {0,1}
        with guaranteed marginal coverage >= (1 - alpha)
    """

    def __init__(
        self,
        base_model: BaseNBAModel,
        alpha: float = 0.10,
        cal_fraction: float = 0.20,
        **kw,
    ):
        super().__init__(alpha=alpha, cal_fraction=cal_fraction, **kw)
        self.base_model = base_model
        self.alpha = alpha
        self.cal_fraction = cal_fraction
        self._qhat: Optional[float] = None
        self._cal_scores: Optional[np.ndarray] = None

    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
    ) -> "ConformalPredictionWrapper":
        """
        Split data into proper-training and calibration sets.
        Train base model on proper-training, compute conformal scores on calibration.
        """
        n = len(X_train)
        cal_size = int(n * self.cal_fraction)
        # Use the LAST cal_size samples for calibration (time-ordered)
        X_proper = X_train[: n - cal_size]
        y_proper = y_train[: n - cal_size]
        X_cal = X_train[n - cal_size :]
        y_cal = y_train[n - cal_size :]

        # Train base model
        self.base_model.fit(X_proper, y_proper, X_val, y_val)

        # Compute non-conformity scores on calibration set
        cal_probs = self.base_model.predict_proba(X_cal)
        # Score = 1 - P(true_class)
        scores = np.where(y_cal == 1, 1.0 - cal_probs, cal_probs)
        self._cal_scores = np.sort(scores)

        # Quantile for desired coverage
        n_cal = len(self._cal_scores)
        level = np.ceil((1.0 - self.alpha) * (n_cal + 1)) / n_cal
        level = min(level, 1.0)
        self._qhat = np.quantile(self._cal_scores, level, method="higher")

        self._is_fitted = True
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """Return point predictions from base model."""
        assert self._is_fitted, "Model not fitted yet"
        return self.base_model.predict_proba(X)

    def predict_sets(self, X: np.ndarray) -> List[set]:
        """
        Return prediction sets with guaranteed (1-alpha) coverage.

        Each set is one of:
          - {1}    — confident home win
          - {0}    — confident away win
          - {0, 1} — uncertain (both plausible)
        """
        assert self._is_fitted, "Model not fitted yet"
        probs = self.base_model.predict_proba(X)
        sets = []
        for p in probs:
            s = set()
            # Include class 1 if score would be <= qhat
            if 1.0 - p <= self._qhat:
                s.add(1)
            # Include class 0 if score would be <= qhat
            if p <= self._qhat:
                s.add(0)
            if not s:
                # Shouldn't happen, but include most likely
                s.add(1 if p >= 0.5 else 0)
            sets.append(s)
        return sets

    def predict_intervals(self, X: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Return (lower_bound, upper_bound) calibrated probability intervals.

        Width of interval reflects model uncertainty after conformal calibration.
        """
        assert self._is_fitted, "Model not fitted yet"
        probs = self.base_model.predict_proba(X)
        lower = np.clip(probs - self._qhat, 0.0, 1.0)
        upper = np.clip(probs + self._qhat, 0.0, 1.0)
        return lower, upper

    def get_params(self) -> Dict[str, Any]:
        base_params = self.base_model.get_params()
        return {
            "wrapper": "conformal",
            "alpha": self.alpha,
            "cal_fraction": self.cal_fraction,
            "qhat": float(self._qhat) if self._qhat is not None else None,
            "base_model": base_params,
        }


# ===========================================================================
# 7. AutoGluon Ensemble
# ===========================================================================

class AutoGluonEnsemble(BaseNBAModel):
    """
    AutoGluon Tabular — auto-search and stack hundreds of model configurations.

    Time-budgeted: runs for *max_time* seconds, tries GBMs, neural nets,
    linear models, k-NN, then stacks the best ones.

    Presets: "best_quality" = maximum stacking/bagging (slow but best),
             "good_quality" = reasonable speed/quality trade-off,
             "medium_quality" = fastest.
    """

    def __init__(
        self,
        max_time: int = 3600,
        preset: str = "best_quality",
        eval_metric: str = "log_loss",
        num_bag_folds: int = 5,
        num_stack_levels: int = 1,
        verbosity: int = 1,
        **kw,
    ):
        super().__init__(
            max_time=max_time, preset=preset, eval_metric=eval_metric,
            num_bag_folds=num_bag_folds, num_stack_levels=num_stack_levels,
            verbosity=verbosity, **kw,
        )
        self.max_time = max_time
        self.preset = preset
        self.eval_metric = eval_metric
        self.num_bag_folds = num_bag_folds
        self.num_stack_levels = num_stack_levels
        self.verbosity = verbosity
        self._predictor = None

    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
    ) -> "AutoGluonEnsemble":
        try:
            from autogluon.tabular import TabularPredictor
            import pandas as pd
        except ImportError:
            raise ImportError(
                "autogluon.tabular not installed. Install with: "
                "pip install autogluon.tabular"
            )

        X_train = self._impute(X_train, fit=True)

        # Build DataFrame with feature columns + label
        n_features = X_train.shape[1]
        col_names = [f"f_{i}" for i in range(n_features)]
        df_train = pd.DataFrame(X_train, columns=col_names)
        df_train["label"] = y_train.astype(int)

        # Validation data (optional tuning set)
        df_val = None
        if X_val is not None and y_val is not None:
            X_val = self._impute(X_val)
            df_val = pd.DataFrame(X_val, columns=col_names)
            df_val["label"] = y_val.astype(int)

        self._col_names = col_names

        self._predictor = TabularPredictor(
            label="label",
            eval_metric=self.eval_metric,
            problem_type="binary",
            verbosity=self.verbosity,
        )

        fit_kwargs = {
            "train_data": df_train,
            "time_limit": self.max_time,
            "presets": self.preset,
            "num_bag_folds": self.num_bag_folds,
            "num_stack_levels": self.num_stack_levels,
        }
        if df_val is not None:
            fit_kwargs["tuning_data"] = df_val

        self._predictor.fit(**fit_kwargs)
        self._is_fitted = True
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        import pandas as pd

        assert self._is_fitted, "Model not fitted yet"
        X = self._impute(X)
        df = pd.DataFrame(X, columns=self._col_names)
        proba = self._predictor.predict_proba(df)
        # Returns DataFrame with columns 0, 1 — we want P(class=1)
        if isinstance(proba, pd.DataFrame):
            return proba[1].values
        return proba

    def leaderboard(self):
        """Return AutoGluon model leaderboard."""
        assert self._is_fitted, "Model not fitted yet"
        return self._predictor.leaderboard(silent=True)

    def feature_importance(self, X: np.ndarray, y: np.ndarray) -> "pd.DataFrame":
        """Return permutation feature importance."""
        import pandas as pd

        X = self._impute(X)
        df = pd.DataFrame(X, columns=self._col_names)
        df["label"] = y.astype(int)
        return self._predictor.feature_importance(df)

    def save(self, path: Union[str, Path]) -> None:
        """AutoGluon has its own save mechanism."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)
        if self._predictor is not None:
            self._predictor.save(str(path / "autogluon_predictor"))
        # Save wrapper state
        state = {
            "params": self.params,
            "_col_names": getattr(self, "_col_names", None),
            "_feature_medians": self._feature_medians.tolist() if self._feature_medians is not None else None,
            "_is_fitted": self._is_fitted,
        }
        with open(path / "wrapper_state.json", "w") as f:
            json.dump(state, f)

    @classmethod
    def load(cls, path: Union[str, Path]) -> "AutoGluonEnsemble":
        from autogluon.tabular import TabularPredictor

        path = Path(path)
        with open(path / "wrapper_state.json") as f:
            state = json.load(f)

        obj = cls(**state["params"])
        obj._col_names = state["_col_names"]
        if state["_feature_medians"] is not None:
            obj._feature_medians = np.array(state["_feature_medians"], dtype=np.float32)
        obj._predictor = TabularPredictor.load(str(path / "autogluon_predictor"))
        obj._is_fitted = state["_is_fitted"]
        return obj


# ===========================================================================
# Utilities
# ===========================================================================

def _is_jsonable(v: Any) -> bool:
    """Check if a value is JSON serialisable."""
    try:
        json.dumps(v)
        return True
    except (TypeError, OverflowError, ValueError):
        return False


# ===========================================================================
# Model Registry — maps names to classes for the genetic algorithm
# ===========================================================================

NEURAL_MODEL_REGISTRY: Dict[str, type] = {
    "lstm": LSTMSequenceModel,
    "transformer": TransformerAttentionModel,
    "tabnet": TabNetModel,
    "ft_transformer": FTTransformerModel,
    "deep_ensemble": DeepEnsemble,
    "conformal": ConformalPredictionWrapper,
    "autogluon": AutoGluonEnsemble,
}


def build_neural_model(model_type: str, **params) -> BaseNBAModel:
    """
    Factory function to build a neural model by name.

    Usage:
        model = build_neural_model("ft_transformer", n_tokens=128, d_token=64)
        model.fit(X_train, y_train)
        probs = model.predict_proba(X_test)

    For conformal wrapper, pass base_model_type and base_model_params:
        model = build_neural_model(
            "conformal",
            base_model_type="deep_ensemble",
            base_model_params={"n_members": 5},
            alpha=0.1,
        )
    """
    if model_type == "conformal":
        base_type = params.pop("base_model_type", "deep_ensemble")
        base_params = params.pop("base_model_params", {})
        base_model = build_neural_model(base_type, **base_params)
        return ConformalPredictionWrapper(base_model=base_model, **params)

    cls = NEURAL_MODEL_REGISTRY.get(model_type)
    if cls is None:
        raise ValueError(
            f"Unknown model type '{model_type}'. "
            f"Available: {list(NEURAL_MODEL_REGISTRY.keys())}"
        )
    return cls(**params)


# ===========================================================================
# Quick smoke test (runs if executed directly)
# ===========================================================================

if __name__ == "__main__":
    print("=" * 60)
    print("NBA Quant AI — Neural Models Smoke Test")
    print("=" * 60)

    np.random.seed(42)
    N_TRAIN, N_TEST, N_FEAT = 500, 100, 200

    X_train = np.random.randn(N_TRAIN, N_FEAT).astype(np.float32)
    # Inject some NaNs to test imputation
    mask = np.random.random(X_train.shape) < 0.05
    X_train[mask] = np.nan
    y_train = (np.random.random(N_TRAIN) > 0.5).astype(np.float32)

    X_test = np.random.randn(N_TEST, N_FEAT).astype(np.float32)
    y_test = (np.random.random(N_TEST) > 0.5).astype(np.float32)

    # Test each model (with small configs for speed)
    tests = [
        ("FT-Transformer", FTTransformerModel(
            n_tokens=32, d_token=16, n_heads=2, n_layers=1,
            epochs=5, patience=3, batch_size=128,
        )),
        ("Deep Ensemble (3 members)", DeepEnsemble(
            n_members=3, hidden_dims=(64, 32),
            epochs=5, patience=3, batch_size=128,
        )),
        ("LSTM Sequence", LSTMSequenceModel(
            seq_len=5, hidden1=32, hidden2=16, dense_dim=16,
            epochs=5, patience=3, batch_size=128,
        )),
        ("Transformer Attention", TransformerAttentionModel(
            seq_len=5, d_model=32, n_heads=2, n_layers=1,
            dim_ff=64, epochs=5, patience=3, batch_size=128,
        )),
    ]

    for name, model in tests:
        print(f"\n--- {name} ---")
        try:
            model.fit(X_train, y_train)
            probs = model.predict_proba(X_test)
            print(f"  Predictions shape: {probs.shape}")
            print(f"  Mean pred: {probs.mean():.4f}, Std: {probs.std():.4f}")
            print(f"  Min: {probs.min():.4f}, Max: {probs.max():.4f}")
            print(f"  Params: {list(model.get_params().keys())}")
        except Exception as e:
            print(f"  ERROR: {e}")

    # Test conformal wrapper
    print("\n--- Conformal Prediction Wrapper ---")
    try:
        base = DeepEnsemble(
            n_members=2, hidden_dims=(64, 32),
            epochs=5, patience=3, batch_size=128,
        )
        conformal = ConformalPredictionWrapper(base_model=base, alpha=0.1)
        conformal.fit(X_train, y_train)
        probs = conformal.predict_proba(X_test)
        sets = conformal.predict_sets(X_test)
        lower, upper = conformal.predict_intervals(X_test)
        print(f"  Point preds shape: {probs.shape}")
        print(f"  Prediction sets (first 5): {sets[:5]}")
        print(f"  Intervals: [{lower[:3]}] - [{upper[:3]}]")
        print(f"  Avg interval width: {(upper - lower).mean():.4f}")
    except Exception as e:
        print(f"  ERROR: {e}")

    # Test TabNet (may fail if pytorch_tabnet not installed)
    print("\n--- TabNet ---")
    try:
        tab = TabNetModel(
            n_d=8, n_a=8, n_steps=3, epochs=5, patience=3, batch_size=128,
        )
        tab.fit(X_train, y_train)
        probs = tab.predict_proba(X_test)
        print(f"  Predictions shape: {probs.shape}")
        print(f"  Mean pred: {probs.mean():.4f}")
        fi = tab.get_feature_importances()
        if fi is not None:
            print(f"  Feature importances shape: {fi.shape}")
    except ImportError:
        print("  SKIPPED (pytorch_tabnet not installed)")
    except Exception as e:
        print(f"  ERROR: {e}")

    # Test factory
    print("\n--- Factory: build_neural_model ---")
    try:
        m = build_neural_model("ft_transformer", n_tokens=32, d_token=16,
                               n_heads=2, n_layers=1, epochs=3, batch_size=128)
        m.fit(X_train, y_train)
        print(f"  Factory FT-Transformer OK, preds mean: {m.predict_proba(X_test).mean():.4f}")
    except Exception as e:
        print(f"  ERROR: {e}")

    print("\n" + "=" * 60)
    print("Smoke test complete.")
    print("=" * 60)