import math
import torch
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


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 STFTBlock(nn.Module):
    """
    Short-time Fourier transform block for converting waveform to spectrogram.
    Matches the iSTFT parameters in the decoder.
    """
    def __init__(self):
        super().__init__()
        self.n_fft = 960
        self.hop_length = 480
        self.win_length = 960
        # Register window as buffer so it moves with .to(device)
        self.register_buffer("window", torch.hann_window(self.win_length))

    def forward(self, x):
        """
        Input: (B, 2, T) stereo waveform
        Output: (B*2, 2, F, T_frames) where F=480 (n_fft//2, Nyquist removed)
        """
        B, C, T = x.shape
        assert C == 2, f"Expected stereo input, got {C} channels"
        
        # Process each channel separately
        # Flatten batch and channels for STFT
        x = rearrange(x, "b c t -> (b c) t")  # (B*2, T)
        
        # Compute STFT
        X = torch.stft(
            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,
            return_complex=True
        )  # (B*2, F+1, T_frames) where F+1 = 481 (includes Nyquist)
        
        # Remove Nyquist bin to get (B*2, 480, T_frames)
        X = X[:, :-1, :]
        
        # Split into real and imaginary
        real = X.real  # (B*2, 480, T_frames)
        imag = X.imag  # (B*2, 480, T_frames)
        
        # Stack real and imag as channels: (B*2, 2, 480, T_frames)
        out = torch.stack([real, imag], dim=1)
        
        return out


def WNConv2d(*args, **kwargs):
    return weight_norm(nn.Conv2d(*args, **kwargs))


class EncoderBlock2D(nn.Module):
    """
    2D Encoder block for spectral processing.
    Downsamples in both frequency and time dimensions.
    """
    def __init__(self, inp_channels, oup_channels, strides):
        super().__init__()
        S_f, S_t = strides
        self.elu = nn.ELU()
        self.strides = strides
        
        # Pre-activation convolution
        self.conv1 = weight_norm(
            nn.Conv2d(inp_channels, oup_channels, kernel_size=3, padding=1)
        )
        
        # Strided convolution for downsampling
        # Use kernel_size = stride for clean downsampling
        self.conv2 = weight_norm(
            nn.Conv2d(
                oup_channels,
                oup_channels,
                kernel_size=(S_f if S_f > 1 else 3, S_t if S_t > 1 else 3),
                stride=(S_f, S_t),
                padding=(0 if S_f > 1 else 1, 0 if S_t > 1 else 1)
            )
        )
        
        # Skip connection - match the exact output size
        if S_f > 1 or S_t > 1:
            self.skip_downsample = nn.Conv2d(
                inp_channels,
                oup_channels,
                kernel_size=(S_f, S_t),
                stride=(S_f, S_t),
                padding=0
            )
        else:
            self.skip_downsample = nn.Identity() if inp_channels == oup_channels else \
                nn.Conv2d(inp_channels, oup_channels, kernel_size=1)
    
    def forward(self, x):
        res = self.skip_downsample(x)
        x = self.conv1(self.elu(x))
        x = self.conv2(self.elu(x))
        
        # Ensure shapes match for residual connection
        if x.shape != res.shape:
            # Crop or pad to match
            if x.shape[2] < res.shape[2]:
                res = res[:, :, :x.shape[2], :]
            if x.shape[3] < res.shape[3]:
                res = res[:, :, :, :x.shape[3]]
        
        return x + res


class SpectralEncoder(nn.Module):
    """
    Spectral encoder that processes STFT features with 2D convolutions.
    Simpler design: downsample spatially, then flatten and project.
    """
    def __init__(self, n_channels=256, vae_dim=128):
        super().__init__()
        
        # Initial convolution on spectrogram (2 channels: real, imag)
        self.initial_conv = weight_norm(
            nn.Conv2d(2, n_channels, kernel_size=(7, 7), padding=(3, 3))
        )
        
        # Encoder blocks - progressive downsampling
        # Start: (B*2, n_channels, 480, ~100) for 1-second audio
        # Target: Downsample to manageable size
        blocks = nn.ModuleList([
            EncoderBlock2D(n_channels * 1, n_channels * 1, (2, 1)),   # F: 480→240
            EncoderBlock2D(n_channels * 1, n_channels * 2, (2, 1)),   # F: 240→120
            EncoderBlock2D(n_channels * 2, n_channels * 2, (2, 1)),   # F: 120→60
            EncoderBlock2D(n_channels * 2, n_channels * 4, (2, 2)),   # F: 60→30, T: /2
            EncoderBlock2D(n_channels * 4, n_channels * 4, (2, 2)),   # F: 30→15, T: /2
        ])
        self.encoder_blocks = nn.Sequential(*blocks)
        
        self.elu = nn.ELU()
        
        # After encoding: (B*2, n_channels*4, ~15, ~25)
        # Merge stereo and flatten frequency, then project to VAE dims
        # We'll use adaptive pooling to get consistent size
        self.adaptive_pool = nn.AdaptiveAvgPool2d((5, None))  # Pool frequency to fixed size
        
        # Final projection: flatten and project to vae_dim*2
        # After pool: (B*2, n_channels*4, 5, T_frames)
        # Merge stereo: (B, n_channels*4*2, 5, T_frames)
        # Flatten freq: (B, n_channels*4*2*5, T_frames)
        self.final_proj = nn.Conv1d(n_channels * 4 * 2 * 5, vae_dim * 2, kernel_size=1)
    
    def forward(self, x):
        """
        Input: (B*2, 2, F, T) where F=480, T=time frames (~100 for 1-second)
        Output: (B, vae_dim*2, T_frames) for VAE mean+scale
        """
        BC_in = x.shape[0]  # B*2
        B = BC_in // 2
        
        # Initial conv
        x = self.initial_conv(x)  # (B*2, n_channels, 480, T)
        
        # Downsample through encoder blocks
        x = self.encoder_blocks(x)  # (B*2, n_channels*4, F', T')
        
        # Adaptive pooling in frequency dimension
        x = self.adaptive_pool(x)  # (B*2, n_channels*4, 5, T')
        
        # Merge stereo channels: (B*2, n*4, 5, T') → (B, n*4*2, 5, T')
        x = rearrange(x, "(b c) n f t -> b (c n f) t", c=2, b=B)
        
        # Project to VAE dimensions
        x = self.final_proj(x)  # (B, vae_dim*2, T')
        
        return 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 = (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


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()
        self.limiter = SoftLimiter(init_gain=0.8)

    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)

        # Soft limit
        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 SpectroStreamVAE(nn.Module):
    """
    Spectral audio codec with VAE bottleneck.
    Processes audio in STFT domain throughout (encoder and decoder).
    """
    def __init__(self, n_channels=256, vae_dim=128):
        super().__init__()
        # STFT: waveform → spectrogram
        self.stft = STFTBlock()
        
        # Spectral encoder: spectrogram → latents
        self.encoder = SpectralEncoder(n_channels, vae_dim)
        
        # VAE bottleneck: latents → sampled latents + KL
        self.bottleneck = VAEBottleneck()
        
        # Spectral decoder: latents → spectrogram → waveform (with iSTFT)
        self.decoder = Decoder(n_channels, vae_dim)
        
        self.apply(init_weights)
        
    def encode(self, x):
        """
        Encode waveform to VAE latents.
        
        Args:
            x: (B, 2, T) stereo waveform
        Returns:
            dict with keys: z, kl, mean, scale
        """
        # Waveform → STFT
        spec = self.stft(x)  # (B*2, 2, F=480, T_frames)
        
        # STFT → latents (mean, scale)
        z = self.encoder(spec)  # (B, vae_dim*2, T_latent)
        
        # VAE sampling + KL
        return self.bottleneck(z)
        
    def decode(self, z, target_length):
        """
        Decode VAE latents to waveform.
        
        Args:
            z: (B, vae_dim, T_latent) sampled latents
            target_length: int, desired waveform length
        Returns:
            (B, 2, target_length) stereo waveform
        """
        return self.decoder(z, target_length)
    
    def forward(self, audio_data, sample_rate=None, n_quantizers=None):
        """
        Full forward pass: waveform → STFT → encode → VAE → decode → iSTFT → waveform
        
        Args:
            audio_data: (B, 2, T) stereo waveform
            sample_rate: Ignored (for compatibility)
            n_quantizers: Ignored (for compatibility)
        Returns:
            dict with keys: audio, kl, z, mean, scale
        """
        q_res = self.encode(audio_data)
        x = self.decode(q_res["z"], audio_data.shape[-1])
        return {
            "audio": x,
            **q_res,
        }
    
    def get_num_params(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


if __name__ == "__main__":
    # Test the spectral codec
    print("Testing SpectroStreamVAE...")
    
    # Create test input
    batch_size = 2
    duration_s = 1.0
    sample_rate = 48000
    inp = torch.randn(batch_size, 2, int(sample_rate * duration_s))
    print(f"Input shape: {inp.shape}")
    
    # Create model
    model = SpectroStreamVAE(n_channels=256, vae_dim=64)
    print(f"Model parameters: {model.get_num_params():,}")
    
    # Test forward pass
    out = model(inp)
    print(f"Output shape: {out['audio'].shape}")
    print(f"Latent shape: {out['z'].shape}")
    print(f"KL value: {out['kl'].item():.4f}")
    
    # Verify shapes match
    assert out["audio"].shape == inp.shape, f"Shape mismatch: {out['audio'].shape} vs {inp.shape}"
    print("\n✅ SpectroStreamVAE test passed!")