"""Does a SECOND cross-encoder close part of the 0.0118 separation gap?

Measured facts this targets:
  - oracle on our candidates 0.9942 vs shipped 0.98214 -> 0.0121 scorer loss
  - of that, only 0.00031 is within-entity RANKING (oracle rerank at fixed k)
  - so ~0.0118 is SEPARATION: true pairs sitting below 0.70 with negatives interleaved
Better separation is what a second, architecturally different cross-encoder can supply
(bge-reranker-v2-m3 is XLM-R based; Chinmay's is mDeBERTa-v3). Same pair_text
serialization, verified identical, so his checkpoint applies directly to our pairs.

METHOD, leak-free: split the 20,000 val entities in half by a fixed hash. Fit his score's
isotonic calibration, the blend weight and the threshold on half A; report on half B only.
Half A numbers are shown too, marked in-sample.
"""
import json, sys, os, time
import numpy as np, pandas as pd, joblib
sys.path.insert(0, "src")
import config as C, decide as DC
from evaluate import macro_f05
from cross_encoder import pair_text, predict_cross_encoder
from sklearn.isotonic import IsotonicRegression

CE2 = "/scratch/user4/dai/bge-out/cache/ce_model_c8a595de7fa1"
CACHE = "work/val_scored_20000.parquet"
OUT = "work/val_ce2_scores.npy"
b = joblib.load("work/results_A/model_bundle.joblib")
MM, THR = b["miss_mass"], b["param"]

cand = pd.read_parquet(CACHE)
truths = {int(k): set(v) for k, v in json.load(open(CACHE + ".truth.json")).items()}
ids = sorted(truths)
print(f"val entities {len(ids):,}  scored pairs {len(cand):,}")

p = C.PATHS["train"]
cols = ["entity_id", "business_name", "business_address", "country_norm"]
s1 = pd.read_parquet(p["source1"], columns=cols)
pool = pd.concat([pd.read_parquet(p[k], columns=cols) for k in ("source2", "source3")],
                 ignore_index=True)
print(f"s1 {len(s1):,}  pool {len(pool):,}")

qi, ci = cand["q"].values, cand["c"].values
s1n, s1a, s1c = (s1["business_name"].values, s1["business_address"].values,
                 s1["country_norm"].values)
pn, pa, pc = (pool["business_name"].values, pool["business_address"].values,
              pool["country_norm"].values)

if os.path.exists(OUT):
    p2 = np.load(OUT)
    assert len(p2) == len(cand), "cached ce2 scores do not match the candidate table"
    print(f"reusing cached ce2 scores ({len(p2):,})")
else:
    ta = [pair_text(s1n[r], s1a[r], s1c[r]) for r in qi]
    tb = [pair_text(pn[c], pa[c], pc[c]) for c in ci]
    print("sample A:", ta[0][:110])
    print("sample B:", tb[0][:110])
    cfg = dict(C.CFG)
    cfg["ce_max_len"] = 160
    cfg["ce_eval_batch"] = 512
    cfg["ce_num_workers"] = 0
    t0 = time.time()
    p2 = predict_cross_encoder(ta, tb, CE2, cfg, cache="work/ce2_pred")
    print(f"scored {len(p2):,} pairs with the mDeBERTa CE in {time.time()-t0:.0f}s")
    np.save(OUT, p2)

cand = cand.copy()
cand["p2"] = p2.astype(np.float32)
cand["y"] = [1 if c in truths.get(q, ()) else 0 for q, c in zip(cand["q"], cand["c"])]

# leak-free half split on entities
h = np.array([hash(("split", int(q))) % 2 for q in ids])
A = set(np.array(ids)[h == 0].tolist())
B = set(np.array(ids)[h == 1].tolist())
print(f"\nfit half A: {len(A):,} entities   report half B: {len(B):,} entities")
inA = cand["q"].isin(A).values
ca, cb = cand[inA], cand[~inA]
idsA, idsB = sorted(A), sorted(B)

def ev(df, idl, col):
    pr = DC.predict_sets(df, "q", "c", col, "threshold", THR, MM, True)
    m, _ = macro_f05(pr, truths, idl)
    return m

print("\n--- baseline (our bge reranker stack, shipped threshold 0.70) ---")
mA, mB = ev(ca, idsA, "p"), ev(cb, idsB, "p")
print(f"  half A {mA['f05']:.5f}   half B {mB['f05']:.5f}")

# his CE alone, isotonic-calibrated on A
iso = IsotonicRegression(out_of_bounds="clip", y_min=0.0, y_max=1.0)
iso.fit(ca["p2"].values, ca["y"].values)
cand["p2c"] = iso.predict(cand["p2"].values).astype(np.float32)
ca, cb = cand[inA], cand[~inA]
print("\n--- mDeBERTa CE alone (isotonic fit on A) ---")
print(f"  half A {ev(ca,idsA,'p2c')['f05']:.5f}   half B {ev(cb,idsB,'p2c')['f05']:.5f}")

print("\n--- BLEND: (1-w)*ours + w*theirs, both calibrated ---")
print(f"{'w':>6}{'thr':>7}{'half A (in-sample)':>20}{'half B (reported)':>20}")
best = (-1, None, None)
for w in [0.0, 0.1, 0.15, 0.2, 0.25, 0.3, 0.4, 0.5]:
    cand["pb"] = (1 - w) * cand["p"] + w * cand["p2c"]
    ca, cb = cand[inA], cand[~inA]
    # tune the threshold on A only
    bt, bv = THR, -1
    for t in np.round(np.arange(0.50, 0.93, 0.025), 3):
        pr = DC.predict_sets(ca, "q", "c", "pb", "threshold", t, MM, True)
        v = macro_f05(pr, truths, idsA)[0]["f05"]
        if v > bv:
            bt, bv = t, v
    prB = DC.predict_sets(cb, "q", "c", "pb", "threshold", bt, MM, True)
    vB = macro_f05(prB, truths, idsB)[0]["f05"]
    print(f"{w:>6.2f}{bt:>7.3f}{bv:>20.5f}{vB:>20.5f}")
    if vB > best[0]:
        best = (vB, w, bt)
print(f"\n  BEST on held-out half B: w={best[1]} thr={best[2]} -> {best[0]:.5f}")
print(f"  baseline on half B                             : {mB['f05']:.5f}")
print(f"  DELTA (held-out, honest)                       : {best[0]-mB['f05']:+.5f}")

# does the second CE rescue the pairs we reject?
rej = cand[(cand.y == 1) & (cand.p < THR)]
print(f"\n--- the {len(rej):,} true pairs our stack puts below 0.70 ---")
print(f"  their mDeBERTa calibrated p: median {rej['p2c'].median():.4f}, "
      f"frac >= 0.70: {(rej['p2c']>=0.70).mean():.2%}")
neg = cand[(cand.y == 0) & (cand.p < THR)]
print(f"  for comparison, {len(neg):,} TRUE NEGATIVES below 0.70: "
      f"frac >= 0.70: {(neg['p2c']>=0.70).mean():.4%}")
tp = int((rej['p2c'] >= 0.70).sum()); fp = int((neg['p2c'] >= 0.70).sum())
print(f"  => rescue trade: {tp:,} TP vs {fp:,} FP  = {tp/max(fp,1):.2f}:1 "
      f"(needs > 4:1 to gain)")