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 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 _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)


def _iter_readout_batches(
    n: int,
    batch_size: int,
) -> Iterable[slice]:
    batch_size = max(1, int(batch_size))
    for start in range(0, n, batch_size):
        yield slice(start, min(start + batch_size, n))


@torch.inference_mode()
def _query_hidden_jobs(
    *,
    ctx: Any,
    jobs: list[tuple[dict[str, Any], float]],
) -> torch.Tensor:
    if not jobs:
        return torch.empty(0, ctx.model.n_embd, device=ctx.device)

    batch_size = len(jobs)
    max_event_len = max(int(np.asarray(row["event_seq"]).size) for row, _ in jobs)
    max_other_len = max(int(np.asarray(row["other_type"]).size) for row, _ in jobs)

    event = np.full((batch_size, max_event_len), PAD_IDX, dtype=np.int64)
    time = np.zeros((batch_size, max_event_len), dtype=np.float32)
    other_type = np.zeros((batch_size, max_other_len), dtype=np.int64)
    other_value = np.zeros((batch_size, max_other_len), dtype=np.float32)
    other_value_kind = np.zeros((batch_size, max_other_len), dtype=np.int64)
    other_time = np.zeros((batch_size, max_other_len), dtype=np.float32)
    sex = np.zeros(batch_size, dtype=np.int64)
    query_times = np.zeros(batch_size, dtype=np.float32)

    for i, (row, query_time) in enumerate(jobs):
        event_seq = np.asarray(row["event_seq"], dtype=np.int64)
        time_seq = np.asarray(row["time_seq"], dtype=np.float32)
        other_type_seq = np.asarray(row["other_type"], dtype=np.int64)
        other_value_seq = np.asarray(row["other_value"], dtype=np.float32)
        other_value_kind_seq = np.asarray(row["other_value_kind"], dtype=np.int64)
        other_time_seq = np.asarray(row["other_time"], dtype=np.float32)

        event[i, : event_seq.size] = event_seq
        time[i, : time_seq.size] = time_seq
        other_type[i, : other_type_seq.size] = other_type_seq
        other_value[i, : other_value_seq.size] = other_value_seq
        other_value_kind[i, : other_value_kind_seq.size] = other_value_kind_seq
        other_time[i, : other_time_seq.size] = other_time_seq
        sex[i] = int(row["sex"])
        query_times[i] = np.float32(query_time)

    event_t = torch.from_numpy(event).long().to(ctx.device)
    return ctx.model(
        event_seq=event_t,
        time_seq=torch.from_numpy(time).float().to(ctx.device),
        sex=torch.from_numpy(sex).long().to(ctx.device),
        padding_mask=event_t > PAD_IDX,
        t_query=torch.from_numpy(query_times).float().to(ctx.device),
        other_type=torch.from_numpy(other_type).long().to(ctx.device),
        other_value=torch.from_numpy(other_value).float().to(ctx.device),
        other_value_kind=torch.from_numpy(other_value_kind).long().to(ctx.device),
        other_time=torch.from_numpy(other_time).float().to(ctx.device),
        target_mode="all_future",
    )


def _build_readout_table(
    *,
    rows: list[dict[str, Any]],
    formed_mode: str,
    horizon: float,
) -> dict[str, Any]:
    jobs: list[tuple[dict[str, Any], float]] = []
    row_indices: list[int] = []
    kinds: list[str] = []
    deltas: list[float] = []

    if formed_mode not in {"observed", "model_weighted"}:
        raise ValueError(f"Unknown formed_mode={formed_mode!r}")

    for row_idx, row in enumerate(rows):
        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)]
                )
                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()):
                    jobs.append((row, float(query_time)))
                    row_indices.append(row_idx)
                    kinds.append("formed")
                    deltas.append(float(delta))
        if horizon > 0:
            jobs.append((row, float(row["t_query"])))
            row_indices.append(row_idx)
            kinds.append("future")
            deltas.append(float(horizon))

    return {
        "jobs": jobs,
        "row_indices": np.asarray(row_indices, dtype=np.int64),
        "kinds": np.asarray(kinds, dtype=object),
        "deltas": np.asarray(deltas, dtype=np.float32),
    }


@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(
    *,
    ctx: Any,
    readout_table: dict[str, Any],
    union_disease_ids: np.ndarray,
    readout_batch_size: int,
) -> torch.Tensor:
    jobs = readout_table["jobs"]
    if not jobs:
        return torch.empty((0, union_disease_ids.size), dtype=torch.float32, device=ctx.device)

    out = torch.empty((len(jobs), union_disease_ids.size), dtype=torch.float32, device=ctx.device)
    deltas = np.asarray(readout_table["deltas"], dtype=np.float32)
    for slc in _iter_readout_batches(len(jobs), readout_batch_size):
        hidden = _query_hidden_jobs(ctx=ctx, jobs=jobs[slc])
        out[slc] = _probabilities_from_hidden_torch(
            ctx=ctx,
            hidden=hidden,
            disease_ids=union_disease_ids,
            deltas=deltas[slc],
        ).to(dtype=out.dtype)
    return out


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 _reduce_readout_table_to_bi_rows(
    *,
    rows: list[dict[str, Any]],
    horizon: float,
    matrices: list[dict[str, Any]],
    union_disease_ids: np.ndarray,
    formed_mode: str,
    readout_table: dict[str, Any],
    readout_prob: torch.Tensor,
    ctx: Any,
) -> list[dict[str, Any]]:
    if formed_mode == "observed":
        formed_by_row = torch.as_tensor(
            _observed_formed_for_rows(
                rows=rows,
                union_disease_ids=union_disease_ids,
            ),
            dtype=readout_prob.dtype,
            device=ctx.device,
        )
    elif formed_mode == "model_weighted":
        formed_by_row = torch.zeros(
            (len(rows), union_disease_ids.size),
            dtype=readout_prob.dtype,
            device=ctx.device,
        )
    else:
        raise ValueError(f"Unknown formed_mode={formed_mode!r}")

    future_prob_by_row = torch.zeros(
        (len(rows), union_disease_ids.size),
        dtype=readout_prob.dtype,
        device=ctx.device,
    )
    row_indices = torch.as_tensor(
        np.asarray(readout_table["row_indices"], dtype=np.int64),
        dtype=torch.long,
        device=ctx.device,
    )
    kinds = np.asarray(readout_table["kinds"], dtype=object)
    if formed_mode == "model_weighted":
        survival_by_row = torch.ones(
            (len(rows), union_disease_ids.size),
            dtype=readout_prob.dtype,
            device=ctx.device,
        )
    else:
        survival_by_row = None

    if readout_prob.numel() > 0:
        kind_is_formed = torch.as_tensor(
            np.asarray(kinds == "formed", dtype=np.bool_),
            dtype=torch.bool,
            device=ctx.device,
        )
        kind_is_future = torch.as_tensor(
            np.asarray(kinds == "future", dtype=np.bool_),
            dtype=torch.bool,
            device=ctx.device,
        )
        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]

    if survival_by_row is not None:
        formed_by_row = 1.0 - survival_by_row

    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=readout_prob.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):
        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": float(row["landmark_age"]),
                        "t_query": float(row["t_query"]),
                        "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 _compute_bi_from_readout_table(
    *,
    rows: list[dict[str, Any]],
    horizon: float,
    matrices: list[dict[str, Any]],
    union_disease_ids: np.ndarray,
    formed_mode: str,
    readout_batch_size: int,
    ctx: Any,
) -> 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}")
    t0 = time.perf_counter()
    readout_table = _build_readout_table(
        rows=rows,
        formed_mode=formed_mode,
        horizon=horizon,
    )
    t1 = time.perf_counter()
    readout_prob = _readout_probabilities(
        ctx=ctx,
        readout_table=readout_table,
        union_disease_ids=union_disease_ids,
        readout_batch_size=readout_batch_size,
    )
    if ctx.device.type == "cuda":
        torch.cuda.synchronize(ctx.device)
    t2 = time.perf_counter()
    rows_out = _reduce_readout_table_to_bi_rows(
        rows=rows,
        horizon=horizon,
        matrices=matrices,
        union_disease_ids=union_disease_ids,
        formed_mode=formed_mode,
        readout_table=readout_table,
        readout_prob=readout_prob,
        ctx=ctx,
    )
    if ctx.device.type == "cuda":
        torch.cuda.synchronize(ctx.device)
    t3 = time.perf_counter()
    timings = {
        "build_readout_sec": t1 - t0,
        "forward_sec": t2 - t1,
        "reduce_sec": t3 - t2,
    }
    return rows_out, len(readout_table["jobs"]), timings


def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
    device = payload["device"]
    run_path = Path(payload["run_path"])
    ctx = load_burden_context(run_path, device=device)
    out, readout_jobs, timings = _compute_bi_from_readout_table(
        rows=payload["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"]),
        ctx=ctx,
    )
    return {"rows": out, "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("--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("--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)

    all_rows: list[dict[str, Any]] = []
    total_readout_jobs = 0
    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,
    )
    if len(row_chunks) == 1:
        all_rows, total_readout_jobs, timings = _compute_bi_from_readout_table(
            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),
            ctx=ctx,
        )
        for key, value in timings.items():
            total_timings[key] += float(value)
    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),
                "rows": chunk_rows,
                "horizon": horizon,
                "matrices": matrices,
                "union_disease_ids": union_disease_ids,
                "readout_batch_size": int(args.readout_batch_size),
                "formed_mode": args.formed_mode,
            }
            for device, chunk_rows in row_chunks
        ]
        with ProcessPoolExecutor(
            max_workers=len(payloads),
            mp_context=mp.get_context("spawn"),
        ) 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()
                all_rows.extend(result["rows"])
                total_readout_jobs += int(result["readout_jobs"])
                for key, value in result["timings"].items():
                    total_timings[key] += float(value)

    write_start = time.perf_counter()
    out_df = pd.DataFrame(all_rows)
    out_df.to_csv(output_path, index=False)
    total_timings["write_csv_sec"] = time.perf_counter() - write_start
    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(
        "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()