src/train_transformer_models.py

"""
Training script for transformer-based models (BioGPT, Clinical-T5, PubMedBERT)
on the PPMI dataset with RAG integration, CUDA acceleration, and leak-free patient split.

Usage:
    cd src
    python train_transformer_models.py
"""

import sys
import argparse
import os
from dataclasses import dataclass
from pathlib import Path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import torch
import numpy as np
from torch.utils.data import DataLoader, Subset
from sklearn.utils.class_weight import compute_class_weight
from sklearn.metrics import (
    classification_report, confusion_matrix,
    f1_score, precision_score, recall_score, roc_auc_score
)
from sklearn.model_selection import StratifiedShuffleSplit
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import joblib
import time
import json
from tqdm import tqdm

from data_preprocessing import DataPreprocessor
from models.transformer_models import TabularDataset
from models.medical_transformers import (
    BioMistralClassifier as BioGPTForTabular,
    ClinicalT5Classifier as ClinicalT5ForTabular,
    PubMedBERTClassifier as PubMedBERTForTabular,
)
from document_manager import DocumentManager


# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
ROOT = Path(__file__).resolve().parents[1]
LEAK_FREE_SPLIT_PATH = ROOT / "evaluation_results" / "leak_free_split.npz"
LEAK_FREE_META_PATH = ROOT / "evaluation_results" / "leak_free_split_meta.joblib"
MODEL_DIR = ROOT / "models" / "saved"
RESULTS_DIR = ROOT / "evaluation_results"
PLOTS_DIR = ROOT / "evaluation_results" / "transformer_plots"

# [CONFIG] Set to True if you want to use RAG (slower start), False for faster training
USE_RAG = True
REQUIRE_CUDA = os.getenv("PD_ALLOW_CPU_TRANSFORMERS", "0") != "1"



# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _load_or_create_leak_free_split(preprocessor, file_paths):
    """Load the cached leak-free split or regenerate it if missing."""
    if LEAK_FREE_SPLIT_PATH.exists() and LEAK_FREE_META_PATH.exists():
        split = np.load(LEAK_FREE_SPLIT_PATH)
        meta = joblib.load(LEAK_FREE_META_PATH)
        feature_names = meta.get("feature_names") if isinstance(meta, dict) else None
        class_mapping = meta.get("class_mapping") if isinstance(meta, dict) else None
        print("[DATA] Loaded cached leak-free split from evaluation_results.")
        return (
            split["X_train"], split["X_test"],
            split["y_train"], split["y_test"],
            feature_names, class_mapping,
        )

    print("[DATA] Leak-free split not found – regenerating via DataPreprocessor ...")
    X_train, X_test, y_train, y_test = preprocessor.prepare_data(
        file_paths, test_size=0.2, use_patient_split=True,
    )
    feature_names = preprocessor.get_feature_names()
    class_mapping = preprocessor.get_class_mapping()

    LEAK_FREE_SPLIT_PATH.parent.mkdir(parents=True, exist_ok=True)
    np.savez(
        LEAK_FREE_SPLIT_PATH,
        X_train=X_train, X_test=X_test,
        y_train=y_train, y_test=y_test,
    )
    joblib.dump(
        {"feature_names": feature_names, "class_mapping": class_mapping},
        LEAK_FREE_META_PATH,
    )
    print(f"[DATA] Saved fresh leak-free split → {LEAK_FREE_SPLIT_PATH}")
    return X_train, X_test, y_train, y_test, feature_names, class_mapping


def _stratified_val_indices(y_train, val_fraction=0.15, seed=42):
    """Return (train_idx, val_idx) using a stratified split so every class
    is proportionally represented in the validation set."""
    sss = StratifiedShuffleSplit(n_splits=1, test_size=val_fraction, random_state=seed)
    train_idx, val_idx = next(sss.split(np.zeros(len(y_train)), y_train))
    return train_idx, val_idx


def _prepare_batch(batch, device):
    """Move tensors to the target device and keep optional RAG contexts aligned."""
    if len(batch) == 3:
        data, targets, contexts = batch
    else:
        data, targets = batch
        contexts = None
    data = data.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)
    if contexts is not None:
        contexts = list(contexts)
    return data, targets, contexts


def _build_context_cache(features, build_fn, split_name="dataset", cache_path=None):
    """Pre-compute RAG contexts with parallel processing and caching."""
    if not USE_RAG:
        return [""] * len(features)

    if cache_path and os.path.exists(cache_path):
        print(f"  [RAG] Loading cached contexts from {cache_path}")
        return joblib.load(cache_path)

    print(f"  [RAG] Generating contexts for {len(features)} samples (Parallel)...")
    from joblib import Parallel, delayed
    
    # Run in parallel to speed up regex/cosine-sim
    contexts = Parallel(n_jobs=-1, verbose=5)(
        delayed(build_fn)(row) for row in features
    )
    
    if cache_path:
        joblib.dump(contexts, cache_path)
        print(f"  [RAG] Saved contexts to {cache_path}")
        
    return contexts



def _print_gpu_info(device):
    """Print GPU diagnostics."""
    if device.type != "cuda":
        return
    print(f"  GPU Name       : {torch.cuda.get_device_name(0)}")
    print(f"  CUDA Version   : {torch.version.cuda}")
    cap = torch.cuda.get_device_capability(0)
    print(f"  Compute Cap.   : {cap[0]}.{cap[1]}")
    mem_total = torch.cuda.get_device_properties(0).total_memory / 1024**3
    print(f"  Total VRAM     : {mem_total:.1f} GB")
    print(f"  cuDNN Enabled  : {torch.backends.cudnn.enabled}")
    print(f"  cuDNN Benchmark: {torch.backends.cudnn.benchmark}")


def _ensure_transformer_cuda() -> None:
    if not torch.cuda.is_available():
        if REQUIRE_CUDA:
            raise RuntimeError(
                "CUDA is required for transformer training in this accuracy-oriented configuration. "
                "Install a CUDA-enabled PyTorch build and use a GPU, or set PD_ALLOW_CPU_TRANSFORMERS=1 to explicitly allow CPU fallback."
            )
        return
    torch.set_float32_matmul_precision("highest")
    if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
        torch.backends.cuda.matmul.allow_tf32 = False
    if hasattr(torch.backends, "cudnn"):
        torch.backends.cudnn.allow_tf32 = False
        torch.backends.cudnn.benchmark = True


@dataclass(frozen=True)
class GPUExecutionProfile:
    name: str
    train_batch_by_model: dict
    eval_batch_by_model: dict
    grad_accum_by_model: dict
    num_workers: int
    prefetch_factor: int
    persistent_workers: bool
    notes: str


def _detect_gpu_execution_profile():
    if not torch.cuda.is_available():
        return GPUExecutionProfile(
            name="cpu",
            train_batch_by_model={"pubmedbert": 4, "biogpt": 2, "clinical_t5": 2},
            eval_batch_by_model={"pubmedbert": 8, "biogpt": 4, "clinical_t5": 4},
            grad_accum_by_model={"pubmedbert": 12, "biogpt": 16, "clinical_t5": 16},
            num_workers=0,
            prefetch_factor=2,
            persistent_workers=False,
            notes="CPU fallback profile.",
        )

    gpu_name = torch.cuda.get_device_name(0).lower()
    memory_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3

    if "a4000" in gpu_name or memory_gb >= 15.0:
        return GPUExecutionProfile(
            name="rtx-a4000",
            train_batch_by_model={"pubmedbert": 16, "biogpt": 10, "clinical_t5": 10},
            eval_batch_by_model={"pubmedbert": 48, "biogpt": 24, "clinical_t5": 24},
            grad_accum_by_model={"pubmedbert": 4, "biogpt": 6, "clinical_t5": 6},
            num_workers=4,
            prefetch_factor=2,
            persistent_workers=True,
            notes="Optimized for RTX A4000 / ~16 GB VRAM.",
        )

    if memory_gb >= 11.0:
        return GPUExecutionProfile(
            name="high-vram",
            train_batch_by_model={"pubmedbert": 12, "biogpt": 8, "clinical_t5": 8},
            eval_batch_by_model={"pubmedbert": 32, "biogpt": 16, "clinical_t5": 16},
            grad_accum_by_model={"pubmedbert": 6, "biogpt": 8, "clinical_t5": 8},
            num_workers=2,
            prefetch_factor=2,
            persistent_workers=True,
            notes="Generic 12 GB+ CUDA profile.",
        )

    return GPUExecutionProfile(
        name="compat",
        train_batch_by_model={"pubmedbert": 8, "biogpt": 6, "clinical_t5": 6},
        eval_batch_by_model={"pubmedbert": 16, "biogpt": 8, "clinical_t5": 8},
        grad_accum_by_model={"pubmedbert": 8, "biogpt": 10, "clinical_t5": 10},
        num_workers=0,
        prefetch_factor=2,
        persistent_workers=False,
        notes="Compatibility profile for lower-VRAM GPUs.",
    )


def _build_loader_kwargs(device, profile):
    kwargs = {
        "pin_memory": device.type == "cuda",
        "num_workers": profile.num_workers if device.type == "cuda" else 0,
    }
    if kwargs["num_workers"] > 0:
        kwargs["persistent_workers"] = profile.persistent_workers
        kwargs["prefetch_factor"] = profile.prefetch_factor
    return kwargs

def _parse_selected_models(raw: str):
    raw = (raw or 'all').strip().lower()
    if raw in ('all', '*'):
        return None
    alias = {
        'pubmed': 'pubmedbert',
        'pubmedbert': 'pubmedbert',
        'biogpt': 'biogpt',
        'bio': 'biogpt',
        'clinical': 'clinical_t5',
        'clinical_t5': 'clinical_t5',
        't5': 'clinical_t5',
    }
    out = []
    for part in [p.strip() for p in raw.split(',') if p.strip()]:
        out.append(alias.get(part, part))
    return set(out) if out else None



# ---------------------------------------------------------------------------
# Training loop
# ---------------------------------------------------------------------------
def train_one_model(
    model, optimizer, scheduler, criterion, scaler,
    train_loader, val_loader, device, model_name,
    num_epochs=25, patience=8, grad_accum_steps=2,
    checkpoint_dir=None,
):
    """Train a single model with mixed precision, gradient accumulation, and
    early stopping.  Returns the best model state dict and training history."""

    checkpoint_path = checkpoint_dir / f"{model_name}_ckpt.pth" if checkpoint_dir else None
    history = {"train_loss": [], "val_loss": [], "val_acc": [], "val_f1": [], "lr": []}
    best_val_loss = float("inf")
    early_stop_counter = 0
    start_epoch = 0

    # Resume from checkpoint if available
    if checkpoint_path and checkpoint_path.exists():
        print(f"  [CKPT] Found checkpoint at {checkpoint_path}")
        try:
            ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
            model.load_state_dict(ckpt["model_state_dict"])
            optimizer.load_state_dict(ckpt["optimizer_state_dict"])
            start_epoch = ckpt["epoch"] + 1
            best_val_loss = ckpt["best_val_loss"]
            history = ckpt.get("history", history)
            print(f"  [CKPT] Resuming from epoch {start_epoch} (best val loss {best_val_loss:.4f})")
        except Exception as e:
            print(f"  [CKPT] Could not load checkpoint: {e}. Starting fresh.")
            start_epoch = 0

    use_amp = device.type == "cuda"

    for epoch in range(start_epoch, num_epochs):
        t0 = time.time()
        model.train()
        running_loss = 0.0
        optimizer.zero_grad(set_to_none=True)

        # Progress bar for training
        pbar = tqdm(enumerate(train_loader), total=len(train_loader), 
                    desc=f"Epoch {epoch+1}/{num_epochs}", unit="batch", ncols=100, ascii=True)
        
        for batch_idx, batch in pbar:
            data, targets, contexts = _prepare_batch(batch, device)

            with torch.amp.autocast(device_type=device.type, enabled=use_amp):
                outputs = model(data, contexts)
                loss = criterion(outputs, targets) / grad_accum_steps

            scaler.scale(loss).backward()

            if (batch_idx + 1) % grad_accum_steps == 0 or (batch_idx + 1) == len(train_loader):
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)

            current_loss = loss.item() * grad_accum_steps
            running_loss += current_loss
            
            # Update progress bar every few batches to reduce overhead
            if batch_idx % 10 == 0:
                pbar.set_postfix(loss=f"{current_loss:.4f}", lr=f"{optimizer.param_groups[0]['lr']:.2e}")

        avg_train_loss = running_loss / len(train_loader)
        history["train_loss"].append(avg_train_loss)
        history["lr"].append(optimizer.param_groups[0]["lr"])

        # ---- Validation ----
        model.eval()
        val_loss = 0.0
        all_preds, all_targets = [], []

        with torch.no_grad():
            for batch in tqdm(val_loader, desc=f"Val  {epoch+1}/{num_epochs}", unit="batch", leave=False, ncols=100, ascii=True):
                data, targets, contexts = _prepare_batch(batch, device)
                with torch.amp.autocast(device_type=device.type, enabled=use_amp):
                    outputs = model(data, contexts)
                    loss = criterion(outputs, targets)
                val_loss += loss.item()
                _, predicted = outputs.max(1)
                all_preds.extend(predicted.cpu().numpy())
                all_targets.extend(targets.cpu().numpy())

        avg_val_loss = val_loss / len(val_loader)
        val_acc = 100 * np.mean(np.array(all_preds) == np.array(all_targets))
        val_f1 = f1_score(all_targets, all_preds, average="weighted")

        history["val_loss"].append(avg_val_loss)
        history["val_acc"].append(val_acc)
        history["val_f1"].append(val_f1)
        scheduler.step(avg_val_loss)

        elapsed = time.time() - t0
        gpu_mem = torch.cuda.memory_allocated(0) / 1024**2 if device.type == "cuda" else 0

        print(
            f"  Epoch {epoch+1:02d}/{num_epochs} │ "
            f"Train Loss {avg_train_loss:.4f} │ Val Loss {avg_val_loss:.4f} │ "
            f"Val Acc {val_acc:.2f}% │ Val F1 {val_f1:.4f} │ "
            f"LR {optimizer.param_groups[0]['lr']:.2e} │ "
            f"GPU {gpu_mem:.0f}MB │ {elapsed:.1f}s"
        )

        # ---- Checkpointing & early stopping ----
        # Save every epoch so interrupted runs can resume from latest epoch.
        if checkpoint_path:
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "best_val_loss": best_val_loss,
                "history": history,
            }, checkpoint_path)

        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            if checkpoint_path:
                torch.save({
                    "epoch": epoch,
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "best_val_loss": best_val_loss,
                    "history": history,
                }, checkpoint_path)
            early_stop_counter = 0
            print(f"  [*] New best (val loss {best_val_loss:.4f})")
        else:
            early_stop_counter += 1
            if early_stop_counter >= patience:
                print(f"  [X] Early stopping after {epoch+1} epochs")
                break

    # Load best weights
    if checkpoint_path and checkpoint_path.exists():
        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model_state_dict"])
        checkpoint_path.unlink()

    return model, history, best_val_loss


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------
def evaluate_on_test(model, test_loader, criterion, device, model_name, class_names):
    """Full evaluation on the held-out test set."""
    model.eval()
    use_amp = device.type == "cuda"
    all_preds, all_targets, all_probs = [], [], []
    test_loss = 0.0

    with torch.no_grad():
        for batch in tqdm(test_loader, desc=f"Evaluating {model_name}", unit="batch"):
            data, targets, contexts = _prepare_batch(batch, device)
            with torch.amp.autocast(device_type=device.type, enabled=use_amp):
                outputs = model(data, contexts)
                loss = criterion(outputs, targets)
            test_loss += loss.item()
            probs = torch.softmax(outputs, dim=1)
            _, predicted = outputs.max(1)
            all_preds.extend(predicted.cpu().numpy())
            all_targets.extend(targets.cpu().numpy())
            all_probs.extend(probs.cpu().numpy())

    all_preds = np.array(all_preds)
    all_targets = np.array(all_targets)
    all_probs = np.array(all_probs)

    accuracy = np.mean(all_preds == all_targets)
    f1 = f1_score(all_targets, all_preds, average="weighted")
    precision = precision_score(all_targets, all_preds, average="weighted")
    recall = recall_score(all_targets, all_preds, average="weighted")

    try:
        auroc = roc_auc_score(all_targets, all_probs, multi_class="ovr", average="weighted")
    except Exception:
        auroc = 0.0

    report = classification_report(all_targets, all_preds, target_names=class_names)
    cm = confusion_matrix(all_targets, all_preds)

    print(f"\n{'='*70}")
    print(f"  {model_name.upper()} — TEST SET RESULTS")
    print(f"{'='*70}")
    print(f"  Accuracy  : {accuracy:.4f}  ({accuracy*100:.2f}%)")
    print(f"  F1 Score  : {f1:.4f}")
    print(f"  Precision : {precision:.4f}")
    print(f"  Recall    : {recall:.4f}")
    print(f"  AUROC     : {auroc:.4f}")
    print(f"\n{report}")

    return {
        "accuracy": accuracy, "f1": f1, "precision": precision,
        "recall": recall, "auroc": auroc,
        "classification_report": report, "confusion_matrix": cm,
        "predictions": all_preds, "targets": all_targets, "probabilities": all_probs,
    }


# ---------------------------------------------------------------------------
# Plotting
# ---------------------------------------------------------------------------
def save_plots(results, history_dict, class_names, plots_dir):
    """Save confusion matrices, training curves, and comparison charts."""
    plots_dir = Path(plots_dir)
    plots_dir.mkdir(parents=True, exist_ok=True)

    for name, res in results.items():
        # Confusion matrix
        plt.figure(figsize=(8, 6))
        sns.heatmap(res["confusion_matrix"], annot=True, fmt="d", cmap="Blues",
                    xticklabels=class_names, yticklabels=class_names)
        plt.title(f"{name} — Confusion Matrix (Leak-Free Split)")
        plt.xlabel("Predicted")
        plt.ylabel("True")
        plt.tight_layout()
        plt.savefig(plots_dir / f"{name}_confusion_matrix.png", dpi=200)
        plt.close()

        # Training curves
        hist = history_dict[name]
        fig, axes = plt.subplots(1, 3, figsize=(18, 5))

        axes[0].plot(hist["train_loss"], label="Train", color="#3498db")
        axes[0].plot(hist["val_loss"], label="Val", color="#e74c3c")
        axes[0].set_title(f"{name} — Loss")
        axes[0].set_xlabel("Epoch")
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)

        axes[1].plot(hist["val_acc"], color="#2ecc71")
        axes[1].set_title(f"{name} — Val Accuracy (%)")
        axes[1].set_xlabel("Epoch")
        axes[1].grid(True, alpha=0.3)

        axes[2].plot(hist["lr"], color="#9b59b6")
        axes[2].set_title(f"{name} — Learning Rate")
        axes[2].set_xlabel("Epoch")
        axes[2].grid(True, alpha=0.3)

        plt.tight_layout()
        plt.savefig(plots_dir / f"{name}_training_curves.png", dpi=200)
        plt.close()

    # Comparison bar chart
    model_names = list(results.keys())
    metric_names = ["accuracy", "f1", "precision", "recall", "auroc"]
    fig, axes = plt.subplots(1, len(metric_names), figsize=(5 * len(metric_names), 5))
    colors = ["#3498db", "#2ecc71", "#e74c3c"]

    for ax, metric in zip(axes, metric_names):
        values = [results[m][metric] for m in model_names]
        bars = ax.bar(model_names, values, color=colors[:len(model_names)])
        ax.set_title(metric.upper())
        ax.set_ylim(0, 1.05)
        for bar, val in zip(bars, values):
            ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
                    f"{val:.3f}", ha="center", va="bottom", fontsize=9)

    plt.tight_layout()
    plt.savefig(plots_dir / "transformer_comparison.png", dpi=200)
    plt.close()

    print(f"[PLOT] Saved all plots → {plots_dir}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    print("=" * 70)
    print("  TRANSFORMER MODEL TRAINING — LEAK-FREE SPLIT + CUDA")
    print("=" * 70)

    # ---- Seed everything ----
    torch.manual_seed(42)
    np.random.seed(42)
    _ensure_transformer_cuda()
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(42)

    # ---- Device ----
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\n[DEVICE] Using: {device}")
    if device.type == "cuda":
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
        torch.cuda.empty_cache()
        _print_gpu_info(device)
    else:
        print("[WARNING] CUDA not available -- CPU fallback is enabled by PD_ALLOW_CPU_TRANSFORMERS=1")

    # ---- Data ----
    preprocessor = DataPreprocessor()
    base_dir = str(ROOT)
    file_paths = [
        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20240129.csv"),
        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20241211.csv"),
        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20250321.csv"),
        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20250714.csv"),
    ]

    X_train, X_test, y_train, y_test, feature_names, class_mapping = \
        _load_or_create_leak_free_split(preprocessor, file_paths)

    X_train = np.asarray(X_train, dtype=np.float32)
    X_test = np.asarray(X_test, dtype=np.float32)
    y_train = np.asarray(y_train, dtype=np.int64)
    y_test = np.asarray(y_test, dtype=np.int64)

    print(f"[DATA] Train: {X_train.shape}  Test: {X_test.shape}")
    print(f"[DATA] Classes: {len(np.unique(y_train))}  Distribution: {dict(zip(*np.unique(y_train, return_counts=True)))}")

    # ---- Class weights ----
    cw = compute_class_weight("balanced", classes=np.unique(y_train), y=y_train)
    class_weights_tensor = torch.FloatTensor(cw).to(device)
    print(f"[DATA] Class weights: {dict(zip(np.unique(y_train), np.round(cw, 3)))}")

    # ---- Stratified validation split (preserves class ratios) ----
    train_idx, val_idx = _stratified_val_indices(y_train, val_fraction=0.15)
    print(f"[DATA] Stratified split: {len(train_idx)} train / {len(val_idx)} val")

    # ---- Feature names ----
    if feature_names is None:
        feature_names = preprocessor.get_feature_names()

    # ---- RAG context ----
    docs_path = str(ROOT / "medical_docs")
    doc_manager = DocumentManager(docs_dir=docs_path)
    doc_count = doc_manager.get_document_count()
    print(f"[RAG] Loaded {doc_count.get('total', doc_count)} documents for context enrichment")

    def get_rag_context(sample_features):
        feature_desc = {name: float(val) for name, val in zip(feature_names, sample_features)}
        query_parts = []
        symptoms = {
            "tremor": feature_desc.get("sym_tremor", 0),
            "rigidity": feature_desc.get("sym_rigid", 0),
            "bradykinesia": feature_desc.get("sym_brady", 0),
            "postural instability": feature_desc.get("sym_posins", 0),
        }
        for symptom, severity in symptoms.items():
            if severity > 0:
                query_parts.append(f"{symptom} severity:{severity}")
        moca = feature_desc.get("moca", 30)
        if moca < 26:
            query_parts.append("cognitive impairment")
        age = feature_desc.get("age", 0)
        if age:
            query_parts.append(f"age {int(age)}")
        if feature_desc.get("fampd", 0) > 0:
            query_parts.append("family history Parkinson's disease")
        query = "Parkinson's disease " + " ".join(query_parts)
        passages = doc_manager.extract_relevant_passages(query, top_k=2)
        if not passages:
            return ""
        return " ".join(
            f"From '{p['doc_title']}' {p['text'][:300]}..." for p in passages
        )

    print(f"\n[RAG] RAG Enabled: {USE_RAG}")
    if USE_RAG:
        print("[RAG] Pre-computing context for train + test splits ...")
    
    train_cache = RESULTS_DIR / "rag_contexts_train.pkl"
    test_cache = RESULTS_DIR / "rag_contexts_test.pkl"
    
    train_contexts = _build_context_cache(X_train, get_rag_context, "train", train_cache)
    test_contexts = _build_context_cache(X_test, get_rag_context, "test", test_cache)


    # ---- Datasets ----
    full_train_ds = TabularDataset(X_train, y_train, feature_names, contexts=train_contexts)
    test_ds = TabularDataset(X_test, y_test, feature_names, contexts=test_contexts)

    train_subset = Subset(full_train_ds, train_idx)
    val_subset = Subset(full_train_ds, val_idx)

    gpu_profile = _detect_gpu_execution_profile()
    loader_kwargs = _build_loader_kwargs(device, gpu_profile)
    print(f"\n[GPU PROFILE] {gpu_profile.name} -> {gpu_profile.notes}")

    # ---- Model definitions (lazy — created one at a time to fit in available VRAM) ----
    input_dim = X_train.shape[1]
    num_classes = len(np.unique(y_train))
    class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]

    print(f"\n[MODEL] Input dim: {input_dim}  Num classes: {num_classes}")

    # Each entry: (display_name, save_name, model_factory)
    # Models are created lazily inside the loop to avoid GPU OOM.
    selected_models = _parse_selected_models(os.getenv("PD_TRAIN_MODELS", "all"))

    model_configs = [
        (
            "PubMedBERT (Encoder-Only)", "pubmedbert",
            lambda: PubMedBERTForTabular(input_dim, num_classes, dropout=0.10, freeze_bert=False),
            {"lr": 1.5e-5, "weight_decay": 0.02},
        ),
        (
            "BioGPT", "biogpt",
            lambda: BioGPTForTabular(input_dim, num_classes, dropout=0.12, train_decoder_layers=8),
            {"lr": 2e-5, "weight_decay": 0.02},
        ),
        (
            "Clinical-T5", "clinical_t5",
            lambda: ClinicalT5ForTabular(input_dim, num_classes, dropout=0.10, freeze_encoder=False),
            {"lr": 1.5e-5, "weight_decay": 0.02},
        ),
    ]

    if selected_models is not None:
        model_configs = [cfg for cfg in model_configs if cfg[1] in selected_models]
        print(f"[MODEL] Filter active -> {sorted(selected_models)}")
    if not model_configs:
        raise ValueError("No valid models selected for training.")

    # ---- Training config ----
    NUM_EPOCHS = 30
    PATIENCE = 10
    DEFAULT_GRAD_ACCUM = 8

    criterion = torch.nn.CrossEntropyLoss(weight=class_weights_tensor)
    checkpoint_dir = MODEL_DIR / "_checkpoints"
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    all_results = {}
    all_histories = {}

    for display_name, save_name, model_factory, opt_kwargs in model_configs:
        final_path = MODEL_DIR / f"{display_name}_best.pth"
        train_bs = gpu_profile.train_batch_by_model.get(save_name, 8)
        eval_bs = gpu_profile.eval_batch_by_model.get(save_name, max(train_bs * 2, 8))
        grad_accum = gpu_profile.grad_accum_by_model.get(save_name, DEFAULT_GRAD_ACCUM)

        train_loader = DataLoader(train_subset, batch_size=train_bs, shuffle=True, **loader_kwargs)
        val_loader = DataLoader(val_subset, batch_size=eval_bs, shuffle=False, **loader_kwargs)
        test_loader = DataLoader(test_ds, batch_size=eval_bs, shuffle=False, **loader_kwargs)

        print(f"\n[LOADER] {display_name}: train_bs={train_bs} eval_bs={eval_bs} grad_accum={grad_accum} workers={loader_kwargs.get('num_workers', 0)}")

        # ---- Skip if already trained ----
        if final_path.exists():
            print(f"\n{'='*70}")
            print(f"  SKIPPING: {display_name}  (already trained)")
            print(f"  Loading saved weights from: {final_path}")
            print(f"{'='*70}")
            try:
                # Create the model on CPU first, load weights, then move to GPU
                print(f"\n[MODEL] Initializing {display_name} for evaluation ...")
                model = model_factory()
                model.load_state_dict(torch.load(final_path, map_location="cpu", weights_only=True))
                model.to(device)

                # Evaluate on test set
                result = evaluate_on_test(
                    model, test_loader, criterion, device, display_name, class_names,
                )
                all_results[display_name] = result
                all_histories[display_name] = {"train_loss": [], "val_loss": [], "val_acc": [], "val_f1": [], "lr": []}

                # Free GPU memory before next model
                del model
                if device.type == "cuda":
                    torch.cuda.empty_cache()

                continue
            except RuntimeError as e:
                print(f"  [WARN] Cannot load saved weights (architecture changed?): {e}")
                print(f"  [WARN] Deleting stale checkpoint and retraining ...")
                final_path.unlink(missing_ok=True)

        # ---- Create model fresh for training ----
        print(f"\n[MODEL] Initializing {display_name} for training ...")
        if device.type == "cuda":
            torch.cuda.empty_cache()

        model = model_factory().to(device)
        optimizer = torch.optim.AdamW(
            filter(lambda p: p.requires_grad, model.parameters()),
            **opt_kwargs,
        )
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode="min", factor=0.5, patience=3, min_lr=1e-7,
        )

        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        total_params = sum(p.numel() for p in model.parameters())

        print(f"\n{'='*70}")
        print(f"  TRAINING: {display_name}")
        print(f"  Trainable params: {trainable_params:,} / {total_params:,}")
        print(f"  Epochs: {NUM_EPOCHS}  Patience: {PATIENCE}  Grad Accum: {grad_accum}")
        print(f"{'='*70}")

        if device.type == "cuda":
            torch.cuda.reset_peak_memory_stats()
            torch.cuda.empty_cache()

        scaler = torch.amp.GradScaler(device=device.type, enabled=device.type == "cuda")

        trained_model, history, best_val = train_one_model(
            model, optimizer, scheduler, criterion, scaler,
            train_loader, val_loader, device, save_name,
            num_epochs=NUM_EPOCHS, patience=PATIENCE, grad_accum_steps=grad_accum,
            checkpoint_dir=checkpoint_dir,
        )

        # Save final model
        torch.save(trained_model.state_dict(), final_path)
        alias_paths = {
            "pubmedbert": [
                MODEL_DIR / "pubmedbert_transformer.pth",
                MODEL_DIR / "pubmedbert.pth",
            ],
            "biogpt": [
                MODEL_DIR / "biogpt_transformer.pth",
                MODEL_DIR / "biogpt.pth",
                MODEL_DIR / "biomistral.pth",
            ],
            "clinical_t5": [
                MODEL_DIR / "clinical_t5_transformer.pth",
                MODEL_DIR / "clinicalt5_transformer.pth",
                MODEL_DIR / "clinical_t5.pth",
            ],
        }
        for alias_path in alias_paths.get(save_name, []):
            torch.save(trained_model.state_dict(), alias_path)
        print(f"  [SAVE] Model saved → {final_path}")

        if device.type == "cuda":
            peak = torch.cuda.max_memory_allocated(0) / 1024**3
            print(f"  [GPU] Peak VRAM usage: {peak:.2f} GB")

        # Evaluate on test set
        result = evaluate_on_test(
            trained_model, test_loader, criterion, device, display_name, class_names,
        )
        all_results[display_name] = result
        all_histories[display_name] = history

        # Free GPU memory before next model
        del model, trained_model, optimizer, scheduler, scaler
        if device.type == "cuda":
            torch.cuda.empty_cache()

    # NOTE: We intentionally keep the checkpoint dir so that interrupted
    # training can resume from the last saved epoch on the next run.

    # ---- Save metrics to JSON + CSV ----
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    summary_rows = []
    for name, res in all_results.items():
        summary_rows.append({
            "Model": name,
            "Type": "Transformer",
            "Accuracy": round(res["accuracy"], 4),
            "F1_Score": round(res["f1"], 4),
            "Precision": round(res["precision"], 4),
            "Recall": round(res["recall"], 4),
            "AUROC": round(res["auroc"], 4),
        })
    df = pd.DataFrame(summary_rows)
    df.to_csv(RESULTS_DIR / "transformer_metrics_latest.csv", index=False)
    with open(RESULTS_DIR / "transformer_metrics_latest.json", "w") as f:
        json.dump(summary_rows, f, indent=2)
    print(f"\n[SAVE] Metrics  → {RESULTS_DIR / 'transformer_metrics_latest.csv'}")

    # ---- Plots ----
    save_plots(all_results, all_histories, class_names, PLOTS_DIR)

    # ---- Final summary ----
    print(f"\n{'='*70}")
    print("  FINAL COMPARISON")
    print(f"{'='*70}")
    print(f"{'Model':<30} {'Accuracy':>10} {'F1':>10} {'Precision':>10} {'Recall':>10} {'AUROC':>10}")
    print("-" * 80)
    for row in summary_rows:
        print(f"{row['Model']:<30} {row['Accuracy']:>10.4f} {row['F1_Score']:>10.4f} "
              f"{row['Precision']:>10.4f} {row['Recall']:>10.4f} {row['AUROC']:>10.4f}")

    best_by_f1 = max(all_results, key=lambda k: all_results[k]["f1"])
    print(f"\n  [*] Best model (by F1): {best_by_f1} -- F1 {all_results[best_by_f1]['f1']:.4f}")
    print(f"\n{'='*70}")
    print("  TRAINING COMPLETE!")
    print(f"{'='*70}\n")


if __name__ == "__main__":
    main()