"""Verify the converted ONNX graph reproduces the true PyTorch cross-encoder score.

Compares the concat workaround (query + sep*2 + passage, fed through a single-text-column
tokenization -- what ONNXEmbeddings does) against the model's real pair-encoded score. See
README.md for why this concat trick is valid for this specific model (its graph never uses
token_type_ids, so a real query/passage boundary signal isn't needed to reproduce the score).
"""

import numpy as np
import onnxruntime as ort
import torch
from tokenizers import Tokenizer
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MODEL_ID = "jinaai/jina-reranker-v1-tiny-en"
SEP_TOKEN = "</s>"
SEP_REPEAT = 2

PAIRS = [
    ("What is the capital of France?", "Paris is the capital and most populous city of France."),
    ("How do airplanes fly?", "Airplanes generate lift through the shape and angle of their wings."),
    (
        "What causes rain?",
        "Rain forms when water vapor in clouds condenses into droplets heavy enough to fall.",
    ),
]


def sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-x))


def hf_baseline_scores(pairs: list[tuple[str, str]]) -> np.ndarray:
    """True PyTorch pair-encoded score, no ONNX involved."""
    # trust_remote_code=True is required -- this repo ships custom modeling code
    # (modeling_bert.py) that differs from the generic transformers BertModel; without it,
    # AutoModel silently falls back to the generic architecture and loads garbage weights.
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    try:
        model = AutoModelForSequenceClassification.from_pretrained(
            MODEL_ID, trust_remote_code=True, dtype=torch.float32
        )
    except RuntimeError:
        # This repo's config.json declares only the base architecture, no num_labels --
        # AutoModelForSequenceClassification then defaults to num_labels=2 and fails to
        # load the checkpoint's actual 1-unit classifier head.
        model = AutoModelForSequenceClassification.from_pretrained(
            MODEL_ID, trust_remote_code=True, num_labels=1, dtype=torch.float32
        )
    model.eval()

    queries, passages = zip(*pairs)
    inputs = tokenizer(
        list(queries), list(passages), padding=True, truncation=True, max_length=512, return_tensors="pt"
    )
    with torch.no_grad():
        logits = model(**inputs).logits.view(-1)
    return torch.sigmoid(logits).numpy()


def onnx_concat_scores(pairs: list[tuple[str, str]]) -> np.ndarray:
    """Simulates ONNXEmbeddings: one pre-joined string through a single-sequence encode."""
    session = ort.InferenceSession("onnx/model.onnx", providers=["CPUExecutionProvider"])
    tokenizer = Tokenizer.from_file("tokenizer.json")
    tokenizer.enable_padding(length=None)
    tokenizer.enable_truncation(max_length=512)

    joined = [f"{q}{SEP_TOKEN * SEP_REPEAT}{p}" for q, p in pairs]
    encodings = tokenizer.encode_batch(joined)
    input_names = {i.name for i in session.get_inputs()}
    feed = {
        "input_ids": np.array([e.ids for e in encodings], dtype=np.int64),
        "attention_mask": np.array([e.attention_mask for e in encodings], dtype=np.int64),
        "token_type_ids": np.array([e.type_ids for e in encodings], dtype=np.int64),
    }
    feed = {k: v for k, v in feed.items() if k in input_names}
    logits = session.run(["sentence_embedding"], feed)[0].reshape(-1)
    return sigmoid(logits)


if __name__ == "__main__":
    hf_scores = hf_baseline_scores(PAIRS)
    onnx_scores = onnx_concat_scores(PAIRS)
    delta = np.abs(onnx_scores - hf_scores)
    print("hf scores:  ", np.round(hf_scores, 6))
    print("onnx scores:", np.round(onnx_scores, 6))
    print("mean/max delta:", delta.mean(), delta.max())