"""int8 (gather-only) ONNX check for a reranker.   python rerank/quant_check.py MODEL_DIR EVAL_DIR
Quantizes onnx/model.onnx -> onnx/model_quantized.onnx (embedding Gather only) and compares with fp32 ONNX on (query,
passage) pairs from the first candidates of 4 eval sets: mean / max |sigmoid difference| and nDCG@10 on those sets.
Writes MODEL_DIR/onnx_sweep.json (same keys as train/quant_sweep.py: gather_only.{mean_diff,max_diff,mb}, n)."""
import glob, json, math, os, sys
import numpy as np, pandas as pd, onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic, QuantType
from transformers import AutoTokenizer

d, ev = sys.argv[1:3]
tok = AutoTokenizer.from_pretrained(d)
sets = [p for p in sorted(glob.glob(f"{ev}/*.parquet")) if os.path.basename(p).split(".")[0] in ("miracl_de", "miracl_zh", "wiki_en", "esci_jp", "rubq")]
data = [(pd.read_parquet(p).head(12)) for p in sets]
so = ort.SessionOptions(); so.intra_op_num_threads = 8


def run(path):
    s = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"]); allp = []; nd = []
    for df in data:
        for q, docs, rels in zip(df["query"], df.docs, df.rels):
            docs = list(docs)[:30]; rels = list(rels)[:30]
            e = tok([q] * len(docs), docs, truncation="only_second", max_length=384, padding=True, return_tensors="np")
            lo = s.run(None, {"input_ids": e["input_ids"], "attention_mask": e["attention_mask"]})[0][:, 0]
            allp += list(1 / (1 + np.exp(-lo)))
            order = np.argsort(-lo)[:10]; ideal = sorted(rels, reverse=True)[:10]
            idcg = sum((2 ** r - 1) / math.log2(k + 2) for k, r in enumerate(ideal))
            nd.append(sum((2 ** rels[i] - 1) / math.log2(k + 2) for k, i in enumerate(order)) / idcg if idcg else 0)
    return np.array(allp), float(np.mean(nd))


ref, nd_ref = run(f"{d}/onnx/model.onnx")
quantize_dynamic(f"{d}/onnx/model.onnx", f"{d}/onnx/model_quantized.onnx", weight_type=QuantType.QInt8, op_types_to_quantize=["Gather"])
q, nd_q = run(f"{d}/onnx/model_quantized.onnx")
r = dict(gather_only=dict(mean_diff=float(np.abs(q - ref).mean()), max_diff=float(np.abs(q - ref).max()), ndcg_fp32=nd_ref, ndcg_int8=nd_q,
                          mb=os.path.getsize(f"{d}/onnx/model_quantized.onnx") / 1e6), n=int(len(ref)), best="gather_only")
json.dump(r, open(f"{d}/onnx_sweep.json", "w"), indent=1); print(r, flush=True)