"""A/B any embedding model on REAL data before committing hours to it.

    python scripts/bench_models.py --gpu 1 --models e5-small e5-large bge-m3

What this measures, exactly
---------------------------
This scores the BLOCKING/RETRIEVAL stage, not the final F0.5. That is deliberate:
blocking recall is the hard ceiling on F0.5 (you cannot match what you never
retrieve), and it is the only thing the embedding model influences. The final
F0.5 is produced by src/run_train_eval.py on a held-out split.

Data
  queries : --n-india / --n-us non-singleton TRAIN Source-1 entities, seed 0.
            Train only, because test has no labels. Countries are read from the
            data, never hard-coded.
  pool    : every Source-2/Source-3 record that is a TRUE match of a sampled query,
            PLUS `--negs` random same-country distractors per query. The negatives
            come from the same country so the test reflects the country-gated
            search the pipeline actually performs.
  truth   : dataset/train/train_ground_truth.tsv

Metric
  recall@k -- the share of true pairs whose match is in the query's top-k, overall
  and on the two slices that decide this problem: non-Latin matched names and
  empty matched addresses. Plus throughput (texts/s) and peak VRAM.

Honest limits
  * The pool is ~`negs`/query, not the full ~4M per country, so ABSOLUTE recall is
    optimistic. Use it to RANK models, not to predict production recall.
  * France is test-only and unlabelled, so no unseen-country signal here.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
import unicodedata as ud

import numpy as np
import pandas as pd

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src"))
from config import EMB, PATHS, model_spec        # noqa: E402

KS = (10, 20, 50)
T0 = time.time()


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


def script_of(name):
    for c in name:
        if c.isalpha() and not ud.name(c, "").startswith("LATIN"):
            return "nonlatin"
    return "latin"


def build_sample(args):
    cols = ["entity_id", "business_name", "business_address", "country"]
    gt = pd.read_csv(PATHS["train"]["ground_truth"], sep="\t", dtype=str, keep_default_na=False)
    gt["ids"] = gt["matched_entity_ids"].str.split(",")
    gt = gt[gt["matched_entity_ids"] != ""]
    s1 = pd.read_parquet(PATHS["train"]["source1"], columns=cols)
    s1 = s1[s1.entity_id.isin(set(gt.source1_entity_id))]
    rng = np.random.default_rng(0)
    picks = []
    for c, n in args.per_country:
        sub = s1[s1.country == c]
        if len(sub) == 0:
            log(f"WARNING: no train records with country={c!r}; skipping")
            continue
        picks.append(sub.iloc[rng.choice(len(sub), min(n, len(sub)), replace=False)])
    Q = pd.concat(picks).reset_index(drop=True)
    truth = gt[gt.source1_entity_id.isin(set(Q.entity_id))]
    need = {m for l in truth.ids for m in l}
    log(f"queries {len(Q):,}  true pairs {sum(len(l) for l in truth.ids):,}")

    # full read of the pool: 10.3M rows x 4 short string cols fits comfortably
    P = pd.concat([pd.read_parquet(PATHS["train"]["source2"], columns=cols),
                   pd.read_parquet(PATHS["train"]["source3"], columns=cols)], ignore_index=True)
    pos = P[P.entity_id.isin(need)]
    negs = []
    for c in Q.country.unique():
        pc = P[(P.country == c) & (~P.entity_id.isin(need))]
        k = min(len(pc), args.negs * int((Q.country == c).sum()))
        negs.append(pc.iloc[rng.choice(len(pc), k, replace=False)])
    pool = pd.concat([pos] + negs, ignore_index=True).drop_duplicates("entity_id").reset_index(drop=True)
    log(f"pool {len(pool):,} ({len(pos):,} true matches + random same-country distractors)")
    return Q, pool, truth


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gpu", default=os.environ.get("M3_GPU"))
    ap.add_argument("--models", nargs="+", default=["e5-small"])
    ap.add_argument("--n-india", type=int, default=2000)
    ap.add_argument("--n-us", type=int, default=1000)
    ap.add_argument("--negs", type=int, default=250, help="random distractors per query")
    ap.add_argument("--out", default=None)
    a = ap.parse_args()
    if a.gpu is None:
        sys.exit("pass --gpu 0 / --gpu 1 / --gpu cpu (or set M3_GPU)")
    dev = "cpu" if str(a.gpu) == "cpu" else f"cuda:{int(a.gpu)}"
    # country labels are data, not code: discover them, then map the two CLI sizes on
    a.per_country = []
    s1c = pd.read_parquet(PATHS["train"]["source1"], columns=["country"])["country"].value_counts()
    for i, c in enumerate(s1c.index.tolist()):
        a.per_country.append((c, a.n_india if i == 0 else a.n_us))
    log(f"sampling {a.per_country} (countries read from the data)")

    Q, pool, truth = build_sample(a)
    qtext = [f"{n} | {ad} | {c}" for n, ad, c in
             zip(Q.business_name, Q.business_address, Q.country)]
    ptext = [f"{n} | {ad} | {c}" for n, ad, c in
             zip(pool.business_name, pool.business_address, pool.country)]
    qpos = {e: i for i, e in enumerate(Q.entity_id)}
    ppos = {e: i for i, e in enumerate(pool.entity_id)}
    pairs = [(qpos[s], ppos[m]) for s, l in zip(truth.source1_entity_id, truth.ids)
             for m in l if m in ppos]
    tq = np.array([x[0] for x in pairs]); tp = np.array([x[1] for x in pairs])
    pinfo = pool.set_index("entity_id")
    mids = [m for s, l in zip(truth.source1_entity_id, truth.ids) for m in l if m in ppos]
    nonlatin = np.array([script_of(pinfo.business_name[m]) == "nonlatin" for m in mids])
    addrempty = np.array([str(pinfo.business_address[m]).strip() == "" for m in mids])
    qc, pc = Q.country.to_numpy(), pool.country.to_numpy()

    from encoder import Encoder
    rows = {}
    for name in a.models:
        hf, qpre, dpre, _d, lic = model_spec(name)
        log(f"--- {name} ({hf}, {lic})")
        enc = Encoder(hf, dev, EMB["max_len"], EMB["batch"])
        enc.encode(ptext[:2000], dpre)                   # warm-up
        t = time.time()
        PE = enc.encode(ptext, dpre).astype(np.float32)
        rate = len(ptext) / (time.time() - t)
        QE = enc.encode(qtext, qpre).astype(np.float32)
        rank = np.full(len(tq), np.inf)
        for c in np.unique(qc):
            qi, pi = np.flatnonzero(qc == c), np.flatnonzero(pc == c)
            for s in range(0, len(qi), 512):
                qq = qi[s:s + 512]
                S = QE[qq] @ PE[pi].T
                kk = min(max(KS), S.shape[1])
                part = np.argpartition(-S, kk - 1, axis=1)[:, :kk]
                order = np.take_along_axis(S, part, 1).argsort(1)[:, ::-1]
                top = np.take_along_axis(part, order, 1)
                for r, q in enumerate(qq):
                    loc = {pi[j]: z + 1 for z, j in enumerate(top[r])}
                    for t_ in np.flatnonzero(tq == q):
                        rank[t_] = loc.get(tp[t_], np.inf)
        row = {"params": f"{sum(p.numel() for p in enc.model.parameters())/1e6:.0f}M",
               "dim": enc.dim, "texts/s": round(rate), "vram_gb": round(enc.peak_vram_gb(), 2)}
        for k in KS:
            h = rank <= k
            row[f"R@{k}"] = round(float(h.mean()), 4)
            row[f"R@{k}_nonLatin"] = round(float(h[nonlatin].mean()), 4) if nonlatin.any() else None
            row[f"R@{k}_addrEmpty"] = round(float(h[addrempty].mean()), 4) if addrempty.any() else None
        rows[name] = row
        log(f"{name}: {json.dumps(row)}")
        del enc, PE, QE
        import gc; gc.collect()
        if dev != "cpu":
            import torch; torch.cuda.empty_cache()

    df = pd.DataFrame(rows).T
    print("\n" + df.to_string())
    out = a.out or os.path.join(PATHS["work_dir"], "bench_models.tsv")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    df.to_csv(out, sep="\t")
    log(f"wrote {out}")


if __name__ == "__main__":
    main()