DeepHealth/train_all_future.py

"""
Train DeepHealth with query-conditioned all-future supervision.

Training samples are patient-level. For each patient and each __getitem__ call,
AllFutureHealthDataset randomly samples a query time t_query, uses events at or
before t_query as history, and uses events after t_query as the future target set.

Validation/test samples are deterministic query points built from future event
times, then split by patient.
"""
from __future__ import annotations

import argparse
import json
import logging
import math
import time
from pathlib import Path
from typing import Any, Dict

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

from dataset import (
    DISEASE_HISTORY_MODES,
    DISEASE_HISTORY_MODE_TIMED,
    AllFutureHealthDataset,
    all_future_collate_fn,
)
from losses import build_loss
from model_architectures import (
    DEFAULT_MODEL_ARCHITECTURE,
    SUPPORTED_MODEL_ARCHITECTURES,
)
from models import DeepHealth
from targets import PAD_IDX, RESERVED_IDX
from train_util import (
    ContinuousRobustScalerStats,
    configure_torch_for_training,
    create_unique_run_dir,
    format_extra_info_types,
    fit_continuous_robust_scaler,
    get_lr,
    get_model_parameter_counts,
    load_extra_info_types_file,
    move_batch_to_device,
    resolve_device,
    save_checkpoint,
    save_config,
    set_optimizer_lr,
    set_seed,
    setup_logging,
    split_all_future_datasets,
    split_all_future_datasets_by_eid_files,
)


MODEL_INPUT_KEYS = (
    "event_seq",
    "time_seq",
    "sex",
    "padding_mask",
    "t_query",
    "other_type",
    "other_value",
    "other_value_kind",
    "other_time",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train DeepHealth with all-future supervision")

    parser.add_argument("--data_prefix", type=str, default="ukb")
    parser.add_argument("--labels_file", type=str, default="labels.csv")
    parser.add_argument("--runs_root", type=str, default="runs")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--extra_info_types_file", type=str, default=None)
    parser.add_argument(
        "--disease_history_mode",
        type=str,
        default=DISEASE_HISTORY_MODE_TIMED,
        choices=DISEASE_HISTORY_MODES,
        help=(
            "timed=real disease times; ordered=chronological disease order with "
            "ordinal positions; set=unordered disease set with no disease time"
        ),
    )

    parser.add_argument("--train_ratio", type=float, default=0.7)
    parser.add_argument("--val_ratio", type=float, default=0.15)
    parser.add_argument("--test_ratio", type=float, default=0.15)
    parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
    parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
    parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
    parser.add_argument("--min_history_events", type=int, default=1)
    parser.add_argument("--min_future_events", type=int, default=1)
    parser.add_argument("--validation_query_seed", type=int, default=None)

    parser.add_argument("--n_embd", type=int, default=120)
    parser.add_argument("--n_head", type=int, default=10)
    parser.add_argument("--n_layer", type=int, default=12)
    parser.add_argument("--n_bins", type=int, default=16)
    parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                        choices=["mean", "sum"])
    parser.add_argument("--time_mode", type=str, default="relative",
                        choices=["relative", "absolute"])
    parser.add_argument("--dist_mode", type=str, default="exponential",
                        choices=["exponential", "weibull"])
    parser.add_argument("--dropout", type=float, default=0.0)
    parser.add_argument(
        "--model_architecture",
        type=str,
        default=DEFAULT_MODEL_ARCHITECTURE,
        choices=SUPPORTED_MODEL_ARCHITECTURES,
    )

    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--base_lr", type=float, default=3e-4)
    parser.add_argument("--weight_decay", type=float, default=0.1)
    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--max_epochs", type=int, default=200)
    parser.add_argument("--warmup_epochs", type=int, default=10)
    parser.add_argument("--patience", type=int, default=15)
    parser.add_argument("--min_lr_ratio", type=float, default=0.1)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--progress_interval", type=int, default=20)

    args = parser.parse_args()
    if args.min_history_events < 1:
        raise ValueError("min_history_events must be >= 1")
    if args.min_future_events < 1:
        raise ValueError("min_future_events must be >= 1")
    use_eid_split = all(
        getattr(args, name)
        for name in ("train_eid_file", "val_eid_file", "test_eid_file")
    )
    if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
    if args.validation_query_seed is None:
        args.validation_query_seed = int(args.seed)
    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
    )
    if args.disease_history_mode != DISEASE_HISTORY_MODE_TIMED:
        expected = {
            "model_architecture": "traj_mixer_v5",
            "time_mode": "relative",
            "dist_mode": "weibull",
        }
        mismatches = [
            f"{name}={getattr(args, name)!r} (expected {value!r})"
            for name, value in expected.items()
            if getattr(args, name) != value
        ]
        if args.extra_info_types != []:
            mismatches.append(
                "extra_info_types must be [] via extra_info_types_none.txt"
            )
        if mismatches:
            raise ValueError(
                f"disease_history_mode={args.disease_history_mode!r} is reserved "
                "for the no-extra TrajMixer + all_future + relative + Weibull "
                "ablation; " + "; ".join(mismatches)
            )
    return args


def build_model(
    args: argparse.Namespace,
    dataset: AllFutureHealthDataset,
    scaler_stats: ContinuousRobustScalerStats,
) -> DeepHealth:
    if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids:
        raise ValueError(
            "RobustScale statistics are not aligned with dataset.cont_type_ids"
        )
    center = scaler_stats.center if dataset.n_cont_types > 0 else None
    scale = scaler_stats.scale if dataset.n_cont_types > 0 else None
    return DeepHealth(
        vocab_size=dataset.vocab_size,
        n_embd=args.n_embd,
        n_head=args.n_head,
        n_layer=args.n_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,
        continuous_value_center=center,
        continuous_value_scale=scale,
        extra_pool_reduce=args.extra_pool_reduce,
        target_mode="all_future",
        time_mode=args.time_mode,
        dist_mode=args.dist_mode,
        dropout=args.dropout,
        model_architecture=args.model_architecture,
    )


def build_criterion(args: argparse.Namespace):
    ignored_idx = {PAD_IDX, RESERVED_IDX}
    if args.dist_mode == "exponential":
        return build_loss("exponential", ignored_idx=ignored_idx)
    if args.dist_mode == "weibull":
        return build_loss("weibull", ignored_idx=ignored_idx)
    raise ValueError(f"Unknown dist_mode: {args.dist_mode}")


def compute_all_future_loss(
    args: argparse.Namespace,
    model: DeepHealth,
    criterion,
    batch: Dict[str, torch.Tensor],
    device: torch.device,
) -> torch.Tensor:
    required_keys = set(MODEL_INPUT_KEYS)
    required_keys.update(("future_targets", "exposure"))
    if args.dist_mode == "weibull":
        required_keys.add("future_dt")
    batch = move_batch_to_device(
        {key: batch[key] for key in required_keys},
        device,
    )

    hidden = model(
        event_seq=batch["event_seq"],
        time_seq=batch["time_seq"],
        sex=batch["sex"],
        padding_mask=batch["padding_mask"],
        t_query=batch["t_query"],
        other_type=batch["other_type"],
        other_value=batch["other_value"],
        other_value_kind=batch["other_value_kind"],
        other_time=batch["other_time"],
    )
    logits = model.calc_risk(hidden)

    if args.dist_mode == "exponential":
        loss = criterion(
            logits=logits,
            targets=batch["future_targets"],
            exposure=batch["exposure"],
        )
    elif args.dist_mode == "weibull":
        loss = criterion(
            logits=logits,
            weibull_rho=model.calc_weibull_rho(hidden),
            targets=batch["future_targets"],
            dt=batch["future_dt"],
            exposure=batch["exposure"],
        )
    else:
        raise ValueError(f"Unknown dist_mode: {args.dist_mode}")

    if not torch.isfinite(loss):
        raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
    return loss


def run_epoch(
    logger: logging.Logger,
    args: argparse.Namespace,
    model: DeepHealth,
    criterion,
    loader: DataLoader,
    optimizer: AdamW | None,
    device: torch.device,
    is_train: bool,
) -> float:
    model.train(is_train)
    total = torch.zeros((), device=device)
    n_batches = 0
    skipped = 0
    desc = "train" if is_train else "val"
    progress_interval = max(1, int(args.progress_interval))

    progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
    for batch_idx, batch in enumerate(progress):
        try:
            loss = compute_all_future_loss(args, model, criterion, batch, device)
            if is_train:
                if optimizer is None:
                    raise ValueError("optimizer is required for training")
                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 = total + loss.detach()
            n_batches += 1
            if (batch_idx + 1) % progress_interval == 0:
                avg = total / max(1, n_batches)
                progress.set_postfix(
                    loss=f"{float(loss.detach().cpu()):.4f}",
                    avg=f"{float(avg.detach().cpu()):.4f}",
                    skipped=skipped,
                )
        except RuntimeError as exc:
            if "Loss is not finite" not in str(exc):
                raise
            skipped += 1
            logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}")

    if skipped:
        logger.info(f"Skipped {skipped} batches due to non-finite loss")
    return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")


def build_metadata(
    args: argparse.Namespace,
    dataset: AllFutureHealthDataset,
    run_name: str,
    train_subset,
    val_subset,
    test_subset,
    scaler_stats: ContinuousRobustScalerStats,
) -> Dict[str, Any]:
    scaler_metadata = scaler_stats.as_metadata()
    return {
        "run_name": run_name,
        "dataset_class": "AllFutureHealthDataset",
        "collate_fn": "all_future_collate_fn",
        "model_class": "DeepHealth",
        "model_architecture": args.model_architecture,
        "model_target_mode": "all_future",
        "target_mode": "all_future",
        "event_stream_version": "disease_death_only_v1",
        "uses_assessment_event_token": False,
        "dist_mode": args.dist_mode,
        "disease_history_mode": args.disease_history_mode,
        "all_future_min_history_events": int(args.min_history_events),
        "all_future_min_future_events": int(args.min_future_events),
        "all_future_validation_query_seed": int(args.validation_query_seed),
        "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": [int(x) for x in dataset.extra_info_types],
        "continuous_value_scaling": "robust",
        "continuous_value_scaler": scaler_metadata,
        "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],
            "event_stream_version": "disease_death_only_v1",
            "uses_assessment_event_token": False,
        },
        "split_sizes": {
            "train": int(len(train_subset)),
            "val": int(len(val_subset)),
            "test": int(len(test_subset)),
        },
        "resolved_readout_name": "none",
        "resolved_loss_name": args.dist_mode,
    }


def main() -> None:
    args = parse_args()
    set_seed(args.seed)
    device = resolve_device(args.device)
    configure_torch_for_training(device)

    run_dir, run_name = create_unique_run_dir(
        lambda timestamp: (
            (
                ""
                if args.disease_history_mode == DISEASE_HISTORY_MODE_TIMED
                else f"{args.disease_history_mode}_"
            )
            + f"{args.time_mode}_{args.dist_mode}_"
            f"all_future_pure_disease_{timestamp}"
        ),
        runs_root=Path(args.runs_root) / args.model_architecture,
    )
    logger = setup_logging(run_dir)

    logger.info(f"Starting all-future training run: {run_name}")
    logger.info(f"Device: {device}")
    logger.info(f"Model architecture: {args.model_architecture}")
    logger.info(f"Disease history mode: {args.disease_history_mode}")
    logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
    logger.info("Continuous value scaling: RobustScale (required)")

    logger.info("Loading all-future datasets...")
    train_dataset = AllFutureHealthDataset(
        data_prefix=args.data_prefix,
        labels_file=args.labels_file,
        split="train",
        min_history_events=args.min_history_events,
        min_future_events=args.min_future_events,
        validation_query_seed=args.validation_query_seed,
        extra_info_types=args.extra_info_types,
        disease_history_mode=args.disease_history_mode,
    )
    val_dataset = AllFutureHealthDataset(
        data_prefix=args.data_prefix,
        labels_file=args.labels_file,
        split="valid",
        min_history_events=args.min_history_events,
        min_future_events=args.min_future_events,
        validation_query_seed=args.validation_query_seed,
        extra_info_types=args.extra_info_types,
        disease_history_mode=args.disease_history_mode,
    )
    test_dataset = AllFutureHealthDataset(
        data_prefix=args.data_prefix,
        labels_file=args.labels_file,
        split="test",
        min_history_events=args.min_history_events,
        min_future_events=args.min_future_events,
        validation_query_seed=args.validation_query_seed,
        extra_info_types=args.extra_info_types,
        disease_history_mode=args.disease_history_mode,
    )
    if args.train_eid_file and args.val_eid_file and args.test_eid_file:
        train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files(
            train_dataset=train_dataset,
            val_dataset=val_dataset,
            test_dataset=test_dataset,
            train_eid_file=args.train_eid_file,
            val_eid_file=args.val_eid_file,
            test_eid_file=args.test_eid_file,
        )
        logger.info(
            "Using eid split files: "
            f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
        )
    else:
        train_subset, val_subset, test_subset = split_all_future_datasets(
            train_dataset=train_dataset,
            val_dataset=val_dataset,
            test_dataset=test_dataset,
            train_ratio=args.train_ratio,
            val_ratio=args.val_ratio,
            test_ratio=args.test_ratio,
            seed=args.seed,
        )
        logger.info(
            f"Using random ratio split: train={args.train_ratio}, "
            f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
        )
    logger.info(
        f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
    )

    if train_dataset.n_cont_types > 0:
        logger.info(
            "Fitting continuous RobustScaler on the complete training subset: "
            f"patients={len(train_subset):,}, features={train_dataset.n_cont_types}"
        )
    scaler_stats = fit_continuous_robust_scaler(
        train_dataset,
        train_subset,
    )
    if train_dataset.n_cont_types > 0:
        logger.info(
            "Continuous RobustScaler fitted: "
            f"observations={int(scaler_stats.observation_count.sum()):,}, "
            f"min_per_feature={int(scaler_stats.observation_count.min()):,}, "
            f"max_per_feature={int(scaler_stats.observation_count.max()):,}"
        )

    train_loader = DataLoader(
        train_subset,
        batch_size=args.batch_size,
        sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
        collate_fn=all_future_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=all_future_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=all_future_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,
    )

    model = build_model(args, train_dataset, scaler_stats=scaler_stats).to(device)
    parameter_counts = get_model_parameter_counts(model)
    logger.info(
        "Model parameters: "
        f"total={parameter_counts['model_parameter_count']:,}, "
        f"trainable={parameter_counts['trainable_parameter_count']:,}"
    )
    optimizer = AdamW(
        model.parameters(),
        lr=args.base_lr,
        betas=tuple(args.betas),
        weight_decay=args.weight_decay,
    )
    criterion = build_criterion(args)
    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)

    train_metadata = build_metadata(
        args,
        train_dataset,
        run_name,
        train_subset,
        val_subset,
        test_subset,
        scaler_stats,
    )
    train_metadata.update(parameter_counts)
    save_config(
        args,
        run_dir / "train_config.json",
        extra=train_metadata,
    )

    best_val = float("inf")
    patience = 0
    history = []
    best_model_path = run_dir / "best_model.pt"
    start = time.time()

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

        train_loss = run_epoch(logger, args, model, criterion, train_loader, optimizer, device, True)
        with torch.no_grad():
            val_loss = run_epoch(logger, args, model, criterion, val_loader, None, device, False)

        is_best = val_loss < best_val
        if is_best:
            best_val = val_loss
            patience = 0
            save_checkpoint(model, best_model_path)
        else:
            patience += 1

        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:.6f} | patience={patience}/{args.patience} | "
            f"elapsed={time.time() - start:.1f}s"
        )
        history.append({
            "epoch": epoch + 1,
            "lr": lr,
            "train_loss": train_loss,
            "val_loss": val_loss,
            "best_val_loss": best_val,
            "is_best": int(is_best),
        })
        if patience >= args.patience:
            logger.info(f"Early stopping triggered at epoch {epoch + 1}")
            break

    with (run_dir / "history.json").open("w", encoding="utf-8") as f:
        json.dump(history, f, indent=2)

    logger.info("Evaluating best model on all-future test queries...")
    model.load_state_dict(torch.load(best_model_path, map_location=device))
    with torch.no_grad():
        test_loss = run_epoch(logger, args, model, criterion, test_loader, None, device, False)
    logger.info(f"Test loss: {test_loss:.6f}")
    logger.info(f"Best checkpoint: {best_model_path}")


if __name__ == "__main__":
    main()