# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn.functional as F
import torch.nn as nn
from modules.vocoder_blocks import *

LRELU_SLOPE = 0.1


class ISTFT(nn.Module):
    """
    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
    See issue: https://github.com/pytorch/pytorch/issues/62323
    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
    The NOLA constraint is met as we trim padded samples anyway.

    Args:
        n_fft (int): Size of Fourier transform.
        hop_length (int): The distance between neighboring sliding window frames.
        win_length (int): The size of window frame and STFT filter.
        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
    """

    def __init__(
        self,
        n_fft: int,
        hop_length: int,
        win_length: int,
        padding: str = "same",
    ):
        super().__init__()
        if padding not in ["center", "same"]:
            raise ValueError("Padding must be 'center' or 'same'.")
        self.padding = padding
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length

    def forward(self, spec: torch.Tensor, window) -> torch.Tensor:
        """
        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.

        Args:
            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
                            N is the number of frequency bins, and T is the number of time frames.

        Returns:
            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
        """
        if self.padding == "center":
            # Fallback to pytorch native implementation
            return torch.istft(
                spec,
                self.n_fft,
                self.hop_length,
                self.win_length,
                window,
                center=True,
            )
        elif self.padding == "same":
            pad = (self.win_length - self.hop_length) // 2
        else:
            raise ValueError("Padding must be 'center' or 'same'.")

        assert spec.dim() == 3, "Expected a 3D tensor as input"
        B, N, T = spec.shape

        # Inverse FFT
        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
        ifft = ifft * window[None, :, None]

        # Overlap and Add
        output_size = (T - 1) * self.hop_length + self.win_length
        y = torch.nn.functional.fold(
            ifft,
            output_size=(1, output_size),
            kernel_size=(1, self.win_length),
            stride=(1, self.hop_length),
        )[:, 0, 0, pad:-pad]

        # Window envelope
        window_sq = window.square().expand(1, T, -1).transpose(1, 2)
        window_envelope = torch.nn.functional.fold(
            window_sq,
            output_size=(1, output_size),
            kernel_size=(1, self.win_length),
            stride=(1, self.hop_length),
        ).squeeze()[pad:-pad]

        # Normalize
        assert (window_envelope > 1e-11).all()
        y = y / window_envelope

        return y
    
# The ASP and PSP Module are adopted from APNet under the MIT License
# https://github.com/YangAi520/APNet/blob/main/models.py

class ASPResBlock(torch.nn.Module):
    def __init__(self, cfg, channels, kernel_size=3, dilation=(1, 3, 5)):
        super(ASPResBlock, self).__init__()
        self.cfg = cfg
        self.convs1 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[0],
                        padding=get_padding(kernel_size, dilation[0]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[1],
                        padding=get_padding(kernel_size, dilation[1]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[2],
                        padding=get_padding(kernel_size, dilation[2]),
                    )
                ),
            ]
        )
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
            ]
        )
        self.convs2.apply(init_weights)

    def forward(self, x):
        for c1, c2 in zip(self.convs1, self.convs2):
            xt = F.leaky_relu(x, LRELU_SLOPE)
            xt = c1(xt)
            xt = F.leaky_relu(xt, LRELU_SLOPE)
            xt = c2(xt)
            x = xt + x
        return x


class PSPResBlock(torch.nn.Module):
    def __init__(self, cfg, channels, kernel_size=3, dilation=(1, 3, 5)):
        super(PSPResBlock, self).__init__()
        self.cfg = cfg
        self.convs1 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[0],
                        padding=get_padding(kernel_size, dilation[0]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[1],
                        padding=get_padding(kernel_size, dilation[1]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[2],
                        padding=get_padding(kernel_size, dilation[2]),
                    )
                ),
            ]
        )
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
            ]
        )
        self.convs2.apply(init_weights)

    def forward(self, x):
        for c1, c2 in zip(self.convs1, self.convs2):
            xt = F.leaky_relu(x, LRELU_SLOPE)
            xt = c1(xt)
            xt = F.leaky_relu(xt, LRELU_SLOPE)
            xt = c2(xt)
            x = xt + x
        return x


class APNet(torch.nn.Module):
    def __init__(self, cfg):
        super(APNet, self).__init__()
        self.cfg = cfg
        self.ASP_num_kernels = len(cfg.model.apnet.ASP_resblock_kernel_sizes)
        self.PSP_num_kernels = len(cfg.model.apnet.PSP_resblock_kernel_sizes)

        self.ASP_input_conv = weight_norm(
            Conv1d(
                cfg.preprocess.n_mel,
                cfg.model.apnet.ASP_channel,
                cfg.model.apnet.ASP_input_conv_kernel_size,
                1,
                padding=get_padding(cfg.model.apnet.ASP_input_conv_kernel_size, 1),
            )
        )
        self.PSP_input_conv = weight_norm(
            Conv1d(
                cfg.preprocess.n_mel,
                cfg.model.apnet.PSP_channel,
                cfg.model.apnet.PSP_input_conv_kernel_size,
                1,
                padding=get_padding(cfg.model.apnet.PSP_input_conv_kernel_size, 1),
            )
        )

        self.ASP_ResNet = nn.ModuleList()
        for j, (k, d) in enumerate(
            zip(
                cfg.model.apnet.ASP_resblock_kernel_sizes,
                cfg.model.apnet.ASP_resblock_dilation_sizes,
            )
        ):
            self.ASP_ResNet.append(ASPResBlock(cfg, cfg.model.apnet.ASP_channel, k, d))

        self.PSP_ResNet = nn.ModuleList()
        for j, (k, d) in enumerate(
            zip(
                cfg.model.apnet.PSP_resblock_kernel_sizes,
                cfg.model.apnet.PSP_resblock_dilation_sizes,
            )
        ):
            self.PSP_ResNet.append(PSPResBlock(cfg, cfg.model.apnet.PSP_channel, k, d))

        self.ASP_output_conv = weight_norm(
            Conv1d(
                cfg.model.apnet.ASP_channel,
                cfg.preprocess.n_fft // 2 + 1,
                cfg.model.apnet.ASP_output_conv_kernel_size,
                1,
                padding=get_padding(cfg.model.apnet.ASP_output_conv_kernel_size, 1),
            )
        )
        self.PSP_output_R_conv = weight_norm(
            Conv1d(
                cfg.model.apnet.PSP_channel,
                cfg.preprocess.n_fft // 2 + 1,
                cfg.model.apnet.PSP_output_R_conv_kernel_size,
                1,
                padding=get_padding(cfg.model.apnet.PSP_output_R_conv_kernel_size, 1),
            )
        )
        self.PSP_output_I_conv = weight_norm(
            Conv1d(
                cfg.model.apnet.PSP_channel,
                cfg.preprocess.n_fft // 2 + 1,
                cfg.model.apnet.PSP_output_I_conv_kernel_size,
                1,
                padding=get_padding(cfg.model.apnet.PSP_output_I_conv_kernel_size, 1),
            )
        )

        self.iSTFT = ISTFT(
            self.cfg.preprocess.n_fft,
            hop_length=self.cfg.preprocess.hop_size,
            win_length=self.cfg.preprocess.win_size,
        )

        self.ASP_output_conv.apply(init_weights)
        self.PSP_output_R_conv.apply(init_weights)
        self.PSP_output_I_conv.apply(init_weights)

    def forward(self, mel):
        logamp = self.ASP_input_conv(mel)
        logamps = None
        for j in range(self.ASP_num_kernels):
            if logamps is None:
                logamps = self.ASP_ResNet[j](logamp)
            else:
                logamps += self.ASP_ResNet[j](logamp)
        logamp = logamps / self.ASP_num_kernels
        logamp = F.leaky_relu(logamp)
        logamp = self.ASP_output_conv(logamp)

        pha = self.PSP_input_conv(mel)
        phas = None
        for j in range(self.PSP_num_kernels):
            if phas is None:
                phas = self.PSP_ResNet[j](pha)
            else:
                phas += self.PSP_ResNet[j](pha)
        pha = phas / self.PSP_num_kernels
        pha = F.leaky_relu(pha)
        R = self.PSP_output_R_conv(pha)
        I = self.PSP_output_I_conv(pha)

        pha = torch.atan2(I, R)

        rea = torch.exp(logamp) * torch.cos(pha)
        imag = torch.exp(logamp) * torch.sin(pha)

        spec = torch.cat((rea.unsqueeze(-1), imag.unsqueeze(-1)), -1)

        spec = torch.view_as_complex(spec)

        audio = self.iSTFT.forward(
            spec, torch.hann_window(self.cfg.preprocess.win_size).to(mel.device)
        )

        return logamp, pha, rea, imag, audio.unsqueeze(1)