import math
from typing import List

import numpy as np
import torch
from einops import rearrange
from torch import nn
from rotary_embedding_torch import RotaryEmbedding

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


class TransformerBlock(nn.Module):
    def __init__(self, dim=512, mlp_dim=2048, num_heads=8):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, eps=1e-2)
        self.norm2 = nn.LayerNorm(dim, eps=1e-2)
        self.attention = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(dim, mlp_dim), nn.GELU(), nn.Linear(mlp_dim, dim))

        self._reset_parameters()

        self.rope = RotaryEmbedding(dim=(dim // num_heads // 2))
        self.num_heads = num_heads
        self.head_dim = dim // num_heads

    def _reset_parameters(self):
        # Initialize with smaller weights
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p, gain=0.1)

    def forward(self, x):
        # x: [B, T, D]
        B, T, D = x.shape

        # Apply normalization
        normed_x = self.norm1(x)

        # Project to queries, keys, values and reshape to [B, H, T, D/H]
        q = k = normed_x.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = normed_x.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)

        # Apply RoPE to queries and keys
        q = self.rope.rotate_queries_or_keys(q)
        k = self.rope.rotate_queries_or_keys(k)

        # Reshape back to [B, T, D] for MultiheadAttention
        q = q.transpose(1, 2).reshape(B, T, D)
        k = k.transpose(1, 2).reshape(B, T, D)
        v = v.transpose(1, 2).reshape(B, T, D)

        # Self attention
        attn_out, _ = self.attention(q, k, v, need_weights=False)
        x = x + attn_out

        # MLP
        x = x + self.mlp(self.norm2(x))
        return x


class TransformerEncoder(nn.Module):
    def __init__(self, dim=512, depth=8, mlp_dim=2048, num_heads=8):
        super().__init__()
        self.layers = nn.ModuleList([TransformerBlock(dim, mlp_dim, num_heads) for _ in range(depth)])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


@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 = (mean * mean + 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 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 MusicFMDAC(nn.Module):
    def __init__(
        self,
        latent_dim: int = 1024,
        decoder_dim: int = 2048,
        decoder_rates: List[int] = [8, 8, 5, 3, 2],
        sample_rate: int = 48000,
        n_transformer_layers: int = 0,
        **kwargs,
    ):
        super().__init__()

        self.decoder_dim = decoder_dim
        self.decoder_rates = decoder_rates
        self.sample_rate = sample_rate

        self.latent_dim = latent_dim
        self.n_transformer_layers = n_transformer_layers
        self.hop_length = 25
        if n_transformer_layers > 0:
            self.transformer_encoder = TransformerEncoder(dim=latent_dim, depth=n_transformer_layers)

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


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

    def get_model_hyperparameters(self):
        return {
            "kwargs": {
                "latent_dim": self.latent_dim,
                "decoder_dim": self.decoder_dim,
                "decoder_rates": self.decoder_rates,
                "sample_rate": self.sample_rate,
                "n_transformer_layers": self.n_transformer_layers,
            }
        }

    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 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 1024 x T]
            Encoded MuiscFM embeddings
        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 = int(audio_data.shape[-1] / 25 * 48000)
        if self.n_transformer_layers > 0:
            audio_data = rearrange(audio_data, "b c t -> b t c")
            audio_data = self.transformer_encoder(audio_data)
            audio_data = rearrange(audio_data, "b t c -> b c t")

        x = self.decoder(audio_data)
        return {
            "audio": x[..., :length],
        }

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