# Source code for archai.datasets.distributed_stratified_sampler

    # Copyright (c) Microsoft Corporation.
    # Licensed under the MIT license.

    import math
    from typing import List, Optional, Tuple

    import torch
    from torch.utils.data import Sampler
    import torch.distributed as dist
    from torch.utils.data.dataset import Dataset

    import numpy as np
    from sklearn.model_selection import StratifiedKFold, StratifiedShuffleSplit

    class DistributedStratifiedSampler(Sampler):
        def __init__(self, dataset:Dataset, world_size:Optional[int]=None,
                     rank:Optional[int]=None, shuffle=True,
                     val_ratio:Optional[float]=0.0, is_val=False,
                     max_items:Optional[int]=None):
            """Performs stratified sampling of dataset for each replica in the distributed as well as non-distributed setting. If validation split is needed then yet another stratified sampling within replica's split is performed to further obtain the train/validation splits.

            This sampler works in distributed as well as non-distributed setting with no panelty in either mode and is replacement for built-in torch.util.data.DistributedSampler. In distributed setting, many instances of the same code runs as process known as replicas. Each replica has sequential number assigned by the launcher, starting from 0 to uniquely identify it. This is known as global rank or simply rank. The number of replicas is known as the world size. For non-distributed setting, world_size=1 and rank=0.

            To perform stratified sampling we need labels. This sampler assumes that labels for each datapoint is available in dataset.targets property which should be array like containing as many values as length of the dataset. This is availalble already for many popular datasets such as cifar and, with newer PyTorch versions, ImageFolder as well as DatasetFolder. If you are using custom dataset, you can usually create this property with one line of code such as `dataset.targets = [yi for _, yi in dataset]`.

            Generally, to do distributed sampling, each replica must shuffle with same seed as all other replicas with every epoch and then chose some subset of dataset for itself. Traditionally, we use epoch number as seed for shuffling for each replica. However, this then requires that training code calls sampler.set_epoch(epoch) to set seed at every epoch.

            Arguments:
                dataset -- PyTorch dataset like object

            Keyword Arguments:
                world_size -- Total number of replicas running in distributed setting, if None then auto-detect, 1 for non distributed setting (default: {None})
                rank -- Global rank of this replica, if None then auto-detect, 0 for non distributed setting (default: {None})
                shuffle {bool} -- If True then suffle at every epoch (default: {True})
                val_ratio {float} -- If you want to create validation split then set to > 0 (default: {0.0})
                is_val {bool} -- If True then validation split is returned set to val_ratio otherwise main split is returned (default: {False})
                max_items -- if >= 0 then dataset will be trimmed to these many items for each replica (useful to test on smaller dataset)
            """

            # cifar10 amd DatasetFolder has this attribute, for others it may be easy to add from outside
            assert hasattr(dataset, 'targets') and dataset.targets is not None, 'dataset needs to have targets attribute to work with this sampler'

            if world_size is None:
                if dist.is_available() and dist.is_initialized():
                    world_size = dist.get_world_size()
                else:
                    world_size = 1
            if rank is None:
                if dist.is_available() and dist.is_initialized():
                    rank = dist.get_rank()
                else:
                    rank = 0
            if val_ratio is None:
                val_ratio = 0.0

            assert world_size >= 1
            assert rank >= 0 and rank < world_size
            assert val_ratio < 1.0 and val_ratio >= 0.0

            self.dataset = dataset
            self.world_size = world_size
            self.rank = rank
            self.epoch = 0 # this will be used as seed so cannot be < 0

            self.shuffle = shuffle
            self.data_len = len(self.dataset)
            self.max_items = max_items if max_items is not None and max_items >= 0 else None
            assert self.data_len == len(dataset.targets)
            self.val_ratio = val_ratio
            self.is_val = is_val

            # computing duplications we needs
            self.replica_len = self.replica_len_full =  int(math.ceil(float(self.data_len)/self.world_size))
            self.total_size = self.replica_len_full * self.world_size
            assert self.total_size >= self.data_len

            if self.max_items is not None:
                self.replica_len = min(self.replica_len_full, self.max_items)

            self.main_split_len = int(math.floor(self.replica_len*(1-val_ratio)))
            self.val_split_len = self.replica_len - self.main_split_len
            self._len = self.val_split_len if self.is_val else self.main_split_len

        def __iter__(self):
            # get shuffled indices, dataset is extended if needed to divide equally
            # between replicas
            indices, targets = self._indices()

            # get the fold which we will assign to current replica
            indices, targets = self._replica_fold(indices, targets)

            indices, targets = self._limit(indices, targets, self.max_items)

            # split current replica's fold between train and val
            # return indices depending on if we are val or train split
            indices, _ = self._split(indices, targets, self.val_split_len, self.is_val)
            assert len(indices) == self._len

            # when val fold is needed and shuffle is on, for epoch > 0 we can
            # shuffle only val fold. The seed for other epochs is 0 so that we don't
            # mix val with other folds
            if self.shuffle and self.val_ratio > 0.0 and self.epoch > 0:
                np.random.shuffle(indices)

            return iter(indices)

        def _replica_fold(self, indices:np.ndarray, targets:np.ndarray)\
                ->Tuple[np.ndarray, np.ndarray]:

            if self.world_size > 1:
                replica_fold_idxs = None
                # we don't need shuffling here as it has already been done in _indices()
                rfolder = StratifiedKFold(n_splits=self.world_size, shuffle=False)
                folds = rfolder.split(indices, targets)
                # walk to the split for our rank
                for _ in range(self.rank + 1):
                    other_fold_idxs, replica_fold_idxs = next(folds)

                assert replica_fold_idxs is not None and \
                        len(replica_fold_idxs)==self.replica_len_full

                return indices[replica_fold_idxs], targets[replica_fold_idxs]
            else:
                assert self.world_size == 1
                return indices, targets

        def _indices(self)->Tuple[np.ndarray, np.ndarray]:
            if self.shuffle:
                g = torch.Generator()
                g.manual_seed(self._get_seed())
                indices = torch.randperm(self.data_len, generator=g).numpy()
            else:
                indices = np.arange(self.data_len)

            # add extra samples to make it evenly divisible
            # this is neccesory because we have __len__ which must return same
            # number consistently
            if self.total_size > self.data_len:
                indices = np.append(indices, indices[:(self.total_size - self.data_len)])
            else:
                assert self.total_size == self.data_len, 'total_size cannot be less than dataset size!'

            targets = np.array(list(self.dataset.targets[i] for i in indices))
            assert len(indices) == self.total_size

            return indices, targets

        def _limit(self, indices:np.ndarray, targets:np.ndarray, max_items:Optional[int])\
                ->Tuple[np.ndarray, np.ndarray]:
            # this will limit the items to specified max value
            if max_items is not None:
                return self._split(indices, targets, len(indices)-max_items, False)
            return indices, targets

        def _get_seed(self)->int:
            # if val fold is needed then only do the first shuffle
            # otherwise deterministically shuffle on every epoch
            return self.epoch if self.val_ratio==0.0 else 0

        def _split(self, indices:np.ndarray, targets:np.ndarray, test_size:int,
                   return_test_split:bool)->Tuple[np.ndarray, np.ndarray]:
            if test_size:
                assert isinstance(test_size, int) # othewise next call assumes ratio instead of count
                vfolder = StratifiedShuffleSplit(n_splits=1,
                                                 test_size=test_size,
                                                 random_state=self._get_seed())
                vfolder = vfolder.split(indices, targets)
                train_idx, valid_idx = next(vfolder)

                idxs = valid_idx if return_test_split else train_idx
                return indices[idxs], targets[idxs]
            else:
                return indices, targets

        def __len__(self):
            return self._len

        def set_epoch(self, epoch):
            self.epoch = epoch