DeepHealth/evaluate_auc_v2.py

"""Evaluate landmark fixed-horizon incident disease AUC for DeepHealth.

This script supports DeepHealth fixed-horizon risk scores for exponential,
Weibull, and mixed all-future distributions.

The default horizons are 0.1, 1, 5, and 10 years.  As in Delphi2M, 0.1 years
is reported as the no-gap evaluation.

Landmark querying depends on the model target mode saved in train_config.json:
    - next_token: insert a <NO_EVENT> token at landmark age and read it out;
    - all_future: pass landmark age directly as t_query.
"""
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
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm

from dataset import (
    DISEASE_HISTORY_MODE_TIMED,
    HealthDataset,
    normalize_disease_history_mode,
    transform_disease_history,
)
from delphi2m_auc_report import (
    DEFAULT_DELPHI2M_PERIODS_YEARS,
    build_delphi2m_auc_report,
)
from eval_data import (
    build_first_occurrence_map,
    build_model_from_dataset,
    cfg_get,
    load_json_config,
    load_sequence_eval_dataset,
    resolve_eval_device,
    split_indices,
    validate_training_mode_config,
    validate_dataset_metadata,
)
from model_architectures import resolve_model_architecture
from models import DeepHealth
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX

SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}


def parse_int_list(value: Any) -> Optional[List[int]]:
    if value is None:
        return None
    if isinstance(value, (list, tuple, np.ndarray)):
        return [int(x) for x in value]
    text = str(value).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(value: Any) -> Optional[List[float]]:
    if value is None:
        return None
    if isinstance(value, (list, tuple, np.ndarray)):
        return [float(x) for x in value]
    text = str(value).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_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dict[str, Any]) -> 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()
    if eval_split in {"valid", "validation"}:
        eval_split = "val"

    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,
        "test": test_idx,
        "all": np.arange(len(dataset), dtype=np.int64),
    }
    if eval_split not in split_map:
        raise ValueError(f"Unsupported eval_split={eval_split!r}")
    indices = split_map[eval_split]

    subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
    if subset_size is not None and int(subset_size) > 0:
        indices = indices[: int(subset_size)]
    return np.asarray(indices, dtype=np.int64)


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


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:
        if mode != "weibull":
            print(
                "[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
        return "weibull"
    if has_rho_death_head:
        if mode != "mixed":
            print(
                "[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
        return "mixed"
    if mode == "weibull":
        print(
            "[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
        return "exponential"
    if 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 if mode in {"exponential", "weibull", "mixed"} else "exponential"


def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
    resolve_model_architecture(model.model_architecture, state_dict)
    model.load_state_dict(state_dict, strict=True)


# ---------------------------------------------------------------------------
# DeLong AUC utilities
# ---------------------------------------------------------------------------


def compute_midrank(x: np.ndarray) -> np.ndarray:
    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]:
    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

    v01 = (tz[:, :m] - tx) / n
    v10 = 1.0 - (tz[:, m:] - ty) / m

    if k == 1:
        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:
        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(case_scores: np.ndarray, control_scores: np.ndarray) -> Tuple[float, float]:
    case_scores = np.asarray(case_scores, dtype=np.float64)
    control_scores = np.asarray(control_scores, dtype=np.float64)
    if case_scores.size == 0 or control_scores.size == 0:
        return np.nan, np.nan

    y_true = np.array([1] * len(case_scores) + [0] *
                      len(control_scores), dtype=np.int8)
    y_score = np.concatenate([case_scores, control_scores]).astype(
        np.float64, copy=False)
    order = (-y_true).argsort()

    aucs, cov = fast_delong(y_score[np.newaxis, order], int(y_true.sum()))
    var = float(np.asarray(cov).reshape(-1)[0])
    return float(aucs[0]), var


# ---------------------------------------------------------------------------
# Metadata and token selection
# ---------------------------------------------------------------------------


def _first_existing_column(df: pd.DataFrame, candidates: Sequence[str]) -> Optional[str]:
    for col in candidates:
        if col in df.columns:
            return col
    return None


def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> Dict[int, float]:
    if labels_meta is None or labels_meta.empty or "count" not in labels_meta.columns:
        return {}

    out: Dict[int, float] = {}
    meta = labels_meta.copy()
    count_series = pd.to_numeric(meta["count"], errors="coerce")

    code_col = _first_existing_column(
        meta, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"])
    if code_col is not None:
        for code_text, count in zip(meta[code_col].astype(str).tolist(), count_series.tolist()):
            code = code_text.split()[0].strip()
            if code in dataset.label_code_to_id and pd.notna(count):
                out[int(dataset.label_code_to_id[code])] = float(count)
        if out:
            return out

    if "index" in meta.columns:
        idx = pd.to_numeric(meta["index"], errors="coerce")
        has_no_event = (
            NO_EVENT_IDX in dataset.label_id_to_code
            and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
        )
        if has_no_event:
            idx = idx.where(idx < NO_EVENT_IDX, idx + 1)
        for token, count in zip(idx.tolist(), count_series.tolist()):
            if pd.notna(token) and pd.notna(count):
                out[int(token)] = float(count)
    return out


def _get_death_token_ids(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> List[int]:
    return [int(dataset.vocab_size) - 1]


def select_disease_tokens(
    dataset: HealthDataset,
    labels_meta: Optional[pd.DataFrame],
    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)

    if requested_tokens is not None:
        return sorted(set(int(x) for x in requested_tokens if int(x) in base_set))

    disease_ids = sorted(base)
    if int(filter_min_total) <= 0:
        return disease_ids

    meta_counts = _metadata_count_map(dataset, labels_meta)
    if meta_counts:
        return [token for token in disease_ids if float(meta_counts.get(token, 0.0)) > float(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]))
    return [token for token in disease_ids if int(split_counts.get(token, 0)) > int(filter_min_total)]


# ---------------------------------------------------------------------------
# Landmark dataset
# ---------------------------------------------------------------------------


class LandmarkDataset(Dataset):
    def __init__(
        self,
        dataset: HealthDataset,
        subset_indices: np.ndarray,
        landmark_ages: np.ndarray,
        model_target_mode: str,
        min_history_events: int,
        first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
        death_token_ids: Sequence[int],
        disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
    ) -> None:
        self.dataset = dataset
        self.subset_indices = np.asarray(subset_indices, dtype=np.int64)
        self.landmark_ages = np.asarray(landmark_ages, dtype=np.float32)
        self.model_target_mode = str(model_target_mode).lower()
        if self.model_target_mode not in {"next_token", "all_future"}:
            raise ValueError(
                "model_target_mode must be next_token or all_future, got "
                f"{self.model_target_mode!r}"
            )
        self.disease_history_mode = normalize_disease_history_mode(
            disease_history_mode
        )
        if (
            self.model_target_mode != "all_future"
            and self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED
        ):
            raise ValueError(
                "ordered/set disease history is only supported for all_future models"
            )
        self.min_history_events = int(min_history_events)

        self.first_occurrence_by_token = first_occurrence_by_token
        self.death_token_ids = sorted(set(int(x) for x in death_token_ids))

        rows: List[Dict[str, Any]] = []

        self.patient_followup_end = np.full(
            len(self.subset_indices), -np.inf, dtype=np.float32)
        self.patient_death_time = np.full(
            len(self.subset_indices), np.inf, dtype=np.float32)
        self.patient_sex = np.full(len(self.subset_indices), -1, dtype=np.int8)

        for patient_id, dataset_index in enumerate(self.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:]])

            self.patient_sex[patient_id] = int(s["sex"])
            followup_end = float(np.max(full_time))
            self.patient_followup_end[patient_id] = np.float32(followup_end)

            d_time = np.inf
            for death_token in self.death_token_ids:
                hit = np.where(full_event == int(death_token))[0]
                if hit.size > 0:
                    d_time = min(d_time, float(full_time[int(hit[0])]))
            self.patient_death_time[patient_id] = np.float32(d_time)

            for landmark_age in self.landmark_ages.tolist():
                landmark_age = float(landmark_age)

                if not (followup_end > landmark_age):
                    continue
                if not (float(self.patient_death_time[patient_id]) > landmark_age):
                    continue

                prefix_mask = full_time <= landmark_age
                if not np.any(prefix_mask):
                    continue

                prefix_events = full_event[prefix_mask]
                prefix_times = full_time[prefix_mask]

                valid_history_mask = ~np.isin(prefix_events, np.array(
                    [PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64))
                if valid_history_mask.sum() < self.min_history_events:
                    continue

                if self.model_target_mode == "next_token":
                    event_seq_landmark = np.concatenate(
                        [
                            prefix_events.astype(np.int64, copy=False),
                            np.array([NO_EVENT_IDX], dtype=np.int64),
                        ]
                    )
                    time_seq_landmark = np.concatenate(
                        [
                            prefix_times.astype(np.float32, copy=False),
                            np.array([np.float32(landmark_age)], dtype=np.float32),
                        ]
                    )
                    time_seq_landmark[-1] = np.nextafter(
                        np.float32(landmark_age),
                        np.float32(np.inf),
                        dtype=np.float32,
                    )
                    landmark_pos = int(len(event_seq_landmark) - 1)
                    readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
                    readout_mask[-1] = True
                else:
                    (
                        event_seq_landmark,
                        time_seq_landmark,
                        model_t_query,
                    ) = transform_disease_history(
                        event_seq=prefix_events,
                        actual_time_seq=prefix_times,
                        actual_t_query=landmark_age,
                        disease_history_mode=self.disease_history_mode,
                    )
                    landmark_pos = int(len(event_seq_landmark) - 1)
                    readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)

                rows.append(
                    {
                        "patient_id": int(patient_id),
                        "dataset_index": int(dataset_index),
                        "sex": int(s["sex"]),
                        "landmark_age": np.float32(landmark_age),
                        "followup_end_time": np.float32(followup_end),
                        "death_time": np.float32(self.patient_death_time[patient_id]),
                        "landmark_pos": landmark_pos,
                        "t_query": (
                            np.float32(landmark_age)
                            if self.model_target_mode == "next_token"
                            else model_t_query
                        ),
                        "event_seq": event_seq_landmark,
                        "time_seq": time_seq_landmark,
                        "readout_mask": readout_mask,
                        "other_type": np.asarray(s["other_type"], dtype=np.int64),
                        "other_value": np.asarray(s["other_value"], dtype=np.float32),
                        "other_value_kind": np.asarray(s["other_value_kind"], dtype=np.int64),
                        "other_time": np.asarray(s["other_time"], dtype=np.float32),
                    }
                )

        if not rows:
            raise RuntimeError(
                "No eligible landmark query samples were produced. Check eval split, landmark ages, and min_history_events."
            )

        self.rows = rows

    def __len__(self) -> int:
        return len(self.rows)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        s = self.rows[idx]
        return {
            "event_seq": torch.from_numpy(s["event_seq"]).long(),
            "time_seq": torch.from_numpy(s["time_seq"]).float(),
            "readout_mask": torch.from_numpy(s["readout_mask"]),
            "sex": torch.tensor(s["sex"], dtype=torch.long),
            "other_type": torch.from_numpy(s["other_type"]).long(),
            "other_value": torch.from_numpy(s["other_value"]).float(),
            "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
            "other_time": torch.from_numpy(s["other_time"]).float(),
            "landmark_pos": torch.tensor(s["landmark_pos"], dtype=torch.long),
            "t_query": torch.tensor(float(s["t_query"]), dtype=torch.float32),
            "patient_id": torch.tensor(s["patient_id"], dtype=torch.long),
            "landmark_age": torch.tensor(float(s["landmark_age"]), dtype=torch.float32),
            "followup_end_time": torch.tensor(float(s["followup_end_time"]), dtype=torch.float32),
            "death_time": torch.tensor(float(s["death_time"]), dtype=torch.float32),
        }


def collate_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
    event_seq = pad_sequence(
        [x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX)
    time_seq = pad_sequence([x["time_seq"]
                            for x in batch], batch_first=True, padding_value=0.0)
    readout_mask = pad_sequence(
        [x["readout_mask"] for x in batch], batch_first=True, padding_value=False)
    other_type = pad_sequence(
        [x["other_type"] for x in batch], batch_first=True, padding_value=0)
    other_value = pad_sequence(
        [x["other_value"] for x in batch], batch_first=True, padding_value=0.0)
    other_value_kind = pad_sequence(
        [x["other_value_kind"] for x in batch], batch_first=True, padding_value=0)
    other_time = pad_sequence(
        [x["other_time"] for x in batch], batch_first=True, padding_value=0.0)

    return {
        "event_seq": event_seq,
        "time_seq": time_seq,
        "padding_mask": event_seq > PAD_IDX,
        "readout_mask": readout_mask,
        "sex": torch.stack([x["sex"] for x in batch]),
        "other_type": other_type,
        "other_value": other_value,
        "other_value_kind": other_value_kind,
        "other_time": other_time,
        "landmark_pos": torch.stack([x["landmark_pos"] for x in batch]),
        "t_query": torch.stack([x["t_query"] for x in batch]),
        "patient_id": torch.stack([x["patient_id"] for x in batch]),
        "landmark_age": torch.stack([x["landmark_age"] for x in batch]),
        "followup_end_time": torch.stack([x["followup_end_time"] for x in batch]),
        "death_time": torch.stack([x["death_time"] for x in batch]),
    }


def _numpy_hidden_dtype(name: str) -> np.dtype:
    key = str(name).lower()
    if key in {"float16", "fp16", "half", "bfloat16", "bf16"}:
        return np.float16
    if key 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_landmark_hidden(
    model: DeepHealth,
    loader: DataLoader,
    device: torch.device,
    model_target_mode: str,
    use_amp: bool,
    hidden_cache_dtype: str,
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
    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}"
        )

    hidden_parts: List[np.ndarray] = []
    arrays = {
        "patient_id": [],
        "sex": [],
        "landmark_age": [],
        "followup_end_time": [],
        "death_time": [],
    }

    out_dtype = _numpy_hidden_dtype(hidden_cache_dtype)
    amp_enabled = bool(use_amp and device.type == "cuda")

    for batch in tqdm(loader, desc="Landmark 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()
        }

        amp_ctx = (
            torch.autocast(device_type=device.type, dtype=torch.float16)
            if amp_enabled
            else contextlib.nullcontext()
        )
        with amp_ctx:
            if model_target_mode == "all_future":
                landmark_hidden = model(
                    event_seq=batch_dev["event_seq"],
                    time_seq=batch_dev["time_seq"],
                    sex=batch_dev["sex"],
                    padding_mask=batch_dev["padding_mask"],
                    t_query=batch_dev["t_query"],
                    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"],
                )
            else:
                hidden = model(
                    event_seq=batch_dev["event_seq"],
                    time_seq=batch_dev["time_seq"],
                    sex=batch_dev["sex"],
                    padding_mask=batch_dev["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"],
                )
                landmark_hidden = hidden.gather(
                    1,
                    batch_dev["landmark_pos"].long()[:, None, None].expand(
                        -1, 1, hidden.shape[-1]
                    ),
                ).squeeze(1)

        hidden_parts.append(landmark_hidden.detach(
        ).cpu().numpy().astype(out_dtype, copy=False))
        arrays["patient_id"].append(
            batch["patient_id"].cpu().numpy().astype(np.int32, copy=False))
        arrays["sex"].append(
            batch["sex"].cpu().numpy().astype(np.int8, copy=False))
        arrays["landmark_age"].append(
            batch["landmark_age"].cpu().numpy().astype(np.float32, copy=False))
        arrays["followup_end_time"].append(
            batch["followup_end_time"].cpu().numpy().astype(np.float32, copy=False))
        arrays["death_time"].append(
            batch["death_time"].cpu().numpy().astype(np.float32, copy=False))

    hidden_all = np.concatenate(hidden_parts, axis=0)
    row_arrays = {
        k: np.concatenate(v, axis=0) for k, v in arrays.items()
    }
    return hidden_all, row_arrays


@torch.inference_mode()
def project_distribution_chunk(
    model: DeepHealth,
    hidden_all: np.ndarray,
    disease_ids: Sequence[int],
    dist_mode: str,
    device: torch.device,
    logit_batch_size: int,
    use_amp: bool,
) -> Tuple[np.ndarray, Optional[np.ndarray]]:
    n = int(hidden_all.shape[0])
    logit_batch_size = max(1, int(logit_batch_size))
    disease_ids = [int(x) for x in disease_ids]
    dist_mode = str(dist_mode).lower()

    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)
    rho_weight = None
    rho_bias = None
    death_rho_weight = None
    death_rho_bias = None
    mixed_death_cols: List[int] = []
    death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))

    if dist_mode == "weibull":
        rho_weight = model.rho_head.weight[disease_ids].detach().to(
            device=device, dtype=compute_dtype)
        rho_bias = model.rho_head.bias[disease_ids].detach().to(
            device=device, dtype=compute_dtype)
    elif dist_mode == "mixed":
        mixed_death_cols = [j for j, token in enumerate(disease_ids)
                            if int(token) == death_idx]
        if mixed_death_cols:
            death_rho_weight = model.rho_death_head.weight.detach().to(
                device=device, dtype=compute_dtype)
            death_rho_bias = model.rho_death_head.bias.detach().to(
                device=device, dtype=compute_dtype)

    out_parts: List[np.ndarray] = []
    rho_parts: List[np.ndarray] = []
    for start in tqdm(range(0, n, logit_batch_size), desc="Risk 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
        rho = None
        if dist_mode == "weibull":
            assert rho_weight is not None and rho_bias is not None
            rho = F.softplus(torch.matmul(h, rho_weight.t()) + rho_bias) + 1e-6
        elif dist_mode == "mixed" and mixed_death_cols:
            assert death_rho_weight is not None and death_rho_bias is not None
            rho = torch.ones_like(logits)
            death_rho = F.softplus(
                torch.matmul(h, death_rho_weight.t()).squeeze(-1) + death_rho_bias.squeeze(0)
            ) + 1e-6
            for col in mixed_death_cols:
                rho[:, int(col)] = death_rho

        out_parts.append(logits.float().cpu(
        ).numpy().astype(np.float32, copy=False))
        if rho is not None:
            rho_parts.append(rho.float().cpu(
            ).numpy().astype(np.float32, copy=False))
        del h, logits, rho
    logits_all = np.concatenate(out_parts, axis=0)
    rho_all = np.concatenate(rho_parts, axis=0) if rho_parts else None
    return logits_all, rho_all


# ---------------------------------------------------------------------------
# Parallel AUC workers
# ---------------------------------------------------------------------------


_WORKER: Dict[str, Any] = {}


def _init_worker(
    disease_ids: np.ndarray,
    score_chunk: np.ndarray,
    rho_chunk: Optional[np.ndarray],
    row_patient_id: np.ndarray,
    row_sex: np.ndarray,
    row_landmark_age: np.ndarray,
    row_followup_end: np.ndarray,
    row_death_time: np.ndarray,
    first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
    patient_count: int,
    horizons: np.ndarray,
    min_cases: int,
    exclude_death_competing: bool,
    death_token_ids: np.ndarray,
    dist_mode: str,
    model_death_idx: int,
) -> None:
    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(
        {
            "disease_ids": np.asarray(disease_ids, dtype=np.int64),
            "score_chunk": np.asarray(score_chunk, dtype=np.float32),
            "rho_chunk": None if rho_chunk is None else np.asarray(rho_chunk, dtype=np.float32),
            "row_patient_id": np.asarray(row_patient_id, dtype=np.int32),
            "row_sex": np.asarray(row_sex, dtype=np.int8),
            "row_landmark_age": np.asarray(row_landmark_age, dtype=np.float32),
            "row_followup_end": np.asarray(row_followup_end, dtype=np.float32),
            "row_death_time": np.asarray(row_death_time, dtype=np.float32),
            "first_occurrence_by_token": first_occurrence_by_token,
            "patient_count": int(patient_count),
            "horizons": np.asarray(horizons, dtype=np.float32),
            "min_cases": int(min_cases),
            "exclude_death_competing": bool(exclude_death_competing),
            "death_token_ids": set(int(x) for x in np.asarray(death_token_ids, dtype=np.int64).tolist()),
            "dist_mode": str(dist_mode).lower(),
            "model_death_idx": int(model_death_idx),
            "first_time_cache": {},
        }
    )


def _first_time_by_patient(token: int) -> np.ndarray:
    cache = _WORKER["first_time_cache"]
    if int(token) in cache:
        return cache[int(token)]

    arr = np.full(int(_WORKER["patient_count"]), np.inf, dtype=np.float32)
    pairs = _WORKER["first_occurrence_by_token"].get(int(token))
    if pairs is not None:
        p, t = pairs
        arr[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32)
    cache[int(token)] = arr
    return arr


def _score_to_probability(
    logits: np.ndarray,
    rho: Optional[np.ndarray],
    score_mode: str,
    horizon: float,
    dist_mode: str,
    token: int,
    death_idx: int,
) -> np.ndarray:
    if score_mode == "eta":
        return logits.astype(np.float64, copy=False)
    rate = np.log1p(np.exp(-np.abs(logits))) + np.maximum(logits, 0.0)
    rate = rate + np.float32(1e-8)
    dist_mode = str(dist_mode).lower()
    if dist_mode == "weibull":
        if rho is None:
            raise RuntimeError("Weibull risk scoring requires rho parameters.")
        exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
        return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
    if dist_mode == "mixed" and int(token) == int(death_idx):
        if rho is None:
            raise RuntimeError("Mixed death risk scoring requires death rho parameters.")
        exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
        return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
    return (-np.expm1(-rate * np.float32(horizon))).astype(np.float64, copy=False)


def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
    j, token, score_mode = task

    token = int(token)
    row_patient_id = _WORKER["row_patient_id"]
    row_sex = _WORKER["row_sex"]
    row_landmark_age = _WORKER["row_landmark_age"]
    row_followup_end = _WORKER["row_followup_end"]
    row_death_time = _WORKER["row_death_time"]
    logits_token = _WORKER["score_chunk"][:, int(j)]
    rho_chunk = _WORKER["rho_chunk"]
    rho_token = None if rho_chunk is None else rho_chunk[:, int(j)]
    dist_mode = _WORKER["dist_mode"]
    model_death_idx = int(_WORKER["model_death_idx"])

    first_time_patient = _first_time_by_patient(token)
    is_death_target = token in _WORKER["death_token_ids"]
    horizons = _WORKER["horizons"]

    out_rows: List[Dict[str, Any]] = []

    for sex_value, sex_name in [(0, "female"), (1, "male")]:
        sex_mask_rows = row_sex == sex_value
        if not np.any(sex_mask_rows):
            continue

        lm_values = np.unique(row_landmark_age[sex_mask_rows])
        for a in lm_values.tolist():
            a = float(a)
            stratum_mask = sex_mask_rows & np.isclose(
                row_landmark_age, np.float32(a))
            if not np.any(stratum_mask):
                continue

            idx = np.flatnonzero(stratum_mask)
            pid = row_patient_id[idx]
            followup = row_followup_end[idx]
            death_time = row_death_time[idx]
            first_time = first_time_patient[pid]

            prevalent = first_time <= np.float32(a)
            if np.all(prevalent):
                continue

            for horizon in horizons.tolist():
                horizon = float(horizon)
                h_end = np.float32(a + horizon)

                cases = (first_time > np.float32(a)) & (first_time <= h_end)
                controls = (~prevalent) & (
                    ((~np.isinf(first_time)) & (first_time > h_end))
                    | (np.isinf(first_time) & (followup >= h_end))
                )

                if bool(_WORKER["exclude_death_competing"]) and (not is_death_target):
                    death_in_window = (death_time > np.float32(a)) & (
                        death_time <= h_end)
                    death_before_disease = death_time < first_time
                    competing_death = death_in_window & death_before_disease
                    controls = controls & (~competing_death)

                case_idx = np.flatnonzero(cases)
                control_idx = np.flatnonzero(controls)
                if case_idx.size < int(_WORKER["min_cases"]) or control_idx.size == 0:
                    continue

                case_scores = _score_to_probability(
                    logits_token[idx[case_idx]],
                    None if rho_token is None else rho_token[idx[case_idx]],
                    score_mode=score_mode,
                    horizon=horizon,
                    dist_mode=dist_mode,
                    token=token,
                    death_idx=model_death_idx,
                )
                control_scores = _score_to_probability(
                    logits_token[idx[control_idx]],
                    None if rho_token is None else rho_token[idx[control_idx]],
                    score_mode=score_mode,
                    horizon=horizon,
                    dist_mode=dist_mode,
                    token=token,
                    death_idx=model_death_idx,
                )

                auc, auc_var = get_auc_delong_var(case_scores, control_scores)
                if np.isnan(auc) or np.isnan(auc_var):
                    continue

                out_rows.append(
                    {
                        "token": token,
                        "sex": sex_name,
                        "landmark_age": float(a),
                        "horizon": float(horizon),
                        "auc": float(auc),
                        "auc_delong": float(auc),
                        "auc_variance_delong": float(auc_var),
                        "n_diseased": int(case_scores.size),
                        "n_healthy": int(control_scores.size),
                    }
                )

    return out_rows


def _token_task_block(tasks: Sequence[Tuple[int, int, str]]) -> List[Dict[str, Any]]:
    out: List[Dict[str, Any]] = []
    for t in tasks:
        out.extend(_eval_token(t))
    return out


def _split_tasks(tasks: Sequence[Tuple[int, int, str]], workers: int, chunk_size: int) -> List[List[Tuple[int, int, str]]]:
    if not tasks:
        return []
    if int(chunk_size) <= 0:
        chunk_size = max(1, math.ceil(len(tasks) / max(1, workers * 4)))
    return [list(tasks[i:i + chunk_size]) for i in range(0, len(tasks), chunk_size)]


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


def evaluate_landmark_auc(
    model: DeepHealth,
    loader: DataLoader,
    landmark_dataset: LandmarkDataset,
    output_path: Path,
    disease_ids: Sequence[int],
    disease_chunk_size: int,
    score_mode: str,
    dist_mode: str,
    horizons: np.ndarray,
    device: torch.device,
    model_target_mode: str,
    num_workers_auc: int,
    auc_task_chunk_size: int,
    min_cases: int,
    exclude_death_competing: bool,
    use_amp: bool,
    hidden_cache_dtype: str,
    logit_batch_size: int,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    model.eval().to(device)

    hidden_all, row_arrays = infer_landmark_hidden(
        model=model,
        loader=loader,
        device=device,
        model_target_mode=model_target_mode,
        use_amp=use_amp,
        hidden_cache_dtype=hidden_cache_dtype,
    )
    print(
        f"Cached landmark hidden: shape={hidden_all.shape}, dtype={hidden_all.dtype}")

    disease_ids = [int(x) for x in disease_ids]
    if disease_chunk_size <= 0:
        disease_chunk_size = len(disease_ids)
    chunks = np.array_split(np.asarray(disease_ids, dtype=np.int64), math.ceil(
        len(disease_ids) / disease_chunk_size))

    all_rows: List[Dict[str, Any]] = []
    for chunk_idx, chunk in enumerate(tqdm(chunks, desc="Disease chunks", dynamic_ncols=True)):
        chunk_ids = chunk.tolist()
        logits_chunk, rho_chunk = project_distribution_chunk(
            model=model,
            hidden_all=hidden_all,
            disease_ids=chunk_ids,
            dist_mode=dist_mode,
            device=device,
            logit_batch_size=logit_batch_size,
            use_amp=use_amp,
        )

        tasks = [(j, int(token), score_mode)
                 for j, token in enumerate(chunk_ids)]
        workers = max(1, min(int(num_workers_auc), len(tasks)))

        if workers <= 1:
            _init_worker(
                disease_ids=np.asarray(chunk_ids, dtype=np.int64),
                score_chunk=logits_chunk,
                rho_chunk=rho_chunk,
                row_patient_id=row_arrays["patient_id"],
                row_sex=row_arrays["sex"],
                row_landmark_age=row_arrays["landmark_age"],
                row_followup_end=row_arrays["followup_end_time"],
                row_death_time=row_arrays["death_time"],
                first_occurrence_by_token=landmark_dataset.first_occurrence_by_token,
                patient_count=len(landmark_dataset.subset_indices),
                horizons=horizons,
                min_cases=min_cases,
                exclude_death_competing=exclude_death_competing,
                death_token_ids=np.asarray(
                    landmark_dataset.death_token_ids, dtype=np.int64),
                dist_mode=dist_mode,
                model_death_idx=int(getattr(
                    model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
            )
            nested = [_eval_token(t) for t in tqdm(
                tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)]
        else:
            ctx = mp.get_context("spawn")
            blocks = _split_tasks(tasks, workers, auc_task_chunk_size)
            with ProcessPoolExecutor(
                max_workers=workers,
                mp_context=ctx,
                initializer=_init_worker,
                initargs=(
                    np.asarray(chunk_ids, dtype=np.int64),
                    logits_chunk,
                    rho_chunk,
                    row_arrays["patient_id"],
                    row_arrays["sex"],
                    row_arrays["landmark_age"],
                    row_arrays["followup_end_time"],
                    row_arrays["death_time"],
                    landmark_dataset.first_occurrence_by_token,
                    len(landmark_dataset.subset_indices),
                    horizons,
                    min_cases,
                    exclude_death_competing,
                    np.asarray(landmark_dataset.death_token_ids,
                               dtype=np.int64),
                    dist_mode,
                    int(getattr(
                        model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
                ),
            ) as ex:
                nested = list(
                    tqdm(
                        ex.map(_token_task_block, blocks),
                        total=len(blocks),
                        desc=f"AUC chunk {chunk_idx}",
                        leave=False,
                        dynamic_ncols=True,
                    )
                )

        for rows in nested:
            for r in rows:
                r["disease_chunk_idx"] = int(chunk_idx)
                all_rows.append(r)

        del logits_chunk, rho_chunk

    if not all_rows:
        raise RuntimeError(
            "No AUC rows produced. Check landmark ages, horizons, min_cases, follow-up, or disease token selection."
        )

    df_unpooled = pd.DataFrame(all_rows)
    df_unpooled["label_code"] = df_unpooled["token"].map(
        landmark_dataset.dataset.label_id_to_code)

    print(
        "Building Delphi2M-style report: mean AUC across landmark-age "
        "strata, reported separately for Female and Male."
    )
    df_report = build_delphi2m_auc_report(
        df_unpooled,
        period_col="horizon",
    )

    output_path.mkdir(parents=True, exist_ok=True)
    report_path = output_path / "df_auc_landmark_delphi2m_report.csv"
    df_report.to_csv(report_path, index=False)
    print(f"Saved Delphi2M-style landmark AUC report: {report_path}")

    return df_unpooled, df_report


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Evaluate DeepHealth landmark fixed-horizon incident disease AUC")

    parser.add_argument("--run_path", type=str, required=True)
    parser.add_argument("--output_path", type=str, default=None)
    parser.add_argument("--eval_split", type=str, default="test",
                        choices=["train", "val", "valid", "validation", "test", "all"])
    parser.add_argument("--dataset_subset_size", type=int, default=None)

    parser.add_argument("--batch_size", type=int, default=None)
    parser.add_argument("--num_workers", type=int, default=None)
    parser.add_argument("--device", type=str, default=None,
                        help="Evaluation device, e.g. cpu, cuda, cuda:1. Defaults to cuda if available, else cpu.")
    parser.add_argument("--num_workers_auc", type=int, default=None)
    parser.add_argument("--auc_task_chunk_size", type=int, default=None)
    parser.add_argument("--disease_chunk_size", type=int, default=None)
    parser.add_argument("--filter_min_total", type=int, default=None)
    parser.add_argument("--labels_meta_path", type=str, default=None)

    parser.add_argument("--landmark_start", type=float, default=None)
    parser.add_argument("--landmark_stop", type=float, default=None)
    parser.add_argument("--landmark_step", type=float, default=None)
    parser.add_argument(
        "--horizons",
        type=str,
        default=None,
        help="Comma-separated horizons in years; defaults to 0.1,1,5,10, where 0.1 is Delphi2M no gap.",
    )

    parser.add_argument("--min_cases", type=int, default=None)
    parser.add_argument("--min_history_events", type=int, default=None)

    parser.add_argument("--score_mode", type=str,
                        default=None, choices=["risk", "eta"])
    parser.add_argument("--exclude_death_in_window_without_disease",
                        action=argparse.BooleanOptionalAction, default=None)
    parser.add_argument(
        "--use_amp", action=argparse.BooleanOptionalAction, default=None)
    parser.add_argument("--logit_batch_size", type=int, default=None)
    parser.add_argument("--hidden_cache_dtype", type=str,
                        default=None, choices=["float16", "float32"])

    args = parser.parse_args()

    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(config_path)
    validate_training_mode_config(cfg)

    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)

    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 = str(cfg.get("dist_mode", "exponential"))
    disease_history_mode = normalize_disease_history_mode(
        cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
    )

    output_path = Path(
        cfg_get(args, cfg, "output_path", None)
        or cfg.get("output_path", None)
        or str(run_path)
    )
    output_path.mkdir(parents=True, exist_ok=True)

    labels_meta_path = cfg_get(args, cfg, "labels_meta_path", None)
    if labels_meta_path is None:
        labels_meta_path = cfg.get(
            "labels_meta_path", "delphi_labels_chapters_colours_icd.csv")
    labels_meta = None
    if labels_meta_path is not None and Path(str(labels_meta_path)).exists():
        labels_meta = pd.read_csv(str(labels_meta_path))

    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=float(no_event_interval_years),
        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)),
        disease_history_mode=disease_history_mode,
    )
    validate_dataset_metadata(dataset, cfg)

    has_no_event = (
        NO_EVENT_IDX in dataset.label_id_to_code
        and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>"
        and dataset.vocab_size > NO_EVENT_IDX
    )
    if model_target_mode == "next_token" and not has_no_event:
        print(
            "[SKIP] This checkpoint/run does not support <NO_EVENT> imputation. "
            "Landmark AUC requires inserting a <NO_EVENT> query token. "
            "Please evaluate only runs trained with the current no-event vocabulary."
        )
        return

    subset_indices = make_eval_indices(dataset, args, cfg)

    first_occurrence_by_token = build_first_occurrence_map(
        dataset, subset_indices
    )

    disease_requested = parse_int_list(
        cfg_get(args, cfg, "diseases_of_interest", None))
    disease_ids = select_disease_tokens(
        dataset=dataset,
        labels_meta=labels_meta,
        requested_tokens=disease_requested,
        filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)),
        first_occurrence_by_token=first_occurrence_by_token,
    )

    if not disease_ids:
        raise RuntimeError("No disease tokens selected after filtering.")

    landmark_start = float(cfg_get(args, cfg, "landmark_start", 40.0))
    landmark_stop = float(cfg_get(args, cfg, "landmark_stop", 80.0))
    landmark_step = float(cfg_get(args, cfg, "landmark_step", 5.0))
    if landmark_step <= 0:
        raise ValueError("landmark_step must be > 0")
    landmark_ages = np.arange(
        landmark_start, landmark_stop, landmark_step, dtype=np.float32)
    if landmark_ages.size == 0:
        raise ValueError(
            "Landmark ages are empty. Check landmark_start/landmark_stop/landmark_step.")

    horizons = np.asarray(
        parse_float_list(
            cfg_get(args, cfg, "horizons", "0.1,1,5,10")
        ) or list(DEFAULT_DELPHI2M_PERIODS_YEARS),
        dtype=np.float32,
    )
    if horizons.size == 0:
        raise ValueError("horizons must contain at least one value")

    score_mode = str(cfg_get(args, cfg, "score_mode", "risk")).lower()
    if score_mode not in {"risk", "eta"}:
        raise ValueError("score_mode must be 'risk' or 'eta'")

    state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu")
    dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
    if dist_mode not in {"exponential", "weibull", "mixed"}:
        raise ValueError(
            f"Unsupported dist_mode={dist_mode!r}; expected exponential, weibull, or mixed."
        )

    if score_mode == "eta":
        print(
            "WARNING: eta diagnostic score is not horizon-specific risk and does not use dist_mode-specific rho parameters.")

    cfg_model = dict(cfg)
    cfg_model["dist_mode"] = dist_mode
    model_architecture = resolve_model_architecture(cfg_model, state_dict)
    cfg_model["model_architecture"] = model_architecture
    print(f"Resolved model architecture: {model_architecture}")

    device = resolve_eval_device(args.device)
    if device.type == "cuda":
        torch.backends.cudnn.benchmark = True

    model = build_model_from_dataset(
        args, cfg_model, dataset, state_dict=state_dict
    ).to(device)

    if (
        model_target_mode == "next_token"
        and (
            model.token_embedding.num_embeddings <= NO_EVENT_IDX
            or model.risk_head.out_features <= NO_EVENT_IDX
        )
    ):
        raise RuntimeError(
            "Model vocabulary does not include <NO_EVENT> token index. "
            "Checkpoint/model shape is incompatible with no-event landmark querying."
        )

    try:
        load_model_state(model, state_dict)
    except RuntimeError as exc:
        raise RuntimeError(
            "Checkpoint vocabulary shape is incompatible with the dataset/model setup. "
            "Please ensure this run was trained with the same special-token vocabulary and labels file."
        ) from exc

    if model.token_embedding.num_embeddings != dataset.vocab_size or model.risk_head.out_features != dataset.vocab_size:
        raise RuntimeError(
            "Checkpoint/model vocabulary shape mismatch with dataset vocabulary. "
            "Please use a checkpoint trained with the same no-event vocabulary configuration."
        )

    death_token_ids = _get_death_token_ids(dataset, labels_meta)
    min_history_events = int(cfg_get(args, cfg, "min_history_events", 1))
    landmark_dataset = LandmarkDataset(
        dataset=dataset,
        subset_indices=subset_indices,
        landmark_ages=landmark_ages,
        model_target_mode=model_target_mode,
        min_history_events=min_history_events,
        first_occurrence_by_token=first_occurrence_by_token,
        death_token_ids=death_token_ids,
        disease_history_mode=disease_history_mode,
    )

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

    eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
    if eval_split in {"valid", "validation"}:
        eval_split = "val"

    landmark_query_mode = (
        "insert_no_event_token"
        if model_target_mode == "next_token"
        else "direct_t_query"
    )
    num_workers_auc = int(
        cfg_get(args, cfg, "num_workers_auc", max(1, (os.cpu_count() or 2) - 1)))
    auc_task_chunk_size = int(cfg_get(args, cfg, "auc_task_chunk_size", 0))
    disease_chunk_size = int(cfg_get(args, cfg, "disease_chunk_size", 0))
    min_cases = int(cfg_get(args, cfg, "min_cases", 2))
    exclude_death_competing = bool(
        cfg_get(args, cfg, "exclude_death_in_window_without_disease", True))
    use_amp = bool(cfg_get(args, cfg, "use_amp", False))
    hidden_cache_dtype = str(
        cfg_get(args, cfg, "hidden_cache_dtype", "float16"))
    logit_batch_size = int(cfg_get(args, cfg, "logit_batch_size", batch_size))

    print(f"Eval split: {eval_split}")
    print(f"Number of selected patients: {len(subset_indices)}")
    print(f"No-event support: {bool(has_no_event)}")
    print(f"Model target mode: {model_target_mode}")
    print(f"Disease history mode: {disease_history_mode}")
    print(f"Landmark query mode: {landmark_query_mode}")
    print(
        "Landmark token mode: no_event"
        if model_target_mode == "next_token"
        else "Landmark token mode: none"
    )
    print(f"Dist mode: {dist_mode}")
    print(f"Score mode: {score_mode}")
    print(f"Landmark ages: {landmark_ages.tolist()}")
    print(f"Horizons: {horizons.tolist()}")
    print(f"Number of landmark query samples: {len(landmark_dataset)}")
    print(f"Number of disease tokens: {len(disease_ids)}")
    print(f"AUC workers: {num_workers_auc}")
    print(f"Output path: {output_path}")

    evaluate_landmark_auc(
        model=model,
        loader=loader,
        landmark_dataset=landmark_dataset,
        output_path=output_path,
        disease_ids=disease_ids,
        disease_chunk_size=disease_chunk_size,
        score_mode=score_mode,
        dist_mode=dist_mode,
        horizons=horizons,
        device=device,
        model_target_mode=model_target_mode,
        num_workers_auc=num_workers_auc,
        auc_task_chunk_size=auc_task_chunk_size,
        min_cases=min_cases,
        exclude_death_competing=exclude_death_competing,
        use_amp=use_amp,
        hidden_cache_dtype=hidden_cache_dtype,
        logit_batch_size=logit_batch_size,
    )


if __name__ == "__main__":
    main()