import torch
import torch.nn as nn
import math


class SpeechEmbedding(nn.Module):
    """
    """
    def __init__(self, args, _):
        super(SpeechEmbedding, self).__init__()
        self.conv = Conv1dModule(args)
        self.sinusoidalpos = False
        self.emb_size = args.emb_size
        if "sinusoidalpos" in args.embedding:
            self.sinusoidalpos = True

    def forward(self, src, _):
        """Embed inputs.
        Args:
            src (FloatTensor): Sequence of word vectors
                ``(batch_size, seq_len, self.dim)``
        """
        speech_emb = self.conv(src)
        if self.sinusoidalpos:
            return speech_emb * math.sqrt(self.emb_size)
        else:
            return speech_emb


class Transpose_module(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x.transpose(-2, -1)


class Conv1dModule(nn.Module):
    """
    Convolutional subsampler: a stack of 1D convolution (along temporal dimension) followed by non-linear activation
    via gated linear units (https://arxiv.org/abs/1911.08460)
    """

    def __init__(self, args):
        super(Conv1dModule, self).__init__()
        self.embedding_dim = args.emb_size
        self.norm_mode = None
        self.feature_grad_mult = 1.0
        self.conv_bias = True
        self.dropout_input = 0.0
        self.use_glu = True if args.data_processor == "s2t" else False
        self.padding = True

        self.conv_channels = args.conv_channels
        self.audio_feature_size = args.audio_feature_size
        self.kernel_sizes = args.conv_kernel_sizes
        self.strides = [2 for _ in range(len(self.kernel_sizes))]

        self.conv_layers = nn.ModuleList()
        
        def conv_layer_block(
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            norm_mode=None,
            conv_bias=False,
        ):
            def make_conv(in_channels, out_channels, kernel_size, stride, padding, conv_bias):
                conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=conv_bias)
                nn.init.kaiming_normal_(conv.weight)
                return conv

            if norm_mode == "layer":
                return nn.Sequential(
                    make_conv(in_channels, out_channels, kernel_size, stride, padding, conv_bias),
                    Transpose_module(),
                    nn.LayerNorm(out_channels, eps=1e-5, elementwise_affine=True),
                    Transpose_module(),
                    nn.GELU(),
                )
            elif norm_mode == "group":
                return nn.Sequential(
                    make_conv(in_channels, out_channels, kernel_size, stride, padding, conv_bias),
                    nn.GroupNorm(out_channels, out_channels, eps=1e-5, affine=True),
                    nn.GELU(),
                )
            elif self.use_glu:
                return nn.Sequential(
                    make_conv(in_channels, out_channels, kernel_size, stride, padding, conv_bias),
                )
            else:
                return nn.Sequential(
                    make_conv(in_channels, out_channels, kernel_size, stride, padding, conv_bias),
                    nn.GELU(),
                )
        assert len(self.strides) == len(self.kernel_sizes), "strides and kernel_sizes are not matched"
        assert len(self.strides) == len(self.conv_channels), "strides and conv_channels are not matched"
        in_channel = self.conv_channels[0] // 2
        for i, (k, s, c) in enumerate(zip(self.kernel_sizes, self.strides, self.conv_channels)):
            if self.audio_feature_size == 1:
                in_channel = c
            if self.norm_mode == "group" and i != 0:
                self.norm_mode = None
            if self.padding:
                padding = k // 2
            else:
                padding = 0
            self.conv_layers.append(
                conv_layer_block(
                    self.audio_feature_size if i == 0 else in_channel,
                    c,
                    k,
                    s,
                    padding,
                    norm_mode=self.norm_mode,
                    conv_bias=self.conv_bias,
                )
            )


    def forward(self, input_features, mask_indices=None, mask_channel_indices=None):
        if len(input_features.size()) == 2:
            hidden_states = input_features.unsqueeze(1) # wav B x T -> B x (C x D) x T
        else:
            hidden_states = input_features.transpose(1, 2).contiguous()  #acoustic feature B x T x (C x D) -> B x (C x D) x T

        for conv in self.conv_layers:
            hidden_states = conv(hidden_states)
            if self.use_glu:
                hidden_states = nn.functional.glu(hidden_states, dim=1)

        hidden_states = hidden_states.transpose(1, 2).contiguous()  # -> B x T x (C x D)

        return hidden_states