"""Custom samplers used by MLWiz data loaders.

Includes :class:`~mlwiz.data.sampler.RandomSampler`, which records the applied permutation.
"""

import torch
from torch.utils.data import sampler

from mlwiz.data.dataset import DatasetInterface


class RandomSampler(sampler.RandomSampler):
    """
    This sampler wraps the dataset and saves the random permutation applied
    to the samples, so that it will be available
    for further use (e.g. for saving embeddings in the original samples order).
    The permutation is saved in the 'permutation' attribute.

    Args:
        data_source (:class:`mlwiz.data.DatasetInterface`): the dataset object
    """

    def __init__(self, data_source: DatasetInterface):
        """
        Initialize the sampler and reset the stored permutation.

        Args:
            data_source (DatasetInterface): Dataset to sample from.

        Side effects:
            Sets ``self.permutation`` to ``None``; it will be populated on the
            next call to ``__iter__``.
        """
        super().__init__(data_source)
        self.permutation = None

    def __iter__(self):
        """
        Iterates over the samples according to a pre-determined permutation.

        Returns:
             An iterable with permuted indices.

        """
        n = len(self.data_source)
        self.permutation = torch.randperm(n).tolist()
        return iter(self.permutation)
