"""
Sampler module.

Defines types generating mini-batches of sample indices.
"""

import numpy as np


def num_batches(n, batch_size):
    """Compute the number of mini-batches required to cover a data set of
    size `n` using batches of size `batch_size`.

    Parameters
    ----------
    n: int
        the number of samples in the data set
    batch_size: int
        the mini-batch size

    Returns
    -------
    int: the number of batches required
    """
    b = n // batch_size
    if n % batch_size > 0:
        b += 1
    return b


class AbstractSampler (object):
    def num_indices_generated(self):
        """
        Get the number of indices that would be generated by this
        sampler.

        Returns
        -------
        int, `np.inf` or `None`.
            An int if the number of samples is known, `np.inf` if it is
            infinite or `None` if the number of samples is unknown.
        """
        raise NotImplementedError

    def map_indices(self, indices):
        """
        Apply index mapping

        Returns the indices parameter unchanged.

        Parameters
        ----------
        indices: NumPy array, 1D dtype=int
            An array of indices

        Returns
        -------
        NumPy array, 1D dtype=int
            The indices, with index mapping applied
        """
        return indices

    def in_order_indices_batch_iterator(self, batch_size):
        """
        Create an iterator that generates in-order mini-batches of sample
        indices. The batches will have `batch_size` elements, with the
        exception of the final batch which will have less if there are not
        enough samples left to fill it.

        The generated mini-batches indices take the form of 1D NumPy integer
        arrays.

        Parameters
        ----------
        batch_size: int
            Mini-batch size

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        raise NotImplementedError('Abstract for type {}'.format(type(self)))

    def shuffled_indices_batch_iterator(self, batch_size, shuffle_rng):
        """
        Create an iterator that generates randomly shuffled mini-batches of
        sample indices. The batches will have `batch_size` elements, with the
        exception of the final batch which will have less if there are not
        enough samples left to fill it.

        The generated mini-batches indices take the form of 1D NumPy integer
        arrays.

        Parameters
        ----------
        batch_size: int
            Mini-batch size
        shuffle_rng: a `numpy.random.RandomState` that will be used to
            randomise element order.

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        raise NotImplementedError('Abstract for type {}'.format(type(self)))


class StandardSampler (AbstractSampler):
    """
    Standard sampler that will construct mini-batches of indices of samples
    either in-order via the `in_order_indices_batch_iterator` method
    or in random order via the `shuffled_indices_batch_iterator` method.
    """
    def __init__(self, length, repeats=1):
        """
        Constructor for standard sampler

        Parameters
        ----------
        length: int
            The number of total samples available
        repeats: int (default=1)
            The number of repetitions, or `-1` for infinite. A value of 0 or
            a negative value that is not -1 will cause `ValueError` to be
            raised.
        """
        if repeats == 0 or repeats < -1:
            raise ValueError('Invalid number of repeats; should be >= 1 or '
                             '-1, not {}'.format(repeats))
        self.length = length
        self.repeats = repeats

    def num_indices_generated(self):
        """
        Get the number of indices that would be generated by this
        sampler.

        Returns
        -------
        int, `np.inf` or `None`.
            An int if the number of samples is known, `np.inf` if it is
            infinite or `None` if the number of samples is unknown.
        """
        if self.repeats == -1:
            return np.inf
        else:
            return self.length * self.repeats

    def in_order_indices_batch_iterator(self, batch_size):
        """
        Create an iterator that generates in-order mini-batches of sample
        indices. The batches will have `batch_size` elements, with the
        exception of the final batch which will have less if there are not
        enough samples left to fill it.

        The generated mini-batches indices take the form of 1D NumPy integer
        arrays.

        Parameters
        ----------
        batch_size: int
            Mini-batch size

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        if self.repeats == 1:
            for i in range(0, self.length, batch_size):
                yield np.arange(i, min(i + batch_size, self.length))
        else:
            repeats = self.repeats
            i = 0
            while True:
                j = i + batch_size
                if j <= self.length:
                    # Within size of data
                    yield np.arange(i, j)
                    i = j
                elif j <= self.length * 2:
                    # One restart is required
                    # Reduce the number of remaining repeats
                    if repeats != -1:
                        repeats -= 1
                    if repeats == 0:
                        # Finished; emit remaining elements
                        if i < self.length:
                            yield np.arange(i, self.length)
                        break

                    # Wrap over
                    # Compute number of elements required to make up
                    # the batch
                    k = batch_size - (self.length - i)
                    yield np.append(np.arange(i, self.length),
                                    np.arange(0, k), axis=0)
                    i = k
                else:
                    # Multiple restarts required to fill the batch
                    batch_ndx = np.arange(0)
                    # i = 0
                    while len(batch_ndx) < batch_size:
                        # Wrap over
                        k = min(batch_size - len(batch_ndx), self.length - i)
                        batch_ndx = np.append(
                            batch_ndx, np.arange(i, i + k), axis=0)
                        i += k
                        if i >= self.length:
                            i -= self.length
                            # Reduce the number of remaining repeats
                            if repeats != -1:
                                repeats -= 1
                            if repeats == 0:
                                break

                    if len(batch_ndx) > 0:
                        yield batch_ndx
                    if repeats == 0:
                        break

    def shuffled_indices_batch_iterator(self, batch_size, shuffle_rng):
        """
        Create an iterator that generates randomly shuffled mini-batches of
        sample indices. The batches will have `batch_size` elements, with the
        exception of the final batch which will have less if there are not
        enough samples left to fill it.

        The generated mini-batches indices take the form of 1D NumPy integer
        arrays.

        Parameters
        ----------
        batch_size: int
            Mini-batch size
        shuffle_rng: a `numpy.random.RandomState` that will be used to
            randomise element order.

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        if self.repeats == 1:
            indices = shuffle_rng.permutation(self.length)
            for i in range(0, self.length, batch_size):
                yield indices[i:i + batch_size]
        else:
            repeats = self.repeats
            indices = shuffle_rng.permutation(self.length)
            i = 0
            while True:
                j = i + batch_size
                if j <= self.length:
                    # Within size of data
                    yield indices[i:j]
                    i = j
                else:
                    # Multiple restarts required to fill the batch
                    batch_ndx = np.arange(0)
                    while len(batch_ndx) < batch_size:
                        # Wrap over
                        k = min(batch_size - len(batch_ndx), self.length - i)
                        batch_ndx = np.append(
                            batch_ndx, indices[i:i + k], axis=0)
                        i += k

                        if i >= self.length:
                            # Loop over; new permutation
                            indices = shuffle_rng.permutation(self.length)
                            i -= self.length
                            # Reduce the number of remaining repeats
                            if repeats != -1:
                                repeats -= 1
                            if repeats == 0:
                                break

                    if len(batch_ndx) > 0:
                        yield batch_ndx
                    if repeats == 0:
                        break


class SubsetSampler (AbstractSampler):
    """
    Sampler that will construct mini-batches of indices of samples from
    a subset of a dataset, where the elements in the subset are identified
    by an array of indices passed to the constructor.
    """
    def __init__(self, indices, repeats=1):
        """
        Constructor

        Parameters
        ----------
        indices: NumPy array, 1D dtype=int
            An array of indices that identify the subset of samples drawn
            from data that are to be used
        repeats: int (default=1)
            The number of repetitions, or `-1` for infinite. A value of 0 or
            a negative value that is not -1 will cause `ValueError` to be
            raised.
        """
        if repeats == 0 or repeats < -1:
            raise ValueError('Invalid number of repeats; should be >= 1 or '
                             '-1, not {}'.format(repeats))
        self.indices = indices
        self.length = len(self.indices)
        self.repeats = repeats

    def num_indices_generated(self):
        """
        Get the number of indices that would be generated by this
        sampler.

        Returns
        -------
        int, `np.inf` or `None`.
            An int if the number of samples is known, `np.inf` if it is
            infinite or `None` if the number of samples is unknown.
        """
        if self.repeats == -1:
            return np.inf
        else:
            return self.length * self.repeats

    def map_indices(self, indices):
        """
        Apply index mapping

        Returns the indices parameter unchanged.

        Parameters
        ----------
        indices: NumPy array, 1D dtype=int
            An array of indices

        Returns
        -------
        NumPy array, 1D dtype=int
            The indices, with index mapping applied
        """
        return self.indices[indices]

    def in_order_indices_batch_iterator(self, batch_size):
        """
        Create an iterator that generates in-order mini-batches of sample
        indices. The batches will have `batch_size` elements, with the
        exception of the final batch which will have less if
        `len(self.indices) is not a multiple of `batch_size`.

        The generated mini-batches indices take the form of 1D NumPy integer
        arrays.

        Parameters
        ----------
        batch_size: int
            Mini-batch size

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        if self.repeats == 1:
            for i in range(0, self.length, batch_size):
                yield self.indices[i:i + batch_size]
        else:
            repeats = self.repeats
            i = 0
            while True:
                j = i + batch_size
                if j <= self.length:
                    # Within size of data
                    yield self.indices[i:j]
                    i = j
                else:
                    # Multiple restarts required to fill the batch
                    batch_ndx = np.arange(0)
                    while len(batch_ndx) < batch_size:
                        # Wrap over
                        k = min(batch_size - len(batch_ndx),
                                self.length - i)
                        batch_ndx = np.append(
                            batch_ndx, self.indices[i:i + k], axis=0)
                        i += k
                        if i >= self.length:
                            i -= self.length
                            # Reduce the number of remaining repeats
                            if repeats != -1:
                                repeats -= 1
                            if repeats == 0:
                                break

                    if len(batch_ndx) > 0:
                        yield batch_ndx
                    if repeats == 0:
                        break

    def shuffled_indices_batch_iterator(self, batch_size, shuffle_rng):
        """
        Create an iterator that generates randomly shuffled mini-batches of
        sample indices. The batches will have `batch_size` elements, with the
        exception of the final batch which will have less if
        `len(self.indices)` is not a multiple of `batch_size`.

        The generated mini-batches indices take the form of 1D NumPy integer
        arrays.

        Parameters
        ----------
        batch_size: int
            Mini-batch size
        shuffle_rng: a `numpy.random.RandomState` that will be used to
            randomise element order.

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        if self.repeats == 1:
            indices = shuffle_rng.permutation(self.indices)
            for i in range(0, self.length, batch_size):
                yield indices[i:i + batch_size]
        else:
            repeats = self.repeats
            indices = shuffle_rng.permutation(self.indices)
            i = 0
            while True:
                j = i + batch_size
                if j <= self.length:
                    # Within size of data
                    yield indices[i:j]
                    i = j
                else:
                    # Multiple restarts required to fill the batch
                    batch_ndx = np.arange(0)
                    while len(batch_ndx) < batch_size:
                        # Wrap over
                        k = min(batch_size - len(batch_ndx),
                                self.length - i)
                        batch_ndx = np.append(
                            batch_ndx, indices[i:i + k], axis=0)
                        i += k

                        if i >= self.length:
                            # Loop over; new permutation
                            indices = shuffle_rng.permutation(self.indices)
                            i -= self.length
                            # Reduce the number of remaining repeats
                            if repeats != -1:
                                repeats -= 1
                            if repeats == 0:
                                break

                    if len(batch_ndx) > 0:
                        yield batch_ndx
                    if repeats == 0:
                        break


class WeightedSampler (AbstractSampler):
    """
    Sampler that draws sample indices whose distribution is controlled by
    a weight array. Samples are drawn in-order via the
    `in_order_indices_batch_iterator` method or in random order via the
    `shuffled_indices_batch_iterator` method.
    """
    def __init__(self, weights):
        """
        Constructor for weighted sampler.

        The weights do not have to sum to 1, as they are normalised.

        Parameters
        ----------
        weights: NumPy array, 1D dtype=float
            sample weights
        """
        weights = weights.astype(float)
        weight_sum = weights.sum()
        if weight_sum == 0.0:
            raise ValueError('weights should not sum to 0')
        self.weights = weights / weight_sum
        self.length = len(self.weights)

    def num_indices_generated(self):
        """
        Get the number of indices that would be generated by this
        sampler.

        Always `np.inf` for `WeightedSampler` as it repeats infinitely.

        Returns
        -------
        int, `np.inf` or `None`.
            An int if the number of samples is known, `np.inf` if it is
            infinite or `None` if the number of samples is unknown.
        """
        return np.inf

    def in_order_indices_batch_iterator(self, batch_size):
        """
        Create an iterator that generates in-order mini-batches of sample
        indices. The batches will have `batch_size` elements.

        Raises `NotImplementedError` as this method is not supported by
        `WeightedSampler`.

        Parameters
        ----------
        batch_size: int
            Mini-batch size

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        raise NotImplementedError('In-order sampling not supported by '
                                  'WeightedSampler')

    def shuffled_indices_batch_iterator(self, batch_size, shuffle_rng):
        """
        Create an iterator that generates randomly shuffled mini-batches of
        sample indices. The batches will have `batch_size` elements.

        The generated mini-batches indices take the form of 1D NumPy integer
        arrays.

        Parameters
        ----------
        batch_size: int
            Mini-batch size
        shuffle_rng: a `numpy.random.RandomState` that will be used to
            randomise element order.

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        while True:
            yield shuffle_rng.choice(len(self.weights), size=(batch_size,),
                                     p=self.weights)

    @staticmethod
    def class_balancing_sample_weights(y):
        """
        Compute sample weight given an array of sample classes. The weights
        are assigned on a per-class basis and the per-class weights are
        inversely proportional to their frequency.

        Parameters
        ----------
        y: NumPy array, 1D dtype=int
            sample classes, values must be 0 or positive

        Returns
        -------
        NumPy array, 1D dtype=float
            per sample weight array
        """
        h = np.bincount(y)
        cls_weight = 1.0 / (h.astype(float) * len(np.nonzero(h)[0]))
        cls_weight[np.isnan(cls_weight)] = 0.0
        sample_weight = cls_weight[y]
        return sample_weight

    @staticmethod
    def class_balancing_sampler(y):
        """
        Construct a `WeightedSampler` that compensates for class imbalance.

        Parameters
        ----------
        y: NumPy array, 1D dtype=int
            sample classes, values must be 0 or positive

        Returns
        -------
        WeightedSampler instance
            Sampler
        """
        weights = WeightedSampler.class_balancing_sample_weights(y)
        return WeightedSampler(weights)


class WeightedSubsetSampler (AbstractSampler):
    """
    Sampler that draws sample indices whose distribution is controlled by
    a weight array. Samples are drawn in-order via the
    `in_order_indices_batch_iterator` method or in random order via the
    `shuffled_indices_batch_iterator` method.
    """
    def __init__(self, sub_weights, indices):
        """
        Constructor for weighted sampler. The array `sub_weights` provides
        the weights of the underlying samples indexec by `indices`.

        The `sub_weights` and `indices` parameters

        Parameters
        ----------
        sub_weights: NumPy array, 1D dtype=float
            sample weights
        indices: NumPy array, 1D dtype=int
            An array of indices that identify the subset of samples drawn
            from data that are to be used
        """
        if len(sub_weights) != len(indices):
            raise ValueError(
                'sub_weights and indices should have the same length; '
                '{} != {}'.format(len(sub_weights), len(indices)))
        sub_weights_sum = sub_weights.sum()
        if sub_weights_sum == 0.0:
            raise ValueError('sub_weights should not sum to 0')
        self.sub_weights = sub_weights / sub_weights_sum
        self.indices = indices
        self.length = len(self.indices)

    def num_indices_generated(self):
        """
        Get the number of indices that would be generated by this
        sampler.

        Always `np.inf` for `WeightedSampler` as it repeats infinitely.

        Returns
        -------
        int, `np.inf` or `None`.
            An int if the number of samples is known, `np.inf` if it is
            infinite or `None` if the number of samples is unknown.
        """
        return np.inf

    def map_indices(self, indices):
        """
        Apply index mapping

        Returns the indices parameter unchanged.

        Parameters
        ----------
        indices: NumPy array, 1D dtype=int
            An array of indices

        Returns
        -------
        NumPy array, 1D dtype=int
            The indices, with index mapping applied
        """
        return self.indices[indices]

    def in_order_indices_batch_iterator(self, batch_size):
        """
        Create an iterator that generates in-order mini-batches of sample
        indices. The batches will have `batch_size` elements.

        Raises `NotImplementedError` as this method is not supported by
        `WeightedSubsetSampler`.

        Parameters
        ----------
        batch_size: int
            Mini-batch size

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        raise NotImplementedError('In-order sampling not supported by '
                                  'WeightedSubsetSampler')

    def shuffled_indices_batch_iterator(self, batch_size, shuffle_rng):
        """
        Create an iterator that generates randomly shuffled mini-batches of
        sample indices. The batches will have `batch_size` elements.

        The generated mini-batches indices take the form of 1D NumPy integer
        arrays.

        Parameters
        ----------
        batch_size: int
            Mini-batch size
        shuffle_rng: a `numpy.random.RandomState` that will be used to
            randomise element order.

        Returns
        -------
        iterator
            An iterator that generates mini-batches in the form of 1D NumPy
            integer arrays.
        """
        while True:
            yield shuffle_rng.choice(self.indices, size=(batch_size,),
                                     p=self.sub_weights)

    @staticmethod
    def class_balancing_sampler(y, indices):
        """
        Construct a `WeightedSubsetSampler` that compensates for class
        imbalance.

        Parameters
        ----------
        y: NumPy array, 1D dtype=int
            sample classes, values must be 0 or positive
        indices: NumPy array, 1D dtype=int
            An array of indices that identify the subset of samples drawn
            from data that are to be used

        Returns
        -------
        WeightedSubsetSampler instance
            Sampler
        """
        weights = WeightedSampler.class_balancing_sample_weights(y[indices])
        return WeightedSubsetSampler(weights, indices=indices)
