"""Where did our F0.5 actually go, and what is recoverable?

Uses the cached scored val set (calibrated probabilities already computed), so every
number here is measured, not modelled. Answers:

  1. ORACLE  - best possible F0.5 given the candidates we retrieved (the real headroom)
  2. DECOMPOSITION - of the true pairs we miss, how many were never retrieved vs
     retrieved-and-rejected vs lost to one-owner arbitration
  3. OPERATING POINT - F0.5 / precision / recall across the threshold sweep; is there a
     better precision-recall trade than the one we shipped?
  4. EXACT EXPECTED-F0.5 vs our plug-in approximation.

     F0.5 = 1.25*TP / (1.25TP + 0.25FN + FP).  With FP = k - TP and FN = T - TP the TP
     terms cancel in the denominator:   F0.5 = 1.25*TP / (0.25*T + k)
     Our decide.py evaluates that at E[TP] and E[T] - a plug-in. The exact value is
     E[ 1.25*TP / (0.25*(TP + U + M) + k) ] where TP ~ PoissonBinomial(selected probs),
     U ~ PoissonBinomial(unselected probs), M = expected unretrieved true matches.
     Computed here with two Poisson-binomial DPs per entity.
"""
import json, sys, time
import numpy as np
import pandas as pd
import joblib

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

ART = "work/results_A"
CACHE = "work/val_scored_20000.parquet"
t0 = time.time()
log = lambda m: print(f"[{time.time()-t0:5.0f}s] {m}", flush=True)

b = joblib.load(f"{ART}/model_bundle.joblib")
miss_mass = b["miss_mass"]
cand = pd.read_parquet(CACHE)
truths = {int(k): set(v) for k, v in json.load(open(CACHE + ".truth.json")).items()}
ids = sorted(truths)
log(f"{len(ids):,} val entities, {len(cand):,} scored candidate pairs")
log(f"shipped rule: {b['method']}({b['param']})  miss_mass={miss_mass:.4f}")

by_q = {q: g for q, g in cand.groupby("q")}
ctry = {q: g["ctry"].iloc[0] for q, g in by_q.items()}

# ---------------------------------------------------------------- 1. ORACLE
n_true_total = sum(len(truths[q]) for q in ids)
retr, oracle_preds = 0, {}
for q in ids:
    g = by_q.get(q)
    cs = set(g["c"].values.tolist()) if g is not None else set()
    hit = truths[q] & cs
    retr += len(hit)
    if hit:
        oracle_preds[q] = list(hit)
m_or, _ = macro_f05(oracle_preds, truths, ids)
log("")
print("=" * 78)
print("1. HEADROOM")
print("=" * 78)
print(f"  true pairs total                 : {n_true_total:,}")
print(f"  retrieved by blocking            : {retr:,}  ({retr/n_true_total:.4f} pair recall)")
print(f"  ORACLE F0.5 on our candidates    : {m_or['f05']:.4f}   <- ceiling for the scorer")
print(f"  perfect blocking + perfect scorer: 1.0000")

# ---------------------------------------------------------------- current
shipped = DC.predict_sets(cand, "q", "c", "p", b["method"], b["param"], miss_mass, True)
m_now, per_now = macro_f05(shipped, truths, ids)
print(f"  SHIPPED F0.5                     : {m_now['f05']:.4f}  "
      f"(P {m_now['precision']:.4f}  R {m_now['recall']:.4f})")
print(f"  gap to oracle                    : {m_or['f05']-m_now['f05']:.4f}  <- scorer/decision loss")
print(f"  gap oracle to 1.0                : {1-m_or['f05']:.4f}  <- blocking loss")

# ---------------------------------------------------------------- 2. DECOMPOSITION
never = rejected = 0
fp = 0
for q in ids:
    g = by_q.get(q)
    cs = set(g["c"].values.tolist()) if g is not None else set()
    pred = set(shipped.get(q, []))
    t = truths[q]
    never += len(t - cs)
    rejected += len((t & cs) - pred)
    fp += len(pred - t)
print()
print("=" * 78)
print("2. WHERE THE MISSES ARE")
print("=" * 78)
miss = never + rejected
print(f"  missed true pairs                : {miss:,}")
print(f"    never retrieved (blocking)     : {never:,}  ({never/max(miss,1):.1%})")
print(f"    retrieved but rejected (scorer): {rejected:,}  ({rejected/max(miss,1):.1%})")
print(f"  false positives predicted        : {fp:,}")

# ---------------------------------------------------------------- 3. OPERATING POINT
print()
print("=" * 78)
print("3. OPERATING POINT  (is precision over-weighted in practice?)")
print("=" * 78)
print(f"{'threshold':>10} {'F0.5':>9} {'precision':>10} {'recall':>9} {'avg preds':>10}")
best = (0, None)
for thr in np.round(np.arange(0.30, 0.96, 0.05), 2):
    p = DC.predict_sets(cand, "q", "c", "p", "threshold", thr, miss_mass, True)
    m, _ = macro_f05(p, truths, ids)
    n_pred = np.mean([len(p.get(q, [])) for q in ids])
    star = ""
    if m["f05"] > best[0]:
        best = (m["f05"], thr); star = ""
    print(f"{thr:>10.2f} {m['f05']:>9.5f} {m['precision']:>10.4f} {m['recall']:>9.4f} {n_pred:>10.2f}")
print(f"  best threshold {best[1]} -> {best[0]:.5f}")

# ------------------------------------------------- 4. EXACT EXPECTED F0.5 (DP)
def pb_dist(ps):
    """Poisson-binomial pmf over sum of independent Bernoulli(ps)."""
    d = np.zeros(len(ps) + 1); d[0] = 1.0
    for p in ps:
        d[1:] = d[1:] * (1 - p) + d[:-1] * p
        d[0] *= (1 - p)
    return d


def exact_select(probs, M, kmax=None):
    """Choose k maximising the EXACT E[F0.5] = E[1.25*TP / (0.25*(TP+U+M) + k)]."""
    o = np.argsort(-probs)
    ps = probs[o]
    n = len(ps)
    kmax = n if kmax is None else min(kmax, n)
    # suffix Poisson-binomial for the unselected tail, built incrementally
    best_k, best_v = 0, 0.0
    # k = 0 -> F0.5 = 0 unless the entity truly has no matches; the empty set scores
    # 1.0 only when T == 0, whose probability is P(all candidates false) * P(M==0).
    tail = pb_dist(ps)                      # all unselected when k = 0
    p_T0 = tail[0] * (1.0 if M <= 0 else np.exp(-M))
    best_v = p_T0                            # E[F0.5 | k=0] = P(T = 0) * 1
    head = np.array([1.0])
    for k in range(1, kmax + 1):
        head = np.convolve(head, [1 - ps[k - 1], ps[k - 1]])
        tail = pb_dist(ps[k:]) if k < n else np.array([1.0])
        v = 0.0
        for tp in range(len(head)):
            if head[tp] <= 1e-12:
                continue
            for u in range(len(tail)):
                if tail[u] <= 1e-12:
                    continue
                denom = 0.25 * (tp + u + M) + k
                if denom > 0:
                    v += head[tp] * tail[u] * (1.25 * tp / denom)
        if v > best_v:
            best_v, best_k = v, k
    return o[:best_k]


print()
print("=" * 78)
print("4. EXACT EXPECTED-F0.5 vs OUR PLUG-IN")
print("=" * 78)
d = DC.resolve_conflicts(cand, "p")
preds_exact = {}
t1 = time.time()
for q, g in d.groupby("q"):
    pr = g["p"].values
    if len(pr) > 30:
        keep = np.argsort(-pr)[:30]
        g = g.iloc[keep]; pr = g["p"].values
    idx = exact_select(pr, miss_mass)
    if len(idx):
        preds_exact[q] = list(g["c"].values[idx])
m_ex, _ = macro_f05(preds_exact, truths, ids)
log(f"exact DP over {len(ids):,} entities took {time.time()-t1:.0f}s")
m_ef, _ = macro_f05(DC.predict_sets(cand, "q", "c", "p", "ef05", 0.0, miss_mass, True), truths, ids)
print(f"  plug-in ef05(0.0)   F0.5 {m_ef['f05']:.5f}  P {m_ef['precision']:.4f}  R {m_ef['recall']:.4f}")
print(f"  shipped threshold   F0.5 {m_now['f05']:.5f}  P {m_now['precision']:.4f}  R {m_now['recall']:.4f}")
print(f"  EXACT DP            F0.5 {m_ex['f05']:.5f}  P {m_ex['precision']:.4f}  R {m_ex['recall']:.4f}")
print(f"  exact vs shipped    {m_ex['f05']-m_now['f05']:+.5f}")

# ---------------------------------------------------------------- 5. BY COUNTRY
print()
print("=" * 78)
print("5. BY COUNTRY (shipped rule)")
print("=" * 78)
for c in sorted(set(ctry.values())):
    sub = [q for q in ids if ctry.get(q) == c]
    m, _ = macro_f05({q: shipped.get(q, []) for q in sub}, truths, sub)
    print(f"  {c:<8} n={len(sub):>6}  F0.5 {m['f05']:.4f}  P {m['precision']:.4f}  R {m['recall']:.4f}")