"""Evaluate disease AUC at date of assessment (DOA).

Cases are patients whose first occurrence of a disease is after DOA and within
the requested horizon. Controls are patients who never have that disease in the
full observed record. Patients prevalent at/before DOA or incident after the
horizon are not used for that disease-horizon AUC.

The script adapts automatically to checkpoint target mode:
  - next_token: use the CHECKUP token position at DOA;
  - all_future: query the model directly with t_query=DOA. The history includes
    the CHECKUP token at DOA.
"""
from __future__ import annotations

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

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

from dataset import _ExpoBaseDataset
from evaluate_auc_v2 import (
    build_metadata_for_merge,
    build_model_from_dataset,
    get_auc_delong_var,
    load_checkpoint_state_dict,
    load_json_config,
    load_model_state,
    parse_float_list,
    parse_int_list,
    project_distribution_chunk,
    resolve_dist_mode_for_checkpoint,
    select_disease_tokens,
    validate_dataset_metadata,
    _score_to_probability,
)
from readouts import build_readout
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX, DAYS_PER_YEAR


SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}


def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: Any) -> Any:
    value = getattr(args, name, None)
    if value is not None:
        return value
    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 = float(train_ratio) + float(val_ratio) + float(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}")
    indices = np.random.RandomState(int(seed)).permutation(int(n))
    n_train = int(n * train_ratio)
    n_val = int(n * val_ratio)
    return indices[:n_train], indices[n_train:n_train + n_val], indices[n_train + n_val:]


def make_eval_indices(
    dataset: Dataset,
    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 = np.asarray(split_map[eval_split], dtype=np.int64)
    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 indices


def subset_first_occurrence_map(
    first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
    selected_patient_ids: np.ndarray,
) -> Dict[int, Tuple[np.ndarray, np.ndarray]]:
    selected = set(int(x) for x in np.asarray(selected_patient_ids, dtype=np.int64).tolist())
    out: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
    for token, pairs in first_occurrence_by_token.items():
        p, t = pairs
        keep = np.array([int(x) in selected for x in p], dtype=bool)
        if np.any(keep):
            out[int(token)] = (
                np.asarray(p, dtype=np.int32)[keep],
                np.asarray(t, dtype=np.float32)[keep],
            )
    return out


class DOAStatusDataset(_ExpoBaseDataset):
    def __init__(
        self,
        data_prefix: str,
        labels_file: str,
        model_target_mode: str,
        extra_info_types: Iterable[int] | None = None,
    ) -> None:
        super().__init__(
            data_prefix=data_prefix,
            labels_file=labels_file,
            no_event_interval_years=5.0,
            include_no_event_in_uts_target=False,
            extra_info_types=extra_info_types,
        )
        self.model_target_mode = str(model_target_mode).lower()
        if self.model_target_mode not in {"next_token", "all_future"}:
            raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")

        self.records: List[Dict[str, Any]] = []
        self.first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}

        unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True)
        ends = np.concatenate([starts[1:], [len(self.event_data)]])
        first_lists: Dict[int, List[Tuple[int, float]]] = {}

        for eid_raw, start, end in zip(unique_eids, starts, ends):
            eid = int(eid_raw)
            rows = self.event_data[start:end]
            checkup_rows = rows[rows[:, 2].astype(np.int64) == CHECKUP_IDX]
            if len(checkup_rows) == 0:
                continue

            features = self._split_features(eid)
            if features is None:
                continue

            doa_days = float(np.min(checkup_rows[:, 1].astype(np.float32)))
            doa_years = np.float32(doa_days / DAYS_PER_YEAR)

            raw_times = rows[:, 1].astype(np.float32) / DAYS_PER_YEAR
            raw_labels = rows[:, 2].astype(np.int64)
            shifted_labels = np.where(
                raw_labels >= NO_EVENT_IDX,
                raw_labels + 1,
                raw_labels,
            ).astype(np.int64)
            order = np.lexsort((shifted_labels, raw_times))
            event_times = raw_times[order].astype(np.float32)
            event_labels = shifted_labels[order].astype(np.int64)

            disease_mask = event_labels != CHECKUP_IDX
            disease_times = event_times[disease_mask]
            disease_labels = event_labels[disease_mask]

            patient_id = len(self.records)
            for token in np.unique(disease_labels).tolist():
                token = int(token)
                if token in SPECIAL_TOKENS:
                    continue
                hit = np.where(disease_labels == token)[0]
                if hit.size:
                    first_lists.setdefault(token, []).append(
                        (patient_id, float(disease_times[int(hit[0])]))
                    )

            hist = event_times <= doa_years
            hist_events = event_labels[hist]
            hist_times = event_times[hist]

            if self.model_target_mode == "next_token":
                checkup_at_doa = (
                    (hist_events == CHECKUP_IDX)
                    & np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6)
                )
                if not np.any(checkup_at_doa):
                    raise RuntimeError(f"Missing CHECKUP token at DOA for eid={eid}")
                event_seq = hist_events
                time_seq = hist_times
                readout_pos = int(np.where(checkup_at_doa)[0][-1])
            else:
                event_seq = hist_events
                time_seq = hist_times
                readout_pos = -1

            self.records.append(
                {
                    "patient_id": patient_id,
                    "eid": eid,
                    "doa": doa_years,
                    "event_seq": event_seq.astype(np.int64),
                    "time_seq": time_seq.astype(np.float32),
                    "readout_pos": readout_pos,
                    "full_events": disease_labels,
                    "full_times": disease_times,
                    "sex": int(features["sex"]),
                    "other_type": features["other_type"],
                    "other_value": features["other_value"],
                    "other_value_kind": features["other_value_kind"],
                    "other_time": features["other_time"],
                }
            )

        for token, pairs in first_lists.items():
            self.first_occurrence_by_token[int(token)] = (
                np.asarray([p for p, _ in pairs], dtype=np.int32),
                np.asarray([t for _, t in pairs], dtype=np.float32),
            )

        if not self.records:
            raise RuntimeError("No DOA records were built from checkup events.")

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

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        s = self.records[idx]
        return {
            "event_seq": torch.from_numpy(s["event_seq"]).long(),
            "time_seq": torch.from_numpy(s["time_seq"]).float(),
            "readout_pos": torch.tensor(s["readout_pos"], dtype=torch.long),
            "t_query": torch.tensor(float(s["doa"]), dtype=torch.float32),
            "patient_id": torch.tensor(s["patient_id"], dtype=torch.long),
            "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(),
        }


def collate_doa_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
    )
    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
    )

    readout_mask = torch.zeros_like(event_seq, dtype=torch.bool)
    readout_pos = torch.stack([x["readout_pos"] for x in batch])
    for i, pos in enumerate(readout_pos.tolist()):
        if pos >= 0:
            readout_mask[i, int(pos)] = True

    return {
        "event_seq": event_seq,
        "time_seq": time_seq,
        "padding_mask": event_seq > PAD_IDX,
        "readout_mask": readout_mask,
        "readout_pos": readout_pos,
        "t_query": torch.stack([x["t_query"] for x in batch]),
        "patient_id": torch.stack([x["patient_id"] for x in batch]),
        "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,
    }


@torch.inference_mode()
def infer_doa_hidden(
    model,
    loader: DataLoader,
    device: torch.device,
    model_target_mode: str,
    readout_name: str,
    readout_reduce: str,
    use_amp: bool,
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
    model_target_mode = str(model_target_mode).lower()
    readout = None
    if model_target_mode == "next_token":
        if readout_name == "same_time_group_end":
            readout = build_readout("same_time_group_end", reduce=readout_reduce).to(device)
        else:
            readout = build_readout(readout_name).to(device)
        readout.eval()

    hidden_parts: List[np.ndarray] = []
    patient_parts: List[np.ndarray] = []
    sex_parts: List[np.ndarray] = []
    autocast_enabled = bool(use_amp and device.type == "cuda")

    for batch in tqdm(loader, desc="DOA 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_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":
                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"],
                    target_mode="all_future",
                )
            else:
                hidden_raw = 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"],
                    target_mode="next_token",
                )
                ro = readout(
                    hidden=hidden_raw,
                    time_seq=batch_dev["time_seq"],
                    padding_mask=batch_dev["padding_mask"],
                    readout_mask=batch_dev["readout_mask"],
                )
                if ro.hidden.dim() == 2:
                    hidden = ro.hidden
                else:
                    hidden = ro.hidden[batch_dev["readout_mask"]]

        hidden_parts.append(hidden.detach().float().cpu().numpy().astype(np.float32, copy=False))
        patient_parts.append(batch["patient_id"].cpu().numpy().astype(np.int32, copy=False))
        sex_parts.append(batch["sex"].cpu().numpy().astype(np.int8, copy=False))

    return (
        np.concatenate(hidden_parts, axis=0),
        {
            "patient_id": np.concatenate(patient_parts, axis=0),
            "sex": np.concatenate(sex_parts, axis=0),
        },
    )


def first_time_array(
    first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
    token: int,
    patient_count: int,
) -> np.ndarray:
    out = np.full(patient_count, np.inf, dtype=np.float32)
    pairs = first_occurrence_by_token.get(int(token))
    if pairs is not None:
        p, t = pairs
        out[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32)
    return out


_DOA_WORKER: Dict[str, Any] = {}


def _init_doa_worker(
    disease_ids: np.ndarray,
    logits_all: np.ndarray,
    rho_all: Optional[np.ndarray],
    row_patient_id: np.ndarray,
    row_sex: np.ndarray,
    row_doa: np.ndarray,
    first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
    patient_count: int,
    horizons: np.ndarray,
    min_cases: int,
    dist_mode: str,
    score_mode: str,
    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")
    _DOA_WORKER.clear()
    _DOA_WORKER.update(
        {
            "disease_ids": np.asarray(disease_ids, dtype=np.int64),
            "logits_all": np.asarray(logits_all, dtype=np.float32),
            "rho_all": None if rho_all is None else np.asarray(rho_all, dtype=np.float32),
            "row_patient_id": np.asarray(row_patient_id, dtype=np.int32),
            "row_sex": np.asarray(row_sex, dtype=np.int8),
            "row_doa": np.asarray(row_doa, 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),
            "dist_mode": str(dist_mode).lower(),
            "score_mode": str(score_mode).lower(),
            "death_idx": int(death_idx),
            "first_time_cache": {},
        }
    )


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

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


def _eval_doa_token(task: Tuple[int, int]) -> List[Dict[str, Any]]:
    col, token = task
    col = int(col)
    token = int(token)

    patient_ids = _DOA_WORKER["row_patient_id"]
    sex = _DOA_WORKER["row_sex"]
    doa = _DOA_WORKER["row_doa"]
    logits = _DOA_WORKER["logits_all"][:, col]
    rho_all = _DOA_WORKER["rho_all"]
    rho = None if rho_all is None else rho_all[:, col]
    first_time = _doa_first_time_by_patient(token)[patient_ids]
    never = np.isinf(first_time)
    incident_after_doa = first_time > doa

    rows: List[Dict[str, Any]] = []
    for horizon in _DOA_WORKER["horizons"].tolist():
        horizon = float(horizon)
        case_mask = incident_after_doa & (first_time <= doa + np.float32(horizon))
        control_mask = never
        if int(case_mask.sum()) < int(_DOA_WORKER["min_cases"]) or int(control_mask.sum()) < int(_DOA_WORKER["min_cases"]):
            continue

        scores = _score_to_probability(
            logits=logits,
            rho=rho,
            score_mode=_DOA_WORKER["score_mode"],
            horizon=horizon,
            dist_mode=_DOA_WORKER["dist_mode"],
            token=token,
            death_idx=int(_DOA_WORKER["death_idx"]),
        )

        for sex_value, sex_name in [(0, "female"), (1, "male"), (-1, "all")]:
            sex_mask = np.ones_like(case_mask, dtype=bool) if sex_value == -1 else sex == sex_value
            cm = case_mask & sex_mask
            nm = control_mask & sex_mask
            if int(cm.sum()) < int(_DOA_WORKER["min_cases"]) or int(nm.sum()) < int(_DOA_WORKER["min_cases"]):
                continue
            auc, var = get_auc_delong_var(scores[cm], scores[nm])
            rows.append(
                {
                    "token": token,
                    "horizon": horizon,
                    "sex": sex_name,
                    "n_case": int(cm.sum()),
                    "n_control": int(nm.sum()),
                    "auc": auc,
                    "auc_var": var,
                    "auc_se": float(np.sqrt(max(var, 0.0))) if np.isfinite(var) else np.nan,
                }
            )
    return rows


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


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


def evaluate_doa_auc_chunk(
    dataset: DOAStatusDataset,
    hidden_all: np.ndarray,
    row_arrays: Dict[str, np.ndarray],
    model,
    disease_ids: Sequence[int],
    horizons: np.ndarray,
    dist_mode: str,
    score_mode: str,
    min_cases: int,
    device: torch.device,
    logit_batch_size: int,
    use_amp: bool,
    num_workers_auc: int,
    auc_task_chunk_size: int,
) -> pd.DataFrame:
    logits_all, rho_all = project_distribution_chunk(
        model=model,
        hidden_all=hidden_all,
        disease_ids=disease_ids,
        dist_mode=dist_mode,
        device=device,
        logit_batch_size=logit_batch_size,
        use_amp=use_amp,
    )
    patient_ids = row_arrays["patient_id"].astype(np.int32)
    doa = np.asarray([r["doa"] for r in dataset.records], dtype=np.float32)[patient_ids]
    patient_count = len(dataset.records)
    death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))
    disease_ids_arr = np.asarray([int(x) for x in disease_ids], dtype=np.int64)
    tasks = [(j, int(token)) for j, token in enumerate(disease_ids_arr.tolist())]

    init_args = (
        disease_ids_arr,
        logits_all,
        rho_all,
        patient_ids,
        row_arrays["sex"].astype(np.int8),
        doa,
        dataset.first_occurrence_by_token,
        patient_count,
        horizons,
        min_cases,
        dist_mode,
        score_mode,
        death_idx,
    )

    if int(num_workers_auc) <= 1 or len(tasks) <= 1:
        _init_doa_worker(*init_args)
        rows = _doa_task_block(tasks)
        return pd.DataFrame(rows)

    rows: List[Dict[str, Any]] = []
    task_blocks = _split_tasks(tasks, int(auc_task_chunk_size))
    with ProcessPoolExecutor(
        max_workers=int(num_workers_auc),
        initializer=_init_doa_worker,
        initargs=init_args,
    ) as pool:
        futures = [pool.submit(_doa_task_block, block) for block in task_blocks]
        for fut in tqdm(futures, desc="DOA AUC workers", leave=False, dynamic_ncols=True):
            rows.extend(fut.result())
    return pd.DataFrame(rows)


def iter_chunks(values: Sequence[int], chunk_size: int) -> Iterable[List[int]]:
    values = [int(x) for x in values]
    if chunk_size <= 0:
        yield values
        return
    for start in range(0, len(values), chunk_size):
        yield values[start:start + chunk_size]


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate DOA fixed-horizon 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=None,
                        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("--num_workers_auc", type=int, default=None)
    parser.add_argument("--auc_task_chunk_size", type=int, default=None)
    parser.add_argument("--logit_batch_size", type=int, default=None)
    parser.add_argument("--disease_chunk_size", type=int, default=None)
    parser.add_argument("--horizons", type=str, default=None)
    parser.add_argument("--score_mode", type=str, choices=["risk", "eta"], default=None)
    parser.add_argument("--filter_min_total", type=int, default=None)
    parser.add_argument("--min_cases", type=int, default=None)
    parser.add_argument("--labels_meta_path", type=str, default=None)
    parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None)
    args = parser.parse_args()

    run_path = Path(args.run_path)
    cfg = load_json_config(run_path / "train_config.json")
    ckpt_path = run_path / "best_model.pt"
    if not ckpt_path.exists():
        raise FileNotFoundError(f"best_model.pt not found in {run_path}")

    output_path = Path(args.output_path or run_path)
    output_path.mkdir(parents=True, exist_ok=True)

    model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
    if model_target_mode not in {"next_token", "all_future"}:
        raise ValueError(f"Unsupported model_target_mode={model_target_mode!r}")

    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 = pd.read_csv(labels_meta_path) if labels_meta_path and Path(labels_meta_path).exists() else None

    dataset = DOAStatusDataset(
        data_prefix=cfg.get("data_prefix", "ukb"),
        labels_file=cfg.get("labels_file", "labels.csv"),
        model_target_mode=model_target_mode,
        extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
    )
    validate_dataset_metadata(dataset, cfg)
    eval_indices = make_eval_indices(dataset, args, cfg)
    eval_patient_ids = np.asarray(
        [dataset.records[int(i)]["patient_id"] for i in eval_indices],
        dtype=np.int32,
    )
    eval_first_occurrence = subset_first_occurrence_map(
        dataset.first_occurrence_by_token,
        eval_patient_ids,
    )

    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=eval_first_occurrence,
    )
    if not disease_ids:
        raise RuntimeError("No disease tokens selected after filtering.")

    horizons = np.asarray(
        parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [1.0, 5.0, 10.0],
        dtype=np.float32,
    )
    score_mode = str(cfg_get(args, cfg, "score_mode", "risk")).lower()
    min_cases = int(cfg_get(args, cfg, "min_cases", 2))

    state_dict = load_checkpoint_state_dict(ckpt_path, map_location="cpu")
    dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
    cfg_model = dict(cfg)
    cfg_model["dist_mode"] = dist_mode

    device = torch.device(cfg.get("device", "cuda") if torch.cuda.is_available() else "cpu")
    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
    load_model_state(model, state_dict)
    model.eval()

    if model_target_mode == "next_token" and (
        model.token_embedding.num_embeddings <= CHECKUP_IDX
        or model.risk_head.out_features <= CHECKUP_IDX
    ):
        raise RuntimeError("Next-token DOA evaluation requires <CHECKUP> in the model vocabulary.")

    eval_dataset = Subset(dataset, eval_indices)
    loader = DataLoader(
        eval_dataset,
        batch_size=int(cfg_get(args, cfg, "batch_size", 128)),
        shuffle=False,
        collate_fn=collate_doa_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,
    )

    target_mode = cfg.get("target_mode", "uts")
    readout_name = str(cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token"))
    readout_reduce = str(cfg.get("readout_reduce", "mean"))

    eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
    if eval_split in {"valid", "validation"}:
        eval_split = "val"
    print(f"DOA records: total={len(dataset)}, eval_{eval_split}={len(eval_dataset)}")
    print(f"Model target mode: {model_target_mode}")
    print(f"Dist mode: {dist_mode}")
    print(f"Score mode: {score_mode}")
    print(f"Horizons: {horizons.tolist()}")
    print(f"Disease tokens: {len(disease_ids)}")

    hidden_all, row_arrays = infer_doa_hidden(
        model=model,
        loader=loader,
        device=device,
        model_target_mode=model_target_mode,
        readout_name=readout_name,
        readout_reduce=readout_reduce,
        use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
    )
    chunk_size = int(cfg_get(args, cfg, "disease_chunk_size", 256))
    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))
    print(f"Disease chunk size: {chunk_size}")
    print(f"AUC workers: {num_workers_auc}")
    result_parts = []
    for disease_chunk in tqdm(
        iter_chunks(disease_ids, chunk_size),
        desc="Disease chunks",
        leave=True,
        dynamic_ncols=True,
    ):
        result_parts.append(
            evaluate_doa_auc_chunk(
                dataset=dataset,
                hidden_all=hidden_all,
                row_arrays=row_arrays,
                model=model,
                disease_ids=disease_chunk,
                horizons=horizons,
                dist_mode=dist_mode,
                score_mode=score_mode,
                min_cases=min_cases,
                device=device,
                logit_batch_size=int(cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))),
                use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
                num_workers_auc=num_workers_auc,
                auc_task_chunk_size=auc_task_chunk_size,
            )
        )
    result = pd.concat(result_parts, ignore_index=True) if result_parts else pd.DataFrame()
    if result.empty:
        raise RuntimeError("No DOA AUC rows produced. Check disease selection and min_cases.")

    meta = build_metadata_for_merge(dataset, labels_meta)
    result = result.merge(meta, on="token", how="left")
    result.insert(0, "eval_split", eval_split)
    out_file = output_path / f"doa_auc_{eval_split}.csv"
    result.to_csv(out_file, index=False)

    summary = result.groupby(["token", "label_code", "horizon"], dropna=False, as_index=False).agg(
        auc_mean=("auc", "mean"),
        n_case=("n_case", "sum"),
        n_control=("n_control", "sum"),
    )
    summary.insert(0, "eval_split", eval_split)
    summary.to_csv(output_path / f"doa_auc_{eval_split}_summary.csv", index=False)
    print(f"Wrote {out_file}")


if __name__ == "__main__":
    main()
