DeepHealth/evaluate_event_free_survival.py

"""Compute landmark future death and incident system-disease risks.

For each selected patient and landmark age, this script computes:

* future death risk within tau years;
* future incident disease risk for each ICD-10 chapter-derived system;
* model attribution of each historical organ/system disease set to predicted
  mortality risk, computed by deleting that system's historical disease tokens
  and re-querying the model;
* historical modeled-disease count;
* historical modeled-disease count within each ICD-10 chapter-derived system.

Death is always token vocab_size - 1. Disease groups are read from
icd10_chapter_organ_mapping.csv.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence

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
from tqdm.auto import tqdm

from dataset import HealthDataset
from eval_data import load_sequence_eval_dataset
from evaluate_auc_v2 import (
    LandmarkDataset,
    build_model_from_dataset,
    cfg_get,
    load_checkpoint_state_dict,
    load_json_config,
    load_model_state,
    make_eval_indices,
    resolve_dist_mode_for_checkpoint,
    resolve_eval_device,
    validate_dataset_metadata,
)
from future_risk import (
    death_risk_from_probabilities,
    new_disease_risk_from_probabilities,
    probabilities_from_logits,
)
from models import DeepHealth
from readouts import build_readout
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
from train_util import load_eid_file, load_extra_info_types_file


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("["):
        values = json.loads(text)
        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 load_extra_info_types(value: Any) -> Optional[List[int]]:
    if value is None:
        return None
    text = str(value)
    path = Path(text)
    if path.exists():
        return load_extra_info_types_file(text)
    return parse_int_list(value)


def make_landmark_ages(start: float, stop: float, step: float) -> np.ndarray:
    if step <= 0:
        raise ValueError("landmark_step must be positive")
    if stop < start:
        raise ValueError("landmark_stop must be >= landmark_start")
    # Include stop when it lands on the grid, e.g. 40,45,...,80.
    return np.arange(start, stop + step * 0.5, step, dtype=np.float32)


def build_first_occurrence_maps_for_landmarks(
    dataset: HealthDataset,
    subset_indices: np.ndarray,
) -> Dict[int, tuple[np.ndarray, np.ndarray]]:
    first_lists: Dict[int, list[tuple[int, float]]] = {}
    for patient_id, dataset_index in enumerate(np.asarray(subset_indices, dtype=np.int64).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:]])
        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)
            if token in SPECIAL_TOKENS:
                continue
            first_lists.setdefault(token, []).append((patient_id, float(full_time[int(idx)])))

    return {
        int(token): (
            np.asarray([p for p, _ in pairs], dtype=np.int32),
            np.asarray([t for _, t in pairs], dtype=np.float32),
        )
        for token, pairs in first_lists.items()
        if pairs
    }


def normalize_eval_split(args: argparse.Namespace, cfg: Dict[str, Any]) -> str:
    eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
    if eval_split in {"valid", "validation"}:
        return "val"
    if eval_split not in {"train", "val", "test", "all"}:
        raise ValueError(f"Unsupported eval_split={eval_split!r}")
    return eval_split


def load_eval_sequence_dataset(
    args: argparse.Namespace,
    cfg: Dict[str, Any],
) -> tuple[Any, np.ndarray, str, str]:
    eval_split = normalize_eval_split(args, cfg)
    model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
    data_prefix = str(cfg.get("data_prefix", "ukb"))
    labels_file = str(cfg.get("labels_file", "labels.csv"))
    no_event_interval_years = float(cfg.get("no_event_interval_years", 5.0))
    include_no_event_in_uts_target = bool(cfg.get("include_no_event_in_uts_target", False))
    extra_info_types = load_extra_info_types(args.extra_info_types)
    if extra_info_types is None:
        extra_info_types = parse_int_list(cfg.get("extra_info_types", None))

    print("Loading one sequence eval 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_in_uts_target,
        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=extra_info_types,
    )

    train_eid_file = cfg_get(args, cfg, "train_eid_file", "ukb_train_eid.csv")
    val_eid_file = cfg_get(args, cfg, "val_eid_file", "ukb_val_eid.csv")
    test_eid_file = cfg_get(args, cfg, "test_eid_file", "ukb_test_eid.csv")
    split_files_exist = all(
        Path(str(path)).exists()
        for path in (train_eid_file, val_eid_file, test_eid_file)
    )

    if eval_split != "all" and split_files_exist:
        split_files = {
            "train": train_eid_file,
            "val": val_eid_file,
            "test": test_eid_file,
        }
        selected_eids = load_eid_file(split_files[eval_split])
        out = np.asarray(
            [
                idx
                for idx, sample in enumerate(dataset.samples)
                if int(sample["eid"]) in selected_eids
            ],
            dtype=np.int64,
        )
        if out.size == 0:
            raise ValueError(
                f"No samples found for eval_split={eval_split!r} using {split_files[eval_split]}"
            )
        split_source = "eid_files"
    else:
        if eval_split == "all":
            out = np.arange(len(dataset.samples), dtype=np.int64)
            split_source = "all"
        else:
            out = make_eval_indices(dataset, args, cfg)
            split_source = "ratio_split"

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


def load_organ_groups(
    path: Path,
    *,
    vocab_size: int,
) -> tuple[dict[str, list[int]], dict[str, str], dict[int, str]]:
    table = pd.read_csv(path)
    required = {"token_id", "organ_system", "organ_system_label", "is_death"}
    missing = required - set(table.columns)
    if missing:
        raise ValueError(f"{path} is missing columns: {sorted(missing)}")

    death_idx = int(vocab_size) - 1
    groups: dict[str, list[int]] = {}
    labels: dict[str, str] = {}
    token_to_group: dict[int, str] = {}
    for row in table.itertuples(index=False):
        token = int(getattr(row, "token_id"))
        if token in SPECIAL_TOKENS or token == death_idx:
            continue
        if token < 0 or token >= int(vocab_size):
            continue
        if int(getattr(row, "is_death")) == 1:
            continue
        group = str(getattr(row, "organ_system"))
        label = str(getattr(row, "organ_system_label"))
        groups.setdefault(group, []).append(token)
        labels[group] = label
        token_to_group[token] = group

    groups = {k: sorted(set(v)) for k, v in groups.items() if v}
    return groups, labels, token_to_group


class IndexedLandmarkDataset(Dataset):
    def __init__(self, base: LandmarkDataset) -> None:
        self.base = base

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

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        item = dict(self.base[idx])
        item["row_idx"] = torch.tensor(int(idx), dtype=torch.long)
        return item


def collate_indexed_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]),
        "row_idx": torch.stack([x["row_idx"] for x in batch]),
    }


def build_group_ablated_slice(
    batch: Dict[str, torch.Tensor],
    token_ids: Sequence[int],
    row_indices: torch.Tensor,
) -> Dict[str, torch.Tensor]:
    """Build one fixed-width ablated slice without rebuilding variable-length rows."""
    event_seq = batch["event_seq"]

    out: Dict[str, torch.Tensor] = {}
    out["event_seq"] = event_seq[row_indices].clone()
    out["time_seq"] = batch["time_seq"][row_indices]
    out["readout_mask"] = batch["readout_mask"][row_indices].clone()
    out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
    out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()

    seq_len = int(event_seq.shape[1])
    positions = torch.arange(seq_len, device=event_seq.device)[None, :]
    ids = torch.as_tensor(token_ids, dtype=event_seq.dtype, device=event_seq.device)
    remove = torch.isin(out["event_seq"], ids) & out["padding_mask"]
    out["event_seq"] = torch.where(
        remove,
        torch.full_like(out["event_seq"], PAD_IDX),
        out["event_seq"],
    )
    out["padding_mask"] &= ~remove
    out["readout_mask"] &= ~remove

    has_valid = out["padding_mask"].any(dim=1)
    if not bool(has_valid.all().item()):
        empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
        out["event_seq"][empty_rows, 0] = CHECKUP_IDX
        out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
            dtype=out["time_seq"].dtype
        )
        out["padding_mask"][empty_rows, 0] = True
        out["readout_mask"][empty_rows, 0] = True
        out["landmark_pos"][empty_rows] = 0

    has_readout = out["readout_mask"].any(dim=1)
    if not bool(has_readout.all().item()):
        rows = torch.nonzero(~has_readout, as_tuple=False).flatten()
        local_valid = out["padding_mask"][rows]
        last_pos = torch.where(
            local_valid,
            positions.expand(local_valid.shape[0], -1),
            torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
        ).amax(dim=1)
        out["readout_mask"][rows] = False
        out["readout_mask"][rows, last_pos] = True
        out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype)

    repeated_keys = (
        "sex",
        "other_type",
        "other_value",
        "other_value_kind",
        "other_time",
        "t_query",
        "patient_id",
        "landmark_age",
        "followup_end_time",
        "death_time",
        "row_idx",
    )
    for key in repeated_keys:
        out[key] = batch[key][row_indices]
    return out


def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
    return {
        key: torch.cat([chunk[key] for chunk in chunks], dim=0)
        for key in chunks[0]
    }


def iter_group_ablated_batches(
    batch: Dict[str, torch.Tensor],
    group_names: Sequence[str],
    organ_groups: dict[str, list[int]],
    occurred: torch.Tensor,
    max_batch_size: int,
):
    """Yield ablated chunks as soon as enough rows are available for a forward pass."""
    pending_batches: list[Dict[str, torch.Tensor]] = []
    pending_groups: list[str] = []
    pending_rows: list[int] = []
    pending_n = 0

    for group in group_names:
        ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device)
        if ids.numel() == 0:
            continue
        active_rows = torch.nonzero(occurred[:, ids].any(dim=1), as_tuple=False).flatten()
        if active_rows.numel() == 0:
            continue

        row_offset = 0
        while row_offset < int(active_rows.numel()):
            capacity = int(max_batch_size) - pending_n
            row_stop = min(int(active_rows.numel()), row_offset + capacity)
            row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
            chunk = build_group_ablated_slice(
                batch=batch,
                token_ids=organ_groups[group],
                row_indices=row_indices,
            )
            chunk_n = int(row_indices.numel())
            pending_batches.append(chunk)
            pending_groups.extend([group] * chunk_n)
            pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist())
            pending_n += chunk_n
            row_offset = row_stop

            if pending_n >= int(max_batch_size):
                yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
                pending_batches = []
                pending_groups = []
                pending_rows = []
                pending_n = 0

    if pending_batches:
        yield concat_tensor_batches(pending_batches), pending_groups, pending_rows


@torch.no_grad()
def infer_landmark_hidden(
    *,
    model: DeepHealth,
    batch: Dict[str, torch.Tensor],
    device: torch.device,
    model_target_mode: str,
    readout_name: str,
    readout_reduce: str,
) -> torch.Tensor:
    batch_dev = {
        k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
        for k, v in batch.items()
    }
    if model_target_mode == "all_future":
        return model(
            event_seq=batch_dev["event_seq"].long(),
            time_seq=batch_dev["time_seq"].float(),
            sex=batch_dev["sex"].long(),
            padding_mask=batch_dev["padding_mask"].bool(),
            t_query=batch_dev["t_query"].float(),
            other_type=batch_dev["other_type"].long(),
            other_value=batch_dev["other_value"].float(),
            other_value_kind=batch_dev["other_value_kind"].long(),
            other_time=batch_dev["other_time"].float(),
            target_mode="all_future",
        )

    hidden = model(
        event_seq=batch_dev["event_seq"].long(),
        time_seq=batch_dev["time_seq"].float(),
        sex=batch_dev["sex"].long(),
        padding_mask=batch_dev["padding_mask"].bool(),
        other_type=batch_dev["other_type"].long(),
        other_value=batch_dev["other_value"].float(),
        other_value_kind=batch_dev["other_value_kind"].long(),
        other_time=batch_dev["other_time"].float(),
        target_mode="next_token",
    )
    readout = build_readout(readout_name, reduce=readout_reduce)
    readout_out = readout(
        hidden=hidden,
        time_seq=batch_dev["time_seq"].float(),
        padding_mask=batch_dev["padding_mask"].bool(),
        readout_mask=batch_dev["readout_mask"].bool(),
    )
    return readout_out.hidden.gather(
        1,
        batch_dev["landmark_pos"].long()[:, None, None].expand(
            -1, 1, readout_out.hidden.shape[-1]
        ),
    ).squeeze(1)


def make_occurred_mask(
    event_seq: torch.Tensor,
    *,
    vocab_size: int,
    device: torch.device,
) -> torch.Tensor:
    occurred = torch.zeros(event_seq.shape[0], int(vocab_size), dtype=torch.bool, device=device)
    valid = (event_seq >= 0) & (event_seq < int(vocab_size))
    safe = event_seq.clamp(min=0, max=int(vocab_size) - 1).to(device)
    occurred.scatter_(1, safe, valid.to(device))
    return occurred


def mortality_hazard_from_risk(risk: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
    return -torch.log1p(-risk.clamp(0.0, 1.0 - float(eps)))


def death_risk_for_batch(
    *,
    model: DeepHealth,
    batch: Dict[str, torch.Tensor],
    device: torch.device,
    model_target_mode: str,
    readout_name: str,
    readout_reduce: str,
    dist_mode: str,
    tau: float,
) -> torch.Tensor:
    hidden = infer_landmark_hidden(
        model=model,
        batch=batch,
        device=device,
        model_target_mode=model_target_mode,
        readout_name=readout_name,
        readout_reduce=readout_reduce,
    )
    logits = model.calc_risk(hidden)
    rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
    death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None
    probabilities = probabilities_from_logits(
        logits,
        tau,
        dist_mode=dist_mode,
        rho=rho,
        death_rho=death_rho,
    )
    return death_risk_from_probabilities(probabilities)


def historical_counts_by_group(
    tokens: np.ndarray,
    *,
    death_idx: int,
    token_to_group: dict[int, str],
    group_names: Sequence[str],
) -> tuple[int, dict[str, int]]:
    unique_tokens = {
        int(token)
        for token in np.asarray(tokens, dtype=np.int64).tolist()
        if int(token) not in SPECIAL_TOKENS and int(token) != int(death_idx)
    }
    total = len(unique_tokens)
    out = {group: 0 for group in group_names}
    for token in unique_tokens:
        group = token_to_group.get(token)
        if group in out:
            out[group] += 1
    return total, out


def output_name_for_run(run_path: Path, eval_split: str, tau: float) -> Path:
    return run_path / f"future_risk_{eval_split}_tau{tau:g}y.csv"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compute landmark death and incident system-disease risks."
    )
    parser.add_argument("--run_path", type=str, required=True)
    parser.add_argument("--output_path", type=str, default=None)
    parser.add_argument("--organ_mapping_path", type=str, default="icd10_chapter_organ_mapping.csv")
    parser.add_argument("--eval_split", type=str, default=None)
    parser.add_argument("--dataset_subset_size", type=int, default=None)
    parser.add_argument("--train_eid_file", type=str, default=None)
    parser.add_argument("--val_eid_file", type=str, default=None)
    parser.add_argument("--test_eid_file", type=str, default=None)
    parser.add_argument("--landmark_start", type=float, default=40.0)
    parser.add_argument("--landmark_stop", type=float, default=80.0)
    parser.add_argument("--landmark_step", type=float, default=5.0)
    parser.add_argument("--tau", type=float, default=5.0)
    parser.add_argument("--min_history_events", type=int, default=None)
    parser.add_argument("--batch_size", type=int, default=None)
    parser.add_argument(
        "--attribution_batch_size",
        type=int,
        default=None,
        help="Forward batch size for expanded organ/system ablation queries.",
    )
    parser.add_argument("--num_workers", type=int, default=None)
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--extra_info_types", type=str, default=None)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    run_path = Path(args.run_path)
    config_path = run_path / "train_config.json"
    checkpoint_path = run_path / "best_model.pt"
    if not config_path.exists():
        raise FileNotFoundError(f"train_config.json not found: {config_path}")
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")

    cfg = load_json_config(config_path)
    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}")

    target_mode = str(cfg.get("target_mode", "uts"))
    attn_mask_mode = str(
        cfg.get("attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware")
    )
    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"))

    dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(
        args,
        cfg,
    )
    validate_dataset_metadata(dataset, cfg)

    landmark_ages = make_landmark_ages(
        float(args.landmark_start),
        float(args.landmark_stop),
        float(args.landmark_step),
    )
    tau = float(args.tau)
    if tau < 0:
        raise ValueError("tau must be non-negative")

    first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
        dataset,
        subset_indices,
    )
    death_idx = int(dataset.vocab_size) - 1
    landmark_dataset = LandmarkDataset(
        dataset=dataset,
        subset_indices=subset_indices,
        landmark_ages=landmark_ages,
        attn_mask_mode=attn_mask_mode,
        model_target_mode=model_target_mode,
        min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)),
        first_occurrence_by_token=first_occurrence_by_token,
        death_token_ids=[death_idx],
    )

    organ_groups, organ_labels, token_to_group = load_organ_groups(
        Path(args.organ_mapping_path),
        vocab_size=int(dataset.vocab_size),
    )
    group_names = sorted(organ_groups)

    state_dict = load_checkpoint_state_dict(checkpoint_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 = resolve_eval_device(args.device)
    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
    load_model_state(model, state_dict)
    model.eval()

    batch_size = int(cfg_get(args, cfg, "batch_size", 128))
    attribution_batch_size = int(
        cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 4, batch_size))
    )
    if attribution_batch_size <= 0:
        raise ValueError("attribution_batch_size must be positive")
    num_workers = int(cfg_get(args, cfg, "num_workers", 4))
    loader = DataLoader(
        IndexedLandmarkDataset(landmark_dataset),
        batch_size=batch_size,
        shuffle=False,
        collate_fn=collate_indexed_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,
    )

    output_path = Path(args.output_path) if args.output_path else output_name_for_run(run_path, eval_split, tau)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"Eval split: {eval_split}")
    print(f"Split source: {split_source}")
    print(f"Selected patients: {len(subset_indices)}")
    print(f"Landmark ages: {landmark_ages.tolist()}")
    print(f"Tau: {tau:g} years")
    print(f"Dist mode: {dist_mode}")
    print(f"Device: {device}")
    print(f"Death token: {death_idx}")
    print(f"Organ/system groups: {len(group_names)}")
    print(f"Landmark rows: {len(landmark_dataset)}")
    print(f"Attribution batch size: {attribution_batch_size}")
    print(f"Output: {output_path}")

    rows: list[dict[str, Any]] = []
    for batch in tqdm(loader, desc="Future risks", dynamic_ncols=True):
        hidden = infer_landmark_hidden(
            model=model,
            batch=batch,
            device=device,
            model_target_mode=model_target_mode,
            readout_name=readout_name,
            readout_reduce=readout_reduce,
        )
        logits = model.calc_risk(hidden)
        rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
        death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None
        probabilities = probabilities_from_logits(
            logits,
            tau,
            dist_mode=dist_mode,
            rho=rho,
            death_rho=death_rho,
        )
        occurred = make_occurred_mask(
            batch["event_seq"].to(device),
            vocab_size=int(dataset.vocab_size),
            device=device,
        )

        death_risk_tensor = death_risk_from_probabilities(probabilities)
        death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
        death_risk = death_risk_tensor.detach().cpu().numpy()

        group_risk: dict[str, np.ndarray] = {}
        for group in group_names:
            group_risk[group] = new_disease_risk_from_probabilities(
                probabilities,
                occurred,
                organ_groups[group],
            ).detach().cpu().numpy()

        group_mortality_attr_prob: dict[str, np.ndarray] = {}
        group_mortality_attr_hazard: dict[str, np.ndarray] = {}
        batch_n = int(batch["event_seq"].shape[0])
        zeros = np.zeros(batch_n, dtype=np.float32)
        for group in group_names:
            group_mortality_attr_prob[group] = zeros.copy()
            group_mortality_attr_hazard[group] = zeros.copy()

        for ablated_chunk, chunk_groups, chunk_rows in iter_group_ablated_batches(
            batch=batch,
            group_names=group_names,
            organ_groups=organ_groups,
            occurred=occurred,
            max_batch_size=attribution_batch_size,
        ):
            ablated_death_risk = death_risk_for_batch(
                model=model,
                batch=ablated_chunk,
                device=device,
                model_target_mode=model_target_mode,
                readout_name=readout_name,
                readout_reduce=readout_reduce,
                dist_mode=dist_mode,
                tau=tau,
            )
            row_tensor = torch.as_tensor(chunk_rows, dtype=torch.long, device=device)
            ablated_death_hazard = mortality_hazard_from_risk(ablated_death_risk)
            attr_prob = (
                death_risk_tensor[row_tensor] - ablated_death_risk
            ).detach().cpu().numpy()
            attr_hazard = (
                death_hazard_tensor[row_tensor] - ablated_death_hazard
            ).detach().cpu().numpy()
            for local_idx, (group, row_idx) in enumerate(zip(chunk_groups, chunk_rows)):
                group_mortality_attr_prob[group][row_idx] = attr_prob[local_idx]
                group_mortality_attr_hazard[group][row_idx] = attr_hazard[local_idx]

        row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
        for j, row_idx in enumerate(row_indices.tolist()):
            meta = landmark_dataset.rows[int(row_idx)]
            dataset_index = int(meta["dataset_index"])
            sample = dataset.samples[dataset_index]
            hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
            total_count, group_counts = historical_counts_by_group(
                hist_tokens,
                death_idx=death_idx,
                token_to_group=token_to_group,
                group_names=group_names,
            )

            out: dict[str, Any] = {
                "patient_id": int(meta["patient_id"]),
                "dataset_index": dataset_index,
                "eid": int(sample.get("eid", -1)),
                "sex": int(meta["sex"]),
                "landmark_age": float(meta["landmark_age"]),
                "tau": tau,
                "followup_end_time": float(meta["followup_end_time"]),
                "history_disease_count": int(total_count),
                "death_risk": float(death_risk[j]),
            }
            for group in group_names:
                out[f"history_count__{group}"] = int(group_counts[group])
                out[f"new_disease_risk__{group}"] = float(group_risk[group][j])
                if int(group_counts[group]) == 0:
                    group_mortality_attr_prob[group][j] = 0.0
                    group_mortality_attr_hazard[group][j] = 0.0
                out[f"mortality_attribution_probability__{group}"] = float(
                    group_mortality_attr_prob[group][j]
                )
                out[f"mortality_attribution_hazard__{group}"] = float(
                    group_mortality_attr_hazard[group][j]
                )
            rows.append(out)

    df = pd.DataFrame(rows)
    df.to_csv(output_path, index=False)
    print(f"Wrote {len(df)} rows to {output_path}")


if __name__ == "__main__":
    main()