from typing import List
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import Union
import logging
import torch
import torch.nn as nn
from torch.nn import functional as F
import numpy as np
from funasr.models.transformer.utils.nets_utils import make_pad_mask
from funasr.models.transformer.layer_norm import LayerNorm
from funasr.models.encoder.abs_encoder import AbsEncoder
import math
from funasr.models.transformer.utils.repeat import repeat
from funasr.models.transformer.utils.multi_layer_conv import FsmnFeedForward


class FsmnBlock(torch.nn.Module):
    def __init__(
        self,
        n_feat,
        dropout_rate,
        kernel_size,
        fsmn_shift=0,
    ):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout_rate)
        self.fsmn_block = nn.Conv1d(
            n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False
        )
        # padding
        left_padding = (kernel_size - 1) // 2
        if fsmn_shift > 0:
            left_padding = left_padding + fsmn_shift
        right_padding = kernel_size - 1 - left_padding
        self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0)

    def forward(self, inputs, mask, mask_shfit_chunk=None):
        b, t, d = inputs.size()
        if mask is not None:
            mask = torch.reshape(mask, (b, -1, 1))
            if mask_shfit_chunk is not None:
                mask = mask * mask_shfit_chunk

        inputs = inputs * mask
        x = inputs.transpose(1, 2)
        x = self.pad_fn(x)
        x = self.fsmn_block(x)
        x = x.transpose(1, 2)
        x = x + inputs
        x = self.dropout(x)
        return x * mask


class EncoderLayer(torch.nn.Module):
    def __init__(self, in_size, size, feed_forward, fsmn_block, dropout_rate=0.0):
        super().__init__()
        self.in_size = in_size
        self.size = size
        self.ffn = feed_forward
        self.memory = fsmn_block
        self.dropout = nn.Dropout(dropout_rate)

    def forward(
        self, xs_pad: torch.Tensor, mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # xs_pad in Batch, Time, Dim

        context = self.ffn(xs_pad)[0]
        memory = self.memory(context, mask)

        memory = self.dropout(memory)
        if self.in_size == self.size:
            return memory + xs_pad, mask

        return memory, mask


class FsmnEncoder(AbsEncoder):
    """Encoder using Fsmn"""

    def __init__(
        self,
        in_units,
        filter_size,
        fsmn_num_layers,
        dnn_num_layers,
        num_memory_units=512,
        ffn_inner_dim=2048,
        dropout_rate=0.0,
        shift=0,
        position_encoder=None,
        sample_rate=1,
        out_units=None,
        tf2torch_tensor_name_prefix_torch="post_net",
        tf2torch_tensor_name_prefix_tf="EAND/post_net",
    ):
        """Initializes the parameters of the encoder.

        Args:
          filter_size: the total order of memory block
          fsmn_num_layers: The number of fsmn layers.
          dnn_num_layers: The number of dnn layers
          num_units: The number of memory units.
          ffn_inner_dim: The number of units of the inner linear transformation
            in the feed forward layer.
          dropout_rate: The probability to drop units from the outputs.
          shift: left padding, to control delay
          position_encoder: The :class:`opennmt.layers.position.PositionEncoder` to
            apply on inputs or ``None``.
        """
        super(FsmnEncoder, self).__init__()
        self.in_units = in_units
        self.filter_size = filter_size
        self.fsmn_num_layers = fsmn_num_layers
        self.dnn_num_layers = dnn_num_layers
        self.num_memory_units = num_memory_units
        self.ffn_inner_dim = ffn_inner_dim
        self.dropout_rate = dropout_rate
        self.shift = shift
        if not isinstance(shift, list):
            self.shift = [shift for _ in range(self.fsmn_num_layers)]
        self.sample_rate = sample_rate
        if not isinstance(sample_rate, list):
            self.sample_rate = [sample_rate for _ in range(self.fsmn_num_layers)]
        self.position_encoder = position_encoder
        self.dropout = nn.Dropout(dropout_rate)
        self.out_units = out_units
        self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf

        self.fsmn_layers = repeat(
            self.fsmn_num_layers,
            lambda lnum: EncoderLayer(
                in_units if lnum == 0 else num_memory_units,
                num_memory_units,
                FsmnFeedForward(
                    in_units if lnum == 0 else num_memory_units,
                    ffn_inner_dim,
                    num_memory_units,
                    1,
                    dropout_rate,
                ),
                FsmnBlock(num_memory_units, dropout_rate, filter_size, self.shift[lnum]),
            ),
        )

        self.dnn_layers = repeat(
            dnn_num_layers,
            lambda lnum: FsmnFeedForward(
                num_memory_units,
                ffn_inner_dim,
                num_memory_units,
                1,
                dropout_rate,
            ),
        )
        if out_units is not None:
            self.conv1d = nn.Conv1d(num_memory_units, out_units, 1, 1)

    def output_size(self) -> int:
        return self.num_memory_units

    def forward(
        self, xs_pad: torch.Tensor, ilens: torch.Tensor, prev_states: torch.Tensor = None
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        inputs = xs_pad
        if self.position_encoder is not None:
            inputs = self.position_encoder(inputs)

        inputs = self.dropout(inputs)
        masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
        inputs = self.fsmn_layers(inputs, masks)[0]
        inputs = self.dnn_layers(inputs)[0]

        if self.out_units is not None:
            inputs = self.conv1d(inputs.transpose(1, 2)).transpose(1, 2)

        return inputs, ilens, None