"""Deterministic padded sequence samples for transformer models."""

from __future__ import annotations

from collections.abc import Iterator, Sequence
from dataclasses import dataclass
from typing import Any

import torch

from gnn4colliders.data.events import PackedEvents
from gnn4colliders.data.metadata import BatchMetadata, EventMetadata
from gnn4colliders.data.sample import EventSample
from gnn4colliders.features import build_node_features


@dataclass(frozen=True)
class SequenceSample:
    """One event represented as an ordered sequence of object tokens."""

    tokens: torch.Tensor
    label: torch.Tensor
    global_features: torch.Tensor | None
    metadata: EventMetadata


@dataclass(frozen=True)
class SequenceBatch:
    """A padded batch with a boolean mask of valid tokens."""

    tokens: torch.Tensor
    token_mask: torch.Tensor
    labels: torch.Tensor
    global_features: torch.Tensor | None
    metadata: BatchMetadata

    def to(
        self, device: torch.device | str, *, non_blocking: bool = False
    ) -> "SequenceBatch":
        return self._map(lambda value: value.to(device, non_blocking=non_blocking))

    def pin_memory(self) -> "SequenceBatch":
        return self._map(lambda value: value.pin_memory())

    def _map(self, move: Any) -> "SequenceBatch":
        metadata = BatchMetadata(
            fold=move(self.metadata.fold),
            weight=move(self.metadata.weight),
            sample_id=self.metadata.sample_id,
            extra=self.metadata.extra,
        )
        return SequenceBatch(
            tokens=move(self.tokens),
            token_mask=move(self.token_mask),
            labels=move(self.labels),
            global_features=(
                None if self.global_features is None else move(self.global_features)
            ),
            metadata=metadata,
        )


def sequence_sample_from_event(
    event: EventSample,
    feature_branches: Sequence[Any],
    object_types: Sequence[str],
    scales: Sequence[float] | torch.Tensor,
) -> SequenceSample:
    """Build tokens from shared collider features without graph topology."""
    tokens, _ = build_node_features(
        event.objects, feature_branches, object_types, scales
    )
    globals_ = torch.as_tensor(event.global_features, dtype=torch.float32)
    return SequenceSample(
        tokens=tokens,
        label=torch.as_tensor(event.label),
        global_features=globals_ if globals_.numel() else None,
        metadata=event.event_metadata,
    )


def collate_sequence_batch(
    events: PackedEvents, indices: torch.Tensor
) -> SequenceBatch:
    """Pad the object rows of ``indices`` into a token batch, in order."""
    if indices.numel() == 0:
        raise ValueError("cannot batch an empty sequence")
    rows, counts = events.gather_objects(indices)
    max_length = int(counts.max()) if counts.numel() else 0
    batch_index = torch.repeat_interleave(
        torch.arange(counts.numel()), counts, output_size=rows.shape[0]
    )
    position = (
        torch.arange(rows.shape[0]) - (torch.cumsum(counts, 0) - counts)[batch_index]
    )
    tokens = rows.new_zeros((counts.numel(), max_length, rows.shape[1]))
    tokens[batch_index, position] = rows
    token_mask = torch.zeros((counts.numel(), max_length), dtype=torch.bool)
    token_mask[batch_index, position] = True
    return SequenceBatch(
        tokens=tokens,
        token_mask=token_mask,
        labels=events.labels[indices],
        global_features=(
            None if events.global_features is None else events.global_features[indices]
        ),
        metadata=events.batch_metadata(indices),
    )


def batch_sequence_samples(samples: Sequence[SequenceSample]) -> SequenceBatch:
    """Pad a sequence of variable-length events in deterministic order."""
    if not samples:
        raise ValueError("cannot batch an empty sequence")
    feature_widths = {int(sample.tokens.shape[1]) for sample in samples}
    if any(sample.tokens.ndim != 2 for sample in samples):
        raise ValueError("sequence tokens must have shape [objects, features]")
    if len(feature_widths) != 1:
        raise ValueError("all sequence samples must have the same feature width")
    max_length = max(int(sample.tokens.shape[0]) for sample in samples)
    feature_width = next(iter(feature_widths))
    tokens = samples[0].tokens.new_zeros((len(samples), max_length, feature_width))
    token_mask = torch.zeros(
        (len(samples), max_length), dtype=torch.bool, device=tokens.device
    )
    for index, sample in enumerate(samples):
        length = int(sample.tokens.shape[0])
        tokens[index, :length] = sample.tokens
        token_mask[index, :length] = True
    labels = torch.stack([torch.as_tensor(sample.label) for sample in samples])
    globals_ = [sample.global_features for sample in samples]
    if any(value is None for value in globals_):
        if not all(value is None for value in globals_):
            raise ValueError("global_features must be present for every sample or none")
        batched_globals = None
    else:
        batched_globals = torch.stack(
            [value for value in globals_ if value is not None]
        )
    return SequenceBatch(
        tokens=tokens,
        token_mask=token_mask,
        labels=labels,
        global_features=batched_globals,
        metadata=BatchMetadata.from_events([sample.metadata for sample in samples]),
    )


class SequenceDataLoader:
    """Small deterministic loader for variable-length sequence samples."""

    def __init__(
        self,
        dataset: Sequence[SequenceSample],
        batch_size: int,
        *,
        shuffle: bool = False,
        seed: int = 12345,
    ) -> None:
        if batch_size < 1:
            raise ValueError("batch_size must be positive")
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.seed = seed
        self.epoch = 0

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

    def __len__(self) -> int:
        return (len(self.dataset) + self.batch_size - 1) // self.batch_size

    def __iter__(self) -> Iterator[SequenceBatch]:
        indices = torch.arange(len(self.dataset), dtype=torch.long)
        if self.shuffle:
            generator = torch.Generator().manual_seed(self.seed + self.epoch)
            indices = indices[torch.randperm(len(indices), generator=generator)]
        for start in range(0, len(indices), self.batch_size):
            selected = [
                self.dataset[int(index)]
                for index in indices[start : start + self.batch_size]
            ]
            yield batch_sequence_samples(selected)