# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import OrderedDict

import torch
import torch.distributed
import torch.nn as nn

from shruti.nemo.collections.asr.parts.submodules.subsampling import ConvSubsampling, StackingSubsampling
from shruti.nemo.core.classes.common import typecheck
from shruti.nemo.core.classes.exportable import Exportable
from shruti.nemo.core.classes.module import NeuralModule
from shruti.nemo.core.neural_types import AcousticEncodedRepresentation, LengthsType, NeuralType, SpectrogramType

__all__ = ['RNNEncoder']


class RNNEncoder(NeuralModule, Exportable):
    """
    The RNN-based encoder for ASR models.
    Followed the architecture suggested in the following paper:
    'STREAMING END-TO-END SPEECH RECOGNITION FOR MOBILE DEVICES' by Yanzhang He et al.
    https://arxiv.org/pdf/1811.06621.pdf


    Args:
        feat_in (int): the size of feature channels
        n_layers (int): number of layers of RNN
        d_model (int): the hidden size of the model
        proj_size (int): the size of the output projection after each RNN layer
        rnn_type (str): the type of the RNN layers, choices=['lstm, 'gru', 'rnn']
        bidirectional (float): specifies whether RNN layers should be bidirectional or not
            Defaults to True.
        feat_out (int): the size of the output features
            Defaults to -1 (means feat_out is d_model)
        subsampling (str): the method of subsampling, choices=['stacking, 'vggnet', 'striding']
            Defaults to stacking.
        subsampling_factor (int): the subsampling factor
            Defaults to 4.
        subsampling_conv_channels (int): the size of the convolutions in the subsampling module for vggnet and striding
            Defaults to -1 which would set it to d_model.
        dropout (float): the dropout rate used between all layers
            Defaults to 0.2.
    """

    def input_example(self):
        """
        Generates input examples for tracing etc.
        Returns:
            A tuple of input examples.
        """
        input_example = torch.randn(16, self._feat_in, 256).to(next(self.parameters()).device)
        input_example_length = torch.randint(0, 256, (16,)).to(next(self.parameters()).device)
        return tuple([input_example, input_example_length])

    @property
    def input_types(self):
        """Returns definitions of module input ports.
        """
        return OrderedDict(
            {
                "audio_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()),
                "length": NeuralType(tuple('B'), LengthsType()),
            }
        )

    @property
    def output_types(self):
        """Returns definitions of module output ports.
        """
        return OrderedDict(
            {
                "outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
                "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
            }
        )

    def __init__(
        self,
        feat_in: int,
        n_layers: int,
        d_model: int,
        proj_size: int = -1,
        rnn_type: str = 'lstm',
        bidirectional: bool = True,
        subsampling: str = 'striding',
        subsampling_factor: int = 4,
        subsampling_conv_channels: int = -1,
        dropout: float = 0.2,
    ):
        super().__init__()

        self.d_model = d_model
        self._feat_in = feat_in

        if subsampling_conv_channels == -1:
            subsampling_conv_channels = proj_size
        if subsampling and subsampling_factor > 1:
            if subsampling in ['stacking', 'stacking_norm']:
                self.pre_encode = StackingSubsampling(
                    subsampling_factor=subsampling_factor,
                    feat_in=feat_in,
                    feat_out=proj_size,
                    norm=True if 'norm' in subsampling else False,
                )
            else:
                self.pre_encode = ConvSubsampling(
                    subsampling=subsampling,
                    subsampling_factor=subsampling_factor,
                    feat_in=feat_in,
                    feat_out=proj_size,
                    conv_channels=subsampling_conv_channels,
                    activation=nn.ReLU(),
                )
        else:
            self.pre_encode = nn.Linear(feat_in, proj_size)

        self._feat_out = proj_size

        self.layers = nn.ModuleList()

        SUPPORTED_RNN = {"lstm": nn.LSTM, "gru": nn.GRU, "rnn": nn.RNN}
        if rnn_type not in SUPPORTED_RNN:
            raise ValueError(f"rnn_type can be one from the following:{SUPPORTED_RNN.keys()}")
        else:
            rnn_module = SUPPORTED_RNN[rnn_type]

        for i in range(n_layers):
            rnn_proj_size = proj_size // 2 if bidirectional else proj_size
            if rnn_type == "lstm":
                layer = rnn_module(
                    input_size=self._feat_out,
                    hidden_size=d_model,
                    num_layers=1,
                    batch_first=True,
                    bidirectional=bidirectional,
                    proj_size=rnn_proj_size,
                )
            self.layers.append(layer)
            self.layers.append(nn.LayerNorm(proj_size))
            self.layers.append(nn.Dropout(p=dropout))
            self._feat_out = proj_size

    @typecheck()
    def forward(self, audio_signal, length=None):
        max_audio_length: int = audio_signal.size(-1)

        if length is None:
            length = audio_signal.new_full(
                audio_signal.size(0), max_audio_length, dtype=torch.int32, device=self.seq_range.device
            )

        audio_signal = torch.transpose(audio_signal, 1, 2)

        if isinstance(self.pre_encode, nn.Linear):
            audio_signal = self.pre_encode(audio_signal)
        else:
            audio_signal, length = self.pre_encode(audio_signal, length)

        for lth, layer in enumerate(self.layers):
            audio_signal = layer(audio_signal)
            if isinstance(audio_signal, tuple):
                audio_signal, _ = audio_signal

        audio_signal = torch.transpose(audio_signal, 1, 2)
        return audio_signal, length
