# --------------------------------------------------------
# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
# Github source: https://github.com/mbzuai-nlp/ArTST

# Based on speecht5, fairseq and espnet code bases
# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
# --------------------------------------------------------

import logging
import math
import torch
import contextlib
from typing import List, Tuple
import torch.nn as nn

from fairseq.data.data_utils import lengths_to_padding_mask
from fairseq.data.data_utils import compute_mask_indices
from fairseq.modules import (
    PositionalEmbedding,
    Fp32GroupNorm,
    FairseqDropout,
    SamePad,
    GradMultiply,
    LayerNorm,
    Fp32LayerNorm,
    TransposeLast,
)
import numpy as np

logger = logging.getLogger(__name__)


class LinearLayer(nn.Module):
    def __init__(self, idim, odom, dropout=0):
        super(LinearLayer, self).__init__()
        self.linear = nn.Sequential(
            nn.Linear(idim, odom),
            nn.LayerNorm(odom),
            nn.Dropout(dropout),
            nn.ReLU(),
        )

    def get_out_seq_lens_tensor(self, in_seq_lens_tensor):
        out = in_seq_lens_tensor.clone()
        return out
    
    def forward(self, src_tokens, src_lengths):
        """
        src_tokens: [B, T, C]
        src_lengths: [B]
        """
        x = self.linear(src_tokens)
        x = x.transpose(0, 1).contiguous() # -> T x B x C
        return x, src_lengths


class SpeechEncoderPrenet(nn.Module):
    """

    Args:
        in_channels (int): the number of input channels
        mid_channels (int): the number of intermediate channels
        out_channels (int): the number of output channels
        kernel_sizes (List[int]): the kernel size for each convolutional layer
    """

    def __init__(self, args):
        super(SpeechEncoderPrenet, self).__init__()
        self.dropout_module = FairseqDropout(
            p=args.dropout, module_name=self.__class__.__name__
        )
        self.embed_scale = math.sqrt(args.encoder_embed_dim)
        if args.no_scale_embedding:
            self.embed_scale = 1.0
        self.padding_idx = 1
        self.freeze_encoder_updates = args.freeze_encoder_updates
        self.num_updates = 0
        assert args.encoder_speech_prenet in ["conv", "linear"], args.encoder_speech_prenet
        feature_enc_layers = eval(args.conv_feature_layers)  # noqa
        self.embed = feature_enc_layers[-1][0]

        self.feature_extractor = ConvFeatureExtractionModel(
            conv_layers=feature_enc_layers,
            dropout=0.0,
            mode=args.extractor_mode,
            conv_bias=args.conv_bias,
        )
        feature_ds_rate = np.prod([s for _, _, s in feature_enc_layers])
        self.feat2tar_ratio = (
            args.label_rates * feature_ds_rate / args.sample_rate
        )

        self.post_extract_proj = (
            nn.Linear(self.embed, args.encoder_embed_dim)
            if self.embed != args.encoder_embed_dim
            else None
        )

        self.use_conv_pos = args.use_conv_pos
        self.use_sinc_pos = args.use_sinc_pos
        self.use_abs_pos = getattr(args, "use_abs_pos", False)

        self.feature_grad_mult = args.feature_grad_mult
        if self.use_conv_pos:
            self.layer_norm = LayerNorm(self.embed)
            self.pos_conv = nn.Conv1d(
                args.encoder_embed_dim,
                args.encoder_embed_dim,
                kernel_size=args.conv_pos,
                padding=args.conv_pos // 2,
                groups=args.conv_pos_groups,
            )
            dropout = 0
            std = math.sqrt((4 * (1.0 - dropout)) / (args.conv_pos * args.encoder_embed_dim))
            nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
            nn.init.constant_(self.pos_conv.bias, 0)
            self.pos_conv = nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)
            self.pos_conv = nn.Sequential(self.pos_conv, SamePad(args.conv_pos), nn.GELU())

        assert not (self.use_sinc_pos and self.use_abs_pos), f"sinc pos: {self.use_sinc_pos} abs pos: {self.use_abs_pos}"
        if self.use_sinc_pos:
            self.embed_positions = PositionalEmbedding(
                args.max_speech_positions, args.encoder_embed_dim, self.padding_idx
            )
        if self.use_abs_pos:
            self.embed_positions = PositionalEmbedding(
                args.max_speech_positions, args.encoder_embed_dim, self.padding_idx, learned=True
            )
        
        # Hubert
        self.mask_prob = args.mask_prob
        self.mask_selection = args.mask_selection
        self.mask_other = args.mask_other
        self.hubert_mask_length = args.hubert_mask_length
        self.no_mask_overlap = args.no_mask_overlap
        self.mask_min_space = args.mask_min_space

        self.mask_channel_prob = args.mask_channel_prob
        self.mask_channel_selection = args.mask_channel_selection
        self.mask_channel_other = args.mask_channel_other
        self.mask_channel_length = args.mask_channel_length
        self.no_mask_channel_overlap = args.no_mask_channel_overlap
        self.mask_channel_min_space = args.mask_channel_min_space

        self.mask_emb = nn.Parameter(
            torch.FloatTensor(args.encoder_embed_dim).uniform_()
        )

    def forward(self, src_tokens, require_feat_pen=False, target_list=None, padding_mask=None, mask=True):
        ft = self.freeze_encoder_updates <= self.num_updates
        with torch.no_grad() if not ft else contextlib.ExitStack():
            return self._forward(src_tokens, require_feat_pen, target_list, padding_mask, mask)

    def _forward(self, src_tokens, require_feat_pen=False, target_list=None, padding_mask=None, mask=True):
        if self.feature_grad_mult > 0:
            x = self.feature_extractor(src_tokens)
            x = x.transpose(1, 2).transpose(0, 1)  # [length, batch, hidden_size]
            if self.feature_grad_mult != 1.0:
                x = GradMultiply.apply(x, self.feature_grad_mult)
        else:
            with torch.no_grad():
                x = self.feature_extractor(src_tokens)
                x = x.transpose(1, 2).transpose(0, 1)  # [length, batch, hidden_size]
        x = x.transpose(0, 1) # [batch, length, hidden_size]

        encoder_padding_mask = padding_mask

        x = x.transpose(1, 2) # [batch, hidden_size, length]
        if target_list is not None:
            x, target_list = self.forward_targets(x, target_list)
        features_pen = x.float().pow(2).mean()
        x = x.transpose(1, 2) # [batch, length, hidden_size]
        x = self.layer_norm(x)
        encoder_padding_mask = self.forward_padding_mask(x, encoder_padding_mask)
        if self.post_extract_proj is not None:
            x = self.post_extract_proj(x)
        x = self.dropout_module(x)
        if mask:
            x, mask_indices = self.apply_hubert_mask(
                x, encoder_padding_mask
            )
        else:
            x = x
            mask_indices = None

        if self.use_conv_pos:
            positions = self.pos_conv(x.transpose(1, 2))
            positions = positions.transpose(1, 2)
        #else:
        #    positions = self.embed_positions(encoder_padding_mask)
            x = x + positions

        if self.use_sinc_pos:
            positions = self.embed_positions(encoder_padding_mask)
            x = x + positions

        # x = self.dropout_module(x)

        if require_feat_pen:
            return (x, features_pen, mask_indices, target_list), encoder_padding_mask
        else:
            # For consistence with encoder
            return x, encoder_padding_mask

    def forward_targets(
        self, features: torch.Tensor, target_list: List[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Trim features to ensure labels exist and then get aligned labels
        feat_tsz = features.size(2)
        targ_tsz = min([t.size(1) for t in target_list])
        if self.feat2tar_ratio * feat_tsz > targ_tsz:
            feat_tsz = int(targ_tsz / self.feat2tar_ratio)
            features = features[..., :feat_tsz]
        target_inds = torch.arange(feat_tsz).float() * self.feat2tar_ratio
        target_list = [t[:, target_inds.long()] for t in target_list]
        return features, target_list

    def forward_padding_mask(
        self, features: torch.Tensor, padding_mask: torch.Tensor,
    ) -> torch.Tensor:
        extra = padding_mask.size(1) % features.size(1)
        if extra > 0:
            padding_mask = padding_mask[:, :-extra]
        padding_mask = padding_mask.view(
            padding_mask.size(0), features.size(1), -1
        )
        padding_mask = padding_mask.all(-1)
        return padding_mask

    def get_src_lengths(self, src_lengths):
        return self.feature_extractor.get_out_seq_lens_tensor(src_lengths)

    def apply_hubert_mask(self, x, padding_mask):
        B, T, C = x.shape
        if self.mask_prob > 0:
            mask_indices = compute_mask_indices(
                (B, T),
                padding_mask,
                self.mask_prob,
                self.hubert_mask_length,
                self.mask_selection,
                self.mask_other,
                min_masks=2,
                no_overlap=self.no_mask_overlap,
                min_space=self.mask_min_space,
            )
            mask_indices = torch.from_numpy(mask_indices).to(x.device)
            x[mask_indices] = self.mask_emb
        else:
            mask_indices = None

        if self.mask_channel_prob > 0:
            mask_channel_indices = compute_mask_indices(
                (B, C),
                None,
                self.mask_channel_prob,
                self.mask_channel_length,
                self.mask_channel_selection,
                self.mask_channel_other,
                no_overlap=self.no_mask_channel_overlap,
                min_space=self.mask_channel_min_space,
            )
            mask_channel_indices = (
                torch.from_numpy(mask_channel_indices)
                .to(x.device)
                .unsqueeze(1)
                .expand(-1, T, -1)
            )
            x[mask_channel_indices] = 0

        return x, mask_indices

    def set_num_updates(self, num_updates):
        """Set the number of parameters updates."""
        self.num_updates = num_updates

class ConvFeatureExtractionModel(nn.Module):
    def __init__(
        self,
        conv_layers: List[Tuple[int, int, int]],
        dropout: float = 0.0,
        mode: str = "default",
        conv_bias: bool = False,
    ):
        super().__init__()

        assert mode in {"default", "layer_norm"}

        def block(
            n_in,
            n_out,
            k,
            stride,
            is_layer_norm=False,
            is_group_norm=False,
            conv_bias=False,
        ):
            def make_conv():
                conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
                nn.init.kaiming_normal_(conv.weight)
                return conv

            assert (
                is_layer_norm and is_group_norm
            ) == False, "layer norm and group norm are exclusive"

            if is_layer_norm:
                return nn.Sequential(
                    make_conv(),
                    nn.Dropout(p=dropout),
                    nn.Sequential(
                        TransposeLast(),
                        Fp32LayerNorm(dim, elementwise_affine=True),
                        TransposeLast(),
                    ),
                    nn.GELU(),
                )
            elif is_group_norm:
                return nn.Sequential(
                    make_conv(),
                    nn.Dropout(p=dropout),
                    Fp32GroupNorm(dim, dim, affine=True),
                    nn.GELU(),
                )
            else:
                return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU())

        in_d = 1
        self.conv_layers = nn.ModuleList()
        self.conv_layers_infos = conv_layers
        for i, cl in enumerate(conv_layers):
            assert len(cl) == 3, "invalid conv definition: " + str(cl)
            (dim, k, stride) = cl

            self.conv_layers.append(
                block(
                    in_d,
                    dim,
                    k,
                    stride,
                    is_layer_norm=mode == "layer_norm",
                    is_group_norm=mode == "default" and i == 0,
                    conv_bias=conv_bias,
                )
            )
            in_d = dim

    def forward(self, x):
        # BxT -> BxCxT
        x = x.unsqueeze(1)
        for conv in self.conv_layers:
            x = conv(x)
        return x

    def get_out_seq_lens_nonmask_after_a_layer(self, in_seq_lens_tensor, i):
        """Returns the out_seq_lens_nonmask 0/1 tensor after a layer.

        Args:
            in_seq_lens_tensor (LongTensor): length

        Returns:
            LongTensor: length
        """
        out_lengths = in_seq_lens_tensor.clone()
        out_lengths = ((out_lengths.float() - (self.conv_layers_infos[i][1] - 1) - 1) / self.conv_layers_infos[i][-1] + 1).floor().long()
        out_nonmask = (~lengths_to_padding_mask(out_lengths)).float()
        return out_nonmask, out_lengths

    def get_out_seq_lens_tensor(self, in_seq_lens_tensor):
        out = in_seq_lens_tensor.clone()
        for i in range(len(self.conv_layers)):
            out = ((out.float() - (self.conv_layers_infos[i][1] - 1) - 1) / self.conv_layers_infos[i][-1] + 1).floor().long()
        return out