# dataset.py
from __future__ import annotations

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_next_token_targets,
)


ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR

DISEASE_HISTORY_MODE_TIMED = "timed"
DISEASE_HISTORY_MODE_ORDERED = "ordered"
DISEASE_HISTORY_MODE_SET = "set"
DISEASE_HISTORY_MODES = (
    DISEASE_HISTORY_MODE_TIMED,
    DISEASE_HISTORY_MODE_ORDERED,
    DISEASE_HISTORY_MODE_SET,
)


def normalize_disease_history_mode(mode: str | None) -> str:
    value = DISEASE_HISTORY_MODE_TIMED if mode is None else str(mode).lower()
    if value not in DISEASE_HISTORY_MODES:
        raise ValueError(
            "disease_history_mode must be one of "
            f"{list(DISEASE_HISTORY_MODES)}, got {mode!r}"
        )
    return value


def transform_disease_history(
    event_seq: np.ndarray,
    actual_time_seq: np.ndarray,
    actual_t_query: float,
    disease_history_mode: str,
) -> Tuple[np.ndarray, np.ndarray, np.float32]:
    """
    Convert an already-truncated disease history into its model representation.

    ``timed`` keeps the real event/query times. ``ordered`` preserves the
    chronological event order but replaces calendar time with ordinal event-time
    groups. Diseases first recorded on the same day share one ordinal position.
    ``set`` removes both time and order by sorting the unique disease codes and
    assigning every disease and the query the same model time.
    """
    mode = normalize_disease_history_mode(disease_history_mode)
    events = np.asarray(event_seq, dtype=np.int64)
    times = np.asarray(actual_time_seq, dtype=np.float32)
    if events.ndim != 1 or times.ndim != 1 or events.shape != times.shape:
        raise ValueError(
            "event_seq and actual_time_seq must be aligned 1D arrays, got "
            f"{events.shape} and {times.shape}"
        )

    if mode == DISEASE_HISTORY_MODE_TIMED:
        return events, times, np.float32(actual_t_query)

    special = events <= NO_EVENT_IDX
    if np.any(special):
        raise ValueError(
            f"{mode} disease history must contain only disease events; "
            f"found special token ids {np.unique(events[special]).tolist()}"
        )

    if mode == DISEASE_HISTORY_MODE_ORDERED:
        _, ordinal_groups = np.unique(times, return_inverse=True)
        model_times = ordinal_groups.astype(np.float32, copy=False)
        n_groups = int(model_times.max()) + 1 if model_times.size else 0
        return events, model_times, np.float32(n_groups)

    set_events = np.unique(events)
    model_times = np.zeros(set_events.size, dtype=np.float32)
    return set_events, model_times, np.float32(0.0)


def transform_disease_history_batch_at_position(
    event_seq: torch.Tensor,
    actual_time_seq: torch.Tensor,
    padding_mask: torch.Tensor,
    query_position: int,
    disease_history_mode: str,
    vocab_size: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Build a model-visible prefix for token-position all-future evaluation.

    Actual event times remain outside this return value for AUC bookkeeping.
    For ordered/set modes, events after ``query_position`` are explicitly
    masked so collapsing time cannot expose future diseases.
    """
    mode = normalize_disease_history_mode(disease_history_mode)
    if event_seq.ndim != 2 or actual_time_seq.shape != event_seq.shape:
        raise ValueError(
            "event_seq and actual_time_seq must be aligned 2D tensors, got "
            f"{tuple(event_seq.shape)} and {tuple(actual_time_seq.shape)}"
        )
    if padding_mask.shape != event_seq.shape:
        raise ValueError(
            "padding_mask must match event_seq, got "
            f"{tuple(padding_mask.shape)} and {tuple(event_seq.shape)}"
        )
    if query_position < 0 or query_position >= event_seq.size(1):
        raise ValueError(
            f"query_position={query_position} is outside sequence length "
            f"{event_seq.size(1)}"
        )

    padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
    if not torch.all(padding_mask[:, query_position]):
        raise ValueError("query_position must be valid for every batch row")

    if mode == DISEASE_HISTORY_MODE_TIMED:
        return (
            event_seq,
            actual_time_seq,
            padding_mask,
            actual_time_seq[:, query_position],
        )

    positions = torch.arange(
        event_seq.size(1),
        device=event_seq.device,
    )[None, :]
    history_mask = padding_mask & (positions <= query_position)
    visible_events = event_seq.masked_select(history_mask)
    if torch.any(visible_events <= NO_EVENT_IDX):
        special_ids = torch.unique(
            visible_events[visible_events <= NO_EVENT_IDX]
        ).detach().cpu().tolist()
        raise ValueError(
            f"{mode} disease history must contain only disease events; "
            f"found special token ids {special_ids}"
        )

    if mode == DISEASE_HISTORY_MODE_ORDERED:
        model_times = torch.zeros_like(actual_time_seq)
        model_t_query = torch.zeros(
            event_seq.size(0),
            device=actual_time_seq.device,
            dtype=actual_time_seq.dtype,
        )
        for row_idx in range(event_seq.size(0)):
            row_mask = history_mask[row_idx]
            _, ordinal_groups = torch.unique(
                actual_time_seq[row_idx, row_mask],
                sorted=True,
                return_inverse=True,
            )
            model_times[row_idx, row_mask] = ordinal_groups.to(
                dtype=actual_time_seq.dtype
            )
            model_t_query[row_idx] = float(
                int(ordinal_groups.max().item()) + 1
                if ordinal_groups.numel()
                else 0
            )
        return event_seq, model_times, history_mask, model_t_query

    sentinel = torch.full_like(event_seq, int(vocab_size))
    sortable = torch.where(history_mask, event_seq, sentinel)
    set_events = torch.sort(sortable, dim=1).values
    set_mask = set_events != int(vocab_size)
    set_events = set_events.masked_fill(~set_mask, PAD_IDX)
    model_times = torch.zeros_like(actual_time_seq)
    model_t_query = torch.zeros(
        event_seq.size(0),
        device=actual_time_seq.device,
        dtype=actual_time_seq.dtype,
    )
    return set_events, model_times, set_mask, model_t_query


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,
        extra_info_types: Iterable[int] | None = None,
    ) -> None:
        self.data_prefix = data_prefix
        self.labels_file = labels_file
        self.no_event_interval_years = float(no_event_interval_years)
        self.requested_extra_info_types = (
            None
            if extra_info_types is None
            else list(dict.fromkeys(int(t) for t in extra_info_types))
        )

        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)
        other_info = np.load(f"{data_prefix}_other_info.npy")
        if other_info.ndim != 2 or other_info.shape[1] != 5:
            raise ValueError(
                f"other_info must have shape (N, 5), got {other_info.shape}"
            )
        cate_types = pd.read_csv("cate_types.csv")
        required_cate_cols = {"type", "name", "n_categories"}
        missing_cate_cols = required_cate_cols - set(cate_types.columns)
        if missing_cate_cols:
            raise ValueError(
                f"cate_types.csv is missing columns: {sorted(missing_cate_cols)}"
            )

        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_other_info(other_info, cate_types, unique_eids)

        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

    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_other_info(
        self,
        other_info: np.ndarray,
        cate_types: pd.DataFrame,
        unique_eids: np.ndarray,
    ) -> None:
        other_info = other_info.copy()
        other_info[:, 0] = other_info[:, 0].astype(np.int64)
        other_info[:, 1] = other_info[:, 1].astype(np.int64)
        other_info[:, 3] = other_info[:, 3].astype(np.int64)

        available_types = sorted(
            int(t) for t in np.unique(other_info[:, 1]) if int(t) > 0
        )
        if self.requested_extra_info_types is None:
            selected_types = available_types
        else:
            selected_types = self.requested_extra_info_types
            missing = sorted(set(selected_types) - set(available_types))
            if missing:
                raise ValueError(f"Requested extra_info_types not found: {missing}")

        keep = np.isin(other_info[:, 0].astype(np.int64), unique_eids)
        keep &= np.isin(other_info[:, 1].astype(np.int64), selected_types)
        other_info = other_info[keep]

        cate_counts = {
            int(row["type"]): int(row["n_categories"])
            for _, row in cate_types.iterrows()
        }
        cate_offsets: Dict[int, int] = {}
        next_offset = 0
        for type_id in selected_types:
            if type_id in cate_counts:
                cate_offsets[type_id] = next_offset
                next_offset += cate_counts[type_id]

        kinds = other_info[:, 3].astype(np.int64)
        types = other_info[:, 1].astype(np.int64)
        cate_rows = kinds == 2
        for type_id in np.unique(types[cate_rows]):
            type_id = int(type_id)
            if type_id not in cate_offsets:
                raise ValueError(
                    f"type {type_id} appears categorical but is missing from cate_types.csv"
                )
            row_mask = cate_rows & (types == type_id)
            local_value = other_info[row_mask, 2].astype(np.int64)
            other_info[row_mask, 2] = local_value + cate_offsets[type_id]

        cont_type_ids = [
            int(t)
            for t in selected_types
            if np.any((types == int(t)) & (kinds == 1))
        ]

        self.extra_info_types = selected_types
        self.cate_type_offsets = cate_offsets
        self.n_types = (max(selected_types) + 1) if selected_types else 1
        self.cont_type_ids = cont_type_ids
        self.n_cont_types = len(cont_type_ids)
        self.n_categories = next_offset + 1

        order = np.lexsort((other_info[:, 4], other_info[:, 1], other_info[:, 0]))
        other_info = other_info[order]
        self.other_info_by_eid: Dict[int, Dict[str, np.ndarray]] = {}

        for eid in unique_eids.astype(np.int64):
            self.other_info_by_eid[int(eid)] = {
                "other_type": np.zeros(0, dtype=np.int64),
                "other_value": np.zeros(0, dtype=np.float32),
                "other_value_kind": np.zeros(0, dtype=np.int64),
                "other_time": np.zeros(0, dtype=np.float32),
            }

        if len(other_info) == 0:
            return

        eids, starts = np.unique(other_info[:, 0].astype(np.int64), return_index=True)
        ends = np.concatenate([starts[1:], [len(other_info)]])
        for eid_raw, start, end in zip(eids, starts, ends):
            rows = other_info[start:end]
            self.other_info_by_eid[int(eid_raw)] = {
                "other_type": rows[:, 1].astype(np.int64),
                "other_value": rows[:, 2].astype(np.float32),
                "other_value_kind": rows[:, 3].astype(np.int64),
                "other_time": (rows[:, 4].astype(np.float32) / DAYS_PER_YEAR),
            }

    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)

            # CHECKUP is the assessment landmark for selected extra-info tokens.
            # An explicitly empty selection represents a disease-only history,
            # so retaining CHECKUP in that case would introduce an empty
            # landmark token that is not part of the disease sequence.
            if not self.extra_info_types:
                keep = labels_raw != CHECKUP_IDX
                times_days_raw = times_days_raw[keep]
                labels_raw = labels_raw[keep]

            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]:
        other_info = self.other_info_by_eid.get(eid)
        if other_info is None:
            return None
        return {
            "sex": self.sex_mapping[eid],
            **other_info,
        }

class NextStepHealthDataset(_ExpoBaseDataset):
    """
    Delphi2M next-token dataset with unified other-info tokens.
    """

    CACHE_VERSION = 3

    def __init__(
        self,
        data_prefix: str = "ukb",
        labels_file: str = "labels.csv",
        no_event_interval_years: float = 5.0,
        extra_info_types: Iterable[int] | None = None,
    ) -> None:
        super().__init__(
            data_prefix=data_prefix,
            labels_file=labels_file,
            no_event_interval_years=no_event_interval_years,
            extra_info_types=extra_info_types,
        )

        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

            targets = build_next_token_targets(
                labels=labels,
                times_days=times_days,
                require_sorted=True,
            )

            self.samples.append({
                "eid": eid,
                "event_seq": targets.input_events,
                "time_seq": targets.input_times_years,
                "target_event_seq": targets.target_events,
                "target_time_seq": targets.target_times_years,
                **features,
            })

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

    def __getitem__(self, idx: int) -> Dict:
        s = self.samples[idx]
        return {
            "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),
            "other_type": torch.from_numpy(s["other_type"]).long(),
            "other_value": torch.from_numpy(s["other_value"]).float(),
            "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
            "other_time": torch.from_numpy(s["other_time"]).float(),
            "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
            "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
        }


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,
        min_history_events: int = 1,
        min_future_events: int = 1,
        validation_query_seed: int = 42,
        extra_info_types: Iterable[int] | None = None,
        disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
    ) -> 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,
            extra_info_types=extra_info_types,
        )

        self.disease_history_mode = normalize_disease_history_mode(
            disease_history_mode
        )
        if (
            self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED
            and self.extra_info_types
        ):
            raise ValueError(
                f"disease_history_mode={self.disease_history_mode!r} is only "
                "supported with an explicitly empty extra-info selection"
            )

        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,
                "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"]
        labels = patient["labels"]
        hist = times <= t_query
        fut = times > t_query
        event_seq, model_time_seq, model_t_query = transform_disease_history(
            event_seq=labels[hist],
            actual_time_seq=times[hist],
            actual_t_query=t_query,
            disease_history_mode=self.disease_history_mode,
        )

        return {
            "event_seq": torch.from_numpy(event_seq).long(),
            "time_seq": torch.from_numpy(model_time_seq).float(),
            "t_query": torch.tensor(model_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),
            "other_type": torch.from_numpy(patient["other_type"]).long(),
            "other_value": torch.from_numpy(patient["other_value"]).float(),
            "other_value_kind": torch.from_numpy(patient["other_value_kind"]).long(),
            "other_time": torch.from_numpy(patient["other_time"]).float(),
        }

    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]),
        "other_type": pad_sequence(
            [s["other_type"] for s in batch],
            batch_first=True,
            padding_value=0,
        ),
        "other_value": pad_sequence(
            [s["other_value"] for s in batch],
            batch_first=True,
            padding_value=0.0,
        ),
        "other_value_kind": pad_sequence(
            [s["other_value_kind"] for s in batch],
            batch_first=True,
            padding_value=0,
        ),
        "other_time": pad_sequence(
            [s["other_time"] for s in batch],
            batch_first=True,
            padding_value=0.0,
        ),
    }


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,
    )
    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,
    }
    out.update(_collate_common_static(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))
    return out


HealthDataset = NextStepHealthDataset
collate_fn = next_step_collate_fn