"""Blocking only: generate LABELLED candidate pairs for an 80/20 train/test split.

    python src/run_blocking.py --gpu 1                      # all 2.21M train S1
    python src/run_blocking.py --gpu 1 --n-queries 500000   # subset, faster

Does NOT train anything. Produces the candidate-pair dataset the matcher will
later consume, plus a blocking-quality report.

Split
  80/20 over Source-1 ENTITIES (not pairs), seeded, so no entity appears in both
  sides and there is no leakage through the shared S2/S3 pool.

Output (work/blocking/<model>/)
  pairs_train/part_XXXXX.parquet   q_id, c_id, emb_score, emb_rank,
  pairs_test/part_XXXXX.parquet    lex_score, lex_rank, n_keys, y
  report.json                      recall, candidates/query, reduction ratio,
                                   per-slice recall, F0.5 oracle ceiling

Resumable: finished chunks are skipped on re-run.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np
import pandas as pd

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import blocking as B                                   # noqa: E402
import data as D                                       # noqa: E402
import text_utils as T                                 # noqa: E402
from config import CFG, EMB, PATHS, emb_dir            # noqa: E402
from encoder import load_matrix, read_ids              # noqa: E402

T0 = time.time()


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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gpu", default=os.environ.get("M3_GPU"))
    ap.add_argument("--n-queries", type=int, default=0, help="0 = all train S1")
    ap.add_argument("--chunk", type=int, default=100_000, help="queries per resumable chunk")
    ap.add_argument("--test-frac", type=float, default=0.2)
    ap.add_argument("--seed", type=int, default=42)
    a = ap.parse_args()
    if a.gpu is None:
        sys.exit("pass --gpu 0 / --gpu 1 / --gpu cpu")
    dev = "cpu" if str(a.gpu) == "cpu" else f"cuda:{int(a.gpu)}"

    out = os.path.join(PATHS["work_dir"], "blocking", EMB["model"])
    for sub in ("pairs_train", "pairs_test"):
        os.makedirs(os.path.join(out, sub), exist_ok=True)

    # ---------------------------------------------------------------- data
    log("loading train sources + ground truth")
    s1, s23, gt = D.load_data(PATHS["train"])
    s1 = s1.merge(gt.rename(columns={"source1_entity_id": "entity_id"}), on="entity_id",
                  how="left", validate="one_to_one")
    s1["matched_entity_ids"] = s1["matched_entity_ids"].fillna("").astype(str)
    del gt
    log(f"S1 {len(s1):,}  S2+S3 {len(s23):,}")

    # abbreviation map mined from TRAIN-side positives only (no leakage into test)
    rng = np.random.default_rng(a.seed)
    perm = rng.permutation(len(s1))
    n_test = int(len(perm) * a.test_frac)
    test_rows = np.sort(perm[:n_test])
    train_rows = np.sort(perm[n_test:])
    is_test = np.zeros(len(s1), bool)
    is_test[test_rows] = True
    log(f"split: train {len(train_rows):,} entities | test {len(test_rows):,} entities "
        f"({a.test_frac:.0%})")

    s23_index = pd.Index(s23["entity_id"])
    lists = s1["matched_entity_ids"].map(D.split_ids)
    owner = np.repeat(np.arange(len(s1)), lists.map(len).values)
    pos = s23_index.get_indexer([x for l in lists for x in l])
    ok = pos >= 0
    truth = {}
    for q, c in zip(owner[ok].tolist(), pos[ok].tolist()):
        truth.setdefault(q, set()).add(c)
    n_true_total = int(ok.sum())
    log(f"true pairs: {n_true_total:,} ({int((~ok).sum()):,} unresolved ids)")

    log("mining abbreviations from TRAIN-side positives only")
    pa, pb = [], []
    for q in train_rows[:200_000]:
        for c in truth.get(q, ()):
            pa.append(s1["name_core"].values[q]); pb.append(s23["name_core"].values[c])
    if len(pa) > CFG["abbrev_max_pos_pairs"]:
        sel = rng.choice(len(pa), CFG["abbrev_max_pos_pairs"], replace=False)
        pa, pb = [pa[i] for i in sel], [pb[i] for i in sel]
    nb = [pb[i] for i in rng.permutation(len(pb))]
    amap = T.mine_abbreviations(zip(pa, pb), zip(pa, nb),
                                CFG["abbrev_min_support"], CFG["abbrev_min_precision"])
    json.dump(amap, open(os.path.join(out, "abbrev_map.json"), "w"), indent=1, ensure_ascii=False)
    log(f"{len(amap)} abbreviation rules")
    del pa, pb, nb

    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]
    s1_nca = s1["name_core_abbr"].values.astype(object)
    s1_addr = s1["address_abbrev_norm"].values.astype(object)
    s23_nca = s23["name_core_abbr"].values.astype(object)
    s23_addr = s23["address_abbrev_norm"].values.astype(object)
    q_country = s1["country_norm"].values.astype(object)
    c_country = s23["country_norm"].values.astype(object)
    s1_ids = s1["entity_id"].values.astype(object)
    s23_ids = s23["entity_id"].values.astype(object)

    # ---------------------------------------------------------------- embeddings
    import emb_source
    c_emb, q_emb, cmeta = emb_source.load("train", list(s1["entity_id"]),
                                          list(s23["entity_id"]))
    log(f"embeddings: pool {c_emb.shape} queries {q_emb.shape} "
        f"({cmeta['model']}, source={cmeta.get('source', 'internal')})")

    # ---------------------------------------------------------------- lexical index
    log("building lexical index over the 10.3M pool (one-off, slowest single step)")
    t = time.time()
    lidx = B.LexicalIndex([B.lexical_tokens(n, ad) for n, ad in zip(s23_nca, s23_addr)], CFG)
    log(f"lexical index built in {time.time() - t:.0f}s")

    # ---------------------------------------------------------------- query order
    all_q = np.arange(len(s1)) if a.n_queries <= 0 else np.sort(
        rng.choice(len(s1), min(a.n_queries, len(s1)), replace=False))
    log(f"blocking {len(all_q):,} queries in chunks of {a.chunk:,}")

    stats = {"train": [], "test": []}
    for ci, lo in enumerate(range(0, len(all_q), a.chunk)):
        qs = all_q[lo:lo + a.chunk]
        done = [os.path.join(out, f"pairs_{s}", f"part_{ci:05d}.parquet") for s in ("train", "test")]
        if all(os.path.exists(p) for p in done):
            continue
        t = time.time()
        emb_s, emb_i = B.dense_topk(np.asarray(q_emb[qs]), c_emb, CFG, dev)
        lex = B.lexical_topk(lidx, [B.lexical_tokens(s1_nca[i], s1_addr[i]) for i in qs],
                             CFG["lex_n_jobs"])
        cand = B.build_candidates(qs, emb_s, emb_i, lex, q_country, c_country,
                                  country_filter=True)
        cand = B.fill_missing_emb(cand, lambda rows, _m={r: j for j, r in enumerate(qs.tolist())}:
                                  q_emb[[qs[_m[r]] for r in rows]], c_emb)
        cand["y"] = np.fromiter((c in truth.get(q, ()) for q, c in
                                 zip(cand["q"].values, cand["c"].values)), np.int8, len(cand))
        cand["q_id"] = s1_ids[cand["q"].values]
        cand["c_id"] = s23_ids[cand["c"].values]
        side = is_test[cand["q"].values]
        for name, mask in (("train", ~side), ("test", side)):
            sub = cand[mask]
            cols = ["q_id", "c_id", "emb_score", "emb_rank", "lex_score", "lex_rank", "n_keys", "y"]
            p = os.path.join(out, f"pairs_{name}", f"part_{ci:05d}.parquet")
            sub[cols].to_parquet(p + ".tmp", index=False)
            os.replace(p + ".tmp", p)
            rows = qs[is_test[qs]] if name == "test" else qs[~is_test[qs]]
            stats[name].append({"n_q": int(len(rows)), "n_cand": int(len(sub)),
                                "n_hit": int(sub["y"].sum()),
                                "n_true": int(sum(len(truth.get(q, ())) for q in rows))})
        log(f"chunk {ci + 1}/{(len(all_q) + a.chunk - 1) // a.chunk}: {len(qs):,} queries, "
            f"{len(cand):,} pairs, {time.time() - t:.0f}s")
        json.dump(stats, open(os.path.join(out, "stats_partial.json"), "w"))

    # ---------------------------------------------------------------- report
    rep = {"model": cmeta["model"], "dim": cmeta["dim"], "emb_k": CFG["emb_k"],
           "lex_k": CFG["lex_k"], "country_filter": True, "seed": a.seed,
           "test_frac": a.test_frac, "n_queries_blocked": int(len(all_q)),
           "abbrev_rules": len(amap), "pool_size": int(len(s23))}
    for name in ("train", "test"):
        S = stats[name]
        nq = sum(x["n_q"] for x in S); nc = sum(x["n_cand"] for x in S)
        nh = sum(x["n_hit"] for x in S); nt = sum(x["n_true"] for x in S)
        rep[name] = {
            "n_entities": nq, "n_candidates": nc, "n_true_pairs": nt, "n_true_retrieved": nh,
            "pair_recall": nh / max(nt, 1),
            "candidates_per_entity": nc / max(nq, 1),
            "reduction_ratio": 1 - nc / max(nq * len(s23), 1),
        }
    rep["seconds"] = round(time.time() - T0)
    json.dump(rep, open(os.path.join(out, "report.json"), "w"), indent=2)
    log("---- blocking report ----")
    for name in ("train", "test"):
        r = rep[name]
        log(f"{name}: {r['n_entities']:,} entities | {r['n_candidates']:,} candidates "
            f"({r['candidates_per_entity']:.1f}/entity) | pair recall {r['pair_recall']:.4f} "
            f"| reduction {r['reduction_ratio']:.6f}")
    log(f"-> {out}")


if __name__ == "__main__":
    main()