DeepHealth/evaluate_auc.py

"""
Evaluate disease-specific AUC for DeepHealth models.

This script follows the logic of the Delphi evaluation script supplied by the user:
    1. choose disease tokens of interest from dataset vocabulary;
  2. precompute, for each patient/target occurrence, the latest prediction token
     at least `offset` years before the target time;
  3. run model inference by disease chunks to avoid materializing all logits;
  4. compute AUC separately by sex and age bracket;
  5. aggregate age brackets with DeLong variance.

Efficiency notes:
  - transformer/readout inference is executed once and cached;
  - disease chunks reuse the cached hidden states and only recompute selected risk-head columns;
  - AUC work is parallelized on CPU across disease task blocks using process workers;
  - per-sex data are compacted into an event-level table before multiprocessing;
  - large per-sex arrays are installed once per worker with fork-style globals on Linux,
    avoiding repeated pickling of arrays for every disease.

Run from the DeepHealth code directory containing dataset.py, models.py,
readouts.py, and train_config.json-compatible checkpoints/configs.
"""
from __future__ import annotations

import argparse
import contextlib
import json
import math
import multiprocessing as mp
import os
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, Subset
from tqdm.auto import tqdm

from dataset import HealthDataset
from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
from models import DeepHealth
from readouts import build_readout
from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX


# ---------------------------------------------------------------------------
# DeLong AUC utilities, adapted from the supplied Delphi evaluation script
# ---------------------------------------------------------------------------

def compute_midrank(x: np.ndarray) -> np.ndarray:
    """Compute midranks for DeLong variance."""
    x = np.asarray(x, dtype=np.float64)
    order = np.argsort(x)
    sorted_x = x[order]
    n = len(x)
    ranks = np.zeros(n, dtype=np.float64)
    i = 0
    while i < n:
        j = i
        while j < n and sorted_x[j] == sorted_x[i]:
            j += 1
        ranks[i:j] = 0.5 * (i + j - 1)
        i = j
    out = np.empty(n, dtype=np.float64)
    out[order] = ranks + 1.0
    return out


def fast_delong(predictions_sorted_transposed: np.ndarray, label_1_count: int) -> Tuple[np.ndarray, np.ndarray]:
    """
    Fast DeLong covariance for ROC AUC.

    This evaluation uses one classifier at a time, so k is normally 1.
    In that case np.cov(v01) and np.cov(v10) must be scalar variances.
    Do NOT expand them to (m,m)/(n,n) covariance matrices; that is the
    source of the broadcast error:
        shapes (n_case,n_case) and (n_control,n_control)
    """
    predictions_sorted_transposed = np.asarray(
        predictions_sorted_transposed, dtype=np.float64)
    m = int(label_1_count)
    n = int(predictions_sorted_transposed.shape[1] - m)
    if m <= 0 or n <= 0:
        return np.array([np.nan], dtype=np.float64), np.array([[np.nan]], dtype=np.float64)

    positive_examples = predictions_sorted_transposed[:, :m]
    negative_examples = predictions_sorted_transposed[:, m:]
    k = int(predictions_sorted_transposed.shape[0])

    tx = np.empty((k, m), dtype=np.float64)
    ty = np.empty((k, n), dtype=np.float64)
    tz = np.empty((k, m + n), dtype=np.float64)
    for r in range(k):
        tx[r] = compute_midrank(positive_examples[r])
        ty[r] = compute_midrank(negative_examples[r])
        tz[r] = compute_midrank(predictions_sorted_transposed[r])

    aucs = tz[:, :m].sum(axis=1) / m / n - float(m + 1.0) / 2.0 / n

    # DeLong structural components:
    # v01: classifier x cases
    # v10: classifier x controls
    v01 = (tz[:, :m] - tx) / n
    v10 = 1.0 - (tz[:, m:] - ty) / m

    if k == 1:
        # np.cov on a single row is easy to misuse. The correct covariance
        # for one classifier is simply the scalar sample variance over cases
        # plus the scalar sample variance over controls.
        sx = np.var(v01[0], ddof=1) if m > 1 else 0.0
        sy = np.var(v10[0], ddof=1) if n > 1 else 0.0
        delong_cov = np.array([[sx / m + sy / n]], dtype=np.float64)
    else:
        # Multiple-classifier general case. rowvar=True: rows are classifiers.
        sx = np.cov(v01, rowvar=True) if m > 1 else np.zeros(
            (k, k), dtype=np.float64)
        sy = np.cov(v10, rowvar=True) if n > 1 else np.zeros(
            (k, k), dtype=np.float64)
        delong_cov = np.atleast_2d(sx) / m + np.atleast_2d(sy) / n

    return aucs, delong_cov


def get_auc_delong_var(control_scores: np.ndarray, case_scores: np.ndarray) -> Tuple[float, float]:
    """Return AUC and DeLong variance for controls/class-0 and cases/class-1."""
    control_scores = np.asarray(control_scores, dtype=np.float64)
    case_scores = np.asarray(case_scores, dtype=np.float64)
    if len(control_scores) == 0 or len(case_scores) == 0:
        return np.nan, np.nan

    ground_truth = np.array([1] * len(case_scores) +
                            [0] * len(control_scores), dtype=np.int8)
    predictions = np.concatenate(
        [case_scores, control_scores]).astype(np.float64, copy=False)
    order = (-ground_truth).argsort()
    label_1_count = int(ground_truth.sum())

    aucs, delong_cov = fast_delong(
        predictions[np.newaxis, order], label_1_count)
    var = float(np.asarray(delong_cov).reshape(-1)[0])
    return float(aucs[0]), var


# ---------------------------------------------------------------------------
# Disease selection
# ---------------------------------------------------------------------------

SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}


def _build_first_occurrence_maps(
    dataset: HealthDataset,
    subset_indices: np.ndarray,
) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]:
    patient_count = len(subset_indices)
    followup_end = np.full(patient_count, -np.inf, dtype=np.float32)
    death_time = np.full(patient_count, np.inf, dtype=np.float32)
    sex = np.full(patient_count, -1, dtype=np.int8)

    first_lists: Dict[int, List[Tuple[int, float]]] = {}

    for patient_id, dataset_index in enumerate(subset_indices.tolist()):
        s = dataset.samples[int(dataset_index)]

        seq_event = np.asarray(s["event_seq"], dtype=np.int64)
        seq_time = np.asarray(s["time_seq"], dtype=np.float32)
        tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
        tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)

        if seq_event.size == 0 or tgt_event.size == 0:
            continue

        full_event = np.concatenate([seq_event, tgt_event[-1:]])
        full_time = np.concatenate([seq_time, tgt_time[-1:]])

        sex[patient_id] = int(s["sex"])
        followup_end[patient_id] = np.max(full_time).astype(np.float32)

        uniq_tokens, first_idx = np.unique(full_event, return_index=True)

        for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
            token = int(token)
            event_time = float(full_time[int(idx)])
            if token not in first_lists:
                first_lists[token] = []
            first_lists[token].append((patient_id, event_time))

    packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
    for token, pairs in first_lists.items():
        if not pairs:
            continue
        packed[int(token)] = (
            np.asarray([p for p, _ in pairs], dtype=np.int32),
            np.asarray([t for _, t in pairs], dtype=np.float32),
        )

    return packed, followup_end, death_time, sex


def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
    ids: List[int] = []

    exact_codes = {"death", "<death>", "dth", "deceased", "mortality"}
    for token, code in dataset.label_id_to_code.items():
        token = int(token)
        if token in SPECIAL_TOKENS:
            continue
        text = str(code).strip().lower()
        if text in exact_codes or ("death" in text) or ("mortality" in text):
            ids.append(token)

    death_ids = sorted(set(int(x)
                           for x in ids if int(x) not in SPECIAL_TOKENS))
    print(f"[INFO] death token ids: {death_ids}")
    return death_ids


def select_disease_tokens(
    dataset: HealthDataset,
    requested_tokens: Optional[Sequence[int]],
    filter_min_total: int,
    first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
) -> List[int]:
    base = [
        int(token)
        for token, code in dataset.label_id_to_code.items()
        if int(token) not in SPECIAL_TOKENS and not str(code).startswith("<")
    ]
    base_set = set(base)
    print(f"[INFO] Valid disease tokens in current vocabulary: {len(base)}")

    if requested_tokens is not None:
        selected = sorted(
            set(int(x) for x in requested_tokens if int(x) in base_set and int(x) not in SPECIAL_TOKENS))
        print(
            "[INFO] Requested disease tokens provided: "
            f"input={len(list(requested_tokens))}, selected={len(selected)} (vocab/SPECIAL filtered).")
        print(f"[INFO] Final disease_ids count: {len(selected)}")
        return selected

    disease_ids = sorted(base)
    if int(filter_min_total) <= 0:
        print(
            f"[INFO] filter_min_total={int(filter_min_total)} <= 0; keeping all {len(disease_ids)} disease tokens.")
        print(f"[INFO] Final disease_ids count: {len(disease_ids)}")
        return disease_ids

    print(
        f"[INFO] Applying filter_min_total={int(filter_min_total)}: before={len(disease_ids)} tokens.")
    print(
        "[INFO] Using split first-occurrence patient counts for filter_min_total.")
    split_counts = {}
    for token, pairs in first_occurrence_by_token.items():
        token = int(token)
        if token not in base_set:
            continue
        split_counts[token] = len(np.unique(pairs[0]))
    filtered = [token for token in disease_ids if int(
        split_counts.get(token, 0)) > int(filter_min_total)]
    print(
        "[INFO] First-occurrence count filtering complete: "
        f"after={len(filtered)} tokens.")
    print(f"[INFO] Final disease_ids count: {len(filtered)}")
    return filtered


# ---------------------------------------------------------------------------
# Dataset/split/model helpers
# ---------------------------------------------------------------------------

def load_json_config(path: Optional[str]) -> Dict[str, Any]:
    if path is None:
        return {}
    p = Path(path)
    if not p.exists():
        return {}
    with p.open("r", encoding="utf-8") as f:
        return json.load(f)


def cfg_get(args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any], name: str, default: Any) -> Any:
    """Get a value from CLI args first, then train_config.json, then default.

    This helper intentionally accepts either an argparse.Namespace or a dict.
    The earlier version passed cfg as both args and cfg, then tried to access
    args.eval_split, which fails because dict has no attributes.
    """
    val = None
    if args is not None:
        if isinstance(args, dict):
            val = args.get(name, None)
        else:
            val = getattr(args, name, None)
    if val is not None:
        return val
    return cfg.get(name, default)


def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    total = train_ratio + val_ratio + test_ratio
    if not np.isclose(total, 1.0, atol=1e-6):
        raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
    rng = np.random.RandomState(seed)
    idx = rng.permutation(n)
    n_train = int(n * train_ratio)
    n_val = int(n * val_ratio)
    return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]


def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
    model_target_mode = str(cfg_get(
        args, cfg, "model_target_mode", "next_token")).lower()
    if model_target_mode not in {"next_token", "all_future"}:
        raise ValueError(
            f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
        )
    return DeepHealth(
        vocab_size=dataset.vocab_size,
        n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
        n_head=int(cfg_get(args, cfg, "n_head", 10)),
        n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
        n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
        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=int(cfg_get(args, cfg, "n_bins", 16)),
        extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")),
        target_mode=model_target_mode,
        time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
        dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
        dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
    )


def _extract_state_dict(ckpt: Any) -> Dict[str, Any]:
    if isinstance(ckpt, dict) and "model" in ckpt:
        return ckpt["model"]
    elif isinstance(ckpt, dict) and "state_dict" in ckpt:
        return ckpt["state_dict"]
    return ckpt


def load_checkpoint_state_dict(checkpoint_path: str, map_location: str | torch.device = "cpu") -> Dict[str, Any]:
    ckpt = torch.load(checkpoint_path, map_location=map_location)
    state = _extract_state_dict(ckpt)
    if not isinstance(state, dict):
        raise TypeError(
            f"Unsupported checkpoint payload type: {type(state)}")
    return state


def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
    mode = str(cfg_dist_mode).lower()
    has_rho_head = any(str(k).startswith("rho_head.")
                       for k in state_dict.keys())
    has_rho_death_head = any(str(k).startswith("rho_death_head.")
                             for k in state_dict.keys())

    if has_rho_head and mode != "weibull":
        print(
            "[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
        return "weibull"
    if has_rho_death_head and mode != "mixed":
        print(
            "[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
        return "mixed"
    if (not has_rho_head) and mode == "weibull":
        print(
            "[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
        return "exponential"
    if (not has_rho_death_head) and mode == "mixed":
        print(
            "[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
        return "exponential"
    return mode


def load_model_state(
    model: torch.nn.Module,
    checkpoint_path: str,
    device: torch.device,
    state_dict: Optional[Dict[str, Any]] = None,
) -> None:
    state = state_dict if state_dict is not None else load_checkpoint_state_dict(
        checkpoint_path, map_location=device)

    model.load_state_dict(state, strict=True)


def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any]) -> Tuple[Subset, np.ndarray]:
    train_ratio = float(cfg_get(args, cfg, "train_ratio", 0.7))
    val_ratio = float(cfg_get(args, cfg, "val_ratio", 0.15))
    test_ratio = float(cfg_get(args, cfg, "test_ratio", 0.15))
    seed = int(cfg_get(args, cfg, "seed", 42))
    eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
    dataset_subset_size = cfg_get(args, cfg, "dataset_subset_size", None)

    train_idx, val_idx, test_idx = split_indices(
        len(dataset), train_ratio, val_ratio, test_ratio, seed)
    split_map = {
        "train": train_idx,
        "val": val_idx,
        "valid": val_idx,
        "validation": val_idx,
        "test": test_idx,
        "all": np.arange(len(dataset)),
    }
    if eval_split not in split_map:
        raise ValueError(
            f"eval_split must be one of {sorted(split_map)}, got {eval_split!r}")

    indices = split_map[eval_split]
    if dataset_subset_size is not None and int(dataset_subset_size) > 0:
        indices = indices[: int(dataset_subset_size)]
    return Subset(dataset, indices.tolist()), np.asarray(indices, dtype=np.int64)


def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None:
    meta = cfg.get("dataset_metadata")
    if not isinstance(meta, dict):
        return

    actual: Dict[str, Any] = {
        "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],
    }
    mismatches = [
        f"{key}: train_config={meta.get(key)!r}, current_dataset={value!r}"
        for key, value in actual.items()
        if key in meta and meta.get(key) != value
    ]
    if mismatches:
        raise RuntimeError(
            "Current dataset metadata does not match train_config.json. "
            "Use the same prepared data and extra_info_types as training. "
            + "; ".join(mismatches)
        )


# ---------------------------------------------------------------------------
# Batched inference + cached hidden states
# ---------------------------------------------------------------------------


def _numpy_hidden_dtype(name: str) -> np.dtype:
    name = str(name).lower()
    if name in {"float16", "fp16", "half"}:
        return np.float16
    if name in {"bfloat16", "bf16"}:
        # NumPy has limited bfloat16 support; store as fp16 for compact CPU cache.
        return np.float16
    if name in {"float32", "fp32", "single"}:
        return np.float32
    raise ValueError(
        f"hidden_cache_dtype must be float16 or float32, got {name!r}")


@torch.inference_mode()
def infer_readout_hidden(
    model: DeepHealth,
    loader: DataLoader,
    device: torch.device,
    model_target_mode: str,
    readout_name: str,
    readout_reduce: str,
    use_amp: bool,
    hidden_cache_dtype: str = "float16",
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
    """Cache per-position hidden states used by the unchanged AUC logic."""
    model_target_mode = str(model_target_mode).lower()
    if model_target_mode not in {"next_token", "all_future"}:
        raise ValueError(
            f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
        )

    readout = None
    if model_target_mode == "next_token" and readout_name == "same_time_group_end":
        readout = build_readout("same_time_group_end",
                                reduce=readout_reduce).to(device)
    elif model_target_mode == "next_token":
        readout = build_readout(readout_name).to(device)
    if readout is not None:
        readout.eval()

    hidden_parts: List[np.ndarray] = []
    arrays: Dict[str, List[np.ndarray]] = {
        "event_seq": [],
        "time_seq": [],
        "target_event_seq": [],
        "target_time_seq": [],
        "padding_mask": [],
        "readout_mask": [],
        "sex": [],
    }
    max_len = 0
    out_dtype = _numpy_hidden_dtype(hidden_cache_dtype)
    autocast_enabled = bool(use_amp and device.type == "cuda")

    for batch in tqdm(loader, desc="Model/readout inference", leave=False, dynamic_ncols=True):
        batch_dev = {
            k: (v.to(device, non_blocking=True)
                if isinstance(v, torch.Tensor) else v)
            for k, v in batch.items()
        }
        event_seq = batch_dev["event_seq"]
        time_seq = batch_dev["time_seq"]
        padding_mask = batch_dev["padding_mask"]

        amp_context = (
            torch.autocast(device_type=device.type, dtype=torch.float16)
            if autocast_enabled else contextlib.nullcontext()
        )
        with amp_context:
            if model_target_mode == "all_future":
                batch_size, seq_len = event_seq.shape
                hidden = torch.zeros(
                    batch_size,
                    seq_len,
                    model.n_embd,
                    device=event_seq.device,
                    dtype=torch.float32,
                )
                for pos in range(seq_len):
                    active = padding_mask[:, pos].bool()
                    if not active.any():
                        continue
                    hidden_pos = model(
                        event_seq=event_seq[active],
                        time_seq=time_seq[active],
                        sex=batch_dev["sex"][active],
                        padding_mask=padding_mask[active],
                        t_query=time_seq[active, pos],
                        other_type=batch_dev["other_type"][active],
                        other_value=batch_dev["other_value"][active],
                        other_value_kind=batch_dev["other_value_kind"][active],
                        other_time=batch_dev["other_time"][active],
                        target_mode="all_future",
                    )
                    hidden[active, pos, :] = hidden_pos.float()
                readout_mask_np = batch["padding_mask"].cpu().numpy()
            else:
                hidden_raw = model(
                    event_seq=event_seq,
                    time_seq=time_seq,
                    sex=batch_dev["sex"],
                    padding_mask=padding_mask,
                    other_type=batch_dev["other_type"],
                    other_value=batch_dev["other_value"],
                    other_value_kind=batch_dev["other_value_kind"],
                    other_time=batch_dev["other_time"],
                    target_mode="next_token",
                )
                ro = readout(
                    hidden=hidden_raw,
                    time_seq=time_seq,
                    padding_mask=padding_mask,
                    readout_mask=batch_dev["readout_mask"],
                )
                hidden = ro.hidden
                readout_mask_np = ro.readout_mask.detach().cpu().numpy()

        h = hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
        hidden_parts.append(h)
        max_len = max(max_len, h.shape[1])

        for k in arrays:
            if k == "sex":
                arrays[k].append(
                    batch[k].cpu().numpy().astype(np.int8, copy=False))
            else:
                arrays[k].append(batch[k].cpu().numpy())
        arrays["readout_mask"][-1] = readout_mask_np

    def pad_3d(parts: List[np.ndarray], fill: float = 0.0) -> np.ndarray:
        out = np.full(
            (sum(x.shape[0] for x in parts), max_len, parts[0].shape[2]),
            fill,
            dtype=out_dtype,
        )
        s = 0
        for x in parts:
            out[s:s + x.shape[0], :x.shape[1], :] = x
            s += x.shape[0]
        return out

    def pad_2d(parts: List[np.ndarray], fill: Any = 0, dtype: Optional[np.dtype] = None) -> np.ndarray:
        dtype = dtype or parts[0].dtype
        out = np.full((sum(x.shape[0]
                           for x in parts), max_len), fill, dtype=dtype)
        s = 0
        for x in parts:
            out[s:s + x.shape[0], :x.shape[1]] = x
            s += x.shape[0]
        return out

    hidden_all = pad_3d(hidden_parts)
    arr_out = {
        "event_seq": pad_2d(arrays["event_seq"], PAD_IDX, np.int64),
        "time_seq": pad_2d(arrays["time_seq"], 0.0, np.float32),
        "target_event_seq": pad_2d(arrays["target_event_seq"], PAD_IDX, np.int64),
        "target_time_seq": pad_2d(arrays["target_time_seq"], 0.0, np.float32),
        "padding_mask": pad_2d(arrays["padding_mask"], False, bool),
        "readout_mask": pad_2d(arrays["readout_mask"], False, bool),
        "sex": np.concatenate(arrays["sex"], axis=0),
    }
    return hidden_all, arr_out


@torch.inference_mode()
def compute_logits_for_disease_chunk(
    model: DeepHealth,
    hidden_all: np.ndarray,
    disease_ids: Sequence[int],
    device: torch.device,
    logit_batch_size: int,
    use_amp: bool,
) -> np.ndarray:
    """Project cached hidden states to only the requested disease columns."""
    n = int(hidden_all.shape[0])
    logit_batch_size = max(1, int(logit_batch_size))
    disease_ids = [int(x) for x in disease_ids]

    compute_dtype = torch.float16 if (
        device.type == "cuda" and use_amp) else torch.float32
    weight = model.risk_head.weight[disease_ids].detach().to(
        device=device, dtype=compute_dtype)
    bias = None
    if model.risk_head.bias is not None:
        bias = model.risk_head.bias[disease_ids].detach().to(
            device=device, dtype=compute_dtype)

    parts: List[np.ndarray] = []
    for start in tqdm(range(0, n, logit_batch_size), desc="Risk-head projection", leave=False, dynamic_ncols=True):
        end = min(start + logit_batch_size, n)
        h = torch.from_numpy(hidden_all[start:end]).to(
            device=device, dtype=compute_dtype, non_blocking=True)
        logits = torch.matmul(h, weight.t())
        if bias is not None:
            logits = logits + bias
        parts.append(logits.float().cpu().numpy().astype(
            np.float32, copy=False))
        del h, logits
    return np.concatenate(parts, axis=0)


# ---------------------------------------------------------------------------
# CPU-parallel calibration AUC
# ---------------------------------------------------------------------------

_WORKER: Dict[str, Any] = {}


def _build_flat_eval_table(
    p_sex: np.ndarray,
    time_seq: np.ndarray,
    target_time_seq: np.ndarray,
    target_event_seq: np.ndarray,
    padding_mask: np.ndarray,
    readout_mask: np.ndarray,
    offset: float,
    valid_target_min_id: int,
    age_groups: np.ndarray,
) -> Dict[str, np.ndarray]:
    """
    Build a compact event-level table for AUC evaluation.

    The previous version kept 2D patient x position arrays and each disease task
    repeatedly scanned them. This table keeps only eligible target occurrences:
      - target event is a real disease/event target;
      - a valid readout token exists at least `offset` years before target time;
      - the prediction token itself is valid and readout-valid;
      - predicted age falls inside the requested age brackets.

    The disease-specific logic is unchanged; this only changes the physical layout
    of the data to make CPU multiprocessing cheaper and more cache-friendly.
    """
    if len(age_groups) < 2:
        raise ValueError("age_groups must contain at least two values")
    age_start = float(age_groups[0])
    age_step = float(age_groups[1] - age_groups[0])
    n_age = int(len(age_groups))

    # Raw valid-target table is used only to decide which patients are cases for
    # a disease. This intentionally happens before offset/age filtering, matching
    # the supplied Delphi logic: a patient who ever has disease k is not a control
    # for k, even if that occurrence later lacks a valid prediction point.
    raw_patient_idx, raw_target_idx = np.where(
        target_event_seq > int(valid_target_min_id))
    raw_target_event = target_event_seq[raw_patient_idx, raw_target_idx].astype(
        np.int32, copy=False)
    raw_sort_order = np.argsort(raw_target_event, kind="mergesort")
    raw_sorted_target_event = raw_target_event[raw_sort_order]
    raw_patient_idx = raw_patient_idx.astype(np.int32, copy=False)

    valid_pred_pos = (
        (time_seq[:, :, None] <= (target_time_seq[:, None, :] - float(offset)))
        & readout_mask[:, :, None]
    )
    pos_index = np.arange(time_seq.shape[1], dtype=np.int32)[None, :, None]
    pred_idx_precompute = np.where(
        valid_pred_pos, pos_index, -1).max(axis=1).astype(np.int32)

    candidate = (target_event_seq > int(valid_target_min_id)) & (
        pred_idx_precompute >= 0)
    patient_idx, target_idx = np.where(candidate)
    if patient_idx.size == 0:
        return {
            "patient": np.empty(0, dtype=np.int32),
            "target_event": np.empty(0, dtype=np.int32),
            "pred_idx": np.empty(0, dtype=np.int32),
            "age_bin": np.empty(0, dtype=np.int16),
            "target_time": np.empty(0, dtype=np.float32),
            "sort_order": np.empty(0, dtype=np.int64),
            "sorted_target_event": np.empty(0, dtype=np.int32),
            "raw_patient": raw_patient_idx,
            "raw_sort_order": raw_sort_order.astype(np.int64, copy=False),
            "raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False),
            "p_sex": p_sex,
            "age_groups": age_groups.astype(np.float32, copy=False),
            "n_patients": np.int32(time_seq.shape[0]),
        }

    pred_idx = pred_idx_precompute[patient_idx, target_idx]
    pred_ok = padding_mask[patient_idx,
                           pred_idx] & readout_mask[patient_idx, pred_idx]
    if not np.any(pred_ok):
        return {
            "patient": np.empty(0, dtype=np.int32),
            "target_event": np.empty(0, dtype=np.int32),
            "pred_idx": np.empty(0, dtype=np.int32),
            "age_bin": np.empty(0, dtype=np.int16),
            "target_time": np.empty(0, dtype=np.float32),
            "sort_order": np.empty(0, dtype=np.int64),
            "sorted_target_event": np.empty(0, dtype=np.int32),
            "raw_patient": raw_patient_idx,
            "raw_sort_order": raw_sort_order.astype(np.int64, copy=False),
            "raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False),
            "p_sex": p_sex,
            "age_groups": age_groups.astype(np.float32, copy=False),
            "n_patients": np.int32(time_seq.shape[0]),
        }

    patient_idx = patient_idx[pred_ok].astype(np.int32, copy=False)
    target_idx = target_idx[pred_ok]
    pred_idx = pred_idx[pred_ok].astype(np.int32, copy=False)

    pred_age = time_seq[patient_idx, pred_idx].astype(np.float32, copy=False)
    age_bin = np.floor((pred_age - age_start) / age_step).astype(np.int16)
    age_ok = (age_bin >= 0) & (age_bin < n_age)
    if not np.any(age_ok):
        return {
            "patient": np.empty(0, dtype=np.int32),
            "target_event": np.empty(0, dtype=np.int32),
            "pred_idx": np.empty(0, dtype=np.int32),
            "age_bin": np.empty(0, dtype=np.int16),
            "target_time": np.empty(0, dtype=np.float32),
            "sort_order": np.empty(0, dtype=np.int64),
            "sorted_target_event": np.empty(0, dtype=np.int32),
            "raw_patient": raw_patient_idx,
            "raw_sort_order": raw_sort_order.astype(np.int64, copy=False),
            "raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False),
            "p_sex": p_sex,
            "age_groups": age_groups.astype(np.float32, copy=False),
            "n_patients": np.int32(time_seq.shape[0]),
        }

    patient_idx = patient_idx[age_ok]
    target_idx = target_idx[age_ok]
    pred_idx = pred_idx[age_ok]
    age_bin = age_bin[age_ok]
    target_event = target_event_seq[patient_idx,
                                    target_idx].astype(np.int32, copy=False)
    target_time = target_time_seq[patient_idx,
                                  target_idx].astype(np.float32, copy=False)

    sort_order = np.argsort(target_event, kind="mergesort")
    sorted_target_event = target_event[sort_order]

    return {
        "patient": patient_idx.astype(np.int32, copy=False),
        "target_event": target_event,
        "pred_idx": pred_idx.astype(np.int32, copy=False),
        "age_bin": age_bin.astype(np.int16, copy=False),
        "target_time": target_time,
        "sort_order": sort_order.astype(np.int64, copy=False),
        "sorted_target_event": sorted_target_event.astype(np.int32, copy=False),
        "raw_patient": raw_patient_idx,
        "raw_sort_order": raw_sort_order.astype(np.int64, copy=False),
        "raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False),
        "p_sex": p_sex,
        "age_groups": age_groups.astype(np.float32, copy=False),
        "n_patients": np.int32(time_seq.shape[0]),
    }


def _init_auc_worker_flat(
    patient: np.ndarray,
    target_event: np.ndarray,
    pred_idx: np.ndarray,
    age_bin: np.ndarray,
    target_time: np.ndarray,
    sort_order: np.ndarray,
    sorted_target_event: np.ndarray,
    raw_patient: np.ndarray,
    raw_sort_order: np.ndarray,
    raw_sorted_target_event: np.ndarray,
    p_sex: np.ndarray,
    age_groups: np.ndarray,
    n_patients: int,
):
    # Prevent BLAS/OpenMP oversubscription when many worker processes are active.
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")
    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
    os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")

    _WORKER.clear()
    _WORKER.update({
        "patient": patient,
        "target_event": target_event,
        "pred_idx": pred_idx,
        "age_bin": age_bin,
        "target_time": target_time,
        "sort_order": sort_order,
        "sorted_target_event": sorted_target_event,
        "raw_patient": raw_patient,
        "raw_sort_order": raw_sort_order,
        "raw_sorted_target_event": raw_sorted_target_event,
        "p_sex": p_sex,
        "age_groups": age_groups,
        "n_patients": int(n_patients),
    })


def _case_indices_for_token(token: int) -> np.ndarray:
    sorted_target_event = _WORKER["sorted_target_event"]
    sort_order = _WORKER["sort_order"]
    left = np.searchsorted(sorted_target_event, int(token), side="left")
    right = np.searchsorted(sorted_target_event, int(token), side="right")
    if right <= left:
        return np.empty(0, dtype=np.int64)
    return sort_order[left:right]


def _raw_case_patients_for_token(token: int) -> np.ndarray:
    raw_sorted_target_event = _WORKER["raw_sorted_target_event"]
    raw_sort_order = _WORKER["raw_sort_order"]
    raw_patient = _WORKER["raw_patient"]
    left = np.searchsorted(raw_sorted_target_event, int(token), side="left")
    right = np.searchsorted(raw_sorted_target_event, int(token), side="right")
    if right <= left:
        return np.empty(0, dtype=np.int32)
    return np.unique(raw_patient[raw_sort_order[left:right]])


def _calibration_auc_one_disease_flat(task: Tuple[int, int]) -> List[Dict[str, Any]]:
    j, token = task
    patient = _WORKER["patient"]
    pred_idx = _WORKER["pred_idx"]
    age_bin = _WORKER["age_bin"]
    target_time = _WORKER["target_time"]
    p_sex = _WORKER["p_sex"]
    age_groups = _WORKER["age_groups"]
    n_patients = _WORKER["n_patients"]

    case_idx = _case_indices_for_token(int(token))
    if case_idx.size < 2:
        return []

    case_patients = _raw_case_patients_for_token(int(token))
    if case_patients.size == 0:
        return []

    patient_has_case = np.zeros(n_patients, dtype=bool)
    patient_has_case[case_patients] = True

    # Controls follow the supplied Delphi logic: any eligible target occurrence from
    # a patient who never has this disease token in the evaluated target table.
    control_idx = np.flatnonzero(~patient_has_case[patient])
    if control_idx.size == 0:
        return []

    out: List[Dict[str, Any]] = []
    for b, aa in enumerate(age_groups):
        case_b = case_idx[age_bin[case_idx] == b]
        control_b = control_idx[age_bin[control_idx] == b]
        if case_b.size == 0 or control_b.size == 0:
            continue

        # Match previous deterministic one-occurrence-per-patient behavior within
        # each age bracket, separately for cases and controls. This avoids letting
        # high-utilization patients dominate the AUC.
        _, case_first = np.unique(patient[case_b], return_index=True)
        _, control_first = np.unique(patient[control_b], return_index=True)
        case_keep = case_b[case_first]
        control_keep = control_b[control_first]
        if case_keep.size == 0 or control_keep.size == 0:
            continue

        # Delphi2M-aligned AUC score: use disease-specific eta/logit only.
        # Prediction offset filters eligible prediction tokens but does not enter the score.
        case_scores = p_sex[patient[case_keep], pred_idx[case_keep], j].astype(
            np.float64, copy=False)
        control_scores = p_sex[patient[control_keep], pred_idx[control_keep], j].astype(
            np.float64, copy=False)
        if case_scores.size == 0 or control_scores.size == 0:
            continue

        auc_value, auc_var = get_auc_delong_var(control_scores, case_scores)

        out.append({
            "token": int(token),
            "auc": float(auc_value),
            "auc_delong": float(auc_value),
            "auc_variance_delong": float(auc_var),
            "age": float(aa),
            "age_right": float(aa + (age_groups[1] - age_groups[0])),
            "n_healthy": int(control_scores.size),
            "n_diseased": int(case_scores.size),
            "mean_target_time": float(np.mean(target_time[case_keep])) if case_keep.size else np.nan,
        })
    return out


def _calibration_auc_task_block(tasks: Sequence[Tuple[int, int]]) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    for task in tasks:
        rows.extend(_calibration_auc_one_disease_flat(task))
    return rows


def _split_tasks_for_workers(
    tasks: Sequence[Tuple[int, int]],
    effective_workers: int,
    task_chunk_size: int,
) -> List[List[Tuple[int, int]]]:
    if not tasks:
        return []
    if task_chunk_size <= 0:
        # Enough chunks to keep workers busy without creating one Future per disease.
        task_chunk_size = max(1, math.ceil(
            len(tasks) / max(1, effective_workers * 4)))
    return [list(tasks[i:i + task_chunk_size]) for i in range(0, len(tasks), task_chunk_size)]


def compute_auc_chunk_parallel(
    p_chunk: np.ndarray,
    arrays: Dict[str, np.ndarray],
    disease_ids: Sequence[int],
    sex_value: int,
    sex_name: str,
    age_groups: np.ndarray,
    offset: float,
    valid_target_min_id: int,
    num_workers: int,
    auc_task_chunk_size: int = 0,
) -> List[Dict[str, Any]]:
    sex_mask = arrays["sex"] == sex_value
    if not np.any(sex_mask):
        return []

    flat = _build_flat_eval_table(
        p_sex=p_chunk[sex_mask],
        time_seq=arrays["time_seq"][sex_mask],
        target_time_seq=arrays["target_time_seq"][sex_mask],
        target_event_seq=arrays["target_event_seq"][sex_mask],
        padding_mask=arrays["padding_mask"][sex_mask],
        readout_mask=arrays["readout_mask"][sex_mask],
        offset=offset,
        valid_target_min_id=valid_target_min_id,
        age_groups=age_groups,
    )
    if flat["patient"].size == 0:
        return []

    # Skip diseases with no cases in this sex before sending tasks to workers.
    sorted_events = flat["sorted_target_event"]
    tasks = []
    for j, token in enumerate(disease_ids):
        left = np.searchsorted(sorted_events, int(token), side="left")
        right = np.searchsorted(sorted_events, int(token), side="right")
        if right - left >= 2:
            tasks.append((j, int(token)))
    if not tasks:
        return []

    effective_workers = max(1, min(int(num_workers), len(tasks)))
    if effective_workers <= 1:
        _init_auc_worker_flat(
            flat["patient"], flat["target_event"], flat["pred_idx"], flat["age_bin"],
            flat["target_time"], flat["sort_order"], flat["sorted_target_event"],
            flat["raw_patient"], flat["raw_sort_order"], flat["raw_sorted_target_event"],
            flat["p_sex"], flat["age_groups"], int(flat["n_patients"]),
        )
        nested = [_calibration_auc_one_disease_flat(t) for t in tqdm(
            tasks, desc=f"AUC {sex_name}", leave=False, dynamic_ncols=True)]
    else:
        ctx = mp.get_context("fork") if hasattr(
            os, "fork") else mp.get_context()
        task_blocks = _split_tasks_for_workers(
            tasks, effective_workers, int(auc_task_chunk_size))
        with ProcessPoolExecutor(
            max_workers=effective_workers,
            mp_context=ctx,
            initializer=_init_auc_worker_flat,
            initargs=(
                flat["patient"], flat["target_event"], flat["pred_idx"], flat["age_bin"],
                flat["target_time"], flat["sort_order"], flat["sorted_target_event"],
                flat["raw_patient"], flat["raw_sort_order"], flat["raw_sorted_target_event"],
                flat["p_sex"], flat["age_groups"], int(flat["n_patients"]),
            ),
        ) as ex:
            nested = list(tqdm(
                ex.map(_calibration_auc_task_block, task_blocks),
                total=len(task_blocks),
                desc=f"AUC {sex_name}",
                leave=False,
                dynamic_ncols=True,
            ))

    out: List[Dict[str, Any]] = []
    for rows in nested:
        for r in rows:
            r["sex"] = sex_name
            r["offset"] = float(offset)
            out.append(r)
    return out


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

def evaluate_auc_pipeline(
    model: DeepHealth,
    loader: DataLoader,
    dataset: HealthDataset,
    output_path: Optional[str],
    diseases_of_interest: Optional[Sequence[int]],
    filter_min_total: int,
    first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
    include_death: bool,
    exclude_death: bool,
    disease_chunk_size: int,
    age_groups: np.ndarray,
    offsets: Sequence[float],
    device: torch.device,
    model_target_mode: str,
    readout_name: str,
    readout_reduce: str,
    num_workers_auc: int,
    use_amp: bool,
    auc_task_chunk_size: int = 0,
    hidden_cache_dtype: str = "float16",
    logit_batch_size: int = 256,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    model.eval().to(device)

    disease_ids = select_disease_tokens(
        dataset=dataset,
        requested_tokens=diseases_of_interest,
        filter_min_total=filter_min_total,
        first_occurrence_by_token=first_occurrence_by_token,
    )
    disease_ids = [int(k) for k in disease_ids if 0 <=
                   int(k) < dataset.vocab_size]

    death_token_ids = _get_death_token_ids(dataset)
    if (not bool(include_death)) or bool(exclude_death):
        before = len(disease_ids)
        death_set = set(int(x) for x in death_token_ids)
        disease_ids = [int(x) for x in disease_ids if int(x) not in death_set]
        print(
            "[INFO] Death exclusion applied on final disease_ids: "
            f"include_death={bool(include_death)}, exclude_death={bool(exclude_death)}, "
            f"before={before}, after={len(disease_ids)}.")

    if not disease_ids:
        raise ValueError("No diseases selected for evaluation.")

    if disease_chunk_size is None or int(disease_chunk_size) <= 0:
        disease_chunk_size = len(disease_ids)
    disease_chunk_size = max(1, int(disease_chunk_size))

    num_chunks = math.ceil(len(disease_ids) / disease_chunk_size)
    chunks = np.array_split(np.asarray(
        disease_ids, dtype=np.int64), num_chunks)
    print(
        f"Evaluating {len(disease_ids)} disease tokens in {len(chunks)} chunk(s).")
    print("Using Delphi2M-aligned rate/logit score for AUC.")
    print("AUC score = disease-specific eta at the latest eligible prediction token.")
    print("Prediction offset controls eligibility only and does not enter the score.")
    print(f"Evaluating prediction offsets: {', '.join(f'{x:g}' for x in offsets)} years.")

    # In current dataset sex is normalized to 0/1. UKB convention after normalization: 0=female, 1=male.
    sex_items = [("female", 0), ("male", 1)]
    all_rows: List[Dict[str, Any]] = []

    valid_target_min_id = CHECKUP_IDX if NO_EVENT_IDX >= dataset.vocab_size else CHECKUP_IDX
    # If NO_EVENT exists and should not be a disease/control target, require target > NO_EVENT_IDX.
    if NO_EVENT_IDX in dataset.label_id_to_code and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>":
        valid_target_min_id = NO_EVENT_IDX

    hidden_all, arrays = infer_readout_hidden(
        model=model,
        loader=loader,
        device=device,
        model_target_mode=model_target_mode,
        readout_name=readout_name,
        readout_reduce=readout_reduce,
        use_amp=use_amp,
        hidden_cache_dtype=hidden_cache_dtype,
    )
    print(
        f"Cached readout hidden: shape={hidden_all.shape}, dtype={hidden_all.dtype}")

    for chunk_idx, chunk in enumerate(tqdm(chunks, desc="Processing disease chunks", dynamic_ncols=True)):
        p_chunk = compute_logits_for_disease_chunk(
            model=model,
            hidden_all=hidden_all,
            disease_ids=chunk.tolist(),
            device=device,
            logit_batch_size=logit_batch_size,
            use_amp=use_amp,
        )
        for offset in offsets:
            for sex_name, sex_value in sex_items:
                rows = compute_auc_chunk_parallel(
                    p_chunk=p_chunk,
                    arrays=arrays,
                    disease_ids=chunk.tolist(),
                    sex_value=sex_value,
                    sex_name=sex_name,
                    age_groups=age_groups,
                    offset=float(offset),
                    valid_target_min_id=valid_target_min_id,
                    num_workers=num_workers_auc,
                    auc_task_chunk_size=auc_task_chunk_size,
                )
                for r in rows:
                    r["disease_chunk_idx"] = int(chunk_idx)
                all_rows.extend(rows)
        del p_chunk

    del hidden_all, arrays
    df_auc_unpooled = pd.DataFrame(all_rows)
    if df_auc_unpooled.empty:
        raise RuntimeError(
            "No AUC rows were produced. Check offset, age_groups, eval split, and disease ids.")

    # Keep outputs self-contained with only evaluation fields and token->code mapping.
    df_auc_unpooled["label_code"] = df_auc_unpooled["token"].map(
        dataset.label_id_to_code)

    print("Using DeLong method to calculate AUC confidence intervals.")
    grouped = df_auc_unpooled.groupby(
        ["token", "label_code", "offset"], dropna=False, as_index=False)
    df_auc = grouped.agg(
        auc=("auc_delong", "mean"),
        n_strata=("auc_delong", "size"),
        n_diseased=("n_diseased", "sum"),
        n_healthy=("n_healthy", "sum"),
        auc_variance_sum=("auc_variance_delong", "sum"),
    )
    df_auc["auc_variance_delong"] = (
        df_auc["auc_variance_sum"]
        / (df_auc["n_strata"].clip(lower=1).astype(np.float64) ** 2)
    )
    df_auc = df_auc.drop(columns=["auc_variance_sum"])

    if output_path is not None:
        out_dir = Path(output_path)
        out_dir.mkdir(parents=True, exist_ok=True)
        df_auc.to_csv(out_dir / "df_both.csv", index=False)
        df_auc_unpooled.to_csv(
            out_dir / "df_auc_unpooled.csv", index=False)

    return df_auc_unpooled, df_auc


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_int_list(s: Any) -> Optional[List[int]]:
    if s is None:
        return None
    if isinstance(s, (list, tuple, np.ndarray)):
        return [int(x) for x in s]
    text = str(s).strip()
    if text == "":
        return None
    if text.startswith("["):
        try:
            values = json.loads(text)
        except json.JSONDecodeError as exc:
            raise ValueError(f"Invalid integer list: {text!r}") from exc
        if not isinstance(values, list):
            raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
        return [int(x) for x in values]
    return [int(x.strip()) for x in text.split(",") if x.strip()]


def parse_float_list(s: Any) -> Optional[List[float]]:
    if s is None:
        return None
    if isinstance(s, (list, tuple, np.ndarray)):
        return [float(x) for x in s]
    text = str(s).strip()
    if text == "":
        return None
    if text.startswith("["):
        try:
            values = json.loads(text)
        except json.JSONDecodeError as exc:
            raise ValueError(f"Invalid float list: {text!r}") from exc
        if not isinstance(values, list):
            raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
        return [float(x) for x in values]
    return [float(x.strip()) for x in text.split(",") if x.strip()]


def make_auc_offsets(args: argparse.Namespace, cfg: Dict[str, Any]) -> List[float]:
    explicit_offsets = parse_float_list(cfg_get(args, cfg, "offsets", None))
    if explicit_offsets is not None:
        base_offsets = explicit_offsets
    else:
        next_token_offset = float(cfg_get(args, cfg, "offset", 0.1))
        base_offsets = [next_token_offset, 1.0, 5.0, 10.0]

    offsets: List[float] = []
    seen = set()
    for value in base_offsets:
        value = float(value)
        key = round(value, 10)
        if key not in seen:
            offsets.append(value)
            seen.add(key)
    if not offsets:
        raise ValueError("At least one AUC offset is required.")
    return offsets


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate DeepHealth AUC")

    # Simplify arguments to only include run_path and output_path
    parser.add_argument("--run_path", type=str, required=True,
                        help="Path containing train_config.json and best_model.pt")
    parser.add_argument("--output_path", type=str,
                        default=None, help="Defaults to run_path")
    parser.add_argument("--eval_split", type=str, default=None,
                        choices=["train", "val", "valid",
                                 "validation", "test", "all"],
                        help="Evaluation split. Defaults to 'test' unless cfg contains eval_split.")
    parser.add_argument("--dataset_subset_size", type=int, default=None,
                        help="Optional number of patients from the selected split.")
    parser.add_argument("--batch_size", type=int, default=None,
                        help="Inference batch size; overrides train_config.json.")
    parser.add_argument("--num_workers", type=int, default=None,
                        help="DataLoader workers; overrides train_config.json.")
    parser.add_argument("--num_workers_auc", type=int, default=None,
                        help="CPU processes for AUC computation.")
    parser.add_argument("--auc_task_chunk_size", type=int, default=None,
                        help="Diseases per submitted CPU task block. 0/None auto-tunes.")
    parser.add_argument("--hidden_cache_dtype", type=str, default=None, choices=["float16", "float32"],
                        help="CPU dtype for cached readout hidden states. float16 saves memory and is usually enough for AUC.")
    parser.add_argument("--logit_batch_size", type=int, default=None,
                        help="Patient batch size for projecting cached hidden states to disease logits.")
    parser.add_argument("--disease_chunk_size", type=int, default=None,
                        help="Number of disease logits to materialize per inference pass. <=0 means one chunk (all diseases).")
    parser.add_argument("--filter_min_total", type=int, default=None,
                        help="Minimum metadata count for disease selection; default 0.")
    parser.add_argument("--offset", type=float, default=None,
                        help="Next-token prediction offset in years; preserved and evaluated alongside 1, 5, and 10 years by default.")
    parser.add_argument("--offsets", type=str, default=None,
                        help="Comma-separated prediction offsets in years. Overrides the default set of offset,1,5,10.")
    parser.add_argument("--age_start", type=float, default=None)
    parser.add_argument("--age_stop", type=float, default=None)
    parser.add_argument("--age_step", type=float, default=None)
    parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None,
                        help="Use CUDA autocast during inference.")

    args = parser.parse_args()

    # Extract paths from run_path
    run_path = Path(args.run_path)
    config_path = run_path / "train_config.json"
    model_ckpt_path = run_path / "best_model.pt"

    if not config_path.exists():
        raise FileNotFoundError(f"train_config.json not found in {run_path}")
    if not model_ckpt_path.exists():
        raise FileNotFoundError(f"best_model.pt not found in {run_path}")

    cfg = load_json_config(str(config_path))

    if args.output_path is None:
        args.output_path = str(run_path)

    # Load configurations from train_config.json
    data_prefix = cfg.get("data_prefix", "ukb")
    labels_file = cfg.get("labels_file", "labels.csv")
    no_event_interval_years = cfg.get("no_event_interval_years", 5.0)
    include_no_event = cfg.get("include_no_event_in_uts_target", False)

    target_mode = cfg.get("target_mode", "uts")
    model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
    if model_target_mode not in {"next_token", "all_future"}:
        raise ValueError(
            "train_config.json model_target_mode must be next_token or all_future, "
            f"got {model_target_mode!r}"
        )
    dist_mode_cfg = cfg.get("dist_mode", "exponential")
    readout_name = cfg.get(
        "readout_name", "same_time_group_end" if target_mode == "uts" else "token")
    readout_reduce = cfg.get("readout_reduce", "mean")

    device = torch.device(cfg.get("device", "cuda")
                          if torch.cuda.is_available() else "cpu")
    if device.type == "cuda":
        torch.backends.cudnn.benchmark = True

    print("Loading dataset...")
    dataset = load_sequence_eval_dataset(
        model_target_mode=model_target_mode,
        data_prefix=data_prefix,
        labels_file=labels_file,
        no_event_interval_years=no_event_interval_years,
        include_no_event_in_uts_target=include_no_event,
        min_history_events=int(cfg.get("all_future_min_history_events", 1)),
        min_future_events=int(cfg.get("all_future_min_future_events", 1)),
        extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
    )
    validate_dataset_metadata(dataset, cfg)

    subset, subset_indices = make_eval_subset(dataset, args, cfg)
    print(f"Dataset: {len(dataset)} samples, vocab_size={dataset.vocab_size}")

    loader = DataLoader(
        subset,
        batch_size=int(cfg_get(args, cfg, "batch_size", 128)),
        shuffle=False,
        collate_fn=sequence_eval_collate_fn,
        num_workers=int(cfg_get(args, cfg, "num_workers", 4)),
        pin_memory=device.type == "cuda",
        persistent_workers=int(cfg_get(args, cfg, "num_workers", 4)) > 0,
        prefetch_factor=2 if int(
            cfg_get(args, cfg, "num_workers", 4)) > 0 else None,
    )

    print("Building/loading model...")
    state_dict = load_checkpoint_state_dict(
        str(model_ckpt_path), map_location="cpu")
    dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
    cfg = dict(cfg)
    cfg["dist_mode"] = dist_mode
    cfg["model_target_mode"] = model_target_mode
    print(f"Resolved dist_mode for evaluation: {dist_mode}")
    print(f"Model target mode for AUC: {model_target_mode}")
    print(
        "AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
        "dist_mode affects model loading but is not converted to horizon-specific risk probability."
    )

    model = build_model_from_dataset(args, cfg, dataset).to(device)
    load_model_state(model, str(model_ckpt_path),
                     device, state_dict=state_dict)
    model.eval()

    age_groups = np.arange(
        float(cfg_get(args, cfg, "age_start", 40.0)),
        float(cfg_get(args, cfg, "age_stop", 80.0)),
        float(cfg_get(args, cfg, "age_step", 5.0)),
        dtype=np.float32,
    )
    disease_spec = cfg_get(args, cfg, "diseases_of_interest", None)
    if disease_spec is None:
        disease_spec = cfg.get("disease_tokens", None)
    diseases = parse_int_list(disease_spec)
    first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(
        dataset, subset_indices)
    include_death = bool(cfg_get(args, cfg, "include_death", True))
    exclude_death = bool(cfg_get(args, cfg, "exclude_death", False))
    auc_offsets = make_auc_offsets(args, cfg)

    evaluate_auc_pipeline(
        model=model,
        loader=loader,
        dataset=dataset,
        output_path=args.output_path,
        diseases_of_interest=diseases,
        filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)),
        first_occurrence_by_token=first_occurrence_by_token,
        include_death=include_death,
        exclude_death=exclude_death,
        disease_chunk_size=int(cfg_get(args, cfg, "disease_chunk_size", 0)),
        age_groups=age_groups,
        offsets=auc_offsets,
        device=device,
        model_target_mode=model_target_mode,
        readout_name=readout_name,
        readout_reduce=readout_reduce,
        num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max(
            1, (os.cpu_count() or 2) - 1))),
        use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
        auc_task_chunk_size=int(cfg_get(args, cfg, "auc_task_chunk_size", 0)),
        hidden_cache_dtype=str(
            cfg_get(args, cfg, "hidden_cache_dtype", "float16")),
        logit_batch_size=int(
            cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))),
    )


if __name__ == "__main__":
    main()