"""
End-to-end training + validation for Business Entity Resolution.

    python src/run_train_eval.py                   # Split A (random S1 entities)
    python src/run_train_eval.py --split B         # train on one country, validate on the other
    python src/run_train_eval.py --no-ce           # skip the cross-encoder (fast baseline)

Outputs (in PATHS["work_dir"]):
    train_source1_augmented.parquet   S1 + ground-truth column `matched_entity_ids`
    results_<split>/metrics.json, report.md, errors.tsv, feature_importance.tsv
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import pandas as pd
import lightgbm as lgb
from sklearn.metrics import roc_auc_score, average_precision_score
from sklearn.model_selection import GroupKFold

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from config import PATHS, CFG, EMB, emb_dir                     # noqa: E402
from encoder import Encoder, load_matrix, read_ids               # noqa: E402
import data as D                                                # noqa: E402
import text_utils as T                                          # noqa: E402
import blocking as B                                            # noqa: E402
from features import compute_features                           # noqa: E402
from evaluate import macro_f05                                  # noqa: E402
import decide as DC                                             # noqa: E402

T0 = time.time()


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


def lgb_oof_and_final(X, y, groups, cfg):
    """5-fold grouped OOF predictions on the stack part + a final model on all of it."""
    oof = np.zeros(len(y))
    iters = []
    for tr, va in GroupKFold(n_splits=cfg["n_folds_stack"]).split(X, y, groups):
        m = lgb.train(cfg["lgbm_params"], lgb.Dataset(X.iloc[tr], y[tr]), cfg["lgbm_rounds"],
                      valid_sets=[lgb.Dataset(X.iloc[va], y[va])],
                      callbacks=[lgb.early_stopping(cfg["lgbm_early_stopping"], verbose=False)])
        oof[va] = m.predict(X.iloc[va], num_iteration=m.best_iteration)
        iters.append(m.best_iteration)
    final = lgb.train(cfg["lgbm_params"], lgb.Dataset(X, y), int(np.mean(iters) * 1.1) + 1)
    return oof, final


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--split", choices=["A", "B"], default=CFG["split"])
    ap.add_argument("--train-country", default=CFG["split_b_train_country"])
    ap.add_argument("--val-country", default=CFG["split_b_val_country"])
    ap.add_argument("--no-ce", action="store_true")
    ap.add_argument("--gpu", default=os.environ.get("M3_GPU"),
                    help="GPU index or 'cpu' (required; never guessed)")
    ap.add_argument("--n-ce", type=int); ap.add_argument("--n-stack", type=int)
    ap.add_argument("--n-val", type=int)
    args = ap.parse_args()
    if args.gpu is None:
        sys.exit("pass --gpu 0 / --gpu 1 / --gpu cpu (or set M3_GPU)")
    args.gpu = "cpu" if str(args.gpu) == "cpu" else f"cuda:{int(args.gpu)}"
    cfg = dict(CFG)
    cfg["split"] = args.split
    cfg["split_b_train_country"], cfg["split_b_val_country"] = args.train_country, args.val_country
    if args.no_ce:
        cfg["use_cross_encoder"] = False
    for k, a in [("n_s1_ce", args.n_ce), ("n_s1_stack", args.n_stack), ("n_s1_val", args.n_val)]:
        if a:
            cfg[k] = a

    rng = np.random.default_rng(cfg["seed"])
    work = D.ensure_dir(PATHS["work_dir"])
    tag = "A" if cfg["split"] == "A" else f"B_{cfg['split_b_train_country']}2{cfg['split_b_val_country']}"
    res_dir = D.ensure_dir(os.path.join(work, f"results_{tag}"))
    metrics = {"config": cfg}

    # ------------------------------------------------------------------ 1. data
    log("loading data")
    s1, s23, gt = D.load_data(PATHS["train"])
    s1 = D.augment_s1(s1, gt, os.path.join(work, "train_source1_augmented.parquet"))
    del gt
    log(f"S1 {len(s1):,}  S2+S3 {len(s23):,}")
    checks = D.sanity_checks(s1, s23, rng)
    metrics["sanity"] = checks
    log(f"sanity: {json.dumps(checks, default=str)}")
    if checks["gt_ids_missing_in_s23"] > 0:
        log("WARNING: some ground-truth IDs are missing from S2/S3 (incomplete files?)")

    if cfg["country_filter"] == "auto":
        rate = checks["true_pairs_same_country_rate"] or 0
        cfg["country_filter"] = rate >= 0.999
    log(f"country filter: {cfg['country_filter']}")

    # truths as row indices
    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_rows = {i: set() for i in range(len(s1))}
    for q, c in zip(owner[ok].tolist(), pos[ok].tolist()):
        truth_rows[q].add(c)
    del lists, owner, pos

    ce_rows, st_rows, va_rows = D.make_split(s1, cfg, rng)
    log(f"split {tag}: ce {len(ce_rows):,} stack {len(st_rows):,} val {len(va_rows):,}")

    # ------------------------------------------------------------------ 2. abbreviations
    log("mining abbreviations (ce part only)")
    pa, pb = [], []
    for q in ce_rows:
        for c in truth_rows[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(res_dir, "abbrev_map.json"), "w"), indent=1, ensure_ascii=False)
    metrics["abbrev_rules"] = len(amap)
    metrics["abbrev_examples"] = dict(list(sorted(amap.items()))[:40])
    log(f"{len(amap)} abbreviation rules")
    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)

    # ------------------------------------------------------------------ 3. blocking
    q_rows = np.concatenate([ce_rows, st_rows, va_rows])
    part = np.array(["ce"] * len(ce_rows) + ["stack"] * len(st_rows) + ["val"] * len(va_rows))
    part_of = dict(zip(q_rows.tolist(), part.tolist()))

    log("loading cached S2+S3 corpus embeddings")
    pool_dir = emb_dir("train_pool")
    c_emb, cmeta = load_matrix(pool_dir)
    assert list(s23["entity_id"]) == read_ids(os.path.join(pool_dir, "ids.txt")), \
        "train pool embedding row order != S2+S3 order (re-run src/embed.py)"
    log(f"pool embeddings {c_emb.shape} from {cmeta['model']}")

    # Only the sampled queries are ever searched, so encode just those (200k, not 2.21M).
    # This is the dead work that cost model-2 ~50 min per run.
    qcache = os.path.join(work, f"q_emb_{tag}_{EMB['model']}_{len(q_rows)}.npy")
    if os.path.exists(qcache):
        q_emb = np.load(qcache, mmap_mode="r")
        log(f"query embeddings from cache {q_emb.shape}")
    else:
        log(f"encoding {len(q_rows):,} sampled queries on {args.gpu}")
        enc = Encoder(cfg["emb_model"], args.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 q_rows],
                           cfg["emb_query_prefix"])
        np.save(qcache + ".tmp.npy", q_emb); os.replace(qcache + ".tmp.npy", qcache)
        del enc
    log("dense search")
    emb_s, emb_i = B.dense_topk(np.asarray(q_emb), c_emb, cfg, args.gpu)

    # Second dense channel (pre-computed e5-large). Lifts the blocking ceiling from
    # 0.9744 to 0.9834 at k=20 because it fails on different records than e5-small.
    # Its S1 vectors are already on disk in row order, so nothing is encoded here.
    emb2_s = emb2_i = None
    if cfg.get("emb2_enabled"):
        import emb_source
        c2_emb, q2_emb = emb_source.load_emb2("train", len(s1), len(s23))
        log(f"second dense channel: pool {c2_emb.shape} (e5-large, pre-computed)")
        cfg2 = dict(cfg); cfg2["emb_k"] = cfg["emb2_k"]
        emb2_s, emb2_i = B.dense_topk(np.asarray(q2_emb[q_rows]), c2_emb, cfg2, args.gpu)
        del c2_emb, q2_emb

    log("building lexical index")
    lidx = B.LexicalIndex([B.lexical_tokens(n, a) for n, a in
                           zip(s23v["name_core_abbr"], s23v["address_abbrev_norm"])], cfg)
    log("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"])
    del lidx

    cand = B.build_candidates(q_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(q_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[q] for q, c in zip(cand["q"].values, cand["c"].values)),
                            np.int8, len(cand))
    cand["part"] = [part_of[q] for q in cand["q"].values]
    log(f"{len(cand):,} candidate pairs")

    blk = {}
    for p, rows in [("ce", ce_rows), ("stack", st_rows), ("val", va_rows)]:
        n_true = sum(len(truth_rows[q]) for q in rows)
        sub = cand[cand["part"] == p]
        blk[p] = {"recall_ceiling": float(sub["y"].sum() / max(n_true, 1)),
                  "avg_candidates": float(len(sub) / len(rows)),
                  "reduction_ratio": float(1 - len(sub) / (len(rows) * len(s23)))}
    sub = cand[cand["part"] == "val"].copy()
    for k in [1, 5, 10, 20]:
        m = (sub["emb_rank"] < k) | (sub["lex_rank"] < k)
        if cfg.get("emb2_enabled"):
            m = m | (sub["emb2_rank"] < k)
        blk.setdefault("val_recall_at_k", {})[k] = float(sub.loc[m, "y"].sum() /
                                                        max(sum(len(truth_rows[q]) for q in va_rows), 1))
    metrics["blocking"] = blk
    log(f"blocking: {json.dumps(blk)}")
    n_true_st = sum(len(truth_rows[q]) for q in st_rows)
    miss_mass = (n_true_st - cand.loc[cand["part"] == "stack", "y"].sum()) / len(st_rows)

    # ------------------------------------------------------------------ 4. features
    log("computing features")
    feats = compute_features(cand, s1v, s23v, cfg["feat_n_jobs"])
    feat_cols = list(feats.columns) + ["emb_score", "emb_rank", "lex_score", "lex_rank", "n_keys"]
    if cfg.get("emb2_enabled"):
        feat_cols += ["emb2_score", "emb2_rank"]
    cand = pd.concat([cand, feats], axis=1)
    del feats

    # ------------------------------------------------------------------ 5. cross-encoder
    def texts(rows, df):
        from cross_encoder import pair_text
        n, a, c = df["business_name"].values, df["business_address"].values, df["country"].values
        return [pair_text(n[r], a[r], c[r]) for r in rows]

    if cfg["use_cross_encoder"]:
        from cross_encoder import train_cross_encoder, predict_cross_encoder
        log("building cross-encoder training pairs")
        ce_c = cand[cand["part"] == "ce"]
        tq, tc, ty = [], [], []
        for q in ce_rows:                                     # all true pairs (incl. missed by blocking)
            for c in truth_rows[q]:
                tq.append(q); tc.append(c); ty.append(1.0)
        neg = ce_c[ce_c["y"] == 0].copy()
        neg["hard"] = neg["n_keys"] * 10 - np.minimum(neg["emb_rank"], neg["lex_rank"])
        neg = neg.sort_values(["q", "hard"], ascending=[True, False])
        hard = neg.groupby("q").head(cfg["ce_hard_negs"])
        rest = neg.drop(hard.index)
        rnd = rest.groupby("q").sample(frac=1.0, random_state=cfg["seed"]).groupby("q").head(cfg["ce_rand_negs"])
        for d in (hard, rnd):
            tq += d["q"].tolist(); tc += d["c"].tolist(); ty += [0.0] * len(d)
        tq, tc, ty = np.array(tq), np.array(tc), np.array(ty, np.float32)
        if len(ty) > cfg["ce_max_train_pairs"]:
            sel = rng.choice(len(ty), cfg["ce_max_train_pairs"], replace=False)
            tq, tc, ty = tq[sel], tc[sel], ty[sel]
        log(f"cross-encoder train pairs: {len(ty):,} (pos rate {ty.mean():.3f})")
        ce_dir = os.path.join(work, f"ce_model_{tag}")
        train_cross_encoder(texts(tq, s1), texts(tc, s23), ty, cfg, ce_dir, cfg["seed"])
        log("scoring stack + val pairs with cross-encoder")
        m = cand["part"].isin(["stack", "val"]).values
        cand["ce_prob"] = np.nan
        cand.loc[m, "ce_prob"] = predict_cross_encoder(texts(cand["q"].values[m], s1),
                                                       texts(cand["c"].values[m], s23), ce_dir, cfg)

    # ------------------------------------------------------------------ 6. scorers
    st = cand[cand["part"] == "stack"].reset_index(drop=True)
    va = cand[cand["part"] == "val"].reset_index(drop=True)
    del cand
    y_st, y_va = st["y"].values, va["y"].values
    truths_st = {q: truth_rows[q] for q in st_rows}
    truths_va = {q: truth_rows[q] for q in va_rows}
    ids_st, ids_va = list(st_rows), list(va_rows)

    scores_st, scores_va, aucs = {}, {}, {}
    # rule baseline
    scores_st["rule"] = 0.7 * st["name_jw"].values + 0.3 * st["addr_tset"].values
    scores_va["rule"] = 0.7 * va["name_jw"].values + 0.3 * va["addr_tset"].values
    variants = [("lgbm_noce", feat_cols)]
    if cfg["use_cross_encoder"]:
        scores_st["ce_only"] = st["ce_prob"].values
        scores_va["ce_only"] = va["ce_prob"].values
        variants.append(("lgbm_ce", feat_cols + ["ce_prob"]))
    fi = None
    models = {}
    for name, cols in variants:
        log(f"training LightGBM {name}")
        oof, final = lgb_oof_and_final(st[cols], y_st, st["q"].values, cfg)
        models[name] = (final, cols)
        scores_st[name] = oof
        scores_va[name] = final.predict(va[cols])
        fi = pd.DataFrame({"feature": cols, "gain": final.feature_importance("gain")}) \
            .sort_values("gain", ascending=False)
    fi.to_csv(os.path.join(res_dir, "feature_importance.tsv"), sep="\t", index=False)
    for k in scores_va:
        aucs[k] = {"roc_auc": float(roc_auc_score(y_va, scores_va[k])) if 0 < y_va.sum() < len(y_va) else None,
                   "pr_auc": float(average_precision_score(y_va, scores_va[k])) if y_va.sum() else None}
    metrics["pair_auc_val"] = aucs

    # ------------------------------------------------------------------ 7. decision + ablation
    thr_grid = np.round(np.arange(0.05, 0.96, 0.025), 3)
    floor_grid = [0.0, 0.05, 0.1, 0.15, 0.2, 0.3, 0.4]
    ablation = {}
    final_key = "lgbm_ce" if cfg["use_cross_encoder"] else "lgbm_noce"
    final_preds = None
    for k in scores_va:
        # calibrate on stack OOF (rule/ce scores are also mapped to probabilities)
        iso = DC.fit_calibrator(scores_st[k], y_st)
        st["p"], va["p"] = iso.predict(scores_st[k]), iso.predict(scores_va[k])
        for method, grid in [("threshold", thr_grid), ("ef05", floor_grid)]:
            for conf in ([True, False] if k == final_key else [True]):
                v, f_st = DC.tune(st, truths_st, ids_st, "q", "c", "p", method, grid, miss_mass, conf,
                                  lambda P, Tr, I: macro_f05(P, Tr, I))
                preds = DC.predict_sets(va, "q", "c", "p", method, v, miss_mass, conf)
                m, per = macro_f05(preds, truths_va, ids_va)
                key = f"{k}|{method}" + ("" if conf else "|no_conflict_resolution")
                ablation[key] = {**m, "tuned_param": float(v), "stack_oof_f05": float(f_st)}
                log(f"{key:45s} val F0.5 {m['f05']:.4f}  P {m['precision']:.4f}  R {m['recall']:.4f}")
                if k == final_key and method == "ef05" and conf:
                    final_preds, final_per, final_m = preds, per, m
                    va["p_final"] = va["p"]
                    final_bundle = {"calibrator": iso, "method": method, "param": float(v)}
    # oracle: perfect matcher on top of our blocking (upper bound)
    hit = va[va["y"] == 1].groupby("q")["c"].apply(list).to_dict()
    oracle = {q: hit.get(q, []) for q in ids_va}
    ablation["oracle_on_candidates"] = macro_f05(oracle, truths_va, ids_va)[0]
    metrics["ablation_val"] = ablation

    # ------------------------------------------------------------------ 8. breakdowns
    per = pd.DataFrame({"q": ids_va, "f05": final_per})
    per["singleton"] = [len(truths_va[q]) == 0 for q in ids_va]
    per["country"] = s1v["country_norm"][per["q"].values]
    metrics["final"] = {
        "variant": f"{final_key}|ef05",
        **final_m,
        "f05_singletons": float(per.loc[per["singleton"], "f05"].mean()) if per["singleton"].any() else None,
        "f05_non_singletons": float(per.loc[~per["singleton"], "f05"].mean()),
        "singleton_share": float(per["singleton"].mean()),
        "f05_by_country": per.groupby("country")["f05"].mean().to_dict(),
    }

    # error analysis
    pred_pairs = {(q, c) for q, cs in final_preds.items() for c in cs}
    va["pred"] = [(q, c) in pred_pairs for q, c in zip(va["q"].values, va["c"].values)]
    fp = va[va["pred"] & (va["y"] == 0)].nlargest(20, "p_final")
    fn = va[~va["pred"] & (va["y"] == 1)].nlargest(20, "p_final")
    rows = []
    for kind, d in [("FP", fp), ("FN", fn)]:
        for _, r in d.iterrows():
            rows.append({"type": kind, "prob": r["p_final"],
                         "s1_id": s1v["entity_id"][r["q"]], "s1_name": s1["business_name"].values[r["q"]],
                         "s1_addr": s1["business_address"].values[r["q"]],
                         "cand_id": s23v["entity_id"][r["c"]], "cand_name": s23["business_name"].values[r["c"]],
                         "cand_addr": s23["business_address"].values[r["c"]]})
    pd.DataFrame(rows).to_csv(os.path.join(res_dir, "errors.tsv"), sep="\t", index=False)

    # ------------------------------------------------------------------ 9. write report
    # ---------------------------------------------------------------- persist for run_test.py
    import joblib
    _lgbm, _cols = models[final_key]
    joblib.dump({"emb2_enabled": bool(cfg.get("emb2_enabled")),
                 "lgbm": _lgbm, "feat_cols": _cols,
                 "calibrator": final_bundle["calibrator"],
                 "method": final_bundle["method"], "param": final_bundle["param"],
                 "miss_mass": float(miss_mass),
                 "country_filter": bool(cfg["country_filter"]),
                 "use_cross_encoder": bool(cfg["use_cross_encoder"]),
                 "ce_dir": (ce_dir if cfg["use_cross_encoder"] else None),
                 "emb_model": cfg["emb_model"], "val_f05": float(final_m["f05"])},
                os.path.join(res_dir, "model_bundle.joblib"))
    log(f"saved model_bundle.joblib -> {res_dir}")

    metrics["runtime_sec"] = round(time.time() - T0)
    json.dump(metrics, open(os.path.join(res_dir, "metrics.json"), "w"), indent=2, default=str)
    L = [f"# Validation report — split {tag}", "",
         f"**Final ({metrics['final']['variant']}) macro F0.5: {final_m['f05']:.4f}**  "
         f"(P {final_m['precision']:.4f}, R {final_m['recall']:.4f}, n={final_m['n']:,})", "",
         f"- singletons: {metrics['final']['f05_singletons']}  |  non-singletons: "
         f"{metrics['final']['f05_non_singletons']:.4f}  |  singleton share: {metrics['final']['singleton_share']:.3f}",
         f"- by country: {metrics['final']['f05_by_country']}", "",
         "## Blocking", "```", json.dumps(blk, indent=1), "```", "",
         "## Pair-level AUC (val)", "| scorer | ROC-AUC | PR-AUC |", "|---|---|---|"]
    L += [f"| {k} | {v['roc_auc']} | {v['pr_auc']} |" for k, v in aucs.items()]
    L += ["", "## Ablation (val macro F0.5)", "| variant | F0.5 | P | R | tuned param |", "|---|---|---|---|---|"]
    for k, v in ablation.items():
        L.append(f"| {k} | {v['f05']:.4f} | {v['precision']:.4f} | {v['recall']:.4f} | {v.get('tuned_param', '')} |")
    L += ["", f"Abbreviation rules learned: {len(amap)} (see abbrev_map.json)",
          f"Country filter used: {cfg['country_filter']}", f"Runtime: {metrics['runtime_sec']}s"]
    open(os.path.join(res_dir, "report.md"), "w").write("\n".join(L))
    log(f"done -> {res_dir}")


if __name__ == "__main__":
    main()