"""Sizing an ensemble with Chinmay's bge-m3 run (leaderboard 0.966) against ours (0.976423).

F0.5 break-even, derived exactly: for one entity F0.5 = 1.25*TP/(0.25*T + k). Adding a
candidate changes the denominator D = 0.25T + k by +1. Adding a TP gains
1.25(D-TP)/(D(D+1)); adding an FP loses 1.25*TP/(D(D+1)). So an extra pair is worth taking
iff its TP:FP ratio beats TP/(D-TP), which at P=R=1 (T=k, D=1.25T) is exactly 4:1 - i.e.
the added pairs need precision > 80%.

This measures the SIZE of each disagreement region on the real test set, per country, so
the break-even requirement can be applied. It cannot measure their precision (no test
labels) - that is what a submission buys.
"""
import sys
import numpy as np, pandas as pd
sys.path.insert(0, "src")
import config as C

OURS = "output/matching_results.tsv"
HIS = "/scratch/user4/dai/bge-out/runs/full_20260926T001535Z_f7bdb9/output/matching_results.tsv"

def load(p):
    d = pd.read_csv(p, sep="\t", dtype=str, keep_default_na=False)
    print(f"  {p.split('/')[-1]:<22} cols={list(d.columns)} rows={len(d):,}")
    a, b = d.columns[0], d.columns[1]
    out = {}
    for q, v in zip(d[a], d[b]):
        s = {x for x in str(v).split(",") if x}
        if s:
            out[q] = s
    return out

print("loading both submissions")
A, B = load(OURS), load(HIS)
s1 = pd.read_parquet(C.PATHS["test"]["source1"], columns=["entity_id", "country_norm"])
ctry = dict(zip(s1["entity_id"], s1["country_norm"].astype(str)))
allq = list(s1["entity_id"])
print(f"\ntest entities {len(allq):,}   ours non-empty {len(A):,}   his non-empty {len(B):,}")

rows = []
for q in allq:
    a, b = A.get(q, set()), B.get(q, set())
    rows.append((ctry.get(q, "?"), len(a), len(b), len(a & b), len(a - b), len(b - a)))
df = pd.DataFrame(rows, columns=["ctry", "na", "nb", "both", "only_a", "only_b"])

print("\n" + "=" * 78)
print("AGREEMENT STRUCTURE (pairs, per country)")
print("=" * 78)
g = df.groupby("ctry")[["na", "nb", "both", "only_a", "only_b"]].sum()
g["entities"] = df.groupby("ctry").size()
g["jaccard"] = g["both"] / (g["both"] + g["only_a"] + g["only_b"])
g["onlyB_per_entity"] = g["only_b"] / g["entities"]
print(g.to_string())
tot = g.sum(numeric_only=True)
print(f"\n  TOTAL ours {int(tot['na']):,}  his {int(tot['nb']):,}  agree {int(tot['both']):,}")
print(f"  ours-only {int(tot['only_a']):,}   HIS-ONLY {int(tot['only_b']):,}")
print(f"  overall jaccard {tot['both']/(tot['both']+tot['only_a']+tot['only_b']):.4f}")

print("\n" + "=" * 78)
print("WHAT A UNION WOULD ADD, and the precision it must clear")
print("=" * 78)
for c in g.index:
    add = g.at[c, "only_b"]
    cur = g.at[c, "na"]
    n = g.at[c, "entities"]
    print(f"  {c:<8} adds {add:>9,} pairs to {cur:>9,} ({add/max(cur,1):>6.2%} more), "
          f"{add/n:>5.3f}/entity  -> must exceed 80% precision to help")

print("\n" + "=" * 78)
print("HOW MANY OF HIS EXCLUSIVE PAIRS WERE EVEN IN OUR CANDIDATE SET?")
print("=" * 78)
print("  (pairs we never retrieved are the ones an ensemble uniquely unlocks; pairs we")
print("   retrieved and rejected we already priced at p < 0.70 and measured as a losing")
print("   trade at 1.9 TP per FP)")
cp = pd.read_csv("output/candidate_pairs.tsv", sep="\t", dtype=str,
                 keep_default_na=False, nrows=None)
print(f"  our candidate_pairs.tsv: cols={list(cp.columns)} rows={len(cp):,}")
ca, cb = cp.columns[0], cp.columns[1]
cand = {}
for q, v in zip(cp[ca], cp[cb]):
    s = {x for x in str(v).split(",") if x}
    if s:
        cand[q] = s
inc = nev = 0
for q in allq:
    ex = B.get(q, set()) - A.get(q, set())
    if not ex:
        continue
    cs = cand.get(q, set())
    inc += len(ex & cs)
    nev += len(ex - cs)
print(f"  his-exclusive pairs INSIDE our candidates : {inc:,}  (we scored and rejected)")
print(f"  his-exclusive pairs we NEVER retrieved    : {nev:,}  (genuinely new)")