# 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_all_targets,
)


ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR


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,
        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.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
        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

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

            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):
    """
    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,
        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,
            include_no_event_in_uts_target=include_no_event_in_uts_target,
            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

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

            self.samples.append({
                "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,
            })

    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(),
            "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(),
        }


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,
        extra_info_types: Iterable[int] | 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,
            extra_info_types=extra_info_types,
        )

        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

        return {
            "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),
            "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,
    )
    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))
    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