DeepHealth/train_next_step.py

"""
Train DeepHealth with next-token / next-time-point supervision.

The next-step dataset uses observed event histories, including CHECKUP state
tokens, plus optional gap <NO_EVENT> imputation. UTS training reads out only
same-time group ends.
"""
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 HealthDataset, collate_fn
from losses import build_loss
from models import DeepHealth, DeepHealthOutput
from readouts import build_readout
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
from train_util import (
    configure_torch_for_training,
    create_unique_run_dir,
    load_extra_info_types_file,
    resolve_device,
    save_checkpoint,
    save_config,
    set_optimizer_lr,
    set_seed,
    setup_logging,
    split_dataset,
    split_dataset_by_eid_files,
)


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train DeepHealth with next-token/point supervision")

    parser.add_argument("--data_prefix", type=str, default="ukb")
    parser.add_argument("--labels_file", type=str, default="labels.csv")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--extra_info_types_file", type=str, default=None)
    parser.add_argument("--no_event_interval_years", type=float, default=5.0)
    parser.add_argument("--include_no_event_in_uts_target", action="store_true")

    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("--n_embd", type=int, default=120)
    parser.add_argument("--n_head", type=int, default=10)
    parser.add_argument("--n_hist_layer", type=int, default=12)
    parser.add_argument("--n_tab_layer", type=int, default=4)
    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("--dropout", type=float, default=0.0)

    parser.add_argument("--target_mode", type=str, default="uts",
                        choices=["delphi2m", "uts"])
    parser.add_argument("--readout_name", type=str, default=None,
                        choices=["token", "same_time_group_end", "last_valid"])
    parser.add_argument("--readout_reduce", type=str, default="mean",
                        choices=["mean", "sum"])
    parser.add_argument("--t_min", type=float, default=0.0027378507871321013)
    parser.add_argument("--max_exp_input", type=float, default=60.0)
    parser.add_argument("--ce_weight", type=float, default=1.0)
    parser.add_argument("--time_weight", type=float, default=1.0)
    parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true")

    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()
    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.target_mode == "uts":
        args.readout_name = args.readout_name or "same_time_group_end"
        args.include_no_event_in_uts_target = True
    else:
        args.readout_name = args.readout_name or "token"
    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
    )
    return args


def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
    if epoch < args.warmup_epochs:
        return adaptive_lr * (epoch + 1) / args.warmup_epochs
    progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
    cosine = 0.5 * (1 + math.cos(math.pi * progress))
    return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))


def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
    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 build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
    return 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,
        extra_pool_reduce=args.extra_pool_reduce,
        target_mode="next_token",
        time_mode=args.time_mode,
        dist_mode="exponential",
        dropout=args.dropout,
    )


def build_next_step_readout(args: argparse.Namespace):
    if args.readout_name == "same_time_group_end":
        return build_readout("same_time_group_end", reduce=args.readout_reduce)
    return build_readout(args.readout_name)


def build_next_step_loss(args: argparse.Namespace):
    if args.target_mode == "delphi2m":
        ignored_tokens = {PAD_IDX, CHECKUP_IDX}
        if args.ignore_no_event_in_delphi2m:
            ignored_tokens.add(NO_EVENT_IDX)
        return 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,
        )
    return build_loss(
        "uts",
        ignored_idx={PAD_IDX, CHECKUP_IDX},
        t_min=args.t_min,
        max_exp_input=args.max_exp_input,
    )


def build_augmented_next_step_targets(
    batch_cpu: Dict[str, torch.Tensor],
    model_out: DeepHealthOutput,
    include_uts_targets: bool,
) -> Dict[str, torch.Tensor]:
    hidden_len = model_out.hidden.size(1)
    event_len = int(model_out.event_len)
    extra_len = hidden_len - event_len
    device = model_out.hidden.device
    non_blocking = device.type == "cuda"
    if extra_len <= 0:
        targets = {
            "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
            "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
            "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking),
        }
        if include_uts_targets:
            targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
                device, non_blocking=non_blocking
            )
            targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
                device, non_blocking=non_blocking
            )
        return targets

    bsz = batch_cpu["target_event_seq"].size(0)
    vocab_size = (
        batch_cpu["target_multi_hot"].size(2)
        if include_uts_targets
        else None
    )
    other_valid = batch_cpu["other_type"] > 0
    extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
    extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
    for b in range(bsz):
        unique_time = torch.unique(batch_cpu["other_time"][b, other_valid[b]], sorted=True)
        n_time = min(int(unique_time.numel()), extra_len)
        if n_time > 0:
            extra_time[b, :n_time] = unique_time[:n_time]
            extra_mask[b, :n_time] = True

    target_event_seq = torch.cat(
        [
            batch_cpu["target_event_seq"],
            torch.full(
                (bsz, extra_len),
                PAD_IDX,
                dtype=batch_cpu["target_event_seq"].dtype,
            ),
        ],
        dim=1,
    )
    target_time_seq = torch.cat(
        [
            batch_cpu["target_time_seq"],
            torch.zeros(
                bsz,
                extra_len,
                dtype=batch_cpu["target_time_seq"].dtype,
            ),
        ],
        dim=1,
    )
    readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1)
    target_dt_unique = None
    target_multi_hot = None
    if include_uts_targets:
        target_dt_unique = torch.cat(
            [
                batch_cpu["target_dt_unique"],
                torch.zeros(
                    bsz,
                    extra_len,
                    dtype=batch_cpu["target_dt_unique"].dtype,
                ),
            ],
            dim=1,
        )
        target_multi_hot = torch.cat(
            [
                batch_cpu["target_multi_hot"],
                torch.zeros(
                    bsz,
                    extra_len,
                    vocab_size,
                    dtype=batch_cpu["target_multi_hot"].dtype,
                ),
            ],
            dim=1,
        )

    for b in range(bsz):
        valid_event = batch_cpu["padding_mask"][b].bool()
        if not valid_event.any():
            continue
        n_event = int(valid_event.sum().item())
        events = torch.cat(
            [
                batch_cpu["event_seq"][b, :n_event],
                batch_cpu["target_event_seq"][b, n_event - 1:n_event],
            ]
        )
        times = torch.cat(
            [
                batch_cpu["time_seq"][b, :n_event],
                batch_cpu["target_time_seq"][b, n_event - 1:n_event],
            ]
        )
        valid_full = events > PAD_IDX
        events = events[valid_full]
        times = times[valid_full]
        if events.numel() == 0:
            continue

        for j in range(extra_len):
            if not bool(extra_mask[b, j]):
                continue
            pos = event_len + j
            t = extra_time[b, j]
            future = times > t
            if not future.any():
                readout_mask[b, pos] = False
                continue

            first_idx = int(torch.nonzero(future, as_tuple=False)[0].item())
            next_time = times[first_idx]
            next_event = events[first_idx]
            target_event_seq[b, pos] = next_event
            target_time_seq[b, pos] = next_time

            if not include_uts_targets:
                continue

            same_next_time = times == next_time
            next_events = events[same_next_time]
            valid_next_events = next_events[
                (next_events > PAD_IDX) & (next_events < vocab_size)
            ].long()
            if valid_next_events.numel() == 0:
                readout_mask[b, pos] = False
                continue
            target_multi_hot[b, pos, valid_next_events] = True
            target_dt_unique[b, pos] = next_time - t

    targets = {
        "target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
        "target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
        "readout_mask": readout_mask.to(device, non_blocking=non_blocking),
    }
    if include_uts_targets:
        targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking)
        targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking)
    return targets


def compute_next_step_loss(
    args: argparse.Namespace,
    model: DeepHealth,
    readout,
    criterion,
    batch: Dict[str, torch.Tensor],
    device: torch.device,
) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
    batch_cpu = batch
    batch = move_batch_to_device(
        {key: batch_cpu[key] for key in MODEL_INPUT_KEYS},
        device,
    )
    model_out = model(
        event_seq=batch["event_seq"],
        time_seq=batch["time_seq"],
        sex=batch["sex"],
        padding_mask=batch["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",
        return_output=True,
    )
    if not isinstance(model_out, DeepHealthOutput):
        raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
    targets = build_augmented_next_step_targets(
        batch_cpu=batch_cpu,
        model_out=model_out,
        include_uts_targets=args.target_mode == "uts",
    )
    readout_out = readout(
        hidden=model_out.hidden,
        time_seq=model_out.time_seq,
        padding_mask=model_out.padding_mask,
        readout_mask=targets["readout_mask"]
        if args.readout_name == "same_time_group_end"
        else None,
    )
    logits = model.calc_risk(readout_out.hidden)

    if args.target_mode == "delphi2m":
        loss, parts = criterion(
            logits=logits,
            target_events=targets["target_event_seq"],
            target_times=targets["target_time_seq"],
            current_times=model_out.time_seq,
            padding_mask=readout_out.readout_mask,
            return_components=True,
        )
    else:
        loss, parts = criterion(
            logits=logits,
            target_multi_hot=targets["target_multi_hot"],
            target_dt_unique=targets["target_dt_unique"],
            readout_mask=readout_out.readout_mask,
            return_components=True,
        )
    if not torch.isfinite(loss):
        raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
    return loss, parts


def run_epoch(
    logger: logging.Logger,
    args: argparse.Namespace,
    model: DeepHealth,
    readout,
    criterion,
    loader: DataLoader,
    optimizer: AdamW | None,
    device: torch.device,
    is_train: bool,
) -> float:
    model.train(is_train)
    readout.train(is_train)
    total = torch.zeros((), device=device)
    n_batches = 0
    skipped = 0
    parts_sum: Dict[str, torch.Tensor] = {}
    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, parts = compute_next_step_loss(args, model, readout, 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
            for name, value in parts.items():
                parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach()
            if (batch_idx + 1) % progress_interval == 0:
                avg = total / max(1, n_batches)
                postfix = {
                    "loss": f"{float(loss.detach().cpu()):.4f}",
                    "avg": f"{float(avg.detach().cpu()):.4f}",
                    "skipped": skipped,
                }
                for name, value in parts_sum.items():
                    postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}"
                progress.set_postfix(postfix)
        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: HealthDataset,
    run_name: str,
    train_subset,
    val_subset,
    test_subset,
) -> Dict[str, Any]:
    return {
        "run_name": run_name,
        "dataset_class": "NextStepHealthDataset",
        "collate_fn": "next_step_collate_fn",
        "model_class": "DeepHealth",
        "model_target_mode": "next_token",
        "target_mode": args.target_mode,
        "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": [int(x) for x in 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.target_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: (
            f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
            f"gap_{args.no_event_interval_years:g}y_{timestamp}"
        )
    )
    logger = setup_logging(run_dir)

    logger.info(f"Starting next-step training run: {run_name}")
    logger.info(f"Device: {device}")
    logger.info(f"extra_info_types: {args.extra_info_types or 'all'}")
    logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")

    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,
    )
    if args.train_eid_file and args.val_eid_file and args.test_eid_file:
        train_subset, val_subset, test_subset = split_dataset_by_eid_files(
            dataset=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_dataset(
            dataset=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"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
    )

    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,
    )

    model = build_model(args, dataset).to(device)
    readout = build_next_step_readout(args).to(device)
    criterion = build_next_step_loss(args)
    optimizer = AdamW(
        model.parameters(),
        lr=args.base_lr,
        betas=tuple(args.betas),
        weight_decay=args.weight_decay,
    )
    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)

    save_config(
        args,
        run_dir / "train_config.json",
        extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset),
    )

    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, readout, criterion, train_loader, optimizer, device, True)
        with torch.no_grad():
            val_loss = run_epoch(logger, args, model, readout, 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 next-step test split...")
    model.load_state_dict(torch.load(best_model_path, map_location=device))
    with torch.no_grad():
        test_loss = run_epoch(logger, args, model, readout, 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()