from __future__ import annotations

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

        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)
            if labels.size < 2:
                continue
            input_len = int(labels.size - 1)
            self.samples.append(
                {
                    "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"]),
                }
            )

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

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        s = self.samples[idx]
        return {
            "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),
        }


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