# ==============================================================================
# model.py — Conformer ASR for Burmese (Myanmar)
# Standalone inference-ready architecture.
# Includes open-vocab support via ConformerASR.resize_vocab().
# ==============================================================================
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path


# ==============================================================================
# Architecture constants (matching training)
# ==============================================================================
D_MODEL      = 256
N_HEADS      = 4
N_LAYERS     = 12
FF_EXPANSION = 4
CONV_KERNEL  = 31
DROPOUT      = 0.1
DEC_LAYERS   = 2
N_MELS       = 80

# Special token ids
BLANK_ID = 0
UNK_ID   = 1
SOS_ID   = BLANK_ID
EOS_ID   = BLANK_ID
PAD_ID   = UNK_ID


def calc_subsampled_lens(lens):
    """Two Conv2d(k=3, s=2, p=1) layers: floor((floor((L+1)/2)+1)/2)."""
    if not isinstance(lens, torch.Tensor):
        lens = torch.tensor(lens, dtype=torch.long)
    l1 = torch.div(lens + 1, 2, rounding_mode="floor")
    l2 = torch.div(l1 + 1, 2, rounding_mode="floor")
    return torch.clamp(l2, min=1)


# ==============================================================================
# RoPE
# ==============================================================================
class RoPE(nn.Module):
    def __init__(self, dim, base=10000.0):
        super().__init__()
        inv = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv)

    def _rotate_half(self, x):
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)

    def forward(self, q, k):
        T = q.shape[2]
        t = torch.arange(T, device=q.device, dtype=torch.float32)
        freqs = torch.outer(t, self.inv_freq.float())
        emb = torch.cat([freqs, freqs], dim=-1)
        cos = emb.cos()[None, None, :, :].to(q.dtype)
        sin = emb.sin()[None, None, :, :].to(q.dtype)
        return (q * cos + self._rotate_half(q) * sin,
                k * cos + self._rotate_half(k) * sin)


class AttentionWithRoPE(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)
        self.dropout_p = dropout
        self.rope = RoPE(self.head_dim)

    def forward(self, x, key_pad_mask=None):
        B, T, D = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim)
        q, k, v = qkv.unbind(dim=2)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        q, k = self.rope(q, k)
        attn_mask = None
        if key_pad_mask is not None:
            attn_mask = (~key_pad_mask)[:, None, None, :].expand(B, 1, T, T)
        out = F.scaled_dot_product_attention(
            q, k, v, attn_mask=attn_mask,
            dropout_p=self.dropout_p if self.training else 0.0)
        out = out.transpose(1, 2).reshape(B, T, D)
        return self.out(out)


class FeedForward(nn.Module):
    def __init__(self, d_model, expansion=4, dropout=0.1):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_model * expansion)
        self.fc2 = nn.Linear(d_model * expansion, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.dropout(self.fc2(F.silu(self.fc1(x))))


class ConformerConvModule(nn.Module):
    def __init__(self, d_model, kernel=31, dropout=0.1):
        super().__init__()
        self.pw1 = nn.Conv1d(d_model, 2 * d_model, 1)
        self.dw = nn.Conv1d(d_model, d_model, kernel,
                            padding=kernel // 2, groups=d_model)
        self.norm = nn.LayerNorm(d_model)
        self.pw2 = nn.Conv1d(d_model, d_model, 1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = x.transpose(1, 2)
        x = self.pw1(x)
        x = F.glu(x, dim=1)
        x = self.dw(x)
        x = x.transpose(1, 2)
        x = self.norm(x)
        x = x.transpose(1, 2)
        x = F.silu(x)
        x = self.pw2(x)
        x = self.dropout(x)
        return x.transpose(1, 2)


class ConformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, ff_exp=4, conv_k=31, dropout=0.1):
        super().__init__()
        self.ff1 = FeedForward(d_model, ff_exp, dropout)
        self.ff2 = FeedForward(d_model, ff_exp, dropout)
        self.norm_ff1 = nn.LayerNorm(d_model)
        self.norm_ff2 = nn.LayerNorm(d_model)
        self.norm_attn = nn.LayerNorm(d_model)
        self.norm_conv = nn.LayerNorm(d_model)
        self.norm_out = nn.LayerNorm(d_model)
        self.attn = AttentionWithRoPE(d_model, n_heads, dropout)
        self.conv = ConformerConvModule(d_model, conv_k, dropout)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, key_pad_mask=None):
        x = x + 0.5 * self.ff1(self.norm_ff1(x))
        x = x + self.attn(self.norm_attn(x), key_pad_mask=key_pad_mask)
        x = x + self.conv(self.norm_conv(x))
        x = x + 0.5 * self.ff2(self.norm_ff2(x))
        return self.norm_out(x)


class ConvSubsampling(nn.Module):
    def __init__(self, n_mels, d_model):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
        )
        self.proj = nn.Linear(64 * (n_mels // 4), d_model)

    def forward(self, x):
        x = x.unsqueeze(1)
        x = self.conv(x)
        b, c, t, f = x.shape
        x = x.transpose(1, 2).reshape(b, t, c * f)
        return self.proj(x)


class AttnDecoder(nn.Module):
    def __init__(self, d_model, n_heads, n_layers, vocab_size, dropout=0.1):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos = nn.Embedding(4096, d_model)
        layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=n_heads,
            dim_feedforward=d_model * 4, dropout=dropout,
            batch_first=True, norm_first=True)
        self.body = nn.TransformerDecoder(layer, num_layers=n_layers)
        self.fc = nn.Linear(d_model, vocab_size)
        self.d_model = d_model

    def forward(self, tgt_ids, memory, memory_key_pad_mask=None):
        B, T = tgt_ids.shape
        pos = torch.arange(T, device=tgt_ids.device).unsqueeze(0).expand(B, T)
        x = self.embed(tgt_ids) * math.sqrt(self.d_model) + self.pos(pos)
        causal = torch.triu(torch.ones(T, T, device=tgt_ids.device,
                                       dtype=torch.bool), 1)
        out = self.body(x, memory,
                        tgt_mask=causal,
                        memory_key_padding_mask=memory_key_pad_mask)
        return self.fc(out)


class ConformerASR(nn.Module):
    def __init__(self, vocab_size=2566):
        super().__init__()
        self.frontend = ConvSubsampling(N_MELS, D_MODEL)
        self.blocks = nn.ModuleList([
            ConformerBlock(D_MODEL, N_HEADS, FF_EXPANSION,
                           CONV_KERNEL, DROPOUT)
            for _ in range(N_LAYERS)
        ])
        self.enc_norm = nn.LayerNorm(D_MODEL)
        self.ctc_fc = nn.Linear(D_MODEL, vocab_size)
        self.decoder = AttnDecoder(D_MODEL, N_HEADS, DEC_LAYERS,
                                   vocab_size, DROPOUT)

    # --------------------------------------------------------------------------
    # Open-vocab support
    # --------------------------------------------------------------------------
    def resize_vocab(self, new_size, init_std=0.02):
        """Grow every vocab-shaped head to `new_size`.

        Preserves rows [0 .. old_size-1] exactly.
        Random-initializes rows [old_size .. new_size-1] with N(0, init_std).
        Returns self (in place).
        """
        old_size = self.ctc_fc.out_features
        if new_size <= old_size:
            return self
        dev = self.ctc_fc.weight.device
        dt  = self.ctc_fc.weight.dtype

        def _grow_linear(lin, new_out):
            w_old = lin.weight.data
            b_old = lin.bias.data
            w_new = torch.empty(new_out, lin.in_features, device=dev, dtype=dt)
            b_new = torch.zeros(new_out, device=dev, dtype=dt)
            nn.init.normal_(w_new, std=init_std)
            with torch.no_grad():
                w_new[:old_size] = w_old
                b_new[:old_size] = b_old
            new_lin = nn.Linear(lin.in_features, new_out,
                                bias=True, device=dev, dtype=dt)
            new_lin.weight.data = w_new
            new_lin.bias.data = b_new
            return new_lin

        # ctc head
        self.ctc_fc = _grow_linear(self.ctc_fc, new_size)

        # attention decoder linear head
        self.decoder.fc = _grow_linear(self.decoder.fc, new_size)

        # attention decoder embedding table
        e_old = self.decoder.embed.weight.data
        e_new = torch.empty(new_size, e_old.shape[1], device=dev, dtype=dt)
        nn.init.normal_(e_new, std=init_std)
        with torch.no_grad():
            e_new[:old_size] = e_old
        new_embed = nn.Embedding(new_size, e_old.shape[1], device=dev, dtype=dt)
        new_embed.weight.data = e_new
        self.decoder.embed = new_embed

        return self

    # --------------------------------------------------------------------------
    # Forward
    # --------------------------------------------------------------------------
    def encode(self, xs, ilens=None):
        """Encodes log-mel spectrogram features.
        Args:
            xs: Tensor of shape (B, T, N_MELS)
            ilens: Tensor or list of input lengths. Defaults to full length T.
        """
        B, T, _ = xs.shape
        if ilens is None:
            ilens = torch.full((B,), T, dtype=torch.long, device="cpu")
        elif isinstance(ilens, torch.Tensor):
            ilens = ilens.to("cpu")

        x = self.frontend(xs)
        sub_T = x.shape[1]
        enc_lens = calc_subsampled_lens(ilens)
        pad_mask = (torch.arange(sub_T, device=x.device).unsqueeze(0) >=
                    enc_lens.unsqueeze(1).to(x.device))

        for blk in self.blocks:
            x = blk(x, key_pad_mask=pad_mask)
        return self.enc_norm(x), pad_mask, enc_lens

    def forward(self, xs, ilens=None):
        """Default forward pass runs encoder + CTC head."""
        enc_out, _, enc_lens = self.encode(xs, ilens)
        ctc_logits = self.ctc_fc(enc_out)
        return ctc_logits, enc_lens


# ==============================================================================
# Loading helper (open-vocab aware)
# ==============================================================================
def load_model(weights_path="/workspace/myanmar_asr/model.safetensors",
               vocab_size=None,
               device="cpu"):
    """Load pretrained model from safetensors.

    If `vocab_size` is None, uses the checkpoint's size (2566).
    If `vocab_size` is larger, vocab-shaped heads are grown on load and
    rows [old .. new-1] are random-initialized. All other tensors must
    match exactly — a mismatch raises.
    """
    from safetensors.torch import load_file

    sd = load_file(str(weights_path))

    # infer vocab size from the checkpoint itself
    ckpt_vocab = sd["ctc_fc.weight"].shape[0]
    if vocab_size is None:
        vocab_size = ckpt_vocab

    model = ConformerASR(vocab_size=ckpt_vocab)

    # Load with strict=True first for the original architecture.
    # If new vocab is larger, grow heads after loading.
    if vocab_size == ckpt_vocab:
        model.load_state_dict(sd, strict=True)
    else:
        # build a temp model of the checkpoint's shape, load, then resize
        model.load_state_dict(sd, strict=True)
        model.resize_vocab(vocab_size)

    model.to(device)
    model.eval()
    return model