DeepHealth/evaluate_single_disease_mortality_attribution.py

"""Compute per-disease attribution to predicted mortality risk.

For each selected patient and landmark age, this script keeps only rows where
each scanned disease token has already occurred in the history. It then deletes
that historical disease token, re-queries the model, and reports both
differences and ratios on probability and cumulative-hazard scales. If
--disease is omitted, all disease tokens in the mapping are scanned.

Death is always token vocab_size - 1.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict

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

from evaluate_auc_v2 import (
    build_model_from_dataset,
    cfg_get,
    load_checkpoint_state_dict,
    load_json_config,
    load_model_state,
    resolve_dist_mode_for_checkpoint,
    resolve_eval_device,
    validate_dataset_metadata,
)
from evaluate_event_free_survival import (
    IndexedLandmarkDataset,
    LandmarkDataset,
    build_first_occurrence_maps_for_landmarks,
    collate_indexed_landmark_fn,
    death_risk_for_batch,
    historical_counts_by_group,
    infer_landmark_hidden,
    load_eval_sequence_dataset,
    load_organ_groups,
    make_landmark_ages,
    make_occurred_mask,
    mortality_hazard_from_risk,
)
from future_risk import death_risk_from_probabilities, probabilities_from_logits
from targets import CHECKUP_IDX, PAD_IDX


OUTPUT_COLUMNS = [
    "patient_id",
    "dataset_index",
    "eid",
    "sex",
    "landmark_age",
    "tau",
    "followup_end_time",
    "history_disease_count",
    "selected_disease_history_count",
    "selected_disease_token_id",
    "selected_disease_code",
    "selected_disease_name",
    "selected_disease_organ_system",
    "selected_disease_organ_system_label",
    "history_count__selected_organ_system",
    "death_risk",
    "death_hazard",
    "ablated_death_risk",
    "ablated_death_hazard",
    "mortality_attribution_probability",
    "mortality_attribution_hazard",
    "mortality_attribution_probability_ratio",
    "mortality_attribution_hazard_ratio",
]

SUMMARY_KEY_COLUMNS = [
    "selected_disease_token_id",
    "selected_disease_code",
    "selected_disease_name",
    "selected_disease_organ_system",
    "selected_disease_organ_system_label",
    "landmark_age",
    "sex",
]

SUMMARY_MEAN_COLUMNS = [
    "history_disease_count",
    "selected_disease_history_count",
    "history_count__selected_organ_system",
    "death_risk",
    "death_hazard",
    "ablated_death_risk",
    "ablated_death_hazard",
    "mortality_attribution_probability",
    "mortality_attribution_hazard",
    "mortality_attribution_probability_ratio",
    "mortality_attribution_hazard_ratio",
]


def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int:
    table = table.reindex(columns=OUTPUT_COLUMNS)
    arrays: dict[str, np.ndarray] = {
        "__columns__": np.asarray(OUTPUT_COLUMNS, dtype="U"),
    }
    for column in OUTPUT_COLUMNS:
        values = table[column] if column in table else pd.Series([], dtype=object)
        if values.dtype == object:
            arrays[column] = values.fillna("").astype(str).to_numpy(dtype="U")
        else:
            arrays[column] = values.to_numpy()
    np.savez_compressed(path, **arrays)
    return int(len(table))


def normalize_output_dir(path: Path) -> Path:
    if path.suffix:
        return path.with_suffix(path.suffix + "_shards")
    return path


def write_manifest(
    output_dir: Path,
    *,
    rows: int,
    shards: list[dict[str, Any]],
    summary_file: str,
    scanned_diseases: list[dict[str, Any]],
    eval_split: str,
    tau: float,
    landmark_start: float,
    landmark_stop: float,
    landmark_step: float,
) -> None:
    payload = {
        "format": "compressed_npz_shards",
        "columns": OUTPUT_COLUMNS,
        "rows": int(rows),
        "shards": shards,
        "summary_file": summary_file,
        "scanned_diseases": scanned_diseases,
        "eval_split": eval_split,
        "tau": float(tau),
        "landmark_start": float(landmark_start),
        "landmark_stop": float(landmark_stop),
        "landmark_step": float(landmark_step),
    }
    with (output_dir / "manifest.json").open("w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)


def update_summary_accumulator(
    summary: dict[tuple[Any, ...], dict[str, float]],
    table: pd.DataFrame,
) -> None:
    if table.empty:
        return
    grouped = table.groupby(SUMMARY_KEY_COLUMNS, dropna=False, sort=False)
    for key, group in grouped:
        if not isinstance(key, tuple):
            key = (key,)
        acc = summary.setdefault(
            key,
            {"n": 0.0, **{column: 0.0 for column in SUMMARY_MEAN_COLUMNS}},
        )
        n = int(len(group))
        acc["n"] += float(n)
        for column in SUMMARY_MEAN_COLUMNS:
            acc[column] += float(pd.to_numeric(group[column], errors="coerce").sum())


def write_summary_csv(
    path: Path,
    summary: dict[tuple[Any, ...], dict[str, float]],
) -> int:
    rows: list[dict[str, Any]] = []
    for key, acc in summary.items():
        n = int(acc["n"])
        out = {column: value for column, value in zip(SUMMARY_KEY_COLUMNS, key)}
        out["n"] = n
        for column in SUMMARY_MEAN_COLUMNS:
            out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan
        rows.append(out)

    columns = [
        *SUMMARY_KEY_COLUMNS,
        "n",
        *[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS],
    ]
    pd.DataFrame(rows, columns=columns).sort_values(
        ["selected_disease_token_id", "landmark_age", "sex"],
        kind="mergesort",
    ).to_csv(path, index=False)
    return len(rows)


def concat_padded_tensor_batches(chunks: list[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
    if not chunks:
        raise ValueError("Cannot concatenate an empty chunk list")

    fill_values = {
        "event_seq": PAD_IDX,
        "time_seq": 0.0,
        "readout_mask": False,
        "padding_mask": False,
        "other_type": 0,
        "other_value": 0.0,
        "other_value_kind": 0,
        "other_time": 0.0,
    }
    out: Dict[str, torch.Tensor] = {}
    for key in chunks[0]:
        tensors = [chunk[key] for chunk in chunks]
        shapes = [tuple(t.shape) for t in tensors]
        if len(set(shapes)) == 1:
            out[key] = torch.cat(tensors, dim=0)
            continue

        if any(t.ndim == 0 for t in tensors):
            raise ValueError(f"Cannot concatenate scalar tensor key={key!r} with mismatched shapes")
        max_shape = list(shapes[0])
        for shape in shapes[1:]:
            if len(shape) != len(max_shape):
                raise ValueError(f"Cannot concatenate key={key!r} with shapes {shapes}")
            max_shape = [max(a, b) for a, b in zip(max_shape, shape)]

        padded: list[torch.Tensor] = []
        fill = fill_values.get(key, 0)
        for tensor in tensors:
            target_shape = [int(tensor.shape[0]), *max_shape[1:]]
            padded_tensor = tensor.new_full(target_shape, fill)
            slices = tuple(slice(0, int(size)) for size in tensor.shape)
            padded_tensor[slices] = tensor
            padded.append(padded_tensor)
        out[key] = torch.cat(padded, dim=0)
    return out


def build_disease_ablated_slice(
    batch: Dict[str, torch.Tensor],
    row_indices: torch.Tensor,
    token_ids: torch.Tensor,
) -> Dict[str, torch.Tensor]:
    """Build an ablated slice for aligned (row, disease_token) pairs."""
    event_seq = batch["event_seq"]
    row_indices = row_indices.to(device=event_seq.device, dtype=torch.long)
    token_ids = token_ids.to(device=event_seq.device, dtype=event_seq.dtype)

    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, :]
    remove = (out["event_seq"] == token_ids[:, None]) & 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 load_disease_metadata(
    mapping_path: Path,
    *,
    vocab_size: int,
) -> dict[int, dict[str, Any]]:
    if not mapping_path.exists():
        raise FileNotFoundError(f"Disease mapping file not found: {mapping_path}")
    table = pd.read_csv(mapping_path)
    required = {"token_id", "code", "name", "is_death"}
    missing = required - set(table.columns)
    if missing:
        raise ValueError(f"{mapping_path} is missing columns: {sorted(missing)}")

    death_idx = int(vocab_size) - 1
    out: dict[int, dict[str, Any]] = {}
    for row in table.itertuples(index=False):
        token = int(getattr(row, "token_id"))
        if token < 0 or token >= int(vocab_size) or token == death_idx:
            continue
        if int(getattr(row, "is_death")) == 1:
            continue
        meta = {
            "token_id": token,
            "code": str(getattr(row, "code")),
            "name": str(getattr(row, "name")),
        }
        for column in (
            "icd10_chapter",
            "icd10_chapter_title",
            "organ_system",
            "organ_system_label",
        ):
            if hasattr(row, column):
                meta[column] = str(getattr(row, column))
        out[token] = meta
    return out


def resolve_disease_token(
    value: str,
    metadata: dict[int, dict[str, Any]],
) -> tuple[int, dict[str, Any]]:
    text = str(value).strip()
    if text == "":
        raise ValueError("--disease must not be empty")

    if text.isdigit() or (text.startswith("-") and text[1:].isdigit()):
        token = int(text)
        if token not in metadata:
            raise ValueError(f"Disease token_id {token} was not found in the mapping")
        return token, metadata[token]

    lower = text.lower()
    exact = [
        (token, meta)
        for token, meta in metadata.items()
        if str(meta.get("code", "")).lower() == lower
        or str(meta.get("name", "")).lower() == lower
    ]
    if len(exact) == 1:
        return exact[0]
    if len(exact) > 1:
        raise ValueError(f"--disease={value!r} matched multiple diseases exactly")

    contains = [
        (token, meta)
        for token, meta in metadata.items()
        if lower in str(meta.get("code", "")).lower()
        or lower in str(meta.get("name", "")).lower()
    ]
    if len(contains) == 1:
        return contains[0]
    if not contains:
        raise ValueError(f"--disease={value!r} did not match any disease token")
    preview = ", ".join(
        f"{token}:{meta.get('code')} ({meta.get('name')})"
        for token, meta in contains[:10]
    )
    raise ValueError(
        f"--disease={value!r} matched {len(contains)} diseases; use token_id or code. "
        f"First matches: {preview}"
    )


def resolve_disease_tokens(
    value: str | None,
    metadata: dict[int, dict[str, Any]],
) -> list[tuple[int, dict[str, Any]]]:
    if value is None or str(value).strip() == "":
        return [(token, metadata[token]) for token in sorted(metadata)]
    out: list[tuple[int, dict[str, Any]]] = []
    seen: set[int] = set()
    for part in str(value).split(","):
        token, meta = resolve_disease_token(part, metadata)
        if token not in seen:
            out.append((token, meta))
            seen.add(token)
    return out


def safe_ratio(
    numerator: torch.Tensor,
    denominator: torch.Tensor,
    *,
    eps: float,
) -> torch.Tensor:
    return numerator / denominator.clamp_min(float(eps))


def output_name_for_run(run_path: Path, eval_split: str, tau: float, *, all_diseases: bool) -> Path:
    scope = "all_diseases" if all_diseases else "selected_diseases"
    return run_path / f"single_disease_mortality_attribution_{eval_split}_{scope}_tau{tau:g}y"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compute per-disease model attribution to mortality risk."
    )
    parser.add_argument("--run_path", type=str, required=True)
    parser.add_argument(
        "--disease",
        type=str,
        default=None,
        help=(
            "Optional disease token_id, ICD-10 code, exact name, unambiguous name "
            "substring, or comma-separated list. If omitted, scan all disease tokens."
        ),
    )
    parser.add_argument(
        "--output_path",
        type=str,
        default=None,
        help="Output directory for compressed .npz shards.",
    )
    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 disease-token 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)
    parser.add_argument(
        "--ratio_eps",
        type=float,
        default=1e-7,
        help="Small lower bound for ratio denominators.",
    )
    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)

    metadata = load_disease_metadata(
        Path(args.organ_mapping_path),
        vocab_size=int(dataset.vocab_size),
    )
    scanned_disease_items = resolve_disease_tokens(args.disease, metadata)
    if not scanned_disease_items:
        raise ValueError("No diseases selected for attribution")
    scanned_disease_tokens = [token for token, _meta in scanned_disease_items]

    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)
    scanned_disease_tensor = torch.as_tensor(
        scanned_disease_tokens,
        dtype=torch.long,
        device=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 * 8, batch_size))
    )
    if attribution_batch_size <= 0:
        raise ValueError("attribution_batch_size must be positive")
    if float(args.ratio_eps) <= 0:
        raise ValueError("--ratio_eps 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,
            all_diseases=args.disease is None or str(args.disease).strip() == "",
        )
    )
    output_dir = normalize_output_dir(output_path)
    output_dir.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}")
    if len(scanned_disease_items) == len(metadata):
        print(f"Diseases: all mapped diseases ({len(scanned_disease_items)})")
    else:
        preview = ", ".join(
            f"{token}:{meta.get('code')}" for token, meta in scanned_disease_items[:10]
        )
        print(f"Diseases: {len(scanned_disease_items)} selected ({preview})")
    print(f"Landmark rows: {len(landmark_dataset)}")
    print(f"Attribution batch size: {attribution_batch_size}")
    print(f"Output directory: {output_dir}")

    written_rows = 0
    shard_index = 0
    shards: list[dict[str, Any]] = []
    summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
    row_base_cache: dict[int, dict[str, Any]] = {}
    pending_batch_chunks: list[Dict[str, torch.Tensor]] = []
    pending_meta_chunks: list[list[dict[str, Any]]] = []
    pending_n = 0

    def flush_pending() -> None:
        nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks, pending_n
        if pending_n == 0:
            return

        ablated_batch = concat_padded_tensor_batches(pending_batch_chunks)
        meta_rows = [row for chunk in pending_meta_chunks for row in chunk]
        with torch.no_grad():
            ablated_risk = death_risk_for_batch(
                model=model,
                batch=ablated_batch,
                device=device,
                model_target_mode=model_target_mode,
                readout_name=readout_name,
                readout_reduce=readout_reduce,
                dist_mode=dist_mode,
                tau=tau,
            )
        ablated_hazard = mortality_hazard_from_risk(ablated_risk)
        orig_risk = torch.as_tensor(
            [row.pop("_death_risk") for row in meta_rows],
            dtype=ablated_risk.dtype,
            device=ablated_risk.device,
        )
        orig_hazard = torch.as_tensor(
            [row.pop("_death_hazard") for row in meta_rows],
            dtype=ablated_hazard.dtype,
            device=ablated_hazard.device,
        )

        attr_prob = orig_risk - ablated_risk
        attr_hazard = orig_hazard - ablated_hazard
        ratio_prob = safe_ratio(orig_risk, ablated_risk, eps=float(args.ratio_eps))
        ratio_hazard = safe_ratio(orig_hazard, ablated_hazard, eps=float(args.ratio_eps))

        for i, row in enumerate(meta_rows):
            row["death_risk"] = float(orig_risk[i].detach().cpu())
            row["death_hazard"] = float(orig_hazard[i].detach().cpu())
            row["ablated_death_risk"] = float(ablated_risk[i].detach().cpu())
            row["ablated_death_hazard"] = float(ablated_hazard[i].detach().cpu())
            row["mortality_attribution_probability"] = float(attr_prob[i].detach().cpu())
            row["mortality_attribution_hazard"] = float(attr_hazard[i].detach().cpu())
            row["mortality_attribution_probability_ratio"] = float(
                ratio_prob[i].detach().cpu()
            )
            row["mortality_attribution_hazard_ratio"] = float(
                ratio_hazard[i].detach().cpu()
            )

        table = pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)
        update_summary_accumulator(summary_accumulator, table)
        shard_name = f"part-{shard_index:06d}.npz"
        shard_path = output_dir / shard_name
        shard_rows = write_compressed_npz_table(shard_path, table)
        shards.append({"file": shard_name, "rows": int(shard_rows)})
        shard_index += 1
        written_rows += len(meta_rows)
        pending_batch_chunks = []
        pending_meta_chunks = []
        pending_n = 0

    def get_row_base(row_idx: int) -> dict[str, Any]:
        cached = row_base_cache.get(row_idx)
        if cached is not None:
            return cached

        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)
        unique_tokens, token_counts = np.unique(hist_tokens, return_counts=True)
        total_count, group_counts = historical_counts_by_group(
            hist_tokens,
            death_idx=death_idx,
            token_to_group=token_to_group,
            group_names=group_names,
        )
        cached = {
            "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),
            "_hist_tokens": hist_tokens,
            "_token_counts": {
                int(token): int(count)
                for token, count in zip(unique_tokens.tolist(), token_counts.tolist())
            },
            "_group_counts": group_counts,
        }
        row_base_cache[row_idx] = cached
        return cached

    for batch in tqdm(loader, desc="Per-disease mortality attribution", 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()
        }
        hidden = infer_landmark_hidden(
            model=model,
            batch=batch_dev,
            device=device,
            model_target_mode=model_target_mode,
            readout_name=readout_name,
            readout_reduce=readout_reduce,
        )
        with torch.no_grad():
            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,
            )
        death_risk_tensor = death_risk_from_probabilities(probabilities)
        death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
        occurred = make_occurred_mask(
            batch_dev["event_seq"],
            vocab_size=int(dataset.vocab_size),
            device=device,
        )
        pair_indices = torch.nonzero(
            occurred[:, scanned_disease_tensor],
            as_tuple=False,
        )
        if pair_indices.numel() == 0:
            continue

        pair_offset = 0
        while pair_offset < int(pair_indices.shape[0]):
            capacity = int(attribution_batch_size) - pending_n
            pair_stop = min(int(pair_indices.shape[0]), pair_offset + capacity)
            pair_chunk = pair_indices[pair_offset:pair_stop]
            local_rows = pair_chunk[:, 0].long()
            disease_token_ids = scanned_disease_tensor[pair_chunk[:, 1]].long()
            ablated_chunk = build_disease_ablated_slice(
                batch=batch_dev,
                row_indices=local_rows,
                token_ids=disease_token_ids,
            )

            meta_chunk: list[dict[str, Any]] = []
            row_ids = batch_dev["row_idx"][local_rows].detach().cpu().numpy().astype(np.int64)
            disease_tokens_list = disease_token_ids.detach().cpu().numpy().astype(np.int64)
            for local_pos, (row_idx, disease_token) in enumerate(
                zip(row_ids.tolist(), disease_tokens_list.tolist())
            ):
                disease_token = int(disease_token)
                disease_meta = metadata[disease_token]
                row_base = get_row_base(int(row_idx))
                group_counts = row_base["_group_counts"]
                disease_history_count = int(row_base["_token_counts"].get(disease_token, 0))
                if disease_history_count <= 0:
                    raise RuntimeError(
                        "Internal mismatch: occurred mask selected disease "
                        f"{disease_token} for row {row_idx}, but cached history has count 0"
                    )

                orig_row = int(local_rows[local_pos].detach().cpu().item())
                meta_chunk.append(
                    {
                        "patient_id": row_base["patient_id"],
                        "dataset_index": row_base["dataset_index"],
                        "eid": row_base["eid"],
                        "sex": row_base["sex"],
                        "landmark_age": row_base["landmark_age"],
                        "tau": row_base["tau"],
                        "followup_end_time": row_base["followup_end_time"],
                        "history_disease_count": row_base["history_disease_count"],
                        "selected_disease_history_count": disease_history_count,
                        "selected_disease_token_id": int(disease_token),
                        "selected_disease_code": str(disease_meta.get("code", "")),
                        "selected_disease_name": str(disease_meta.get("name", "")),
                        "selected_disease_organ_system": str(
                            disease_meta.get("organ_system", "")
                        ),
                        "selected_disease_organ_system_label": str(
                            disease_meta.get("organ_system_label", "")
                        ),
                        "history_count__selected_organ_system": int(
                            group_counts.get(str(disease_meta.get("organ_system", "")), 0)
                        ),
                        "_death_risk": float(death_risk_tensor[orig_row].detach().cpu()),
                        "_death_hazard": float(death_hazard_tensor[orig_row].detach().cpu()),
                    }
                )

            if meta_chunk:
                pending_batch_chunks.append(ablated_chunk)
                pending_meta_chunks.append(meta_chunk)
                pending_n += len(meta_chunk)
            pair_offset = pair_stop

            if pending_n >= int(attribution_batch_size):
                flush_pending()

    flush_pending()
    if not shards:
        empty_path = output_dir / "part-000000.npz"
        write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))
        shards.append({"file": empty_path.name, "rows": 0})
    summary_path = output_dir / "summary_by_disease_age_sex.csv"
    summary_rows = write_summary_csv(summary_path, summary_accumulator)
    write_manifest(
        output_dir,
        rows=written_rows,
        shards=shards,
        summary_file=summary_path.name,
        scanned_diseases=[
            {"token_id": int(token), **{k: v for k, v in meta.items() if k != "token_id"}}
            for token, meta in scanned_disease_items
        ],
        eval_split=eval_split,
        tau=tau,
        landmark_start=float(args.landmark_start),
        landmark_stop=float(args.landmark_stop),
        landmark_step=float(args.landmark_step),
    )
    print(f"Wrote {written_rows} rows in {len(shards)} shard(s) to {output_dir}")
    print(f"Wrote {summary_rows} summary rows to {summary_path}")


if __name__ == "__main__":
    main()