"""TinyConformer-CTC model, matching the notebook-03 architecture exactly.

Presets:
  nano  (96,  5): ~1.4M params, ~1.4MB int8
  micro (112, 6): ~2.1M params, ~2.1MB int8
  tiny  (128, 6): ~2.8M params, ~2.8MB int8
  small (256, 9): ~16M params, ~16MB int8  ← v2 target

CTC (not RNN-T/seq2seq) so the decode stays trivial. QKV is explicit Linear layers
(not nn.MultiheadAttention) so int8 dynamic quantization covers every big matmul.
Pre-LN on every sublayer (ff1, attn, conv, ff2) for training stability at larger sizes.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class FFN(nn.Module):
    def __init__(self, d, mult=4, p=0.1):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(d, mult*d), nn.SiLU(), nn.Dropout(p),
                                 nn.Linear(mult*d, d), nn.Dropout(p))
    def forward(self, x): return self.net(x)


class MHSA(nn.Module):
    """Self-attention with explicit Linear QKV/O so int8 dynamic quant covers it."""
    def __init__(self, d, heads=4, p=0.1):
        super().__init__()
        assert d % heads == 0
        self.h, self.dh = heads, d // heads
        self.qkv = nn.Linear(d, 3*d)
        self.o = nn.Linear(d, d)
        self.drop = nn.Dropout(p)
    def forward(self, x, pad_mask=None):
        B, T, D = x.shape
        q, k, v = self.qkv(x).chunk(3, -1)
        q = q.view(B, T, self.h, self.dh).transpose(1, 2)
        k = k.view(B, T, self.h, self.dh).transpose(1, 2)
        v = v.view(B, T, self.h, self.dh).transpose(1, 2)
        s = q @ k.transpose(-2, -1) / math.sqrt(self.dh)
        if pad_mask is not None:
            s = s.masked_fill(pad_mask[:, None, None, :], float("-inf"))
        w = self.drop(s.softmax(-1))
        return self.o((w @ v).transpose(1, 2).reshape(B, T, D))


class ConvModule(nn.Module):
    def __init__(self, d, k=15, p=0.1):
        super().__init__()
        self.pw1 = nn.Conv1d(d, 2*d, 1)
        self.dw  = nn.Conv1d(d, d, k, padding=k//2, groups=d)
        self.bn  = nn.BatchNorm1d(d)
        self.pw2 = nn.Conv1d(d, d, 1)
        self.drop = nn.Dropout(p)
    def forward(self, x):  # (B,T,D)
        x = x.transpose(1, 2)
        x = self.pw1(x)
        a, b = x.chunk(2, 1)
        x = a * torch.sigmoid(b)
        x = self.bn(self.dw(x))
        x = self.drop(self.pw2(F.silu(x)))
        return x.transpose(1, 2)


class ConformerBlock(nn.Module):
    def __init__(self, d=112, heads=4, p=0.1):
        super().__init__()
        self.ff1_norm  = nn.LayerNorm(d)
        self.ff1       = FFN(d, p=p)
        self.attn_norm = nn.LayerNorm(d)
        self.attn      = MHSA(d, heads, p)
        self.conv_norm = nn.LayerNorm(d)
        self.conv      = ConvModule(d, p=p)
        self.ff2_norm  = nn.LayerNorm(d)
        self.ff2       = FFN(d, p=p)
        self.final     = nn.LayerNorm(d)

    def forward(self, x, pad_mask=None):
        x = x + 0.5 * self.ff1(self.ff1_norm(x))
        x = x + self.attn(self.attn_norm(x), pad_mask)
        x = x + self.conv(self.conv_norm(x))
        x = x + 0.5 * self.ff2(self.ff2_norm(x))
        return self.final(x)


class TinyConformerCTC(nn.Module):
    N_MELS = 80
    def __init__(self, d=112, n_blocks=6, n_classes=29, p=0.1):
        super().__init__()
        self.c1     = nn.Conv2d(1, d, 3, stride=2, padding=1)
        self.c2     = nn.Conv2d(d, d, 3, stride=2, padding=1)
        self.proj   = nn.Linear(d * (self.N_MELS // 4), d)
        self.blocks = nn.ModuleList([ConformerBlock(d, p=p) for _ in range(n_blocks)])
        self.head   = nn.Linear(d, n_classes)

    @staticmethod
    def out_len(l):  # two stride-2 k3p1 convs
        return ((l + 2 - 3)//2 + 1 + 2 - 3)//2 + 1

    def forward(self, mel, lens):
        x = F.relu(self.c1(mel)); x = F.relu(self.c2(x))   # (B,d,T/4,20)
        B, C, T, F_ = x.shape
        x = self.proj(x.permute(0, 2, 1, 3).reshape(B, T, C*F_))
        # tensor arithmetic keeps out_len ONNX-traceable (avoids baking a constant)
        ol = (lens.long().to(x.device) - 1) // 2 + 1
        ol = (ol - 1) // 2 + 1
        pad = torch.arange(T, device=x.device)[None, :] >= ol[:, None]
        for blk in self.blocks:
            x = blk(x, pad)
        return self.head(x).log_softmax(-1), ol


def count_params(m):
    n = sum(p.numel() for p in m.parameters())
    print(f"params: {n/1e6:.2f}M | fp32 {n*4/1e6:.1f}MB | int8 ~{n/1e6:.1f}MB (linears)")
    return n


def greedy_decode(log_probs, alphabet, blank=0):
    """log_probs: (T, C) -> (text, confidence). Collapses repeats, drops blanks."""
    probs = log_probs.exp()
    ids = probs.argmax(-1).tolist()
    conf = probs.amax(-1)
    out, prev, cs = [], None, []
    for t, i in enumerate(ids):
        if i != blank:
            cs.append(float(conf[t]))
            if i != prev:
                out.append(alphabet[i])
        prev = i
    text = "".join(out)
    return text, (sum(cs)/len(cs) if cs else 0.0)


# confidence < threshold -> "uncertain" (tongue's tie-reporting analog)
UNCERTAIN_TH = 0.35