# dataset.py
from __future__ import annotations

import hashlib
import os
import pickle
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]


def _cache_file_path(
    data_prefix: str,
    labels_file: str,
    no_event_interval_years: float,
    include_no_event_in_uts_target: bool,
    dataset_kind: str,
    extra_info_types: Iterable[int] | None = None,
    split: str | None = None,
    min_history_events: int | None = None,
    min_future_events: int | None = None,
) -> str:
    event_path = f"{data_prefix}_event_data.npy"
    basic_path = f"{data_prefix}_basic_info.csv"
    other_path = f"{data_prefix}_other_info.npy"
    cate_types_path = "cate_types.csv"
    selected_types = ""
    if extra_info_types is not None:
        seen_types: set[int] = set()
        selected = []
        for raw_type in extra_info_types:
            type_id = int(raw_type)
            if type_id not in seen_types:
                seen_types.add(type_id)
                selected.append(type_id)
        selected_types = ",".join(str(t) for t in selected)
    signature_parts = [
        "deephealthnew_dataset_cache_v2",
        dataset_kind,
        split or "",
        event_path,
        basic_path,
        other_path,
        cate_types_path,
        selected_types,
        labels_file,
        f"{no_event_interval_years:.8f}",
        str(int(include_no_event_in_uts_target)),
        "" if min_history_events is None else str(int(min_history_events)),
        "" if min_future_events is None else str(int(min_future_events)),
    ]

    for path in (event_path, basic_path, other_path, cate_types_path, labels_file):
        try:
            stat = os.stat(path)
            signature_parts.append(f"{path}:{stat.st_mtime_ns}:{stat.st_size}")
        except OSError:
            signature_parts.append(f"{path}:missing")

    digest = hashlib.sha1("|".join(signature_parts).encode("utf-8")).hexdigest()
    cache_dir = os.path.dirname(event_path) or "."
    return os.path.join(cache_dir, f"{data_prefix}_{dataset_kind}_cache_{digest}.pkl")


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) -> 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)
            labels_raw = np.where(labels_raw >= NO_EVENT_IDX, labels_raw + 1, labels_raw)
            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,
        }

    @staticmethod
    def _load_cache(cache_path: str, cache_version: int) -> Optional[Dict]:
        try:
            with open(cache_path, "rb") as f:
                payload = pickle.load(f)
        except OSError:
            return None
        except Exception:
            return None

        if not isinstance(payload, dict):
            return None
        if payload.get("_cache_version") != cache_version:
            return None
        state = payload.get("state")
        if not isinstance(state, dict):
            return None
        return state

    def _save_cache(self, cache_path: str, cache_version: int) -> None:
        payload = {
            "_cache_version": cache_version,
            "state": {key: value for key, value in self.__dict__.items()},
        }
        try:
            cache_dir = os.path.dirname(cache_path)
            if cache_dir:
                os.makedirs(cache_dir, exist_ok=True)
            with open(cache_path, "wb") as f:
                pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
        except OSError:
            return


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 = 1

    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:
        cache_path = _cache_file_path(
            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,
            dataset_kind="next_step",
            extra_info_types=extra_info_types,
        )
        cached_state = self._load_cache(cache_path, self.CACHE_VERSION)
        if cached_state is not None:
            self.__dict__.update(cached_state)
            return

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

        self._save_cache(cache_path, self.CACHE_VERSION)

    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 deterministic pre-event query points.
    """

    CACHE_VERSION = 1

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

        cache_path = _cache_file_path(
            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,
            dataset_kind="all_future",
            extra_info_types=extra_info_types,
            split=split,
            min_history_events=min_history_events,
            min_future_events=min_future_events,
        )
        cached_state = self._load_cache(cache_path, self.CACHE_VERSION)
        if cached_state is not None:
            self.__dict__.update(cached_state)
            return

        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.patients: List[Dict] = []
        self.valid_queries: List[Tuple[int, float]] = []

        for eid, times_days, labels in self._iter_patient_events():
            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"}:
                for target_time in unique_times[1:]:
                    t_query = float(target_time) - ONE_DAY_YEARS
                    if self._is_valid_query(patient, t_query):
                        self.valid_queries.append((pidx, t_query))

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

        self._save_cache(cache_path, self.CACHE_VERSION)

    def _is_valid_query(self, patient: Dict, t_query: float) -> bool:
        times = patient["times"]
        n_hist = int((times <= t_query).sum())
        n_future = int((times > t_query).sum())
        return (
            n_hist >= self.min_history_events
            and n_future >= self.min_future_events
            and patient["t_obs"] > t_query
        )

    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