#!/usr/bin/env python3
"""
Shared encoder used by 13 (build index) and 14 (profile/eval).

The layer names here match what 12_train_encoder.py saved, so a checkpoint trained by
12 loads straight into this class.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer


class Encoder(nn.Module):
    def __init__(self, model_name, emb_dim):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(model_name, trust_remote_code=True)
        h = self.backbone.config.hidden_size
        self.attn = nn.Linear(h, 1)                          # attention pooling
        self.proj = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, emb_dim))

    def forward(self, input_ids, attention_mask):
        hidden = self.backbone(input_ids=input_ids, attention_mask=attention_mask)[0]
        scores = self.attn(hidden).squeeze(-1).masked_fill(attention_mask == 0, -1e9)
        w = scores.softmax(-1).unsqueeze(-1)
        pooled = (hidden * w).sum(1)
        return F.normalize(self.proj(pooled), dim=-1)


def load_encoder(ckpt_path, device):
    ck = torch.load(ckpt_path, map_location=device)
    cfg = ck["config"]
    model = Encoder(cfg["MODEL"], cfg["EMB_DIM"]).to(device)
    model.load_state_dict(ck["model"])
    model.eval()
    tok = AutoTokenizer.from_pretrained(cfg["MODEL"], trust_remote_code=True)
    return model, tok, cfg


@torch.no_grad()
def embed(model, tok, seqs, device, max_length=128, bf16=True):
    """Return an (N, emb_dim) float32 numpy array of L2-normalized embeddings."""
    t = tok(seqs, padding=True, truncation=True, max_length=max_length, return_tensors="pt")
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16,
                        enabled=(bf16 and device == "cuda")):
        z = model(t["input_ids"].to(device), t["attention_mask"].to(device))
    return z.float().cpu().numpy()