import numbers
import random

import numpy as np
import torch
from torchvision.transforms import RandomCrop, RandomResizedCrop


def _is_tensor_video_clip(clip):
    if not torch.is_tensor(clip):
        raise TypeError("clip should be Tesnor. Got %s" % type(clip))

    if not clip.ndimension() == 4:
        raise ValueError("clip should be 4D. Got %dD" % clip.dim())

    return True


def crop(clip, i, j, h, w):
    """
    Args:
        clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W)
    """
    assert len(clip.size()) == 4, "clip should be a 4D tensor"
    return clip[..., i: i + h, j: j + w]


def temporal_center_crop(clip, clip_len):
    """
    Args:
        clip (torch.tensor): Video clip to be
        cropped along the temporal axis. Size is (C, T, H, W)
    """
    assert len(clip.size()) == 4, "clip should be a 4D tensor"
    assert clip.size(1) >= clip_len, "clip is shorter than the proposed lenght"
    middle = int(clip.size(1) // 2)
    start = middle - clip_len // 2
    return clip[:, start: start + clip_len, ...]


def resize(clip, target_size, interpolation_mode):
    assert len(target_size) == 2, "target size should be tuple (height, width)"
    return torch.nn.functional.interpolate(clip, size=target_size, mode=interpolation_mode)


def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"):
    """
    Do spatial cropping and resizing to the video clip
    Args:
        clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W)
        i (int): i in (i,j) i.e coordinates of the upper left corner.
        j (int): j in (i,j) i.e coordinates of the upper left corner.
        h (int): Height of the cropped region.
        w (int): Width of the cropped region.
        size (tuple(int, int)): height and width of resized clip
    Returns:
        clip (torch.tensor): Resized and cropped clip. Size is (C, T, H, W)
    """
    assert _is_tensor_video_clip(clip), "clip should be a 4D torch.tensor"
    clip = crop(clip, i, j, h, w)
    clip = resize(clip, size, interpolation_mode)
    return clip


def center_crop(clip, crop_size):
    assert _is_tensor_video_clip(clip), "clip should be a 4D torch.tensor"
    h, w = clip.size(-2), clip.size(-1)
    th, tw = crop_size
    assert h >= th and w >= tw, "height and width must be >= than crop_size"

    i = int(round((h - th) / 2.0))
    j = int(round((w - tw) / 2.0))
    return crop(clip, i, j, th, tw)


def to_tensor(clip):
    """
    Convert tensor data type from uint8 to float, divide value by 255.0 and
    permute the dimenions of clip tensor
    Args:
        clip (torch.tensor, dtype=torch.uint8): Size is (T, H, W, C)
    Return:
        clip (torch.tensor, dtype=torch.float): Size is (C, T, H, W)
    """
    _is_tensor_video_clip(clip)
    if not clip.dtype == torch.uint8:
        raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype))
    return clip.float().permute(3, 0, 1, 2) / 255.0


def normalize(clip, mean, std, inplace=False):
    """
    Args:
        clip (torch.tensor): Video clip to be normalized. Size is (C, T, H, W)
        mean (tuple): pixel RGB mean. Size is (3)
        std (tuple): pixel standard deviation. Size is (3)
    Returns:
        normalized clip (torch.tensor): Size is (C, T, H, W)
    """
    assert _is_tensor_video_clip(clip), "clip should be a 4D torch.tensor"
    if not inplace:
        clip = clip.clone()
    mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device)
    std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device)
    clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None])
    return clip


def hflip(clip):
    """
    Args:
        clip (torch.tensor): Video clip to be normalized. Size is (C, T, H, W)
    Returns:
        flipped clip (torch.tensor): Size is (C, T, H, W)
    """
    assert _is_tensor_video_clip(clip), "clip should be a 4D torch.tensor"
    return clip.flip((-1))


class RandomCropVideo(RandomCrop):
    def __init__(self, size):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size

    def __call__(self, clip):
        """
        Args:
            clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W)
        Returns:
            torch.tensor: randomly cropped/resized video clip.
                size is (C, T, OH, OW)
        """
        i, j, h, w = self.get_params(clip, self.size)
        return crop(clip, i, j, h, w)

    def __repr__(self):
        return self.__class__.__name__ + "(size={0})".format(self.size)


class RandomResizedCropVideo(RandomResizedCrop):
    def __init__(
            self, size, scale=(0.08, 1.0), ratio=(3.0 / 4.0, 4.0 / 3.0), interpolation_mode="bilinear",
    ):
        if isinstance(size, tuple):
            assert len(size) == 2, "size should be tuple (height, width)"
            self.size = size
        else:
            self.size = (size, size)

        self.interpolation_mode = interpolation_mode
        self.scale = scale
        self.ratio = ratio

    def __call__(self, clip):
        """
        Args:
            clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W)
        Returns:
            torch.tensor: randomly cropped/resized video clip.
                size is (C, T, H, W)
        """
        i, j, h, w = self.get_params(clip, self.scale, self.ratio)
        return resized_crop(clip, i, j, h, w, self.size, self.interpolation_mode)

    def __repr__(self):
        return self.__class__.__name__ + "(size={0}, interpolation_mode={1}, scale={2}, ratio={3})".format(
            self.size, self.interpolation_mode, self.scale, self.ratio
        )


class CenterCropVideo(object):
    def __init__(self, crop_size):
        if isinstance(crop_size, numbers.Number):
            self.crop_size = (int(crop_size), int(crop_size))
        else:
            self.crop_size = crop_size

    def __call__(self, clip):
        """
        Args:
            clip (torch.tensor): Video clip to be cropped. Size is (C, T, H, W)
        Returns:
            torch.tensor: central cropping of video clip. Size is
            (C, T, crop_size, crop_size)
        """
        return center_crop(clip, self.crop_size)

    def __repr__(self):
        r = self.__class__.__name__ + "(crop_size={0})".format(self.crop_size)
        return r


class TemporalCenterCrop(object):
    def __init__(self, clip_len):
        self.clip_len = clip_len

    def __call__(self, clip):
        return temporal_center_crop(clip, self.clip_len)


class UnfoldClips(object):
    def __init__(self, clip_len, overlap):
        self.clip_len = clip_len
        assert overlap > 0 and overlap <= 1
        self.step = round(clip_len * overlap)

    def __call__(self, clip):
        if clip.size(1) < self.clip_len:
            return clip.unfold(1, clip.size(1), clip.size(1)).permute(1, 0, 4, 2, 3)

        results = clip.unfold(1, self.clip_len, self.clip_len).permute(1, 0, 4, 2, 3)
        return results


class TempPadClip(object):
    def __init__(self, clip_len):
        self.num_frames = clip_len

    def __call__(self, clip):
        if clip.size(1) == 0:
            return clip
        if clip.size(1) < self.num_frames:
            # do something and return
            step = clip.size(1) / self.num_frames
            idxs = torch.arange(self.num_frames, dtype=torch.float32) * step
            idxs = idxs.floor().to(torch.int64)
            return clip[:, idxs, ...]
        step = clip.size(1) / self.num_frames
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            return clip[:, slice(None, None, step), ...]
        idxs = torch.arange(self.num_frames, dtype=torch.float32) * step
        idxs = idxs.floor().to(torch.int64)
        return clip[:, idxs, ...]


class NormalizeVideo(object):
    """
    Normalize the video clip by mean subtraction
    and division by standard deviation
    Args:
        mean (3-tuple): pixel RGB mean
        std (3-tuple): pixel RGB standard deviation
        inplace (boolean): whether do in-place normalization
    """

    def __init__(self, mean, std, inplace=False):
        self.mean = mean
        self.std = std
        self.inplace = inplace

    def __call__(self, clip):
        """
        Args:
            clip (torch.tensor): video clip to be
                                normalized. Size is (C, T, H, W)
        """
        return normalize(clip, self.mean, self.std, self.inplace)

    def __repr__(self):
        return self.__class__.__name__ + "(mean={0}, std={1}, inplace={2})".format(self.mean, self.std, self.inplace)


class ToTensorVideo(object):
    """
    Convert tensor data type from uint8 to float, divide value by 255.0 and
    permute the dimensions of clip tensor
    """

    def __init__(self):
        pass

    def __call__(self, clip):
        """
        Args:
            clip (torch.tensor, dtype=torch.uint8): Size is (T, H, W, C)
        Return:
            clip (torch.tensor, dtype=torch.float): Size is (C, T, H, W)
        """
        return to_tensor(clip)

    def __repr__(self):
        return self.__class__.__name__


class RandomHorizontalFlipVideo(object):
    """
    Flip the video clip along the horizonal direction with a given probability
    Args:
        p (float): probability of the clip being flipped. Default value is 0.5
    """

    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, clip):
        """
        Args:
            clip (torch.tensor): Size is (C, T, H, W)
        Return:
            clip (torch.tensor): Size is (C, T, H, W)
        """
        if random.random() < self.p:
            clip = hflip(clip)
        return clip

    def __repr__(self):
        return self.__class__.__name__ + "(p={0})".format(self.p)


class Resize(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, vid):
        return resize(vid, self.size, interpolation_mode="bilinear")


class TimeMask:
    """time mask for 1d specAug"""

    def __init__(self, T=8, n_mask=1, replace_with_zero=False):
        self.n_mask = n_mask
        self.T = T
        self.replace_with_zero = replace_with_zero

    def __call__(self, x, idxs=None):
        cloned = x.clone()
        len_raw = cloned.size(1)

        if idxs is None:
            idxs = []
            for i in range(self.n_mask):
                t = random.randrange(self.T)
                if len_raw - t <= 0:
                    continue
                t_zero = random.randrange(0, len_raw - t)
                # avoids randrange error if values are equal and range is empty
                if t_zero == t_zero + t:
                    continue
                mask_end = t_zero + t
                idxs.extend(list(range(t_zero, mask_end)))

        if self.replace_with_zero:
            cloned[:, idxs] = 0
        else:
            cloned[:, idxs] = cloned.mean()
        return cloned


class TimeMaskAudio:
    """time mask for 1d specAug"""

    def __init__(self, T=8, n_mask=1, downsample=4, replace_with_zero=False):
        self.n_mask = n_mask
        self.T = T
        self.downsample = downsample  # Only used when idxs is supplied in __call__
        self.replace_with_zero = replace_with_zero

    def __call__(self, x, idxs=None):
        cloned = x.clone()
        len_raw = cloned.size(1)

        if idxs is None:
            idxs = []
            for i in range(self.n_mask):
                t = random.randrange(self.T)
                if len_raw - t <= 0:
                    continue
                t_zero = random.randrange(0, len_raw - t)
                # avoids randrange error if values are equal and range is empty
                if t_zero == t_zero + t:
                    continue
                mask_end = t_zero + t
                idxs.extend(list(range(t_zero, mask_end)))
        else:
            idxs = [idx * self.downsample + j for idx in idxs for j in range(self.downsample)]

        if self.replace_with_zero:
            cloned[:, idxs] = 0
        else:
            cloned[:, idxs] = cloned.mean()
        return cloned
    

class FrequencyMask:
    """time mask for 1d specAug"""

    def __init__(self, F=8, n_mask=1, replace_with_zero=False):
        self.n_mask = n_mask
        self.F = F
        self.replace_with_zero = replace_with_zero

    def __call__(self, x):
        cloned = x.clone()
        len_raw = cloned.size(2)
        for i in range(self.n_mask):
            f = random.randrange(self.F)
            if len_raw - f <= 0:
                continue
            f_zero = random.randrange(0, len_raw - f)
            # avoids randrange error if values are equal and range is empty
            if f_zero == f_zero + f:
                continue
            mask_end = f_zero + f
            if self.replace_with_zero:
                cloned[:, :, f_zero:mask_end] = 0
            else:
                cloned[:, :, f_zero:mask_end] = cloned.mean()
        return cloned


class AdaptiveLengthTimeMask:
    def __init__(self, window, stride, replace_with_zero=False, inplace=False):
        self.window = window
        self.stride = stride

        self.replace_with_zero = replace_with_zero
        self.inplace = inplace

    def __call__(self, x):
        if self.inplace:
            cloned = x
        else:
            cloned = x.clone()

        len_raw = cloned.size(1)

        n_mask = int((len_raw + self.stride-0.1) // self.stride)

        ts = torch.randint(0, self.window, size=(n_mask, 2))
        for t, t_end in ts:
            if len_raw - t <= 0:
                continue
            t_start = random.randrange(0, len_raw - t)

            # avoids randrange error if values are equal and range is empty
            if t_start == t_start + t:
                continue

            t_end += t_start
            if self.replace_with_zero:
                cloned[:, t_start:t_end] = 0
            else:
                cloned[:, t_start:t_end] = cloned.mean()
        return cloned


class AddNoise(object):
    """Add SNR noise [-1, 1]
    """

    def __init__(self, noise=np.load('/data/home/pingchuanma/babble_noise/babbleNoise_resample_16K.npy'), snr_target=None, snr_levels=[-5, 0, 5, 10, 15, 20, 9999]):
        assert noise.dtype in [np.float32, np.float64], "noise only supports float data type"
        self.noise = noise
        self.snr_levels = snr_levels
        self.snr_target = snr_target

    def get_power(self, clip):
        clip2 = clip.copy()
        clip2 = clip2 **2
        return np.sum(clip2) / (len(clip2) * 1.0)

    def __call__(self, signal):
        device = signal.device
        signal = signal[0].numpy()
        assert signal.dtype in [np.float32, np.float64], "signal only supports float32 data type"
        snr_target = random.choice(self.snr_levels) if not self.snr_target else self.snr_target
        if snr_target == 9999:
            return torch.tensor(signal, device=device)
        else:
            # -- get noise
            start_idx = random.randint(0, len(self.noise)-len(signal))
            noise_clip = self.noise[start_idx:start_idx+len(signal)]
            sig_power = self.get_power(signal)
            noise_clip_power = self.get_power(noise_clip)
            factor = (sig_power / noise_clip_power ) / (10**(snr_target / 10.0))
            desired_signal = (signal + noise_clip*np.sqrt(factor)).astype(np.float32)
            torch_desired_signal = torch.tensor(desired_signal, device=device)
            return torch.unsqueeze(torch_desired_signal, 0)
