from __future__ import annotations

import random
from typing import Any, Iterator

from torch.utils.data import Dataset, IterableDataset


class LazyDataLoaderIterator:

    def __init__(self, dataloader: Any):
        self.dataloader = dataloader
        self.iterator: Iterator[Any] | None = None

    def __next__(self) -> Any:
        if self.iterator is None:
            self.iterator = iter(self.dataloader)
        return next(self.iterator)


class MixedDatasetSampler:

    def __init__(
        self,
        datasets: dict[str, Dataset | IterableDataset],
        weights: dict[str, float],
        iterators: dict[str, Iterator[Any]],
        seed: int | None = None,
    ):
        self.datasets = datasets
        self.weights = weights
        self.iterators = iterators
        self._rng = random.Random(seed)

        if len(weights) == 0:
            raise ValueError("weights is empty")
        for name, w in weights.items():
            if float(w) <= 0.0:
                raise ValueError(f"Dataset weight must be > 0, got {name}={float(w)}")
            if name not in datasets:
                raise ValueError(f"Unknown dataset in weights: {name}")
            if name not in iterators:
                raise ValueError(f"Missing iterator for dataset: {name}")

        total_weight = float(sum(float(v) for v in weights.values()))
        self.probs = {name: float(w) / total_weight for name, w in weights.items()}
        self.dataset_names = list(datasets.keys())
        self.prob_list = [self.probs[name] for name in self.dataset_names]

    def sample(self) -> tuple[str, Any]:
        dataset_name = self.choose_dataset_name()
        batch = self.next_batch(dataset_name)
        return dataset_name, batch

    def choose_dataset_name(self, allowed_dataset_names: list[str] | None = None) -> str:
        if allowed_dataset_names is None:
            names = self.dataset_names
            probs = self.prob_list
        else:
            names = [name for name in self.dataset_names if name in set(allowed_dataset_names)]
            if len(names) == 0:
                raise ValueError("No allowed dataset names available for sampling.")
            probs = [self.probs[name] for name in names]
        return self._rng.choices(names, weights=probs, k=1)[0]

    def next_batch(self, dataset_name: str) -> Any:
        if dataset_name not in self.iterators:
            raise ValueError(f"Unknown dataset iterator: {dataset_name}")
        try:
            batch = next(self.iterators[dataset_name])
        except StopIteration as exc:
            raise StopIteration(f"Dataset {dataset_name} exhausted") from exc
        return batch

    def get_sampling_stats(self) -> dict[str, float]:
        return {
            "probabilities": self.probs.copy(),
            "sampling": self.weights.copy(),
        }