"""Audio front-end for DiffusionGemma.

Two designs were tried:
  * encoder-free raw-waveform projection (Gemma-4-12B-Unified style) — FAILED to
    ground on content with a frozen/LoRA backbone: a shallow projector can't do
    acoustic feature extraction, and a frozen LLM (never trained on audio) can't
    either. WER plateaued at ~94% (fluent but audio-ignoring).
  * frozen Whisper encoder + lightweight projector (this file) — the encoder
    supplies acoustic→linguistic features the frozen LLM can actually read; only
    the projector (+ decoder LoRA) trains. This is the working path.

The Whisper encoder (open, MIT, frozen) is a feature extractor, NOT a decoder —
the diffusion LLM still performs recognition via its denoising mechanism.
"""
from __future__ import annotations
import torch
import torch.nn as nn

SAMPLE_RATE = 16000
WHISPER_HOP = 160              # 10 ms mel hop @ 16 kHz
WHISPER_MEL_FRAMES = 3000     # Whisper always pads/trims to 30 s
WHISPER_ENC_FRAMES = 1500     # encoder downsamples mel by 2


def audio_out_len(n_frames: int, subsample_factor: int) -> int:
    """Audio-token count the projector emits for `n_frames` encoder frames.

    Shared by projector and collator so the number of <audio> placeholders always
    matches the number of projected embeddings.
    """
    n_conv = subsample_factor.bit_length() - 1  # log2
    t = n_frames
    for _ in range(n_conv):
        t = (t - 1) // 2 + 1  # Conv1d L_out, k=3 s=2 p=1
    return t


def real_encoder_frames(num_samples: int) -> int:
    """Whisper encoder frames that correspond to real (non-silence-pad) audio."""
    mel = min(WHISPER_MEL_FRAMES, -(-num_samples // WHISPER_HOP))  # ceil
    return min(WHISPER_ENC_FRAMES, mel // 2)


class AudioProjector(nn.Module):
    """Lightweight projector: Whisper features [B, T_in, in_dim] -> [B, T, d_model].

    Conv1d stride-2 stack downsamples Whisper's 50 fps to ~50/subsample_factor fps
    (factor 8 -> 6.25 tok/s, 1500 enc frames -> 188 audio tokens for a 30 s window).
    """

    def __init__(self, d_model: int = 2816, in_dim: int = 768, hidden: int = 1280,
                 subsample_factor: int = 8, dropout: float = 0.0):
        super().__init__()
        assert subsample_factor >= 1 and (subsample_factor & (subsample_factor - 1)) == 0
        self.in_dim = in_dim
        self.subsample_factor = subsample_factor
        self.in_proj = nn.Linear(in_dim, hidden)
        n_conv = subsample_factor.bit_length() - 1
        self.convs = nn.ModuleList(
            [nn.Conv1d(hidden, hidden, kernel_size=3, stride=2, padding=1) for _ in range(n_conv)]
        )
        self.act = nn.GELU()
        self.drop = nn.Dropout(dropout)
        self.out_proj = nn.Linear(hidden, d_model)
        self.out_norm = nn.LayerNorm(d_model)

    def out_len(self, n_frames: int) -> int:
        return audio_out_len(n_frames, self.subsample_factor)

    def forward(self, feats: torch.Tensor) -> torch.Tensor:
        """feats: [B, T_in, in_dim] Whisper encoder states -> [B, T, d_model]."""
        x = self.act(self.in_proj(feats))    # [B, T_in, H]
        x = x.transpose(1, 2)                # [B, H, T_in]
        for conv in self.convs:
            x = self.act(conv(x))            # downsample by 2
        x = x.transpose(1, 2)                # [B, T, H]
        x = self.drop(x)
        return self.out_norm(self.out_proj(x))