#!/usr/bin/env python3
"""
Slice 19 model zoo: three genuinely sequence-sensitive tiny models, trained
from scratch on Maven's narrow binary pragmatics task.

    A. CharCNN      — codepoint ids → char embedding → parallel small 1D convs
                      (several kernel widths) → global max-pool → linear head
    B. BiGRU        — subword ids → token embedding → 1-layer BiGRU →
                      maxpool[final] → linear head
    C. TinyTransformer — subword ids → token embedding + sine position →
                      N self-attention encoder blocks (heads, FFN 4x, PreNorm) →
                      CLS → linear head

All expose :forward(ids) returning the binary logit, plus .n_params().
Deterministic: everything is plain torch ops.
"""

import math

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


class CharCNN(nn.Module):
    def __init__(self, vocab_size, embed_dim, filters, widths, pad_idx=0, dropout=0.3):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
        self.convs = nn.ModuleList([
            nn.Conv1d(embed_dim, filters, k, padding=(k - 1) // 2)
            for k in widths
        ])
        self.dropout = nn.Dropout(dropout)
        self.head = nn.Linear(filters * len(widths), 1)

    def forward(self, ids):
        # ids: (B, T)
        x = self.embed(ids).transpose(1, 2)          # (B, D, T)
        hiddens = [F.relu(conv(x)) for conv in self.convs]  # each (B, F, T)
        pooled = torch.cat([h.max(dim=2).values for h in hiddens], dim=1)  # (B, F*W)
        return self.head(self.dropout(pooled)).squeeze(-1)

    def n_params(self):
        return sum(p.numel() for p in self.parameters())


class BiGRU(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden, pad_idx=0, dropout=0.3):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
        self.encoder = nn.GRU(embed_dim, hidden, num_layers=1, bidirectional=True,
                              batch_first=True)
        self.dropout = nn.Dropout(dropout)
        self.head = nn.Linear(hidden * 2, 1)

    def forward(self, ids):
        mask = (ids != 0).float()                    # (B, T)
        x = self.embed(ids)
        lens = mask.sum(dim=1).clamp(min=1).long()
        x_p = nn.utils.rnn.pack_padded_sequence(x, lens.cpu(), batch_first=True,
                                                enforce_sorted=False)
        out, _ = self.encoder(x_p)
        out, _ = nn.utils.rnn.pad_packed_sequence(out, batch_first=True,
                                                  total_length=mask.size(1))
        out = out * mask.unsqueeze(-1)
        maxed = out.max(dim=1).values                 # (B, 2H)
        return self.head(self.dropout(maxed)).squeeze(-1)

    def n_params(self):
        return sum(p.numel() for p in self.parameters())


class TinyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, n_layers, n_heads, ff_mult=4,
                 max_len=64, pad_idx=0, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.embed = nn.Embedding(vocab_size, d_model, padding_idx=pad_idx)
        self.dropout = nn.Dropout(dropout)
        self.pos = nn.Parameter(torch.empty(1, max_len, d_model))
        nn.init.normal_(self.pos, std=0.02)
        blocks = []
        for _ in range(n_layers):
            blocks.append(TransformerBlock(d_model, n_heads, ff_mult, dropout))
        self.blocks = nn.ModuleList(blocks)
        self.ln_out = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, 1)

    def forward(self, ids):
        B, T = ids.shape
        mask = (ids != 0)
        x = self.embed(ids) * math.sqrt(self.d_model) + self.pos[:, :T, :]
        x = self.dropout(x)
        for blk in self.blocks:
            x = blk(x, mask)
        x = self.ln_out(x)
        pooled = x.masked_fill(~mask.unsqueeze(-1), float("-inf")).max(dim=1).values
        return self.head(pooled).squeeze(-1)

    def n_params(self):
        return sum(p.numel() for p in self.parameters())


class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, ff_mult, dropout):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout,
                                          batch_first=True)
        self.ln1 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_model * ff_mult),
            nn.GELU(),
            nn.Linear(d_model * ff_mult, d_model),
        )
        self.ln2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        # mask: (B, T) bool of non-pad; MultiheadAttention uses key_padding_mask
        h = self.ln1(x)
        h = self.attn(h, h, h, key_padding_mask=~mask,
                      need_weights=False, is_causal=False)[0]
        x = x + self.dropout(h)
        h = self.ln2(x)
        x = x + self.dropout(self.ff(h))
        return x


# ─── Sizes ladder ───────────────────────────────────────────────────────────

def make_model(arch, size, char_vocab, bpe_vocab):
    if arch == "char_cnn":
        configs = {
            "tiny":  dict(embed_dim=32,  filters=64,  widths=[3, 4, 5]),
            "medium": dict(embed_dim=64, filters=160, widths=[2, 3, 4, 5]),
        }
        c = configs[size]
        return CharCNN(char_vocab, c["embed_dim"], c["filters"], c["widths"])
    if arch == "bigru":
        configs = {
            "tiny":  dict(embed_dim=64,  hidden=64),
            "medium": dict(embed_dim=128, hidden=128),
            "large": dict(embed_dim=256, hidden=256),
        }
        c = configs[size]
        return BiGRU(bpe_vocab, c["embed_dim"], c["hidden"])
    if arch == "tiny_transformer":
        configs = {
            "small":  dict(d_model=128, n_layers=2, n_heads=4),
            "medium": dict(d_model=192, n_layers=4, n_heads=4),
        }
        c = configs[size]
        return TinyTransformer(bpe_vocab, c["d_model"], c["n_layers"], c["n_heads"])
    raise ValueError(arch)


def n_params_of(arch, size, char_vocab, bpe_vocab):
    return make_model(arch, size, char_vocab, bpe_vocab).n_params()