import math
import torch
import funcy
from torch import nn
from torch.nn import functional as F
try:
    from torch.nn.utils.parametrizations import weight_norm
except ImportError:
    from torch.nn.utils import weight_norm
from einops import rearrange
from suno_utils.audio import Audio
from suno_utils.utils.s3 import read_from_s3


def init_weights(m):
    # Works for plain Conv and parametrizations.weight_norm-wrapped Conv
    if isinstance(m, (nn.Conv1d, nn.Conv2d)):
        # If weight_norm was applied via parametrizations, module has weight_v/weight_g
        if hasattr(m, "weight_v") and hasattr(m, "weight_g"):
            nn.init.trunc_normal_(m.weight_v, std=0.02)
            nn.init.ones_(m.weight_g)
        else:
            nn.init.trunc_normal_(m.weight, std=0.02)
        if m.bias is not None:
            nn.init.constant_(m.bias, 0)
        

@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))


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)


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


class SoftLimiter(nn.Module):
    def __init__(self, init_gain=0.8):
        super().__init__()
        # gain in (0, 2) via sigmoid
        self.logit_gain = nn.Parameter(torch.logit(torch.tensor(init_gain)))
    def forward(self, x):
        g = torch.sigmoid(self.logit_gain) * 2.0
        return torch.tanh(g * x)


class BottleneckBlock(nn.Module):
    def __init__(self, inp_channels, oup_channels):
        super().__init__()
        self.conv1 = weight_norm(nn.Conv1d(inp_channels, oup_channels, kernel_size=1))
        self.conv2 = weight_norm(nn.Conv1d(oup_channels, oup_channels, kernel_size=1))
        self.conv3 = nn.Conv1d(inp_channels, oup_channels, kernel_size=1)
        self.elu = nn.ELU()

    def forward(self, x):
        inp = x
        x = self.conv1(self.elu(x))
        x = self.conv2(self.elu(x))
        inp = self.conv3(inp)
        x = x + inp
        return x


class DecoderBlock(nn.Module):
    """Upsample → ConvTranspose → Residual skip, exact ×stride output"""
    def __init__(self, inp_channels, oup_channels, strides):
        super().__init__()
        S_f, S_t = strides
        self.elu = nn.ELU()

        self.upsample = nn.Upsample(scale_factor=(S_f, S_t), mode="nearest") \
            if (S_f > 1 or S_t > 1) else nn.Identity()

        self.conv1 = weight_norm(
            nn.Conv2d(inp_channels, oup_channels, kernel_size=3, padding=1)
        )

        self.conv2 = weight_norm(
            nn.ConvTranspose2d(
                oup_channels,
                oup_channels,
                kernel_size=3,
                stride=(S_f, S_t),
                padding=1,
                output_padding=(S_f - 1, S_t - 1),
            )
        )

        self.skip = nn.Identity() if inp_channels == oup_channels else \
            nn.Conv2d(inp_channels, oup_channels, kernel_size=1)

    def forward(self, x):
        res = self.skip(self.upsample(x))
        x = self.conv1(self.elu(x))
        x = self.conv2(self.elu(x))
        return x + res


class Decoder(nn.Module):
    def __init__(self, n_channels=256, vae_dim=128):
        super().__init__()
        # Map VAE dim back to flattened (n * f)
        self.bottleneck_block = BottleneckBlock(vae_dim, n_channels * 40)

        # Expect to reshape to (B, n=n_channels*8, f=?, t)
        self.pre_decoder_block = DecoderBlock(n_channels * 8, n_channels * 16, (1, 1))

        blocks = nn.ModuleList([
            # Mirrors encoder order: (1,2) → (2,2) → (2,1) → (2,1) → (3,1) → (2,1) → (2,1)
            DecoderBlock(n_channels * 8, n_channels * 8, (1, 2)),
            DecoderBlock(n_channels * 8, n_channels * 4, (2, 2)),
            DecoderBlock(n_channels * 4, n_channels * 4, (2, 1)),
            DecoderBlock(n_channels * 4, n_channels * 4, (2, 1)),
            DecoderBlock(n_channels * 4, n_channels * 2, (3, 1)),
            DecoderBlock(n_channels * 2, n_channels * 2, (2, 1)),
            DecoderBlock(n_channels * 2, n_channels * 1, (2, 1)),
        ])
        self.decoder_blocks = nn.Sequential(*blocks)

        # Keep the final head plain (no weight norm) for stable amplitude/phase
        self.final_conv = nn.Conv2d(n_channels, 2, kernel_size=(7, 7), padding=(3, 3), bias=False)
        self.istft = iSTFTBlock()

    def forward(self, x, target_length):
        # x: (B, vae_dim, T_lat) → (B, n_channels*40, T_lat)
        x = self.bottleneck_block(x)

        # Infer f from expected 'n' of pre_decoder_block input
        B, C, T = x.shape
        n_expected = self.pre_decoder_block.conv1.in_channels  # = n_channels*8
        assert C % n_expected == 0, f"Flattened channels {C} not divisible by n={n_expected}"
        f = C // n_expected

        # Reshape (B, n*f, t) → (B, n, f, t)
        x = rearrange(x, "b (n f) t -> b n f t", n=n_expected, f=f)

        # Pre-decoder mixes features across channels (no up/downsample)
        x = self.pre_decoder_block(x)  # (B, n*2, f, t)

        # Split back into stereo by folding channel-dim into batch
        x = rearrange(x, "b (c n) f t -> (b c) n f t", c=2)

        # Upsample back to (F=480, frames=original)
        x = self.decoder_blocks(x)

        # Project to (real, imag)
        x = self.final_conv(x)  # (B*C, 2, F, T_frames)

        # Inverse STFT to waveform; crop to exact target_length (e.g., 48000)
        x = self.istft(x, length=target_length)  # (B*C, time)

        # tanh
        x = torch.tanh(x)

        # Restore stereo
        x = rearrange(x, "(b c) t -> b c t", c=2)
        return x


class iSTFTBlock(nn.Module):
    """
    Inverse short-time Fourier transform block.
    Input:  stacked real/imag spectrogram with Nyquist removed
            shape: (batch, 2, frequency, time) where frequency == n_fft//2
            for n_fft=960 -> frequency=480
    Output: mono waveform (batch, time)

    Notes:
      - Matches STFTBlock(n_fft=960, hop=480, win=960, center=True, onesided=True).
      - Restores the Nyquist bin as zeros before istft.
      - Optionally pass `length` to trim padding introduced by center=True.
    """
    def __init__(self):
        super().__init__()
        self.n_fft = 960
        self.hop_length = 480
        self.win_length = 960
        # register as buffer so it moves with .to(device)
        self.register_buffer("window", torch.hann_window(self.win_length))

    def forward(self, x, length: int | None = None):
        """
        x: (B, 2, F, T) with F == n_fft//2 (Nyquist removed)
        length: optional target waveform length to trim padding from center=True
        """
        # split real/imag
        real = x[:, 0]                   # (B, F, T)
        imag = x[:, 1]                   # (B, F, T)

        B, F, T = real.shape
        assert F == self.n_fft // 2, f"Expected F={self.n_fft//2}, got {F}"

        # restore Nyquist bin (zeros) to make onesided length = n_fft//2 + 1
        zero_nyq_r = torch.zeros(B, 1, T, dtype=real.dtype, device=real.device)
        zero_nyq_i = torch.zeros_like(zero_nyq_r)

        real_full = torch.cat([real, zero_nyq_r], dim=1)   # (B, F+1, T)
        imag_full = torch.cat([imag, zero_nyq_i], dim=1)   # (B, F+1, T)

        X = torch.complex(real_full.float(), imag_full.float())            # (B, F+1, T)

        # inverse STFT (matches STFTBlock settings)
        y = torch.istft(
            X,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window,
            center=True,
            normalized=False,
            onesided=True,
            length=length,   # None -> no explicit trim; provide original N to crop
        )
        # y: (B, time)
        return y


class DACSpectroStreamVAE(nn.Module):
    def __init__(self, encoder_dim: int = 128,
        encoder_rates: List[int] = [2, 3, 5, 8, 8], 
        n_channels=256, 
        vae_dim=128, 
        is_frozen_encoder=False, 
        **kwargs,
    ):
        super().__init__()
        # self.encoder = Encoder(n_channels, vae_dim)
        self.encoder = Encoder(encoder_dim, encoder_rates, vae_dim * 2)
        self.quantizer = VAEBottleneck()
        self.decoder = Decoder(n_channels, vae_dim)
        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"]
            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

            # Convert frozen parts to float32 for better quality
            self.encoder = self.encoder.float()
            self.quantizer = self.quantizer.float()
        
    def encode(self, x):
        if self.is_frozen_encoder:
            audio_data_fp32 = x.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(x)
            return self.quantizer(z)
        
    def decode(self, z, target_length):
        return self.decoder(z, target_length)
    
    def forward(self, audio_data, is_mask=None):
        q_res = self.encode(audio_data)
        x = self.decode(q_res["z"], audio_data.shape[-1])
        return {
            "audio": x,
            "mask": None,
            "masked": audio_data,
            **q_res,
        }
    
    def get_num_params(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


if __name__ == "__main__":
    sample_audio = Audio.from_file("/home/minz/temp/country_road.mp3", sample_rate=48000, n_channels=2)
    inp = torch.from_numpy(sample_audio.array_float[:, :48000*1]).unsqueeze(0)

    model = SpectroStreamVAE(n_channels=4, vae_dim=64, is_frozen_encoder=True)
    out = model(inp)
    print(out["audio"].shape)