"""THE LAST UNTESTED LEVER: S2<->S3 corroboration.

Two independently-trained cross-encoders - bge-reranker-v2-m3 (XLM-R 568M) and Chinmay's
mDeBERTa-v3 (278M) - reject the SAME true pairs. Blending them gained +0.00032 held out, and of
the true pairs the first rejects the second promotes 7.71% while also promoting 0.0602% of the
negatives: a 0.52:1 trade against a ~3.76:1 break-even.

I concluded from that those pairs were irreducibly ambiguous. That was WRONG, and the public
leaderboard maximum of 0.990788 proves it. Both models read ONLY "name | addr | country" for two
records. What was shown is that the PAIRWISE signal is exhausted at that scale. Collective
evidence is structurally invisible to a pairwise scorer.

THE IDEA (Model_Harsh.md s10 / Model-5 s10, never implemented). True matches come in groups:
S1-q matches some S2-x AND some S3-y, and x and y describe the SAME business, so they resemble
EACH OTHER. A false candidate that merely resembles q is unlikely to also resemble q's confident
match from the other source.

    corrob(q,c) = max cosine( emb[c], emb[a] )  over q's ACCEPTED records from the OTHER source

Label-free (anchors are our own predictions), feature-only (never invents a candidate), and
self-excluding (a record is never corroborated against itself - that circularity once made a true
match its own distractor and read recall@1 as 53.87% instead of 95.01%).

THE GATE: among pairs we currently REJECT, does corroboration buy true pairs at better than
~3.76 TP per FP? Below that it lowers macro F0.5 no matter how sensible it looks.
"""
import json
import sys
import time

import joblib
import numpy as np
import pandas as pd

sys.path.insert(0, "src")
import config as C
import decide as DC
from evaluate import macro_f05

N_S2 = 5_034_616
THR = 0.70
t0 = time.time()


def log(m):
    print(f"[{time.time()-t0:6.0f}s] {m}", flush=True)


b = joblib.load("work/results_A/model_bundle.joblib")
MM = b["miss_mass"]
CACHE = "work/val_scored_20000.parquet"
cand = pd.read_parquet(CACHE)
truths = {int(k): set(v) for k, v in json.load(open(CACHE + ".truth.json")).items()}
ids = sorted(truths)
shipped = DC.predict_sets(cand, "q", "c", "p", b["method"], THR, MM, True)
m0, _ = macro_f05(shipped, truths, ids)
log(f"baseline F0.5 {m0['f05']:.5f}  P {m0['precision']:.4f} R {m0['recall']:.4f}")

cand = cand.copy()
cand["y"] = [1 if c in truths.get(q, ()) else 0 for q, c in zip(cand["q"], cand["c"])]
sel = {(q, c) for q, cs in shipped.items() for c in cs}
cand["sel"] = [1 if (q, c) in sel else 0 for q, c in zip(cand["q"], cand["c"])]

need = np.unique(cand["c"].to_numpy())
log(f"distinct pool records needed: {len(need):,}")
# np.load, NOT np.memmap: these are .npy files with a 128-byte header, which memmap maps as
# data and which made a raw reshape fail by exactly 64 float16 elements. np.load parses the
# header and returns the true shape and dtype instead of us assuming them.
s2 = np.load(C.HYBRID["train_s2"], mmap_mode="r")
s3 = np.load(C.HYBRID["train_s3"], mmap_mode="r")
log(f"s2 {s2.shape} {s2.dtype}   s3 {s3.shape} {s3.dtype}")
assert s2.ndim == 2 and s2.shape[1] == 1024, s2.shape
assert s2.shape[0] == N_S2, s2.shape
E = np.empty((len(need), 1024), np.float32)
CH = 50_000                          # small blocks: this box is shared and heavily loaded
for i in range(0, len(need), CH):
    nd = need[i:i + CH]              # already sorted -> sequential-ish page access
    blk = E[i:i + CH]
    m2 = nd < N_S2
    if m2.any():
        blk[m2] = np.asarray(s2[nd[m2]], np.float32)
    if (~m2).any():
        blk[~m2] = np.asarray(s3[nd[~m2] - N_S2], np.float32)
    if (i // CH) % 4 == 0:
        log(f"  loaded {min(i+CH, len(need)):,}/{len(need):,}")
E /= np.maximum(np.linalg.norm(E, axis=1, keepdims=True), 1e-9)
pos = {int(c): i for i, c in enumerate(need)}
log("embeddings loaded, unit-normalised")

corr = np.full(len(cand), np.nan, np.float32)
qv = cand["q"].to_numpy()
cv = cand["c"].to_numpy()
sv = cand["sel"].to_numpy()
pv = cand["p"].to_numpy()
order = np.argsort(qv, kind="stable")
uq = np.unique(qv)
bounds = np.r_[np.searchsorted(qv[order], uq), len(order)]
n_anchor = 0
for i in range(len(uq)):
    idx = order[bounds[i]:bounds[i + 1]]
    acc = idx[sv[idx] == 1]
    if not len(acc):
        continue
    # confident anchors only, highest first: corroborating against a shaky acceptance
    # propagates its error to everything it touches
    acc = acc[np.argsort(-pv[acc])]
    acc = acc[pv[acc] >= 0.90][:4]
    if not len(acc):
        continue
    a2 = np.array([pos[int(c)] for c in cv[acc] if int(c) < N_S2], np.int64)
    a3 = np.array([pos[int(c)] for c in cv[acc] if int(c) >= N_S2], np.int64)
    if not len(a2) and not len(a3):
        continue
    n_anchor += 1
    for j in idx:
        A = a3 if cv[j] < N_S2 else a2          # the OTHER source's accepted records
        if not len(A):
            continue
        me = pos[int(cv[j])]
        A = A[A != me]
        if not len(A):
            continue
        corr[j] = float((E[A] @ E[me]).max())
cand["corr"] = corr
log(f"entities with a usable cross-source anchor: {n_anchor:,}/{len(uq):,}")

un = cand[(cand.sel == 0) & cand["corr"].notna()]
rt, rn = un[un.y == 1], un[un.y == 0]
print(f"\nunselected pairs with a corroboration value: {len(un):,} "
      f"({len(rt):,} true, {len(rn):,} false)")
print(f"  rejected TRUE median corr {rt['corr'].median():.4f}   "
      f"rejected FALSE median corr {rn['corr'].median():.4f}")

print("\n=== THE GATE: exchange rate of accepting by corroboration ===")
print(f"{'corr>=':>8}{'TP':>8}{'FP':>9}{'ratio':>8}   (needs > 3.76)")
for t in (0.80, 0.85, 0.88, 0.90, 0.92, 0.94, 0.95, 0.96, 0.97, 0.98, 0.99):
    tp = int((rt["corr"] >= t).sum())
    fp = int((rn["corr"] >= t).sum())
    r = tp / fp if fp else float("inf")
    flag = "   <-- CLEARS BREAK-EVEN" if r > 3.76 else ""
    print(f"{t:>8.2f}{tp:>8,}{fp:>9,}{r:>8.2f}{flag}")

print("\n=== combined with a pairwise-score floor ===")
print(f"{'p>=':>6}{'corr>=':>8}{'TP':>7}{'FP':>8}{'ratio':>7}")
for pt in (0.2, 0.3, 0.4, 0.5, 0.6):
    for ct in (0.90, 0.94, 0.96, 0.97, 0.98):
        tp = int(((rt["corr"] >= ct) & (rt.p >= pt)).sum())
        fp = int(((rn["corr"] >= ct) & (rn.p >= pt)).sum())
        if tp + fp == 0:
            continue
        r = tp / fp if fp else float("inf")
        flag = "   <-- CLEARS" if r > 3.76 else ""
        print(f"{pt:>6.2f}{ct:>8.2f}{tp:>7,}{fp:>8,}{r:>7.2f}{flag}")

print("\n=== END TO END, held out on half the entities ===")
h = np.array([hash(("half", int(q))) % 2 for q in ids])
HA = set(np.array(ids)[h == 0].tolist())
HB = [q for q in ids if q not in HA]
HAl = sorted(HA)


def ev(pr, idl):
    return macro_f05({q: pr.get(q, []) for q in idl}, truths, idl)[0]["f05"]


bA, bB = ev(shipped, HAl), ev(shipped, HB)
print(f"baseline  half A {bA:.5f}   half B {bB:.5f}")
best = (-1, None, None)
for pt in (0.2, 0.3, 0.4, 0.5, 0.6):
    for ct in (0.90, 0.94, 0.96, 0.97, 0.98):
        add = cand[(cand.sel == 0) & (cand["corr"] >= ct) & (cand.p >= pt)]
        if not len(add):
            continue
        aug = DC.resolve_conflicts(pd.concat([cand[cand.sel == 1], add]), "p")
        pr = {}
        for q, c in zip(aug["q"], aug["c"]):
            pr.setdefault(int(q), []).append(int(c))
        a, bb = ev(pr, HAl), ev(pr, HB)
        if a > best[0]:
            best = (a, (pt, ct, len(add)), bb)
        print(f"  p>={pt:.2f} corr>={ct:.2f}: +{len(add):>6,} -> A {a:.5f} B {bb:.5f}")
if best[1]:
    pt, ct, n = best[1]
    print(f"\nbest on half A: p>={pt} corr>={ct} (+{n:,} pairs)")
    print(f"  half A (in-sample) {best[0]:.5f}  ({best[0]-bA:+.5f})")
    print(f"  half B (HELD OUT)  {best[2]:.5f}  ({best[2]-bB:+.5f})   <-- the honest number")