DeepHealth/train.py

"""
Training script for DeepHealth model.

Implements the complete training pipeline:
1. Load data and split train/val/test
2. Build model, optimizer, readout, and loss
3. Train with adaptive learning rate (warmup + cosine annealing)
4. Early stopping based on validation loss
5. Save checkpoints and metrics
"""

from __future__ import annotations

import argparse
import json
import logging
import math
import os
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Tuple

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, RandomSampler, Subset
from torch.optim import AdamW
from torch.nn.utils import clip_grad_norm_
from tqdm.auto import tqdm

from dataset import HealthDataset, collate_fn
from models import DeepHealth
from readouts import build_readout
from losses import build_loss
from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX


# ---------------------------------------------------------------------------
# Setup Logging
# ---------------------------------------------------------------------------

def setup_logging(run_dir: Path) -> logging.Logger:
    """Configure logging to both console and file."""
    run_dir.mkdir(parents=True, exist_ok=True)
    log_file = run_dir / "train.log"

    logger = logging.getLogger("DeepHealth")
    logger.setLevel(logging.INFO)
    logger.handlers.clear()

    # Console handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_formatter = logging.Formatter(
        "%(asctime)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )
    console_handler.setFormatter(console_formatter)
    logger.addHandler(console_handler)

    # File handler
    file_handler = logging.FileHandler(log_file, mode="w")
    file_handler.setLevel(logging.INFO)
    file_formatter = logging.Formatter(
        "%(asctime)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)

    return logger


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def set_seed(seed: int) -> None:
    """Set random seed for reproducibility."""
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)


def load_extra_info_types_file(path: str) -> list[int]:
    """
    Load other-information type ids from a text or JSON file.

    Text files may use whitespace, commas, or semicolons as separators. Lines may
    contain comments after '#'. JSON files should contain a top-level list of ids.
    """
    file_path = Path(path)
    if not file_path.exists():
        raise FileNotFoundError(f"extra_info_types_file not found: {path}")
    if not file_path.is_file():
        raise ValueError(f"extra_info_types_file is not a file: {path}")

    text = file_path.read_text(encoding="utf-8").strip()
    if not text:
        return []

    if text.startswith("["):
        try:
            raw_items = json.loads(text)
        except json.JSONDecodeError as exc:
            raise ValueError(
                f"Invalid JSON in extra_info_types_file: {path}"
            ) from exc
        if not isinstance(raw_items, list):
            raise ValueError(
                f"extra_info_types_file JSON must be a list, got {type(raw_items).__name__}"
            )
    else:
        tokens: list[str] = []
        for line in text.splitlines():
            line = line.split("#", 1)[0].strip()
            if not line:
                continue
            tokens.extend(line.replace(",", " ").replace(";", " ").split())
        raw_items = tokens

    parsed: list[int] = []
    for item in raw_items:
        try:
            type_id = int(item)
        except (TypeError, ValueError) as exc:
            raise ValueError(
                f"Invalid extra info type id {item!r} in {path}; expected integers."
            ) from exc
        parsed.append(type_id)

    return parsed


def configure_torch_for_training(device: torch.device) -> None:
    """Enable backend settings that can improve training throughput on CUDA."""
    if device.type == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        if hasattr(torch, "set_float32_matmul_precision"):
            torch.set_float32_matmul_precision("high")


def resolve_device(device_arg: str) -> torch.device:
    """Resolve and validate requested device string."""
    requested = device_arg.strip().lower()

    if requested == "cpu":
        return torch.device("cpu")

    if requested == "cuda":
        if not torch.cuda.is_available():
            return torch.device("cpu")
        return torch.device("cuda")

    if requested.startswith("cuda:"):
        if not torch.cuda.is_available():
            return torch.device("cpu")
        try:
            index = int(requested.split(":", 1)[1])
        except ValueError as exc:
            raise ValueError(
                f"Invalid CUDA device format '{device_arg}'. Use 'cuda' or 'cuda:<index>'."
            ) from exc

        n_cuda = torch.cuda.device_count()
        if index < 0 or index >= n_cuda:
            raise ValueError(
                f"Requested device '{device_arg}' is out of range. "
                f"Available CUDA devices: 0..{max(0, n_cuda - 1)}"
            )
        return torch.device(f"cuda:{index}")

    raise ValueError(
        f"Unsupported device '{device_arg}'. Use 'cpu', 'cuda', or 'cuda:<index>'."
    )


def split_dataset(
    dataset: HealthDataset,
    train_ratio: float,
    val_ratio: float,
    test_ratio: float,
    seed: int,
) -> Tuple[Subset, Subset, Subset]:
    """
    Split dataset into train/val/test subsets.

    Parameters
    ----------
    train_ratio, val_ratio, test_ratio : float
        Ratios must sum to 1.0 (within 1e-6 tolerance).
    seed : int
        Random seed for splitting.

    Returns
    -------
    train_subset, val_subset, test_subset
    """
    total = train_ratio + val_ratio + test_ratio
    if not np.isclose(total, 1.0, atol=1e-6):
        raise ValueError(
            f"train_ratio + val_ratio + test_ratio must equal 1.0, "
            f"got {total}"
        )

    n = len(dataset)
    rng = np.random.RandomState(seed)
    indices = rng.permutation(n)

    n_train = int(n * train_ratio)
    n_val = int(n * val_ratio)

    train_indices = indices[:n_train]
    val_indices = indices[n_train: n_train + n_val]
    test_indices = indices[n_train + n_val:]

    return (
        Subset(dataset, train_indices),
        Subset(dataset, val_indices),
        Subset(dataset, test_indices),
    )


def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
    """
    Build DeepHealth model using metadata from dataset.

    Uses unified other-information metadata computed during dataset initialization.
    """
    model = DeepHealth(
        vocab_size=dataset.vocab_size,
        n_embd=args.n_embd,
        n_head=args.n_head,
        n_hist_layer=args.n_hist_layer,
        n_tab_layer=args.n_tab_layer,
        n_types=dataset.n_types,
        n_cont_types=dataset.n_cont_types,
        n_categories=dataset.n_categories,
        cont_type_ids=dataset.cont_type_ids,
        n_bins=args.n_bins,
        target_mode="next_token",
        time_mode=args.time_mode,
        dist_mode="exponential",
        dropout=args.dropout,
    )

    return model


def build_optimizer(args: argparse.Namespace, model: nn.Module) -> AdamW:
    """Build AdamW optimizer."""
    return AdamW(
        model.parameters(),
        lr=args.base_lr,
        betas=args.betas,
        weight_decay=args.weight_decay,
    )


def get_lr(
    epoch: int,
    args: argparse.Namespace,
    adaptive_lr: float,
    total_epochs: int,
) -> float:
    """
    Calculate learning rate for the given epoch.

    Warmup: linear from 0 to adaptive_lr over warmup_epochs
    After: cosine annealing from adaptive_lr to adaptive_lr * min_lr_ratio
    """
    if epoch < args.warmup_epochs:
        # Linear warmup
        return adaptive_lr * (epoch + 1) / args.warmup_epochs
    else:
        # Cosine annealing
        progress = (epoch - args.warmup_epochs) / \
            (total_epochs - args.warmup_epochs)
        return adaptive_lr * (args.min_lr_ratio + 0.5 * (1 + math.cos(math.pi * progress)) * (1 - args.min_lr_ratio))


def set_optimizer_lr(optimizer: AdamW, lr: float) -> None:
    """Set learning rate for all parameter groups."""
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr


def move_batch_to_device(batch: Dict, device: torch.device) -> Dict:
    """Move all tensors in batch to specified device."""
    non_blocking = device.type == "cuda"
    return {
        key: value.to(device, non_blocking=non_blocking) if isinstance(
            value, torch.Tensor) else value
        for key, value in batch.items()
    }


def compute_loss(
    args: argparse.Namespace,
    model: DeepHealth,
    readout: nn.Module,
    criterion: nn.Module,
    batch: Dict[str, torch.Tensor],
    device: torch.device,
) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
    """
    Compute loss for one batch.

    Flow: forward model, apply readout, compute risk logits, compute loss.

    Parameters
    ----------
    args : argparse.Namespace
        Training configuration with target_mode and loss options.
    model : DeepHealth
        The model.
    readout : nn.Module
        Readout module.
    criterion : nn.Module
        Loss criterion.
    batch : Dict
        Batch data from DataLoader.
    device : torch.device
        Device to compute on.

    Returns
    -------
    loss : torch.Tensor
        Scalar loss tensor.
    """
    # Move batch to device
    batch = move_batch_to_device(batch, device)

    event_seq = batch["event_seq"]           # (B, L)
    time_seq = batch["time_seq"]             # (B, L)
    padding_mask = batch["padding_mask"]     # (B, L)
    sex = batch["sex"]                       # (B,)

    hidden = model(
        event_seq=event_seq,
        time_seq=time_seq,
        sex=sex,
        padding_mask=padding_mask,
        other_type=batch["other_type"],
        other_value=batch["other_value"],
        other_value_kind=batch["other_value_kind"],
        other_time=batch["other_time"],
        target_mode="next_token",
    )

    # Apply readout
    readout_mask = (
        batch["readout_mask"]
        if args.readout_name == "same_time_group_end"
        else None
    )
    readout_out = readout(
        hidden=hidden,
        time_seq=time_seq,
        padding_mask=padding_mask,
        readout_mask=readout_mask,
    )

    # Compute risk logits
    logits = model.calc_risk(readout_out.hidden)

    # Compute loss based on target_mode
    if args.target_mode == "delphi2m":
        loss_out = criterion(
            logits=logits,
            target_events=batch["target_event_seq"],
            target_times=batch["target_time_seq"],
            current_times=batch["time_seq"],
            padding_mask=readout_out.readout_mask,
            return_components=True,
        )
    elif args.target_mode == "uts":
        loss_out = criterion(
            logits=logits,
            target_multi_hot=batch["target_multi_hot"],
            target_dt_unique=batch["target_dt_unique"],
            readout_mask=readout_out.readout_mask,
            return_components=True,
        )
    else:
        raise ValueError(f"Unknown target_mode: {args.target_mode}")

    loss, loss_parts = loss_out

    # Check for NaN/Inf
    if not torch.isfinite(loss):
        raise RuntimeError(
            f"Loss is not finite: {loss.item()}. "
            f"batch_idx info: event_seq shape {event_seq.shape}, "
            f"logits shape {logits.shape}, logits range [{logits.min():.4f}, {logits.max():.4f}]"
        )

    return loss, loss_parts


def run_one_epoch(
    logger: logging.Logger,
    args: argparse.Namespace,
    model: DeepHealth,
    readout: nn.Module,
    criterion: nn.Module,
    train_loader: DataLoader,
    optimizer: AdamW,
    device: torch.device,
    is_train: bool = True,
) -> float:
    """
    Run one epoch of training or validation.

    Parameters
    ----------
    logger : logging.Logger
    args : argparse.Namespace
    model : DeepHealth
    readout : nn.Module
    criterion : nn.Module
    train_loader : DataLoader
    optimizer : AdamW
        Unused if is_train=False.
    device : torch.device
    is_train : bool
        If True, perform training updates. Otherwise just evaluate.

    Returns
    -------
    avg_loss : float
        Average loss over the epoch.
    """
    if is_train:
        model.train()
    else:
        model.eval()

    total_loss = 0.0
    n_batches = 0
    n_skipped = 0
    component_sums: Dict[str, float] = {}

    epoch_desc = "train" if is_train else "val"
    progress = tqdm(
        train_loader,
        desc=epoch_desc,
        total=len(train_loader),
        leave=False,
        dynamic_ncols=True,
    )

    for batch_idx, batch in enumerate(progress):
        try:
            loss, loss_parts = compute_loss(
                args=args,
                model=model,
                readout=readout,
                criterion=criterion,
                batch=batch,
                device=device,
            )

            if is_train:
                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                if args.grad_clip > 0:
                    clip_grad_norm_(model.parameters(), args.grad_clip)
                optimizer.step()

            total_loss += loss.item()
            n_batches += 1

            for name, value in loss_parts.items():
                component_sums[name] = component_sums.get(name, 0.0) + float(
                    value.detach().item()
                )

            current_loss = loss.item()
            avg_loss = total_loss / max(1, n_batches)
            postfix = {
                "loss": f"{current_loss:.4f}",
                "avg": f"{avg_loss:.4f}",
                "skipped": n_skipped,
            }
            for name, sum_value in component_sums.items():
                postfix[name] = f"{sum_value / max(1, n_batches):.4f}"
            progress.set_postfix(postfix)

        except RuntimeError as e:
            if "Loss is not finite" in str(e):
                n_skipped += 1
                logger.warning(f"Batch {batch_idx} skipped: {str(e)[:100]}")
                progress.set_postfix(
                    loss="nan",
                    avg=f"{total_loss / max(1, n_batches):.4f}" if n_batches > 0 else "inf",
                    skipped=n_skipped,
                )
                continue
            else:
                raise

    if n_skipped > 0:
        logger.info(f"Skipped {n_skipped} batches due to non-finite loss")

    if component_sums and n_batches > 0:
        component_summary = ", ".join(
            f"{name}={sum_value / n_batches:.4f}"
            for name, sum_value in sorted(component_sums.items())
        )
        logger.info(f"Epoch loss breakdown: {component_summary}")

    avg_loss = total_loss / max(1, n_batches) if n_batches > 0 else float("inf")
    return avg_loss


def evaluate(
    logger: logging.Logger,
    args: argparse.Namespace,
    model: DeepHealth,
    readout: nn.Module,
    criterion: nn.Module,
    val_loader: DataLoader,
    device: torch.device,
) -> float:
    """
    Evaluate model on validation set.

    Returns
    -------
    val_loss : float
    """
    with torch.no_grad():
        val_loss = run_one_epoch(
            logger=logger,
            args=args,
            model=model,
            readout=readout,
            criterion=criterion,
            train_loader=val_loader,
            optimizer=None,
            device=device,
            is_train=False,
        )
    return val_loss


def save_checkpoint(
    model: DeepHealth,
    checkpoint_path: Path,
) -> None:
    """Save model state_dict."""
    torch.save(model.state_dict(), checkpoint_path)


def save_config(
    args: argparse.Namespace,
    config_path: Path,
    extra: Dict[str, Any] | None = None,
) -> None:
    """Save training config as JSON."""
    config_dict: Dict[str, Any] = {}
    for key, value in vars(args).items():
        if isinstance(value, tuple):
            config_dict[key] = list(value)
        elif isinstance(value, list):
            config_dict[key] = value
        elif isinstance(value, (int, float, str, bool, type(None))):
            config_dict[key] = value
        else:
            config_dict[key] = str(value)
    if extra:
        config_dict.update(extra)
    with open(config_path, "w") as f:
        json.dump(config_dict, f, indent=2)


def build_run_metadata(
    args: argparse.Namespace,
    dataset: HealthDataset,
    train_subset: Subset,
    val_subset: Subset,
    test_subset: Subset,
    run_name: str,
) -> Dict[str, Any]:
    """Collect resolved training facts needed to rebuild the model for evaluation."""
    return {
        "run_name": run_name,
        "dataset_class": "NextStepHealthDataset",
        "collate_fn": "next_step_collate_fn",
        "model_class": "DeepHealth",
        "model_target_mode": "next_token",
        "dist_mode": "exponential",
        "extra_info_types_file": (
            Path(args.extra_info_types_file).name
            if args.extra_info_types_file is not None
            else None
        ),
        "extra_info_types": dataset.extra_info_types,
        "dataset_metadata": {
            "vocab_size": int(dataset.vocab_size),
            "n_types": int(dataset.n_types),
            "n_cont_types": int(dataset.n_cont_types),
            "n_categories": int(dataset.n_categories),
            "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
            "extra_info_types": [int(x) for x in dataset.extra_info_types],
        },
        "split_sizes": {
            "train": int(len(train_subset)),
            "val": int(len(val_subset)),
            "test": int(len(test_subset)),
        },
        "resolved_readout_name": args.readout_name,
        "resolved_loss_name": args.loss_name,
    }


def normalize_training_config(args: argparse.Namespace) -> None:
    """Fill in and validate training options that depend on other flags."""
    if args.target_mode not in {"delphi2m", "uts"}:
        raise ValueError(f"Unknown target_mode: {args.target_mode}")

    # gap_5y is always enabled, so preserve NO_EVENT target behavior.
    args.ignore_no_event_in_delphi2m = False
    if args.target_mode == "uts":
        args.include_no_event_in_uts_target = True


def normalize_loss_and_distribution_config(args: argparse.Namespace) -> None:
    """Validate and resolve loss/distribution options after auto-selection."""
    if args.loss_name not in {"delphi2m", "uts"}:
        raise ValueError(
            "Unknown loss_name. Supported values: delphi2m, uts."
        )

    if args.loss_name == "delphi2m" and args.target_mode != "delphi2m":
        raise ValueError(
            "loss_name=delphi2m requires target_mode=delphi2m."
        )

    if args.loss_name == "uts" and args.target_mode != "uts":
        raise ValueError(
            "loss_name=uts requires target_mode=uts."
        )


# ---------------------------------------------------------------------------
# Main Training Loop
# ---------------------------------------------------------------------------

def main():
    """Main training function."""
    parser = argparse.ArgumentParser(description="DeepHealth Training")

    # ---- Data & Output ----
    parser.add_argument("--data_prefix", type=str, default="ukb",
                        help="Prefix for data files")
    parser.add_argument("--labels_file", type=str, default="labels.csv",
                        help="Path to labels file")
    parser.add_argument("--seed", type=int, default=42,
                        help="Random seed")

    # ---- Dataset ----
    parser.add_argument("--no_event_interval_years", type=float, default=5.0,
                        help="Interval in years for no-event insertion")
    parser.add_argument("--include_no_event_in_uts_target", action="store_true",
                        help="Include NO_EVENT in UTS target multi-hot")

    # ---- Split ----
    parser.add_argument("--train_ratio", type=float, default=0.7,
                        help="Training set ratio")
    parser.add_argument("--val_ratio", type=float, default=0.15,
                        help="Validation set ratio")
    parser.add_argument("--test_ratio", type=float, default=0.15,
                        help="Test set ratio")

    # ---- Model ----
    parser.add_argument("--n_embd", type=int, default=120,
                        help="Embedding dimension")
    parser.add_argument("--n_head", type=int, default=10,
                        help="Number of attention heads")
    parser.add_argument("--n_hist_layer", type=int, default=12,
                        help="Number of history encoder layers")
    parser.add_argument("--n_tab_layer", type=int, default=4,
                        help="Number of self-attention layers for other-token encoder")
    parser.add_argument("--n_bins", type=int, default=16,
                        help="Number of bins for continuous other-token values")
    parser.add_argument("--time_mode", type=str, default="relative",
                        choices=["relative", "absolute"],
                        help="Time encoding mode for disease history")
    parser.add_argument("--dropout", type=float, default=0.0,
                        help="Dropout rate")
    parser.add_argument("--extra_info_types_file", type=str, default=None,
                        help="Optional file containing other-information type ids to include")

    # ---- Training Protocol ----
    parser.add_argument("--target_mode", type=str, default="uts",
                        choices=["delphi2m", "uts"],
                        help="Target supervision mode")
    parser.add_argument("--readout_name", type=str, default=None,
                        help="Readout name (auto-selected if None)")
    parser.add_argument("--readout_reduce", type=str, default="mean",
                        choices=["mean", "sum"],
                        help="Readout reduction for SameTimeGroupEndReadout")

    # ---- Loss ----
    parser.add_argument("--loss_name", type=str, default=None,
                        help="Loss name (auto-selected if None): delphi2m, uts")
    parser.add_argument("--t_min", type=float, default=0.0027378507871321013,
                        help="Minimum time for loss (1/365.25)")
    parser.add_argument("--max_exp_input", type=float, default=60.0,
                        help="Max exponent input for loss")
    parser.add_argument("--ce_weight", type=float, default=1.0,
                        help="Cross-entropy weight in delphi2m loss")
    parser.add_argument("--time_weight", type=float, default=1.0,
                        help="Time loss weight in delphi2m loss")
    parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true",
                        help="Ignore NO_EVENT in delphi2m loss")

    # ---- Optimization ----
    parser.add_argument("--batch_size", type=int, default=128,
                        help="Batch size")
    parser.add_argument("--base_lr", type=float, default=3e-4,
                        help="Base learning rate")
    parser.add_argument("--weight_decay", type=float, default=0.1,
                        help="Weight decay (L2 regularization)")
    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99),
                        help="AdamW betas")
    parser.add_argument("--grad_clip", type=float, default=1.0,
                        help="Gradient clipping norm")
    parser.add_argument("--max_epochs", type=int, default=200,
                        help="Maximum number of epochs")
    parser.add_argument("--warmup_epochs", type=int, default=10,
                        help="Number of warmup epochs")
    parser.add_argument("--patience", type=int, default=15,
                        help="Early stopping patience")
    parser.add_argument("--min_lr_ratio", type=float, default=0.1,
                        help="Minimum LR as ratio of adaptive_lr")
    parser.add_argument("--num_workers", type=int, default=4,
                        help="Number of DataLoader workers")
    parser.add_argument("--device", type=str, default="cuda",
                        help="Device to use for training: cpu, cuda, or cuda:<index>")

    args = parser.parse_args()
    args.extra_info_types = (
        load_extra_info_types_file(args.extra_info_types_file)
        if args.extra_info_types_file is not None
        else None
    )

    # ---- Setup ----
    set_seed(args.seed)
    device = resolve_device(args.device)
    configure_torch_for_training(device)

    normalize_training_config(args)
    runs_root = Path("runs")
    while True:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        run_name = (
            f"{args.time_mode}_exponential_{args.target_mode}_"
            f"other_tokens_gap_5y_{timestamp}"
        )
        run_dir = runs_root / run_name
        if not run_dir.exists():
            break
        time.sleep(1)

    run_dir.mkdir(parents=True, exist_ok=False)

    logger = setup_logging(run_dir)
    logger.info(f"Starting training run: {run_name}")
    logger.info(f"Device: {device}")
    logger.info(f"extra_info_types: {args.extra_info_types or 'all'}")
    logger.info(
        f"Resolved no_event config: ignore_no_event_in_delphi2m={args.ignore_no_event_in_delphi2m}, "
        f"include_no_event_in_uts_target={args.include_no_event_in_uts_target}"
    )

    # ---- Validate Arguments ----
    total_ratio = args.train_ratio + args.val_ratio + args.test_ratio
    if not np.isclose(total_ratio, 1.0, atol=1e-6):
        raise ValueError(
            f"train_ratio + val_ratio + test_ratio must equal 1.0, got {total_ratio}"
        )

    # Auto-select readout if not specified
    if args.readout_name is None:
        args.readout_name = (
            "token" if args.target_mode == "delphi2m"
            else "same_time_group_end"
        )

    # Auto-select loss if not specified
    if args.loss_name is None:
        args.loss_name = (
            "delphi2m" if args.target_mode == "delphi2m"
            else "uts"
        )

    normalize_loss_and_distribution_config(args)

    logger.info(f"Auto-selected readout: {args.readout_name}")
    logger.info(f"Auto-selected loss: {args.loss_name}")

    # ---- Load Dataset ----
    logger.info("Loading dataset...")
    dataset = HealthDataset(
        data_prefix=args.data_prefix,
        labels_file=args.labels_file,
        no_event_interval_years=args.no_event_interval_years,
        include_no_event_in_uts_target=args.include_no_event_in_uts_target,
        extra_info_types=args.extra_info_types,
    )
    logger.info(
        f"Dataset loaded: {len(dataset)} samples, vocab_size={dataset.vocab_size}")

    # ---- Split Dataset ----
    train_subset, val_subset, test_subset = split_dataset(
        dataset=dataset,
        train_ratio=args.train_ratio,
        val_ratio=args.val_ratio,
        test_ratio=args.test_ratio,
        seed=args.seed,
    )
    logger.info(
        f"Dataset split: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
    )

    # ---- Build DataLoaders ----
    train_loader = DataLoader(
        train_subset,
        batch_size=args.batch_size,
        sampler=RandomSampler(
            train_subset, generator=torch.Generator().manual_seed(args.seed)),
        collate_fn=collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )
    val_loader = DataLoader(
        val_subset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )
    test_loader = DataLoader(
        test_subset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )

    # ---- Build Model ----
    logger.info("Building model...")
    model = build_model(args, dataset)
    model.to(device)
    n_params = sum(p.numel() for p in model.parameters())
    logger.info(f"Model built: {n_params:,} parameters")

    # ---- Build Optimizer ----
    optimizer = build_optimizer(args, model)
    logger.info(
        f"Optimizer: AdamW, base_lr={args.base_lr}, weight_decay={args.weight_decay}")

    # ---- Compute Adaptive LR ----
    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
    logger.info(
        f"Adaptive LR: {adaptive_lr:.6f} (base_lr * sqrt(batch_size/128))")

    # ---- Build Readout ----
    if args.readout_name == "token":
        readout = build_readout("token")
    elif args.readout_name == "same_time_group_end":
        readout = build_readout("same_time_group_end",
                                reduce=args.readout_reduce)
    elif args.readout_name == "last_valid":
        readout = build_readout("last_valid")
    else:
        raise ValueError(f"Unknown readout: {args.readout_name}")
    logger.info(f"Readout: {args.readout_name}")

    # ---- Build Loss ----
    if args.loss_name == "delphi2m":
        ignored_tokens = {PAD_IDX, CHECKUP_IDX}
        if args.ignore_no_event_in_delphi2m:
            ignored_tokens.add(NO_EVENT_IDX)
        criterion = build_loss(
            "delphi2m",
            ignored_tokens=ignored_tokens,
            t_min=args.t_min,
            max_exp_input=args.max_exp_input,
            ce_weight=args.ce_weight,
            time_weight=args.time_weight,
        )
        logger.info(f"Loss: delphi2m, ignored_tokens={ignored_tokens}")
    elif args.loss_name == "uts":
        ignored_idx = {PAD_IDX, CHECKUP_IDX}
        criterion = build_loss(
            "uts",
            ignored_idx=ignored_idx,
            t_min=args.t_min,
            max_exp_input=args.max_exp_input,
        )
        logger.info(f"Loss: uts, ignored_idx={ignored_idx}")
    else:
        raise ValueError(f"Unknown loss: {args.loss_name}")

    # ---- Save Config ----
    save_config(
        args,
        run_dir / "train_config.json",
        extra=build_run_metadata(
            args=args,
            dataset=dataset,
            train_subset=train_subset,
            val_subset=val_subset,
            test_subset=test_subset,
            run_name=run_name,
        ),
    )
    logger.info(f"Config saved to {run_dir / 'train_config.json'}")

    # ---- Training Loop ----
    logger.info("Starting training...")
    best_val_loss = float("inf")
    patience_counter = 0
    metrics = []

    best_model_path = run_dir / "best_model.pt"
    history_path = run_dir / "history.json"

    start_time = time.time()

    for epoch in range(args.max_epochs):
        lr = get_lr(epoch, args, adaptive_lr, args.max_epochs)
        set_optimizer_lr(optimizer, lr)

        # Train
        train_loss = run_one_epoch(
            logger=logger,
            args=args,
            model=model,
            readout=readout,
            criterion=criterion,
            train_loader=train_loader,
            optimizer=optimizer,
            device=device,
            is_train=True,
        )

        # Validate
        val_loss = evaluate(
            logger=logger,
            args=args,
            model=model,
            readout=readout,
            criterion=criterion,
            val_loader=val_loader,
            device=device,
        )

        # Early stopping
        is_best = False
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            is_best = True
            save_checkpoint(model, best_model_path)
        else:
            patience_counter += 1

        elapsed = time.time() - start_time
        logger.info(
            f"Epoch {epoch+1}/{args.max_epochs} | lr={lr:.6f} | "
            f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
            f"best_val_loss={best_val_loss:.6f} | patience={patience_counter}/{args.patience} | "
            f"elapsed={elapsed:.1f}s"
        )

        metrics.append({
            "epoch": epoch + 1,
            "lr": lr,
            "train_loss": train_loss,
            "val_loss": val_loss,
            "best_val_loss": best_val_loss,
            "is_best": int(is_best),
        })

        if patience_counter >= args.patience:
            logger.info(f"Early stopping triggered at epoch {epoch+1}")
            break

    # ---- Save Training History ----
    with open(history_path, "w") as f:
        json.dump(metrics, f, indent=2)
    logger.info(f"History saved to {history_path}")

    # ---- Test on Best Model ----
    logger.info("Evaluating best model on test set...")
    best_state_dict = torch.load(best_model_path, map_location=device)
    model.load_state_dict(best_state_dict)

    test_loss = evaluate(
        logger=logger,
        args=args,
        model=model,
        readout=readout,
        criterion=criterion,
        val_loader=test_loader,
        device=device,
    )
    logger.info(f"Test loss: {test_loss:.6f}")

    total_time = time.time() - start_time
    logger.info(f"Training completed in {total_time:.1f}s")
    logger.info(f"Best checkpoint: {best_model_path}")


if __name__ == "__main__":
    main()