import torch
import torch.nn as nn
import torch.nn.functional as F
from fireredtts2.codec.whisper import WhisperEncoderLayer
from fireredtts2.codec.utils import make_nonpad_mask, make_block_causal_mask


class ResnetBlock(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int = None,
        conv_shortcut: bool = False,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels
        self.use_conv_shortcut = conv_shortcut

        self.block1 = nn.Sequential(
            nn.GroupNorm(
                num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
            ),
            nn.SiLU(),
            nn.Conv1d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
        )

        self.block2 = nn.Sequential(
            nn.GroupNorm(
                num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
            ),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Conv1d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
        )

        if self.in_channels != self.out_channels:
            if self.use_conv_shortcut:
                self.conv_shortcut = torch.nn.Conv1d(
                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
                )
            else:
                self.nin_shortcut = torch.nn.Conv1d(
                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
                )

    def forward(self, x: torch.Tensor):
        """
        Args:
            x: shape (b, c, t)
        """
        h = x
        h = self.block1(h)
        h = self.block2(h)

        if self.in_channels != self.out_channels:
            if self.use_conv_shortcut:
                x = self.conv_shortcut(x)
            else:
                x = self.nin_shortcut(x)
        return x + h


class Transpose(torch.nn.Module):
    def __init__(self, dim0: int, dim1: int):
        super().__init__()
        self.dim0 = dim0
        self.dim1 = dim1

    def forward(self, x: torch.Tensor):
        x = torch.transpose(x, self.dim0, self.dim1)
        return x


# A causal variant of Conv1d
class CausalConv1d(torch.nn.Conv1d):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
    ) -> None:
        super(CausalConv1d, self).__init__(in_channels, out_channels, kernel_size)
        self.causal_padding = (kernel_size - 1, 0)

    def forward(self, x: torch.Tensor):
        x = F.pad(x, self.causal_padding)
        x = super(CausalConv1d, self).forward(x)
        return x

    def forward_chunk(self, x: torch.Tensor, cnn_cache: torch.Tensor = None):
        if cnn_cache is None:
            cnn_cache = x.new_zeros(
                (x.shape[0], self.in_channels, self.causal_padding[0])
            )
        x = torch.cat([cnn_cache, x], dim=2)
        new_cnn_cache = x[..., -self.causal_padding[0] :]
        x = super(CausalConv1d, self).forward(x)
        return x, new_cnn_cache


# A causal variant of ResnetBlock
class CausalResnetBlock(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int = None,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels

        self.block1 = nn.Sequential(
            Transpose(1, 2),
            nn.LayerNorm(in_channels),
            Transpose(1, 2),
            nn.SiLU(),
            CausalConv1d(in_channels, out_channels, kernel_size=3),
        )

        self.block2 = nn.Sequential(
            Transpose(1, 2),
            nn.LayerNorm(out_channels),
            Transpose(1, 2),
            nn.SiLU(),
            nn.Dropout(dropout),
            CausalConv1d(out_channels, out_channels, kernel_size=3),
        )
        if self.in_channels != self.out_channels:
            self.nin_shortcut = torch.nn.Conv1d(
                in_channels, out_channels, kernel_size=1, stride=1, padding=0
            )

    def forward(self, x: torch.Tensor):
        """
        Args:
            x: shape (b, c, t)
        """
        h = x
        h = self.block1(h)
        h = self.block2(h)
        if self.in_channels != self.out_channels:
            x = self.nin_shortcut(x)
        return x + h

    def forward_chunk(self, x: torch.Tensor, cache: torch.Tensor = None):
        """
        Args:
            x: shape (b, c, t)
            cache: shape (b, c_in+c_out, t=2)
        """
        cache1, cache2 = (
            (None, None)
            if cache is None
            else cache.split((self.in_channels, self.out_channels), dim=1)
        )
        h = x
        # block1
        h = self.block1[:4](h)
        h, new_cache1 = self.block1[4].forward_chunk(h, cache1)
        # block2
        h = self.block2[:5](h)
        h, new_cache2 = self.block2[5].forward_chunk(h, cache2)
        if self.in_channels != self.out_channels:
            x = self.nin_shortcut(x)
        new_cache = torch.cat([new_cache1, new_cache2], dim=1)
        return x + h, new_cache


# Nonstreaming Vocos backbone based on Transformer layers
class VocosBackbone(nn.Module):
    def __init__(
        self,
        embed_dim: int = 1024,
        num_layers: int = 12,
        num_heads: int = 16,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.in_proj = nn.Conv1d(embed_dim, embed_dim, kernel_size=7, padding=3)
        self.prior_net = nn.Sequential(
            ResnetBlock(embed_dim, embed_dim, dropout=dropout),
            ResnetBlock(embed_dim, embed_dim, dropout=dropout),
        )
        self.transformers = nn.ModuleList(
            [WhisperEncoderLayer(embed_dim, num_heads) for _ in range(num_layers)]
        )
        self.post_net = nn.Sequential(
            ResnetBlock(embed_dim, embed_dim, dropout=dropout),
            ResnetBlock(embed_dim, embed_dim, dropout=dropout),
        )
        self.final_norm = nn.LayerNorm(embed_dim, eps=1e-6)

    def forward(
        self,
        x: torch.Tensor,
        x_lens: torch.Tensor,
    ):
        """
        Args:
            x: shape (b, t, c)
            x_lens: shape (b,)
        """
        x = x.transpose(1, 2)
        x = self.in_proj(x)
        x = self.prior_net(x)
        x = x.transpose(1, 2)

        attention_mask = make_nonpad_mask(x_lens).unsqueeze(1)  # (b, 1, t)
        # NOTE(sfy): I think positional embedding is unnecessary
        for layer in self.transformers:
            x = layer(x, attention_mask)
        x = x.transpose(1, 2)
        x = self.post_net(x)
        x = x.transpose(1, 2)
        x = self.final_norm(x)
        return x


# Streaming Vocos backbone based on Transformer layers
class CausalVocosBackbone(nn.Module):
    def __init__(
        self,
        embed_dim: int = 1024,
        num_layers: int = 12,
        num_heads: int = 16,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.in_proj = CausalConv1d(embed_dim, embed_dim, kernel_size=7)
        self.prior_net = nn.Sequential(
            CausalResnetBlock(embed_dim, embed_dim, dropout=dropout),
            CausalResnetBlock(embed_dim, embed_dim, dropout=dropout),
        )
        self.transformers = nn.ModuleList(
            [WhisperEncoderLayer(embed_dim, num_heads) for _ in range(num_layers)]
        )
        self.post_net = nn.Sequential(
            CausalResnetBlock(embed_dim, embed_dim, dropout=dropout),
            CausalResnetBlock(embed_dim, embed_dim, dropout=dropout),
        )
        self.final_norm = nn.LayerNorm(embed_dim, eps=1e-6)

    def forward(
        self,
        x: torch.Tensor,
        x_lens: torch.Tensor,
    ):
        """
        Args:
            x: shape (b, t, c)
            x_lens: shape (b,)
        """
        x = x.transpose(1, 2)
        x = self.in_proj(x)
        x = self.prior_net(x)
        x = x.transpose(1, 2)

        # NOTE(sfy): We have no padding in training, so safe for sdpa attention, no Nan.
        # Also, 1 token(12.5Hz) -> 4 latents(50Hz) -> 8 latents(100Hz),
        # so we design a 8 block causal attention mask instead of fully causal to improve performance
        attention_mask = make_block_causal_mask(x_lens, chunk_size=8)
        for layer in self.transformers:
            x = layer(x, attention_mask)

        x = x.transpose(1, 2)
        x = self.post_net(x)
        x = x.transpose(1, 2)
        x = self.final_norm(x)
        return x

    def forward_chunk(
        self,
        x: torch.Tensor,
        conv_cache1: torch.Tensor = None,
        conv_cache2: torch.Tensor = None,
        kv_cache: torch.Tensor = None,
    ):
        # Unpack cache
        cache1 = conv_cache1
        cache2, cache3, cache4, cache5 = (
            (None, None, None, None)
            if conv_cache2 is None
            else conv_cache2.chunk(4, dim=1)
        )

        # cache1: shape (b, c=embed_dim, t=6)
        x = x.transpose(1, 2)
        x, new_cache1 = self.in_proj.forward_chunk(x, cache1)
        # cache2: shape (b, c=embed_dim*2, t=2)
        x, new_cache2 = self.prior_net[0].forward_chunk(x, cache2)
        # cache3: shape (b, c=embed_dim*2, t=2)
        x, new_cache3 = self.prior_net[1].forward_chunk(x, cache3)
        x = x.transpose(1, 2)

        # k,v-cache: shape (b, nlayer, nh, t, c*2)
        new_kv_cache = []
        for idx, layer in enumerate(self.transformers):
            kv_cache_i = None if kv_cache is None else kv_cache[:, idx]
            x, new_kv_cache_i = layer.forward_chunk(x, kv_cache=kv_cache_i)
            new_kv_cache.append(new_kv_cache_i)
        new_kv_cache = torch.stack(new_kv_cache, dim=1)

        x = x.transpose(1, 2)
        # cache4: shape (b, c=embed_dim*2, t=2)
        x, new_cache4 = self.post_net[0].forward_chunk(x, cache4)
        # cache5: shape (b, c=embed_dim*2, t=2)
        x, new_cache5 = self.post_net[1].forward_chunk(x, cache5)
        x = x.transpose(1, 2)
        x = self.final_norm(x)

        new_conv_cache1 = new_cache1
        new_conv_cache2 = torch.cat(
            [new_cache2, new_cache3, new_cache4, new_cache5], dim=1
        )
        return x, new_conv_cache1, new_conv_cache2, new_kv_cache


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__()
        assert padding in ["center", "same"], "Padding must be 'center' or 'same'."
        self.padding = padding
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        window = torch.hann_window(win_length)
        self.register_buffer("window", window)

    def forward(self, spec: torch.Tensor) -> 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,
                self.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 * self.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 = self.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

    def forward_chunk(
        self, spec: torch.Tensor, cache: torch.Tensor = None, last_chunk: bool = False
    ):
        """Forward only one frame.

        Args:
            spec: shape (B, N, T=chunk_size)
            cache: previous chunk's last ifft frame, shape (B, N, T=3)
            last_chunk: if last_chunk, will not trim the last (win-hop) segment
        Returns:
            y: shape (B, T=effective_length)
        """
        assert self.padding == "same", "Padding must be same."
        assert (
            self.win_length % self.hop_length == 0
        ), f"{self.win_length} {self.hop_length}"
        pad = (self.win_length - self.hop_length) // 2

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

        # Append previous cache
        if cache is not None:
            ifft = torch.cat([cache, ifft], dim=-1)
        new_cache_t = self.win_length // self.hop_length - 1
        new_cache = ifft[..., -new_cache_t:]

        # Overlap and Add
        output_size = (ifft.shape[-1] - 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, :]

        # Window envelope
        window_sq = (
            self.window.square().expand(1, ifft.shape[-1], -1).transpose(1, 2)
        )  # (B=1, N, T)
        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()

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

        # Only take effective part
        if cache is None:
            y = y[:, pad:]
        else:
            y = y[:, (self.win_length - self.hop_length) :]
        if last_chunk:
            y = y[:, :-pad]
        else:
            y = y[:, : -(self.win_length - self.hop_length)]
        return y, new_cache


class ISTFTHead(nn.Module):
    """
    ISTFT Head module for predicting STFT complex coefficients.

    Args:
        dim (int): Hidden dimension of the model.
        n_fft (int): Size of Fourier transform.
        hop_length (int): The distance between neighboring sliding window frames, which should align with
                          the resolution of the input features.
        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
    """

    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
        super().__init__()
        self.hop_length = hop_length
        out_dim = n_fft + 2
        self.out = torch.nn.Linear(dim, out_dim)
        self.istft = ISTFT(
            n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding
        )

    def forward(self, x: torch.Tensor, x_len: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of the ISTFTHead module.

        Args:
            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
                        L is the sequence length, and H denotes the model dimension.

        Returns:
            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
        """
        x_pred = self.out(x)
        x_pred = x_pred.transpose(1, 2)
        mag, p = x_pred.chunk(2, dim=1)
        mag = torch.exp(mag)
        mag = torch.clip(
            mag, max=1e2
        )  # safeguard to prevent excessively large magnitudes
        # wrapping happens here. These two lines produce real and imaginary value
        x = torch.cos(p)
        y = torch.sin(p)
        # recalculating phase here does not produce anything new
        # only costs time
        # phase = torch.atan2(y, x)
        # S = mag * torch.exp(phase * 1j)
        # better directly produce the complex value
        S = mag * (x + 1j * y)
        audio = self.istft(S)
        audio_length = x_len * self.hop_length
        return audio, audio_length

    def forward_chunk(
        self, x: torch.Tensor, cache: torch.Tensor = None, last_chunk: bool = False
    ):
        """ISTFTHead can be adapted in streaming inference without retraining.

        Args:
            x: shape (B, T, C)
            cache: shape (B, N, T=3), istft cache
        Returns:
            audio: shape (B, t)
        """
        x_pred = self.out(x)
        x_pred = x_pred.transpose(1, 2)
        mag, p = x_pred.chunk(2, dim=1)
        mag = torch.exp(mag)  # (B, C, T)
        mag = torch.clip(
            mag, max=1e2
        )  # safeguard to prevent excessively large magnitudes
        # wrapping happens here. These two lines produce real and imaginary value
        x = torch.cos(p)
        y = torch.sin(p)
        S = mag * (x + 1j * y)  # (B, C, T)
        audio, new_cache = self.istft.forward_chunk(S, cache, last_chunk)
        return audio, new_cache


# UpsampleConv(50->100Hz) + VocosBackbone + ISTFTHead
class AcousticDecoder(nn.Module):
    def __init__(
        self,
        # Transformer
        embed_dim: int,
        num_layers: int,
        num_heads: int,
        dropout: float = 0.0,
        # iSTFT
        hop_length: int = 240,
        # Causal
        causal: bool = False,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.hop_length = hop_length
        self.causal = causal

        # Output upsample
        self.upsample_conv = nn.Sequential(
            nn.ConvTranspose1d(
                embed_dim,
                embed_dim,
                kernel_size=3,
                stride=2,
                padding=0,  # Do not fill input side
                output_padding=0,  # Can be adjusted to precisely control length
            ),
            nn.GELU(),
            nn.ConvTranspose1d(
                embed_dim,
                embed_dim,
                kernel_size=3,
                stride=1,
                padding=0,  # Do not fill input side
            ),
            nn.GELU(),
        )
        self.backbone = (
            CausalVocosBackbone(embed_dim, num_layers, num_heads, dropout)
            if causal
            else VocosBackbone(embed_dim, num_layers, num_heads, dropout)
        )
        self.isift = ISTFTHead(embed_dim, hop_length * 4, hop_length, padding="same")
        # Init weights
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Conv1d):
            nn.init.trunc_normal_(m.weight, std=0.02)
            nn.init.constant_(m.bias, 0)

    def forward(self, x: torch.Tensor, x_lens: torch.Tensor):
        """
        Args:
            x: shape (b, t, c)
            x_lens: shape (b,)
        """
        # Upsample
        target_length = x.shape[1] * 2
        x = x.transpose(1, 2)
        x = self.upsample_conv(x)
        x = x.transpose(1, 2)
        # NOTE strict upsampling, trim the last 3 elements
        x = x[:, :target_length]
        x_lens = x_lens * 2
        # Backbone
        x = self.backbone(x, x_lens)
        # iSTFT
        y, y_lens = self.isift(x, x_lens)
        return y, y_lens

    def forward_upsample_conv_chunk(self, x: torch.Tensor, cache: torch.Tensor = None):
        """Stream forward upsample_conv module with previous block cache.

        Args:
            x: shape (B, C, T)
            cache: shape (B, C, 3), where 3 denotes 1 history state for 1st conv and 2 for the rest conv.
        """
        # Unpack cache
        cache1, cache2 = (
            (None, None) if cache is None else torch.split(cache, [1, 2], dim=2)
        )
        # 1st conv cache
        if cache1 is not None:
            x = torch.cat([cache1, x], dim=2)
        new_cache1 = x[..., -1:]
        # 1st conv
        x = self.upsample_conv[0](x)[..., :-1]  # remove extra 1 frame
        if cache1 is not None:
            x = x[..., 2:]  # remove cache1 part
        x = self.upsample_conv[1](x)
        # 2nd conv cache
        if cache2 is not None:
            x = torch.cat([cache2, x], dim=2)
        new_cache2 = x[..., -2:]
        # 2nd conv
        x = self.upsample_conv[2](x)[..., :-2]  # remove extra 2 frame
        if cache2 is not None:
            x = x[..., 2:]  # remove cache2 part
        x = self.upsample_conv[3](x)

        new_cache = torch.cat([new_cache1, new_cache2], dim=2)
        return x, new_cache

    def forward_chunk(
        self,
        x: torch.Tensor,
        # Upsample conv cache
        up_conv_cache: torch.Tensor = None,
        # Backbone conv cache
        bb_conv_cache1: torch.Tensor = None,
        bb_conv_cache2: torch.Tensor = None,
        # Backbone attention cache
        bb_kv_cache: torch.Tensor = None,
        # iSTFT cache
        is_cache: torch.Tensor = None,
        last_chunk: bool = False,
    ):
        """
        Args:
            x: input sequence at 50Hz, length should be multiples of 4
        """
        assert (
            self.causal
        ), "Only AcousticDecoder with causal=True supports forward_chunk method."

        x = x.transpose(1, 2)
        x, new_up_conv_cache = self.forward_upsample_conv_chunk(x, up_conv_cache)
        x = x.transpose(1, 2)
        # Backbone
        x, new_bb_conv_cache1, new_bb_conv_cache2, new_bb_kv_cache = (
            self.backbone.forward_chunk(
                x,
                bb_conv_cache1,
                bb_conv_cache2,
                bb_kv_cache,
            )
        )
        # iSTFT
        y, new_is_cache = self.isift.forward_chunk(x, is_cache, last_chunk)
        return (
            y,
            new_up_conv_cache,
            new_bb_conv_cache1,
            new_bb_conv_cache2,
            new_bb_kv_cache,
            new_is_cache,
        )