"""Embedding / rerank quality evaluation (tokenflow.dev/notes/chinese-retrieval-benchmark), one fixed protocol for every candidate.

  python eval.py prep                              # build the four sets under $EVAL_WORKDIR/sets
  python eval.py embed <model> <backend> [sets]    # backend: vllm | st; sets: comma list (default all four)
                                                   # writes results/embed.<model>.<sets joined by +>.json
                                                   # and results/runs/<model>.<set>.json (top 100 per query)
  python eval.py rerank <model> <retriever>        # reranks <retriever>'s top 50 with a vLLM cross-encoder
                                                   # writes results/rerank.<model>.<retriever>.json

Sets: C-MTEB T2Retrieval and DuRetrieval dev (300 queries sampled with seed 0, full C-MTEB corpus),
MTEB SciFact and NFCorpus test (full). Embedding inputs, prefixes included, truncated to 512 tokens;
reranker documents clipped to 512 tokens before the prompt template.
nDCG uses exponential gain (2^rel - 1). Only NFCorpus has graded labels (1, 2); the other sets are binary,
where this equals the linear gain used by trec_eval / MTEB.
$EVAL_WORKDIR defaults to /workspace.
"""
import json, math, os, random, sys, time

import numpy as np

W = os.environ.get("EVAL_WORKDIR", "/workspace")
SETS = ["t2retrieval", "duretrieval", "scifact", "nfcorpus"]
SEED, N_QUERIES, DOC_TOKENS, RERANK_K = 0, 300, 512, 50

# Retrieval instructions as used in the Qwen3-Embedding / E5 evaluation prompts.
INSTRUCT = {
    "t2retrieval": "Given a Chinese search query, retrieve web passages that answer the question",
    "duretrieval": "Given a Chinese search query, retrieve web passages that answer the question",
    "scifact": "Given a scientific claim, retrieve documents that support or refute the claim",
    "nfcorpus": "Given a question, retrieve relevant documents that best answer the question",
}
# EmbeddingGemma 2 model-card task prefixes.
GEMMA_TASK = {"t2retrieval": "search result", "duretrieval": "search result",
              "scifact": "fact checking", "nfcorpus": "question answering"}

MODELS = {
    "Qwen3-Embedding-0.6B": {"query": "qwen", "doc": "plain"},
    "Qwen3-Embedding-4B": {"query": "qwen", "doc": "plain"},
    "bge-m3": {"query": "plain", "doc": "plain"},
    "embeddinggemma-2": {"query": "gemma", "doc": "gemma"},
}


def prep():
    import pyarrow.parquet as pq
    from glob import glob

    def parquet(pattern):
        files = glob(f"{W}/data/{pattern}")
        assert len(files) == 1, f"expected one shard for {pattern}, found {files}"  # true at the pinned revisions
        return pq.read_table(files[0]).to_pylist()

    def jsonl(path):
        with open(path) as f:
            return [json.loads(l) for l in f]

    for name, repo in [("t2retrieval", "C-MTEB__T2Retrieval"), ("duretrieval", "C-MTEB__DuRetrieval")]:
        queries = {r["id"]: r["text"] for r in parquet(f"{repo}/data/queries-*.parquet")}
        corpus = {r["id"]: {"title": "", "text": r["text"]} for r in parquet(f"{repo}/data/corpus-*.parquet")}
        qrels = {}
        for r in parquet(f"{repo}-qrels/data/dev-*.parquet"):
            if r["score"] > 0:
                qrels.setdefault(r["qid"], {})[r["pid"]] = int(r["score"])
        rng = random.Random(SEED)
        qids = sorted(q for q in qrels if q in queries)
        qids = rng.sample(qids, N_QUERIES)
        write(name, {q: queries[q] for q in qids}, corpus, {q: {p: g for p, g in qrels[q].items() if p in corpus} for q in qids})

    for name, repo in [("scifact", "mteb__scifact"), ("nfcorpus", "mteb__nfcorpus")]:
        queries = {r["_id"]: r["text"] for r in jsonl(f"{W}/data/{repo}/queries.jsonl")}
        corpus = {r["_id"]: {"title": r.get("title") or "", "text": r["text"]} for r in jsonl(f"{W}/data/{repo}/corpus.jsonl")}
        qrels = {}
        for r in jsonl(f"{W}/data/{repo}/qrels/test.jsonl"):
            if int(r["score"]) > 0:
                qrels.setdefault(str(r["query-id"]), {})[str(r["corpus-id"])] = int(r["score"])
        qids = sorted(q for q in qrels if q in queries)
        write(name, {q: queries[q] for q in qids}, corpus, {q: qrels[q] for q in qids})


def write(name, queries, corpus, qrels):
    os.makedirs(f"{W}/sets", exist_ok=True)
    with open(f"{W}/sets/{name}.json", "w") as f:
        json.dump({"queries": queries, "corpus": corpus, "qrels": qrels}, f, ensure_ascii=False)
    print(name, "queries", len(queries), "corpus", len(corpus), "qrels", sum(len(v) for v in qrels.values()), flush=True)


def load(name):
    with open(f"{W}/sets/{name}.json") as f:
        return json.load(f)


def query_text(style, task, q):
    if style == "qwen":
        return f"Instruct: {INSTRUCT[task]}\nQuery:{q}"
    if style == "gemma":
        return f"task: {GEMMA_TASK[task]} | query: {q}"
    return q


def doc_text(style, d):
    if style == "gemma":
        return f"title: {d['title'] or 'none'} | text: {d['text']}"
    return f"{d['title']} {d['text']}".strip()


def ndcg_at(ranked, rel, k=10):
    dcg = sum((2 ** rel.get(p, 0) - 1) / math.log2(i + 2) for i, p in enumerate(ranked[:k]))
    ideal = sorted(rel.values(), reverse=True)[:k]
    idcg = sum((2 ** g - 1) / math.log2(i + 2) for i, g in enumerate(ideal))
    return dcg / idcg if idcg else 0.0


def recall_at(ranked, rel, k=100):
    return len(set(ranked[:k]) & set(rel)) / len(rel)


def score_runs(runs, qrels):
    n = [ndcg_at(runs[q], qrels[q]) for q in qrels]
    r = [recall_at(runs[q], qrels[q]) for q in qrels]
    return {"ndcg@10": round(100 * float(np.mean(n)), 2), "recall@100": round(100 * float(np.mean(r)), 2), "queries": len(n)}


class VllmEmbedder:
    def __init__(self, path):
        from vllm import LLM
        self.llm = LLM(model=path, runner="pooling", max_model_len=8192, gpu_memory_utilization=0.85,
                       dtype="bfloat16", enforce_eager=False, seed=0)
        self.tok = self.llm.get_tokenizer()

    def __call__(self, texts, is_doc):
        # Tokenize like sentence-transformers does (truncate to DOC_TOKENS, special tokens kept).
        from vllm.inputs import TokensPrompt
        ids = self.tok(texts, truncation=True, max_length=DOC_TOKENS)["input_ids"]
        outs = self.llm.embed([TokensPrompt(prompt_token_ids=i) for i in ids], use_tqdm=False)
        return np.array([o.outputs.embedding for o in outs], dtype=np.float32), sum(len(o.prompt_token_ids) for o in outs)


class StEmbedder:
    def __init__(self, path):
        import torch
        from sentence_transformers import SentenceTransformer
        self.m = SentenceTransformer(path, device="cuda", model_kwargs={"torch_dtype": torch.bfloat16})
        self.m.max_seq_length = DOC_TOKENS

    def __call__(self, texts, is_doc):
        e = self.m.encode(texts, batch_size=64, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)
        toks = sum(len(x) for x in self.m.tokenizer(texts, truncation=True, max_length=DOC_TOKENS)["input_ids"])
        return e.astype(np.float32), toks


def embed(model, backend, sets=SETS):
    spec = MODELS[model]
    path = f"{W}/models/{model}"
    emb = VllmEmbedder(path) if backend == "vllm" else StEmbedder(path)
    out = {"model": model, "backend": backend, "doc_tokens": DOC_TOKENS, "sets": {}}
    os.makedirs(f"{W}/results/runs", exist_ok=True)
    for name in sets:
        s = load(name)
        pids = list(s["corpus"])
        qids = list(s["qrels"])
        # doc_encode_s: wall time of one offline call, including text formatting and tokenization (twice for
        # the sentence-transformers backend). Kept for reference; the published note does not compare speed.
        t0 = time.time()
        D, dtok = emb([doc_text(spec["doc"], s["corpus"][p]) for p in pids], True)
        t_doc = time.time() - t0
        Q, _ = emb([query_text(spec["query"], name, s["queries"][q]) for q in qids], False)
        D /= np.linalg.norm(D, axis=1, keepdims=True)
        Q /= np.linalg.norm(Q, axis=1, keepdims=True)
        sims = Q @ D.T
        top = np.argsort(-sims, axis=1)[:, :100]
        runs = {q: [pids[j] for j in top[i]] for i, q in enumerate(qids)}
        res = score_runs(runs, s["qrels"])
        res.update({"corpus": len(pids), "dim": int(D.shape[1]), "doc_tokens_total": int(dtok),
                    "doc_encode_s": round(t_doc, 1)})
        out["sets"][name] = res
        with open(f"{W}/results/runs/{model}.{name}.json", "w") as f:
            json.dump(runs, f)
        print(model, name, res, flush=True)
    with open(f"{W}/results/embed.{model}.{'+'.join(sets)}.json", "w") as f:
        json.dump(out, f, indent=1)


QWEN_RR_PREFIX = ('<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the '
                  'Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n')
QWEN_RR_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"


def rerank(model, retriever):
    from vllm import LLM
    path = f"{W}/models/{model}"
    qwen = model.startswith("Qwen3-Reranker")
    kw = {}
    if qwen:
        kw["hf_overrides"] = {"architectures": ["Qwen3ForSequenceClassification"],
                              "classifier_from_token": ["no", "yes"], "is_original_qwen3_reranker": True}
    llm = LLM(model=path, runner="pooling", max_model_len=8192, gpu_memory_utilization=0.85, dtype="bfloat16", seed=0, **kw)
    tok = llm.get_tokenizer()

    def clip(text):
        # Truncate the document itself (same 512 tokens as embedding), never the pair: the Qwen3
        # reranker scores the token after its template suffix, so a truncated pair loses it.
        ids = tok(text, add_special_tokens=False, truncation=True, max_length=DOC_TOKENS)["input_ids"]
        return tok.decode(ids)

    out = {"model": model, "retriever": retriever, "k": RERANK_K, "doc_tokens": DOC_TOKENS, "sets": {}}
    for name in ["scifact", "nfcorpus", "duretrieval", "t2retrieval"]:
        s = load(name)
        with open(f"{W}/results/runs/{retriever}.{name}.json") as f:
            runs = json.load(f)
        base = score_runs({q: runs[q][:RERANK_K] for q in s["qrels"]}, s["qrels"])
        reranked = {}
        t0 = time.time()
        for q in s["qrels"]:
            cands = runs[q][:RERANK_K]
            docs = [clip(doc_text("plain", s["corpus"][p])) for p in cands]
            if qwen:
                qt = f"{QWEN_RR_PREFIX}<Instruct>: {INSTRUCT[name]}\n<Query>: {s['queries'][q]}\n"
                docs = [f"<Document>: {d}{QWEN_RR_SUFFIX}" for d in docs]
            else:
                qt = s["queries"][q]
            outs = llm.score(qt, docs, use_tqdm=False)
            sc = [o.outputs.score for o in outs]
            reranked[q] = [cands[i] for i in np.argsort(-np.array(sc))]
        res = score_runs(reranked, s["qrels"])
        out["sets"][name] = {"ndcg@10": res["ndcg@10"], "retriever_ndcg@10": base["ndcg@10"],
                             "queries": res["queries"], "seconds": round(time.time() - t0, 1)}
        print(model, name, out["sets"][name], flush=True)
    with open(f"{W}/results/rerank.{model}.{retriever}.json", "w") as f:
        json.dump(out, f, indent=1)


if __name__ == "__main__":
    cmd = sys.argv[1]
    if cmd == "prep":
        prep()
    elif cmd == "embed":
        embed(sys.argv[2], sys.argv[3], sys.argv[4].split(",") if len(sys.argv) > 4 else SETS)
    elif cmd == "rerank":
        rerank(sys.argv[2], sys.argv[3])