"""Motif-Audio: dual-stream audio autoencoder.

Self-contained inference model. Depends only on PyTorch, diffusers,
``torchaudio``, and ``timm``.

Both mel spectrograms are computed inside the model, so it takes a raw 16 kHz
mono waveform directly. Run the model in bfloat16 to match the precision the
weights were trained under, and pad the input so the mel frame count is
divisible by the encoder patch sizes.

Usage::

    import torch, torchaudio

    model = MotifAudio.from_pretrained("Motif-Technologies/Motif-Audio",
                                      low_cpu_mem_usage=False)
    model = model.to("cuda", torch.bfloat16).eval()

    wav, sr = torchaudio.load("input.wav")   # mono 16 kHz, (1, T)
    out = model(wav.unsqueeze(0).cuda())     # (1, 1, T)

    torchaudio.save("output.wav", out["waveform"].squeeze(0).float().cpu(), sr)
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio.transforms
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from timm.models.vision_transformer import DropPath, PatchEmbed
from torch import Tensor


# =============================================================================
#  RoPE helpers (axial 2D)
# =============================================================================


def axial_cos_sin(
    head_dim: int,
    grid_h: int,
    grid_w: int,
    theta: float = 100.0,
    device: torch.device | None = None,
) -> tuple[Tensor, Tensor]:
    assert head_dim % 4 == 0
    half = head_dim // 2
    freq_idx = torch.arange(0, half, 2, device=device, dtype=torch.float32)
    inv_freq = 1.0 / (theta ** (freq_idx / half))
    n = grid_h * grid_w
    pos = torch.arange(n, device=device)
    rows = (pos // grid_w).to(torch.float32)
    cols = (pos % grid_w).to(torch.float32)
    ang_r = torch.outer(rows, inv_freq)
    ang_c = torch.outer(cols, inv_freq)
    ang = torch.cat([ang_r, ang_c], dim=-1)
    ang = torch.cat([ang, ang], dim=-1)
    return ang.cos(), ang.sin()


def rotate_half(x: Tensor) -> Tensor:
    d = x.shape[-1]
    x1 = x[..., : d // 2]
    x2 = x[..., d // 2 :]
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    return x * cos.to(x.dtype) + rotate_half(x) * sin.to(x.dtype)


def prefix_pad(cos: Tensor, sin: Tensor, n: int) -> tuple[Tensor, Tensor]:
    if n <= 0:
        return cos, sin
    b, _, d = cos.shape
    one = torch.ones(b, n, d, dtype=cos.dtype, device=cos.device)
    zero = torch.zeros(b, n, d, dtype=sin.dtype, device=sin.device)
    return torch.cat([one, cos], dim=1), torch.cat([zero, sin], dim=1)


# =============================================================================
#  Semantic encoder sub-modules (SwiGLU, RoPE attention, block)
# =============================================================================


class SwiGLUFFN(nn.Module):
    def __init__(self, dim: int, mlp_ratio: float = 4.0):
        super().__init__()
        hidden = int(mlp_ratio * dim * 2 / 3)
        hidden = (hidden + 7) // 8 * 8
        self.w12 = nn.Linear(dim, 2 * hidden, bias=False)
        self.w3 = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        gate, value = self.w12(x).chunk(2, dim=-1)
        return self.w3(F.silu(gate) * value)


class RoPEAttentionKBiasZero(nn.Module):
    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = True,
                 attn_drop: float = 0.0, proj_drop: float = 0.0, qk_norm: bool = True):
        super().__init__()
        assert dim % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qkv = nn.Linear(dim, dim * 3, bias=False)
        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(dim))
            self.v_bias = nn.Parameter(torch.zeros(dim))
        else:
            self.q_bias = None
            self.v_bias = None
        self.q_norm = nn.RMSNorm(self.head_dim, eps=1e-6) if qk_norm else nn.Identity()
        self.k_norm = nn.RMSNorm(self.head_dim, eps=1e-6) if qk_norm else nn.Identity()
        self.attn_drop_p = attn_drop
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
        B, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(x, self.qkv.weight, qkv_bias)
        qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        q = self.q_norm(q)
        k = self.k_norm(k)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        dropout_p = self.attn_drop_p if self.training else 0.0
        x = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
        x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class LayerScale(nn.Module):
    def __init__(self, dim: int, init_values: float = 1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.full((dim,), init_values))

    def forward(self, x: Tensor) -> Tensor:
        return x * self.gamma


class RoPEBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0,
                 qkv_bias: bool = True, drop: float = 0.0, attn_drop: float = 0.0,
                 drop_path: float = 0.0, norm_layer=nn.RMSNorm, layer_scale_init: float = 0.0):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = RoPEAttentionKBiasZero(dim, num_heads=num_heads, qkv_bias=qkv_bias,
                                           attn_drop=attn_drop, proj_drop=drop)
        if DropPath is not None:
            self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        else:
            self.drop_path = nn.Identity()
        self.norm2 = norm_layer(dim)
        self.mlp = SwiGLUFFN(dim, mlp_ratio=mlp_ratio)
        self.ls1 = LayerScale(dim, layer_scale_init) if layer_scale_init > 0.0 else nn.Identity()
        self.ls2 = LayerScale(dim, layer_scale_init) if layer_scale_init > 0.0 else nn.Identity()

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
        x = x + self.drop_path(self.ls1(self.attn(self.norm1(x), cos, sin)))
        x = x + self.drop_path(self.ls2(self.mlp(self.norm2(x))))
        return x


# =============================================================================
#  Semantic encoder
# =============================================================================


class SemanticEncoder(nn.Module):
    """Semantic encoder (ViT with RoPE) for MotifAudio inference."""

    def __init__(self, config: dict):
        super().__init__()
        self.embed_dim = config["sem_dim"]
        enc_config = config["semantic_encoder"]
        self.use_cls_token = enc_config["use_cls_token"]
        self.num_register_tokens = enc_config["num_register_tokens"]
        self.num_heads = enc_config["num_heads"]
        self.depth = enc_config["depth"]
        self.mlp_ratio = enc_config["mlp_ratio"]
        self.rope_theta = enc_config["rope_theta"]
        self.use_rope = enc_config["use_rope"]
        self.layer_scale_init = enc_config["layer_scale_init"]
        self._n_prefix = (1 if self.use_cls_token else 0) + self.num_register_tokens

        img_size = tuple(enc_config["img_size"])
        patch_size = tuple(enc_config["patch_size"])
        in_chans = enc_config["in_chans"]

        if PatchEmbed is not None:
            self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, self.embed_dim, strict_img_size=False)
        else:
            raise RuntimeError("timm is required for SemanticEncoder")

        if self.use_cls_token:
            self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
        if self.num_register_tokens > 0:
            self.register_tokens = nn.Parameter(torch.zeros(1, self.num_register_tokens, self.embed_dim))

        if not self.use_rope:
            _h, _w = img_size
            _ph, _pw = patch_size
            _num_patches = (_h // _ph) * (_w // _pw)
            self.register_buffer("pos_embed", torch.zeros(1, _num_patches, self.embed_dim))

        self.blocks = nn.ModuleList([
            RoPEBlock(self.embed_dim, self.num_heads, self.mlp_ratio, qkv_bias=True,
                      norm_layer=nn.RMSNorm, layer_scale_init=self.layer_scale_init)
            for _ in range(self.depth)
        ])
        self.norm = nn.RMSNorm(self.embed_dim)

    def forward(self, lms: Tensor, gh: int, gw: int) -> Tensor:
        x = self.patch_embed(lms)
        B = x.shape[0]
        head_dim = x.shape[-1] // self.num_heads

        if not self.use_rope:
            x = x + self.pos_embed[:, :gh * gw, :].to(x.dtype)
            cos = torch.ones(B, 1, x.shape[1], head_dim, dtype=x.dtype, device=x.device)
            sin = torch.zeros(B, 1, x.shape[1], head_dim, dtype=x.dtype, device=x.device)
        else:
            cos, sin = axial_cos_sin(head_dim, gh, gw, self.rope_theta, x.device)
            cos = cos.unsqueeze(0).expand(B, -1, -1)
            sin = sin.unsqueeze(0).expand(B, -1, -1)

        if self._n_prefix > 0:
            prefix = []
            if self.use_cls_token:
                prefix.append(self.cls_token.expand(B, -1, -1))
            if self.num_register_tokens > 0:
                prefix.append(self.register_tokens.expand(B, -1, -1))
            x = torch.cat(prefix + [x], dim=1)
            cos, sin = prefix_pad(cos, sin, self._n_prefix)

        cos = cos.unsqueeze(1)
        sin = sin.unsqueeze(1)

        for blk in self.blocks:
            x = blk(x, cos, sin)
        x = self.norm(x)
        x = x[:, self._n_prefix:, :]
        return x


# =============================================================================
#  Acoustic encoder
# =============================================================================


class AcousticEncoder(nn.Module):
    def __init__(self, config: dict):
        super().__init__()
        self.n_mels = config["n_mels"]
        self.patch_f = config["patch_f"]
        self.patch_t = config["patch_t"]
        self.embed_dim = config["embed_dim"]
        self.norm_mean = config["norm_mean"]
        self.norm_std = config["norm_std"]

        self.patch_embed = nn.Conv2d(
            1, self.embed_dim,
            kernel_size=(self.patch_f, self.patch_t),
            stride=(self.patch_f, self.patch_t),
        )
        self.norm = nn.LayerNorm(self.embed_dim, eps=1e-6)

    def forward(self, mel: Tensor) -> Tensor:
        mel = mel.to(self.patch_embed.weight.dtype)
        mel = mel.clamp(min=1e-7).log()
        mel = (mel - self.norm_mean) / self.norm_std
        mel = mel.unsqueeze(1)
        x = self.patch_embed(mel)
        x = x.flatten(2).transpose(1, 2)
        x = self.norm(x)
        return x


# =============================================================================
#  Fusion (cross-attention)
# =============================================================================


class Fusion(nn.Module):
    def __init__(self, dim: int, num_heads: int = 8):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.norm = nn.RMSNorm(dim)

    def forward(self, sem: Tensor, acou: Tensor) -> Tensor:
        # The two streams may arrive in different dtypes; align them with the
        # module's parameters before attending.
        param_dtype = self.attn.in_proj_weight.dtype
        sem = sem.to(param_dtype)
        acou = acou.to(param_dtype)
        attn_out, _ = self.attn(acou, sem, sem)
        return self.norm(acou + attn_out)


# =============================================================================
#  Decoder sub-modules
# =============================================================================


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

    ``alpha`` (frequency) and ``beta`` (magnitude) are per-channel learned
    parameters stored in the log domain, so both are exponentiated here.
    """

    def __init__(self, in_features: int):
        super().__init__()
        self.in_features = in_features
        self.alpha = nn.Parameter(torch.zeros(in_features))
        self.beta = nn.Parameter(torch.zeros(in_features))
        self.no_div_by_zero = 0.000000001

    def forward(self, x: Tensor) -> Tensor:
        alpha = torch.exp(self.alpha.unsqueeze(0).unsqueeze(0))
        beta = torch.exp(self.beta.unsqueeze(0).unsqueeze(0))
        x = x + (1.0 / (beta + self.no_div_by_zero)) * (torch.sin(x * alpha) ** 2)
        return x


class ConvNeXtBlock(nn.Module):
    def __init__(self, dim: int, intermediate_dim: int, layer_scale_init_value: float):
        super().__init__()
        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim)
        self.norm = nn.RMSNorm(dim)
        self.pwconv1 = nn.Linear(dim, intermediate_dim)
        self.act = Snake(intermediate_dim)
        self.pwconv2 = nn.Linear(intermediate_dim, dim)
        self.gamma = (
            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
            if layer_scale_init_value > 0
            else None
        )

    def forward(self, x: Tensor) -> Tensor:
        residual = x
        x = self.dwconv(x)
        x = x.transpose(1, 2)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.transpose(1, 2)
        x = residual + x
        return x


class ISTFTHead(nn.Module):
    def __init__(self, dim: int, n_fft: int, hop_length: int):
        super().__init__()
        self.out = nn.Linear(dim, n_fft + 2)
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.register_buffer("window", torch.hann_window(n_fft))

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        x = self.out(x).transpose(1, 2)
        mag_pre, phase = x.chunk(2, dim=1)
        mag = torch.exp(mag_pre.float())
        mag = torch.clip(mag, max=1e2)
        S = mag * (torch.cos(phase.float()) + 1j * torch.sin(phase.float()))
        audio = torch.istft(S, self.n_fft, self.hop_length, self.n_fft, self.window, center=True)
        return audio.unsqueeze(1), mag_pre, phase, S.real, S.imag


class DecoderBackbone(nn.Module):
    def __init__(self, input_channels: int, dim: int, intermediate_dim: int,
                 num_layers: int, layer_scale_init_value: float | None = None):
        super().__init__()
        self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
        self.norm = nn.RMSNorm(dim)
        layer_scale_init_value = layer_scale_init_value or 1 / num_layers
        self.convnext = nn.ModuleList([
            ConvNeXtBlock(dim=dim, intermediate_dim=intermediate_dim,
                          layer_scale_init_value=layer_scale_init_value)
            for _ in range(num_layers)
        ])
        self.final_layer_norm = nn.RMSNorm(dim)

    def forward(self, x: Tensor) -> Tensor:
        x = self.embed(x)
        x = self.norm(x.transpose(1, 2))
        x = x.transpose(1, 2)
        for block in self.convnext:
            x = block(x)
        x = self.final_layer_norm(x.transpose(1, 2))
        return x


class Decoder(nn.Module):
    def __init__(self, input_channels: int, hidden_dim: int, intermediate_dim: int,
                 num_layers: int, n_fft: int, hop_length: int, upsample_tokens: int = 1):
        super().__init__()
        self.backbone = DecoderBackbone(input_channels, hidden_dim, intermediate_dim, num_layers)
        if upsample_tokens > 1:
            kernel_size = 7
            output_padding = 1 if (kernel_size - upsample_tokens) % 2 else 0
            padding = (kernel_size - upsample_tokens + output_padding) // 2
            self.upsampler = nn.ConvTranspose1d(
                input_channels, input_channels,
                kernel_size=kernel_size, stride=upsample_tokens,
                padding=padding, output_padding=output_padding,
            )
        else:
            self.upsampler = None
        self.head = ISTFTHead(dim=hidden_dim, n_fft=n_fft, hop_length=hop_length)

    def forward(self, z: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        if self.upsampler is not None:
            z = self.upsampler(z)
        x = self.backbone(z)
        return self.head(x)


# =============================================================================
#  MotifAudio top-level model
# =============================================================================


_LOG_EPS = 1e-5


class MotifAudio(ModelMixin, ConfigMixin):
    """Motif-Audio: dual-stream audio autoencoder.

    Encodes a 16 kHz waveform with a semantic and an acoustic stream, fuses
    them into a continuous latent, and reconstructs the waveform through
    an inverse STFT. Both mel spectrograms are computed internally.

    Config keys carry a ``_config`` suffix so they never collide with the
    submodule attribute names (``fusion``, ``decoder``).

    Args:
        sample_rate: Audio sample rate.
        semantic_config: Semantic encoder / mel config dict.
        acoustic_config: Acoustic encoder / mel config dict.
        fusion_config: Fusion module config dict.
        decoder_config: Decoder config dict.
        auto_map: Repo metadata used by ``diffusers.AutoModel`` to locate this
            class when loading with ``trust_remote_code=True``. Accepted and
            ignored so it can live in ``config.json`` without warnings.
    """

    config_name = "config.json"

    @register_to_config
    def __init__(
        self,
        sample_rate: int = 16000,
        *,
        semantic_config: dict,
        acoustic_config: dict,
        fusion_config: dict,
        decoder_config: dict,
        auto_map: dict | None = None,
    ):
        super().__init__()

        del auto_map  # repo metadata for AutoModel remote-code loading; not used here

        sem_cfg = semantic_config
        acou_cfg = acoustic_config
        dec_cfg = decoder_config
        fusion_cfg = fusion_config

        self.semantic_encoder = SemanticEncoder(sem_cfg)
        self.acoustic_encoder = AcousticEncoder(acou_cfg)
        self.fusion = Fusion(
            dim=fusion_cfg["dim"],
            num_heads=fusion_cfg["num_heads"],
        )
        self.decoder = Decoder(
            input_channels=dec_cfg["input_channels"],
            hidden_dim=dec_cfg["hidden_dim"],
            intermediate_dim=dec_cfg["intermediate_dim"],
            num_layers=dec_cfg["num_layers"],
            n_fft=dec_cfg["n_fft"],
            hop_length=dec_cfg["hop_length"],
            upsample_tokens=dec_cfg["upsample_tokens"],
        )

        # The mel filterbanks and windows are derived constants, not learned
        # weights. Build them in float32 regardless of the ambient default
        # dtype: constructing them under a bf16 default (as `from_pretrained`
        # does when passed `torch_dtype`) computes the filterbank itself in
        # bf16, which corrupts it and yields NaNs downstream.
        _default_dtype = torch.get_default_dtype()
        torch.set_default_dtype(torch.float32)
        try:
            self._build_mel_front_end(sample_rate, sem_cfg, acou_cfg)
        finally:
            torch.set_default_dtype(_default_dtype)

        self.sem_norm_mean = sem_cfg["norm_mean"]
        self.sem_norm_std = sem_cfg["norm_std"]
        self.sem_patch_f = sem_cfg["patch_f"]
        self.sem_patch_t = sem_cfg["patch_t"]

    def _build_mel_front_end(self, sample_rate: int, sem_cfg: dict, acou_cfg: dict) -> None:
        if torchaudio is not None:
            self.sem_mel = torchaudio.transforms.MelSpectrogram(
                sample_rate=sample_rate,
                n_fft=sem_cfg["n_fft"], win_length=sem_cfg["n_fft"],
                hop_length=sem_cfg["hop_length"],
                f_min=sem_cfg["f_min"], f_max=sem_cfg["f_max"],
                n_mels=sem_cfg["n_mels"], power=2.0, center=True, norm=None,
            )
            self.acou_mel = torchaudio.transforms.MelSpectrogram(
                sample_rate=sample_rate,
                n_fft=acou_cfg["n_fft"], win_length=acou_cfg["n_fft"],
                hop_length=acou_cfg["hop_length"],
                f_min=acou_cfg["f_min"], f_max=acou_cfg["f_max"],
                n_mels=acou_cfg["n_mels"], power=2.0, center=True, norm=None,
            )
        else:
            self.sem_mel = None
            self.acou_mel = None
            raise RuntimeError("torchaudio is required for MotifAudio")

    def encode_semantic(self, mel: Tensor) -> Tensor:
        lms = (mel + _LOG_EPS).log().unsqueeze(1)
        lms = (lms - self.sem_norm_mean) / self.sem_norm_std
        gh = lms.shape[2] // self.sem_patch_f
        gw = lms.shape[3] // self.sem_patch_t
        param = next(self.semantic_encoder.parameters())
        lms = lms.to(dtype=param.dtype)
        return self.semantic_encoder(lms, gh, gw)

    @torch.inference_mode()
    def forward(self, waveform: Tensor) -> dict[str, Tensor]:
        wav = waveform.squeeze(1).float()
        with torch.autocast(device_type=wav.device.type, enabled=False):
            mel_sem = self.sem_mel.float()(wav)
            mel_acou = self.acou_mel.float()(wav)

        sem = self.encode_semantic(mel_sem)
        acou = self.acoustic_encoder(mel_acou)
        z = self.fusion(sem, acou)
        y = self.decoder(z.transpose(1, 2))[0]

        # Inference exposes the reconstructed waveform plus the fused and
        # per-stream latents (usable as downstream representations). The
        # decoder's STFT intermediates are internal to reconstruction and are
        # not returned.
        return {
            "waveform": y,
            "z_fused": z,
            "z_sem": sem,
            "z_acou": acou,
        }