DeepHealthExpo/dataset.py

# dataset.py
from __future__ import annotations

from pathlib import Path
from typing import Dict, Iterable, List, Literal, Optional, Tuple

import numpy as np
import pandas as pd
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset

from targets import (
    CHECKUP_IDX,
    DAYS_PER_YEAR,
    NO_EVENT_IDX,
    PAD_IDX,
    build_all_targets,
)


ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR
DAILY_EXPOSURE_SHAPE = (1826, 4)
MONTHLY_EXPOSURE_SHAPE = (241, 2)


class ExposureCache:
    """Random-access view over files produced by prepare_exposure_cache.py."""

    def __init__(self, cache_dir: str | Path):
        cache_dir = Path(cache_dir)
        self.cache_dir = cache_dir
        eid_path = cache_dir / "exposure_eid.npy"
        token_path = cache_dir / "exposure_token.npy"
        onset_date_path = cache_dir / "exposure_onset_date.npy"
        if not (eid_path.is_file() and token_path.is_file() and onset_date_path.is_file()):
            raise FileNotFoundError(
                "Exposure cache must contain exposure_eid.npy, "
                "exposure_token.npy, and exposure_onset_date.npy. "
                "Regenerate it with the current prepare_exposure_cache.py."
            )
        self.eids = np.load(eid_path, mmap_mode="r")
        self.raw_tokens = np.load(token_path, mmap_mode="r")
        self.onset_dates = np.load(onset_date_path, mmap_mode="r")
        self.daily = np.load(cache_dir / "exposure_daily.npy", mmap_mode="r")
        self.monthly = np.load(cache_dir / "exposure_monthly.npy", mmap_mode="r")
        quality_path = cache_dir / "exposure_quality.npy"
        self.quality = np.load(quality_path, mmap_mode="r") if quality_path.is_file() else None

        if self.daily.ndim != 3 or self.daily.shape[1:] != DAILY_EXPOSURE_SHAPE:
            raise ValueError(
                f"exposure_daily.npy must have shape (N, {DAILY_EXPOSURE_SHAPE[0]}, "
                f"{DAILY_EXPOSURE_SHAPE[1]}), got {self.daily.shape}"
            )
        if self.monthly.ndim != 3 or self.monthly.shape[1:] != MONTHLY_EXPOSURE_SHAPE:
            raise ValueError(
                f"exposure_monthly.npy must have shape (N, {MONTHLY_EXPOSURE_SHAPE[0]}, "
                f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}"
            )
        n_rows = len(self.eids)
        if (
            len(self.raw_tokens) != n_rows
            or len(self.onset_dates) != n_rows
            or self.daily.shape[0] != n_rows
            or self.monthly.shape[0] != n_rows
        ):
            raise ValueError("Exposure cache metadata/daily/monthly row counts do not match")

        self._key_to_index: dict[tuple[int, int, int], int] | None = None

    def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None:
        keys: dict[tuple[int, int, int], int] = {}
        eids = np.asarray(self.eids, dtype=np.int64)
        tokens = np.asarray(self.raw_tokens, dtype=np.int64)
        onset_dates = np.asarray(self.onset_dates, dtype="datetime64[D]")
        for idx, (eid, token, onset_date) in enumerate(zip(eids, tokens, onset_dates)):
            birth_date = birth_date_by_eid.get(int(eid))
            if birth_date is None or np.isnat(onset_date) or np.isnat(birth_date):
                continue
            age_days = int((onset_date - birth_date).astype("timedelta64[D]").astype(np.int64))
            if age_days < 0:
                continue
            keys[(int(eid), int(token), age_days)] = idx
        self._key_to_index = keys

    def lookup_indices(self, eid: int, raw_tokens: np.ndarray, age_days: np.ndarray) -> np.ndarray:
        if self._key_to_index is None:
            raise RuntimeError("ExposureCache.build_age_index must be called before lookup")
        out = np.full(len(raw_tokens), -1, dtype=np.int64)
        real = raw_tokens > 1
        if not np.any(real):
            return out
        real_pos = np.nonzero(real)[0]
        out[real_pos] = [
            self._key_to_index.get((int(eid), int(raw_tokens[pos]), int(round(float(age_days[pos])))), -1)
            for pos in real_pos
        ]
        return out

    def daily_window(self, index: int) -> np.ndarray:
        if index < 0:
            return np.full(DAILY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
        return np.asarray(self.daily[index], dtype=np.float32)

    def monthly_window(self, index: int) -> np.ndarray:
        if index < 0:
            return np.full(MONTHLY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
        return np.asarray(self.monthly[index], dtype=np.float32)


def load_label_vocab(
    labels_file: str,
    include_no_event: bool = True,
) -> Tuple[Dict[str, int], Dict[int, str]]:
    label_id_to_code: Dict[int, str] = {
        PAD_IDX: "<PAD>",
        CHECKUP_IDX: "<CHECKUP>",
    }
    if include_no_event:
        label_id_to_code[NO_EVENT_IDX] = "<NO_EVENT>"

    offset = NO_EVENT_IDX + 1 if include_no_event else CHECKUP_IDX + 1
    label_code_to_id: Dict[str, int] = {}
    with open(labels_file, encoding="utf-8") as f:
        for i, line in enumerate(f):
            parts = line.strip().split()
            if not parts:
                continue
            idx = offset + i
            code = parts[0]
            label_code_to_id[code] = idx
            label_id_to_code[idx] = code
    return label_code_to_id, label_id_to_code


def _insert_gap_no_event_tokens(
    times_days: np.ndarray,
    labels: np.ndarray,
    interval_years: float = 5.0,
) -> Tuple[np.ndarray, np.ndarray]:
    if len(times_days) < 2:
        return times_days, labels

    step_days = interval_years * DAYS_PER_YEAR
    unique_times = np.unique(times_days.astype(np.float64))
    extra_times: List[float] = []

    for i in range(len(unique_times) - 1):
        t_left = float(unique_times[i])
        t_right = float(unique_times[i + 1])
        if t_right - t_left <= step_days:
            continue
        first = np.ceil((t_left + 1e-6) / step_days) * step_days
        t = first
        while t < t_right - 1e-6:
            extra_times.append(t)
            t += step_days

    if not extra_times:
        return times_days, labels

    extra_arr = np.array(extra_times, dtype=np.float32)
    no_event_labels = np.full(len(extra_arr), NO_EVENT_IDX, dtype=np.int64)
    all_times = np.concatenate([times_days.astype(np.float32), extra_arr])
    all_labels = np.concatenate([labels.astype(np.int64), no_event_labels])
    order = np.lexsort((all_labels, all_times))
    return all_times[order], all_labels[order]


class _ExpoBaseDataset(Dataset):
    def __init__(
        self,
        data_prefix: str = "ukb",
        labels_file: str = "labels.csv",
        no_event_interval_years: float = 5.0,
        include_no_event_in_uts_target: bool = False,
        exposure_cache_dir: str | Path | None = None,
        mask_onset_exposure: bool = False,
    ) -> None:
        self.data_prefix = data_prefix
        self.labels_file = labels_file
        self.no_event_interval_years = float(no_event_interval_years)
        self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
        self.exposure_cache = (
            ExposureCache(exposure_cache_dir)
            if exposure_cache_dir is not None
            else None
        )
        self.mask_onset_exposure = bool(mask_onset_exposure)

        self.label_code_to_id, self.label_id_to_code = load_label_vocab(
            labels_file,
            include_no_event=True,
        )

        event_data = np.load(f"{data_prefix}_event_data.npy")
        if event_data.ndim != 2 or event_data.shape[1] < 3:
            raise ValueError(f"event_data must have shape (N, 3+), got {event_data.shape}")
        event_data = event_data[:, :3].copy()
        order = np.lexsort((event_data[:, 2], event_data[:, 1], event_data[:, 0]))
        self.event_data = event_data[order]

        basic_table = pd.read_csv(f"{data_prefix}_basic_info.csv", index_col=0)
        basic_table.index = basic_table.index.astype(np.int64)

        unique_eids = np.unique(self.event_data[:, 0].astype(np.int64))
        basic_table = basic_table.loc[unique_eids]

        self._prepare_sex(basic_table, unique_eids)
        self._prepare_birth_dates(basic_table, unique_eids)
        if self.exposure_cache is not None:
            self.exposure_cache.build_age_index(self.birth_date_mapping)

        max_id_in_vocab = max(self.label_id_to_code.keys())
        max_id_in_data = int(self.event_data[:, 2].max()) if len(self.event_data) > 0 else 0
        max_id_in_data += 1
        self.vocab_size = max(max_id_in_vocab, max_id_in_data) + 1

        if not self.include_no_event_in_uts_target:
            self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
        else:
            self.ignored_uts_target_ids = {PAD_IDX, CHECKUP_IDX}

    def _prepare_sex(self, basic_table: pd.DataFrame, unique_eids: np.ndarray) -> None:
        sex_values = pd.to_numeric(basic_table["sex"], errors="coerce").to_numpy()
        if np.isnan(sex_values).any():
            raise ValueError("sex column contains missing or non-numeric values")

        sex_values = sex_values.astype(np.int64)
        sex_unique = np.unique(sex_values)
        if np.all(np.isin(sex_unique, [0, 1])):
            sex01 = sex_values
        elif np.all(np.isin(sex_unique, [1, 2])):
            sex01 = sex_values - 1
        else:
            raise ValueError(
                f"Unexpected sex values: {sex_unique.tolist()}. Expected {{0,1}} or {{1,2}}."
            )
        self.sex_mapping = {int(eid): int(s) for eid, s in zip(unique_eids, sex01)}

    def _prepare_birth_dates(self, basic_table: pd.DataFrame, unique_eids: np.ndarray) -> None:
        if "date_of_birth" not in basic_table.columns:
            if self.exposure_cache is None:
                self.birth_date_mapping = {}
                return
            raise ValueError(
                "Exposure alignment requires ukb_basic_info.csv to contain "
                "'date_of_birth'. Regenerate it with the current prepare_data.py."
            )
        birth = pd.to_datetime(basic_table["date_of_birth"], errors="coerce")
        if birth.isna().any() and self.exposure_cache is not None:
            raise ValueError("date_of_birth contains missing or invalid values")
        birth_np = birth.to_numpy(dtype="datetime64[D]")
        self.birth_date_mapping = {
            int(eid): np.datetime64(date, "D")
            for eid, date in zip(unique_eids, birth_np)
            if not np.isnat(date)
        }

    def _iter_patient_events(
        self,
        *,
        impute_no_event_gaps: bool,
    ) -> Iterable[tuple[int, np.ndarray, np.ndarray]]:
        unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True)
        ends = np.concatenate([starts[1:], [len(self.event_data)]])
        for eid_raw, start, end in zip(unique_eids, starts, ends):
            eid = int(eid_raw)
            rows = self.event_data[start:end]
            times_days_raw = rows[:, 1].astype(np.float32)
            labels_raw = rows[:, 2].astype(np.int64)

            if len(labels_raw) == 0:
                yield eid, times_days_raw, labels_raw
                continue

            labels_raw = np.where(labels_raw >= NO_EVENT_IDX, labels_raw + 1, labels_raw)
            if not impute_no_event_gaps:
                yield eid, times_days_raw, labels_raw
                continue

            times_days, labels = _insert_gap_no_event_tokens(
                times_days_raw,
                labels_raw,
                interval_years=self.no_event_interval_years,
            )
            yield eid, times_days, labels

    def _split_features(self, eid: int) -> Optional[Dict]:
        if eid not in self.sex_mapping:
            return None
        return {
            "sex": self.sex_mapping[eid],
        }

    def _raw_tokens_from_model_tokens(self, model_tokens: np.ndarray) -> np.ndarray:
        raw_tokens = np.full(len(model_tokens), -1, dtype=np.int64)
        real = model_tokens > NO_EVENT_IDX
        raw_tokens[real] = model_tokens[real].astype(np.int64) - 1
        return raw_tokens

    def _exposure_indices_for_inputs(
        self,
        eid: int,
        input_events: np.ndarray,
        input_times_days: np.ndarray,
    ) -> np.ndarray | None:
        if self.exposure_cache is None:
            return None
        raw_tokens = self._raw_tokens_from_model_tokens(input_events)
        return self.exposure_cache.lookup_indices(
            eid=eid,
            raw_tokens=raw_tokens,
            age_days=input_times_days,
        )

    def _load_exposure_windows(self, exposure_index: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]:
        if self.exposure_cache is None:
            raise RuntimeError("Exposure cache is not enabled")

        daily = np.stack(
            [self.exposure_cache.daily_window(int(idx)) for idx in exposure_index],
            axis=0,
        ).astype(np.float32, copy=True)
        monthly = np.stack(
            [self.exposure_cache.monthly_window(int(idx)) for idx in exposure_index],
            axis=0,
        ).astype(np.float32, copy=True)

        if self.mask_onset_exposure:
            daily[:, 0, :] = np.nan
            monthly[:, 0, :] = np.nan

        return torch.from_numpy(daily).float(), torch.from_numpy(monthly).float()

class NextStepHealthDataset(_ExpoBaseDataset):
    """
    Dataset for next-token and next-time-point losses with unified other-info
    tokens.

    Returned targets cover both:
        - Delphi2MLoss: target_event_seq, target_time_seq
        - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot
    """

    CACHE_VERSION = 3

    def __init__(
        self,
        data_prefix: str = "ukb",
        labels_file: str = "labels.csv",
        no_event_interval_years: float = 5.0,
        include_no_event_in_uts_target: bool = False,
        exposure_cache_dir: str | Path | None = None,
        mask_onset_exposure: bool = False,
    ) -> None:
        super().__init__(
            data_prefix=data_prefix,
            labels_file=labels_file,
            no_event_interval_years=no_event_interval_years,
            include_no_event_in_uts_target=include_no_event_in_uts_target,
            exposure_cache_dir=exposure_cache_dir,
            mask_onset_exposure=mask_onset_exposure,
        )

        self.samples: List[Dict] = []
        for eid, times_days, labels in self._iter_patient_events(
            impute_no_event_gaps=True,
        ):
            if len(labels) < 2:
                continue

            features = self._split_features(eid)
            if features is None:
                continue

            target_pack = build_all_targets(
                labels=labels,
                times_days=times_days,
                vocab_size=self.vocab_size,
                ignored_uts_target_ids=self.ignored_uts_target_ids,
                require_sorted=True,
            )

            sample = {
                "eid": eid,
                "event_seq": target_pack.next_token.input_events,
                "time_seq": target_pack.next_token.input_times_years,
                "target_event_seq": target_pack.next_token.target_events,
                "target_time_seq": target_pack.next_token.target_times_years,
                "readout_mask": target_pack.unique_time_set.readout_mask,
                "target_dt_unique": target_pack.unique_time_set.target_dt_unique,
                "target_multi_hot": target_pack.unique_time_set.target_multi_hot,
                **features,
            }
            exposure_index = self._exposure_indices_for_inputs(
                eid=eid,
                input_events=target_pack.next_token.input_events,
                input_times_days=times_days[:-1],
            )
            if exposure_index is not None:
                sample["exposure_index"] = exposure_index
            self.samples.append(sample)

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

    def __getitem__(self, idx: int) -> Dict:
        s = self.samples[idx]
        out = {
            "event_seq": torch.from_numpy(s["event_seq"]).long(),
            "time_seq": torch.from_numpy(s["time_seq"]).float(),
            "sex": torch.tensor(s["sex"], dtype=torch.long),
            "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
            "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
            "readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
            "target_dt_unique": torch.from_numpy(s["target_dt_unique"]).float(),
            "target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(),
        }
        if "exposure_index" in s:
            daily, monthly = self._load_exposure_windows(s["exposure_index"])
            out["exposure_daily"] = daily
            out["exposure_monthly"] = monthly
        return out


class AllFutureHealthDataset(_ExpoBaseDataset):
    """
    Dataset with unified other-info tokens and DeepHealthV2-style all-future
    targets.

    Train samples one query time per patient at each __getitem__ call.
    Valid/test use random-but-fixed query points. For each patient with N real
    disease events, N - 2 query points are sampled from the eligible observed
    time range, with at least one future event after every query.
    """

    CACHE_VERSION = 5

    def __init__(
        self,
        data_prefix: str = "ukb",
        labels_file: str = "labels.csv",
        split: Literal["train", "valid", "test"] = "train",
        no_event_interval_years: float = 5.0,
        include_no_event_in_uts_target: bool = False,
        min_history_events: int = 1,
        min_future_events: int = 1,
        validation_query_seed: int = 42,
        exposure_cache_dir: str | Path | None = None,
        mask_onset_exposure: bool = False,
    ) -> None:
        if split not in {"train", "valid", "test"}:
            raise ValueError(f"split must be train/valid/test, got {split!r}")

        super().__init__(
            data_prefix=data_prefix,
            labels_file=labels_file,
            no_event_interval_years=no_event_interval_years,
            include_no_event_in_uts_target=include_no_event_in_uts_target,
            exposure_cache_dir=exposure_cache_dir,
            mask_onset_exposure=mask_onset_exposure,
        )

        self.split = split
        self.min_history_events = int(min_history_events)
        self.min_future_events = int(min_future_events)
        self.validation_query_seed = int(validation_query_seed)
        self.patients: List[Dict] = []
        self.valid_queries: List[Tuple[int, float]] = []
        validation_rng = None
        if split in {"valid", "test"}:
            split_offset = 0 if split == "valid" else 1_000_003
            validation_rng = np.random.RandomState(self.validation_query_seed + split_offset)

        for eid, times_days, labels in self._iter_patient_events(
            impute_no_event_gaps=False,
        ):
            times_years = (times_days / DAYS_PER_YEAR).astype(np.float32)
            unique_times = np.unique(times_years)
            if len(labels) < 2 or len(unique_times) < 2:
                continue

            features = self._split_features(eid)
            if features is None:
                continue

            patient = {
                "eid": eid,
                "times": times_years,
                "times_days": times_days.astype(np.float32),
                "labels": labels.astype(np.int64),
                "t_obs": float(times_years.max()),
                **features,
            }

            pidx = len(self.patients)
            self.patients.append(patient)

            if split in {"valid", "test"}:
                if validation_rng is None:
                    raise RuntimeError("validation_rng was not initialized")
                self.valid_queries.extend(
                    (pidx, t_query)
                    for t_query in self._sample_fixed_validation_queries(
                        patient,
                        validation_rng,
                    )
                )

        if split in {"valid", "test"} and not self.valid_queries:
            raise ValueError("No random-but-fixed validation query points were built.")

    def _is_valid_query(self, patient: Dict, t_query: float) -> bool:
        times = patient["times"]
        labels = patient["labels"]
        real_event_mask = ~np.isin(
            labels,
            np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
        )
        n_hist = int((times <= t_query).sum())
        n_future = int(((times > t_query) & real_event_mask).sum())
        return (
            n_hist >= self.min_history_events
            and n_future >= self.min_future_events
            and patient["t_obs"] > t_query
        )

    def _sample_fixed_validation_queries(
        self,
        patient: Dict,
        rng: np.random.RandomState,
    ) -> List[float]:
        times = np.asarray(patient["times"], dtype=np.float32)
        labels = np.asarray(patient["labels"], dtype=np.int64)
        real_event_mask = ~np.isin(
            labels,
            np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64),
        )
        real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False))
        n_real_events = int(real_times.size)
        n_queries = max(0, n_real_events - 2)
        if n_queries == 0:
            return []

        min_hist = int(self.min_history_events)
        min_future = int(self.min_future_events)
        if n_real_events < min_hist + min_future:
            return []

        left = float(real_times[min_hist - 1])
        right_event_time = float(real_times[n_real_events - min_future])
        right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf))
        if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left:
            return []

        queries: List[float] = []
        max_attempts = max(100, n_queries * 50)
        for _ in range(max_attempts):
            if len(queries) >= n_queries:
                break
            t_query = float(rng.uniform(left, float(right)))
            if self._is_valid_query(patient, t_query):
                queries.append(t_query)

        return queries

    def _sample_train_query(self, patient: Dict) -> float:
        unique_times = np.unique(patient["times"])
        if len(unique_times) < 2:
            raise RuntimeError("Training patient has fewer than two unique times.")

        j = np.random.randint(1, len(unique_times))
        left = float(unique_times[j - 1])
        right = float(unique_times[j])

        if right - left <= ONE_DAY_YEARS:
            t_query = right - ONE_DAY_YEARS
        else:
            t_query = np.random.uniform(left, right - ONE_DAY_YEARS)

        if not self._is_valid_query(patient, t_query):
            t_query = right - 1e-6
        return float(t_query)

    def _build_item(self, patient: Dict, t_query: float) -> Dict:
        times = patient["times"]
        times_days = patient["times_days"]
        labels = patient["labels"]
        hist = times <= t_query
        fut = times > t_query

        out = {
            "event_seq": torch.from_numpy(labels[hist]).long(),
            "time_seq": torch.from_numpy(times[hist]).float(),
            "t_query": torch.tensor(t_query, dtype=torch.float32),
            "future_targets": torch.from_numpy(labels[fut]).long(),
            "future_dt": torch.from_numpy(times[fut] - np.float32(t_query)).float(),
            "exposure": torch.tensor(np.float32(patient["t_obs"] - t_query), dtype=torch.float32),
            "sex": torch.tensor(patient["sex"], dtype=torch.long),
        }
        if self.exposure_cache is not None:
            exposure_index = self._exposure_indices_for_inputs(
                eid=int(patient["eid"]),
                input_events=labels[hist],
                input_times_days=times_days[hist],
            )
            if exposure_index is not None:
                daily, monthly = self._load_exposure_windows(exposure_index)
                out["exposure_daily"] = daily
                out["exposure_monthly"] = monthly
        return out

    def __len__(self) -> int:
        if self.split == "train":
            return len(self.patients)
        return len(self.valid_queries)

    def __getitem__(self, idx: int) -> Dict:
        if self.split == "train":
            patient = self.patients[idx]
            t_query = self._sample_train_query(patient)
        else:
            pidx, t_query = self.valid_queries[idx]
            patient = self.patients[pidx]
        return self._build_item(patient, t_query)


def _collate_common_static(batch: List[Dict]) -> Dict:
    return {
        "sex": torch.stack([s["sex"] for s in batch]),
    }


def _pad_exposure(batch: List[Dict], key: str, shape: tuple[int, int]) -> torch.Tensor:
    max_len = max(int(s["event_seq"].numel()) for s in batch)
    out = torch.full(
        (len(batch), max_len, shape[0], shape[1]),
        float("nan"),
        dtype=torch.float32,
    )
    for idx, sample in enumerate(batch):
        value = sample.get(key)
        if value is None:
            continue
        seq_len = int(value.size(0))
        out[idx, :seq_len] = value
    return out


def next_step_collate_fn(batch: List[Dict]) -> Dict:
    event_seq = pad_sequence(
        [s["event_seq"] for s in batch],
        batch_first=True,
        padding_value=PAD_IDX,
    )
    time_seq = pad_sequence(
        [s["time_seq"] for s in batch],
        batch_first=True,
        padding_value=0.0,
    )
    target_event_seq = pad_sequence(
        [s["target_event_seq"] for s in batch],
        batch_first=True,
        padding_value=PAD_IDX,
    )
    target_time_seq = pad_sequence(
        [s["target_time_seq"] for s in batch],
        batch_first=True,
        padding_value=0.0,
    )
    readout_mask = pad_sequence(
        [s["readout_mask"] for s in batch],
        batch_first=True,
        padding_value=False,
    )
    target_dt_unique = pad_sequence(
        [s["target_dt_unique"] for s in batch],
        batch_first=True,
        padding_value=0.0,
    )
    target_multi_hot = pad_sequence(
        [s["target_multi_hot"] for s in batch],
        batch_first=True,
        padding_value=False,
    )

    out = {
        "event_seq": event_seq,
        "time_seq": time_seq,
        "padding_mask": event_seq > PAD_IDX,
        "target_event_seq": target_event_seq,
        "target_time_seq": target_time_seq,
        "readout_mask": readout_mask,
        "target_dt_unique": target_dt_unique,
        "target_multi_hot": target_multi_hot,
    }
    out.update(_collate_common_static(batch))
    if any("exposure_daily" in s for s in batch):
        out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE)
        out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE)
    return out


def all_future_collate_fn(batch: List[Dict]) -> Dict:
    event_seq = pad_sequence(
        [s["event_seq"] for s in batch],
        batch_first=True,
        padding_value=PAD_IDX,
    )
    time_seq = pad_sequence(
        [s["time_seq"] for s in batch],
        batch_first=True,
        padding_value=0.0,
    )
    future_targets = pad_sequence(
        [s["future_targets"] for s in batch],
        batch_first=True,
        padding_value=PAD_IDX,
    )
    future_dt = pad_sequence(
        [s["future_dt"] for s in batch],
        batch_first=True,
        padding_value=0.0,
    )

    out = {
        "event_seq": event_seq,
        "time_seq": time_seq,
        "padding_mask": event_seq > PAD_IDX,
        "t_query": torch.stack([s["t_query"] for s in batch]),
        "future_targets": future_targets,
        "future_dt": future_dt,
        "exposure": torch.stack([s["exposure"] for s in batch]),
    }
    out.update(_collate_common_static(batch))
    if any("exposure_daily" in s for s in batch):
        out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE)
        out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE)
    return out


HealthDataset = NextStepHealthDataset
collate_fn = next_step_collate_fn