"""Does reranking only the top-M candidates per entity cost any F0.5?

The test stage reranks all ~45 candidates/entity (80M pairs, ~12 h). But the answer
averages 3.46 matches and expected-F0.5 selection stops early on a precision-weighted
metric, so deep-rank candidates may never be selected. If so, reranking only the
top-M by a RERANKER-FREE ranker is free speed.

Method (paired - same entities, same models, only the reranking scope differs):
  full      : p = iso(lgbm_ce(features + ce_prob))            for every candidate
  cascade-M : rank by out-of-fold lgbm_noce; keep top-M/entity;
              those keep p; the rest are dropped (never rerankable, so never chosen)
  then identical calibration / conflict resolution / expected-F0.5 selection.

Paired on the same entities, so the comparison is far more sensitive than two
independent runs. Ship only at delta == 0.
"""
import os, sys, time, json
import numpy as np
import pandas as pd
import joblib
import lightgbm as lgb
from sklearn.isotonic import IsotonicRegression

sys.path.insert(0, "src")
import config as cfgmod
from config import CFG, EMB, PATHS, emb_dir
import blocking as B
import data as D
import decide as DC
import text_utils as T
import emb_source
from encoder import load_matrix, read_ids, Encoder
from features import compute_features
from evaluate import macro_f05

N_VAL = int(sys.argv[1]) if len(sys.argv) > 1 else 10000
GPU = sys.argv[2] if len(sys.argv) > 2 else "cuda:0"
if not str(GPU).startswith("cuda") and GPU != "cpu":
    GPU = f"cuda:{int(GPU)}"
MS = [4, 6, 8, 10, 12, 16, 20, 25]
ART = "work/results_A"
t0 = time.time()
log = lambda m: print(f"[{time.time()-t0:6.0f}s] {m}", flush=True)

bundle = joblib.load(os.path.join(ART, "model_bundle.joblib"))
method, param, miss_mass = bundle["method"], bundle["param"], bundle["miss_mass"]
feat_cols, lgbm_ce = bundle["feat_cols"], bundle["lgbm"]
amap = json.load(open(os.path.join(ART, "abbrev_map.json")))
ce_dir = bundle["ce_dir"]
log(f"bundle: {method}({param}) miss_mass={miss_mass:.4f} feats={len(feat_cols)}")

log("loading train data")
s1, s23, gt = D.load_data(PATHS["train"], with_gt=True)
s1["name_core_abbr"] = [T.apply_abbrev(x, amap) for x in s1["name_core"].values]
s23["name_core_abbr"] = [T.apply_abbrev(x, amap) for x in s23["name_core"].values]

def vals(df):
    v = {c: df[c].values.astype(object) for c in
         ["entity_id", "name_core_abbr", "name_norm", "address_abbrev_norm",
          "legal_suffix", "country_norm"]}
    nums = df["address_numbers"].values.astype(object)
    v["nums"] = np.where(nums == "", df["business_address"].values.astype(object), nums)
    return v
s1v, s23v = vals(s1), vals(s23)

pool_pos = {e: i for i, e in enumerate(s23["entity_id"].values)}
s1 = D.augment_s1(s1, gt)            # gt is a DataFrame; this is how the pipeline does it
lists = s1["matched_entity_ids"].map(D.split_ids)
truth_rows = {}
for r, ms in enumerate(lists):
    if ms:
        truth_rows[r] = set(pool_pos[m] for m in ms if m in pool_pos)
_npairs = sum(len(v) for v in truth_rows.values())
log(f"ground truth: {len(truth_rows):,} entities with matches, {_npairs:,} true pairs")
assert _npairs > 5_000_000, f"ground truth looks wrong ({_npairs} pairs) - expected ~7.6M"

rng = np.random.RandomState(CFG["seed"])
ce_rows, st_rows, va_rows = D.make_split(s1, CFG, rng)
va_rows = np.sort(rng.choice(va_rows, min(N_VAL, len(va_rows)), replace=False))
log(f"using {len(va_rows):,} val entities (same split/seed as the real run)")

# ------------------------------------------------------------------ candidates
c_emb, cmeta = load_matrix(emb_dir("train_pool"))
log(f"pool emb {c_emb.shape}")
enc = Encoder(CFG["emb_model"], GPU, EMB["max_len"], EMB["batch"])
q_emb = enc.encode([B.record_text(s1["business_name"].values[i],
                                  s1["business_address"].values[i],
                                  s1["country"].values[i]) for i in va_rows],
                   CFG["emb_query_prefix"])
del enc
emb_s, emb_i = B.dense_topk(np.asarray(q_emb), c_emb, CFG, GPU)

emb2_s = emb2_i = None
if CFG.get("emb2_enabled"):
    c2, q2 = emb_source.load_emb2("train", len(s1), len(s23))
    cfg2 = dict(CFG); cfg2["emb_k"] = CFG["emb2_k"]
    emb2_s, emb2_i = B.dense_topk(np.asarray(q2[va_rows]), c2, cfg2, GPU)
    del c2, q2
log("dense channels done")

lidx = B.LexicalIndex([B.lexical_tokens(n, a) for n, a in
                       zip(s23v["name_core_abbr"], s23v["address_abbrev_norm"])], CFG)
lex = B.lexical_topk(lidx, [B.lexical_tokens(s1v["name_core_abbr"][i],
                                             s1v["address_abbrev_norm"][i])
                            for i in va_rows], CFG["lex_n_jobs"])
del lidx
cand = B.build_candidates(va_rows, emb_s, emb_i, lex, s1v["country_norm"],
                          s23v["country_norm"], CFG["country_filter"],
                          emb2_s=emb2_s, emb2_i=emb2_i)
qpos = {r: i for i, r in enumerate(va_rows.tolist())}
cand = B.fill_missing_emb(cand, lambda rows: q_emb[[qpos[r] for r in rows]], c_emb)
cand["y"] = np.fromiter((c in truth_rows.get(q, ()) for q, c in
                         zip(cand["q"].values, cand["c"].values)), np.int8, len(cand))
log(f"{len(cand):,} candidates ({len(cand)/len(va_rows):.1f}/entity)")

feats = compute_features(cand, s1v, s23v, CFG["feat_n_jobs"])
cand = pd.concat([cand, feats], axis=1); del feats
log("features done")

# ------------------------------------------------------------------ reranker
from cross_encoder import pair_text, predict_cross_encoder
ta = [pair_text(s1["business_name"].values[r], s1["business_address"].values[r],
                s1["country"].values[r]) for r in cand["q"].values]
tb = [pair_text(s23["business_name"].values[c], s23["business_address"].values[c],
                s23["country"].values[c]) for c in cand["c"].values]
cand["ce_prob"] = predict_cross_encoder(ta, tb, ce_dir, CFG)
del ta, tb
log("reranker scoring done")

# ------------------------------------------------------ CE-free pre-ranker (OOF)
noce_cols = [c for c in feat_cols if c != "ce_prob"]
oof = np.zeros(len(cand))
groups = cand["q"].values
uq = np.unique(groups)
fold_of = {q: i % 5 for i, q in enumerate(rng.permutation(uq))}
fold = np.array([fold_of[q] for q in groups])
for f in range(5):
    tr, te = fold != f, fold == f
    m = lgb.train({**CFG["lgbm_params"], "verbose": -1},
                  lgb.Dataset(cand.loc[tr, noce_cols], cand.loc[tr, "y"]),
                  num_boost_round=300)
    oof[te] = m.predict(cand.loc[te, noce_cols])
cand["noce"] = oof
log("out-of-fold CE-free pre-ranker built")

# ------------------------------------------------------------------ compare
cand["p"] = lgbm_ce.predict(cand[feat_cols])
cand["p"] = bundle["calibrator"].predict(cand["p"].values)   # the real calibrator

truths = {q: truth_rows.get(q, set()) for q in va_rows}
ids = list(va_rows)

full_preds = DC.predict_sets(cand, "q", "c", "p", method, param, miss_mass, True)
m_full, per_full = macro_f05(full_preds, truths, ids)
log(f"FULL      F0.5 {m_full['f05']:.6f}  P {m_full['precision']:.4f}  R {m_full['recall']:.4f}")
if not (0.93 <= m_full["f05"] <= 0.995):
    sys.exit(f"SANITY FAIL: full-pipeline F0.5 {m_full['f05']:.4f} is nowhere near the "
             f"0.9812 the real run reported - the harness is wrong, not the cascade")

cand["rk"] = cand.groupby("q")["noce"].rank(ascending=False, method="first")
print()
print(f"{'M':>4} {'F0.5':>10} {'delta':>11} {'P':>8} {'R':>8} {'rerank%':>9} {'selected outside top-M':>24}")
print(f"{'full':>4} {m_full['f05']:>10.6f} {0.0:>11.6f} {m_full['precision']:>8.4f} "
      f"{m_full['recall']:>8.4f} {100.0:>8.1f}% {0:>24}")
for M in MS:
    sub = cand[cand["rk"] <= M]
    preds = DC.predict_sets(sub, "q", "c", "p", method, param, miss_mass, True)
    mm, per = macro_f05(preds, truths, ids)
    lost = sum(1 for q, cs in full_preds.items() for c in cs
               if c not in set(preds.get(q, [])))
    print(f"{M:>4} {mm['f05']:>10.6f} {mm['f05']-m_full['f05']:>+11.6f} "
          f"{mm['precision']:>8.4f} {mm['recall']:>8.4f} "
          f"{100*len(sub)/len(cand):>8.1f}% {lost:>24}")
print()
print(f"paired on {len(ids):,} identical entities; ship only at delta == 0")