import math
import torch  

import torch.nn as nn
import torch.nn.functional as F

def length_to_mask(length, max_len=None, dtype=None, device=None):
    assert len(length.shape) == 1

    if max_len is None: max_len = length.max().long().item() 

    mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand(len(length), max_len) < length.unsqueeze(1)

    if dtype is None: dtype = length.dtype
    if device is None: device = length.device

    return torch.as_tensor(mask, dtype=dtype, device=device)

def get_padding_elem(L_in, stride, kernel_size, dilation):
    if stride > 1: padding = [math.floor(kernel_size / 2), math.floor(kernel_size / 2)]
    else:
        L_out = (math.floor((L_in - dilation * (kernel_size - 1) - 1) / stride) + 1)
        padding = [math.floor((L_in - L_out) / 2), math.floor((L_in - L_out) / 2)]

    return padding

class _BatchNorm1d(nn.Module):
    def __init__(self, input_shape=None, input_size=None, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True, combine_batch_time=False, skip_transpose=False):
        super().__init__()
        self.combine_batch_time = combine_batch_time
        self.skip_transpose = skip_transpose

        if input_size is None and skip_transpose: input_size = input_shape[1]
        elif input_size is None: input_size = input_shape[-1]

        self.norm = nn.BatchNorm1d(input_size, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)

    def forward(self, x):
        shape_or = x.shape

        if self.combine_batch_time:x = x.reshape(shape_or[0] * shape_or[1], shape_or[2]) if x.ndim == 3 else x.reshape(shape_or[0] * shape_or[1], shape_or[3], shape_or[2])
        elif not self.skip_transpose: x = x.transpose(-1, 1)

        x_n = self.norm(x)

        if self.combine_batch_time: x_n = x_n.reshape(shape_or)
        elif not self.skip_transpose: x_n = x_n.transpose(1, -1)

        return x_n

class _Conv1d(nn.Module):
    def __init__(self, out_channels, kernel_size, input_shape=None, in_channels=None, stride=1, dilation=1, padding="same", groups=1, bias=True, padding_mode="reflect", skip_transpose=False, weight_norm=False, conv_init=None, default_padding=0):
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.dilation = dilation
        self.padding = padding
        self.padding_mode = padding_mode
        self.unsqueeze = False
        self.skip_transpose = skip_transpose

        if input_shape is None and in_channels is None: raise ValueError
        if in_channels is None: in_channels = self._check_input_shape(input_shape)

        self.in_channels = in_channels
        self.conv = nn.Conv1d(in_channels, out_channels, self.kernel_size, stride=self.stride, dilation=self.dilation, padding=default_padding, groups=groups, bias=bias)

        if conv_init == "kaiming": nn.init.kaiming_normal_(self.conv.weight)
        elif conv_init == "zero": nn.init.zeros_(self.conv.weight)
        elif conv_init == "normal": nn.init.normal_(self.conv.weight, std=1e-6)

        if weight_norm: self.conv = nn.utils.weight_norm(self.conv)

    def forward(self, x):
        if not self.skip_transpose: x = x.transpose(1, -1)
        if self.unsqueeze: x = x.unsqueeze(1)

        if self.padding == "same": x = self._manage_padding(x, self.kernel_size, self.dilation, self.stride)
        elif self.padding == "causal": x = F.pad(x, ((self.kernel_size - 1) * self.dilation, 0))
        elif self.padding == "valid": pass
        else: raise ValueError

        wx = self.conv(x)

        if self.unsqueeze: wx = wx.squeeze(1)
        if not self.skip_transpose: wx = wx.transpose(1, -1)

        return wx

    def _manage_padding(self, x, kernel_size, dilation, stride):
        return F.pad(x, get_padding_elem(self.in_channels, stride, kernel_size, dilation), mode=self.padding_mode)

    def _check_input_shape(self, shape):
        if len(shape) == 2:
            self.unsqueeze = True
            in_channels = 1
        elif self.skip_transpose: in_channels = shape[1]
        elif len(shape) == 3: in_channels = shape[2]
        else: raise ValueError

        if not self.padding == "valid" and self.kernel_size % 2 == 0: raise ValueError
        return in_channels

    def remove_weight_norm(self):
        self.conv = nn.utils.remove_weight_norm(self.conv)

class Linear(torch.nn.Module):
    def __init__(self, n_neurons, input_shape=None, input_size=None, bias=True, max_norm=None, combine_dims=False):
        super().__init__()
        self.max_norm = max_norm
        self.combine_dims = combine_dims

        if input_shape is None and input_size is None: raise ValueError
        if input_size is None:
            input_size = input_shape[-1]
            if len(input_shape) == 4 and self.combine_dims: input_size = input_shape[2] * input_shape[3]

        self.w = nn.Linear(input_size, n_neurons, bias=bias)

    def forward(self, x):
        if x.ndim == 4 and self.combine_dims: x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3])
        if self.max_norm is not None: self.w.weight.data = torch.renorm(self.w.weight.data, p=2, dim=0, maxnorm=self.max_norm)

        return self.w(x)

class Conv1d(_Conv1d):
    def __init__(self, *args, **kwargs):
        super().__init__(skip_transpose=True, *args, **kwargs)

class BatchNorm1d(_BatchNorm1d):
    def __init__(self, *args, **kwargs):
        super().__init__(skip_transpose=True, *args, **kwargs)

class TDNNBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation, activation=nn.ReLU, groups=1, dropout=0.0):
        super().__init__()
        self.conv = Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, dilation=dilation, groups=groups)
        self.activation = activation()
        self.norm = BatchNorm1d(input_size=out_channels)
        self.dropout = nn.Dropout1d(p=dropout)

    def forward(self, x):
        return self.dropout(self.norm(self.activation(self.conv(x))))

class Res2NetBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1, dropout=0.0):
        super().__init__()
        assert in_channels % scale == 0
        assert out_channels % scale == 0
        in_channel = in_channels // scale
        hidden_channel = out_channels // scale
        self.blocks = nn.ModuleList([TDNNBlock(in_channel, hidden_channel, kernel_size=kernel_size, dilation=dilation, dropout=dropout) for _ in range(scale - 1)])
        self.scale = scale

    def forward(self, x):
        y = []

        for i, x_i in enumerate(torch.chunk(x, self.scale, dim=1)):
            if i == 0: y_i = x_i
            elif i == 1: y_i = self.blocks[i - 1](x_i)
            else: y_i = self.blocks[i - 1](x_i + y_i)

            y.append(y_i)

        return torch.cat(y, dim=1)

class SEBlock(nn.Module):
    def __init__(self, in_channels, se_channels, out_channels):
        super().__init__()

        self.conv1 = Conv1d(in_channels=in_channels, out_channels=se_channels, kernel_size=1)
        self.relu = torch.nn.ReLU(inplace=True)
        self.conv2 = Conv1d(in_channels=se_channels, out_channels=out_channels, kernel_size=1)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x, lengths=None):
        L = x.shape[-1]

        if lengths is not None:
            mask = length_to_mask(lengths * L, max_len=L, device=x.device).unsqueeze(1)
            s = (x * mask).sum(dim=2, keepdim=True) / mask.sum(dim=2, keepdim=True)
        else: s = x.mean(dim=2, keepdim=True)

        return self.sigmoid(self.conv2(self.relu(self.conv1(s)))) * x

class AttentiveStatisticsPooling(nn.Module):
    def __init__(self, channels, attention_channels=128, global_context=True):
        super().__init__()
        self.eps = 1e-12
        self.global_context = global_context
        self.tdnn = TDNNBlock(channels * 3, attention_channels, 1, 1) if global_context else TDNNBlock(channels, attention_channels, 1, 1)
        self.tanh = nn.Tanh()
        self.conv = Conv1d(in_channels=attention_channels, out_channels=channels, kernel_size=1)

    def forward(self, x, lengths=None):
        L = x.shape[-1]

        def _compute_statistics(x, m, dim=2, eps=self.eps):
            mean = (m * x).sum(dim)
            return mean, torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(eps))

        if lengths is None: lengths = torch.ones(x.shape[0], device=x.device)
        mask = length_to_mask(lengths * L, max_len=L, device=x.device).unsqueeze(1)

        if self.global_context:
            mean, std = _compute_statistics(x, mask / mask.sum(dim=2, keepdim=True).float())
            attn = torch.cat([x, mean.unsqueeze(2).repeat(1, 1, L), std.unsqueeze(2).repeat(1, 1, L)], dim=1)
        else: attn = x

        mean, std = _compute_statistics(x, F.softmax(self.conv(self.tanh(self.tdnn(attn))).masked_fill(mask == 0, float("-inf")), dim=2))
        return torch.cat((mean, std), dim=1).unsqueeze(2)

class SERes2NetBlock(nn.Module):
    def __init__(self, in_channels, out_channels, res2net_scale=8, se_channels=128, kernel_size=1, dilation=1, activation=torch.nn.ReLU, groups=1, dropout=0.0):
        super().__init__()
        self.out_channels = out_channels
        self.tdnn1 = TDNNBlock(in_channels, out_channels, kernel_size=1, dilation=1, activation=activation, groups=groups, dropout=dropout)
        self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
        self.tdnn2 = TDNNBlock(out_channels, out_channels, kernel_size=1, dilation=1, activation=activation, groups=groups, dropout=dropout)
        self.se_block = SEBlock(out_channels, se_channels, out_channels)

        self.shortcut = None
        if in_channels != out_channels: self.shortcut = Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)

    def forward(self, x, lengths=None):
        residual = x
        if self.shortcut: residual = self.shortcut(x)

        return self.se_block(self.tdnn2(self.res2net_block(self.tdnn1(x))), lengths) + residual

class ECAPA_TDNN(torch.nn.Module):
    def __init__(self, input_size, device="cpu", lin_neurons=192, activation=torch.nn.ReLU, channels=[512, 512, 512, 512, 1536], kernel_sizes=[5, 3, 3, 3, 1], dilations=[1, 2, 3, 4, 1], attention_channels=128, res2net_scale=8, se_channels=128, global_context=True, groups=[1, 1, 1, 1, 1], dropout=0.0):
        super().__init__()
        assert len(channels) == len(kernel_sizes)
        assert len(channels) == len(dilations)

        self.channels = channels
        self.blocks = nn.ModuleList()

        self.blocks.append(TDNNBlock(input_size, channels[0], kernel_sizes[0], dilations[0], activation, groups[0], dropout))

        for i in range(1, len(channels) - 1):
            self.blocks.append(SERes2NetBlock(channels[i - 1], channels[i], res2net_scale=res2net_scale, se_channels=se_channels, kernel_size=kernel_sizes[i], dilation=dilations[i], activation=activation, groups=groups[i], dropout=dropout))

        self.mfa = TDNNBlock(channels[-2] * (len(channels) - 2), channels[-1], kernel_sizes[-1], dilations[-1], activation, groups=groups[-1], dropout=dropout)
        self.asp = AttentiveStatisticsPooling(channels[-1], attention_channels=attention_channels, global_context=global_context)
        self.asp_bn = BatchNorm1d(input_size=channels[-1] * 2)
        self.fc = Conv1d(in_channels=channels[-1] * 2, out_channels=lin_neurons, kernel_size=1)

    def forward(self, x, lengths=None):
        x = x.transpose(1, 2)

        xl = []
        for layer in self.blocks:
            try:
                x = layer(x, lengths=lengths)
            except TypeError:
                x = layer(x)

            xl.append(x)

        return self.fc(self.asp_bn(self.asp(self.mfa(torch.cat(xl[1:], dim=1)), lengths=lengths))).transpose(1, 2)

class Classifier(torch.nn.Module):
    def __init__(self, input_size, device="cpu", lin_blocks=0, lin_neurons=192, out_neurons=1211):
        super().__init__()
        self.blocks = nn.ModuleList()

        for _ in range(lin_blocks):
            self.blocks.extend([_BatchNorm1d(input_size=input_size), Linear(input_size=input_size, n_neurons=lin_neurons)])
            input_size = lin_neurons

        self.weight = nn.Parameter(torch.FloatTensor(out_neurons, input_size, device=device))
        nn.init.xavier_uniform_(self.weight)

    def forward(self, x):
        for layer in self.blocks:
            x = layer(x)

        return F.linear(F.normalize(x.squeeze(1)), F.normalize(self.weight)).unsqueeze(1)