"""Evaluate rerankers on rerank/evals.   python rerank/eval_rerank.py --models A,B --evals rerank/evals --out R.json
Model specs: <hub id or dir> = cross-encoder (AutoModelForSequenceClassification; 1 logit, or logit[1] - logit[0] with 2);
qwen3rr:<hub id> = Qwen3-Reranker (causal LM; P(yes) with the model card's template and default web-search instruction).
Metric: nDCG@10 per set (graded relevance, all candidates of a query reranked), MRR@10; family means (_miracl, _wiki,
_other, _all). Query+document truncated to --max_len tokens (document side).
"""
import argparse, glob, json, math, os, time
import numpy as np, pandas as pd, torch

ap = argparse.ArgumentParser()
ap.add_argument("--models", required=True); ap.add_argument("--evals", required=True); ap.add_argument("--out", required=True)
ap.add_argument("--bs", type=int, default=64); ap.add_argument("--max_len", type=int, default=512)
ap.add_argument("--only", default="")
args = ap.parse_args()
EV = {os.path.basename(p)[:-8]: pd.read_parquet(p) for p in sorted(glob.glob(f"{args.evals}/*.parquet"))
      if not args.only or os.path.basename(p).startswith(tuple(args.only.split(",")))}


def ndcg10(rels, scores):
    order = np.argsort(-np.asarray(scores), kind="stable")[:10]
    dcg = sum((2 ** rels[i] - 1) / math.log2(k + 2) for k, i in enumerate(order))
    ideal = sorted(rels, reverse=True)[:10]
    idcg = sum((2 ** r - 1) / math.log2(k + 2) for k, r in enumerate(ideal))
    return dcg / idcg if idcg > 0 else 0.0


def mrr10(rels, scores):
    order = np.argsort(-np.asarray(scores), kind="stable")[:10]
    return next((1 / (k + 1) for k, i in enumerate(order) if rels[i] > 0), 0.0)


class CE:
    def __init__(self, spec):
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
        self.tok = AutoTokenizer.from_pretrained(spec, trust_remote_code=True)
        self.m = AutoModelForSequenceClassification.from_pretrained(spec, torch_dtype=torch.bfloat16, trust_remote_code=True).cuda().eval()

    @torch.no_grad()
    def __call__(self, pairs):
        out = np.zeros(len(pairs)); order = np.argsort([len(q) + len(d) for q, d in pairs])
        for s in range(0, len(pairs), args.bs):
            b = order[s:s + args.bs]
            enc = self.tok([pairs[i][0] for i in b], [pairs[i][1] for i in b], truncation="only_second", max_length=args.max_len,
                           padding=True, return_tensors="pt")
            lo = self.m(**{k: v.cuda() for k, v in enc.items() if k in ("input_ids", "attention_mask", "token_type_ids")}).logits.float()
            out[b] = (lo[:, 0] if lo.shape[1] == 1 else lo[:, 1] - lo[:, 0]).cpu().numpy()
        return out


class Qwen3RR:
    PRE = ("<|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")
    SUF = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
    INST = "Given a web search query, retrieve relevant passages that answer the query"

    def __init__(self, spec):
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.tok = AutoTokenizer.from_pretrained(spec, padding_side="left")
        self.m = AutoModelForCausalLM.from_pretrained(spec, torch_dtype=torch.bfloat16).cuda().eval()
        self.yes, self.no = self.tok.convert_tokens_to_ids("yes"), self.tok.convert_tokens_to_ids("no")
        self.pre, self.suf = self.tok.encode(self.PRE, add_special_tokens=False), self.tok.encode(self.SUF, add_special_tokens=False)

    @torch.no_grad()
    def __call__(self, pairs):
        out = np.zeros(len(pairs)); order = np.argsort([len(q) + len(d) for q, d in pairs])
        for s in range(0, len(pairs), args.bs):
            b = order[s:s + args.bs]
            body = [self.tok.encode(f"<Instruct>: {self.INST}\n<Query>: {pairs[i][0]}\n<Document>: {pairs[i][1]}", add_special_tokens=False)
                    [: args.max_len - len(self.pre) - len(self.suf)] for i in b]
            ids = [self.pre + x + self.suf for x in body]; L = max(map(len, ids))
            inp = torch.full((len(ids), L), self.tok.pad_token_id); att = torch.zeros((len(ids), L), dtype=torch.long)
            for j, x in enumerate(ids):
                inp[j, L - len(x):] = torch.tensor(x); att[j, L - len(x):] = 1
            lo = self.m(input_ids=inp.cuda(), attention_mask=att.cuda(), logits_to_keep=1).logits[:, -1, :].float()
            out[b] = torch.log_softmax(torch.stack([lo[:, self.no], lo[:, self.yes]], 1), 1)[:, 1].exp().cpu().numpy()
        return out


res = json.load(open(args.out)) if os.path.exists(args.out) else {}
for spec in args.models.split(","):
    kind, mid = spec.split(":", 1) if spec.startswith("qwen3rr:") else ("ce", spec)
    model = Qwen3RR(mid) if kind == "qwen3rr" else CE(mid)
    r, t0, npairs = {}, time.time(), 0
    for name, df in EV.items():
        pairs = [(q, d) for q, ds in zip(df["query"], df.docs) for d in ds]; npairs += len(pairs)
        sc = model(pairs); k = 0; nd, mr = [], []
        for rels in df.rels:
            s = sc[k:k + len(rels)]; k += len(rels)
            nd.append(ndcg10(list(rels), s)); mr.append(mrr10(list(rels), s))
        r[name] = dict(n=len(df), ndcg10=float(np.mean(nd)), mrr10=float(np.mean(mr)))
    fam = lambda p: [v["ndcg10"] for k2, v in r.items() if k2.startswith(p)]
    r["_miracl"] = float(np.mean(fam("miracl_"))) if fam("miracl_") else None
    r["_wiki"] = float(np.mean(fam("wiki_"))) if fam("wiki_") else None
    oth = [v["ndcg10"] for k2, v in r.items() if not k2.startswith(("miracl_", "wiki_", "_"))]
    r["_other"] = float(np.mean(oth)) if oth else None
    r["_all"] = float(np.mean([x for x in (r["_miracl"], r["_wiki"], r["_other"]) if x is not None]))
    r["_pairs_per_s"] = npairs / (time.time() - t0)
    print(spec, {k2: (round(v, 4) if isinstance(v, float) else v) for k2, v in r.items() if k2.startswith("_")}, flush=True)
    res[spec] = r; json.dump(res, open(args.out, "w"), indent=1)
    del model; torch.cuda.empty_cache()