from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List

import numpy as np
import torch
from torch.nn.utils.rnn import pad_sequence

from dataset import (
    AllFutureHealthDataset,
    HealthDataset,
)
from targets import PAD_IDX


class AllFutureSequenceEvalDataset:
    """
    Eval-only sequence view for all-future checkpoints.

    All-future training uses the observed history, including CHECKUP state
    tokens, without reusing the next-step view that contains imputed
    <NO_EVENT> gap tokens.
    """

    def __init__(
        self,
        data_prefix: str,
        labels_file: str,
        min_history_events: int = 1,
        min_future_events: int = 1,
        exposure_cache_dir: str | Path | None = None,
        exposure_embeddings_file: str | Path | None = None,
    ) -> None:
        base = AllFutureHealthDataset(
            data_prefix=data_prefix,
            labels_file=labels_file,
            split="train",
            min_history_events=min_history_events,
            min_future_events=min_future_events,
            exposure_cache_dir=exposure_cache_dir,
            exposure_embeddings_file=exposure_embeddings_file,
        )

        self.base = base
        self.label_code_to_id = base.label_code_to_id
        self.label_id_to_code = base.label_id_to_code
        self.vocab_size = base.vocab_size

        self.samples: List[Dict[str, Any]] = []
        for patient in base.patients:
            labels = np.asarray(patient["labels"], dtype=np.int64)
            times = np.asarray(patient["times"], dtype=np.float32)
            times_days = np.asarray(patient["times_days"], dtype=np.float32)
            if labels.size < 2:
                continue
            input_len = int(labels.size - 1)
            self.samples.append(
                sample := {
                    "eid": int(patient["eid"]),
                    "event_seq": labels[:-1],
                    "time_seq": times[:-1],
                    "target_event_seq": labels[1:],
                    "target_time_seq": times[1:],
                    "readout_mask": np.ones(input_len, dtype=bool),
                    "sex": int(patient["sex"]),
                }
            )
            if base.exposure_cache is not None:
                exposure_index = base._exposure_indices_for_inputs(
                    eid=int(patient["eid"]),
                    input_events=labels[:-1],
                    input_times_days=times_days[:-1],
                )
                if exposure_index is not None:
                    sample["exposure_index"] = exposure_index

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

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        s = self.samples[idx]
        out = {
            "event_seq": torch.from_numpy(s["event_seq"]).long(),
            "time_seq": torch.from_numpy(s["time_seq"]).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(),
            "sex": torch.tensor(s["sex"], dtype=torch.long),
        }
        if "exposure_index" in s:
            out["exposure_embedding"] = self.base._load_exposure_embeddings(
                s["exposure_index"]
            )
        return out


def load_sequence_eval_dataset(
    *,
    model_target_mode: str,
    data_prefix: str,
    labels_file: str,
    no_event_interval_years: float,
    include_no_event_in_uts_target: bool,
    min_history_events: int,
    min_future_events: int,
    exposure_cache_dir: str | Path | None = None,
    exposure_embeddings_file: str | Path | None = None,
):
    mode = str(model_target_mode).lower()
    if mode == "next_token":
        return HealthDataset(
            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,
        )
    if mode == "all_future":
        return AllFutureSequenceEvalDataset(
            data_prefix=data_prefix,
            labels_file=labels_file,
            min_history_events=min_history_events,
            min_future_events=min_future_events,
            exposure_cache_dir=exposure_cache_dir,
            exposure_embeddings_file=exposure_embeddings_file,
        )
    raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")


def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
    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
    )
    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,
        "sex": torch.stack([s["sex"] for s in batch]),
    }
    if any("exposure_embedding" in s for s in batch):
        embedding_dim = next(
            int(s["exposure_embedding"].size(1))
            for s in batch if "exposure_embedding" in s
        )
        encoded = torch.zeros(
            len(batch), event_seq.size(1), embedding_dim, dtype=torch.float32
        )
        for idx, sample in enumerate(batch):
            value = sample.get("exposure_embedding")
            if value is not None:
                encoded[idx, :value.size(0)] = value
        out["exposure_embedding"] = encoded
    return out