import math
from typing import List

import numpy as np
import torch
import funcy

from torch import nn

try:
    from torch.nn.utils.parametrizations import weight_norm
except ImportError:
    from torch.nn.utils import weight_norm

from suno_utils.utils.s3 import read_from_s3


@torch.jit.script
def snake(x, alpha):
    shape = x.shape
    x = x.reshape(shape[0], shape[1], -1)
    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
    x = x.reshape(shape)
    return x


class Snake1d(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.alpha = nn.Parameter(torch.ones(1, channels, 1))

    def forward(self, x):
        return snake(x, self.alpha)


def WNConv1d(*args, **kwargs):
    return weight_norm(nn.Conv1d(*args, **kwargs))


def WNConvTranspose1d(*args, **kwargs):
    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))


def vae_sample(mean, scale):
    stdev = nn.functional.softplus(scale) + 1e-4
    var = stdev * stdev
    logvar = torch.log(var)
    latents = torch.randn_like(mean) * stdev + mean

    kl = 0.5 * (mean.pow(2) + var - logvar - 1).sum(1).mean()

    return latents, kl


class VAEBottleneck(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        # batch, dim, time
        mean, scale = x.chunk(2, dim=1)
        x, kl = vae_sample(mean, scale)

        return {
            "z": x,
            "kl": kl,
            "mean": mean,
            "scale": scale,
        }

    def encode(self, x, return_info=False, **kwargs):
        if return_info:
            out = self.forward(x)
            latents = out["z"]
            bottleneck_info = {"kl": float(out["kl"].item())}
            return latents, bottleneck_info
        return self.forward(x)

    def decode(self, x, **kwargs):
        return x


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


class ResidualUnit(nn.Module):
    def __init__(self, dim: int = 16, dilation: int = 1):
        super().__init__()
        pad = ((7 - 1) * dilation) // 2
        self.block = nn.Sequential(
            Snake1d(dim),
            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
            Snake1d(dim),
            WNConv1d(dim, dim, kernel_size=1),
        )

    def forward(self, x):
        y = self.block(x)
        pad = (x.shape[-1] - y.shape[-1]) // 2
        if pad > 0:
            x = x[..., pad:-pad]
        return x + y


class EncoderBlock(nn.Module):
    def __init__(self, dim: int = 16, stride: int = 1):
        super().__init__()
        self.block = nn.Sequential(
            ResidualUnit(dim // 2, dilation=1),
            ResidualUnit(dim // 2, dilation=3),
            ResidualUnit(dim // 2, dilation=9),
            Snake1d(dim // 2),
            WNConv1d(
                dim // 2,
                dim,
                kernel_size=2 * stride,
                stride=stride,
                padding=math.ceil(stride / 2),
            ),
        )

    def forward(self, x):
        return self.block(x)


class Encoder(nn.Module):
    def __init__(
        self,
        d_model: int = 64,
        strides: list = [2, 4, 8, 8],
        d_latent: int = 64,
    ):
        super().__init__()
        # Create first convolution
        self.block = [WNConv1d(2, d_model, kernel_size=7, padding=3)]

        # Create EncoderBlocks that double channels as they downsample by `stride`
        for stride in strides:
            d_model *= 2
            self.block += [EncoderBlock(d_model, stride=stride)]

        # Create last convolution and groupnorm
        self.block += [
            # nn.GroupNorm(4, d_model, affine=False),
            Snake1d(d_model),
            WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
        ]

        # Wrap black into nn.Sequential
        self.block = nn.Sequential(*self.block)
        self.enc_dim = d_model

    def forward(self, x):
        return self.block(x)


class DecoderBlock(nn.Module):
    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1):
        super().__init__()
        self.block = nn.Sequential(
            Snake1d(input_dim),
            WNConvTranspose1d(
                input_dim,
                output_dim,
                kernel_size=2 * stride,
                stride=stride,
                padding=math.floor(stride / 2),
            ),
            ResidualUnit(output_dim, dilation=1),
            ResidualUnit(output_dim, dilation=3),
            ResidualUnit(output_dim, dilation=9),
        )

    def forward(self, x):
        return self.block(x)


class Decoder(nn.Module):
    def __init__(
        self,
        input_channel,
        channels,
        rates,
        d_out: int = 2,
    ):
        super().__init__()

        # Add first conv layer
        layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)]

        # Add upsampling + MRF blocks
        for i, stride in enumerate(rates):
            input_dim = channels // 2**i
            output_dim = channels // 2 ** (i + 1)
            layers += [DecoderBlock(input_dim, output_dim, stride)]

        # Add final conv layer
        layers += [
            Snake1d(output_dim),
            WNConv1d(output_dim, d_out, kernel_size=7, padding=3),
            nn.Tanh(),
        ]

        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)


class DACVAE(nn.Module):
    def __init__(
        self,
        encoder_dim: int = 128,
        encoder_rates: List[int] = [2, 3, 5, 8, 8],
        vae_dim: int = 128,
        decoder_dim: int = 2048,
        decoder_rates: List[int] = [8, 8, 5, 3, 2],
        sample_rate: int = 48000,
        is_frozen_encoder: bool = False,
        **kwargs,
    ):
        super().__init__()

        self.encoder_dim = encoder_dim
        self.encoder_rates = encoder_rates
        self.decoder_dim = decoder_dim
        self.decoder_rates = decoder_rates
        self.sample_rate = sample_rate

        self.vae_dim = vae_dim

        self.hop_length = np.prod(encoder_rates)
        self.encoder = Encoder(encoder_dim, encoder_rates, vae_dim * 2)

        self.quantizer = VAEBottleneck()

        self.decoder = Decoder(
            vae_dim,
            decoder_dim,
            decoder_rates,
        )
        self.sample_rate = sample_rate
        self.is_frozen_encoder = is_frozen_encoder
        self.apply(init_weights)

        if is_frozen_encoder:
            # preload the weights
            load_f = funcy.partial(torch.load, mmap=True, weights_only=True)
            sd = read_from_s3("s3://suno-data/minz/models/dac_vae_tuned_25hz.pth", read_f=load_f)
            state_dict = sd["state_dict"]
            
            # load all weights
            self.load_state_dict(sd["state_dict"])

            # preload encoder-only
            # encoder_quantizer_dict = {k: v for k, v in state_dict.items() 
            #                           if k.startswith("encoder.") or k.startswith("quantizer.")}
            # self.load_state_dict(encoder_quantizer_dict, strict=False)

            # freeze encoder and quantizer
            for param in self.encoder.parameters():
                param.requires_grad = False
            for param in self.quantizer.parameters():
                param.requires_grad = False

    @property
    def device(self):
        return next(self.parameters()).device

    def get_model_hyperparameters(self):
        return {
            "kwargs": {
                "encoder_dim": self.encoder_dim,
                "encoder_rates": self.encoder_rates,
                "vae_dim": self.vae_dim,
                "decoder_dim": self.decoder_dim,
                "decoder_rates": self.decoder_rates,
                "sample_rate": self.sample_rate,
            }
        }

    def preprocess(self, audio_data, sample_rate):
        if sample_rate is None:
            sample_rate = self.sample_rate
        assert sample_rate == self.sample_rate

        length = audio_data.shape[-1]
        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
        audio_data = nn.functional.pad(audio_data, (0, right_pad))

        return audio_data

    def encode(
        self,
        audio_data: torch.Tensor,
        n_quantizers: int = None,
    ):
        """Encode given audio data and return quantized latent codes

        Parameters
        ----------
        audio_data : Tensor[B x 1 x T]
            Audio data to encode
        n_quantizers : int, optional
            Number of quantizers to use, by default None
            If None, all quantizers are used.

        Returns
        -------
        dict
            A dictionary with the following keys:
            "z" : Tensor[B x D x T]
                Quantized continuous representation of input
            "codes" : Tensor[B x N x T]
                Codebook indices for each codebook
                (quantized discrete representation of input)
            "latents" : Tensor[B x N*D x T]
                Projected latents (continuous representation of input before quantization)
            "vq/commitment_loss" : Tensor[1]
                Commitment loss to train encoder to predict vectors closer to codebook
                entries
            "vq/codebook_loss" : Tensor[1]
                Codebook loss to update the codebook
            "length" : int
                Number of samples in input audio
        """
        # Mixed precision: if frozen encoder, run encoder+quantizer in float32
        if self.is_frozen_encoder:
            audio_data_fp32 = audio_data.float()
            z = self.encoder(audio_data_fp32)
            q_res = self.quantizer(z)
            # Convert output to decoder's dtype (bfloat16) - get dtype from decoder
            decoder_dtype = next(self.decoder.parameters()).dtype
            q_res["z"] = q_res["z"].to(decoder_dtype)
            if "kl" in q_res:
                q_res["kl"] = q_res["kl"].to(decoder_dtype)
            return q_res
        else:
            z = self.encoder(audio_data)
            return self.quantizer(z)

    def decode(self, z: torch.Tensor):
        """Decode given latent codes and return audio data

        Parameters
        ----------
        z : Tensor[B x D x T]
            Quantized continuous representation of input
        length : int, optional
            Number of samples in output audio, by default None

        Returns
        -------
        dict
            A dictionary with the following keys:
            "audio" : Tensor[B x 1 x length]
                Decoded audio data.
        """
        return self.decoder(z)

    def forward(
        self,
        audio_data: torch.Tensor,
        sample_rate: int = None,
        n_quantizers: int = None,
    ):
        """Model forward pass

        Parameters
        ----------
        audio_data : Tensor[B x 1 x T]
            Audio data to encode
        sample_rate : int, optional
            Sample rate of audio data in Hz, by default None
            If None, defaults to `self.sample_rate`
        n_quantizers : int, optional
            Number of quantizers to use, by default None.
            If None, all quantizers are used.

        Returns
        -------
        dict
            A dictionary with the following keys:
            "z" : Tensor[B x D x T]
                Quantized continuous representation of input
            "codes" : Tensor[B x N x T]
                Codebook indices for each codebook
                (quantized discrete representation of input)
            "latents" : Tensor[B x N*D x T]
                Projected latents (continuous representation of input before quantization)
            "vq/commitment_loss" : Tensor[1]
                Commitment loss to train encoder to predict vectors closer to codebook
                entries
            "vq/codebook_loss" : Tensor[1]
                Codebook loss to update the codebook
            "length" : int
                Number of samples in input audio
            "audio" : Tensor[B x 1 x length]
                Decoded audio data.
        """
        length = audio_data.shape[-1]
        audio_data = self.preprocess(audio_data, sample_rate)
        q_res = self.encode(audio_data, n_quantizers)

        x = self.decode(q_res["z"])
        return {
            "audio": x[..., :length],
            **q_res,
        }

    def get_num_params(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)