import numpy as np
import torch
import torch.nn.functional as F
from padertorch.contrib.je.modules.conv import Pad
from torch import nn


class Scale(nn.Module):
    """
    >>> x = torch.ones((3, 4, 5))
    >>> x = Scale(log_truncnormal_sampling_fn)(x)
    """
    def __init__(self, scale_sampling_fn):
        super().__init__()
        self.scale_sampling_fn = scale_sampling_fn

    def forward(self, x):
        if not self.training:
            return x
        scales = self.scale_sampling_fn(x.shape[0])
        scales = torch.from_numpy(scales[(...,) + (x.dim()-1)*(None,)]).float()
        return x * scales.to(x.device)


class Shift(nn.Module):
    """
    >>> x = torch.ones((3, 4, 5))
    >>> Shift(truncnormal_sampling_fn, scale=0.5)(x)
    """
    def __init__(self, shift_sampling_fn):
        super().__init__()
        self.shift_sampling_fn = shift_sampling_fn

    def forward(self, x):
        if not self.training:
            return x
        shifts = self.shift_sampling_fn(x.shape[0])
        shifts = torch.from_numpy(shifts[(...,) + (x.dim()-1)*(None,)]).float()
        return x + shifts.to(x.device)


class TimeWarping(nn.Module):
    """
    >>> x = torch.cumsum(torch.ones((3, 1, 4, 5)), -1) - 1
    >>> resampling_factors = np.array([1,2,.5])
    >>> warping_fn=lambda seq_len: (\
            np.minimum(np.arange(max(seq_len))/resampling_factors[:, None], max(seq_len)-1),\
            np.minimum(resampling_factors*np.array(seq_len), max(seq_len))\
        )
    >>> TimeWarping(warping_fn=warping_fn)(x, seq_len=[3,4,5])
    (tensor([[[[0.0000, 1.0000, 2.0000, 3.0000, 4.0000],
              [0.0000, 1.0000, 2.0000, 3.0000, 4.0000],
              [0.0000, 1.0000, 2.0000, 3.0000, 4.0000],
              [0.0000, 1.0000, 2.0000, 3.0000, 4.0000]]],
    <BLANKLINE>
    <BLANKLINE>
            [[[0.0000, 0.5000, 1.0000, 1.5000, 2.0000],
              [0.0000, 0.5000, 1.0000, 1.5000, 2.0000],
              [0.0000, 0.5000, 1.0000, 1.5000, 2.0000],
              [0.0000, 0.5000, 1.0000, 1.5000, 2.0000]]],
    <BLANKLINE>
    <BLANKLINE>
            [[[0.0000, 0.3333, 0.6667, 1.0000, 1.3333],
              [0.0000, 0.3333, 0.6667, 1.0000, 1.3333],
              [0.0000, 0.3333, 0.6667, 1.0000, 1.3333],
              [0.0000, 0.3333, 0.6667, 1.0000, 1.3333]]]]), array([3, 5, 5]))
    """
    def __init__(self, warping_fn, batch_axis=0, sequence_axis=-1):
        super().__init__()
        self.warping_fn = warping_fn
        self.batch_axis = batch_axis
        self.sequence_axis = sequence_axis

    def forward(self, *tensors, seq_len):
        """

        Args:
            tensors:
            seq_len:

        Returns:

        """
        if self.training:
            assert seq_len is not None
            time_indices, seq_len = self.warping_fn(seq_len)
            time_indices_ceil = np.ceil(time_indices).astype(np.int)
            time_indices_floor = np.floor(time_indices).astype(np.int)
            batch_indices = np.arange(len(seq_len)).astype(np.int)[:, None]
            tensors = list(tensors)
            for i, tensor in enumerate(tensors):
                batch_axis = self.batch_axis
                seq_axis = self.sequence_axis
                if batch_axis < 0:
                    batch_axis = tensor.dim() + batch_axis
                if seq_axis < 0:
                    seq_axis = tensor.dim() + seq_axis
                if batch_axis != 0:
                    tensor = tensor.transpose(0, batch_axis)
                    if seq_axis == 0:
                        seq_axis = batch_axis
                if seq_axis != 1:
                    tensor = tensor.transpose(1, seq_axis)
                ceil_weights = \
                    torch.Tensor(1 - time_indices_ceil + time_indices).to(tensor.device)
                floor_weights = (
                    torch.Tensor(1 - time_indices + time_indices_floor)
                    * torch.Tensor(time_indices_floor != time_indices_ceil)
                ).to(tensor.device)
                for _ in range(tensor.dim() - 2):
                    ceil_weights = ceil_weights.unsqueeze(-1)
                    floor_weights = floor_weights.unsqueeze(-1)
                tensor = (
                    tensor[batch_indices, time_indices_ceil] * ceil_weights
                    + tensor[batch_indices, time_indices_floor] * floor_weights
                )
                if seq_axis != 1:
                    tensor = tensor.transpose(1, seq_axis)
                if batch_axis != 0:
                    tensor = tensor.transpose(0, batch_axis)
                tensors[i] = tensor
        return (*tensors, seq_len)


class Superpose(nn.Module):
    """
    >>> x = torch.cumsum(torch.ones((3, 4, 5)), 0)
    >>> y = torch.eye(3).float()
    >>> superpose = Superpose(p=1.)
    >>> superpose(x, seq_len=[3,4,5], targets=y)
    """
    def __init__(self, p, scale_fn=None):
        super().__init__()
        self.p = p
        self.scale_fn = scale_fn

    def forward(self, x, targets, seq_len=None):
        if self.training:
            B = x.shape[0]
            shuffle_idx = np.roll(np.arange(B, dtype=np.int), np.random.choice(B-1)+1)
            lambda_ = np.random.binomial(1, self.p, B)
            if seq_len is not None:
                seq_len = np.maximum(seq_len, (lambda_ > 0.) * np.array(seq_len)[shuffle_idx])
            lambda_ = torch.from_numpy(lambda_).float().to(x.device)
            assert all(lambda_ >= 0.) and all(lambda_ <= 1.)
            lambda_x = lambda_[(...,) + (x.dim() - 1) * (None,)]
            x_shuffled = x[shuffle_idx]
            if self.scale_fn is not None:
                x = self.scale_fn(x)
                x_shuffled = self.scale_fn(x_shuffled)
            x = x + lambda_x * x_shuffled
            if isinstance(targets, (list, tuple)):
                targets = list(targets)
                for i in range(len(targets)):
                    lambda_t = lambda_[(...,) + (targets[i].dim() - 1) * (None,)]
                    targets[i] = targets[i] + lambda_t * (1 - targets[i]) * targets[i][shuffle_idx]
            elif targets is not None:
                lambda_t = lambda_[(...,) + (targets.dim() - 1) * (None,)]
                targets = targets + lambda_t * (1 - targets) * targets[shuffle_idx]
        return x, targets, seq_len


class Mixup(nn.Module):
    """
    >>> x = torch.cumsum(torch.ones((3, 4, 5)), 0)
    >>> y = torch.eye(3).float()
    >>> mixup = Mixup(p=1.)
    >>> mixup(x, seq_len=[3,4,5], targets=y)
    """
    def __init__(self, p, alpha=1.):
        super().__init__()
        self.p = p
        self.alpha = alpha

    def forward(self, x, targets=None, seq_len=None):
        if self.training:
            B = x.shape[0]
            shuffle_idx = np.roll(np.arange(B, dtype=np.int), np.random.choice(B))
            lambda_ = np.maximum(
                np.random.binomial(1, 1 - self.p, B),
                np.random.beta(self.alpha, self.alpha, B)
            )
            if seq_len is not None:
                seq_len = np.maximum(seq_len, (lambda_ < 1.) * np.array(seq_len)[shuffle_idx])
            lambda_ = torch.from_numpy(lambda_).float().to(x.device)
            assert all(lambda_ >= 0.) and all(lambda_ <= 1.)
            lambda_x = lambda_[(...,) + (x.dim() - 1) * (None,)]
            x = lambda_x * x + (1. - lambda_x) * x[shuffle_idx]
            if isinstance(targets, (list, tuple)):
                targets = list(targets)
                for i in range(len(targets)):
                    lambda_t = lambda_[(...,) + (targets[i].dim() - 1) * (None,)]
                    targets[i] = lambda_t * targets[i] + (1. - lambda_t) * targets[i][shuffle_idx]
            elif targets is not None:
                lambda_t = lambda_[(...,) + (targets.dim() - 1) * (None,)]
                targets = lambda_t * targets + (1. - lambda_t) * targets[shuffle_idx]
        return x, targets, seq_len


class Mask(nn.Module):
    """
    >>> x = torch.ones((3, 4, 5))
    >>> x = Mask(axis=-1, max_masked_rate=1., max_masked_steps=10)(x, seq_len=[1,2,3])
    """
    def __init__(self, axis, n_masks=1, max_masked_steps=None, max_masked_rate=1.):
        super().__init__()
        self.axis = axis
        self.n_masks = n_masks
        self.max_masked_values = max_masked_steps
        self.max_masked_rate = max_masked_rate

    def __call__(self, x, seq_len=None):
        if not self.training:
            return x
        mask = torch.ones_like(x)
        idx = torch.arange(x.shape[self.axis]).float()
        axis = self.axis
        if axis < 0:
            axis = x.dim() + axis
        idx = idx[(...,) + (x.dim() - axis - 1)*(None,)]
        idx = idx.expand(x.shape)
        if seq_len is None:
            seq_len = x.shape[axis] * torch.ones(x.shape[0])
        else:
            seq_len = torch.Tensor(seq_len)
        max_width = self.max_masked_rate/self.n_masks * seq_len
        if self.max_masked_values is not None:
            max_width = torch.min(self.max_masked_values*torch.ones_like(max_width)/self.n_masks, max_width)
        max_width = torch.floor(max_width)
        for i in range(self.n_masks):
            width = torch.floor(torch.rand(x.shape[0]) * (max_width + 1))
            max_onset = seq_len - width
            onset = torch.floor(torch.rand(x.shape[0]) * (max_onset + 1))
            width = width[(...,) + (x.dim()-1)*(None,)]
            onset = onset[(...,) + (x.dim()-1)*(None,)]
            offset = onset + width
            mask = mask * ((idx < onset) + (idx >= offset)).float().to(x.device)
        return x * mask


class AdditiveNoise(nn.Module):
    """
    >>> x = torch.zeros((3, 4, 5))
    >>> AdditiveNoise(1.)(x)
    """
    def __init__(self, max_scale):
        super().__init__()
        self.max_scale = max_scale

    def forward(self, x):
        if self.training:
            B = x.shape[0]
            scale = torch.rand(B).to(x.device) * self.max_scale
            x = x + scale[(...,) + (x.dim()-1)*(None,)] * torch.randn_like(x)
        return x


class GaussianBlur2d(nn.Module):
    r"""Copied (and slightly adapted) from
    https://github.com/kornia/kornia/blob/master/kornia/filters

    Creates an operator that blurs a tensor using a Gaussian filter.
    The operator smooths the given tensor with a gaussian kernel by convolving
    it to each channel. It suports batched operation.
    Arguments:
        kernel_size: the size of the kernel.
        sigma_sampling_fn: the standard deviation of the kernel.
        pad_mode (str): the padding mode to be applied before convolving.
          The expected modes are: ``'constant'``, ``'reflect'``,
          ``'replicate'`` or ``'circular'``. Default: ``'reflect'``.
    Returns:
        Tensor: the blurred tensor.
    Shape:
        - Input: :math:`(B, C, H, W)`
        - Output: :math:`(B, C, H, W)`
    Examples::
        >>> x = torch.zeros(3, 1, 5, 5)
        >>> x[:,:, 2, 2] = 1
        >>> blur = GaussianBlur2d(3, lambda n: [.5, 1., 2.])
        >>> output = blur(x)
        >>> output.shape
        torch.Size([3, 1, 5, 5])
        >>> output
        tensor([[[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                  [0.0000, 0.0113, 0.0838, 0.0113, 0.0000],
                  [0.0000, 0.0838, 0.6193, 0.0838, 0.0000],
                  [0.0000, 0.0113, 0.0838, 0.0113, 0.0000],
                  [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]],
        <BLANKLINE>
        <BLANKLINE>
                [[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                  [0.0000, 0.0751, 0.1238, 0.0751, 0.0000],
                  [0.0000, 0.1238, 0.2042, 0.1238, 0.0000],
                  [0.0000, 0.0751, 0.1238, 0.0751, 0.0000],
                  [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]],
        <BLANKLINE>
        <BLANKLINE>
                [[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                  [0.0000, 0.1019, 0.1154, 0.1019, 0.0000],
                  [0.0000, 0.1154, 0.1308, 0.1154, 0.0000],
                  [0.0000, 0.1019, 0.1154, 0.1019, 0.0000],
                  [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])
    """

    def __init__(
            self, kernel_size, sigma_sampling_fn, pad_mode: str = 'reflect'
    ):
        super().__init__()
        self.kernel_size = kernel_size
        self.sigma_sampling_fn = sigma_sampling_fn
        assert pad_mode in ["constant", "reflect", "replicate", "circular"]
        self.pad_mode = pad_mode

    def __repr__(self) -> str:
        return self.__class__.__name__ +\
            '(kernel_size=' + str(self.kernel_size) + ', ' +\
            'sigma_sampling_fn=' + repr(self.sigma_sampling_fn) + ', ' +\
            'border_type=' + self.pad_mode + ')'

    def forward(self, x):

        if not x.dim() == 4:
            raise ValueError(
                f"Invalid input shape, we expect BxCxHxW. Got: {x.shape}")
        if not self.training:
            return x

        # pad the input tensor
        x = Pad(mode=self.pad_mode, side='both')(x, size=self.kernel_size-1)
        b, c, hp, wp = x.shape
        # convolve the tensor with the kernel.
        sigma = torch.from_numpy(np.array(self.sigma_sampling_fn(b))).float()
        kernel = get_gaussian_kernel2d(self.kernel_size, sigma).unsqueeze(1).to(x.device)
        return F.conv2d(
            x.transpose(0, 1), kernel, groups=b, padding=0, stride=1
        ).transpose(0, 1)


def get_gaussian_kernel2d(kernel_size, sigma, force_even: bool = False):
    r"""Copied (and slightly adapted) from
    https://github.com/kornia/kornia/blob/master/kornia/filters

    Function that returns Gaussian filter matrix coefficients.
    Args:
        kernel_size: filter sizes in the x and y direction.
         Sizes should be odd and positive.
        sigma: gaussian standard deviation in the x and y
         direction.
        force_even (bool): overrides requirement for odd kernel size.
    Returns:
        Tensor: 2D tensor with gaussian filter matrix coefficients.
    Shape:
        - Output: :math:`(\text{kernel_size}_x, \text{kernel_size}_y)`
    Examples::
        >>> get_gaussian_kernel2d(3, 1.5)
        tensor([[0.0947, 0.1183, 0.0947],
                [0.1183, 0.1478, 0.1183],
                [0.0947, 0.1183, 0.0947]])
        >>> get_gaussian_kernel2d(3, [1.5, .000001])
        tensor([[0.0947, 0.1183, 0.0947],
                [0.1183, 0.1478, 0.1183],
                [0.0947, 0.1183, 0.0947]])
    """
    kernel_1d: torch.Tensor = get_gaussian_kernel1d(
        kernel_size, sigma, force_even
    )
    kernel_2d: torch.Tensor = kernel_1d.unsqueeze(-1) @ kernel_1d.unsqueeze(-2)
    return kernel_2d


def get_gaussian_kernel1d(
        kernel_size: int, sigma: float, force_even: bool = False
) -> torch.Tensor:
    r"""Copied (and slightly adapted) from
    https://github.com/kornia/kornia/blob/master/kornia/filters

    Function that returns Gaussian filter coefficients.
    Args:
        kernel_size (int): filter size. It should be odd and positive.
        sigma (float): gaussian standard deviation.
        force_even (bool): overrides requirement for odd kernel size.
    Returns:
        Tensor: 1D tensor with gaussian filter coefficients.
    Shape:
        - Output: :math:`(\text{kernel_size})`
    Examples::
        >>> get_gaussian_kernel1d(3, 2.5)
        tensor([0.3243, 0.3513, 0.3243])
        >>> get_gaussian_kernel1d(5, 1.5)
        tensor([0.1201, 0.2339, 0.2921, 0.2339, 0.1201])
    """
    if (not isinstance(kernel_size, int) or (
            (kernel_size % 2 == 0) and not force_even) or (
            kernel_size <= 0)):
        raise TypeError(
            "kernel_size must be an odd positive integer. "
            f"Got {kernel_size}"
        )
    window_1d: torch.Tensor = gaussian(kernel_size, sigma)
    return window_1d


def gaussian(window_size, sigma):
    """Copied (and slightly adapted) from
    https://github.com/kornia/kornia/blob/master/kornia/filters

    Args:
        window_size:
        sigma:

    Returns:

    """
    x = torch.arange(window_size).float() - window_size // 2
    if torch.is_tensor(sigma) and sigma.dim() > 0:
        sigma = sigma[..., None]
    if window_size % 2 == 0:
        x = x + 0.5
    gauss = torch.exp((-x.pow(2.0) / (2 * sigma ** 2))).float()
    return gauss / gauss.sum(-1, keepdim=True)
