#!/usr/bin/env python3

import logging

import torch
import torch.nn as nn
from fairseq import checkpoint_utils, utils
from fairseq.data.data_utils import lengths_to_padding_mask
from fairseq.models import (
    register_model,
    register_model_architecture,
)
from fairseq.models.transformer import Embedding
from fairseq.models.speech_to_text import (
    S2TTransformerModel,
    S2TTransformerEncoder,
    CTC,
    CTCCompressStrategy,
)
from fairseq.modules import (
    LayerNorm
)

logger = logging.getLogger(__name__)


class Adapter(nn.Module):
    def __init__(self, args, dictionary, embed_tokens):
        super().__init__()

        embed_dim = args.encoder_embed_dim

        self.adapter_type = args.adapter
        if self.adapter_type in ["linear", "league", "gated_league", "gated_league2"]:
            self.linear_adapter = nn.Sequential(
                nn.Linear(embed_dim, embed_dim),
                LayerNorm(args.encoder_embed_dim),
                nn.ReLU(),
            )
        elif self.adapter_type == "linear2":
            self.linear_adapter = nn.Sequential(
                nn.Linear(embed_dim, embed_dim),
            )

        if self.adapter_type in ["embed", "context", "league", "gated_league", "gated_league2"]:
            if embed_tokens is None:
                num_embeddings = len(dictionary)
                self.embed_adapter = Embedding(num_embeddings, embed_dim, self.padding_idx)
            else:
                self.embed_adapter = embed_tokens

        if self.adapter_type == "gated_league":
            self.gate_linear = nn.Linear(2 * embed_dim, embed_dim)
        elif self.adapter_type == "gated_league2":
            self.gate_linear1 = nn.Linear(embed_dim, embed_dim)
            self.gate_linear2 = nn.Linear(embed_dim, embed_dim)

        if self.adapter_type == "shrink":
            self.ctc_compress_method = getattr(CTCCompressStrategy, args.ctc_compress_strategy)

    def forward(self, x, padding):

        representation, distribution = x
        batch, seq_len, embed_dim = representation.size()
        org_distribution = distribution
        if distribution is not None:
            distribution = distribution.view(-1, distribution.size(-1))
        lengths = (~padding).long().sum(-1)

        if self.adapter_type == "linear":
            out = self.linear_adapter(representation)

        elif self.adapter_type == "context":
            out = torch.mm(distribution, self.embed_adapter.weight).view(batch, seq_len, -1)

        elif self.adapter_type == "league":
            linear_out = self.linear_adapter(representation)
            soft_out = torch.mm(distribution, self.embed_adapter.weight).view(batch, seq_len, -1)
            out = linear_out + soft_out

        elif self.adapter_type == "gated_league":
            linear_out = self.linear_adapter(representation)
            soft_out = torch.mm(distribution, self.embed_adapter.weight).view(batch, seq_len, -1)
            coef = (self.gate_linear(torch.cat([linear_out, soft_out], dim=-1))).sigmoid()
            out = coef * linear_out + (1 - coef) * soft_out

        elif self.adapter_type == "none":
            out = representation

        elif self.adapter_type == "shrink":
            from itertools import groupby

            with torch.no_grad():
                batch_predicted = []
                prob_ctc = org_distribution.transpose(0, 1)  # T x B x D -> B x T x D
                for b in range(prob_ctc.shape[0]):
                    predicted = prob_ctc[b][: lengths[b]].argmax(-1).tolist()
                    batch_predicted.append([(p[0], len(list(p[1]))) for p in groupby(predicted)])

                new_lengths = [len(p) for p in batch_predicted]
                weights_matrix = self.ctc_compress_method(prob_ctc, batch_predicted, new_lengths,
                                                          representation.dtype, representation.device)

            # x is T x B x C -> B x C x T; weights_matrix is B x T x T'
            compressed_output = representation.permute(1, 2, 0).bmm(weights_matrix)  # B x C x T'
            out = compressed_output.permute(2, 0, 1)

            out_lengths = lengths.new(new_lengths)
            padding = lengths_to_padding_mask(out_lengths)

        else:
            out = None
            logging.error("Unsupported adapter type: {}.".format(self.adapter_type))

        return out, padding


@register_model("multi_ctc_s2t_transformer")
class MultiCTCS2TTransformerModel(S2TTransformerModel):
    """Speech-to-Text Transformer with multiple CTC Loss in different layers"""

    def __init__(self, encoder, decoder):
        super().__init__(encoder, decoder)

    @staticmethod
    def add_args(parser):
        S2TTransformerModel.add_args(parser)
        parser.add_argument(
            "--multi-ctc-layers",
            default=None,
            type=str,
            help="the position of the ctc loss, separated by ",
        )
        parser.add_argument(
            "--adapter",
            default="league",
            type=str,
            help="adapter type",
        )
        parser.add_argument(
            "--ctc-compress-strategy",
            default="avg",
            type=str,
            help="compress strategy, such as avg, weighted, and softmax",
        )
        pass

    @classmethod
    def build_encoder(cls, args, task=None, embed_tokens=None):
        encoder = S2TMultiCTCTransformerEncoder(args, task, embed_tokens)
        if getattr(args, "load_pretrained_encoder_from", None):
            encoder = checkpoint_utils.load_pretrained_component_from_model(
                component=encoder, checkpoint=args.load_pretrained_encoder_from, strict=False
            )
            logger.info(
                f"loaded pretrained encoder from: "
                f"{args.load_pretrained_encoder_from}"
            )
        return encoder


class S2TMultiCTCTransformerEncoder(S2TTransformerEncoder):
    """Speech-to-text Transformer encoder that consists of multiple  input subsampler and
    Conformer encoder."""

    def __init__(self, args, task=None, embed_tokens=None):
        super().__init__(args, task, embed_tokens)

        if self.use_ctc:
            del self.ctc

        self.multi_ctc_layers = []
        if args.multi_ctc_layers is not None:
            multi_ctc_layers = args.multi_ctc_layers.split(",")
            for layer_idx in multi_ctc_layers:
                layer_idx = int(layer_idx)
                if layer_idx <= 0:
                    layer_idx += args.encoder_layers
                self.multi_ctc_layers.append(layer_idx)

                inter_ctc = True if layer_idx != args.encoder_layers else False
                if inter_ctc:
                    logger.info("Intermedia CTC loss in layer %d" % layer_idx)
                ctc = CTC(args.encoder_embed_dim,
                          dictionary_size=len(task.source_dictionary),
                          dropout=args.dropout,
                          need_layernorm=inter_ctc)

                if task.source_dictionary == task.target_dictionary and embed_tokens is not None:
                    ctc.ctc_projection.weight = embed_tokens.weight

                setattr(self, f"ctc{layer_idx}", ctc)

                if inter_ctc:
                    adapter = Adapter(args, task.source_dictionary, ctc.ctc_projection)
                    setattr(self, f"adapter{layer_idx}", adapter)

    def forward(self, src_tokens, src_lengths):

        if self.history is not None:
            self.history.clean()

        # down-sampling
        x, input_lengths = self.subsample(src_tokens, src_lengths)

        if type(x) == list:
            inner_x = x
            x = inner_x[-1]

        # embedding scaling
        x = self.embed_scale * x

        # padding and position embedding
        encoder_padding_mask = lengths_to_padding_mask(input_lengths)
        positions = self.embed_positions(encoder_padding_mask).transpose(0, 1)
        if self.attn_type != "rel_selfattn":
            x += positions
        x = self.dropout_module(x)
        positions = self.dropout_module(positions)

        # add emb into history
        if self.history is not None:
            self.history.push(x)

        layer_idx = 0
        ctc_logit = []
        for layer in self.layers:
            layer_idx += 1

            if self.history is not None:
                x = self.history.pop()

            # encoder layer
            x = layer(x, encoder_padding_mask, pos_emb=positions)

            # interleave CTC
            if self.use_ctc and layer_idx in self.multi_ctc_layers and layer_idx != len(self.layers):
                ctc = getattr(self, f"ctc{layer_idx}")
                adapter = getattr(self, f"adapter{layer_idx}")
                logit = ctc(x)
                prob = ctc.softmax(logit)
                x, encoder_padding_mask = adapter([x, prob], encoder_padding_mask)

                ctc_logit.append(ctc(x))

            if layer_idx != len(self.layers) \
                    and self.interleaved_dropout is not None \
                    and layer_idx % self.interleaved_dropout == 0:
                x = self.dropout_module(x)

            if self.history is not None:
                self.history.push(x)

        if self.history is not None:
            x = self.history.pop()

        if self.layer_norm is not None:
            x = self.layer_norm(x)

        if self.use_ctc and len(self.layers) in self.multi_ctc_layers:
            ctc = getattr(self, f"ctc{len(self.layers)}")
            ctc_logit.append(ctc(x))

        return {
            "encoder_out": [x],  # T x B x C
            "ctc_logit": ctc_logit,  # B x T x C
            "encoder_padding_mask": [encoder_padding_mask],  # B x T
            "encoder_embedding": [],  # B x T x C
            "encoder_states": [],  # List[T x B x C]
            "src_tokens": [],
            "src_lengths": [],
        }


@register_model_architecture(model_name="multi_ctc_s2t_transformer", arch_name="multi_ctc_s2t_transformer")
def base_architecture(args):
    # Convolutional subsampler
    args.conv_kernel_sizes = getattr(args, "conv_kernel_sizes", "5,5")
    args.conv_channels = getattr(args, "conv_channels", 1024)

    # Transformer
    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 2048)
    args.encoder_layers = getattr(args, "encoder_layers", 12)
    args.encoder_attention_type = getattr(args, "encoder_attention_type", "selfattn")
    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 8)
    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", True)
    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", args.encoder_embed_dim)
    args.decoder_ffn_embed_dim = getattr(
        args, "decoder_ffn_embed_dim", args.encoder_ffn_embed_dim
    )
    args.decoder_layers = getattr(args, "decoder_layers", 6)
    args.decoder_attention_type = getattr(args, "decoder_attention_type", "selfattn")
    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", True)
    args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
    args.dropout = getattr(args, "dropout", 0.1)
    args.attention_dropout = getattr(args, "attention_dropout", args.dropout)
    args.activation_dropout = getattr(args, "activation_dropout", args.dropout)
    args.activation_fn = getattr(args, "activation_fn", "relu")
    args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
    args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
    args.tie_adaptive_weights = getattr(args, "tie_adaptive_weights", False)
    args.tie_adaptive_proj = getattr(args, "tie_adaptive_proj", False)
    args.adaptive_softmax_factor = getattr(args, "adaptive_softmax_factor", 4)

    args.share_decoder_input_output_embed = getattr(
        args, "share_decoder_input_output_embed", False
    )
    args.share_all_embeddings = getattr(args, "share_all_embeddings", False)
    args.no_token_positional_embeddings = getattr(
        args, "no_token_positional_embeddings", False
    )
    args.adaptive_input = getattr(args, "adaptive_input", False)
    args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.0)
    args.decoder_output_dim = getattr(
        args, "decoder_output_dim", args.decoder_embed_dim
    )
    args.decoder_input_dim = getattr(args, "decoder_input_dim", args.decoder_embed_dim)
    args.no_scale_embedding = getattr(args, "no_scale_embedding", False)
    args.quant_noise_pq = getattr(args, "quant_noise_pq", 0)

    # CTC
    args.multi_ctc_layers = getattr(args, "multi_ctc_layers", 0)

    # Conformer
    args.macaron_style = getattr(args, "macaron_style", False)
    args.use_cnn_module = getattr(args, "use_cnn_module", False)
    args.cnn_module_kernel = getattr(args, "cnn_module_kernel", 31)

    # settings for DLCL
    args.use_enc_dlcl = getattr(args, "use_enc_dlcl", False)
    args.use_dec_dlcl = getattr(args, "use_dec_dlcl", False)
    args.init_value = getattr(args, 'init_value', 'avg')
    args.weight_type = getattr(args, 'weight_type', 'scalar')
    args.encoder_learnable = getattr(args, 'encoder_learnable', True)
    args.decoder_learnable = getattr(args, 'decoder_learnable', True)
    args.normalize_embed = getattr(args, 'normalize_embed', False)
    args.history_dropout = getattr(args, 'history_dropout', 0.0)
    args.history_window_size = getattr(args, 'history_window_size', -1)

    # Relative position encoding
    args.max_encoder_relative_length = getattr(args, 'max_encoder_relative_length', -1)
    args.max_decoder_relative_length = getattr(args, 'max_decoder_relative_length', -1)
    args.k_only = getattr(args, 'k_only', True)

    # local modeling
    args.hard_mask_window = getattr(args, 'hard_mask_window', 0)
    args.gauss_mask_sigma = getattr(args, 'gauss_mask_sigma', 0)
    args.init_mask_weight = getattr(args, 'init_mask_weight', 0)

    # interleaved dropout
    args.interleave_dropout = getattr(args, "interleave_dropout", None)
    args.cl_dropout = getattr(args, "cl_dropout", False)
    args.cl_dropout_epoch = getattr(args, "cl_dropout_epoch", None)
    args.cl_dropout_strategy = getattr(args, "cl_dropout_strategy", "linear")


@register_model_architecture("multi_ctc_s2t_transformer", "multi_ctc_s2t_transformer_s")
def multi_ctc_s2t_transformer_s(args):
    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 256)
    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 256 * 8)
    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 4)
    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 4)
    args.dropout = getattr(args, "dropout", 0.1)

    args.multi_ctc_layers = getattr(args, "multi_ctc_layers", None)
    base_architecture(args)
