"""Standalone decoder-only MiniMax H3 audio VAE."""

from __future__ import annotations

import math
from pathlib import Path

import torch
from safetensors import safe_open
from torch import nn
from torch.nn import functional as F


def snake(x: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
    t = torch.sin(alpha * x)
    return t.mul_(t).mul_((beta + 1e-9).reciprocal()).add_(x)


class SnakeBeta(nn.Module):
    def __init__(self, channels: int, *, device=None):
        super().__init__()
        self.alpha = nn.Parameter(torch.empty(channels, device=device))
        self.beta = nn.Parameter(torch.empty(channels, device=device))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        alpha = torch.exp(self.alpha.to(device=x.device, dtype=x.dtype)).view(1, -1, 1)
        beta = torch.exp(self.beta.to(device=x.device, dtype=x.dtype)).view(1, -1, 1)
        return snake(x, alpha, beta)


def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
    even = kernel_size % 2 == 0
    half_size = kernel_size // 2
    delta_f = 4 * half_width
    a = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
    if a > 50.0:
        beta = 0.1102 * (a - 8.7)
    elif a >= 21.0:
        beta = 0.5842 * (a - 21) ** 0.4 + 0.07886 * (a - 21.0)
    else:
        beta = 0.0
    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
    time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size
    filt = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
    filt /= filt.sum()
    return filt.view(1, 1, kernel_size)


class UpSample1d(nn.Module):
    def __init__(self, ratio: int = 2, kernel_size: int = 12, *, device=None):
        super().__init__()
        self.ratio = ratio
        self.stride = ratio
        self.pad = kernel_size // ratio - 1
        self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2
        self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2
        self.register_buffer("filter", kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size).to(device=device))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _, channels, _ = x.shape
        x = F.pad(x, (self.pad, self.pad), mode="replicate")
        filt = self.filter.to(device=x.device, dtype=x.dtype).expand(channels, -1, -1)
        x = F.conv_transpose1d(x, filt, stride=self.stride, groups=channels).mul_(self.ratio)
        return x[..., self.pad_left:-self.pad_right]


class LowPassFilter1d(nn.Module):
    def __init__(self, cutoff: float = 0.5, half_width: float = 0.6, stride: int = 1, kernel_size: int = 12, *, device=None):
        super().__init__()
        self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0)
        self.pad_right = kernel_size // 2
        self.stride = stride
        self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size).to(device=device))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _, channels, _ = x.shape
        x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
        filt = self.filter.to(device=x.device, dtype=x.dtype).expand(channels, -1, -1)
        return F.conv1d(x, filt, stride=self.stride, groups=channels)


class DownSample1d(nn.Module):
    def __init__(self, ratio: int = 2, kernel_size: int = 12, *, device=None):
        super().__init__()
        self.lowpass = LowPassFilter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=kernel_size, device=device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.lowpass(x)


class Activation1d(nn.Module):
    def __init__(self, activation: nn.Module, *, device=None):
        super().__init__()
        self.act = activation
        self.upsample = UpSample1d(device=device)
        self.downsample = DownSample1d(device=device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.downsample(self.act(self.upsample(x)))


def get_padding(kernel_size: int, dilation: int = 1) -> int:
    return int((kernel_size * dilation - dilation) / 2)


class AMPBlock1(nn.Module):
    def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int, int] = (1, 3, 5), *, device=None):
        super().__init__()
        self.convs1 = nn.ModuleList([nn.Conv1d(channels, channels, kernel_size, dilation=d, padding=get_padding(kernel_size, d), device=device) for d in dilation])
        self.convs2 = nn.ModuleList([nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=get_padding(kernel_size, 1), device=device) for _ in dilation])
        self.activations = nn.ModuleList([Activation1d(SnakeBeta(channels, device=device), device=device) for _ in range(len(dilation) * 2)])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        acts1, acts2 = self.activations[::2], self.activations[1::2]
        for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2):
            residual = conv2(act2(conv1(act1(x))))
            x = residual.add_(x)
        return x


class BigVGAN(nn.Module):
    def __init__(self, num_mels: int = 2048, upsample_initial_channel: int = 1024, *, device=None):
        super().__init__()
        upsample_rates = (5, 5, 2, 2, 2, 2, 2)
        upsample_kernel_sizes = (9, 9, 4, 4, 4, 4, 4)
        resblock_kernel_sizes = (3, 7, 11)
        resblock_dilation_sizes = ((1, 3, 5), (1, 3, 5), (1, 3, 5))
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = nn.Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3, device=device)
        self.ups = nn.ModuleList()
        for index, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(nn.ModuleList([nn.ConvTranspose1d(upsample_initial_channel // (2 ** index), upsample_initial_channel // (2 ** (index + 1)), kernel, rate, padding=(kernel - rate) // 2, device=device)]))
        self.resblocks = nn.ModuleList()
        for index in range(len(self.ups)):
            channels = upsample_initial_channel // (2 ** (index + 1))
            for kernel, dilation in zip(resblock_kernel_sizes, resblock_dilation_sizes):
                self.resblocks.append(AMPBlock1(channels, kernel, dilation, device=device))
        self.activation_post = Activation1d(SnakeBeta(channels, device=device), device=device)
        self.conv_post = nn.Conv1d(channels, 1, 7, 1, padding=3, bias=False, device=device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv_pre(x)
        for index in range(self.num_upsamples):
            for upsample in self.ups[index]:
                x = upsample(x)
            combined = None
            for kernel_index in range(self.num_kernels):
                value = self.resblocks[index * self.num_kernels + kernel_index](x)
                combined = value if combined is None else combined + value
            x = combined.div_(self.num_kernels)
        return self.conv_post(self.activation_post(x)).clamp_(-1.0, 1.0)


class MiniMaxH3AudioVAE(nn.Module):
    """Decoder-only MiniMax H3 stereo audio VAE at 32 kHz."""

    def __init__(self, latent_channels: int = 32, latent_dim: int = 2048, decoder_dim: int = 1024, *, device=None):
        super().__init__()
        self.sample_rate = 32000
        self.output_sample_rate = self.sample_rate
        self.samples_per_latent = 800
        self.latents_per_second = 40
        self.dec_in_proj = nn.Conv1d(latent_channels, latent_dim, 1, device=device)
        self.decoder = BigVGAN(num_mels=latent_dim, upsample_initial_channel=decoder_dim, device=device)
        self.register_buffer("latents_mean", torch.empty(latent_channels, device=device))
        self.register_buffer("latents_std", torch.empty(latent_channels, device=device))

    @classmethod
    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32) -> "MiniMaxH3AudioVAE":
        model = cls(device="meta")
        expected = model.state_dict()
        with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
            missing = sorted(set(expected) - set(checkpoint.keys()))
            shape_errors = [(name, tuple(expected[name].shape), tuple(checkpoint.get_slice(name).get_shape())) for name in expected if name in checkpoint.keys() and tuple(expected[name].shape) != tuple(checkpoint.get_slice(name).get_shape())]
            if missing or shape_errors:
                details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else [])
                raise ValueError("incompatible H3 audio VAE checkpoint; " + "; ".join(details))
            weights = {name: checkpoint.get_tensor(name).to(dtype=dtype) for name in expected}
        model.load_state_dict(weights, strict=True, assign=True)
        return model

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        """Decode normalized latents `[B,32,2,T]` to waveform `[B,2,L]`."""
        batch, channels, stereo, steps = z.shape
        z = z.permute(0, 2, 1, 3).reshape(batch * stereo, channels, steps)
        mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
        std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
        z = z * std + mean
        waveform = self.decoder(self.dec_in_proj(z))
        return waveform.reshape(batch, stereo, -1)