DeepHealth/compute_burden_index_landmarks.py

from __future__ import annotations

import argparse
import multiprocessing as mp
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Iterable

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

from burden_index import (
    _build_readout_grid,
    _observed_formed_burden,
    load_burden_context,
)
from evaluate_auc_v2 import (
    make_eval_indices,
    parse_float_list,
)
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX


def _parse_landmark_ages(args: argparse.Namespace) -> np.ndarray:
    explicit = parse_float_list(args.landmark_ages)
    if explicit:
        ages = np.asarray(explicit, dtype=np.float32)
    else:
        ages = np.arange(
            float(args.landmark_start),
            float(args.landmark_stop) + 1e-6,
            float(args.landmark_step),
            dtype=np.float32,
        )
    if ages.size == 0:
        raise ValueError("No landmark ages were provided.")
    return ages


def _parse_horizon(value: Any) -> float:
    horizon = float(value)
    if horizon < 0:
        raise ValueError(f"horizon must be non-negative, got {horizon}")
    return horizon


def _format_horizon_for_filename(horizon: float) -> str:
    text = f"{float(horizon):g}".replace("-", "m").replace(".", "p")
    return f"h{text}"


def _parse_devices(args: argparse.Namespace) -> list[str | None]:
    if args.devices is not None and str(args.devices).strip():
        devices = [x.strip() for x in str(args.devices).split(",") if x.strip()]
        if not devices:
            raise ValueError("--devices was provided but no valid devices were parsed.")
        return devices
    return [args.device]


def _build_burden_matrix_from_mapping(
    mapping_csv: Path,
    *,
    token_col: str,
    category_id_col: str,
    category_col: str,
    key_area_col: str,
    weight_col: str,
) -> tuple[np.ndarray, np.ndarray, pd.DataFrame]:
    df = pd.read_csv(mapping_csv)
    required = {token_col, category_id_col, category_col, weight_col}
    missing = sorted(required - set(df.columns))
    if missing:
        raise ValueError(f"{mapping_csv} is missing required columns: {missing}")

    df = df.copy()
    df[token_col] = pd.to_numeric(df[token_col], errors="raise").astype(int)
    raw_category = df[category_id_col]
    numeric_category = pd.to_numeric(raw_category, errors="coerce")
    if numeric_category.notna().all():
        df["_burden_category_key"] = numeric_category.astype(int)
    else:
        df["_burden_category_key"] = raw_category.astype(str)
    df[weight_col] = pd.to_numeric(df[weight_col], errors="raise").astype(float)
    df = df[df[weight_col] != 0].copy()
    if df.empty:
        raise ValueError(f"{mapping_csv} has no non-zero burden weights.")

    disease_ids = np.asarray(sorted(df[token_col].unique().tolist()), dtype=np.int64)
    category_keys = sorted(df["_burden_category_key"].unique().tolist())

    disease_pos = {int(token): j for j, token in enumerate(disease_ids.tolist())}
    category_pos = {cat: i for i, cat in enumerate(category_keys)}
    A = np.zeros((len(category_keys), disease_ids.size), dtype=np.float64)
    for _, row in df.iterrows():
        token = int(row[token_col])
        cat = row["_burden_category_key"]
        weight = float(row[weight_col])
        A[category_pos[cat], disease_pos[token]] += weight

    meta_cols = list(dict.fromkeys(["_burden_category_key", category_id_col, category_col]))
    if key_area_col in df.columns:
        meta_cols.append(key_area_col)
    category_meta = (
        df[meta_cols]
        .drop_duplicates(subset=["_burden_category_key"])
        .sort_values("_burden_category_key")
        .reset_index(drop=True)
    )
    category_meta = category_meta.rename(
        columns={
            "_burden_category_key": "burden_dimension_id",
            category_col: "burden_dimension",
            key_area_col: "burden_key_area",
        }
    )
    if category_id_col in category_meta.columns:
        category_meta = category_meta.drop(columns=[category_id_col])
    if "burden_key_area" not in category_meta.columns:
        category_meta["burden_key_area"] = ""

    return A, disease_ids, category_meta


def _parse_burden_types(value: str) -> list[str]:
    out = [x.strip().lower() for x in str(value).split(",") if x.strip()]
    if not out:
        raise ValueError("burden_types must contain at least one value.")
    valid = {"functional", "organ"}
    unknown = sorted(set(out) - valid)
    if unknown:
        raise ValueError(f"Unknown burden_types {unknown}; valid values are {sorted(valid)}")
    return list(dict.fromkeys(out))


def _load_mapping_specs(args: argparse.Namespace) -> list[dict[str, Any]]:
    specs: list[dict[str, Any]] = []
    for burden_type in _parse_burden_types(args.burden_types):
        if burden_type == "functional":
            path = Path(args.functional_mapping_csv)
            specs.append(
                {
                    "burden_type": "functional",
                    "mapping_csv": path,
                    "token_col": "token_id",
                    "category_id_col": "hfrm_category_id",
                    "category_col": "hfrm_category",
                    "key_area_col": "hfrm_key_area",
                    "weight_col": args.functional_weight_col,
                }
            )
        elif burden_type == "organ":
            path = Path(args.organ_mapping_csv)
            specs.append(
                {
                    "burden_type": "organ",
                    "mapping_csv": path,
                    "token_col": "token_id",
                    "category_id_col": "organ_system",
                    "category_col": "organ_system",
                    "key_area_col": "icd10_chapter_name",
                    "weight_col": "organ_weight",
                }
            )
    for spec in specs:
        if not spec["mapping_csv"].exists():
            raise FileNotFoundError(
                f"{spec['burden_type']} mapping_csv not found: {spec['mapping_csv']}"
            )
    return specs


def _eligible_landmark_rows(
    dataset: Any,
    subset_indices: np.ndarray,
    landmark_ages: np.ndarray,
    *,
    min_history_events: int,
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    special = np.asarray([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64)
    for patient_id, dataset_index in enumerate(subset_indices.tolist()):
        sample = dataset.samples[int(dataset_index)]
        seq_event = np.asarray(sample["event_seq"], dtype=np.int64)
        seq_time = np.asarray(sample["time_seq"], dtype=np.float32)
        tgt_event = np.asarray(sample["target_event_seq"], dtype=np.int64)
        tgt_time = np.asarray(sample["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:]])
        followup_end = float(np.max(full_time))

        for landmark_age in landmark_ages.tolist():
            landmark_age = float(landmark_age)
            if not (followup_end > landmark_age):
                continue
            prefix_mask = full_time <= np.float32(landmark_age)
            if not np.any(prefix_mask):
                continue
            prefix_events = full_event[prefix_mask].astype(np.int64, copy=False)
            prefix_times = full_time[prefix_mask].astype(np.float32, copy=False)
            valid_history = ~np.isin(prefix_events, special)
            if int(valid_history.sum()) < int(min_history_events):
                continue

            rows.append(
                {
                    "patient_id": int(patient_id),
                    "dataset_index": int(dataset_index),
                    "sex": int(sample["sex"]),
                    "landmark_age": np.float32(landmark_age),
                    "t_query": np.float32(landmark_age),
                    "followup_end_time": np.float32(followup_end),
                    "event_seq": prefix_events,
                    "time_seq": prefix_times,
                    "other_type": np.asarray(sample["other_type"], dtype=np.int64),
                    "other_value": np.asarray(sample["other_value"], dtype=np.float32),
                    "other_value_kind": np.asarray(
                        sample["other_value_kind"], dtype=np.int64),
                    "other_time": np.asarray(sample["other_time"], dtype=np.float32),
                }
            )
    return rows


def _row_to_worker_spec(row: dict[str, Any]) -> dict[str, Any]:
    return {
        "patient_id": int(row["patient_id"]),
        "dataset_index": int(row["dataset_index"]),
        "landmark_age": float(row["landmark_age"]),
        "followup_end_time": float(row["followup_end_time"]),
    }


def _materialize_worker_rows(
    dataset: Any,
    row_specs: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for spec in row_specs:
        sample = dataset.samples[int(spec["dataset_index"])]
        seq_event = np.asarray(sample["event_seq"], dtype=np.int64)
        seq_time = np.asarray(sample["time_seq"], dtype=np.float32)
        tgt_event = np.asarray(sample["target_event_seq"], dtype=np.int64)
        tgt_time = np.asarray(sample["target_time_seq"], dtype=np.float32)
        full_event = np.concatenate([seq_event, tgt_event[-1:]])
        full_time = np.concatenate([seq_time, tgt_time[-1:]])
        t_query = np.float32(float(spec["landmark_age"]))
        prefix_mask = full_time <= t_query
        landmark_age = np.float32(float(spec["landmark_age"]))
        if landmark_age != t_query:
            raise RuntimeError(
                f"landmark_age and t_query diverged for dataset_index={spec['dataset_index']}"
            )
        rows.append(
            {
                "patient_id": int(spec["patient_id"]),
                "dataset_index": int(spec["dataset_index"]),
                "sex": int(sample["sex"]),
                "landmark_age": landmark_age,
                "t_query": t_query,
                "followup_end_time": np.float32(float(spec["followup_end_time"])),
                "event_seq": full_event[prefix_mask].astype(np.int64, copy=False),
                "time_seq": full_time[prefix_mask].astype(np.float32, copy=False),
                "other_type": np.asarray(sample["other_type"], dtype=np.int64),
                "other_value": np.asarray(sample["other_value"], dtype=np.float32),
                "other_value_kind": np.asarray(sample["other_value_kind"], dtype=np.int64),
                "other_time": np.asarray(sample["other_time"], dtype=np.float32),
            }
        )
    return rows


def _config_split_indices(
    n: int,
    cfg: dict[str, Any],
    eval_split: str,
    subset_size: int,
) -> np.ndarray:
    args = argparse.Namespace(
        train_ratio=None,
        val_ratio=None,
        test_ratio=None,
        seed=None,
        eval_split=eval_split,
        dataset_subset_size=subset_size if subset_size > 0 else None,
    )
    # make_eval_indices reads len(dataset), so a small shim is enough.
    class _Sized:
        def __len__(self) -> int:
            return n

    return make_eval_indices(_Sized(), args, cfg)


class ReadoutJobIterableDataset(IterableDataset):
    def __init__(
        self,
        *,
        rows: list[dict[str, Any]],
        formed_mode: str,
        horizon: float,
    ) -> None:
        super().__init__()
        self.rows = rows
        self.formed_mode = str(formed_mode)
        self.horizon = float(horizon)
        if self.formed_mode not in {"observed", "model_weighted"}:
            raise ValueError(f"Unknown formed_mode={self.formed_mode!r}")

    def __iter__(self) -> Iterable[dict[str, torch.Tensor]]:
        worker = get_worker_info()
        if worker is None:
            start, step = 0, 1
        else:
            start, step = int(worker.id), int(worker.num_workers)

        for row_idx in range(start, len(self.rows), step):
            row = self.rows[row_idx]
            if self.formed_mode == "model_weighted":
                grid = _build_readout_grid(
                    event_seq=row["event_seq"],
                    time_seq=row["time_seq"],
                    other_type=row["other_type"],
                    other_time=row["other_time"],
                    t_query=float(row["t_query"]),
                )
                if grid.size > 0:
                    end_times = np.concatenate(
                        [grid[1:], np.asarray([row["t_query"]], dtype=np.float32)]
                    )
                    row_deltas = np.maximum(end_times - grid, 0.0).astype(np.float32)
                    valid = row_deltas > 0
                    for query_time, delta in zip(grid[valid].tolist(), row_deltas[valid].tolist()):
                        yield _make_readout_job(row, row_idx, "formed", query_time, delta)
            if self.horizon > 0:
                yield _make_readout_job(
                    row,
                    row_idx,
                    "future",
                    float(row["t_query"]),
                    self.horizon,
                )


def _make_readout_job(
    row: dict[str, Any],
    row_idx: int,
    kind: str,
    query_time: float,
    delta: float,
) -> dict[str, torch.Tensor]:
    return {
        "event_seq": torch.from_numpy(np.asarray(row["event_seq"], dtype=np.int64)).long(),
        "time_seq": torch.from_numpy(np.asarray(row["time_seq"], dtype=np.float32)).float(),
        "sex": torch.tensor(int(row["sex"]), dtype=torch.long),
        "other_type": torch.from_numpy(np.asarray(row["other_type"], dtype=np.int64)).long(),
        "other_value": torch.from_numpy(np.asarray(row["other_value"], dtype=np.float32)).float(),
        "other_value_kind": torch.from_numpy(
            np.asarray(row["other_value_kind"], dtype=np.int64)
        ).long(),
        "other_time": torch.from_numpy(np.asarray(row["other_time"], dtype=np.float32)).float(),
        "query_time": torch.tensor(float(query_time), dtype=torch.float32),
        "delta": torch.tensor(float(delta), dtype=torch.float32),
        "row_idx": torch.tensor(int(row_idx), dtype=torch.long),
        "kind": torch.tensor(0 if kind == "formed" else 1, dtype=torch.long),
    }


def _collate_readout_jobs(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
    )
    return {
        "event_seq": event_seq,
        "time_seq": time_seq,
        "padding_mask": event_seq > PAD_IDX,
        "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,
        "query_time": torch.stack([x["query_time"] for x in batch]),
        "delta": torch.stack([x["delta"] for x in batch]),
        "row_idx": torch.stack([x["row_idx"] for x in batch]),
        "kind": torch.stack([x["kind"] for x in batch]),
    }


@torch.inference_mode()
def _query_hidden_readout_batch(
    *,
    ctx: Any,
    batch: dict[str, torch.Tensor],
) -> torch.Tensor:
    event = batch["event_seq"].long().to(ctx.device, non_blocking=True)
    return ctx.model(
        event_seq=event,
        time_seq=batch["time_seq"].float().to(ctx.device, non_blocking=True),
        sex=batch["sex"].long().to(ctx.device, non_blocking=True),
        padding_mask=event > PAD_IDX,
        t_query=batch["query_time"].float().to(ctx.device, non_blocking=True),
        other_type=batch["other_type"].long().to(ctx.device, non_blocking=True),
        other_value=batch["other_value"].float().to(ctx.device, non_blocking=True),
        other_value_kind=batch["other_value_kind"].long().to(ctx.device, non_blocking=True),
        other_time=batch["other_time"].float().to(ctx.device, non_blocking=True),
        target_mode="all_future",
    )

@torch.inference_mode()
def _probabilities_from_hidden_torch(
    *,
    ctx: Any,
    hidden: torch.Tensor,
    disease_ids: np.ndarray,
    deltas: np.ndarray,
) -> torch.Tensor:
    if hidden.ndim != 2:
        raise ValueError(f"hidden must have shape (N, H), got {tuple(hidden.shape)}")
    if deltas.ndim != 1 or deltas.size != hidden.shape[0]:
        raise ValueError(
            "deltas must be 1D with the same length as hidden rows, got "
            f"{deltas.shape} vs {tuple(hidden.shape)}"
        )

    ids = torch.as_tensor(disease_ids, dtype=torch.long, device=ctx.device)
    logits = ctx.model.calc_risk(hidden)[:, ids]
    rate = torch.nn.functional.softplus(logits).clamp_min(1e-8)
    delta_t = torch.as_tensor(deltas, dtype=rate.dtype, device=ctx.device).clamp_min(0)

    if ctx.dist_mode == "weibull":
        rho = ctx.model.calc_weibull_rho(hidden)[:, ids]
        exposure = torch.pow(delta_t[:, None], rho)
    elif ctx.dist_mode == "mixed":
        exposure = delta_t[:, None].expand_as(rate)
        death_idx = int(getattr(ctx.model, "death_idx", getattr(ctx.model, "vocab_size", 0) - 1))
        death_cols = [j for j, token in enumerate(disease_ids.tolist()) if int(token) == death_idx]
        if death_cols:
            death_rho = ctx.model.calc_death_rho(hidden)
            for col in death_cols:
                exposure[:, int(col)] = torch.pow(delta_t, death_rho)
    else:
        exposure = delta_t[:, None].expand_as(rate)

    return -torch.expm1(-rate * exposure)

@torch.inference_mode()
def _readout_probabilities_from_batch(
    *,
    ctx: Any,
    batch: dict[str, torch.Tensor],
    union_disease_ids: np.ndarray,
) -> torch.Tensor:
    hidden = _query_hidden_readout_batch(ctx=ctx, batch=batch)
    deltas = batch["delta"].detach().cpu().numpy().astype(np.float32, copy=False)
    return _probabilities_from_hidden_torch(
        ctx=ctx,
        hidden=hidden,
        disease_ids=union_disease_ids,
        deltas=deltas,
    ).to(dtype=torch.float32)


def _observed_formed_for_rows(
    *,
    rows: list[dict[str, Any]],
    union_disease_ids: np.ndarray,
) -> np.ndarray:
    formed = np.zeros((len(rows), union_disease_ids.size), dtype=np.float64)
    for row_idx, row in enumerate(rows):
        formed[row_idx] = _observed_formed_burden(
            disease_ids=union_disease_ids,
            event_seq=row["event_seq"],
            time_seq=row["time_seq"],
            t_query=float(row["t_query"]),
        )
    return formed


def _apply_readout_batch_to_accumulators(
    *,
    batch: dict[str, torch.Tensor],
    readout_prob: torch.Tensor,
    survival_by_row: torch.Tensor | None,
    future_prob_by_row: torch.Tensor,
    ctx: Any,
) -> None:
    if readout_prob.numel() == 0:
        return

    row_indices = batch["row_idx"].long().to(ctx.device, non_blocking=True)
    kinds = batch["kind"].long().to(ctx.device, non_blocking=True)
    kind_is_formed = kinds == 0
    kind_is_future = kinds == 1

    if survival_by_row is not None and bool(kind_is_formed.any().item()):
        formed_rows = row_indices[kind_is_formed]
        formed_survival = 1.0 - readout_prob[kind_is_formed].clamp(0.0, 1.0)
        if hasattr(survival_by_row, "scatter_reduce_"):
            survival_by_row.scatter_reduce_(
                dim=0,
                index=formed_rows[:, None].expand_as(formed_survival),
                src=formed_survival,
                reduce="prod",
                include_self=True,
            )
        else:
            for job_idx in torch.nonzero(kind_is_formed, as_tuple=False).flatten().tolist():
                survival_by_row[int(row_indices[job_idx].item())] *= (
                    1.0 - readout_prob[job_idx].clamp(0.0, 1.0)
                )
    if bool(kind_is_future.any().item()):
        future_rows = row_indices[kind_is_future]
        future_prob_by_row[future_rows] = readout_prob[kind_is_future]


def _project_bi_rows(
    *,
    rows: list[dict[str, Any]],
    horizon: float,
    matrices: list[dict[str, Any]],
    formed_mode: str,
    formed_by_row: torch.Tensor,
    future_prob_by_row: torch.Tensor,
    ctx: Any,
) -> list[dict[str, Any]]:
    disease_future_by_row = (1.0 - formed_by_row) * future_prob_by_row
    disease_total_by_row = formed_by_row + disease_future_by_row
    projected: list[dict[str, Any]] = []
    for matrix in matrices:
        A = torch.as_tensor(matrix["A_union"], dtype=formed_by_row.dtype, device=ctx.device)
        projected.append(
            {
                "matrix": matrix,
                "historical": torch.matmul(formed_by_row, A.T).detach().cpu().numpy(),
                "future": torch.matmul(disease_future_by_row, A.T).detach().cpu().numpy(),
                "total": torch.matmul(disease_total_by_row, A.T).detach().cpu().numpy(),
            }
        )

    out: list[dict[str, Any]] = []
    for row_idx, row in enumerate(rows):
        query_time = float(row["t_query"])
        for item in projected:
            matrix = item["matrix"]
            historical = item["historical"][row_idx]
            future = item["future"][row_idx]
            total = item["total"][row_idx]
            for dim_idx, meta in matrix["category_meta"].iterrows():
                out.append(
                    {
                        "patient_id": row["patient_id"],
                        "dataset_index": row["dataset_index"],
                        "sex": row["sex"],
                        "landmark_age": query_time,
                        "t_query": query_time,
                        "followup_end_time": float(row["followup_end_time"]),
                        "horizon": float(horizon),
                        "formed_mode": formed_mode,
                        "burden_type": matrix["burden_type"],
                        "burden_dimension_id": meta["burden_dimension_id"],
                        "burden_dimension": str(meta["burden_dimension"]),
                        "burden_key_area": str(meta.get("burden_key_area", "")),
                        "bi_historical": float(historical[int(dim_idx)]),
                        "bi_future": float(future[int(dim_idx)]),
                        "bi_total": float(total[int(dim_idx)]),
                    }
                )
    return out


def _write_rows_csv(rows: list[dict[str, Any]], output_path: Path) -> int:
    df = pd.DataFrame(rows)
    df.to_csv(output_path, index=False)
    return int(len(df))


def _concat_csv_shards(shard_paths: list[Path], output_path: Path) -> None:
    wrote_header = False
    with output_path.open("w", encoding="utf-8", newline="") as out_f:
        for shard_path in shard_paths:
            with shard_path.open("r", encoding="utf-8", newline="") as in_f:
                header = in_f.readline()
                if not wrote_header:
                    out_f.write(header)
                    wrote_header = True
                for line in in_f:
                    out_f.write(line)
            shard_path.unlink(missing_ok=True)


def _compute_bi_from_streamed_readouts(
    *,
    rows: list[dict[str, Any]],
    horizon: float,
    matrices: list[dict[str, Any]],
    union_disease_ids: np.ndarray,
    formed_mode: str,
    readout_batch_size: int,
    num_workers: int,
    ctx: Any,
    log_prefix: str | None = None,
) -> tuple[list[dict[str, Any]], int, dict[str, float]]:
    horizon = float(horizon)
    if horizon < 0:
        raise ValueError(f"horizon must be non-negative, got {horizon}")

    dtype = torch.float32
    if formed_mode == "observed":
        formed_by_row = torch.as_tensor(
            _observed_formed_for_rows(
                rows=rows,
                union_disease_ids=union_disease_ids,
            ),
            dtype=dtype,
            device=ctx.device,
        )
        survival_by_row = None
    elif formed_mode == "model_weighted":
        survival_by_row = torch.ones(
            (len(rows), union_disease_ids.size),
            dtype=dtype,
            device=ctx.device,
        )
        formed_by_row = None
    else:
        raise ValueError(f"Unknown formed_mode={formed_mode!r}")
    future_prob_by_row = torch.zeros(
        (len(rows), union_disease_ids.size),
        dtype=dtype,
        device=ctx.device,
    )

    readout_jobs = 0
    n_batches = 0
    build_readout_sec = 0.0
    forward_sec = 0.0
    reduce_sec = 0.0
    t_loader0 = time.perf_counter()
    readout_dataset = ReadoutJobIterableDataset(
        rows=rows,
        formed_mode=formed_mode,
        horizon=horizon,
    )
    readout_loader = DataLoader(
        readout_dataset,
        batch_size=max(1, int(readout_batch_size)),
        collate_fn=_collate_readout_jobs,
        num_workers=max(0, int(num_workers)),
        pin_memory=ctx.device.type == "cuda",
        persistent_workers=int(num_workers) > 0,
        prefetch_factor=2 if int(num_workers) > 0 else None,
    )
    build_readout_sec += time.perf_counter() - t_loader0

    for batch in readout_loader:
        t_forward0 = time.perf_counter()
        readout_prob = _readout_probabilities_from_batch(
            ctx=ctx,
            batch=batch,
            union_disease_ids=union_disease_ids,
        )
        if ctx.device.type == "cuda":
            torch.cuda.synchronize(ctx.device)
        forward_sec += time.perf_counter() - t_forward0

        t_reduce0 = time.perf_counter()
        _apply_readout_batch_to_accumulators(
            batch=batch,
            readout_prob=readout_prob,
            survival_by_row=survival_by_row,
            future_prob_by_row=future_prob_by_row,
            ctx=ctx,
        )
        if ctx.device.type == "cuda":
            torch.cuda.synchronize(ctx.device)
        reduce_sec += time.perf_counter() - t_reduce0

        n_batches += 1
        readout_jobs += int(batch["row_idx"].numel())
        if log_prefix and (n_batches == 1 or n_batches % 50 == 0):
            print(
                f"{log_prefix} processed {readout_jobs} readout jobs "
                f"in {n_batches} batches",
                flush=True,
            )

    if survival_by_row is not None:
        formed_by_row = 1.0 - survival_by_row
    assert formed_by_row is not None

    t_project0 = time.perf_counter()
    rows_out = _project_bi_rows(
        rows=rows,
        horizon=horizon,
        matrices=matrices,
        formed_mode=formed_mode,
        formed_by_row=formed_by_row,
        future_prob_by_row=future_prob_by_row,
        ctx=ctx,
    )
    if ctx.device.type == "cuda":
        torch.cuda.synchronize(ctx.device)
    reduce_sec += time.perf_counter() - t_project0
    timings = {
        "build_readout_sec": build_readout_sec,
        "forward_sec": forward_sec,
        "reduce_sec": reduce_sec,
    }
    return rows_out, readout_jobs, timings


def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
    device = payload["device"]
    run_path = Path(payload["run_path"])
    shard_path = Path(payload["shard_path"])
    print(
        f"[BI worker {device}] starting with {len(payload['row_specs'])} rows",
        flush=True,
    )
    load_start = time.perf_counter()
    ctx = load_burden_context(run_path, device=device)
    print(
        f"[BI worker {device}] context loaded in {time.perf_counter() - load_start:.2f}s",
        flush=True,
    )
    materialize_start = time.perf_counter()
    rows = _materialize_worker_rows(ctx.dataset, payload["row_specs"])
    print(
        f"[BI worker {device}] rows materialized in "
        f"{time.perf_counter() - materialize_start:.2f}s",
        flush=True,
    )
    out, readout_jobs, timings = _compute_bi_from_streamed_readouts(
        rows=rows,
        horizon=payload["horizon"],
        matrices=payload["matrices"],
        union_disease_ids=payload["union_disease_ids"],
        formed_mode=payload["formed_mode"],
        readout_batch_size=int(payload["readout_batch_size"]),
        num_workers=int(payload["num_workers"]),
        ctx=ctx,
        log_prefix=f"[BI worker {device}]",
    )
    print(
        f"[BI worker {device}] done: readout_jobs={readout_jobs}, "
        f"build={timings['build_readout_sec']:.2f}s, "
        f"forward={timings['forward_sec']:.2f}s, "
        f"reduce={timings['reduce_sec']:.2f}s",
        flush=True,
    )
    write_start = time.perf_counter()
    row_count = _write_rows_csv(out, shard_path)
    timings["write_csv_sec"] = time.perf_counter() - write_start
    print(
        f"[BI worker {device}] wrote {row_count} rows to {shard_path} "
        f"in {timings['write_csv_sec']:.2f}s",
        flush=True,
    )
    return {
        "shard_path": str(shard_path),
        "row_count": row_count,
        "readout_jobs": readout_jobs,
        "timings": timings,
    }


def _attach_union_projection(
    matrices: list[dict[str, Any]],
) -> tuple[np.ndarray, list[dict[str, Any]]]:
    union_disease_ids = np.asarray(
        sorted(
            {
                int(token)
                for matrix in matrices
                for token in np.asarray(matrix["disease_ids"], dtype=np.int64).tolist()
            }
        ),
        dtype=np.int64,
    )
    if union_disease_ids.size == 0:
        raise ValueError("No disease tokens are covered by the requested burden matrices.")

    union_pos = {int(token): i for i, token in enumerate(union_disease_ids.tolist())}
    projected: list[dict[str, Any]] = []
    for matrix in matrices:
        disease_ids = np.asarray(matrix["disease_ids"], dtype=np.int64)
        A = np.asarray(matrix["A"], dtype=np.float64)
        A_union = np.zeros((A.shape[0], union_disease_ids.size), dtype=np.float64)
        for local_col, token in enumerate(disease_ids.tolist()):
            A_union[:, union_pos[int(token)]] += A[:, int(local_col)]
        item = dict(matrix)
        item["A_union"] = A_union
        projected.append(item)
    return union_disease_ids, projected


def _split_rows_for_devices(
    rows: list[dict[str, Any]],
    devices: list[str | None],
    *,
    formed_mode: str,
    horizon: float,
) -> list[tuple[str | None, list[dict[str, Any]]]]:
    if len(devices) <= 1:
        return [(devices[0], rows)]

    buckets: list[list[dict[str, Any]]] = [[] for _ in devices]
    loads = np.zeros(len(devices), dtype=np.int64)
    weighted_rows = sorted(
        rows,
        key=lambda row: _estimate_readout_jobs_for_row(
            row,
            formed_mode=formed_mode,
            horizon=horizon,
        ),
        reverse=True,
    )
    for row in weighted_rows:
        bucket_idx = int(np.argmin(loads))
        buckets[bucket_idx].append(row)
        loads[bucket_idx] += _estimate_readout_jobs_for_row(
            row,
            formed_mode=formed_mode,
            horizon=horizon,
        )

    chunks: list[tuple[str | None, list[dict[str, Any]]]] = []
    for device, bucket in zip(devices, buckets):
        if not bucket:
            continue
        chunks.append((device, bucket))
    return chunks


def _estimate_readout_jobs_for_row(
    row: dict[str, Any],
    *,
    formed_mode: str,
    horizon: float,
) -> int:
    n_jobs = 0
    if formed_mode == "model_weighted":
        grid = _build_readout_grid(
            event_seq=row["event_seq"],
            time_seq=row["time_seq"],
            other_type=row["other_type"],
            other_time=row["other_time"],
            t_query=float(row["t_query"]),
        )
        if grid.size > 0:
            end_times = np.concatenate([grid[1:], np.asarray([row["t_query"]], dtype=np.float32)])
            n_jobs += int(np.sum(np.maximum(end_times - grid, 0.0) > 0))
    if horizon > 0:
        n_jobs += 1
    return max(n_jobs, 1)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compute DeepHealth Burden Indices at landmark ages."
    )
    parser.add_argument("--run_path", type=str, required=True)
    parser.add_argument("--burden_types", type=str, default="functional,organ",
                        help="Comma-separated burden types to compute: functional,organ.")
    parser.add_argument("--functional_mapping_csv", type=str,
                        default="cihi_hfrm_label_mapping.csv")
    parser.add_argument("--organ_mapping_csv", type=str,
                        default="icd10_organ_label_mapping.csv")
    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("--formed_mode", type=str, default="model_weighted",
                        choices=["observed", "model_weighted"])
    parser.add_argument(
        "--horizon",
        type=float,
        required=True,
        help=(
            "Future horizon in years. Use 0 to compute historical burden only "
            "(bi_future=0 and bi_total=bi_historical)."
        ),
    )
    parser.add_argument("--landmark_ages", 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("--min_history_events", type=int, default=1)
    parser.add_argument("--dataset_subset_size", type=int, default=0)
    parser.add_argument(
        "--readout_batch_size",
        type=int,
        default=8192,
        help=(
            "Number of readout points forwarded together inside each worker. "
            "Increase this to improve GPU utilization if memory allows."
        ),
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=4,
        help="DataLoader workers per GPU process for readout job generation.",
    )
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument(
        "--devices",
        type=str,
        default=None,
        help=(
            "Comma-separated devices for data-parallel BI computation, e.g. "
            "'cuda:0,cuda:1'. Overrides --device when provided."
        ),
    )
    parser.add_argument(
        "--mp_start_method",
        type=str,
        default="fork",
        choices=["fork", "forkserver"],
        help="Multiprocessing start method for Linux multi-GPU runs.",
    )
    parser.add_argument("--functional_weight_col", type=str,
                        default="hfrm_normalized_weight")
    args = parser.parse_args()

    run_path = Path(args.run_path)
    mapping_specs = _load_mapping_specs(args)
    devices = _parse_devices(args)

    initial_device = "cpu" if len(devices) > 1 else devices[0]
    ctx = load_burden_context(run_path, device=initial_device)
    matrices = []
    for spec in mapping_specs:
        A, disease_ids, category_meta = _build_burden_matrix_from_mapping(
            spec["mapping_csv"],
            token_col=spec["token_col"],
            category_id_col=spec["category_id_col"],
            category_col=spec["category_col"],
            key_area_col=spec["key_area_col"],
            weight_col=spec["weight_col"],
        )
        matrices.append(
            {
                "burden_type": spec["burden_type"],
                "A": A,
                "disease_ids": disease_ids,
                "category_meta": category_meta,
            }
        )
    union_disease_ids, matrices = _attach_union_projection(matrices)
    landmark_ages = _parse_landmark_ages(args)
    horizon = _parse_horizon(args.horizon)
    eval_split = str(args.eval_split).lower()
    if eval_split in {"valid", "validation"}:
        eval_split = "val"
    subset_indices = _config_split_indices(
        len(ctx.dataset),
        ctx.cfg,
        eval_split,
        int(args.dataset_subset_size),
    )
    rows = _eligible_landmark_rows(
        ctx.dataset,
        subset_indices,
        landmark_ages,
        min_history_events=int(args.min_history_events),
    )
    if not rows:
        raise RuntimeError(
            "No eligible landmark rows. Check eval split, landmark ages, and min_history_events."
        )

    output_path = Path(args.output_path) if args.output_path else (
        run_path
        / f"burden_index_{eval_split}_{args.formed_mode}_{_format_horizon_for_filename(horizon)}.csv"
    )
    output_path.parent.mkdir(parents=True, exist_ok=True)

    total_readout_jobs = 0
    output_written = False
    shard_paths: list[Path] = []
    total_timings = {
        "build_readout_sec": 0.0,
        "forward_sec": 0.0,
        "reduce_sec": 0.0,
        "write_csv_sec": 0.0,
    }
    row_chunks = _split_rows_for_devices(
        rows,
        devices,
        formed_mode=args.formed_mode,
        horizon=horizon,
    )
    for device, chunk_rows in row_chunks:
        estimated_jobs = sum(
            _estimate_readout_jobs_for_row(
                row,
                formed_mode=args.formed_mode,
                horizon=horizon,
            )
            for row in chunk_rows
        )
        print(
            f"Assigned {len(chunk_rows)} rows / ~{estimated_jobs} readout jobs to {device}",
            flush=True,
        )
    if len(row_chunks) == 1:
        all_rows, total_readout_jobs, timings = _compute_bi_from_streamed_readouts(
            rows=rows,
            horizon=horizon,
            matrices=matrices,
            union_disease_ids=union_disease_ids,
            formed_mode=args.formed_mode,
            readout_batch_size=int(args.readout_batch_size),
            num_workers=int(args.num_workers),
            ctx=ctx,
            log_prefix="[BI main]",
        )
        for key, value in timings.items():
            total_timings[key] += float(value)
        write_start = time.perf_counter()
        _write_rows_csv(all_rows, output_path)
        total_timings["write_csv_sec"] = time.perf_counter() - write_start
        output_written = True
    else:
        # The main-process context is only needed to build the dataset and rows.
        # Workers load their own model copy on the assigned device.
        del ctx
        payloads = [
            {
                "device": device,
                "run_path": str(run_path),
                "shard_path": str(
                    output_path.with_name(
                        f"{output_path.stem}.part{part_idx:03d}{output_path.suffix}"
                    )
                ),
                "row_specs": [_row_to_worker_spec(row) for row in chunk_rows],
                "horizon": horizon,
                "matrices": matrices,
                "union_disease_ids": union_disease_ids,
                "readout_batch_size": int(args.readout_batch_size),
                "num_workers": int(args.num_workers),
                "formed_mode": args.formed_mode,
            }
            for part_idx, (device, chunk_rows) in enumerate(row_chunks)
        ]
        with ProcessPoolExecutor(
            max_workers=len(payloads),
            mp_context=mp.get_context(args.mp_start_method),
        ) as executor:
            futures = [executor.submit(_compute_chunk_worker, p) for p in payloads]
            for future in tqdm(
                as_completed(futures),
                total=len(futures),
                desc="Computing BI chunks",
                dynamic_ncols=True,
            ):
                result = future.result()
                shard_paths.append(Path(result["shard_path"]))
                total_readout_jobs += int(result["readout_jobs"])
                for key, value in result["timings"].items():
                    total_timings[key] += float(value)

        write_start = time.perf_counter()
        _concat_csv_shards(sorted(shard_paths), output_path)
        total_timings["write_csv_sec"] += time.perf_counter() - write_start
        output_written = True

    if not output_written:
        raise RuntimeError("BI output was not written.")
    print(f"Run path: {run_path}")
    print(f"Eval split: {eval_split}")
    print(f"Horizon: {horizon:g}")
    print(f"Landmark rows: {len(rows)}")
    print(f"Readout jobs: {total_readout_jobs}")
    print(f"Readout batch size per worker: {int(args.readout_batch_size)}")
    print(f"DataLoader workers per GPU process: {int(args.num_workers)}")
    print(f"Multiprocessing start method: {args.mp_start_method}")
    print(
        "Timing seconds: "
        f"build_readout={total_timings['build_readout_sec']:.2f}, "
        f"forward={total_timings['forward_sec']:.2f}, "
        f"reduce={total_timings['reduce_sec']:.2f}, "
        f"write_csv={total_timings['write_csv_sec']:.2f}"
    )
    print(f"Devices: {', '.join(str(d) for d, _ in row_chunks)}")
    for matrix in matrices:
        print(
            f"{matrix['burden_type']} dimensions: {matrix['A'].shape[0]}, "
            f"mapped disease tokens: {matrix['A'].shape[1]}"
        )
    print(f"Union disease tokens evaluated once per sample: {union_disease_ids.size}")
    print(f"Output: {output_path}")


if __name__ == "__main__":
    main()