src/train_model_suite.py

from __future__ import annotations

import argparse
import json
import sys
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import joblib
import numpy as np
import pandas as pd
import torch
from lightgbm import LGBMClassifier
from scipy import sparse
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_recall_fscore_support,
    roc_auc_score,
)
from sklearn.model_selection import GroupKFold, GroupShuffleSplit
from sklearn.utils.class_weight import compute_class_weight
from sklearn.svm import SVC
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from xgboost import XGBClassifier

sys.path.append(str(Path(__file__).resolve().parent))

from data_preprocessing import DataPreprocessor  # type: ignore
from training_runtime import PauseRequested, StopRequested, TrainingRunController  # type: ignore


warnings.filterwarnings(
    "ignore",
    message="X does not have valid feature names, but LGBMClassifier was fitted with feature names",
)


ROOT = Path(__file__).resolve().parents[1]
MODEL_DIR = ROOT / "models" / "saved"
RUNS_DIR = MODEL_DIR / "training_runs"
EVAL_DIR = ROOT / "evaluation_results"
MODEL_METRICS_DIR = EVAL_DIR / "model_metrics"
CLASS_NAMES_DEFAULT = ["HC", "PD", "SWEDD", "PRODROMAL"]
TRADITIONAL_MODELS = ("lightgbm", "xgboost", "svm")
TRANSFORMER_MODELS = ("pubmedbert", "biogpt", "clinical_t5")
ALL_BASE_MODELS = TRADITIONAL_MODELS + TRANSFORMER_MODELS
TRANSFORMER_SELECTION_METRIC = "val_f1"
DEFAULT_TRANSFORMER_LOSS = "focal"
DEFAULT_FOCAL_GAMMA = 1.5


TRADITIONAL_SEARCH_SPACES: Dict[str, List[Dict[str, Any]]] = {
    "lightgbm": [
        {"n_estimators": 350, "learning_rate": 0.03, "max_depth": -1, "num_leaves": 31, "min_child_samples": 20, "subsample": 0.90, "colsample_bytree": 0.90},
        {"n_estimators": 500, "learning_rate": 0.02, "max_depth": -1, "num_leaves": 63, "min_child_samples": 25, "subsample": 0.85, "colsample_bytree": 0.85},
        {"n_estimators": 450, "learning_rate": 0.025, "max_depth": 10, "num_leaves": 47, "min_child_samples": 15, "subsample": 0.90, "colsample_bytree": 0.80},
        {"n_estimators": 650, "learning_rate": 0.015, "max_depth": 12, "num_leaves": 95, "min_child_samples": 30, "subsample": 0.80, "colsample_bytree": 0.80},
        {"n_estimators": 800, "learning_rate": 0.012, "max_depth": -1, "num_leaves": 127, "min_child_samples": 20, "subsample": 0.85, "colsample_bytree": 0.80},
        {"n_estimators": 420, "learning_rate": 0.03, "max_depth": 8, "num_leaves": 63, "min_child_samples": 10, "subsample": 0.95, "colsample_bytree": 0.85},
    ],
    "xgboost": [
        {"n_estimators": 300, "learning_rate": 0.05, "max_depth": 6, "min_child_weight": 2, "subsample": 0.90, "colsample_bytree": 0.90, "gamma": 0.0, "reg_lambda": 1.0},
        {"n_estimators": 450, "learning_rate": 0.03, "max_depth": 5, "min_child_weight": 1, "subsample": 0.85, "colsample_bytree": 0.85, "gamma": 0.05, "reg_lambda": 1.0},
        {"n_estimators": 500, "learning_rate": 0.025, "max_depth": 7, "min_child_weight": 3, "subsample": 0.80, "colsample_bytree": 0.80, "gamma": 0.10, "reg_lambda": 1.5},
        {"n_estimators": 350, "learning_rate": 0.04, "max_depth": 4, "min_child_weight": 1, "subsample": 0.95, "colsample_bytree": 0.90, "gamma": 0.0, "reg_lambda": 0.8},
        {"n_estimators": 650, "learning_rate": 0.015, "max_depth": 8, "min_child_weight": 2, "subsample": 0.85, "colsample_bytree": 0.85, "gamma": 0.05, "reg_lambda": 1.2},
        {"n_estimators": 520, "learning_rate": 0.02, "max_depth": 6, "min_child_weight": 4, "subsample": 0.90, "colsample_bytree": 0.80, "gamma": 0.08, "reg_lambda": 1.8},
    ],
    "svm": [
        {"C": 6.0, "gamma": "scale", "kernel": "rbf"},
        {"C": 8.0, "gamma": "scale", "kernel": "rbf"},
        {"C": 10.0, "gamma": 0.01, "kernel": "rbf"},
        {"C": 12.0, "gamma": 0.005, "kernel": "rbf"},
        {"C": 14.0, "gamma": "scale", "kernel": "rbf"},
        {"C": 16.0, "gamma": 0.003, "kernel": "rbf"},
    ],
}


TRANSFORMER_TRIALS: Dict[str, List[Dict[str, Any]]] = {
    "pubmedbert": [
        {"model_kwargs": {"dropout": 0.10, "freeze_bert": False, "train_encoder_layers": 8}, "optimizer": {"lr": 8.0e-6, "weight_decay": 0.02}, "grad_accum": 2},
        {"model_kwargs": {"dropout": 0.08, "freeze_bert": False, "train_encoder_layers": 6}, "optimizer": {"lr": 1.0e-5, "weight_decay": 0.02}, "grad_accum": 2},
        {"model_kwargs": {"dropout": 0.12, "freeze_bert": False, "train_encoder_layers": 4}, "optimizer": {"lr": 1.5e-5, "weight_decay": 0.01}, "grad_accum": 2},
        {"model_kwargs": {"dropout": 0.18, "freeze_bert": True}, "optimizer": {"lr": 2.5e-5, "weight_decay": 0.01}, "grad_accum": 2},
        {"model_kwargs": {"dropout": 0.06, "freeze_bert": False, "train_encoder_layers": 10}, "optimizer": {"lr": 6.0e-6, "weight_decay": 0.03}, "grad_accum": 2},
        {"model_kwargs": {"dropout": 0.10, "freeze_bert": False, "train_encoder_layers": 12}, "optimizer": {"lr": 5.0e-6, "weight_decay": 0.03}, "grad_accum": 1},
    ],
    "biogpt": [
        {"model_kwargs": {"dropout": 0.10, "train_decoder_layers": 4}, "optimizer": {"lr": 3.0e-5, "weight_decay": 0.01}, "grad_accum": 8},
        {"model_kwargs": {"dropout": 0.15, "train_decoder_layers": 6}, "optimizer": {"lr": 2.5e-5, "weight_decay": 0.02}, "grad_accum": 8},
        {"model_kwargs": {"dropout": 0.20, "train_decoder_layers": 8}, "optimizer": {"lr": 2.0e-5, "weight_decay": 0.02}, "grad_accum": 10},
        {"model_kwargs": {"dropout": 0.12, "train_decoder_layers": 10}, "optimizer": {"lr": 1.5e-5, "weight_decay": 0.02}, "grad_accum": 12},
        {"model_kwargs": {"dropout": 0.10, "train_decoder_layers": 12}, "optimizer": {"lr": 1.0e-5, "weight_decay": 0.02}, "grad_accum": 8},
        {"model_kwargs": {"dropout": 0.08, "train_decoder_layers": 10}, "optimizer": {"lr": 1.2e-5, "weight_decay": 0.015}, "grad_accum": 6},
    ],
    "clinical_t5": [
        {"model_kwargs": {"dropout": 0.10, "freeze_encoder": False}, "optimizer": {"lr": 2.0e-5, "weight_decay": 0.01}, "grad_accum": 8},
        {"model_kwargs": {"dropout": 0.15, "freeze_encoder": False}, "optimizer": {"lr": 1.5e-5, "weight_decay": 0.02}, "grad_accum": 8},
        {"model_kwargs": {"dropout": 0.20, "freeze_encoder": True}, "optimizer": {"lr": 2.5e-5, "weight_decay": 0.01}, "grad_accum": 6},
        {"model_kwargs": {"dropout": 0.08, "freeze_encoder": False}, "optimizer": {"lr": 1.0e-5, "weight_decay": 0.02}, "grad_accum": 10},
        {"model_kwargs": {"dropout": 0.12, "freeze_encoder": False}, "optimizer": {"lr": 8.0e-6, "weight_decay": 0.02}, "grad_accum": 8},
        {"model_kwargs": {"dropout": 0.10, "freeze_encoder": False}, "optimizer": {"lr": 1.2e-5, "weight_decay": 0.015}, "grad_accum": 6},
    ],
}


@dataclass
class TrainingBundle:
    X_train_dense: np.ndarray
    X_test_dense: np.ndarray
    y_train: np.ndarray
    y_test: np.ndarray
    train_groups: np.ndarray
    test_groups: np.ndarray
    feature_names: List[str]
    class_mapping: Dict[str, int]
    class_names: List[str]
    preprocessor: Any


@dataclass(frozen=True)
class GPUExecutionProfile:
    name: str
    train_batch_by_model: Dict[str, int]
    eval_batch_by_model: Dict[str, int]
    grad_accum_cap_by_model: Dict[str, int]
    num_workers: int
    prefetch_factor: int
    persistent_workers: bool
    notes: str


class FocalLoss(torch.nn.Module):
    """Class-weighted focal loss for imbalanced multi-class classification."""

    def __init__(
        self,
        class_weights: Optional[torch.Tensor] = None,
        gamma: float = DEFAULT_FOCAL_GAMMA,
        reduction: str = "mean",
    ) -> None:
        super().__init__()
        if reduction not in {"mean", "sum", "none"}:
            raise ValueError(f"Unsupported focal-loss reduction: {reduction}")
        self.gamma = float(gamma)
        self.reduction = reduction
        if class_weights is not None:
            normalized = class_weights.float() / class_weights.float().mean().clamp_min(1e-6)
            self.register_buffer("class_weights", normalized)
        else:
            self.register_buffer("class_weights", None)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        log_probs = torch.nn.functional.log_softmax(logits, dim=1)
        probs = log_probs.exp()
        target_log_probs = log_probs.gather(1, targets.unsqueeze(1)).squeeze(1)
        target_probs = probs.gather(1, targets.unsqueeze(1)).squeeze(1).clamp_min(1e-6)
        ce_loss = -target_log_probs
        focal_factor = (1.0 - target_probs).pow(self.gamma)
        loss = focal_factor * ce_loss
        if self.class_weights is not None:
            loss = self.class_weights[targets] * loss
        if self.reduction == "mean":
            return loss.mean()
        if self.reduction == "sum":
            return loss.sum()
        return loss


def _json_ready(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(key): _json_ready(val) for key, val in value.items()}
    if isinstance(value, (list, tuple)):
        return [_json_ready(item) for item in value]
    if isinstance(value, np.ndarray):
        return value.tolist()
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, Path):
        return str(value)
    return value


def _as_dense(matrix: Any) -> np.ndarray:
    if sparse.issparse(matrix):
        return matrix.toarray()
    return np.asarray(matrix)


def _candidate_file_paths() -> List[Path]:
    return [
        ROOT / "PPMI_Curated_Data_Cut_Public_20240129.csv",
        ROOT / "PPMI_Curated_Data_Cut_Public_20241211.csv",
        ROOT / "PPMI_Curated_Data_Cut_Public_20250321.csv",
        ROOT / "PPMI_Curated_Data_Cut_Public_20250714.csv",
    ]


def _prepare_training_bundle() -> TrainingBundle:
    preprocessor = DataPreprocessor()
    X_train, X_test, y_train, y_test = preprocessor.prepare_data(
        _candidate_file_paths(),
        test_size=0.2,
        use_patient_split=True,
    )
    train_df, test_df = preprocessor.get_split_frames()
    feature_names = preprocessor.get_feature_names()
    class_mapping = preprocessor.get_class_mapping()
    class_names = [None] * len(class_mapping)
    for label, idx in class_mapping.items():
        class_names[idx] = label

    EVAL_DIR.mkdir(parents=True, exist_ok=True)
    np.savez(
        EVAL_DIR / "leak_free_split.npz",
        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,
            "train_groups": train_df["PATNO"].to_numpy(),
            "test_groups": test_df["PATNO"].to_numpy(),
        },
        EVAL_DIR / "leak_free_split_meta.joblib",
    )

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    joblib.dump(preprocessor.get_preprocessor(), MODEL_DIR / "traditional_preprocessor.joblib")
    (MODEL_DIR / "traditional_class_mapping.json").write_text(
        json.dumps(_json_ready(class_mapping), indent=2),
        encoding="utf-8",
    )

    return TrainingBundle(
        X_train_dense=_as_dense(X_train).astype(np.float32),
        X_test_dense=_as_dense(X_test).astype(np.float32),
        y_train=np.asarray(y_train, dtype=np.int64),
        y_test=np.asarray(y_test, dtype=np.int64),
        train_groups=train_df["PATNO"].to_numpy(),
        test_groups=test_df["PATNO"].to_numpy(),
        feature_names=list(feature_names),
        class_mapping=dict(class_mapping),
        class_names=[
            str(name if name is not None else CLASS_NAMES_DEFAULT[idx])
            for idx, name in enumerate(class_names)
        ],
        preprocessor=preprocessor.get_preprocessor(),
    )


def _compute_class_weight_dict(y_train: np.ndarray) -> Dict[int, float]:
    classes = np.unique(y_train)
    weights = compute_class_weight(class_weight="balanced", classes=classes, y=y_train)
    return {int(cls): float(weight) for cls, weight in zip(classes, weights)}


def _build_transformer_criterion(
    loss_name: str,
    class_weights_tensor: torch.Tensor,
    focal_gamma: float,
) -> Tuple[torch.nn.Module, str]:
    normalized_name = (loss_name or DEFAULT_TRANSFORMER_LOSS).strip().lower()
    if normalized_name in {"ce", "cross_entropy", "cross-entropy"}:
        return torch.nn.CrossEntropyLoss(weight=class_weights_tensor), "cross_entropy"
    if normalized_name == "focal":
        return FocalLoss(class_weights=class_weights_tensor, gamma=focal_gamma), f"focal(gamma={focal_gamma:.2f})"
    raise ValueError(f"Unsupported transformer loss: {loss_name}")


def _is_better_validation_epoch(
    val_f1: float,
    val_loss: float,
    best_val_f1: float,
    best_val_loss: float,
    min_delta: float = 1e-4,
) -> bool:
    if val_f1 > (best_val_f1 + min_delta):
        return True
    if abs(val_f1 - best_val_f1) <= min_delta and val_loss < (best_val_loss - min_delta):
        return True
    return False


def _evaluate_predictions(
    model_name: str,
    model_type: str,
    y_true: np.ndarray,
    y_pred: np.ndarray,
    probabilities: Optional[np.ndarray],
    class_names: Sequence[str],
) -> Dict[str, Any]:
    normalized_class_names = [str(name) for name in class_names]
    accuracy = accuracy_score(y_true, y_pred)
    precision, recall, f1, _ = precision_recall_fscore_support(y_true, y_pred, average="weighted", zero_division=0)
    report = classification_report(y_true, y_pred, target_names=normalized_class_names, zero_division=0)
    cm = confusion_matrix(y_true, y_pred)
    auroc = None
    if probabilities is not None:
        try:
            auroc = roc_auc_score(y_true, probabilities, multi_class="ovr", average="weighted")
        except ValueError:
            auroc = None
    return {
        "Model": model_name,
        "Type": model_type,
        "Accuracy": float(accuracy),
        "Precision": float(precision),
        "Recall": float(recall),
        "F1_Score": float(f1),
        "AUROC": None if auroc is None else float(auroc),
        "classification_report": report,
        "confusion_matrix": cm.tolist(),
    }


def _traditional_model_builder(model_name: str, params: Dict[str, Any], class_weight_dict: Dict[int, float], num_classes: int):
    if model_name == "lightgbm":
        return LGBMClassifier(random_state=42, objective="multiclass", num_class=num_classes, verbose=-1, class_weight=class_weight_dict, **params)
    if model_name == "xgboost":
        return XGBClassifier(random_state=42, objective="multi:softprob", num_class=num_classes, eval_metric="mlogloss", verbosity=0, **params)
    if model_name == "svm":
        return SVC(random_state=42, probability=True, decision_function_shape="ovr", class_weight=class_weight_dict, cache_size=2048, **params)
    raise ValueError(f"Unsupported traditional model: {model_name}")


def _run_grouped_traditional_search(
    model_name: str,
    search_space: List[Dict[str, Any]],
    bundle: TrainingBundle,
    controller: TrainingRunController,
    max_trials: int,
) -> Tuple[Dict[str, Any], Dict[str, Any], Path]:
    class_weight_dict = _compute_class_weight_dict(bundle.y_train)
    checkpoint_state = controller.load_checkpoint_state(
        model_name,
        {
            "phase": "search",
            "trials": {},
            "best_trial_index": None,
            "best_score": None,
        },
    )

    search_space = list(search_space[:max_trials])
    n_splits = max(2, min(3, len(np.unique(bundle.train_groups))))
    splitter = GroupKFold(n_splits=n_splits)

    controller.mark_running("traditional_search", model_name=model_name)
    controller.update_model_state(
        model_name,
        status="running",
        family="traditional",
        trials_total=len(search_space),
    )

    for trial_index, params in enumerate(search_space):
        trial_key = str(trial_index)
        trial_state = checkpoint_state["trials"].setdefault(
            trial_key,
            {
                "params": params,
                "fold_scores": [],
                "status": "pending",
            },
        )
        if trial_state.get("status") == "completed":
            continue

        controller.mark_running("traditional_trial", model_name=model_name, extra={"trial_index": trial_index})

        for fold_index, (train_idx, val_idx) in enumerate(
            splitter.split(bundle.X_train_dense, bundle.y_train, groups=bundle.train_groups)
        ):
            if fold_index < len(trial_state["fold_scores"]):
                continue

            train_features = np.asarray(bundle.X_train_dense[train_idx], dtype=np.float32)
            val_features = np.asarray(bundle.X_train_dense[val_idx], dtype=np.float32)
            train_targets = np.asarray(bundle.y_train[train_idx], dtype=np.int64)
            val_targets = np.asarray(bundle.y_train[val_idx], dtype=np.int64)

            estimator = _traditional_model_builder(
                model_name,
                params,
                class_weight_dict=class_weight_dict,
                num_classes=len(bundle.class_names),
            )
            estimator.fit(train_features, train_targets)
            preds = estimator.predict(val_features)
            score = f1_score(val_targets, preds, average="weighted")

            trial_state["fold_scores"].append(float(score))
            trial_state["status"] = "running"
            checkpoint_state["phase"] = "search"
            controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
            controller.update_model_state(
                model_name,
                current_trial=trial_index,
                current_fold=fold_index,
                latest_fold_score=float(score),
            )
            controller.raise_if_requested()

        trial_state["mean_score"] = float(np.mean(trial_state["fold_scores"]))
        trial_state["status"] = "completed"
        controller.append_trial_result(
            model_name,
            {
                "trial_index": trial_index,
                "params": params,
                "mean_score": trial_state["mean_score"],
            },
        )
        if checkpoint_state.get("best_score") is None or trial_state["mean_score"] > checkpoint_state["best_score"]:
            checkpoint_state["best_score"] = float(trial_state["mean_score"])
            checkpoint_state["best_trial_index"] = trial_index
        controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
        controller.raise_if_requested()

    if checkpoint_state.get("best_trial_index") is None:
        raise RuntimeError(f"No completed trials were recorded for {model_name}")

    best_index = int(checkpoint_state["best_trial_index"])
    best_params = dict(search_space[best_index])

    checkpoint_state["phase"] = "fit_final"
    controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
    controller.mark_running("traditional_fit_final", model_name=model_name)

    best_model = _traditional_model_builder(
        model_name,
        best_params,
        class_weight_dict=class_weight_dict,
        num_classes=len(bundle.class_names),
    )
    best_model.fit(
        np.asarray(bundle.X_train_dense, dtype=np.float32),
        np.asarray(bundle.y_train, dtype=np.int64),
    )

    artifact_path = MODEL_DIR / f"{model_name}_model.joblib"
    joblib.dump(best_model, artifact_path)

    test_features = np.asarray(bundle.X_test_dense, dtype=np.float32)
    y_pred = best_model.predict(test_features)
    probabilities = best_model.predict_proba(test_features) if hasattr(best_model, "predict_proba") else None
    metrics = _evaluate_predictions(
        model_name=model_name.replace("_", " ").title().replace("Svm", "SVM").replace("Lightgbm", "LightGBM").replace("Xgboost", "XGBoost"),
        model_type="Traditional ML",
        y_true=bundle.y_test,
        y_pred=y_pred,
        probabilities=probabilities,
        class_names=bundle.class_names,
    )

    checkpoint_state["phase"] = "complete"
    checkpoint_state["artifact_path"] = str(artifact_path)
    checkpoint_state["best_params"] = best_params
    checkpoint_state["metrics"] = {
        key: value for key, value in metrics.items()
        if key in {"Accuracy", "Precision", "Recall", "F1_Score", "AUROC", "Model", "Type"}
    }
    controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
    controller.update_model_state(
        model_name,
        status="completed",
        artifact_path=str(artifact_path),
        best_params=best_params,
        best_cv_score=float(checkpoint_state["best_score"]),
        metrics=checkpoint_state["metrics"],
    )
    return metrics, best_params, artifact_path


def _build_rag_contexts(features: np.ndarray, feature_names: Sequence[str], doc_manager: Any) -> List[str]:
    contexts: List[str] = []
    for row in features:
        feature_desc = {name: float(val) for name, val in zip(feature_names, row)}
        query_parts = []
        for symptom_key, col in {
            "tremor": "sym_tremor",
            "rigidity": "sym_rigid",
            "bradykinesia": "sym_brady",
            "postural instability": "sym_posins",
        }.items():
            severity = feature_desc.get(col, 0)
            if severity > 0:
                query_parts.append(f"{symptom_key} severity:{severity:.2f}")
        if feature_desc.get("moca", 30) < 26:
            query_parts.append("cognitive impairment")
        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:
            contexts.append("")
            continue
        contexts.append(
            " ".join(
                f"From '{item.get('doc_title') or item.get('doc_id') or 'document'}' {item['text'][:240]}..."
                for item in passages
            )
        )
    return contexts


def _encode_contexts_for_model(model_name: str, model: Any, contexts: Optional[List[str]]) -> Optional[Dict[str, torch.Tensor]]:
    if not contexts:
        return None
    if model_name == "clinical_t5":
        texts = [f"classify patient: {text}" for text in contexts]
    else:
        texts = contexts
    encoded = model.tokenizer(
        texts,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=512,
    )
    return {key: value.cpu() for key, value in encoded.items()}


def _prepare_batch(batch: Any, device: torch.device):
    if len(batch) == 3:
        features, targets, contexts = batch
        if isinstance(contexts, dict):
            contexts = {key: value.to(device, non_blocking=True) for key, value in contexts.items()}
        else:
            contexts = list(contexts)
    else:
        features, targets = batch
        contexts = None
    return features.to(device), targets.to(device), contexts


def _create_transformer_model(model_name: str, input_dim: int, num_classes: int, model_kwargs: Dict[str, Any]):
    from models.medical_transformers import (  # type: ignore
        BioMistralClassifier as BioGPTForTabular,
        ClinicalT5Classifier as ClinicalT5ForTabular,
        PubMedBERTClassifier as PubMedBERTForTabular,
    )

    if model_name == "pubmedbert":
        return PubMedBERTForTabular(input_dim=input_dim, num_classes=num_classes, **model_kwargs)
    if model_name == "biogpt":
        return BioGPTForTabular(input_dim=input_dim, num_classes=num_classes, **model_kwargs)
    if model_name == "clinical_t5":
        return ClinicalT5ForTabular(input_dim=input_dim, num_classes=num_classes, **model_kwargs)
    raise ValueError(f"Unsupported transformer model: {model_name}")


def _torch_checkpoint_path(controller: TrainingRunController, model_name: str, trial_index: int) -> Path:
    return controller.paths.checkpoints_dir / f"{model_name}_trial{trial_index}.pth"


def _canonical_transformer_artifacts(model_name: str) -> List[Path]:
    aliases = {
        "pubmedbert": ["pubmedbert_transformer.pth", "pubmedbert.pth"],
        "biogpt": ["biogpt_transformer.pth", "biogpt.pth", "biomistral.pth"],
        "clinical_t5": ["clinical_t5_transformer.pth", "clinicalt5_transformer.pth", "clinical_t5.pth"],
    }
    return [MODEL_DIR / filename for filename in aliases[model_name]]


def _run_transformer_search(
    model_name: str,
    trial_space: List[Dict[str, Any]],
    bundle: TrainingBundle,
    controller: TrainingRunController,
    max_trials: int,
    num_epochs: int,
    patience: int,
    use_rag: bool,
    gpu_profile_name: str,
    transformer_loss_name: str,
    focal_gamma: float,
) -> Tuple[Dict[str, Any], Dict[str, Any], Path]:
    from models.transformer_models import TabularDataset  # type: ignore

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    amp_enabled = device.type == "cuda" and model_name != "clinical_t5"
    gpu_profile = _detect_gpu_execution_profile(requested_profile=gpu_profile_name)
    train_bs = gpu_profile.train_batch_by_model.get(model_name, 8)
    eval_bs = gpu_profile.eval_batch_by_model.get(model_name, max(train_bs * 2, 8))
    loader_kwargs = _build_loader_kwargs(device, gpu_profile)

    class_weights = compute_class_weight("balanced", classes=np.unique(bundle.y_train), y=bundle.y_train)
    class_weights_tensor = torch.FloatTensor(class_weights).to(device)
    criterion, criterion_label = _build_transformer_criterion(
        transformer_loss_name,
        class_weights_tensor,
        focal_gamma,
    )

    split = GroupShuffleSplit(n_splits=1, test_size=0.15, random_state=42)
    train_index, val_index = next(split.split(bundle.X_train_dense, bundle.y_train, groups=bundle.train_groups))

    train_features = bundle.X_train_dense[train_index]
    train_targets = bundle.y_train[train_index]
    val_features = bundle.X_train_dense[val_index]
    val_targets = bundle.y_train[val_index]

    train_contexts = None
    val_contexts = None
    if use_rag:
        from document_manager import DocumentManager  # type: ignore

        doc_manager = DocumentManager(docs_dir=str(ROOT / "medical_docs"))
        train_contexts = _build_rag_contexts(train_features, bundle.feature_names, doc_manager)
        val_contexts = _build_rag_contexts(val_features, bundle.feature_names, doc_manager)

    test_ds = TabularDataset(bundle.X_test_dense, bundle.y_test, bundle.feature_names, contexts=None)
    test_loader = DataLoader(test_ds, batch_size=eval_bs, shuffle=False, **loader_kwargs)

    checkpoint_state = controller.load_checkpoint_state(
        model_name,
        {
            "phase": "search",
            "trials": {},
            "best_trial_index": None,
            "best_f1": None,
            "selection_metric": TRANSFORMER_SELECTION_METRIC,
            "loss_name": criterion_label,
            "focal_gamma": float(focal_gamma),
        },
    )
    if (
        checkpoint_state.get("selection_metric") != TRANSFORMER_SELECTION_METRIC
        or checkpoint_state.get("loss_name") != criterion_label
        or float(checkpoint_state.get("focal_gamma", focal_gamma)) != float(focal_gamma)
    ):
        print(
            f"[TRANSFORMER] {model_name} resetting saved trial state to match "
            f"selection={TRANSFORMER_SELECTION_METRIC} and loss={criterion_label}"
        )
        checkpoint_state = {
            "phase": "search",
            "trials": {},
            "best_trial_index": None,
            "best_f1": None,
            "selection_metric": TRANSFORMER_SELECTION_METRIC,
            "loss_name": criterion_label,
            "focal_gamma": float(focal_gamma),
        }
        for stale_checkpoint in controller.paths.checkpoints_dir.glob(f"{model_name}_trial*.pth"):
            stale_checkpoint.unlink(missing_ok=True)
    trial_space = list(trial_space[:max_trials])

    controller.mark_running("transformer_search", model_name=model_name)
    controller.update_model_state(
        model_name,
        status="running",
        family="transformer",
        trials_total=len(trial_space),
        device=str(device),
        gpu_profile=gpu_profile.name,
        train_batch_size=train_bs,
        eval_batch_size=eval_bs,
        num_workers=loader_kwargs.get("num_workers", 0),
        gpu_notes=gpu_profile.notes,
        selection_metric=TRANSFORMER_SELECTION_METRIC,
        transformer_loss=criterion_label,
    )

    for trial_index, trial_cfg in enumerate(trial_space):
        trial_key = str(trial_index)
        trial_state = checkpoint_state["trials"].setdefault(
            trial_key,
            {
                "config": trial_cfg,
                "status": "pending",
                "history": {"train_loss": [], "val_loss": [], "val_f1": [], "val_acc": [], "lr": []},
                "best_val_loss": None,
                "best_val_f1": None,
                "best_epoch": None,
                "selection_metric": TRANSFORMER_SELECTION_METRIC,
                "loss_name": criterion_label,
            },
        )
        if trial_state.get("status") == "completed":
            continue

        controller.mark_running("transformer_trial", model_name=model_name, extra={"trial_index": trial_index})

        model = _create_transformer_model(
            model_name,
            input_dim=bundle.X_train_dense.shape[1],
            num_classes=len(bundle.class_names),
            model_kwargs=trial_cfg["model_kwargs"],
        ).to(device)
        encoded_train_contexts = _encode_contexts_for_model(model_name, model, train_contexts)
        encoded_val_contexts = _encode_contexts_for_model(model_name, model, val_contexts)
        train_ds = TabularDataset(train_features, train_targets, bundle.feature_names, contexts=encoded_train_contexts)
        val_ds = TabularDataset(val_features, val_targets, bundle.feature_names, contexts=encoded_val_contexts)
        train_loader = DataLoader(train_ds, batch_size=train_bs, shuffle=True, **loader_kwargs)
        val_loader = DataLoader(val_ds, batch_size=eval_bs, shuffle=False, **loader_kwargs)
        optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), **trial_cfg["optimizer"])
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", factor=0.5, patience=2, min_lr=1e-7)
        scaler = torch.amp.GradScaler(device=device.type, enabled=amp_enabled)
        grad_accum = min(
            int(trial_cfg.get("grad_accum", 8)),
            int(gpu_profile.grad_accum_cap_by_model.get(model_name, trial_cfg.get("grad_accum", 8))),
        )
        print(
            f"[TRANSFORMER] {model_name} trial {trial_index + 1}/{len(trial_space)} "
            f"| train_bs={train_bs} eval_bs={eval_bs} grad_accum={grad_accum} "
            f"effective_batch={train_bs * grad_accum} "
            f"train_steps={len(train_loader)} val_steps={len(val_loader)} "
            f"| selection={TRANSFORMER_SELECTION_METRIC} loss={criterion_label}"
        )
        start_epoch = int(trial_state.get("next_epoch", 0))
        early_stop_counter = int(trial_state.get("early_stop_counter", 0))
        best_val_loss = float(trial_state["best_val_loss"]) if trial_state.get("best_val_loss") is not None else float("inf")
        best_val_f1 = float(trial_state["best_val_f1"]) if trial_state.get("best_val_f1") is not None else float("-inf")
        best_model_state_dict = None
        torch_ckpt_path = _torch_checkpoint_path(controller, model_name, trial_index)

        if torch_ckpt_path.exists():
            saved = torch.load(torch_ckpt_path, map_location=device, weights_only=False)
            saved_model_state = saved.get("model_state_dict", {})
            checkpoint_has_non_finite = any(
                torch.is_tensor(tensor) and not torch.isfinite(tensor).all()
                for tensor in saved_model_state.values()
            )
            if (
                saved.get("selection_metric") == TRANSFORMER_SELECTION_METRIC
                and saved.get("loss_name") == criterion_label
                and float(saved.get("focal_gamma", focal_gamma)) == float(focal_gamma)
                and not checkpoint_has_non_finite
            ):
                model.load_state_dict(saved["model_state_dict"])
                optimizer.load_state_dict(saved["optimizer_state_dict"])
                scheduler.load_state_dict(saved["scheduler_state_dict"])
                if saved.get("scaler_state_dict") and device.type == "cuda":
                    scaler.load_state_dict(saved["scaler_state_dict"])
                start_epoch = int(saved.get("epoch", -1)) + 1
                early_stop_counter = int(saved.get("early_stop_counter", early_stop_counter))
                best_val_loss = float(saved.get("best_val_loss", best_val_loss))
                best_val_f1 = float(saved.get("best_val_f1", best_val_f1))
                best_model_state_dict = saved.get("best_model_state_dict")
                trial_state["history"] = saved.get("history", trial_state["history"])
            else:
                print(
                    f"[TRANSFORMER] {model_name} trial {trial_index + 1} ignoring stale "
                    f"checkpoint with incompatible selection/loss settings or non-finite weights"
                )
                start_epoch = 0
                early_stop_counter = 0
                best_val_loss = float("inf")
                best_val_f1 = float("-inf")
                best_model_state_dict = None
                trial_state["history"] = {"train_loss": [], "val_loss": [], "val_f1": [], "val_acc": [], "lr": []}
                torch_ckpt_path.unlink(missing_ok=True)

        for epoch in range(start_epoch, num_epochs):
            model.train()
            running_train_loss = 0.0
            processed_train_batches = 0
            optimizer.zero_grad(set_to_none=True)
            if device.type == "cuda":
                torch.cuda.reset_peak_memory_stats(device)

            progress = tqdm(
                train_loader,
                total=len(train_loader),
                dynamic_ncols=True,
                leave=False,
                desc=f"{model_name} t{trial_index + 1}/{len(trial_space)} e{epoch + 1}/{num_epochs}",
            )
            for batch_index, batch in enumerate(progress):
                features, targets, contexts = _prepare_batch(batch, device)
                with torch.amp.autocast(device_type=device.type, enabled=amp_enabled):
                    outputs = model(features, contexts)
                    loss = criterion(outputs, targets) / grad_accum

                if not torch.isfinite(outputs).all() or not torch.isfinite(loss):
                    optimizer.zero_grad(set_to_none=True)
                    continue

                scaler.scale(loss).backward()

                if (batch_index + 1) % grad_accum == 0 or (batch_index + 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)

                running_train_loss += float(loss.item() * grad_accum)
                processed_train_batches += 1
                current_loss = float(loss.item() * grad_accum)
                current_lr = float(optimizer.param_groups[0]["lr"])
                if device.type == "cuda":
                    gpu_mem_gb = torch.cuda.memory_allocated(device) / 1024**3
                    progress.set_postfix(loss=f"{current_loss:.4f}", lr=f"{current_lr:.2e}", gpu_gb=f"{gpu_mem_gb:.2f}")
                else:
                    progress.set_postfix(loss=f"{current_loss:.4f}", lr=f"{current_lr:.2e}")

            progress.close()

            avg_train_loss = running_train_loss / max(processed_train_batches, 1)

            model.eval()
            running_val_loss = 0.0
            processed_val_batches = 0
            all_preds: List[int] = []
            all_targets: List[int] = []
            with torch.no_grad():
                for batch in val_loader:
                    features, targets, contexts = _prepare_batch(batch, device)
                    with torch.amp.autocast(device_type=device.type, enabled=amp_enabled):
                        outputs = model(features, contexts)
                        val_loss = criterion(outputs, targets)
                    if not torch.isfinite(outputs).all() or not torch.isfinite(val_loss):
                        continue
                    running_val_loss += float(val_loss.item())
                    processed_val_batches += 1
                    preds = torch.argmax(outputs, dim=1)
                    all_preds.extend(preds.cpu().numpy().tolist())
                    all_targets.extend(targets.cpu().numpy().tolist())

            avg_val_loss = running_val_loss / max(processed_val_batches, 1)
            if not all_preds:
                val_acc = 0.0
                val_f1 = 0.0
                avg_val_loss = float("inf")
            else:
                val_acc = float(np.mean(np.asarray(all_preds) == np.asarray(all_targets)))
                val_f1 = float(f1_score(all_targets, all_preds, average="weighted"))
            scheduler.step(val_f1)

            history = trial_state["history"]
            history["train_loss"].append(float(avg_train_loss))
            history["val_loss"].append(float(avg_val_loss))
            history["val_f1"].append(val_f1)
            history["val_acc"].append(val_acc)
            history["lr"].append(float(optimizer.param_groups[0]["lr"]))
            gpu_peak_gb = 0.0
            if device.type == "cuda":
                gpu_peak_gb = torch.cuda.max_memory_allocated(device) / 1024**3
            print(
                f"[TRANSFORMER] {model_name} trial {trial_index + 1}/{len(trial_space)} "
                f"epoch {epoch + 1}/{num_epochs} "
                f"train_loss={avg_train_loss:.4f} val_loss={avg_val_loss:.4f} "
                f"val_f1={val_f1:.4f} val_acc={val_acc:.4f} "
                f"lr={optimizer.param_groups[0]['lr']:.2e} "
                f"gpu_peak_gb={gpu_peak_gb:.2f}"
            )

            if _is_better_validation_epoch(
                val_f1=val_f1,
                val_loss=avg_val_loss,
                best_val_f1=best_val_f1,
                best_val_loss=best_val_loss,
            ):
                best_val_loss = avg_val_loss
                best_val_f1 = val_f1
                trial_state["best_val_loss"] = float(best_val_loss)
                trial_state["best_val_f1"] = float(best_val_f1)
                trial_state["best_epoch"] = epoch
                best_model_state_dict = {
                    key: value.detach().cpu().clone()
                    for key, value in model.state_dict().items()
                }
                early_stop_counter = 0
            else:
                early_stop_counter += 1

            trial_state["status"] = "running"
            trial_state["next_epoch"] = epoch + 1
            trial_state["early_stop_counter"] = early_stop_counter
            checkpoint_state["phase"] = "search"
            controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
            torch.save(
                {
                    "epoch": epoch,
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "scheduler_state_dict": scheduler.state_dict(),
                    "scaler_state_dict": scaler.state_dict() if device.type == "cuda" else None,
                    "history": history,
                    "early_stop_counter": early_stop_counter,
                    "best_val_loss": best_val_loss,
                    "best_val_f1": best_val_f1,
                    "best_model_state_dict": best_model_state_dict,
                    "selection_metric": TRANSFORMER_SELECTION_METRIC,
                    "loss_name": criterion_label,
                    "focal_gamma": float(focal_gamma),
                },
                torch_ckpt_path,
            )
            controller.update_model_state(
                model_name,
                current_trial=trial_index,
                current_epoch=epoch,
                grad_accum=grad_accum,
                best_val_f1=float(best_val_f1),
                best_val_loss=float(best_val_loss),
                selection_metric=TRANSFORMER_SELECTION_METRIC,
                transformer_loss=criterion_label,
            )
            controller.raise_if_requested()

            if early_stop_counter >= patience:
                print(
                    f"[TRANSFORMER] {model_name} trial {trial_index + 1}/{len(trial_space)} "
                    f"early-stopped after epoch {epoch + 1}"
                )
                break

        trial_state["status"] = "completed"
        controller.append_trial_result(
            model_name,
            {
                "trial_index": trial_index,
                "config": trial_cfg,
                "best_val_f1": float(trial_state.get("best_val_f1") or 0.0),
                "best_epoch": trial_state.get("best_epoch"),
            },
        )
        best_so_far = checkpoint_state.get("best_f1")
        if best_so_far is None or float(trial_state.get("best_val_f1") or -1.0) > float(best_so_far):
            checkpoint_state["best_f1"] = float(trial_state["best_val_f1"])
            checkpoint_state["best_trial_index"] = trial_index
        controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
        controller.raise_if_requested()

    if checkpoint_state.get("best_trial_index") is None:
        raise RuntimeError(f"No completed transformer trial found for {model_name}")

    best_trial_index = int(checkpoint_state["best_trial_index"])
    best_trial_cfg = dict(trial_space[best_trial_index])
    checkpoint_state["phase"] = "evaluate"
    controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))

    best_model = _create_transformer_model(
        model_name,
        input_dim=bundle.X_train_dense.shape[1],
        num_classes=len(bundle.class_names),
        model_kwargs=best_trial_cfg["model_kwargs"],
    ).to(device)
    best_state = torch.load(_torch_checkpoint_path(controller, model_name, best_trial_index), map_location=device, weights_only=False)
    best_model.load_state_dict(best_state.get("best_model_state_dict") or best_state["model_state_dict"])
    best_model.eval()

    all_preds: List[int] = []
    all_targets: List[int] = []
    all_probabilities: List[np.ndarray] = []
    with torch.no_grad():
        for batch in test_loader:
            features, targets, contexts = _prepare_batch(batch, device)
            outputs = best_model(features, contexts)
            probs = torch.softmax(outputs, dim=1)
            preds = torch.argmax(outputs, dim=1)
            all_targets.extend(targets.cpu().numpy().tolist())
            all_preds.extend(preds.cpu().numpy().tolist())
            all_probabilities.append(probs.cpu().numpy())

    probabilities = np.vstack(all_probabilities) if all_probabilities else None
    metrics = _evaluate_predictions(
        model_name={"pubmedbert": "PubMedBERT", "biogpt": "BioGPT", "clinical_t5": "Clinical-T5"}[model_name],
        model_type="Transformer",
        y_true=np.asarray(all_targets),
        y_pred=np.asarray(all_preds),
        probabilities=probabilities,
        class_names=bundle.class_names,
    )

    primary_artifact = _canonical_transformer_artifacts(model_name)[0]
    for path in _canonical_transformer_artifacts(model_name):
        torch.save(best_model.state_dict(), path)

    checkpoint_state["phase"] = "complete"
    checkpoint_state["artifact_path"] = str(primary_artifact)
    checkpoint_state["best_trial_index"] = best_trial_index
    checkpoint_state["best_config"] = best_trial_cfg
    checkpoint_state["metrics"] = {
        key: value for key, value in metrics.items()
        if key in {"Accuracy", "Precision", "Recall", "F1_Score", "AUROC", "Model", "Type"}
    }
    controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
    controller.update_model_state(
        model_name,
        status="completed",
        artifact_path=str(primary_artifact),
        best_trial_index=best_trial_index,
        best_config=best_trial_cfg,
        metrics=checkpoint_state["metrics"],
    )
    return metrics, best_trial_cfg, primary_artifact


def _train_ensemble(bundle: TrainingBundle, controller: TrainingRunController) -> Optional[Dict[str, Any]]:
    from models.multimodal_ml import MultimodalEnsemble  # type: ignore

    controller.mark_running("ensemble_training", model_name="ensemble")
    controller.raise_if_requested()

    ensemble = MultimodalEnsemble()
    ensemble.load_traditional_models(model_dir=str(MODEL_DIR))
    ensemble.load_transformer_models(model_dir=str(MODEL_DIR), input_dim=bundle.X_train_dense.shape[1], num_classes=len(bundle.class_names))
    ensemble.train_ensemble(bundle.X_train_dense, bundle.y_train, ensemble_type="stacking")
    ensemble.save_ensemble(str(MODEL_DIR / "multimodal_ensemble.joblib"))
    results = ensemble.evaluate_ensemble(bundle.X_test_dense, bundle.y_test)

    probabilities = results.get("probabilities")
    metrics = _evaluate_predictions(
        model_name="Multimodal Ensemble",
        model_type="Ensemble",
        y_true=bundle.y_test,
        y_pred=np.asarray(results["predictions"]),
        probabilities=np.asarray(probabilities) if probabilities is not None else None,
        class_names=bundle.class_names,
    )
    controller.update_model_state(
        "ensemble",
        status="completed",
        artifact_path=str(MODEL_DIR / "multimodal_ensemble.joblib"),
        metrics={key: value for key, value in metrics.items() if key in {"Accuracy", "Precision", "Recall", "F1_Score", "AUROC", "Model", "Type"}},
    )
    return metrics


def _write_metric_outputs(base_metrics: List[Dict[str, Any]], ensemble_metrics: Optional[Dict[str, Any]]) -> None:
    EVAL_DIR.mkdir(parents=True, exist_ok=True)
    MODEL_METRICS_DIR.mkdir(parents=True, exist_ok=True)

    summary_rows = []
    traditional_rows = []
    transformer_rows = []
    for metrics in base_metrics:
        summary_rows.append({
            "Model": metrics["Model"],
            "Type": metrics["Type"],
            "Accuracy": metrics["Accuracy"],
            "Precision": metrics["Precision"],
            "Recall": metrics["Recall"],
            "F1_Score": metrics["F1_Score"],
            "AUROC": metrics["AUROC"] if metrics["AUROC"] is not None else "",
        })
        if "traditional" in metrics["Type"].lower():
            traditional_rows.append(metrics)
        if "transformer" in metrics["Type"].lower():
            transformer_rows.append(metrics)

    summary_df = pd.DataFrame(summary_rows)
    summary_df.to_csv(EVAL_DIR / "summary_metrics.csv", index=False)
    summary_df.to_csv(MODEL_METRICS_DIR / "model_metrics_summary.csv", index=False)

    (EVAL_DIR / "traditional_metrics_latest.json").write_text(json.dumps(_json_ready(traditional_rows), indent=2), encoding="utf-8")
    (EVAL_DIR / "transformer_metrics_latest.json").write_text(json.dumps(_json_ready(transformer_rows), indent=2), encoding="utf-8")
    if ensemble_metrics is not None:
        (EVAL_DIR / "ensemble_metrics_latest.json").write_text(json.dumps(_json_ready(ensemble_metrics), indent=2), encoding="utf-8")


def _write_run_manifest_snapshot(controller: TrainingRunController, base_metrics: List[Dict[str, Any]], ensemble_metrics: Optional[Dict[str, Any]]) -> None:
    controller.write_metrics_file(
        "final_metrics.json",
        {
            "base_metrics": _json_ready(base_metrics),
            "ensemble_metrics": _json_ready(ensemble_metrics),
        },
    )


def _parse_selected_models(raw: str) -> List[str]:
    raw = (raw or "all").strip().lower()
    if raw in {"all", "*"}:
        return list(ALL_BASE_MODELS)
    aliases = {
        "lgbm": "lightgbm",
        "lightgbm": "lightgbm",
        "xgb": "xgboost",
        "xgboost": "xgboost",
        "svm": "svm",
        "pubmed": "pubmedbert",
        "pubmedbert": "pubmedbert",
        "biogpt": "biogpt",
        "bio": "biogpt",
        "clinical": "clinical_t5",
        "clinical_t5": "clinical_t5",
        "t5": "clinical_t5",
    }
    selected: List[str] = []
    for piece in [part.strip() for part in raw.split(",") if part.strip()]:
        canonical = aliases.get(piece)
        if canonical and canonical not in selected:
            selected.append(canonical)
    if not selected:
        raise ValueError(f"No valid models were selected from '{raw}'")
    return selected


def _detect_gpu_execution_profile(
    requested_profile: str = "auto",
    cuda_available: Optional[bool] = None,
    device_name: Optional[str] = None,
    total_memory_gb: Optional[float] = None,
) -> GPUExecutionProfile:
    cuda_ready = torch.cuda.is_available() if cuda_available is None else cuda_available
    if not cuda_ready:
        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_cap_by_model={"pubmedbert": 12, "biogpt": 16, "clinical_t5": 16},
            num_workers=0,
            prefetch_factor=2,
            persistent_workers=False,
            notes="CPU fallback profile.",
        )

    name = (device_name or torch.cuda.get_device_name(0)).lower()
    memory_gb = float(total_memory_gb if total_memory_gb is not None else (torch.cuda.get_device_properties(0).total_memory / 1024**3))
    requested = (requested_profile or "auto").strip().lower()

    if requested in {"rtx-a4000", "a4000"} or (requested == "auto" and ("a4000" in name or memory_gb >= 15.0)):
        return GPUExecutionProfile(
            name="rtx-a4000",
            train_batch_by_model={"pubmedbert": 32, "biogpt": 12, "clinical_t5": 12},
            eval_batch_by_model={"pubmedbert": 64, "biogpt": 24, "clinical_t5": 24},
            grad_accum_cap_by_model={"pubmedbert": 2, "biogpt": 5, "clinical_t5": 5},
            num_workers=6,
            prefetch_factor=4,
            persistent_workers=True,
            notes="Optimized for RTX A4000 / ~16 GB VRAM with larger per-step batches.",
        )

    if requested in {"high-vram", "16gb"} or (requested == "auto" and 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_cap_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_cap_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: torch.device, profile: GPUExecutionProfile) -> Dict[str, Any]:
    kwargs: Dict[str, Any] = {
        "num_workers": profile.num_workers if device.type == "cuda" else 0,
        "pin_memory": device.type == "cuda",
    }
    if kwargs["num_workers"] > 0:
        kwargs["persistent_workers"] = profile.persistent_workers
        kwargs["prefetch_factor"] = profile.prefetch_factor
    return kwargs


def _configure_transformer_runtime(selected_models: Sequence[str], allow_cpu_transformers: bool, cuda_available: Optional[bool] = None) -> str:
    cuda_ready = torch.cuda.is_available() if cuda_available is None else cuda_available
    transformer_requested = any(model_name in TRANSFORMER_MODELS for model_name in selected_models)
    if not transformer_requested:
        return "not-applicable"
    if not cuda_ready and not allow_cpu_transformers:
        raise RuntimeError(
            "Transformer retraining requires CUDA for this accuracy profile. "
            "Use a CUDA-enabled PyTorch install/GPU, or rerun with --allow-cpu-transformers if you explicitly accept slower, lower-throughput training."
        )
    if cuda_ready:
        torch.set_float32_matmul_precision("high")
        if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
            torch.backends.cuda.matmul.allow_tf32 = True
        if hasattr(torch.backends, "cudnn"):
            torch.backends.cudnn.allow_tf32 = True
            torch.backends.cudnn.benchmark = True
        return "cuda"
    return "cpu"


def _run_training(args: argparse.Namespace) -> int:
    selected_models = _parse_selected_models(args.models)
    transformer_runtime = _configure_transformer_runtime(
        selected_models,
        allow_cpu_transformers=args.allow_cpu_transformers,
    )
    controller = TrainingRunController(RUNS_DIR, args.run_name)
    controller.initialize(
        selected_models=selected_models,
        config={
            "epochs": args.epochs,
            "patience": args.patience,
            "traditional_trials": args.traditional_trials,
            "transformer_trials": args.transformer_trials,
            "use_rag": args.use_rag,
            "train_ensemble": not args.skip_ensemble,
            "allow_cpu_transformers": args.allow_cpu_transformers,
            "transformer_runtime": transformer_runtime,
            "gpu_profile": args.gpu_profile,
            "transformer_loss": args.transformer_loss,
            "focal_gamma": args.focal_gamma,
        },
        resume=args.resume,
    )

    if args.dry_run:
        controller.mark_paused("dry-run initialized")
        print(json.dumps(controller.status_summary(), indent=2))
        return 0

    try:
        controller.clear_stop()
        if args.resume:
            controller.clear_pause()

        bundle = _prepare_training_bundle()
        base_metrics: List[Dict[str, Any]] = []

        for model_name in selected_models:
            if model_name in TRADITIONAL_MODELS:
                metrics, best_config, artifact_path = _run_grouped_traditional_search(
                    model_name=model_name,
                    search_space=TRADITIONAL_SEARCH_SPACES[model_name],
                    bundle=bundle,
                    controller=controller,
                    max_trials=args.traditional_trials,
                )
            else:
                metrics, best_config, artifact_path = _run_transformer_search(
                    model_name=model_name,
                    trial_space=TRANSFORMER_TRIALS[model_name],
                    bundle=bundle,
                    controller=controller,
                    max_trials=args.transformer_trials,
                    num_epochs=args.epochs,
                    patience=args.patience,
                    use_rag=args.use_rag,
                    gpu_profile_name=args.gpu_profile,
                    transformer_loss_name=args.transformer_loss,
                    focal_gamma=args.focal_gamma,
                )
            base_metrics.append(metrics)
            controller.update_model_state(
                model_name,
                best_config=_json_ready(best_config),
                artifact_path=str(artifact_path),
                metrics={key: value for key, value in metrics.items() if key in {"Accuracy", "Precision", "Recall", "F1_Score", "AUROC", "Model", "Type"}},
            )

        ensemble_metrics = None
        should_train_ensemble = (not args.skip_ensemble) and set(selected_models) == set(ALL_BASE_MODELS)
        if should_train_ensemble:
            ensemble_metrics = _train_ensemble(bundle, controller)
        elif not args.skip_ensemble:
            controller.update_model_state(
                "ensemble",
                status="skipped",
                reason="ensemble retraining is only automatic when all six base models are selected",
            )

        _write_metric_outputs(base_metrics, ensemble_metrics)
        _write_run_manifest_snapshot(controller, base_metrics, ensemble_metrics)
        controller.mark_completed()
        print(json.dumps(controller.status_summary(), indent=2))
        return 0
    except PauseRequested as exc:
        controller.mark_paused(str(exc))
        print(json.dumps(controller.status_summary(), indent=2))
        return 0
    except StopRequested as exc:
        controller.mark_stopped(str(exc))
        print(json.dumps(controller.status_summary(), indent=2))
        return 0
    except KeyboardInterrupt:
        controller.mark_paused("keyboard interrupt")
        print(json.dumps(controller.status_summary(), indent=2))
        return 1
    except Exception as exc:
        controller.mark_failed(str(exc))
        raise


def _run_status(args: argparse.Namespace) -> int:
    controller = TrainingRunController(RUNS_DIR, args.run_name)
    print(json.dumps(controller.status_summary(), indent=2))
    return 0


def _run_pause(args: argparse.Namespace) -> int:
    controller = TrainingRunController(RUNS_DIR, args.run_name)
    controller.request_pause()
    controller.mark_paused("pause requested from CLI")
    print(json.dumps(controller.status_summary(), indent=2))
    return 0


def _run_stop(args: argparse.Namespace) -> int:
    controller = TrainingRunController(RUNS_DIR, args.run_name)
    controller.request_stop()
    controller.mark_stopped("stop requested from CLI")
    print(json.dumps(controller.status_summary(), indent=2))
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Retrain the six base Parkinson's models with resumable checkpoints.")
    subparsers = parser.add_subparsers(dest="command", required=True)

    train_parser = subparsers.add_parser("train", help="Start a new training run.")
    train_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
    train_parser.add_argument("--models", default="all", help="Comma-separated list of models or 'all'.")
    train_parser.add_argument("--epochs", type=int, default=30, help="Max epochs per transformer trial.")
    train_parser.add_argument("--patience", type=int, default=8, help="Early-stopping patience for transformers.")
    train_parser.add_argument("--traditional-trials", type=int, default=6, help="Number of traditional-model trials per model.")
    train_parser.add_argument("--transformer-trials", type=int, default=6, help="Number of transformer trials per model.")
    train_parser.add_argument("--use-rag", dest="use_rag", action="store_true", help="Build RAG contexts for transformer training.")
    train_parser.add_argument("--no-rag", dest="use_rag", action="store_false", help="Disable RAG context enrichment for transformer training.")
    train_parser.add_argument("--gpu-profile", default="auto", choices=["auto", "rtx-a4000", "high-vram", "compat"], help="CUDA loader/batch preset for transformer training.")
    train_parser.add_argument("--transformer-loss", default=DEFAULT_TRANSFORMER_LOSS, choices=["cross_entropy", "focal"], help="Loss function for transformer fine-tuning.")
    train_parser.add_argument("--focal-gamma", type=float, default=DEFAULT_FOCAL_GAMMA, help="Gamma value for focal loss when --transformer-loss focal is used.")
    train_parser.add_argument("--allow-cpu-transformers", action="store_true", help="Allow transformer training on CPU when CUDA is unavailable.")
    train_parser.add_argument("--skip-ensemble", action="store_true", help="Skip ensemble retraining after the six base models.")
    train_parser.add_argument("--resume", action="store_true", help="Resume an existing run from saved state.")
    train_parser.add_argument("--dry-run", action="store_true", help="Initialize the run manifest without training.")
    train_parser.set_defaults(func=_run_training, use_rag=True)

    resume_parser = subparsers.add_parser("resume", help="Resume a paused training run.")
    resume_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
    resume_parser.add_argument("--models", default="all", help="Comma-separated list of models or 'all'.")
    resume_parser.add_argument("--epochs", type=int, default=30, help="Max epochs per transformer trial.")
    resume_parser.add_argument("--patience", type=int, default=8, help="Early-stopping patience for transformers.")
    resume_parser.add_argument("--traditional-trials", type=int, default=6, help="Number of traditional-model trials per model.")
    resume_parser.add_argument("--transformer-trials", type=int, default=6, help="Number of transformer trials per model.")
    resume_parser.add_argument("--use-rag", dest="use_rag", action="store_true", help="Build RAG contexts for transformer training.")
    resume_parser.add_argument("--no-rag", dest="use_rag", action="store_false", help="Disable RAG context enrichment for transformer training.")
    resume_parser.add_argument("--gpu-profile", default="auto", choices=["auto", "rtx-a4000", "high-vram", "compat"], help="CUDA loader/batch preset for transformer training.")
    resume_parser.add_argument("--transformer-loss", default=DEFAULT_TRANSFORMER_LOSS, choices=["cross_entropy", "focal"], help="Loss function for transformer fine-tuning.")
    resume_parser.add_argument("--focal-gamma", type=float, default=DEFAULT_FOCAL_GAMMA, help="Gamma value for focal loss when --transformer-loss focal is used.")
    resume_parser.add_argument("--allow-cpu-transformers", action="store_true", help="Allow transformer training on CPU when CUDA is unavailable.")
    resume_parser.add_argument("--skip-ensemble", action="store_true", help="Skip ensemble retraining after the six base models.")
    resume_parser.add_argument("--dry-run", action="store_true", help="Load and print the current run state without training.")
    resume_parser.set_defaults(func=_run_training, resume=True, use_rag=True)

    pause_parser = subparsers.add_parser("pause", help="Request a graceful pause for a running training job.")
    pause_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
    pause_parser.set_defaults(func=_run_pause)

    stop_parser = subparsers.add_parser("stop", help="Request a graceful stop for a running training job.")
    stop_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
    stop_parser.set_defaults(func=_run_stop)

    status_parser = subparsers.add_parser("status", help="Show the current run manifest.")
    status_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
    status_parser.set_defaults(func=_run_status)

    return parser


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    return int(args.func(args))


if __name__ == "__main__":
    raise SystemExit(main())