"""AutoencoderOobleck — Audio VAE ported from diffusers to MLX.

This is a channels-last MLX port of the Oobleck autoencoder used in ACE-Step
and Stable Audio. All convolutions expect input shaped [B, T, C] (MLX convention)
rather than the PyTorch [B, C, T] convention.

Weight norm handling:
    The PyTorch model uses torch.nn.utils.weight_norm on all Conv1d/ConvTranspose1d
    layers, which creates weight_g and weight_v parameters instead of a single weight.
    In this MLX implementation we use plain Conv1d/ConvTranspose1d layers. The weight
    converter (convert.py) must fuse weight_norm parameters before loading:
        weight = weight_g * weight_v / norm(weight_v, dim=...)
    Once fused, the resulting single weight tensor is loaded directly into these layers.
"""

import math
from typing import List, Optional, Tuple

import mlx.core as mx
import mlx.nn as nn

from .config import AceStepVAEConfig


class Snake1d(nn.Module):
    """Learnable periodic activation: x + (1/beta) * sin^2(alpha * x).

    Parameters alpha and beta are stored in log-scale and exponentiated at
    forward time. In PyTorch these are shaped (1, hidden_dim, 1) for [B, C, T]
    broadcasting. In MLX (channels-last [B, T, C]) we store them as (hidden_dim,)
    and rely on standard broadcasting.
    """

    def __init__(self, hidden_dim: int):
        super().__init__()
        # Initialized to zeros in log-scale (so exp(0)=1 initially)
        self.alpha = mx.zeros((hidden_dim,))
        self.beta = mx.zeros((hidden_dim,))

    def __call__(self, x: mx.array) -> mx.array:
        alpha = mx.exp(self.alpha)
        beta = mx.exp(self.beta)
        return x + mx.reciprocal(beta + 1e-9) * mx.square(mx.sin(alpha * x))


class OobleckResidualUnit(nn.Module):
    """Residual block with dilated Conv1d layers and Snake activations.

    Structure: Snake -> Conv1d(k=7, dilated) -> Snake -> Conv1d(k=1) + skip.
    """

    def __init__(self, dimension: int, dilation: int = 1):
        super().__init__()
        pad = ((7 - 1) * dilation) // 2

        self.snake1 = Snake1d(dimension)
        self.conv1 = nn.Conv1d(
            dimension, dimension, kernel_size=7, dilation=dilation, padding=pad
        )
        self.snake2 = Snake1d(dimension)
        self.conv2 = nn.Conv1d(dimension, dimension, kernel_size=1)

    def __call__(self, hidden_state: mx.array) -> mx.array:
        output = self.conv1(self.snake1(hidden_state))
        output = self.conv2(self.snake2(output))

        # Trim hidden_state if conv output is shorter (boundary effects)
        padding = (hidden_state.shape[1] - output.shape[1]) // 2
        if padding > 0:
            hidden_state = hidden_state[:, padding:-padding, :]
        return hidden_state + output


class OobleckEncoderBlock(nn.Module):
    """Encoder block: 3 residual units (dilations 1, 3, 9) + strided downsampling."""

    def __init__(self, input_dim: int, output_dim: int, stride: int = 1):
        super().__init__()
        self.res_unit1 = OobleckResidualUnit(input_dim, dilation=1)
        self.res_unit2 = OobleckResidualUnit(input_dim, dilation=3)
        self.res_unit3 = OobleckResidualUnit(input_dim, dilation=9)
        self.snake1 = Snake1d(input_dim)
        self.conv1 = nn.Conv1d(
            input_dim,
            output_dim,
            kernel_size=2 * stride,
            stride=stride,
            padding=math.ceil(stride / 2),
        )

    def __call__(self, hidden_state: mx.array) -> mx.array:
        hidden_state = self.res_unit1(hidden_state)
        hidden_state = self.res_unit2(hidden_state)
        hidden_state = self.snake1(self.res_unit3(hidden_state))
        hidden_state = self.conv1(hidden_state)
        return hidden_state


class OobleckDecoderBlock(nn.Module):
    """Decoder block: strided upsampling (ConvTranspose1d) + 3 residual units."""

    def __init__(self, input_dim: int, output_dim: int, stride: int = 1):
        super().__init__()
        self.snake1 = Snake1d(input_dim)
        self.conv_t1 = nn.ConvTranspose1d(
            input_dim,
            output_dim,
            kernel_size=2 * stride,
            stride=stride,
            padding=math.ceil(stride / 2),
        )
        self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
        self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)
        self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9)

    def __call__(self, hidden_state: mx.array) -> mx.array:
        hidden_state = self.snake1(hidden_state)
        hidden_state = self.conv_t1(hidden_state)
        hidden_state = self.res_unit1(hidden_state)
        hidden_state = self.res_unit2(hidden_state)
        hidden_state = self.res_unit3(hidden_state)
        return hidden_state


class OobleckEncoder(nn.Module):
    """Oobleck encoder: Conv1d -> N encoder blocks -> Snake -> Conv1d.

    Progressively downsamples the time axis and increases channels according
    to channel_multiples. Output is the raw encoder hidden states (before
    splitting into mean/scale for the Gaussian posterior).
    """

    def __init__(
        self,
        encoder_hidden_size: int,
        audio_channels: int,
        downsampling_ratios: List[int],
        channel_multiples: List[int],
    ):
        super().__init__()
        strides = downsampling_ratios
        # Prepend 1 to channel_multiples so index 0 = base multiplier
        c_mults = [1] + list(channel_multiples)

        self.conv1 = nn.Conv1d(
            audio_channels, encoder_hidden_size, kernel_size=7, padding=3
        )

        self.block = []
        for i, stride in enumerate(strides):
            self.block.append(
                OobleckEncoderBlock(
                    input_dim=encoder_hidden_size * c_mults[i],
                    output_dim=encoder_hidden_size * c_mults[i + 1],
                    stride=stride,
                )
            )

        d_model = encoder_hidden_size * c_mults[-1]
        self.snake1 = Snake1d(d_model)
        self.conv2 = nn.Conv1d(d_model, encoder_hidden_size, kernel_size=3, padding=1)

    def __call__(self, hidden_state: mx.array) -> mx.array:
        # hidden_state: [B, T, audio_channels]
        hidden_state = self.conv1(hidden_state)

        for module in self.block:
            hidden_state = module(hidden_state)

        hidden_state = self.snake1(hidden_state)
        hidden_state = self.conv2(hidden_state)
        return hidden_state


class OobleckDecoder(nn.Module):
    """Oobleck decoder: Conv1d -> N decoder blocks -> Snake -> Conv1d.

    Progressively upsamples the time axis from the latent representation back
    to the original audio sample rate.
    """

    def __init__(
        self,
        channels: int,
        input_channels: int,
        audio_channels: int,
        upsampling_ratios: List[int],
        channel_multiples: List[int],
    ):
        super().__init__()
        strides = upsampling_ratios
        c_mults = [1] + list(channel_multiples)

        self.conv1 = nn.Conv1d(
            input_channels, channels * c_mults[-1], kernel_size=7, padding=3
        )

        self.block = []
        for i, stride in enumerate(strides):
            self.block.append(
                OobleckDecoderBlock(
                    input_dim=channels * c_mults[len(strides) - i],
                    output_dim=channels * c_mults[len(strides) - i - 1],
                    stride=stride,
                )
            )

        self.snake1 = Snake1d(channels)
        self.conv2 = nn.Conv1d(
            channels, audio_channels, kernel_size=7, padding=3, bias=False
        )

    def __call__(self, hidden_state: mx.array) -> mx.array:
        # hidden_state: [B, T_latent, latent_dim]
        hidden_state = self.conv1(hidden_state)

        for layer in self.block:
            hidden_state = layer(hidden_state)

        hidden_state = self.snake1(hidden_state)
        hidden_state = self.conv2(hidden_state)
        return hidden_state


class OobleckDiagonalGaussianDistribution:
    """Diagonal Gaussian posterior for the Oobleck VAE.

    Takes the raw encoder output and splits it along the channel axis into
    mean and scale, where std = softplus(scale) + 1e-4.
    """

    def __init__(self, parameters: mx.array, deterministic: bool = False):
        # parameters: [B, T, 2*latent_dim] — channels-last
        self.parameters = parameters
        # Split along last axis (channels) into mean and scale
        dim = parameters.shape[-1] // 2
        self.mean = parameters[..., :dim]
        self.scale = parameters[..., dim:]
        self.std = nn.softplus(self.scale) + 1e-4
        self.deterministic = deterministic

    def sample(self) -> mx.array:
        noise = mx.random.normal(self.mean.shape)
        return self.mean + self.std * noise

    def mode(self) -> mx.array:
        return self.mean


class AutoencoderOobleck(nn.Module):
    """Audio VAE using the Oobleck architecture (from Stable Audio / ACE-Step).

    Encodes stereo audio waveforms into a compressed latent representation and
    decodes latents back to audio.

    Input/output convention (MLX channels-last):
        - Audio:  [B, T_samples, audio_channels]  (e.g., [1, 96000, 2])
        - Latent: [B, T_latent, latent_dim]        (e.g., [1, 50, 64])

    The total temporal compression factor (hop_length) equals the product of
    all downsampling_ratios: 2*4*4*6*10 = 1920 for the default ACE-Step config.
    """

    def __init__(self, config: AceStepVAEConfig):
        super().__init__()
        self.config = config
        self.hop_length = config.hop_length

        self.encoder = OobleckEncoder(
            encoder_hidden_size=config.encoder_hidden_size,
            audio_channels=config.audio_channels,
            downsampling_ratios=config.downsampling_ratios,
            channel_multiples=config.channel_multiples,
        )

        self.decoder = OobleckDecoder(
            channels=config.decoder_channels,
            input_channels=config.decoder_input_channels,
            audio_channels=config.audio_channels,
            upsampling_ratios=config.upsampling_ratios,
            channel_multiples=config.channel_multiples,
        )

    def encode(
        self, x: mx.array, deterministic: bool = False
    ) -> Tuple[mx.array, OobleckDiagonalGaussianDistribution]:
        """Encode audio to latent space.

        Args:
            x: Audio tensor [B, T, audio_channels].
            deterministic: If True, return the mean (no sampling).

        Returns:
            Tuple of (latent, posterior).
            latent: [B, T/hop_length, latent_dim] sampled (or mode) latent.
            posterior: The full Gaussian distribution for KL if needed.
        """
        h = self.encoder(x)
        # h: [B, T/hop_length, encoder_hidden_size]
        # encoder_hidden_size = 2 * decoder_input_channels (= 2 * latent_dim)
        # so h splits evenly into mean and scale
        posterior = OobleckDiagonalGaussianDistribution(h, deterministic=deterministic)

        if deterministic:
            latent = posterior.mode()
        else:
            latent = posterior.sample()

        return latent, posterior

    def decode(self, z: mx.array) -> mx.array:
        """Decode latent back to audio.

        Args:
            z: Latent tensor [B, T_latent, latent_dim].

        Returns:
            Reconstructed audio [B, T_latent * hop_length, audio_channels].
        """
        return self.decoder(z)

    def __call__(
        self, audio: mx.array, deterministic: bool = False
    ) -> Tuple[mx.array, OobleckDiagonalGaussianDistribution]:
        """Full forward pass: encode then decode.

        Args:
            audio: [B, T, audio_channels].
            deterministic: Use posterior mean instead of sampling.

        Returns:
            Tuple of (reconstructed_audio, posterior).
        """
        latent, posterior = self.encode(audio, deterministic=deterministic)
        reconstructed = self.decode(latent)
        return reconstructed, posterior
