#
#  extensions.py
#  bxtorch
#
#  Created by Oliver Borchert on 09/05/19.
#  Copyright (c) 2019 Oliver Borchert. All rights reserved.
#

import math
import numpy as np
import torch.utils.data as data

###############
### METHODS ###
###############
def loader(self, **kwargs):
    """
    Returns a data loader for this dataset. If the dataset defines a 
    `collate_fn` function, this is automatically set 

    Parameters:
    -----------
    - kwargs: keyword arguments
        Paramaters passed directly to the DataLoader.

    Returns:
    --------
    - torch.utils.data.DataLoader
        The data loader with the specified attributes.
    """
    if hasattr(self, 'collate_fn'):
        kwargs['collate_fn'] = self.collate_fn
    return data.DataLoader(self, **kwargs)


def random_split(self, *sizes, seed=None):
        """
        Splits the dataset randomly into multiple subsets.

        Parameters:
        -----------
        - sizes: variadic argument of float
            The sizes of the splits, given as fraction of the size of the
            dataset. Hence, the sizes must sum to 1.
        - seed: int, default: None
            If given, uses the specified seed to sample the indices for each
            subset.

        Returns:
        --------
        - list of bxtorch.nn.Subset
            The random splits of this dataset.
        """
        assert math.isclose(sum(sizes), 1), \
            "Sizes do not sum to 1."

        randomizer = np.random.RandomState(seed)
        
        # Get subset sizes
        nums = []
        for i, size in enumerate(sizes):
            if i == len(sizes) - 1:
                nums.append(len(self) - sum(nums))
            else:
                nums.append(int(np.round(size * len(self))))

        # Get subset indices
        indices = randomizer.permutation(len(self))
        index_choices = []
        c = 0
        for num in nums:
            index_choices.append(indices[c:c+num])
            c += num

        # Generate subsets
        if hasattr(self, 'collate_fn'):
            # Dynamically define a new subset
            subset_cls = type(
                '_Subset', (data.Subset,), {'collate_fn': self.collate_fn}
            )
        else:
            subset_cls = data.Subset
        
        return [
            subset_cls(self, indices) for indices in index_choices
        ]

##################
### EXTENSIONS ###
##################
data.Dataset.loader = loader
data.Dataset.random_split = random_split
