"""Local PyTorch MERT2 encoder for the unified tokenizer checkpoint."""

from typing import Dict, Sequence

import torch
import torch.nn.functional as F
from torch import nn
from torchaudio.transforms import AmplitudeToDB, MelScale, Spectrogram

SAMPLE_RATE = 24000
HOP = 960


class MelFrontend(nn.Module):
    def __init__(self, cfg: Dict):
        super().__init__()
        bins = cfg["num_mel_bins"]
        self.register_buffer("mel_mean", torch.zeros(bins))
        self.register_buffer("mel_std", torch.ones(bins))
        self.spectrogram = Spectrogram(
            n_fft=cfg["n_fft"],
            win_length=cfg["win_length"],
            hop_length=cfg["hop_length"],
            power=2.0,
        )
        self.mel_scale = MelScale(
            n_mels=bins,
            sample_rate=cfg["sampling_rate"],
            n_stft=cfg["n_fft"] // 2 + 1,
        )
        self.amplitude_to_db = AmplitudeToDB(stype="power", top_db=None)

    @torch.no_grad()
    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        with torch.autocast(device_type=waveform.device.type, enabled=False):
            mel = self.amplitude_to_db(self.mel_scale(self.spectrogram(waveform.float())))
            mel = mel[..., :-1].transpose(-1, -2)
            return (mel - self.mel_mean) / self.mel_std.clamp_min(1e-5)


class Transpose(nn.Module):
    def forward(self, x):
        return x.transpose(1, 2)


class GlobalResponseNorm(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(1, 1, dim))
        self.bias = nn.Parameter(torch.zeros(1, 1, dim))

    def forward(self, x):
        magnitude = torch.norm(x, p=2, dim=1, keepdim=True)
        return self.weight * (x * (magnitude / (magnitude.mean(dim=-1, keepdim=True) + 1e-6))) + self.bias + x


class ConvNextLayer(nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.depthwise_block = nn.Sequential(
            Transpose(),
            nn.Conv1d(dim, dim, 7, padding=3, groups=dim),
            Transpose(),
        )
        self.pointwise_block = nn.Sequential(
            nn.LayerNorm(dim, eps=eps),
            nn.Linear(dim, 4 * dim),
            nn.GELU(),
            GlobalResponseNorm(4 * dim),
            nn.Linear(4 * dim, dim),
        )

    def forward(self, x):
        return x + self.pointwise_block(self.depthwise_block(x))


class ConvNextBlock(nn.Module):
    def __init__(self, cin: int, cout: int, stride: int, depth: int, eps: float):
        super().__init__()
        self.resampling_layer = (
            nn.Sequential(
                nn.LayerNorm(cin, eps=eps),
                Transpose(),
                nn.Conv1d(cin, cout, 2, stride=stride),
                Transpose(),
            )
            if cin != cout or stride > 1
            else nn.Identity()
        )
        self.convnext_layers = nn.Sequential(*[ConvNextLayer(cout, eps) for _ in range(depth)])

    def forward(self, x):
        return self.convnext_layers(self.resampling_layer(x))


class RotaryEmbedding(nn.Module):
    def __init__(self, cfg: Dict):
        super().__init__()
        head_dim = cfg["hidden_size"] // cfg["num_attention_heads"]
        inv = 1.0 / (
            cfg["rotary_embedding_base"]
            ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
        )
        self.register_buffer("inv_freq", inv, persistent=False)

    def forward(self, x):
        positions = torch.arange(x.shape[1], device=x.device, dtype=torch.float32)
        angles = torch.einsum("i,j->ij", positions, self.inv_freq.float())
        angles = torch.cat((angles, angles), dim=-1)[None, :, None, :]
        return angles.cos().to(x.dtype), angles.sin().to(x.dtype)


def rotate_half(x):
    first, second = x.chunk(2, dim=-1)
    return torch.cat((-second, first), dim=-1)


class SelfAttention(nn.Module):
    def __init__(self, cfg: Dict, causal: bool):
        super().__init__()
        width = cfg["hidden_size"]
        self.num_heads = cfg["num_attention_heads"]
        self.head_dim = width // self.num_heads
        self.causal = causal
        self.query_proj = nn.Linear(width, width)
        self.key_proj = nn.Linear(width, width)
        self.value_proj = nn.Linear(width, width)
        self.out_proj = nn.Linear(width, width)

    def forward(self, x, rope):
        batch, frames, width = x.shape
        shape = (batch, frames, self.num_heads, self.head_dim)
        q, k, v = (
            projection(x).reshape(shape)
            for projection in (self.query_proj, self.key_proj, self.value_proj)
        )
        cos, sin = rope
        q = q * cos + rotate_half(q) * sin
        k = k * cos + rotate_half(k) * sin
        out = F.scaled_dot_product_attention(
            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=self.causal
        )
        return self.out_proj(out.transpose(1, 2).reshape(batch, frames, width))


class FeedForward(nn.Module):
    def __init__(self, cfg: Dict):
        super().__init__()
        self.w_1 = nn.Linear(cfg["hidden_size"], cfg["intermediate_size"])
        self.w_2 = nn.Linear(cfg["intermediate_size"], cfg["hidden_size"])

    def forward(self, x):
        return self.w_2(F.gelu(self.w_1(x)))


class ConvolutionModule(nn.Module):
    def __init__(self, cfg: Dict):
        super().__init__()
        width = cfg["hidden_size"]
        kernel = cfg["conv_depthwise_kernel_size"]
        eps = cfg["layer_norm_eps"]
        self.layer_norm = nn.LayerNorm(width, eps=eps)
        self.conv_block = nn.Sequential(
            Transpose(),
            nn.Conv1d(width, 2 * width, 1, bias=False),
            nn.GLU(dim=1),
            nn.Conv1d(
                width,
                width,
                kernel,
                padding=(kernel - 1) // 2,
                groups=width,
                bias=False,
            ),
            nn.Sequential(Transpose(), nn.LayerNorm(width, eps=eps), Transpose()),
            nn.GELU(),
            nn.Conv1d(width, width, 1, bias=False),
            Transpose(),
        )

    def forward(self, x):
        return self.conv_block(self.layer_norm(x))


class ConformerBlock(nn.Module):
    def __init__(self, cfg: Dict, causal: bool):
        super().__init__()
        width, eps = cfg["hidden_size"], cfg["layer_norm_eps"]
        self.ffn1_layer_norm = nn.LayerNorm(width, eps=eps)
        self.ffn1 = FeedForward(cfg)
        self.attn_layer_norm = nn.LayerNorm(width, eps=eps)
        self.attn = SelfAttention(cfg, causal)
        self.conv_module = ConvolutionModule(cfg)
        self.ffn2_layer_norm = nn.LayerNorm(width, eps=eps)
        self.ffn2 = FeedForward(cfg)
        self.final_layer_norm = nn.LayerNorm(width, eps=eps)

    def forward(self, x, rope):
        x = x + 0.5 * self.ffn1(self.ffn1_layer_norm(x))
        x = self.attn(self.attn_layer_norm(x), rope) + x
        x = self.conv_module(x) + x
        x = x + 0.5 * self.ffn2(self.ffn2_layer_norm(x))
        return self.final_layer_norm(x)


class MERT2Encoder(nn.Module):
    """Mono 24 kHz waveform to 25 Hz Conformer block outputs."""

    def __init__(self, cfg: Dict, causal: bool, layers: int):
        super().__init__()
        self.feature_extractor = MelFrontend(cfg)
        channels = [cfg["num_mel_bins"]] + list(cfg["subsampling_channels"])
        self.subsampling_module = nn.Sequential(
            *[
                ConvNextBlock(
                    channels[i],
                    channels[i + 1],
                    (1, 2, 2)[i],
                    cfg["subsampling_depths"][i],
                    cfg["subsampling_layer_norm_eps"],
                )
                for i in range(3)
            ]
        )
        self.layers = nn.ModuleList([ConformerBlock(cfg, causal) for _ in range(layers)])
        self.embed_positions = RotaryEmbedding(cfg)

    def forward(self, waveform: torch.Tensor, output: int) -> torch.Tensor:
        mel = self.feature_extractor(waveform)
        first = self.subsampling_module[0].convnext_layers[0].depthwise_block[1]
        x = self.subsampling_module(mel.to(first.weight.dtype))
        rope = self.embed_positions(x)
        for index, layer in enumerate(self.layers):
            x = layer(x, rope)
            if index == output:
                return x
        raise ValueError(f"Requested encoder block {output}, but only {len(self.layers)} are present")

    @torch.inference_mode()
    def features(self, mono24: torch.Tensor, output: int, chunk: float = 0) -> torch.Tensor:
        """Return fp32 [frames, 1024] features, optionally in fixed-second pieces."""
        if mono24.ndim != 1 or mono24.shape[0] < 1025:
            raise ValueError("The encoder needs at least 1025 mono samples at 24 kHz")
        with torch.autocast(
            device_type=mono24.device.type,
            dtype=torch.bfloat16,
            enabled=mono24.device.type == "cuda",
        ):
            if not chunk:
                return self(mono24[None], output)[0].float()
            samples = int(SAMPLE_RATE * chunk)
            pieces = [
                piece
                for piece in (mono24[start : start + samples] for start in range(0, len(mono24), samples))
                if len(piece) >= SAMPLE_RATE
            ]
            if not pieces:
                raise ValueError(f"Audio is shorter than the minimum one-second {chunk}-second chunk")
            full = [piece for piece in pieces if len(piece) == samples]
            groups = ([torch.stack(full)] if full else []) + [piece[None] for piece in pieces if len(piece) < samples]
            parts = [self(group, output) for group in groups]
        frames = int(round(len(mono24) / HOP))
        joined = torch.cat([part.reshape(-1, part.shape[-1]) for part in parts]).float()
        return F.interpolate(
            joined.T[None], size=frames, mode="linear", align_corners=False
        )[0].T