DeepHealth/burden_index.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, Sequence

import numpy as np
import torch
import torch.nn.functional as F

from evaluate_auc_v2 import (
    build_model_from_dataset,
    load_checkpoint_state_dict,
    load_json_config,
    load_model_state,
    resolve_dist_mode_for_checkpoint,
    validate_dataset_metadata,
)
from eval_data import load_sequence_eval_dataset
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX


FormedBurdenMode = Literal["observed", "model_weighted"]


@dataclass(frozen=True)
class BurdenContext:
    model: torch.nn.Module
    dataset: Any
    cfg: dict[str, Any]
    dist_mode: str
    device: torch.device
    run_path: Path


@dataclass(frozen=True)
class DiseaseBurdenResult:
    disease_id: int
    formed: float
    future: float
    total: float
    formed_mode: str
    horizon: float


@dataclass(frozen=True)
class BurdenIndexResult:
    historical: np.ndarray
    future: np.ndarray
    total: np.ndarray
    disease_ids: np.ndarray
    formed: np.ndarray
    disease_future: np.ndarray
    disease_total: np.ndarray
    formed_mode: str
    horizon: float


def load_burden_context(
    run_path: str | Path,
    *,
    device: str | torch.device | None = None,
) -> BurdenContext:
    run_path = Path(run_path)
    config_path = run_path / "train_config.json"
    model_ckpt_path = run_path / "best_model.pt"
    if not config_path.exists():
        raise FileNotFoundError(f"train_config.json not found in {run_path}")
    if not model_ckpt_path.exists():
        raise FileNotFoundError(f"best_model.pt not found in {run_path}")

    cfg = load_json_config(config_path)
    model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
    if model_target_mode == "next_token":
        raise RuntimeError(
            "Burden Index computation requires an all_future checkpoint because "
            "it uses p_d(h, Delta). The provided run is model_target_mode='next_token'."
        )
    if model_target_mode != "all_future":
        raise ValueError(
            "train_config.json model_target_mode must be all_future for burden "
            f"computation, got {model_target_mode!r}."
        )

    device_obj = torch.device(
        device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
    )
    if device_obj.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError(f"Requested device {device_obj}, but CUDA is not available.")

    data_prefix = cfg.get("data_prefix", "ukb")
    labels_file = cfg.get("labels_file", "labels.csv")
    extra_info_types = cfg.get("extra_info_types", None)
    dataset = load_sequence_eval_dataset(
        model_target_mode="all_future",
        data_prefix=data_prefix,
        labels_file=labels_file,
        no_event_interval_years=float(cfg.get("no_event_interval_years", 5.0)),
        include_no_event_in_uts_target=bool(
            cfg.get("include_no_event_in_uts_target", False)
        ),
        min_history_events=int(cfg.get("all_future_min_history_events", 1)),
        min_future_events=int(cfg.get("all_future_min_future_events", 1)),
        extra_info_types=extra_info_types,
    )
    validate_dataset_metadata(dataset, cfg)

    state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu")
    dist_mode = resolve_dist_mode_for_checkpoint(
        str(cfg.get("dist_mode", "exponential")),
        state_dict,
    )
    if dist_mode not in {"exponential", "weibull", "mixed"}:
        raise ValueError(f"Unsupported dist_mode={dist_mode!r}")

    cfg_model = dict(cfg)
    cfg_model["dist_mode"] = dist_mode
    args = _ConfigNamespace()
    model = build_model_from_dataset(args, cfg_model, dataset)
    load_model_state(model, state_dict)
    model.eval().to(device_obj)

    return BurdenContext(
        model=model,
        dataset=dataset,
        cfg=cfg,
        dist_mode=dist_mode,
        device=device_obj,
        run_path=run_path,
    )


@torch.inference_mode()
def compute_disease_burden(
    *,
    run_path: str | Path,
    disease_id: int,
    event_seq: Sequence[int] | np.ndarray,
    time_seq: Sequence[float] | np.ndarray,
    sex: int,
    other_type: Sequence[int] | np.ndarray,
    other_value: Sequence[float] | np.ndarray,
    other_value_kind: Sequence[int] | np.ndarray,
    other_time: Sequence[float] | np.ndarray,
    t_query: float,
    horizon: float,
    formed_mode: FormedBurdenMode,
    device: str | torch.device | None = None,
    context: BurdenContext | None = None,
) -> DiseaseBurdenResult:
    ctx = context or load_burden_context(run_path, device=device)
    disease_ids = np.asarray([int(disease_id)], dtype=np.int64)
    formed, future_prob = _compute_disease_components(
        ctx=ctx,
        disease_ids=disease_ids,
        event_seq=event_seq,
        time_seq=time_seq,
        sex=sex,
        other_type=other_type,
        other_value=other_value,
        other_value_kind=other_value_kind,
        other_time=other_time,
        t_query=float(t_query),
        horizon=float(horizon),
        formed_mode=formed_mode,
    )
    future = (1.0 - formed) * future_prob
    total = formed + future
    return DiseaseBurdenResult(
        disease_id=int(disease_id),
        formed=float(formed[0]),
        future=float(future[0]),
        total=float(total[0]),
        formed_mode=str(formed_mode),
        horizon=float(horizon),
    )


@torch.inference_mode()
def compute_burden_index(
    *,
    run_path: str | Path,
    burden_matrix: np.ndarray,
    disease_ids: Sequence[int] | np.ndarray,
    event_seq: Sequence[int] | np.ndarray,
    time_seq: Sequence[float] | np.ndarray,
    sex: int,
    other_type: Sequence[int] | np.ndarray,
    other_value: Sequence[float] | np.ndarray,
    other_value_kind: Sequence[int] | np.ndarray,
    other_time: Sequence[float] | np.ndarray,
    t_query: float,
    horizon: float,
    formed_mode: FormedBurdenMode,
    device: str | torch.device | None = None,
    context: BurdenContext | None = None,
) -> BurdenIndexResult:
    ctx = context or load_burden_context(run_path, device=device)
    disease_ids_arr = np.asarray(disease_ids, dtype=np.int64)
    if disease_ids_arr.ndim != 1 or disease_ids_arr.size == 0:
        raise ValueError("disease_ids must be a non-empty 1D sequence.")

    A = np.asarray(burden_matrix, dtype=np.float64)
    if A.ndim != 2:
        raise ValueError(f"burden_matrix must be 2D, got shape {A.shape}")
    if A.shape[1] != disease_ids_arr.size:
        raise ValueError(
            "burden_matrix columns must match disease_ids length, got "
            f"{A.shape[1]} columns vs {disease_ids_arr.size} disease ids."
        )

    formed, future_prob = _compute_disease_components(
        ctx=ctx,
        disease_ids=disease_ids_arr,
        event_seq=event_seq,
        time_seq=time_seq,
        sex=sex,
        other_type=other_type,
        other_value=other_value,
        other_value_kind=other_value_kind,
        other_time=other_time,
        t_query=float(t_query),
        horizon=float(horizon),
        formed_mode=formed_mode,
    )
    disease_future = (1.0 - formed) * future_prob
    disease_total = formed + disease_future

    historical = A @ formed
    future = A @ disease_future
    total = A @ disease_total
    return BurdenIndexResult(
        historical=historical,
        future=future,
        total=total,
        disease_ids=disease_ids_arr.copy(),
        formed=formed,
        disease_future=disease_future,
        disease_total=disease_total,
        formed_mode=str(formed_mode),
        horizon=float(horizon),
    )


class _ConfigNamespace:
    def __getattr__(self, _name: str) -> None:
        return None


def _compute_disease_components(
    *,
    ctx: BurdenContext,
    disease_ids: np.ndarray,
    event_seq: Sequence[int] | np.ndarray,
    time_seq: Sequence[float] | np.ndarray,
    sex: int,
    other_type: Sequence[int] | np.ndarray,
    other_value: Sequence[float] | np.ndarray,
    other_value_kind: Sequence[int] | np.ndarray,
    other_time: Sequence[float] | np.ndarray,
    t_query: float,
    horizon: float,
    formed_mode: FormedBurdenMode,
) -> tuple[np.ndarray, np.ndarray]:
    formed_mode = _validate_formed_mode(formed_mode)
    disease_ids = np.asarray(disease_ids, dtype=np.int64)
    _validate_disease_ids(ctx, disease_ids)
    event_seq_arr, time_seq_arr = _validate_event_inputs(event_seq, time_seq)
    other_type_arr, other_value_arr, other_value_kind_arr, other_time_arr = (
        _validate_other_inputs(other_type, other_value, other_value_kind, other_time)
    )
    if horizon < 0:
        raise ValueError(f"horizon must be non-negative, got {horizon}")

    if formed_mode == "observed":
        formed = _observed_formed_burden(
            disease_ids=disease_ids,
            event_seq=event_seq_arr,
            time_seq=time_seq_arr,
            t_query=t_query,
        )
    else:
        formed = _model_weighted_formed_burden(
            ctx=ctx,
            disease_ids=disease_ids,
            event_seq=event_seq_arr,
            time_seq=time_seq_arr,
            sex=sex,
            other_type=other_type_arr,
            other_value=other_value_arr,
            other_value_kind=other_value_kind_arr,
            other_time=other_time_arr,
            t_query=t_query,
        )

    hidden_query = _query_hidden(
        ctx=ctx,
        event_seq=event_seq_arr,
        time_seq=time_seq_arr,
        sex=sex,
        other_type=other_type_arr,
        other_value=other_value_arr,
        other_value_kind=other_value_kind_arr,
        other_time=other_time_arr,
        query_times=np.asarray([t_query], dtype=np.float32),
    )
    future_prob = _probabilities_from_hidden(
        ctx=ctx,
        hidden=hidden_query,
        disease_ids=disease_ids,
        deltas=np.asarray([horizon], dtype=np.float32),
    )[0]
    return formed.astype(np.float64), future_prob.astype(np.float64)


def _model_weighted_formed_burden(
    *,
    ctx: BurdenContext,
    disease_ids: np.ndarray,
    event_seq: np.ndarray,
    time_seq: np.ndarray,
    sex: int,
    other_type: np.ndarray,
    other_value: np.ndarray,
    other_value_kind: np.ndarray,
    other_time: np.ndarray,
    t_query: float,
) -> np.ndarray:
    grid = _build_readout_grid(
        event_seq=event_seq,
        time_seq=time_seq,
        other_type=other_type,
        other_time=other_time,
        t_query=t_query,
    )
    if grid.size == 0:
        return np.zeros(disease_ids.size, dtype=np.float64)

    end_times = np.concatenate([grid[1:], np.asarray([t_query], dtype=np.float32)])
    deltas = np.maximum(end_times - grid, 0.0).astype(np.float32)
    if np.all(deltas <= 0):
        return np.zeros(disease_ids.size, dtype=np.float64)

    hidden = _query_hidden(
        ctx=ctx,
        event_seq=event_seq,
        time_seq=time_seq,
        sex=sex,
        other_type=other_type,
        other_value=other_value,
        other_value_kind=other_value_kind,
        other_time=other_time,
        query_times=grid.astype(np.float32),
    )
    interval_prob = _probabilities_from_hidden(
        ctx=ctx,
        hidden=hidden,
        disease_ids=disease_ids,
        deltas=deltas,
    ).astype(np.float64)
    survival = np.prod(1.0 - np.clip(interval_prob, 0.0, 1.0), axis=0)
    return 1.0 - survival


def _observed_formed_burden(
    *,
    disease_ids: np.ndarray,
    event_seq: np.ndarray,
    time_seq: np.ndarray,
    t_query: float,
) -> np.ndarray:
    valid = (
        (time_seq <= np.float32(t_query))
        & (event_seq > PAD_IDX)
        & (event_seq != CHECKUP_IDX)
        & (event_seq != NO_EVENT_IDX)
    )
    observed = set(int(x) for x in event_seq[valid].tolist())
    return np.asarray([1.0 if int(d) in observed else 0.0 for d in disease_ids])


def _build_readout_grid(
    *,
    event_seq: np.ndarray,
    time_seq: np.ndarray,
    other_type: np.ndarray,
    other_time: np.ndarray,
    t_query: float,
) -> np.ndarray:
    event_mask = (event_seq > PAD_IDX) & (time_seq <= np.float32(t_query))
    other_mask = (other_type > 0) & (other_time <= np.float32(t_query))
    times = np.concatenate(
        [
            time_seq[event_mask].astype(np.float32, copy=False),
            other_time[other_mask].astype(np.float32, copy=False),
        ]
    )
    if times.size == 0:
        return np.zeros(0, dtype=np.float32)
    return np.unique(times)


def _query_hidden(
    *,
    ctx: BurdenContext,
    event_seq: np.ndarray,
    time_seq: np.ndarray,
    sex: int,
    other_type: np.ndarray,
    other_value: np.ndarray,
    other_value_kind: np.ndarray,
    other_time: np.ndarray,
    query_times: np.ndarray,
) -> torch.Tensor:
    if query_times.ndim != 1:
        raise ValueError("query_times must be 1D.")
    batch_size = int(query_times.size)
    if batch_size == 0:
        return torch.empty(0, ctx.model.n_embd, device=ctx.device)

    event = torch.from_numpy(event_seq[None, :].repeat(batch_size, axis=0)).long()
    times = torch.from_numpy(time_seq[None, :].repeat(batch_size, axis=0)).float()
    other_t = torch.from_numpy(other_type[None, :].repeat(batch_size, axis=0)).long()
    other_v = torch.from_numpy(other_value[None, :].repeat(batch_size, axis=0)).float()
    other_k = torch.from_numpy(
        other_value_kind[None, :].repeat(batch_size, axis=0)
    ).long()
    other_tm = torch.from_numpy(other_time[None, :].repeat(batch_size, axis=0)).float()
    sex_t = torch.full((batch_size,), int(sex), dtype=torch.long)
    tq = torch.from_numpy(query_times.astype(np.float32, copy=False)).float()

    event = event.to(ctx.device)
    times = times.to(ctx.device)
    other_t = other_t.to(ctx.device)
    other_v = other_v.to(ctx.device)
    other_k = other_k.to(ctx.device)
    other_tm = other_tm.to(ctx.device)
    sex_t = sex_t.to(ctx.device)
    tq = tq.to(ctx.device)

    return ctx.model(
        event_seq=event,
        time_seq=times,
        sex=sex_t,
        padding_mask=event > PAD_IDX,
        t_query=tq,
        other_type=other_t,
        other_value=other_v,
        other_value_kind=other_k,
        other_time=other_tm,
        target_mode="all_future",
    )


def _probabilities_from_hidden(
    *,
    ctx: BurdenContext,
    hidden: torch.Tensor,
    disease_ids: np.ndarray,
    deltas: np.ndarray,
) -> np.ndarray:
    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 = F.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)

    prob = -torch.expm1(-rate * exposure)
    return prob.detach().cpu().numpy().astype(np.float64, copy=False)


def _validate_formed_mode(mode: str) -> FormedBurdenMode:
    if mode not in {"observed", "model_weighted"}:
        raise ValueError(
            "formed_mode must be either 'observed' or 'model_weighted', "
            f"got {mode!r}."
        )
    return mode  # type: ignore[return-value]


def _validate_disease_ids(ctx: BurdenContext, disease_ids: np.ndarray) -> None:
    if disease_ids.ndim != 1 or disease_ids.size == 0:
        raise ValueError("disease_ids must be a non-empty 1D array.")
    vocab_size = int(getattr(ctx.model, "vocab_size", ctx.model.risk_head.out_features))
    if np.any(disease_ids < 0) or np.any(disease_ids >= vocab_size):
        raise ValueError(f"disease_ids must be in [0, {vocab_size}), got {disease_ids}")


def _validate_event_inputs(
    event_seq: Sequence[int] | np.ndarray,
    time_seq: Sequence[float] | np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    events = np.asarray(event_seq, dtype=np.int64)
    times = np.asarray(time_seq, dtype=np.float32)
    if events.ndim != 1 or times.ndim != 1:
        raise ValueError("event_seq and time_seq must be 1D.")
    if events.shape != times.shape:
        raise ValueError(
            f"event_seq and time_seq must have the same shape, got {events.shape} vs {times.shape}"
        )
    if events.size == 0:
        raise ValueError("event_seq must contain at least one token.")
    if not np.all(np.isfinite(times)):
        raise ValueError("time_seq contains non-finite values.")
    return events, times


def _validate_other_inputs(
    other_type: Sequence[int] | np.ndarray,
    other_value: Sequence[float] | np.ndarray,
    other_value_kind: Sequence[int] | np.ndarray,
    other_time: Sequence[float] | np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    typ = np.asarray(other_type, dtype=np.int64)
    val = np.asarray(other_value, dtype=np.float32)
    kind = np.asarray(other_value_kind, dtype=np.int64)
    tm = np.asarray(other_time, dtype=np.float32)
    if not (typ.shape == val.shape == kind.shape == tm.shape):
        raise ValueError(
            "other_type, other_value, other_value_kind, and other_time must "
            f"have the same shape, got {typ.shape}, {val.shape}, {kind.shape}, {tm.shape}."
        )
    if typ.ndim != 1:
        raise ValueError("other_* inputs must be 1D.")
    if not np.all(np.isfinite(tm)):
        raise ValueError("other_time contains non-finite values.")
    return typ, val, kind, tm