"""Encoder-only embedding model. Pre-norm, GELU, learned positions, SDPA attention, tied MLM head."""
import json, math
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

PAD_ID, MASK_ID, VOCAB = 50256, 50257, 50258

CFG_57M = dict(vocab=VOCAB, hidden=512, heads=8, ffn=2048, layers=10, max_pos=1024, dropout=0.0)

class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        h = c["hidden"]; self.nh = c["heads"]; self.p = c["dropout"]
        self.ln1 = nn.LayerNorm(h); self.qkv = nn.Linear(h, 3 * h); self.proj = nn.Linear(h, h)
        self.ln2 = nn.LayerNorm(h); self.fc1 = nn.Linear(h, c["ffn"]); self.fc2 = nn.Linear(c["ffn"], h)
    def forward(self, x, mask=None):
        B, L, H = x.shape
        q, k, v = self.qkv(self.ln1(x)).view(B, L, 3, self.nh, H // self.nh).permute(2, 0, 3, 1, 4)
        # mask=None (MLM, no padding) -> flash/mem-efficient SDPA, never a materialised LxL math path.
        a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.p if self.training else 0.0)
        x = x + F.dropout(self.proj(a.transpose(1, 2).reshape(B, L, H)), self.p, self.training)
        return x + F.dropout(self.fc2(F.gelu(self.fc1(self.ln2(x)))), self.p, self.training)

class Encoder(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.c = c
        self.tok_emb = nn.Embedding(c["vocab"], c["hidden"])
        self.pos_emb = nn.Embedding(c["max_pos"], c["hidden"])
        self.blocks = nn.ModuleList([Block(c) for _ in range(c["layers"])])
        self.ln_f = nn.LayerNorm(c["hidden"])
        self.grad_ckpt = False
        self.apply(self._init)
        for n, p in self.named_parameters():       # GPT-2 style scaled residual init
            if n.endswith("proj.weight") or n.endswith("fc2.weight"):
                nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * c["layers"]))
    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, 0.0, 0.02)
            if getattr(m, "bias", None) is not None: nn.init.zeros_(m.bias)
    def forward(self, ids, attn_mask=None):
        B, L = ids.shape
        x = self.tok_emb(ids) + self.pos_emb(torch.arange(L, device=ids.device))[None]
        for blk in self.blocks:
            if self.grad_ckpt and self.training and torch.is_grad_enabled():
                x = checkpoint(blk, x, attn_mask, use_reentrant=False)
            else:
                x = blk(x, attn_mask)
        return self.ln_f(x)

class MLMModel(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.encoder = Encoder(c)
        h = c["hidden"]
        self.head_dense = nn.Linear(h, h); self.head_ln = nn.LayerNorm(h)
        self.head_bias = nn.Parameter(torch.zeros(c["vocab"]))
    def forward(self, ids, labels):
        h = self.encoder(ids)                      # no mask: windows have no padding
        sel = labels != -100
        hs = self.head_ln(F.gelu(self.head_dense(h[sel])))      # logits only for masked positions
        logits = F.linear(hs, self.encoder.tok_emb.weight, self.head_bias)   # tied
        return F.cross_entropy(logits.float(), labels[sel])

class EmbedModel(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.encoder = Encoder(c)
    def forward(self, ids, lens):
        L = ids.shape[1]
        keep = torch.arange(L, device=ids.device)[None] < lens[:, None]          # [B,L]
        h = self.encoder(ids, keep[:, None, None, :])                           # bool mask: True = attend
        m = keep.unsqueeze(-1).to(h.dtype)
        e = (h * m).sum(1) / m.sum(1).clamp_min(1)
        return F.normalize(e.float(), dim=-1)

def count_params(model):
    return sum(p.numel() for p in model.parameters())

def mask_tokens(ids, p=0.15, gen=None):
    """Dynamic BERT masking 80/10/10 on GPU. EOT (pad/doc separator) is never selected."""
    sel = (torch.rand(ids.shape, device=ids.device, generator=gen) < p) & (ids != PAD_ID)
    labels = torch.where(sel, ids, torch.full_like(ids, -100))
    r = torch.rand(ids.shape, device=ids.device, generator=gen)
    inp = ids.clone()
    inp[sel & (r < 0.8)] = MASK_ID
    rnd = sel & (r >= 0.8) & (r < 0.9)
    inp[rnd] = torch.randint(0, PAD_ID, ids.shape, device=ids.device, generator=gen)[rnd]
    return inp, labels

def load_encoder_weights(model, path, map_location="cpu"):
    sd = torch.load(path, map_location=map_location)
    sd = sd.get("model", sd)
    enc = {k[len("encoder."):]: v for k, v in sd.items() if k.startswith("encoder.")} or sd
    missing, unexpected = model.encoder.load_state_dict(enc, strict=True), None
    return missing

class Embedder:
    """Usage: e = Embedder('path/to/57M_embed_Full'); e.encode(['hello world'])"""
    def __init__(self, folder, device=None):
        from safetensors.torch import load_file
        from transformers import AutoTokenizer
        self.c = json.load(open(f"{folder}/config.json"))["encoder_config"]
        self.m = EmbedModel(self.c); self.m.encoder.load_state_dict(load_file(f"{folder}/model.safetensors"))
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.m.to(self.device).eval()
        self.tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
    @torch.no_grad()
    def encode(self, texts, max_len=128, bs=64):
        out = []
        for i in range(0, len(texts), bs):
            ids = self.tok(texts[i:i+bs], truncation=True, max_length=max_len, add_special_tokens=False)["input_ids"]
            L = max(1, max(len(x) for x in ids))
            t = torch.full((len(ids), L), PAD_ID, dtype=torch.long)
            lens = torch.tensor([max(1, len(x)) for x in ids])
            for j, x in enumerate(ids): t[j, :len(x)] = torch.tensor(x, dtype=torch.long)
            out.append(self.m(t.to(self.device), lens.to(self.device)).cpu())
        return torch.cat(out)