from typing import Optional, Iterator
import math

import torch
from torch.utils.data import Dataset
from torch.utils.data.distributed import DistributedSampler, T_co


class RepeatedDistributedSampler(DistributedSampler):
    def __init__(self, dataset: Dataset, batch_size, min_batches, num_replicas: Optional[int] = None,
                 rank: Optional[int] = None, shuffle: bool = True,
                 seed: int = 0, drop_last: bool = False):
        
        assert not drop_last, "drop_last is not supported"
        super(RepeatedDistributedSampler, self).__init__(dataset=dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle, seed=seed, drop_last=drop_last)
        self.num_samples = batch_size * min_batches
        self.total_size = self.num_samples * self.num_replicas

    def __iter__(self) -> Iterator[T_co]:
        if self.shuffle:
            # deterministically shuffle based on epoch and seed
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
        else:
            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]

        # add extra samples to make it evenly divisible
        padding_size = self.total_size - len(indices)
        if padding_size <= len(indices):
            indices += indices[:padding_size]
        else:
            indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
        
        assert len(indices) == self.total_size, f"Error: {len(indices)}, {self.total_size}"

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples
        return iter(indices)
