import abc
from typing import Tuple

import numpy as np
from numpy.random import default_rng

from badgers.core.base import GeneratorMixin
from badgers.core.decorators.time_series import preprocess_inputs
from badgers.core.utils import random_sign


class OutliersGenerator(GeneratorMixin):
    """
    Base class for transformers that generate point outliers in time-series data
    """

    def __init__(self, random_generator=default_rng(seed=0)):
        """
        :param random_generator: a random number generator
        :param n_outliers: the number of outliers to generate
        """
        self.random_generator = random_generator
        self.outliers_indices_ = []

    @abc.abstractmethod
    def generate(self, X, y, **params) -> Tuple:
        pass


class RandomZerosGenerator(OutliersGenerator):
    """
    Randomly set data points to 0
    """

    def __init__(self, random_generator=default_rng(seed=0)):
        """

        :param random_generator: a random number generator

        """
        super().__init__(random_generator=random_generator)

    @preprocess_inputs
    def generate(self, X, y, n_outliers: int = 10) -> Tuple:
        """
        Randomly set values to zero
        :param X:
        :param y:
        :param n_outliers: the number of outliers to generate
        :return:
        """
        # generate extreme values indices and values
        rows = self.random_generator.choice(X.shape[0], size=n_outliers, replace=False, p=None)
        cols = self.random_generator.integers(low=0, high=X.shape[1], size=n_outliers)

        self.outliers_indices_ = list(zip(rows, cols))

        for r, c in self.outliers_indices_:
            X.iloc[r, c] = 0

        return X, y


class LocalZScoreGenerator(OutliersGenerator):
    """
    Randomly generates locally extreme values.
    Given a time interval (a window) of size l, an outlier is generated by setting the value a time t
    """

    def __init__(self, random_generator=default_rng(seed=0)):
        """

        :param random_generator: a random number generator
        """
        super().__init__(random_generator=random_generator)

    @preprocess_inputs
    def generate(self, X, y, n_outliers: int = 10,
                 local_window_size: int = 10):
        """
        Computes indices of extreme values using a uniform distribution.
        Computes the values at random outside the range [mean(X) - 3 sigma(X), mean(X) + 3 sigma(X)].
        The sign of the extreme value is the same as the value being replaced.

        :param X:
        :param y:
        :param n_outliers: the number of outliers to generate
        :param  local_window_size: the width (number of data points) of the local window to compute local Z-Score
        :return: the transformed array
        """
        # generate extreme values indices and values
        delta = int(local_window_size / 2)

        rows = self.random_generator.choice(
            np.arange(delta, X.shape[0] - delta, dtype=int), size=n_outliers, replace=False, p=None)
        cols = self.random_generator.integers(low=0, high=X.shape[1], size=n_outliers)

        self.outliers_indices_ = list(zip(rows, cols))

        for r, c in self.outliers_indices_:
            local_window = X.iloc[r - delta:r + delta, c]
            local_mean = local_window.mean(axis=0)
            local_std = local_window.std(axis=0)
            value = local_mean + random_sign(self.random_generator) * (
                3. * local_std + self.random_generator.exponential())
            # updating with new outliers
            X.iloc[r, c] = value

        return X, y
