"""Infer France's operating point from the ONE France measurement we have: 0.976423.

We know, measured:
  - val (US/India) per-country F0.5 at threshold 0.70
  - the test country mix
  - the full per-pair probability distribution for every test entity (saved chunks)
So France's realised F0.5 is identified by arithmetic. Then, holding France's TRUE match
distribution at the US/India shape (justified: mean matches/entity is close), a single
scalar 'domain-shift' parameter g explains the shortfall via  true_prec(p) = p**g.
g is fitted to reproduce 0.976423 at threshold 0.70 - no free parameters left - and the
optimal France threshold under that g is then read off directly.
"""
import sys, glob, json
import numpy as np, pandas as pd, joblib
sys.path.insert(0, "src")
import config as C, decide as DC

b = joblib.load("work/results_A/model_bundle.joblib")
MM, THR = b["miss_mass"], b["param"]
print(f"shipped rule: {b['method']}({THR})  miss_mass={MM:.4f}")

s1 = pd.read_parquet(C.PATHS["test"]["source1"], columns=["entity_id", "country_norm"])
ctry = s1["country_norm"].astype(str).values
N = len(s1)
print(f"test S1 entities: {N:,}")
u, cnt = np.unique(ctry, return_counts=True)
W = {k: v / N for k, v in zip(u, cnt)}
for k in u:
    print(f"  {k:<8} {dict(zip(u,cnt))[k]:>9,}  weight {W[k]:.5f}")

files = sorted(glob.glob("work/results_A/test_chunks/chunk_*.parquet"))
d = pd.concat([pd.read_parquet(f) for f in files], ignore_index=True)
print(f"\nscored test pairs: {len(d):,}")
d = DC.resolve_conflicts(d, "p")           # same global one-owner arbitration as shipped
print(f"after one-owner arbitration: {len(d):,}")
d["ctry"] = ctry[d["q"].values]

# ---- realised France F0.5 by arithmetic -----------------------------------------
LB = 0.976423
VAL = {"us": 0.9832, "india": 0.9806}      # measured on 20k held-out val entities
wus, win, wfr = W.get("us", 0), W.get("india", 0), W.get("france", 0)
usin = wus * VAL["us"] + win * VAL["india"]
fr_implied = (LB - usin) / wfr
print("\n" + "=" * 74)
print("FRANCE, BY ARITHMETIC")
print("=" * 74)
print(f"  leaderboard                      : {LB:.6f}")
print(f"  US+India contribution (val rates): {usin:.6f}   (weight {wus+win:.5f})")
print(f"  => implied FRANCE F0.5           : {fr_implied:.6f}   (weight {wfr:.5f})")
print(f"  US/India level                   : {usin/(wus+win):.6f}")
print(f"  FRANCE DEFICIT                   : {fr_implied - usin/(wus+win):+.6f}")

# ---- per-country prediction counts across thresholds ----------------------------
print("\n" + "=" * 74)
print("PREDICTION COUNTS PER ENTITY BY THRESHOLD (measured on the real test set)")
print("=" * 74)
grid = [0.50, 0.60, 0.70, 0.75, 0.80, 0.85, 0.90, 0.95]
counts = {}
print(f"{'thr':>6} " + "".join(f"{c:>12}" for c in u) + f"{'  frac-empty(fr)':>18}")
nq = {c: int((ctry == c).sum()) for c in u}
for t in grid:
    s = d[d.p >= t]
    row, cc = {}, s.groupby("ctry").size()
    line = f"{t:>6.2f} "
    for c in u:
        k = cc.get(c, 0) / nq[c]
        row[c] = k
        line += f"{k:>12.4f}"
    # fraction of French entities predicted empty
    fr_nonempty = s[s.ctry == "france"]["q"].nunique()
    line += f"{1 - fr_nonempty/nq['france']:>18.4f}"
    counts[t] = row
    print(line)

# ---- fit the domain-shift exponent g -------------------------------------------
# For each entity: pick candidates with p>=t. Under true_prec(p)=p**g the expected TP is
# sum(p**g) over selected, expected T is that plus the unselected mass plus miss_mass.
print("\n" + "=" * 74)
print("DOMAIN-SHIFT FIT:  true_prec(p) = p**g   (g=1 means calibration transfers)")
print("=" * 74)

fr = d[d.ctry == "france"][["q", "p"]].copy()
fr_p = fr["p"].values.astype(np.float64)
fr_q = fr["q"].values
order = np.argsort(fr_q, kind="stable")
fr_q, fr_p = fr_q[order], fr_p[order]
bounds = np.searchsorted(fr_q, np.unique(fr_q))
uq = np.unique(fr_q)
n_fr = nq["france"]
n_empty_all = n_fr - len(uq)          # French entities with NO candidate at all

def france_f05(g, t):
    """Expected macro F0.5 over French entities under true_prec(p)=p**g."""
    pg = fr_p ** g
    sel = fr_p >= t
    tot = 0.0
    idx = np.r_[bounds, len(fr_q)]
    for i in range(len(uq)):
        a, z = idx[i], idx[i + 1]
        m = sel[a:z]
        k = int(m.sum())
        tp = float(pg[a:z][m].sum())
        unsel = float(pg[a:z][~m].sum())
        T = tp + unsel + MM
        if k == 0:
            # empty prediction scores 1.0 exactly when the entity truly has no match
            tot += float(np.exp(-T)) if T > 0 else 1.0
        else:
            tot += 1.25 * tp / (0.25 * T + k)
    tot += n_empty_all * 1.0          # no candidates at all -> empty -> assume singleton
    return tot / n_fr

lo, hi = 1.0, 6.0
for _ in range(40):
    mid = (lo + hi) / 2
    if france_f05(mid, THR) > fr_implied:
        lo = mid
    else:
        hi = mid
g = (lo + hi) / 2
print(f"  fitted g = {g:.4f}   (reproduces France {france_f05(g, THR):.6f} "
      f"vs implied {fr_implied:.6f})")
print(f"  sanity: at g=1 the model would expect France {france_f05(1.0, THR):.6f}")

print(f"\n  France F0.5 vs threshold under the fitted g={g:.3f}:")
print(f"{'thr':>7}{'France F0.5':>14}{'delta vs 0.70':>15}{'preds/entity':>14}{'LB effect':>12}")
base = france_f05(g, THR)
best = (base, THR)
for t in [0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90, 0.93, 0.95, 0.97]:
    v = france_f05(g, t)
    if v > best[0]:
        best = (v, t)
    print(f"{t:>7.2f}{v:>14.6f}{v-base:>+15.6f}{counts.get(t, {}).get('france', float('nan')):>14.4f}"
          f"{wfr*(v-base):>+12.6f}")
print(f"\n  BEST France threshold {best[1]:.2f} -> {best[0]:.6f} "
      f"({best[0]-base:+.6f} on France, {wfr*(best[0]-base):+.6f} on the leaderboard)")
json.dump({"g": g, "fr_implied": fr_implied, "best_thr": best[1],
           "best_fr": best[0], "lb_gain": wfr * (best[0] - base)},
          open("work/fr_operating.json", "w"), indent=1)