"""Test-set inference -> output/matching_results.tsv + output/candidate_pairs.tsv.

    python src/run_test.py --gpu 1 --artifacts work/results_A

Requires (produced earlier in the queue):
  * work/emb/<model>/test_pool/   and  work/emb/<model>/test_s1/   (src/embed.py)
  * an artifacts dir from run_train_eval.py holding abbrev_map.json, the LightGBM
    model, the isotonic calibrator, the tuned decision parameter and miss_mass.

Output contract (see the organisers' README + utils/validate_submission.py):
  * exactly one row per test Source-1 entity, INCLUDING singletons (empty list)
  * ids may only be S2-/S3- ids that exist in the test set
  * no duplicate ids inside a list, no duplicate source1_entity_id rows
  * matching_results must be a SUBSET of candidate_pairs
"""
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 decide as DC                                     # 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
from features import compute_features                   # noqa: E402

T0 = time.time()


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


def write_tsv(path, rows, col):
    """rows: list of (s1_id, [ids]). Writes the exact organiser format."""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w") as f:
        f.write(f"source1_entity_id\t{col}\n")
        for s1, ids in rows:
            f.write(f"{s1}\t{','.join(ids)}\n")
    os.replace(tmp, path)
    log(f"wrote {path} ({len(rows):,} rows)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gpu", default=os.environ.get("M3_GPU"))
    ap.add_argument("--artifacts", required=True, help="dir written by run_train_eval.py")
    ap.add_argument("--out", default=PATHS["output_dir"])
    ap.add_argument("--chunk", type=int, default=200_000, help="test S1 rows per batch")
    a = ap.parse_args()
    if a.gpu is None:
        sys.exit("pass --gpu 0 / --gpu 1 / --gpu cpu (or set M3_GPU)")
    device = "cpu" if str(a.gpu) == "cpu" else f"cuda:{int(a.gpu)}"
    art = a.artifacts

    # ---------------------------------------------------------------- artifacts
    import joblib
    amap = json.load(open(os.path.join(art, "abbrev_map.json")))
    bundle = joblib.load(os.path.join(art, "model_bundle.joblib"))
    lgbm, iso = bundle["lgbm"], bundle["calibrator"]
    feat_cols, method = bundle["feat_cols"], bundle["method"]
    param, miss_mass = bundle["param"], bundle["miss_mass"]
    country_filter = bundle["country_filter"]
    use_ce = bundle.get("use_cross_encoder", False)
    ce_dir = bundle.get("ce_dir")
    log(f"artifacts: {len(amap)} abbrev rules | {len(feat_cols)} features | "
        f"{method}({param}) | country_filter={country_filter} | ce={use_ce}")

    # ---------------------------------------------------------------- data
    log("loading test data")
    s1, s23, _ = D.load_data(PATHS["test"], with_gt=False)
    log(f"test S1 {len(s1):,}  test S2+S3 {len(s23):,}")
    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)
        raw = df["business_address"].values.astype(object)
        v["nums"] = np.where(nums == "", raw, nums)
        return v
    s1v, s23v = vals(s1), vals(s23)

    # ---------------------------------------------------------------- embeddings
    pool_dir, q_dir = emb_dir("test_pool"), emb_dir("test_s1")
    c_emb, cmeta = load_matrix(pool_dir)
    q_emb, qmeta = load_matrix(q_dir)
    pool_ids, q_ids = read_ids(os.path.join(pool_dir, "ids.txt")), read_ids(os.path.join(q_dir, "ids.txt"))
    # row order of the cached matrices MUST match the dataframes we just loaded
    assert list(s23["entity_id"]) == pool_ids, "test pool embedding row order != S2+S3 order"
    assert list(s1["entity_id"]) == q_ids, "test S1 embedding row order != S1 order"
    log(f"embeddings ok: pool {c_emb.shape} queries {q_emb.shape} (model {cmeta['model']})")

    # ---------------------------------------------------------------- lexical index
    log("building lexical index over the test pool")
    lidx = B.LexicalIndex([B.lexical_tokens(n, ad) for n, ad in
                           zip(s23v["name_core_abbr"], s23v["address_abbrev_norm"])], CFG)

    # ---------------------------------------------------------------- per-chunk inference
    # Each chunk is scored and persisted atomically, so a disconnect, an OOM from
    # another user, or a node reboot costs at most one chunk instead of the whole
    # multi-hour run. Re-running skips finished chunks.
    #
    # NOTE: conflict resolution is deliberately NOT applied here any more. Every
    # S2/S3 record has exactly one true owner (verified on the training labels:
    # 0 multi-owner ids across all 7,638,365 links), but two S1 entities sitting in
    # DIFFERENT chunks can each claim the same record. Resolving per chunk left
    # those duplicate claims in place, which costs precision on a precision-weighted
    # metric. It is now done once, globally, after every chunk exists.
    emb2_on = bool(bundle.get("emb2_enabled", CFG.get("emb2_enabled")))
    if emb2_on:
        import emb_source
        c2_emb, q2_emb = emb_source.load_emb2("test", len(s1), len(s23))
        log(f"second dense channel: pool {c2_emb.shape} (e5-large, pre-computed)")

    # cascade pre-ranker (reranker-free); absent file = cascade off
    cascade_m = int(CFG.get("cascade_m", 0) or 0)
    preranker = None
    if cascade_m:
        pr_path = os.path.join(art, "preranker.joblib")
        if os.path.exists(pr_path):
            preranker = joblib.load(pr_path)
            log(f"cascade ON: reranking top-{cascade_m}/entity using a "
                f"{len(preranker['cols'])}-feature pre-ranker "
                f"(validated delta F0.5 +0.000000)")
        else:
            log(f"cascade requested (M={cascade_m}) but {pr_path} is missing - "
                f"reranking ALL candidates instead")

    ck_dir = os.path.join(art, "test_chunks")
    os.makedirs(ck_dir, exist_ok=True)
    n = len(s1)
    chunk_los_all = list(range(0, n, a.chunk))
    _lo = int(os.environ.get("M3_CHUNK_LO", "0"))
    _hi = int(os.environ.get("M3_CHUNK_HI", str(len(chunk_los_all))))
    chunk_los = chunk_los_all[_lo:_hi]
    skip_decide = os.environ.get("M3_SKIP_DECIDE", "0") == "1"
    if (_lo, _hi) != (0, len(chunk_los_all)) or skip_decide:
        log(f"worker slice: chunks [{_lo}:{_hi}] of {len(chunk_los_all)}"
            f"{' (skipping global decision)' if skip_decide else ''}")
    for lo in chunk_los:
        hi = min(lo + a.chunk, n)
        cpath = os.path.join(ck_dir, f"chunk_{lo:09d}.parquet")
        if os.path.exists(cpath):
            log(f"chunk {lo:,}-{hi:,}: already complete, skipping")
            continue
        q_rows = np.arange(lo, hi)
        log(f"chunk {lo:,}-{hi:,}: dense search")
        emb_s, emb_i = B.dense_topk(np.asarray(q_emb[lo:hi]), c_emb, CFG, device)
        log(f"chunk {lo:,}-{hi:,}: lexical search")
        lex = B.lexical_topk(lidx, [B.lexical_tokens(s1v["name_core_abbr"][i],
                                                     s1v["address_abbrev_norm"][i])
                                    for i in q_rows], CFG["lex_n_jobs"])
        emb2_s = emb2_i = None
        if emb2_on:
            cfg2 = dict(CFG); cfg2["emb_k"] = CFG["emb2_k"]
            emb2_s, emb2_i = B.dense_topk(np.asarray(q2_emb[lo:hi]), c2_emb, cfg2, device)
        cand = B.build_candidates(q_rows, emb_s, emb_i, lex, s1v["country_norm"],
                                  s23v["country_norm"], country_filter,
                                  emb2_s=emb2_s, emb2_i=emb2_i)
        qpos = {r: i for i, r in enumerate(q_rows.tolist())}
        cand = B.fill_missing_emb(cand, lambda rows: q_emb[[lo + qpos[r] for r in rows]], c_emb)
        log(f"chunk {lo:,}-{hi:,}: {len(cand):,} candidates; features")
        feats = compute_features(cand, s1v, s23v, CFG["feat_n_jobs"])
        cand = pd.concat([cand, feats], axis=1)
        del feats

        if preranker is not None:
            before = len(cand)
            sc = preranker["model"].predict(cand[preranker["cols"]])
            rk = pd.Series(sc, index=cand.index).groupby(cand["q"].values) \
                   .rank(ascending=False, method="first")
            cand = cand[rk.values <= cascade_m].reset_index(drop=True)
            log(f"chunk {lo:,}-{hi:,}: cascade kept {len(cand):,}/{before:,} "
                f"({100*len(cand)/max(before,1):.1f}%) for reranking")

        if use_ce:
            from cross_encoder import pair_text, predict_cross_encoder
            log(f"chunk {lo:,}-{hi:,}: 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]
            # resumable inside the chunk too: scoring 2M+ pairs can itself take an hour
            cand["ce_prob"] = predict_cross_encoder(
                ta, tb, ce_dir, CFG, cache=os.path.join(ck_dir, f"ce_{lo:09d}"))
            del ta, tb

        cand["p"] = iso.predict(lgbm.predict(cand[feat_cols]))
        out = pd.DataFrame({"q": cand["q"].values.astype(np.int64),
                            "c": cand["c"].values.astype(np.int64),
                            "p": cand["p"].values.astype(np.float32)})
        out.to_parquet(cpath + ".tmp", index=False)
        os.replace(cpath + ".tmp", cpath)
        for stale in (f"ce_{lo:09d}.part.npy", f"ce_{lo:09d}.done.json"):
            sp = os.path.join(ck_dir, stale)
            if os.path.exists(sp):
                os.remove(sp)
        log(f"chunk {lo:,}-{hi:,}: persisted {len(out):,} scored pairs -> {cpath}")
        del cand, out

    # ---------------------------------------------------------------- global decision
    if skip_decide:
        log(f"worker finished its {len(chunk_los)} chunk(s); global decision skipped "
            f"(run again without M3_SKIP_DECIDE once every chunk exists)")
        return

    # The decision MUST see every chunk, never just this worker's slice, or the
    # one-owner constraint would be resolved against a partial claim set and the
    # submission would be missing entities.
    missing = [lo for lo in chunk_los_all
               if not os.path.exists(os.path.join(ck_dir, f"chunk_{lo:09d}.parquet"))]
    if missing:
        raise SystemExit(f"cannot run the global decision: {len(missing)} chunk(s) "
                         f"still missing (first at row {missing[0]:,}). Wait for the "
                         f"other worker(s) to finish.")
    log("loading every chunk for GLOBAL conflict resolution + selection")
    allc = pd.concat([pd.read_parquet(os.path.join(ck_dir, f"chunk_{lo:09d}.parquet"))
                      for lo in chunk_los_all], ignore_index=True)
    log(f"{len(allc):,} scored pairs over {n:,} entities "
        f"({len(allc)/max(n,1):.1f} candidates/entity)")
    n_claims = len(allc)
    n_owners = allc["c"].nunique()
    log(f"{n_owners:,} distinct pool records claimed by {n_claims:,} candidate edges")
    preds = DC.predict_sets(allc, "q", "c", "p", method, param, miss_mass, conflicts=True)
    log(f"selection done: {len(preds):,} entities received at least one match")

    pool_id_arr = s23v["entity_id"]
    s1_id_arr = s1v["entity_id"]
    by_q = allc.groupby("q")["c"].apply(list).to_dict()
    del allc

    match_rows, cand_rows = [], []
    for r in range(n):
        sid = s1_id_arr[r]
        cand_rows.append((sid, [str(pool_id_arr[c]) for c in by_q.get(r, [])]))
        match_rows.append((sid, [str(pool_id_arr[c]) for c in preds.get(r, [])]))

    # ---------------------------------------------------------------- write + verify
    assert len(match_rows) == n, f"{len(match_rows)} rows for {n} test S1 entities"
    valid = set(s23v["entity_id"].astype(str))
    seen_owner = {}
    for sid, ids in match_rows:
        assert len(ids) == len(set(ids)), f"duplicate ids for {sid}"
        assert all(i in valid for i in ids), f"unknown id in {sid}"
        for i in ids:                       # the constraint we just fixed - prove it
            assert i not in seen_owner, f"{i} claimed by both {seen_owner[i]} and {sid}"
            seen_owner[i] = sid
    cand_map = {s: set(i) for s, i in cand_rows}
    for sid, ids in match_rows:
        assert set(ids) <= cand_map[sid], f"{sid}: match not a subset of candidates"
    log(f"verified: {len(seen_owner):,} matched ids, each owned by exactly one S1 entity")

    write_tsv(os.path.join(a.out, "matching_results.tsv"), match_rows, "matched_entity_ids")
    write_tsv(os.path.join(a.out, "candidate_pairs.tsv"), cand_rows, "candidate_entity_ids")
    nz = sum(1 for _, i in match_rows if i)
    log(f"predicted matches for {nz:,}/{n:,} entities "
        f"({1 - nz / n:.4f} predicted-singleton rate; train singleton rate is 0.0558)")
    log(f"now run: python3 {PATHS['validator']} --matching {a.out}/matching_results.tsv "
        f"--candidate {a.out}/candidate_pairs.tsv --test-dir {os.path.dirname(PATHS['test']['source1'])}")


if __name__ == "__main__":
    main()