# dataset.py
from __future__ import annotations

import json
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)
DAILY_EXPOSURE_CHANNELS = ("tmean", "tmax", "tmin", "rhmean")
MONTHLY_EXPOSURE_CHANNELS = ("tmean", "rhmean")


def _daily_exposure_columns() -> list[str]:
    cols: list[str] = []
    for name in DAILY_EXPOSURE_CHANNELS:
        cols.extend(f"{name}_d{idx:04d}" for idx in range(DAILY_EXPOSURE_SHAPE[0]))
    return cols


def _monthly_exposure_columns() -> list[str]:
    cols: list[str] = []
    for name in MONTHLY_EXPOSURE_CHANNELS:
        cols.extend(f"{name}_m{idx:03d}" for idx in range(MONTHLY_EXPOSURE_SHAPE[0]))
    return cols


def _load_readonly_npy(path: Path) -> np.ndarray:
    arr = np.load(path, mmap_mode="r")
    arr.setflags(write=False)
    return arr


class ExposureCache:
    """Eid-sequence-aligned exposure windows from prepare_exposure_cache.py."""

    def __init__(
        self,
        cache_dir: str | Path,
        embeddings_file: str | Path | None = None,
    ):
        cache_dir = Path(cache_dir)
        self.cache_dir = cache_dir
        manifest_path = cache_dir / "exposure_manifest.json"
        self.manifest = (
            json.loads(manifest_path.read_text(encoding="utf-8"))
            if manifest_path.is_file()
            else {}
        )
        self.storage = self.manifest.get("storage")
        if self.storage != "eid_sequence_npy":
            raise ValueError(
                "Exposure cache must use storage='eid_sequence_npy'. "
                "Regenerate it with the current prepare_exposure_cache.py."
            )

        eid_path = cache_dir / "exposure_eid.npy"
        token_path = cache_dir / "exposure_token.npy"
        age_path = cache_dir / "exposure_age_days.npy"
        onset_date_path = cache_dir / "exposure_onset_date.npy"
        row_index_path = cache_dir / "exposure_row_index.npy"
        eid_index_path = cache_dir / "exposure_eid_index.npy"
        eid_start_path = cache_dir / "exposure_eid_start.npy"
        daily_path = cache_dir / "exposure_daily.npy"
        monthly_path = cache_dir / "exposure_monthly.npy"
        required_paths = [
            eid_path,
            token_path,
            age_path,
            onset_date_path,
            row_index_path,
            eid_index_path,
            eid_start_path,
            daily_path,
            monthly_path,
        ]
        if any(not path.is_file() for path in required_paths):
            raise FileNotFoundError(
                "Exposure cache is missing one or more eid-sequence files. "
                "Regenerate it with the current prepare_exposure_cache.py."
            )

        self.eids = _load_readonly_npy(eid_path)
        self.raw_tokens = _load_readonly_npy(token_path)
        self.age_days = _load_readonly_npy(age_path)
        self.onset_dates = _load_readonly_npy(onset_date_path)
        self.row_index = _load_readonly_npy(row_index_path)
        self.eid_index = _load_readonly_npy(eid_index_path)
        self.eid_start = _load_readonly_npy(eid_start_path)
        self.daily = _load_readonly_npy(daily_path)
        self.monthly = _load_readonly_npy(monthly_path)
        self.embeddings = (
            _load_readonly_npy(Path(embeddings_file))
            if embeddings_file is not None
            else None
        )
        quality_path = cache_dir / "exposure_quality.npy"
        self.quality = _load_readonly_npy(quality_path) 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.age_days) != n_rows
            or len(self.onset_dates) != n_rows
            or len(self.row_index) != n_rows
        ):
            raise ValueError("Exposure cache sequence metadata row counts do not match")
        max_window_index = int(np.max(self.row_index)) if n_rows else -1
        if (
            max_window_index >= self.daily.shape[0]
            or max_window_index >= self.monthly.shape[0]
        ):
            raise ValueError("Exposure row index points past daily/monthly window arrays")
        if len(self.eid_start) != len(self.eid_index) + 1:
            raise ValueError("exposure_eid_start.npy must have len(eid_index) + 1")
        if len(self.eid_start) and int(self.eid_start[-1]) != n_rows:
            raise ValueError("Last exposure eid offset must equal exposure row count")
        if self.embeddings is not None:
            if self.embeddings.ndim != 2:
                raise ValueError(
                    "Exposure embeddings must have shape (N, D), got "
                    f"{self.embeddings.shape}"
                )
            if max_window_index >= len(self.embeddings):
                raise ValueError(
                    "Exposure row index points past exposure embeddings"
                )

        self._eid_to_pos = {
            int(eid): idx
            for idx, eid in enumerate(np.asarray(self.eid_index, dtype=np.int64))
        }

    def locality_key(self, indices: np.ndarray) -> tuple[int, int]:
        """Return a stable locality key for sampler-side batching."""
        indices = np.asarray(indices, dtype=np.int64)
        valid = indices[indices >= 0]
        if len(valid) == 0:
            return (2**31 - 1, 2**31 - 1)
        first = int(valid[0])
        return (first // 1024, first % 1024)

    def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None:
        """Kept for the dataset constructor; the new cache already stores age days."""
        return None

    def lookup_indices(self, eid: int, raw_tokens: np.ndarray, age_days: np.ndarray) -> np.ndarray:
        out = np.full(len(raw_tokens), -1, dtype=np.int64)
        real = raw_tokens > 1
        if not np.any(real):
            return out
        eid_pos = self._eid_to_pos.get(int(eid))
        if eid_pos is None:
            return out

        start = int(self.eid_start[eid_pos])
        end = int(self.eid_start[eid_pos + 1])
        if start == end:
            return out

        real_pos = np.nonzero(real)[0]
        n_take = min(len(real_pos), end - start)
        if n_take == 0:
            return out

        out[real_pos[:n_take]] = np.asarray(
            self.row_index[start:start + n_take],
            dtype=np.int64,
        )

        expected_tokens = np.asarray(self.raw_tokens[start:start + n_take], dtype=np.int64)
        expected_age_days = np.asarray(self.age_days[start:start + n_take], dtype=np.int64)
        actual_tokens = np.asarray(raw_tokens[real_pos[:n_take]], dtype=np.int64)
        actual_age_days = np.rint(
            np.asarray(age_days[real_pos[:n_take]], dtype=np.float64)
        ).astype(np.int64)
        if (
            not np.array_equal(expected_tokens, actual_tokens)
            or not np.array_equal(expected_age_days, actual_age_days)
        ):
            raise ValueError(
                "Exposure cache is not aligned to the disease sequence for "
                f"eid={eid}. Regenerate it with the same data_prefix and labels."
            )
        return out

    def daily_window(self, index: int) -> np.ndarray:
        return self.daily_windows(np.asarray([index], dtype=np.int64))[0]

    def monthly_window(self, index: int) -> np.ndarray:
        return self.monthly_windows(np.asarray([index], dtype=np.int64))[0]

    def daily_windows(self, indices: np.ndarray) -> np.ndarray:
        return self._windows("daily", indices)

    def monthly_windows(self, indices: np.ndarray) -> np.ndarray:
        return self._windows("monthly", indices)

    def embedding_windows(self, indices: np.ndarray) -> np.ndarray:
        if self.embeddings is None:
            raise RuntimeError("Exposure embeddings are not enabled")
        indices = np.asarray(indices, dtype=np.int64)
        out = np.zeros(
            (len(indices), self.embeddings.shape[1]), dtype=np.float32
        )
        valid_pos = np.nonzero(indices >= 0)[0]
        if len(valid_pos):
            out[valid_pos] = np.asarray(
                self.embeddings[indices[valid_pos]], dtype=np.float32
            )
        return out

    def _windows(
        self,
        kind: Literal["daily", "monthly"],
        indices: np.ndarray,
    ) -> np.ndarray:
        indices = np.asarray(indices, dtype=np.int64)
        shape = DAILY_EXPOSURE_SHAPE if kind == "daily" else MONTHLY_EXPOSURE_SHAPE
        out = np.full((len(indices), shape[0], shape[1]), np.nan, dtype=np.float32)
        valid_pos = np.nonzero(indices >= 0)[0]
        if len(valid_pos) == 0:
            return out

        valid_indices = indices[valid_pos]
        source = self.daily if kind == "daily" else self.monthly
        out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32)
        return out


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,
        exposure_embeddings_file: str | Path | None = None,
    ) -> 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, exposure_embeddings_file)
            if exposure_cache_dir is not None
            else None
        )
        if exposure_embeddings_file is not None and exposure_cache_dir is None:
            raise ValueError(
                "exposure_cache_dir is required with exposure_embeddings_file"
            )

        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_embeddings(
        self, exposure_index: np.ndarray
    ) -> torch.Tensor:
        if self.exposure_cache is None:
            raise RuntimeError("Exposure cache is not enabled")
        return torch.from_numpy(
            self.exposure_cache.embedding_windows(exposure_index)
        ).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,
        exposure_embeddings_file: str | Path | None = None,
    ) -> 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,
            exposure_embeddings_file=exposure_embeddings_file,
        )

        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:
            out["exposure_embedding"] = self._load_exposure_embeddings(
                s["exposure_index"]
            )
        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,
        exposure_embeddings_file: str | Path | None = None,
    ) -> 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,
            exposure_embeddings_file=exposure_embeddings_file,
        )

        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:
                out["exposure_embedding"] = self._load_exposure_embeddings(
                    exposure_index
                )
        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_embedding(batch: List[Dict]) -> torch.Tensor:
    max_len = max(int(s["event_seq"].numel()) for s in batch)
    embedding_dim = next(
        int(s["exposure_embedding"].size(1))
        for s in batch
        if "exposure_embedding" in s
    )
    out = torch.zeros(
        len(batch), max_len, embedding_dim, dtype=torch.float32
    )
    for idx, sample in enumerate(batch):
        value = sample.get("exposure_embedding")
        if value is not None:
            out[idx, :value.size(0)] = 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_embedding" in s for s in batch):
        out["exposure_embedding"] = _pad_exposure_embedding(batch)
    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_embedding" in s for s in batch):
        out["exposure_embedding"] = _pad_exposure_embedding(batch)
    return out


HealthDataset = NextStepHealthDataset
collate_fn = next_step_collate_fn