"""Deterministic DDP batches grouped by clean-history length."""

from __future__ import annotations

import math
import random
from collections import defaultdict
from typing import Iterator

from torch.utils.data import Sampler

from .dataset import FRAMES_PER_CHUNK


class DistributedContextBucketBatchSampler(Sampler[list[int]]):
    def __init__(
        self,
        dataset,
        *,
        batch_size: int,
        rank: int,
        world_size: int,
        seed: int = 0,
        drop_last: bool = True,
    ) -> None:
        self.dataset = dataset
        self.batch_size = int(batch_size)
        self.rank = int(rank)
        self.world_size = int(world_size)
        self.seed = int(seed)
        self.drop_last = bool(drop_last)
        self.epoch = 0
        if self.batch_size <= 0 or not 0 <= self.rank < self.world_size:
            raise ValueError("Invalid distributed bucket sampler configuration")
        self._buckets: dict[int, list[int]] = defaultdict(list)
        for record_index, record in enumerate(dataset.records):
            context_frames = int(
                record.get(
                    "context_frames",
                    int(record["chunk_id"]) * FRAMES_PER_CHUNK,
                )
            )
            for pair_index in range(len(dataset.PAIRS)):
                self._buckets[context_frames].append(
                    record_index * len(dataset.PAIRS) + pair_index
                )

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

    def _global_batches(self) -> list[list[int]]:
        rng = random.Random(self.seed + self.epoch)
        global_batch_size = self.batch_size * self.world_size
        batches: list[list[int]] = []
        for bucket in self._buckets.values():
            indices = list(bucket)
            rng.shuffle(indices)
            if not self.drop_last and len(indices) % global_batch_size:
                needed = global_batch_size - len(indices) % global_batch_size
                indices.extend((indices * math.ceil(needed / len(indices)))[:needed])
            usable = len(indices) - len(indices) % global_batch_size
            batches.extend(
                indices[start : start + global_batch_size]
                for start in range(0, usable, global_batch_size)
            )
        rng.shuffle(batches)
        return batches

    def __iter__(self) -> Iterator[list[int]]:
        start = self.rank * self.batch_size
        end = start + self.batch_size
        for global_batch in self._global_batches():
            yield global_batch[start:end]

    def __len__(self) -> int:
        global_batch_size = self.batch_size * self.world_size
        if self.drop_last:
            return sum(
                len(values) // global_batch_size for values in self._buckets.values()
            )
        return sum(
            math.ceil(len(values) / global_batch_size)
            for values in self._buckets.values()
        )