"""baa-ai-Embedding-Reranker-v1 (4-bit) — standalone embedder+reranker.
Group-64 int4-packed Linear weights (fp16 for embeddings/LayerNorm/classifier/pooler), dequantized to fp16 at
load. Pure torch + transformers + safetensors; runs on CPU / Apple MPS / CUDA. API: BaaEmbeddingReranker."""
import os, json, numpy as np, torch, torch.nn.functional as F
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoModelForSequenceClassification, AutoTokenizer
GROUP = 64

def _dequant(st, fk):
    p = st[fk + "::q"].numpy().astype(np.uint8)
    L = int(st[fk + "::len"][0]); out_, in_ = int(st[fk + "::shape"][0]), int(st[fk + "::shape"][1])
    q = np.empty(L, dtype=np.uint8); q[0::2] = p >> 4; q[1::2] = p & 0xF
    s = st[fk + "::s"].float().numpy().reshape(-1, 1); m = st[fk + "::m"].float().numpy().reshape(-1, 1)
    w = (q.reshape(-1, GROUP).astype(np.float32) * s + m).reshape(out_, in_)
    return torch.from_numpy(w)

class BaaEmbeddingReranker:
    def __init__(self, path=None, device=None):
        path = path or os.path.dirname(os.path.abspath(__file__))
        self.device = device or ("mps" if torch.backends.mps.is_available()
                                  else "cuda" if torch.cuda.is_available() else "cpu")
        qc = json.load(open(os.path.join(path, "quant_config.json")))
        st = load_file(os.path.join(path, "weights_q4.safetensors"))
        shared = st["shared::word_embeddings"]
        self.emb = AutoModel.from_config(AutoConfig.from_pretrained(os.path.join(path, "embedder")))
        self.rr = AutoModelForSequenceClassification.from_config(
            AutoConfig.from_pretrained(os.path.join(path, "reranker")))
        self._fill(self.emb, st, "emb", "embeddings.word_embeddings.weight", shared, set(qc["emb_q4"]))
        self._fill(self.rr, st, "rr", "roberta.embeddings.word_embeddings.weight", shared, set(qc["rr_q4"]))
        self.emb = self.emb.half().to(self.device).eval()
        self.rr = self.rr.half().to(self.device).eval()
        self.emb_tok = AutoTokenizer.from_pretrained(os.path.join(path, "embedder"))
        self.rr_tok = AutoTokenizer.from_pretrained(os.path.join(path, "reranker"))

    def _fill(self, model, st, ns, wemb_key, shared, q4):
        sd = dict(model.state_dict())
        for k in list(sd.keys()):
            if k == wemb_key:
                sd[k] = shared.to(sd[k].dtype); continue
            fk = f"{ns}::{k}"
            if k in q4:
                sd[k] = _dequant(st, fk).to(sd[k].dtype)
            elif fk in st:
                sd[k] = st[fk].to(sd[k].dtype)
        model.load_state_dict(sd)

    @torch.no_grad()
    def embed(self, texts, is_query=False, batch_size=32):
        pref = "query: " if is_query else ""
        out = []
        for i in range(0, len(texts), batch_size):
            enc = self.emb_tok([pref + t for t in texts[i:i+batch_size]], padding=True, truncation=True,
                               max_length=512, return_tensors="pt").to(self.device)
            h = self.emb(**enc).last_hidden_state[:, 0]      # CLS
            out.append(F.normalize(h, dim=-1).float().cpu().numpy())
        return np.vstack(out)

    @torch.no_grad()
    def rerank(self, query, docs, top_k=None, batch_size=32):
        scores = []
        for i in range(0, len(docs), batch_size):
            enc = self.rr_tok([(query, d[:2000]) for d in docs[i:i+batch_size]], padding=True,
                              truncation=True, max_length=512, return_tensors="pt").to(self.device)
            scores.extend(self.rr(**enc).logits[:, 0].float().cpu().tolist())
        order = sorted(range(len(docs)), key=lambda j: -scores[j])
        if top_k:
            order = order[:top_k]
        return [(docs[j], scores[j]) for j in order]