"""
src/model.py

Non-Autoregressive On-Device CTC Reader Architecture for EmbeddingGemma-2.
Maps frozen 300M Audio Encoder representations (MRL 256-dim) to the shared multimodal
vector space, matching fixed text embeddings via cosine similarity projection.
"""

from __future__ import annotations

import math
import random
from typing import Optional, Tuple, Dict, Any

import torch
import torch.nn as nn
import torch.nn.functional as F

# Standard PyTorch Execution Mode
MetalFusedConvMidFunction = None
MetalFusedQKVLinearFunction = None
MetalFusedResL2NormFunction = None


# ==============================================================================
# 1. Conformer / Transformer Building Blocks (Mobile-Friendly Standard Ops)
# ==============================================================================

class DepthwiseConv1d(nn.Module):
    """
    Depthwise separable 1D convolution block with GLU and BatchNorm,
    providing temporal context modeling without heavy parameter overhead.
    Designed with standard PyTorch ops for seamless CoreML / ONNX export.
    """
    def __init__(self, channels: int, kernel_size: int = 15, dropout: float = 0.1):
        super().__init__()
        assert kernel_size % 2 == 1, "Kernel size must be odd for same-padding"
        self.padding = kernel_size // 2

        self.pointwise_in = nn.Conv1d(channels, channels * 2, kernel_size=1)
        self.act_glu = nn.GLU(dim=1)
        self.depthwise = nn.Conv1d(
            channels, channels,
            kernel_size=kernel_size,
            padding=self.padding,
            groups=channels,
            bias=False
        )
        self.bn = nn.BatchNorm1d(channels)
        self.act_swish = nn.SiLU()
        self.pointwise_out = nn.Conv1d(channels, channels, kernel_size=1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Fast-path on Apple Silicon MPS with bfloat16: F.linear pointwise + Fused Metal GLU + DWConv15 + BN + SiLU
        if (
            x.device.type == "mps"
            and x.dtype == torch.bfloat16
            and self.depthwise.kernel_size[0] == 15
            and MetalFusedConvMidFunction is not None
            and not torch.jit.is_tracing()
            and not torch.onnx.is_in_onnx_export()
        ):
            res = x
            u = F.linear(x, self.pointwise_in.weight.squeeze(-1), self.pointwise_in.bias)
            y = MetalFusedConvMidFunction.apply(
                u,
                self.depthwise.weight,
                self.bn.weight,
                self.bn.bias,
                self.bn.running_mean,
                self.bn.running_var,
                self.training
            )
            y = F.linear(y, self.pointwise_out.weight.squeeze(-1), self.pointwise_out.bias)
            return res + self.dropout(y)

        # Standard mobile-exportable path: (Batch, Time, Channels) -> transpose to (Batch, Channels, Time)
        res = x
        x = x.transpose(1, 2)
        x = self.pointwise_in(x)
        x = self.act_glu(x)
        x = self.depthwise(x)
        x = self.bn(x)
        x = self.act_swish(x)
        x = self.pointwise_out(x)
        x = self.dropout(x)
        x = x.transpose(1, 2)
        return res + x


class ExplicitMultiHeadAttention(nn.Module):
    """
    Explicit Multi-Head Attention implemented using standard tensor operations
    (MatMul, Scale, Softmax, MatMul) without proprietary PyTorch C++ kernels.
    Guarantees 100% compatibility with ONNX export (opset 17) and CoreML conversion.
    """
    def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0):
        super().__init__()
        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)

        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None,
        attn_bias: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        b, t_q, _ = q.size()
        _, t_k, _ = k.size()

        # Fused QKV self-attention fast-path on MPS bfloat16
        if (
            q is k
            and k is v
            and q.device.type == "mps"
            and not torch.jit.is_tracing()
            and not torch.onnx.is_in_onnx_export()
        ):
            qkv_w_view = getattr(self, "_qkv_w_view", None)
            if qkv_w_view is not None and MetalFusedQKVLinearFunction is not None:
                if torch.is_grad_enabled() and q.requires_grad:
                    qkv_lin = MetalFusedQKVLinearFunction.apply(
                        q, qkv_w_view, self._qkv_b_view, self._qkv_w_grad_view, self._qkv_b_grad_view
                    )
                else:
                    qkv_lin = F.linear(q, qkv_w_view, self._qkv_b_view)
            else:
                w_qkv = torch.cat([self.q_proj.weight, self.k_proj.weight, self.v_proj.weight], dim=0)
                b_qkv = torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias], dim=0)
                qkv_lin = F.linear(q, w_qkv, b_qkv)
            qkv = qkv_lin.view(b, t_q, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
            q_proj, k_proj, v_proj = qkv[0], qkv[1], qkv[2]

            scores = torch.matmul(q_proj, k_proj.transpose(-2, -1)) * self.scale
            if attn_bias is not None:
                scores = scores + attn_bias
            elif key_padding_mask is not None:
                scores = scores.masked_fill(key_padding_mask[:, None, None, :], -10000.0)

            probs = F.softmax(scores.float(), dim=-1).to(dtype=q.dtype)
            if self.training and self.dropout.p > 0.0:
                probs = F.dropout(probs, p=self.dropout.p, training=True)
            context = torch.matmul(probs, v_proj).transpose(1, 2).reshape(b, t_q, self.embed_dim)
            return self.out_proj(context)

        q_proj = self.q_proj(q).view(b, t_q, self.num_heads, self.head_dim).transpose(1, 2)
        k_proj = self.k_proj(k).view(b, t_k, self.num_heads, self.head_dim).transpose(1, 2)
        v_proj = self.v_proj(v).view(b, t_k, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled dot-product attention using MPS-native fused SDPA kernel
        attn_mask = None
        if attn_bias is not None:
            attn_mask = attn_bias
        elif key_padding_mask is not None:
            # key_padding_mask: True for pad positions (b, t_k)
            # In SDPA boolean mask: True means attend (valid), False means mask out (pad)
            attn_mask = (~key_padding_mask).unsqueeze(1).unsqueeze(2)

        dropout_p = self.dropout.p if self.training else 0.0
        context = F.scaled_dot_product_attention(
            q_proj, k_proj, v_proj,
            attn_mask=attn_mask,
            dropout_p=dropout_p
        )  # (b, heads, t_q, head_dim)
        context = context.transpose(1, 2).contiguous().view(b, t_q, self.embed_dim)
        return self.out_proj(context)


class ConformerBlock(nn.Module):
    """
    A lightweight Conformer block combining:
    - Pre-norm Multi-Head Self-Attention (MHSA)
    - Depthwise Convolution for local speech features
    - Feed-Forward Network (FFN) with half-step Macaron-style scaling
    """
    def __init__(
        self,
        d_model: int = 256,
        nhead: int = 4,
        dim_feedforward: int = 1024,
        dropout: float = 0.1,
        conv_kernel_size: int = 15
    ):
        super().__init__()
        # 1. First FFN (Macaron style)
        self.ffn1_norm = nn.LayerNorm(d_model)
        self.ffn1 = nn.Sequential(
            nn.Linear(d_model, dim_feedforward),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, d_model),
            nn.Dropout(dropout)
        )

        # 2. Multi-Head Self-Attention (Explicit ops for mobile deployment)
        self.attn_norm = nn.LayerNorm(d_model)
        self.self_attn = ExplicitMultiHeadAttention(
            embed_dim=d_model,
            num_heads=nhead,
            dropout=dropout
        )
        self.attn_dropout = nn.Dropout(dropout)

        # 3. Convolution Module
        self.conv_norm = nn.LayerNorm(d_model)
        self.conv_module = DepthwiseConv1d(
            channels=d_model,
            kernel_size=conv_kernel_size,
            dropout=dropout
        )

        # 4. Second FFN
        self.ffn2_norm = nn.LayerNorm(d_model)
        self.ffn2 = nn.Sequential(
            nn.Linear(d_model, dim_feedforward),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, d_model),
            nn.Dropout(dropout)
        )

        self.final_norm = nn.LayerNorm(d_model)

    def forward(
        self,
        x: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None,
        attn_bias: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        # FFN 1 with 0.5 residual scale
        x = x + 0.5 * self.ffn1(self.ffn1_norm(x))

        # Multi-Head Attention
        normed = self.attn_norm(x)
        attn_out = self.self_attn(
            normed, normed, normed,
            key_padding_mask=key_padding_mask,
            attn_bias=attn_bias
        )
        x = x + self.attn_dropout(attn_out)

        # Convolution Module
        x = x + self.conv_module(self.conv_norm(x))

        # FFN 2 with 0.5 residual scale
        x = x + 0.5 * self.ffn2(self.ffn2_norm(x))

        return self.final_norm(x)


# ==============================================================================
# 2. Audio-to-Vector Reader Module
# ==============================================================================

class AudioToVectorReader(nn.Module):
    """
    Reader v2 Conformer Architecture (~21.5M parameters).
    Input: Audio Encoder features of shape (Batch, Time, input_dim=256)
    Internal: Conformer latent space (d_model=384, nhead=6, num_layers=10, dim_feedforward=1536)
    Output: Normalized vectors of shape (Batch, Time, output_dim=256) aligned with EmbeddingGemma-2 text space.
    """
    def __init__(
        self,
        input_dim: int = 256,
        output_dim: int = 256,
        d_model: int = 384,
        nhead: int = 6,
        num_layers: int = 10,
        dim_feedforward: int = 1536,
        dropout: float = 0.1,
        conv_kernel_size: int = 15
    ):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.d_model = d_model

        self.input_proj = nn.Linear(input_dim, d_model)
        self.pos_emb = PositionalEncoding(d_model, max_len=8192)

        self.layers = nn.ModuleList([
            ConformerBlock(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
                conv_kernel_size=conv_kernel_size
            )
            for _ in range(num_layers)
        ])

        self.output_proj = nn.Sequential(
            nn.Linear(d_model, output_dim),
            nn.LayerNorm(output_dim)
        )
        # Small std=0.02 normal initialization to enable smooth early gradient flow
        nn.init.normal_(self.output_proj[0].weight, std=0.02)
        nn.init.zeros_(self.output_proj[0].bias)

        # Learnable global residual scale to balance acoustic adaptation and pretrained alignment
        self.res_scale = nn.Parameter(torch.tensor(0.5, dtype=torch.float32))

    def forward(
        self,
        x: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Args:
            x: Tensor of shape (Batch, Time, input_dim)
            key_padding_mask: Optional bool tensor (Batch, Time) where True denotes padded positions
        Returns:
            L2-normalized vectors (Batch, Time, output_dim)
        """
        res = x
        h = self.input_proj(x)
        h = self.pos_emb(h)

        # Precompute additive attention mask once across all Conformer layers
        attn_bias: Optional[torch.Tensor] = None
        if key_padding_mask is not None:
            attn_bias = key_padding_mask[:, None, None, :].to(dtype=h.dtype) * -10000.0

        for layer in self.layers:
            h = layer(h, attn_bias=attn_bias)

        out = self.output_proj(h)

        # Fused Metal Residual Scale + Mask + FP32-Accumulated L2 Normalization fast-path
        if (
            out.device.type == "mps"
            and out.dtype == torch.bfloat16
            and res.dtype == torch.bfloat16
            and res.shape[-1] == out.shape[-1]
            and out.shape[-1] == 256
            and key_padding_mask is not None
            and MetalFusedResL2NormFunction is not None
            and not torch.jit.is_tracing()
            and not torch.onnx.is_in_onnx_export()
        ):
            return MetalFusedResL2NormFunction.apply(out, res, self.res_scale, key_padding_mask)

        # Global residual connection with learnable scale parameter
        if res.shape[-1] == out.shape[-1]:
            combined = self.res_scale * res + out
        else:
            combined = out

        # Zero out padding positions before normalization if key_padding_mask provided
        if key_padding_mask is not None:
            combined = combined.masked_fill(key_padding_mask.unsqueeze(-1), 0.0)

        # Use FP32 accumulation + safe eps=1e-5 to guarantee zero underflow/NaN in bfloat16
        normed = F.normalize(combined.float(), p=2, dim=-1, eps=1e-5).to(dtype=x.dtype)
        if key_padding_mask is not None:
            normed = normed.masked_fill(key_padding_mask.unsqueeze(-1), 0.0)
        return normed


class PositionalEncoding(nn.Module):
    """Sinusoidal Positional Encoding compatible with dynamic sequence lengths."""
    def __init__(self, d_model: int, max_len: int = 8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len].to(dtype=x.dtype, device=x.device)


# ==============================================================================
# 3. Cosine Similarity Matching & Blank Logit Projection Head
# ==============================================================================

class CosineCTCHead(nn.Module):
    """
    Unified Cosine Similarity Matching & Blank Head:
    - Learns a unit vector in the spherical space for <blank>: self.blank_vec
    - Combines [norm_blank, vector_bank] -> all_vectors: (V + 1, embedding_dim)
    - Computes single matrix multiplication: (embeddings @ all_vectors.T) * logit_scale
    - Applies learnable scalar bias offset to <blank> (index 0) and per-token vocab_bias (index 1..V)
    - Supports trainable vector bank mode for end-to-end representation adaptation into acoustic space
    """
    def __init__(
        self,
        embedding_dim: int = 256,
        init_temperature: float = 0.05,
        vector_bank: Optional[torch.Tensor] = None,
        trainable_vector_bank: bool = True
    ):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.trainable_vector_bank = trainable_vector_bank
        # logit_scale initialized to ln(1 / init_temperature), safely bounded in [-5.0, 5.2983]
        init_logit_scale = math.log(1.0 / max(init_temperature, 1e-4))
        init_logit_scale = max(min(init_logit_scale, 5.2983), -5.0)
        self.logit_scale = nn.Parameter(torch.tensor(init_logit_scale, dtype=torch.float32))

        # Learnable parameter vector in R^embedding_dim for <blank> token (Index 0)
        self.blank_vec = nn.Parameter(torch.randn(1, embedding_dim))
        nn.init.normal_(self.blank_vec, std=0.02)

        # Learnable scalar bias for <blank> token (Index 0)
        self.blank_bias = nn.Parameter(torch.zeros(1, dtype=torch.float32))

        # Vector bank buffer or parameter (V, embedding_dim) & token prior bias
        if vector_bank is not None:
            vb_f32 = vector_bank.to(torch.float32)
            norm_vb = F.normalize(vb_f32, p=2, dim=-1)
            if trainable_vector_bank:
                self.vector_bank = nn.Parameter(norm_vb)
            else:
                self.register_buffer("vector_bank", norm_vb, persistent=False)
            self.vocab_bias = nn.Parameter(torch.zeros(norm_vb.size(0), dtype=torch.float32))
        else:
            self.register_buffer("vector_bank", None, persistent=False)
            self.vocab_bias = nn.Parameter(torch.zeros(0, dtype=torch.float32))

    def set_vector_bank(self, vector_bank: torch.Tensor):
        """Sets or updates the target vector bank (V, embedding_dim)."""
        vb_f32 = vector_bank.to(torch.float32)
        norm_vb = F.normalize(vb_f32, p=2, dim=-1)
        if self.trainable_vector_bank:
            self.vector_bank = nn.Parameter(norm_vb)
        else:
            self.register_buffer("vector_bank", norm_vb, persistent=False)
        self.vocab_bias = nn.Parameter(torch.zeros(norm_vb.size(0), dtype=torch.float32))

    def forward(
        self,
        embeddings: torch.Tensor,
        external_vector_bank: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Args:
            embeddings: Normalized Reader output (Batch, Time, embedding_dim)
            external_vector_bank: Optional runtime vector bank (V, embedding_dim)
        Returns:
            Logits of shape (Batch, Time, V + 1), where index 0 is <blank>
        """
        vb = external_vector_bank if external_vector_bank is not None else self.vector_bank
        if vb is None:
            raise ValueError("Vector bank must be provided either in constructor, set_vector_bank(), or as an argument.")

        # On MPS, execute the large (B*T, 256) x (256, V+1) projection strictly in bfloat16
        # (with FP32 L2-normalization and clamped scale) to prevent MPS float32/float16 NaN bugs.
        compute_dtype = torch.bfloat16 if embeddings.device.type == "mps" else torch.float32

        emb_f32 = embeddings.to(dtype=torch.float32)
        norm_blank = F.normalize(self.blank_vec.to(device=embeddings.device, dtype=torch.float32), p=2, dim=-1, eps=1e-5)

        if external_vector_bank is not None or self.trainable_vector_bank:
            vb_matched = vb.to(device=embeddings.device, dtype=torch.float32)
            norm_vb = F.normalize(vb_matched, p=2, dim=-1, eps=1e-5)
        else:
            # Frozen vector_bank is already float32 and L2-normalized in __init__ / set_vector_bank
            norm_vb = vb if (vb.device == embeddings.device and vb.dtype == torch.float32) else vb.to(device=embeddings.device, dtype=torch.float32)

        all_vectors = torch.cat([norm_blank, norm_vb], dim=0).to(dtype=compute_dtype)  # (V + 1, embedding_dim)

        # Safely clamp logit_scale parameter directly BEFORE exp to prevent math overflow
        # ln(200) = ~5.2983. Clamping between -5.0 and 5.2983 bounds scale in [0.0067, 200.0]
        clamped_scale_log = torch.clamp(self.logit_scale, min=-5.0, max=5.2983)
        scale = torch.exp(clamped_scale_log)
        scaled_emb = (emb_f32 * scale).to(dtype=compute_dtype)

        # Fused linear projection + per-token prior bias [blank_bias, vocab_bias]
        blank_b = self.blank_bias.to(device=embeddings.device, dtype=compute_dtype)
        if self.vocab_bias is not None and self.vocab_bias.size(0) == norm_vb.size(0):
            vocab_b = self.vocab_bias.to(device=embeddings.device, dtype=compute_dtype)
        else:
            vocab_b = torch.zeros(norm_vb.size(0), device=embeddings.device, dtype=compute_dtype)
        bias = torch.cat([blank_b, vocab_b], dim=0)

        return F.linear(scaled_emb, all_vectors, bias)


# ==============================================================================
# 4. SpecAugment for 256-dim Conformer Features
# ==============================================================================

class SpecAugment1D(nn.Module):
    """
    Time & Feature Masking for 256-dim feature sequences:
    - Time Masking: 2 masks, random length between 5 and 15 frames
    - Feature Masking: 2 masks, random length between 16 and 32 dimensions
    Applied strictly during model.train() mode.
    """
    def __init__(
        self,
        time_mask_min: int = 5,
        time_mask_max: int = 15,
        num_time_masks: int = 2,
        freq_mask_min: int = 16,
        freq_mask_max: int = 32,
        num_freq_masks: int = 2
    ):
        super().__init__()
        self.time_mask_min = time_mask_min
        self.time_mask_max = time_mask_max
        self.num_time_masks = num_time_masks
        self.freq_mask_min = freq_mask_min
        self.freq_mask_max = freq_mask_max
        self.num_freq_masks = num_freq_masks

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.training:
            return x

        B, T, D = x.shape
        x_aug = x.clone()

        # Time masking (2 times, 5~15 frames)
        for _ in range(self.num_time_masks):
            t_len = random.randint(self.time_mask_min, min(self.time_mask_max, max(T - 1, self.time_mask_min)))
            t_start = random.randint(0, max(0, T - t_len))
            x_aug[:, t_start : t_start + t_len, :] = 0.0

        # Feature masking (2 times, 16~32 dims)
        for _ in range(self.num_freq_masks):
            f_len = random.randint(self.freq_mask_min, min(self.freq_mask_max, max(D - 1, self.freq_mask_min)))
            f_start = random.randint(0, max(0, D - f_len))
            x_aug[:, :, f_start : f_start + f_len] = 0.0

        return x_aug


# ==============================================================================
# 5. Complete Trainable Reader System
# ==============================================================================

class EmbeddingGemmaCTCReader(nn.Module):
    """
    Combined Module for training:
    AudioToVectorReader + CosineCTCHead with SpecAugment and Trainable Vector Bank options.
    """
    def __init__(
        self,
        input_dim: int = 256,
        output_dim: int = 256,
        d_model: int = 384,
        nhead: int = 6,
        num_layers: int = 10,
        dim_feedforward: int = 1536,
        dropout: float = 0.1,
        conv_kernel_size: int = 15,
        init_temperature: float = 0.05,
        vector_bank: Optional[torch.Tensor] = None,
        use_spec_augment: bool = True,
        trainable_vector_bank: bool = True
    ):
        super().__init__()
        self.reader = AudioToVectorReader(
            input_dim=input_dim,
            output_dim=output_dim,
            d_model=d_model,
            nhead=nhead,
            num_layers=num_layers,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            conv_kernel_size=conv_kernel_size
        )
        self.ctc_head = CosineCTCHead(
            embedding_dim=output_dim,
            init_temperature=init_temperature,
            vector_bank=vector_bank,
            trainable_vector_bank=trainable_vector_bank
        )
        self.spec_aug = SpecAugment1D(
            time_mask_min=5,
            time_mask_max=15,
            num_time_masks=2,
            freq_mask_min=16,
            freq_mask_max=32,
            num_freq_masks=2
        ) if use_spec_augment else None

    def enable_mps_bfloat16_fastpath(self) -> None:
        """
        Configures Reader Conformer layers in native bfloat16 while preserving float32
        BatchNorm1d parameters/stats, res_scale, and CosineCTCHead parameters.
        Eliminates 200+ per-step autocast weight-casting kernels on Apple Silicon MPS.
        """
        self.reader.to(dtype=torch.bfloat16)
        for layer in self.reader.layers:
            layer.conv_module.bn.to(dtype=torch.float32)
        self.reader.res_scale.data = self.reader.res_scale.data.to(dtype=torch.float32)
        self.ctc_head.to(dtype=torch.float32)

    def forward_reader(
        self,
        features: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Runs SpecAugment (if training) + AudioToVectorReader without CosineCTCHead."""
        if self.training and self.spec_aug is not None:
            features = self.spec_aug(features)
        return self.reader(features, key_padding_mask=key_padding_mask)

    def forward(
        self,
        features: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None,
        external_vector_bank: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Args:
            features: Cached audio features (Batch, Time, 256)
            key_padding_mask: Padding mask (Batch, Time)
            external_vector_bank: Optional custom vector bank (V, 256)
        Returns:
            Logits: (Batch, Time, V + 1)
        """
        if self.training and self.spec_aug is not None:
            features = self.spec_aug(features)

        embeddings = self.reader(features, key_padding_mask=key_padding_mask)
        logits = self.ctc_head(embeddings, external_vector_bank=external_vector_bank)
        return logits

    def extract_reader_embeddings(
        self,
        features: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Outputs normalized embeddings and blank logit separately
        (Used for Mobile OTA / Plug-and-Play Dictionary mode).
        """
        embeddings = self.reader(features, key_padding_mask=key_padding_mask)
        norm_blank = F.normalize(self.ctc_head.blank_vec.to(device=embeddings.device), p=2, dim=-1)
        scale = torch.clamp(self.ctc_head.logit_scale.exp(), max=200.0)
        blank_logit = torch.matmul(embeddings, norm_blank.t()) * scale + self.ctc_head.blank_bias
        return embeddings, blank_logit

    def count_parameters(self) -> Dict[str, int]:
        """Returns train parameter counts."""
        total = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        reader_params = sum(p.numel() for p in self.reader.parameters())
        head_params = sum(p.numel() for p in self.ctc_head.parameters())
        return {
            "total_params": total,
            "trainable_params": trainable,
            "reader_params": reader_params,
            "ctc_head_params": head_params
        }


# ==============================================================================
# 5. Native Fallback Audio Encoder (Matches EmbeddingGemma-2 Audio Front-End)
# ==============================================================================

class StandardTransformerBlock(nn.Module):
    """Clean, standard Transformer block without proprietary fastpath C++ kernels."""
    def __init__(self, d_model: int = 1024, nhead: int = 8, dim_feedforward: int = 2048, dropout: float = 0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.self_attn = ExplicitMultiHeadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, dim_feedforward),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, d_model),
            nn.Dropout(dropout)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        nx = self.norm1(x)
        attn_out = self.self_attn(nx, nx, nx)
        x = x + attn_out
        x = x + self.ffn(self.norm2(x))
        return x


class RealEmbeddingGemma2AudioEncoder(nn.Module):
    """
    Official 305M Audio Encoder from google/embeddinggemma-2:
    - audio_tower: 304.8M Gemma 4 Audio Conformer Encoder
    - embed_audio: 0.8M Multimodal Projector
    - Matryoshka Slicing to 256 dimensions + L2 Normalization
    Strictly uses real pretrained weights; raises RuntimeError on failure.
    """
    def __init__(self, audio_tower: nn.Module, embed_audio: nn.Module):
        super().__init__()
        self.audio_tower = audio_tower
        self.embed_audio = embed_audio

    @classmethod
    def from_pretrained(
        cls,
        model_name: str = "google/embeddinggemma-2",
        dtype: torch.dtype = torch.bfloat16
    ) -> "RealEmbeddingGemma2AudioEncoder":
        """Loads and freezes the genuine 300M Audio Encoder from Hugging Face checkpoint."""
        try:
            from transformers import AutoModel
            full_model = AutoModel.from_pretrained(model_name, dtype=dtype)
            audio_tower = full_model.audio_tower
            embed_audio = full_model.embed_audio

            # Freeze all parameters
            for p in audio_tower.parameters():
                p.requires_grad = False
            for p in embed_audio.parameters():
                p.requires_grad = False

            encoder = cls(audio_tower=audio_tower, embed_audio=embed_audio)
            encoder.eval()
            return encoder
        except Exception as e:
            raise RuntimeError(
                f"[FATAL ERROR] Failed to load genuine '{model_name}' Audio Encoder: {e}\n"
                f"Make sure transformers is up-to-date: pip install git+https://github.com/huggingface/transformers.git"
            )

    def forward(self, input_features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            input_features: (Batch, Time_Frames, 128) 128-bin log-mel filterbanks
        Returns:
            (Batch, Time_Tokens, 256) MRL-sliced, L2-normalized representations
        """
        tower_out = self.audio_tower(input_features, attention_mask=None)
        audio_hs = tower_out.last_hidden_state  # (B, T, 1536)
        audio_proj = self.embed_audio(audio_hs)  # (B, T, 512)
        mrl_256 = audio_proj[:, :, :256]
        return F.normalize(mrl_256.float(), p=2, dim=-1)


# Quick Self-Test when executed directly
if __name__ == "__main__":
    print("[Testing model.py - Reader v2 Conformer]")
    batch_size = 2
    time_steps = 150
    input_dim = 256
    output_dim = 256
    vocab_size = 5999  # Sliced vocab size (5,999 non-blank tokens -> 6,000 total with <blank>)

    dummy_features = torch.randn(batch_size, time_steps, input_dim)
    dummy_vector_bank = torch.randn(vocab_size, output_dim)

    model = EmbeddingGemmaCTCReader(
        input_dim=input_dim,
        output_dim=output_dim,
        d_model=384,
        nhead=6,
        num_layers=10,
        dim_feedforward=1536,
        conv_kernel_size=15,
        init_temperature=0.05,
        vector_bank=dummy_vector_bank,
        trainable_vector_bank=True
    )

    stats = model.count_parameters()
    print("Parameter Stats:")
    for k, v in stats.items():
        print(f"  {k}: {v:,} ({v / 1e6:.2f}M)")

    logits = model(dummy_features)
    print(f"Input features shape: {dummy_features.shape}")
    print(f"Logits output shape:  {logits.shape} (Expected: ({batch_size}, {time_steps}, {vocab_size + 1}))")
    assert logits.shape == (batch_size, time_steps, vocab_size + 1)
    print("Model test PASSED successfully!")