"""baa.ai unified Embedding+Reranker loader (generic).

A single artifact that does both halves of RAG retrieval — bi-encoder embedding AND cross-encoder reranking —
over ONE shared word-embedding table. The reranker's word-embedding matrix is stored only once (in the
embedder) and injected at load, so the packaged model is smaller than shipping the two components separately,
at no measured quality cost.

Works for BERT-based and XLM-RoBERTa-based stacks alike: the reranker's encoder submodule is resolved
generically via `reranker.base_model` (so `.bert` / `.roberta` are both handled). Optional per-model query/doc
prompts are read from config.json (e.g. some models use a "query: " prefix).

Usage:
    from modeling_baa import BaaEmbeddingReranker
    m = BaaEmbeddingReranker("path/to/model-dir")
    qv = m.embed(["what is a cross-encoder?"], is_query=True)         # normalized bi-encoder vectors
    dv = m.embed(["a cross-encoder scores a (query, doc) pair jointly"])
    ranked = m.rerank("what is a cross-encoder?", ["doc A ...", "doc B ..."])   # [(doc, score), ...] desc
"""
import os, json, torch
from safetensors.torch import load_file
from sentence_transformers import SentenceTransformer
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer


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"))
        cfg = {}
        cfg_path = os.path.join(path, "config.json")
        if os.path.exists(cfg_path):
            cfg = json.load(open(cfg_path))
        self.q_prompt = cfg.get("embed_query_prompt", "") or ""
        self.d_prompt = cfg.get("embed_doc_prompt", "") or ""
        trc = bool(cfg.get("trust_remote_code", False))
        emb_dir, rr_dir = os.path.join(path, "embedder"), os.path.join(path, "reranker")

        # embedder = bi-encoder; holds the canonical shared word-embedding table
        self.embedder = SentenceTransformer(emb_dir, device=self.device, trust_remote_code=trc)
        shared_wemb = self.embedder[0].auto_model.embeddings.word_embeddings.weight.data

        # reranker = cross-encoder seq-classifier; word-embedding stripped on disk -> injected from shared table
        rr_cfg = AutoConfig.from_pretrained(rr_dir, trust_remote_code=trc)
        self.reranker = AutoModelForSequenceClassification.from_config(rr_cfg, trust_remote_code=trc)
        self.reranker.load_state_dict(load_file(os.path.join(rr_dir, "model.safetensors")), strict=False)
        # resolve the encoder submodule generically (.bert for BERT, .roberta for XLM-R, ...)
        base = self.reranker.base_model
        base.embeddings.word_embeddings.weight.data = shared_wemb.to(self.reranker.dtype).clone()
        self.reranker.to(self.device).eval()
        self.rr_tok = AutoTokenizer.from_pretrained(rr_dir, trust_remote_code=trc)
        # Weights may be stored fp16 on disk (smaller artifact); CPU can't compute in half -> upcast to fp32.
        if str(self.device) == "cpu":
            self.embedder = self.embedder.to(torch.float32)
            self.reranker = self.reranker.float()

    def embed(self, texts, is_query=False, batch_size=32):
        """Return L2-normalized bi-encoder vectors. Applies the model's query/doc prompt if configured."""
        prompt = self.q_prompt if is_query else self.d_prompt
        texts = [prompt + t for t in texts] if prompt else list(texts)
        return self.embedder.encode(texts, normalize_embeddings=True,
                                    batch_size=batch_size, show_progress_bar=False)

    @torch.no_grad()
    def rerank(self, query, docs, top_k=None, batch_size=32):
        """Cross-encoder relevance scores for (query, doc) pairs; returns [(doc, score), ...] sorted desc."""
        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.reranker(**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]


if __name__ == "__main__":
    import numpy as np
    m = BaaEmbeddingReranker()
    q = "How does a cross-encoder reranker work?"
    docs = ["A cross-encoder jointly encodes the query and document to score relevance.",
            "The mitochondria is the powerhouse of the cell.",
            "Bi-encoders embed query and document separately for fast retrieval."]
    qv = m.embed([q], is_query=True)[0]; dv = m.embed(docs)
    print("embed cos:", [round(float(np.dot(qv, d)), 3) for d in dv])
    print("rerank   :", [(round(s, 2), d[:45]) for d, s in m.rerank(q, docs)])