from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from giant.data.loader import iter_file_chunks
from giant.data.transforms import Normalizer, build_features


class StepsDataset(Dataset):
    def __init__(
        self,
        cond_cont: np.ndarray,
        cond_cat: np.ndarray,
        target: np.ndarray,
    ) -> None:
        self.cond_cont = torch.from_numpy(cond_cont).float()
        self.cond_cat = torch.from_numpy(cond_cat).long()
        self.target = torch.from_numpy(target).float()

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

    def __getitem__(self, index):
        return self.cond_cont[index], self.cond_cat[index], self.target[index]


def train_val_split(
    data: dict,
    cond_cont: np.ndarray,
    cond_cat: np.ndarray,
    target: np.ndarray,
    val_fraction: float = 0.1,
    seed: int = 42,
) -> tuple[StepsDataset, StepsDataset]:
    rng = np.random.default_rng(seed)
    unique_events = np.unique(data["event_id"])
    rng.shuffle(unique_events)
    n_val = max(1, int(len(unique_events) * val_fraction))
    val_events = set(unique_events[:n_val].tolist())

    val_mask = np.array([e in val_events for e in data["event_id"]])
    train_mask = ~val_mask

    return (
        StepsDataset(cond_cont[train_mask], cond_cat[train_mask], target[train_mask]),
        StepsDataset(cond_cont[val_mask], cond_cat[val_mask], target[val_mask]),
    )


def make_event_split(
    all_event_ids: np.ndarray,
    val_fraction: float = 0.1,
    seed: int = 42,
) -> tuple[set, set]:
    """Assign unique event_ids to train/val sets by event_id, not by row."""
    rng = np.random.default_rng(seed)
    unique = np.unique(all_event_ids)
    rng.shuffle(unique)
    n_val = max(1, int(len(unique) * val_fraction))
    val_set = set(unique[:n_val].tolist())
    train_set = set(unique[n_val:].tolist())
    return train_set, val_set


class StreamingStepsDataset(IterableDataset):
    """Streams parquet files one row-group at a time.

    Never loads more than `shuffle_buffer` rows into RAM simultaneously.
    Files are split evenly across DataLoader workers via worker_info.

    Yields whole batches (use with `DataLoader(..., batch_size=None)`)
    rather than single rows, so the batch is assembled with vectorized
    numpy slicing instead of a per-row Python loop in the default collate.
    """

    def __init__(
        self,
        files: list[Path],
        split_events: set,
        pdg_map: dict[int, int],
        mat_map: dict[str, int],
        cond_normalizer: Normalizer,
        target_normalizer: Normalizer,
        batch_size: int,
        shuffle_buffer: int = 65536,
        shuffle: bool = True,
    ) -> None:
        self.files = list(files)
        self.split_events = split_events
        self._events_arr = np.array(sorted(split_events))  # for np.isin
        self.pdg_map = pdg_map
        self.mat_map = mat_map
        self.cond_normalizer = cond_normalizer
        self.target_normalizer = target_normalizer
        self.batch_size = batch_size
        # Buffer must hold at least one batch or we could never emit one.
        self.shuffle_buffer = max(shuffle_buffer, batch_size)
        self.shuffle = shuffle

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        files = self.files
        if worker_info is not None:
            files = files[worker_info.id :: worker_info.num_workers]

        if self.shuffle:
            files = list(files)
            np.random.default_rng().shuffle(files)

        buf_cont: list[np.ndarray] = []
        buf_cat: list[np.ndarray] = []
        buf_tgt: list[np.ndarray] = []
        buf_n = 0

        for path in files:
            for chunk in iter_file_chunks(path):
                mask = np.isin(chunk["event_id"], self._events_arr)
                if not mask.any():
                    continue
                chunk = {k: v[mask] for k, v in chunk.items()}

                cond_cont, cond_cat, target, _, _ = build_features(
                    chunk,
                    self.pdg_map,
                    self.mat_map,
                    cond_normalizer=self.cond_normalizer,
                    target_normalizer=self.target_normalizer,
                )
                buf_cont.append(cond_cont)
                buf_cat.append(cond_cat)
                buf_tgt.append(target)
                buf_n += len(cond_cont)

                if buf_n >= self.shuffle_buffer:
                    buf_cont, buf_cat, buf_tgt, buf_n = yield from self._flush(
                        buf_cont, buf_cat, buf_tgt, final=False
                    )

        if buf_n > 0:
            yield from self._flush(buf_cont, buf_cat, buf_tgt, final=True)

    def _flush(
        self,
        buf_cont: list[np.ndarray],
        buf_cat: list[np.ndarray],
        buf_tgt: list[np.ndarray],
        final: bool,
    ):
        """Yield full batches of `batch_size`; carry any remainder back to the caller.

        All batching is done via vectorized numpy slicing (no per-row Python loop).
        """
        cont = np.concatenate(buf_cont)
        cat = np.concatenate(buf_cat)
        tgt = np.concatenate(buf_tgt)
        if self.shuffle:
            idx = np.random.permutation(len(cont))
            cont, cat, tgt = cont[idx], cat[idx], tgt[idx]

        bs = self.batch_size
        n = len(cont)
        n_full = n // bs if not final else (n + bs - 1) // bs
        for start in range(0, n_full * bs, bs):
            end = min(start + bs, n)
            yield (
                torch.from_numpy(cont[start:end]).float(),
                torch.from_numpy(cat[start:end]).long(),
                torch.from_numpy(tgt[start:end]).float(),
            )

        if final:
            return [], [], [], 0
        rem = n_full * bs
        return [cont[rem:]], [cat[rem:]], [tgt[rem:]], n - rem
