#!/usr/bin/env python3
"""Paired comparison of two embedding models scored by knn-bench.py.

Both runs must cover the same documents in the same order. For each label field
it reports how the per-document outcomes overlap, an exact two-sided McNemar
test on the documents where exactly one model is right, and how often the two
models pick the same nearest neighbour.

Usage:
    knn-compare.py --corpus-a a.jsonl --vecs-a a.vecs.json \\
                   --corpus-b b.jsonl --vecs-b b.vecs.json

The vectors files are the caches knn-bench.py writes next to its corpus.
Standard library only.
"""
import argparse
import importlib.util
import json
import math
import os

HERE = os.path.dirname(os.path.abspath(__file__))
spec = importlib.util.spec_from_file_location("knn_bench", os.path.join(HERE, "knn-bench.py"))
bench = importlib.util.module_from_spec(spec)
spec.loader.exec_module(bench)


def load(corpus, vecs):
    with open(corpus) as fh:
        docs = [json.loads(line) for line in fh if line.strip()]
    with open(vecs) as fh:
        vectors = [bench.normalise(v) for v in json.load(fh)]
    if len(docs) != len(vectors):
        raise SystemExit(f"{vecs}: {len(vectors)} vectors for {len(docs)} documents")
    return docs, bench.nearest_neighbours(vectors)


def mcnemar_exact(only_a, only_b):
    """Two-sided exact p-value for a fair coin over the discordant pairs."""
    n = only_a + only_b
    if n == 0:
        return 1.0
    k = min(only_a, only_b)
    tail = sum(math.comb(n, i) for i in range(k + 1)) / 2 ** n
    return min(1.0, 2 * tail)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    for side in "ab":
        ap.add_argument(f"--corpus-{side}", required=True)
        ap.add_argument(f"--vecs-{side}", required=True)
    args = ap.parse_args()

    docs_a, nn_a = load(args.corpus_a, args.vecs_a)
    docs_b, nn_b = load(args.corpus_b, args.vecs_b)
    if [d["id"] for d in docs_a] != [d["id"] for d in docs_b]:
        raise SystemExit("the two corpora do not hold the same documents in the same order")

    n = len(docs_a)
    same_nn = sum(a == b for a, b in zip(nn_a, nn_b))
    print(f"documents={n}  same nearest neighbour in both models: {same_nn} ({100 * same_nn / n:.1f}%)")
    for field in ("correspondent", "doctype"):
        ok_a = [docs_a[i][field] == docs_a[nn_a[i]][field] for i in range(n)]
        ok_b = [docs_b[i][field] == docs_b[nn_b[i]][field] for i in range(n)]
        both = sum(a and b for a, b in zip(ok_a, ok_b))
        only_a = sum(a and not b for a, b in zip(ok_a, ok_b))
        only_b = sum(b and not a for a, b in zip(ok_a, ok_b))
        neither = n - both - only_a - only_b
        p = mcnemar_exact(only_a, only_b)
        print(f"{field:<14} both right={both}  only A={only_a}  only B={only_b}  both wrong={neither}  "
              f"McNemar p={p:.3f}")


if __name__ == "__main__":
    main()