# -*- coding: utf-8 -*-

# Copyright 2020 Tomoki Hayashi
#  MIT License (https://opensource.org/licenses/MIT)

"""VQVAE Modules."""

import logging

import torch

import ParallelWaveGAN.parallel_wavegan.models
from ParallelWaveGAN.parallel_wavegan.layers import VQCodebook


class VQVAE(torch.nn.Module):
    """VQVAE module."""

    def __init__(
        self,
        in_channels=1,
        out_channels=1,
        num_embeds=512,
        embed_dim=256,
        num_local_embeds=None,
        local_embed_dim=None,
        num_global_embeds=None,
        global_embed_dim=None,
        encoder_type="MelGANDiscriminator",
        decoder_type="MelGANGenerator",
        encoder_conf={
            "out_channels": 256,
            "downsample_scales": [4, 4, 2, 2],
            "max_downsample_channels": 1024,
        },
        decoder_conf={
            "in_channels": 256,
            "upsample_scales": [4, 4, 2, 2],
            "channels": 512,
            "stacks": 3,
        },
        use_weight_norm=True,
    ):
        """Initialize VQVAE module.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            num_embeds (int): Number of embeddings.
            embed_dim (int): Dimension of each embedding.
            num_local_embeds (int): Number of local embeddings.
            local_embed_dim (int): Dimension of each local embedding.
            num_global_embeds (int): Number of global embeddings.
            global_embed_dim (int): Dimension of each global embedding.
            encoder_type (str): Encoder module name.
            decoder_type (str): Decoder module name.
            encoder_conf (dict): Hyperparameters for the encoder.
            decoder_conf (dict): Hyperparameters for the decoder.
            use_weight_norm (bool): Whether to use weight norm.

        """
        super(VQVAE, self).__init__()
        encoder_class = getattr(parallel_wavegan.models, encoder_type)
        decoder_class = getattr(parallel_wavegan.models, decoder_type)
        encoder_conf.update({"in_channels": in_channels})
        decoder_conf.update({"out_channels": out_channels})
        if not issubclass(decoder_class, parallel_wavegan.models.MelGANGenerator):
            raise NotImplementedError(f"{decoder_class} is not supported yet.")
        if num_local_embeds is not None:
            if local_embed_dim is not None:
                self.local_embed = torch.nn.Conv1d(num_local_embeds, local_embed_dim, 1)
            else:
                self.local_embed = None
        if num_global_embeds is not None:
            self.global_embed = torch.nn.Embedding(num_global_embeds, global_embed_dim)
        self.encoder = encoder_class(**encoder_conf)
        self.codebook = VQCodebook(num_embeds=num_embeds, embed_dim=embed_dim)
        self.decoder = decoder_class(**decoder_conf)

        # apply weight norm
        if use_weight_norm:
            self.remove_weight_norm()  # for duplicated weight norm
            self.apply_weight_norm()

    def forward(self, x, l=None, g=None):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input tensor (B, in_channels, T).
            l (Tensor): Local conditioning tensor (B, num_local_embeds, T).
            g (LongTensor): Global conditioning idx (B, ).

        Return:
            Tensor: Reconstruced input tensor (B, in_channels, T).
            Tensor: Encoder hidden states (B, embed_dim, T // prod(downsample_scales)).
            Tensor: Quantized encoder hidden states (B, embed_dim, T // prod(downsample_scales)).

        """
        z_e = self.encoder(x)
        z_e = z_e[-1] if isinstance(z_e, list) else z_e  # For MelGAN Discriminator
        z_q_st, z_q = self.codebook.straight_through(z_e)
        if l is not None:
            if self.local_embed is not None:
                l = self.local_embed(l)
            z_q_st = torch.cat([z_q_st, l], dim=1)
        if g is not None:
            g = self.global_embed(g).unsqueeze(2).expand(-1, -1, z_q_st.size(2))
            z_q_st = torch.cat([z_q_st, g], dim=1)
        x_bar = self.decoder(z_q_st)

        return x_bar, z_e, z_q

    def encode(self, x):
        """Encode the inputs into the latent codes.

        Args:
            x (Tensor): Input tensor (B, in_channels, T).

        Returns:
            LongTensor: Quantized tensor (B, T).

        """
        z_e = self.encoder(x)[-1]
        z_e = z_e[-1] if isinstance(z_e, list) else z_e  # For MelGAN Discriminator
        return self.codebook(z_e)

    def decode(self, indices, l=None, g=None):
        """Decode the latent codes to the inputs.

        Args:
            indices (LongTensor): Quantized tensor (B, T).
            l (Tensor): Local conditioning tensor (B, num_local_embeds, T).
            g (LongTensor): Global conditioning idx (B, ).

        Return:
            Tensor: Reconstruced tensor (B, 1, T).

        """
        z_q = self.codebook.embedding(indices).transpose(2, 1)
        if l is not None:
            if self.local_embed is not None:
                l = self.local_embed(l)
            z_q = torch.cat([z_q, l], dim=1)
        if g is not None:
            g = self.global_embed(g).unsqueeze(2).expand(-1, -1, z_q.size(2))
            z_q = torch.cat([z_q, g], dim=1)
        return self.decoder(z_q)

    def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv1d) or isinstance(
                m, torch.nn.ConvTranspose1d
            ):
                torch.nn.utils.weight_norm(m)
                logging.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)

    def remove_weight_norm(self):
        """Remove weight normalization module from all of the layers."""

        def _remove_weight_norm(m):
            try:
                logging.debug(f"Weight norm is removed from {m}.")
                torch.nn.utils.remove_weight_norm(m)
            except ValueError:  # this module didn't have weight norm
                return

        self.apply(_remove_weight_norm)