"""Batch collation and dynamic batching for translation training."""

from __future__ import annotations

import random
from typing import Iterator, Sequence

import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Sampler


class TranslationCollator:
    """Pad variable-length translation examples into one batch."""

    def __init__(self, pad_token_id: int = 0, label_pad_token_id: int = -100):
        self.pad_token_id = pad_token_id
        self.label_pad_token_id = label_pad_token_id

    def __call__(self, batch: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
        src_ids = pad_sequence(
            [item["src_ids"] for item in batch],
            batch_first=True,
            padding_value=self.pad_token_id,
        )
        tgt_input_ids = pad_sequence(
            [item["tgt_input_ids"] for item in batch],
            batch_first=True,
            padding_value=self.pad_token_id,
        )
        labels = pad_sequence(
            [item["labels"] for item in batch],
            batch_first=True,
            padding_value=self.label_pad_token_id,
        )

        src_padding_mask = src_ids.eq(self.pad_token_id)
        tgt_padding_mask = tgt_input_ids.eq(self.pad_token_id)

        return {
            "src_ids": src_ids,
            "tgt_input_ids": tgt_input_ids,
            "labels": labels,
            "src_padding_mask": src_padding_mask,
            "tgt_padding_mask": tgt_padding_mask,
            "src_attention_mask": (~src_padding_mask).long(),
            "tgt_attention_mask": (~tgt_padding_mask).long(),
            "src_lens": torch.stack([item["src_len"] for item in batch]),
            "tgt_lens": torch.stack([item["tgt_len"] for item in batch]),
        }


class DynamicBatchSampler(Sampler[list[int]]):
    """Create batches constrained by an approximate max token budget."""

    def __init__(
        self,
        lengths: Sequence[int | tuple[int, int]],
        max_tokens_per_batch: int = 8192,
        shuffle: bool = True,
        drop_last: bool = False,
    ):
        self.lengths = [max(length) if isinstance(length, tuple) else int(length) for length in lengths]
        self.max_tokens_per_batch = max_tokens_per_batch
        self.shuffle = shuffle
        self.drop_last = drop_last

    def __iter__(self) -> Iterator[list[int]]:
        indices = list(range(len(self.lengths)))
        if self.shuffle:
            random.shuffle(indices)

        indices.sort(key=lambda idx: self.lengths[idx])
        batches: list[list[int]] = []
        batch: list[int] = []
        max_len = 0

        for idx in indices:
            candidate_max_len = max(max_len, self.lengths[idx])
            candidate_tokens = candidate_max_len * (len(batch) + 1)

            if batch and candidate_tokens > self.max_tokens_per_batch:
                batches.append(batch)
                batch = []
                max_len = 0

            batch.append(idx)
            max_len = max(max_len, self.lengths[idx])

        if batch and not self.drop_last:
            batches.append(batch)

        if self.shuffle:
            random.shuffle(batches)

        yield from batches

    def __len__(self) -> int:
        count = 0
        batch_size = 0
        max_len = 0

        for length in sorted(self.lengths):
            candidate_max_len = max(max_len, length)
            if batch_size and candidate_max_len * (batch_size + 1) > self.max_tokens_per_batch:
                count += 1
                batch_size = 0
                max_len = 0
            batch_size += 1
            max_len = max(max_len, length)

        if batch_size and not self.drop_last:
            count += 1
        return count