import bisect
import random
from typing import Iterable

from torch.utils.data import Dataset, IterableDataset


class ConcatRepeatDataset(Dataset):
    datasets: list[Dataset]
    cumulative_sizes: list[int]
    repeats: list[int]

    @staticmethod
    def cumsum(sequence, repeats):
        r, s = [], 0
        for dataset, repeat in zip(sequence, repeats):
            l = len(dataset) * repeat
            r.append(l + s)
            s += l
        return r

    def __init__(self, datasets: Iterable[Dataset], repeats: list[int]):
        super().__init__()

        self.datasets = list(datasets)
        self.repeats = repeats

        assert len(self.datasets) > 0, "datasets should not be an empty iterable"
        assert len(self.datasets) == len(
            repeats
        ), "datasets and repeats should have the same length"

        for d in self.datasets:
            assert not isinstance(
                d, IterableDataset
            ), "ConcatRepeatDataset does not support IterableDataset"

        self.cumulative_sizes = self.cumsum(self.datasets, self.repeats)

    def __len__(self):
        return self.cumulative_sizes[-1]

    def __getitem__(self, idx):
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)

        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]

        dataset = self.datasets[dataset_idx]

        return dataset[sample_idx % len(dataset)]