import copy

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
from torchaudio import transforms
from collections import namedtuple
from src.nb_SpecAugment import time_warp, time_mask, freq_mask

# class CMVN(torch.jit.ScriptModule):
class CMVN(nn.Module):

    __constants__ = ["mode", "dim", "eps"]

    def __init__(self, mode="global", dim=2, eps=1e-10):
        # `torchaudio.load()` loads audio with shape [channel, feature_dim, time]
        # so perform normalization on dim=2 by default
        super(CMVN, self).__init__()

        if mode != "global":
            raise NotImplementedError(
                "Only support global mean variance normalization.")

        self.mode = mode
        self.dim = dim
        self.eps = eps

    # @torch.jit.script_method
    def forward(self, x):
        if self.mode == "global":
            return (x - x.mean(self.dim, keepdim=True)) / (self.eps + x.std(self.dim, keepdim=True))

    def extra_repr(self):
        return "mode={}, dim={}, eps={}".format(self.mode, self.dim, self.eps)


# class Delta(torch.jit.ScriptModule):
class Delta(nn.Module):

    __constants__ = ["order", "window_size", "padding"]

    def __init__(self, order=1, window_size=2):
        # Reference:
        # https://kaldi-asr.org/doc/feature-functions_8cc_source.html
        # https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/layers/common_audio.py
        super(Delta, self).__init__()

        self.order = order
        self.window_size = window_size

        self.filters = self._create_filters(order, window_size)
        # self.register_buffer("filters", filters)
        # print(filters.shape)
        self.padding = (0, (self.filters.shape[-1] - 1) // 2)
        # (out_channels, in_channels, kernel_size[0]. kernel_size[1]): (3, 1, 1, 9)
        self.conv = nn.Conv2d(1, 3, (1, 9), padding=self.padding, bias=False)
        self.conv.weight = nn.Parameter(self.filters)

    # @torch.jit.script_method
    def forward(self, x):
        # Unsqueeze batch dim
        # print('conv_device', next(self.conv.parameters()).device)
        # print('x device', x.device)
        x = x.unsqueeze(0)
        return self.conv(x)[0]
        # return F.conv2d(x, weight=self.filters, padding=self.padding)[0]

    # TODO(WindQAQ): find more elegant way to create `scales`
    def _create_filters(self, order, window_size):
        scales = [[1.0]]
        for i in range(1, order + 1):
            prev_offset = (len(scales[i-1]) - 1) // 2
            curr_offset = prev_offset + window_size

            curr = [0] * (len(scales[i-1]) + 2 * window_size)
            normalizer = 0.0
            for j in range(-window_size, window_size + 1):
                normalizer += j * j
                for k in range(-prev_offset, prev_offset + 1):
                    curr[j+k+curr_offset] += (j * scales[i-1][k+prev_offset])
            curr = [x / normalizer for x in curr]
            scales.append(curr)

        max_len = len(scales[-1])
        for i, scale in enumerate(scales[:-1]):
            padding = (max_len - len(scale)) // 2
            scales[i] = [0] * padding + scale + [0] * padding

        return torch.tensor(scales).unsqueeze(1).unsqueeze(1)

    def extra_repr(self):
        return "order={}, window_size={}".format(self.order, self.window_size)


# class Postprocess(torch.jit.ScriptModule):
class Postprocess(nn.Module):
    # @torch.jit.script_method
    def forward(self, x):
        # [channel, feature_dim, time] -> [time, channel, feature_dim]
        x = x.permute(2, 0, 1)
        # [time, channel, feature_dim] -> [time, feature_dim * channel]
        return x.reshape(x.size(0), -1).detach()
# class SpecAug(torch.jit.ScriptModule):
class SpecAug(nn.Module):
    #@torch.jit.script_method
    def forward(self, x):
        # [time, feature_dim * channel]
        x = time_mask(freq_mask(time_warp(x.permute(1, 0).unsqueeze(0)), num_masks=2), num_masks=2)
        #x = time_mask(freq_mask(time_warp(x), num_masks=2), num_masks=2)
        return x.squeeze(0).permute(1,0)

# TODO(Windqaq): make this scriptable
class ExtractAudioFeature(nn.Module):
    def __init__(self, mode="fbank", num_mel_bins=40, sample_rate=16000, **kwargs):
        super(ExtractAudioFeature, self).__init__()
        self.mode = mode
        self.extract_fn = torchaudio.compliance.kaldi.fbank if mode == "fbank" else torchaudio.compliance.kaldi.mfcc
        self.num_mel_bins = num_mel_bins
        self.sample_rate = sample_rate
        self.kwargs = kwargs

    def forward(self, waveform):
        # waveform, sample_rate = torchaudio.load(filepath)
        y = self.extract_fn(waveform,
                            num_mel_bins=self.num_mel_bins,
                            channel=-1,
                            sample_frequency=self.sample_rate,
                            **self.kwargs)
        return y.transpose(0, 1).unsqueeze(0).detach()

    def extra_repr(self):
        return "mode={}, num_mel_bins={}".format(self.mode, self.num_mel_bins)


def create_transform(audio_config):
    audio_config = copy.deepcopy(audio_config)
    feat_type = audio_config.pop("feat_type")
    feat_dim = audio_config.pop("feat_dim")
    sample_rate = audio_config.pop("sample_rate")

    delta_order = audio_config.pop("delta_order", 0)
    delta_window_size = audio_config.pop("delta_window_size", 2)
    apply_cmvn = audio_config.pop("apply_cmvn")

    transforms = [ExtractAudioFeature(feat_type, feat_dim, sample_rate, **audio_config)]

    if delta_order >= 1:
        transforms.append(Delta(delta_order, delta_window_size))

    if apply_cmvn:
        transforms.append(CMVN())

    transforms.append(Postprocess())
    transforms.append(SpecAug())

    return nn.Sequential(*transforms), feat_dim * (delta_order + 1), sample_rate
