"""Stage 1 (cheap LightGBM over all fused candidates), stage-2 candidate set, stage 2 (LightGBM + cross-encoder,
calibration, decision selection), test inference + submission files, and the leave-one-country-out proxy."""
import pickle
import subprocess
import sys
from multiprocessing import Pool

import lightgbm as lgb
import numpy as np
import pandas as pd

from common import get_logger, load_json, save_json, save_parquet, shard_path, n_workers
from config import CFG, OUTPUT, PATHS, WORK
from decide import crossfit_isotonic, decide, fit_isotonic, macro_f05
from features import feature_columns
from prep import load_prep, load_splits
from retrieve_dense import n_shards

log = get_logger("stack")
MODELS = WORK / "models"
PARAMS = dict(objective="binary", learning_rate=0.05, num_leaves=127, min_data_in_leaf=100, feature_fraction=0.8,
              bagging_fraction=0.8, bagging_freq=1, lambda_l2=1.0, max_bin=255, verbose=-1)


def fold_of(q, n):
    return (q.astype(np.int64) * 2654435761 % 2 ** 32) % n


def train_lgb(X, y, q, path, seed):
    """Early stopping on an inner holdout of whole S1 groups; model saved atomically for resume."""
    if path.exists():
        return lgb.Booster(model_file=str(path))
    inner = fold_of(q + 7, 10) == 0
    p = dict(PARAMS, seed=seed, num_threads=CFG["lgb_threads"])
    dtr = lgb.Dataset(X[~inner], y[~inner], free_raw_data=True)
    dva = lgb.Dataset(X[inner], y[inner], reference=dtr)
    m = lgb.train(p, dtr, num_boost_round=3000, valid_sets=[dva],
                  callbacks=[lgb.early_stopping(100, verbose=False), lgb.log_evaluation(250)])
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    m.save_model(str(tmp), num_iteration=m.best_iteration)
    tmp.replace(path)
    log.info(f"saved {path.name} best_iter={m.best_iteration}")
    return lgb.Booster(model_file=str(path))


def load_feats(split, qset=None):
    parts = []
    for i in range(n_shards(split)):
        d = pd.read_parquet(shard_path("feat", split, i))
        if qset is not None:
            d = d[np.isin(d.q.values, qset)]
        parts.append(d)
    return pd.concat(parts, ignore_index=True)


# ----------------------------------------------------------------------------------------------- stage 1

def run_stage1_train():
    sp = load_splits()
    df = load_feats("train")
    feats = feature_columns(df)
    MODELS.mkdir(parents=True, exist_ok=True)
    save_json(feats, MODELS / "stage1_features.json")
    stack = df[np.isin(df.q.values, sp["g_stack"])].reset_index(drop=True)
    val = df[np.isin(df.q.values, sp["g_val"])].reset_index(drop=True)
    del df
    F = CFG["n_folds"]
    fold = fold_of(stack.q.values, F)
    oof = np.zeros(len(stack))
    vp = np.zeros(len(val))
    rng = np.random.default_rng(CFG["seed"])
    for k in range(F):
        tr = stack[fold != k]
        keep = (tr.y.values == 1) | (tr.rrf_rank.values < CFG["s1_neg_keep_rank"]) | (rng.random(len(tr)) < CFG["s1_neg_frac"])
        tr = tr[keep]
        m = train_lgb(tr[feats].values.astype(np.float32), tr.y.values, tr.q.values, MODELS / f"s1_fold{k}.txt", k)
        oof[fold == k] = m.predict(stack.loc[fold == k, feats].values.astype(np.float32))
        vp += m.predict(val[feats].values.astype(np.float32)) / F
    out = pd.concat([stack[["q", "c", "y"]].assign(p1=oof), val[["q", "c", "y"]].assign(p1=vp)], ignore_index=True)
    save_parquet(out, WORK / "stage1" / "train.parquet")
    from sklearn.metrics import roc_auc_score
    log.info(f"stage1 AUC stack-oof={roc_auc_score(stack.y, oof):.5f} val={roc_auc_score(val.y, vp):.5f}")


def fold_predictor(models, X_check):
    """Mean of the fold models' probabilities. Uses the exact GPU tree evaluator (gpu_lgb.py) when a GPU is
    present and it reproduces LightGBM on a check sample (|diff| < 1e-9); otherwise LightGBM on the CPU."""
    cpu = lambda X: sum(m.predict(X, num_threads=CFG["lgb_threads"]) for m in models) / len(models)
    try:
        import torch
        if torch.cuda.is_available():
            from gpu_lgb import GpuForest
            fs = [GpuForest(m, torch.device("cuda:0")) for m in models]
            gpu = lambda X: sum(f.predict(X) for f in fs) / len(fs)
            diff = float(np.abs(gpu(X_check) - cpu(X_check)).max())
            if diff < 1e-9:
                log.info(f"LightGBM scoring on GPU (check max|diff|={diff:.1e})")
                return gpu
            log.warning(f"GPU tree evaluator differs from LightGBM ({diff:.1e}); using CPU")
    except Exception as e:  # noqa: BLE001
        log.warning(f"GPU tree evaluator unavailable ({e}); using CPU")
    return cpu


def run_stage1_test():
    feats = load_json(MODELS / "stage1_features.json")
    models = [lgb.Booster(model_file=str(MODELS / f"s1_fold{k}.txt")) for k in range(CFG["n_folds"])]
    predict = None
    for i in range(n_shards("test")):
        out = shard_path("stage1", "test", i)
        if out.exists():
            continue
        d = pd.read_parquet(shard_path("feat", "test", i), columns=["q", "c"] + feats)
        X = d[feats].values.astype(np.float32)
        if predict is None:
            predict = fold_predictor(models, X[:20_000])
        p = predict(X)
        save_parquet(d[["q", "c"]].assign(p1=p.astype(np.float32)), out)


# ----------------------------------------------------------------------------------------------- stage-2 set

def _s2(df):
    df = df.sort_values(["q", "p1"], ascending=[True, False], kind="stable")
    g = df.groupby("q")
    df["p1_rank"] = g.cumcount().astype(np.float32)
    df["p1_gap"] = (g.p1.transform("max") - df.p1).astype(np.float32)
    return df[(df.p1_rank < CFG["s2_topn"]) & (df.p1 >= CFG["s2_min_prob"])].reset_index(drop=True)


def run_s2set(split):
    if split == "train":
        s2 = _s2(pd.read_parquet(WORK / "stage1" / "train.parquet"))
        sp = load_splits()
        feat_pos = int(pd.read_parquet(WORK / "stage1" / "train.parquet", columns=["y"]).y.sum())
        allq = np.concatenate([sp["g_stack"], sp["g_val"]])
        rep = {"rows": len(s2), "true_in_blocking": feat_pos, "true_in_s2": int(s2.y.sum()),
               "true_total": int(sp["n_true"][allq].sum())}
        rep["s2_recall_vs_blocking"] = rep["true_in_s2"] / max(feat_pos, 1)
        rep["s2_recall_vs_total"] = rep["true_in_s2"] / max(rep["true_total"], 1)
        save_json(rep, WORK / "reports" / "s2set_train.json")
        log.info(f"stage-2 set: {rep}")
    else:
        s2 = _s2(pd.concat([pd.read_parquet(shard_path("stage1", "test", i)) for i in range(n_shards("test"))],
                           ignore_index=True))
    save_parquet(s2, WORK / "s2" / f"{split}.parquet")


# ----------------------------------------------------------------------------------------------- stage 2

def _attach(split, s2):
    """Join pair features and cross-encoder scores onto the stage-2 rows."""
    from cross_encoder import load_ce_scores
    qs = np.unique(s2.q.values)
    parts = []
    key = s2[["q", "c"]]
    for i in range(n_shards(split)):
        d = pd.read_parquet(shard_path("feat", split, i))
        d = d[np.isin(d.q.values, qs)]
        parts.append(key[np.isin(key.q.values, np.unique(d.q.values))].merge(d, on=["q", "c"], how="left"))
    f = pd.concat(parts, ignore_index=True)
    s2 = s2.merge(f.drop(columns=[c for c in ("y",) if c in s2.columns]), on=["q", "c"], how="left")
    ce = load_ce_scores(split)
    s2 = s2.merge(ce, on=["q", "c"], how="left")
    g = s2.groupby("q")
    s2["ce_rank"] = g.ce.rank(ascending=False, method="first").astype(np.float32)
    s2["ce_gap"] = (g.ce.transform("max") - s2.ce).astype(np.float32)
    return s2


def stage2_features(s2, variant):
    base = [c for c in s2.columns if c not in ("q", "c", "y", "ce", "ce_rank", "ce_gap")]
    return base + (["ce", "ce_rank", "ce_gap"] if variant == "ce" else [])


def run_stage2_train():
    sp = load_splits()
    s2 = pd.read_parquet(WORK / "s2" / "train.parquet")
    s2 = _attach("train", s2)
    stack = s2[np.isin(s2.q.values, sp["g_stack"])].reset_index(drop=True)
    val = s2[np.isin(s2.q.values, sp["g_val"])].reset_index(drop=True)
    F = CFG["n_folds"]
    fold = fold_of(stack.q.values, F)
    country = load_prep("train", "s1", ["country"]).country.values
    results, preds = {}, {}
    with Pool(n_workers()) as pool:
        for variant in ("noce", "ce"):
            feats = stage2_features(s2, variant)
            oof, vp = np.zeros(len(stack)), np.zeros(len(val))
            for k in range(F):
                tr = stack[fold != k]
                m = train_lgb(tr[feats].values.astype(np.float32), tr.y.values, tr.q.values,
                              MODELS / f"s2_{variant}_fold{k}.txt", 100 + k)
                oof[fold == k] = m.predict(stack.loc[fold == k, feats].values.astype(np.float32))
                vp += m.predict(val[feats].values.astype(np.float32)) / F
            cal = crossfit_isotonic(oof, stack.y.values, stack.q.values, F)
            preds[variant] = (oof, vp, feats)
            d = stack[["q", "c"]].assign(p=cal)
            fams = [("expf05", None)] + [("threshold", t) for t in CFG["thr_grid"]]
            for fam, thr in fams:
                sel = decide(d, fam, thr, CFG["arb_max_iter"], pool)
                f, _ = macro_f05(sel, sp["g_stack"], sp["n_true"], sp["gt_q"], sp["gt_c"])
                results[(variant, fam, thr)] = f
                log.info(f"stack F0.5 variant={variant} family={fam} thr={thr}: {f:.5f}")
        best = max(results, key=results.get)
        variant, fam, thr = best
        oof, vp, feats = preds[variant]
        iso = fit_isotonic(oof, stack.y.values)
        vcal = iso.predict(vp)
        sel = decide(val[["q", "c"]].assign(p=vcal), fam, thr, CFG["arb_max_iter"], pool)
    f, per = macro_f05(sel, sp["g_val"], sp["n_true"], sp["gt_q"], sp["gt_c"])
    per["country"] = country[per.q.values]
    by_c = per.groupby("country").f.mean().to_dict()
    npred = sel.groupby("q").size()
    val_stats = _pred_stats(sel, sp["g_val"], country)
    bundle = {"variant": variant, "family": fam, "thr": thr, "features": feats,
              "stack_scores": {f"{v}|{fa}|{t}": s for (v, fa, t), s in results.items()},
              "val_f05": f, "val_f05_by_country": by_c, "val_pred_stats": val_stats}
    with open(MODELS / "stage2_iso.pkl", "wb") as fh:
        pickle.dump(iso, fh)
    save_json(bundle, MODELS / "stage2_bundle.json")
    # error analysis + importance
    imp = pd.DataFrame({"feature": feats})
    for k in range(F):
        m = lgb.Booster(model_file=str(MODELS / f"s2_{variant}_fold{k}.txt"))
        imp[f"gain_{k}"] = m.feature_importance("gain")
    imp["gain"] = imp.filter(like="gain_").mean(1)
    imp.sort_values("gain", ascending=False)[["feature", "gain"]].to_csv(WORK / "reports" / "feature_importance.tsv", sep="\t", index=False)
    per.sort_values("f").head(2000).to_csv(WORK / "reports" / "val_errors.tsv", sep="\t", index=False)
    lines = [f"# Stage-2 validation report", "",
             f"- chosen: variant **{variant}**, decision **{fam}**" + (f" (thr {thr})" if thr else ""),
             f"- validation macro F0.5: **{f:.5f}**",
             f"- by country: {', '.join(f'{k}: {v:.5f}' for k, v in by_c.items())}",
             f"- stack F0.5 of all options: see models/stage2_bundle.json",
             f"- predicted matches per S1 (val): {val_stats}"]
    (WORK / "reports" / "report.md").write_text("\n".join(lines) + "\n")
    log.info("\n".join(lines))


def _pred_stats(sel, qs, country):
    n = pd.Series(0, index=qs).add(sel.groupby("q").size(), fill_value=0).loc[qs]
    c = pd.Series(country[qs], index=qs)
    out = {}
    for cc, g in n.groupby(c):
        out[cc] = {"n_s1": int(len(g)), "empty_rate": round(float((g == 0).mean()), 4),
                   "mean_pred": round(float(g.mean()), 3),
                   "dist": {str(k): round(float(v), 4) for k, v in g.clip(upper=6).value_counts(normalize=True).sort_index().items()}}
    return out


# ----------------------------------------------------------------------------------------------- test + submission

def run_stage2_test():
    bundle = load_json(MODELS / "stage2_bundle.json")
    with open(MODELS / "stage2_iso.pkl", "rb") as fh:
        iso = pickle.load(fh)
    s2 = _attach("test", pd.read_parquet(WORK / "s2" / "test.parquet"))
    feats = bundle["features"]
    X = s2[feats].values.astype(np.float32)
    models = [lgb.Booster(model_file=str(MODELS / f"s2_{bundle['variant']}_fold{k}.txt")) for k in range(CFG["n_folds"])]
    raw = fold_predictor(models, X[:20_000])(X)
    s2["p"] = iso.predict(raw)
    save_parquet(s2[["q", "c", "p1", "p"] + (["ce"] if "ce" in s2 else [])], WORK / "stage2" / "test_scores.parquet")
    with Pool(n_workers()) as pool:
        sel = decide(s2[["q", "c", "p"]], bundle["family"], bundle["thr"], CFG["arb_max_iter"], pool)
    write_submission(s2, sel, bundle)


def write_submission(s2, sel, bundle):
    sp = load_splits()
    s1 = load_prep("test", "s1", ["id", "country"])
    pool_ids = load_prep("test", "pool", ["id"]).id.values
    q_all = sp["q_test"]
    OUTPUT.mkdir(parents=True, exist_ok=True)

    def lists(df):
        d = df.assign(cid=pool_ids[df.c.values]).groupby("q").cid.apply(lambda x: ",".join(dict.fromkeys(x)))
        return pd.Series(q_all).map(d).fillna("").values

    cand = pd.DataFrame({"source1_entity_id": s1.id.values[q_all], "candidate_entity_ids": lists(s2)})
    match = pd.DataFrame({"source1_entity_id": s1.id.values[q_all], "matched_entity_ids": lists(sel)})
    for df, name in ((match, "matching_results.tsv"), (cand, "candidate_pairs.tsv")):
        tmp = OUTPUT / f".{name}.tmp"
        df.to_csv(tmp, sep="\t", index=False)
        tmp.replace(OUTPUT / name)
    # in-script checks
    assert match.source1_entity_id.is_unique and len(match) == len(q_all)
    cset = dict(zip(cand.source1_entity_id, cand.candidate_entity_ids))
    for sid, m in zip(match.source1_entity_id, match.matched_entity_ids):
        if m:
            ids = m.split(",")
            assert len(ids) == len(set(ids)) and all(x[:3] in ("S2-", "S3-") for x in ids)
            assert set(ids) <= set(cset[sid].split(",")), f"match outside candidates for {sid}"
    stats = _pred_stats(sel, q_all, s1.country.values)
    rep = {"test_pred_stats": stats, "val_pred_stats": bundle["val_pred_stats"],
           "train_singleton_rate": float((sp["n_true"] == 0).mean()),
           "train_mean_true": float(sp["n_true"].mean())}
    val_mean = np.mean([v["mean_pred"] for v in bundle["val_pred_stats"].values()])
    val_empty = np.mean([v["empty_rate"] for v in bundle["val_pred_stats"].values()])
    flags = []
    for c, v in stats.items():
        if abs(v["mean_pred"] - val_mean) > 0.35 * val_mean or abs(v["empty_rate"] - val_empty) > 0.05:
            flags.append(f"{c}: mean_pred={v['mean_pred']} empty_rate={v['empty_rate']} vs val {val_mean:.3f}/{val_empty:.3f}")
    rep["country_shift_flags"] = flags
    if CFG["test_limit"] is None:
        r = subprocess.run([sys.executable, str(PATHS["validator"]), "--matching", str(OUTPUT / "matching_results.tsv"),
                            "--candidate", str(OUTPUT / "candidate_pairs.tsv"), "--test-dir", str(PATHS["raw_test_dir"]), "--check-ids"],
                           capture_output=True, text=True)
        rep["validator"] = {"returncode": r.returncode, "stdout": r.stdout[-4000:], "stderr": r.stderr[-2000:]}
        if r.returncode != 0:
            save_json(rep, WORK / "reports" / "test_report.json")
            raise RuntimeError(f"validator failed:\n{r.stdout}\n{r.stderr}")
    save_json(rep, WORK / "reports" / "test_report.json")
    log.info(f"submission written to {OUTPUT}; per-country prediction stats: {stats}; flags: {flags}")


# ----------------------------------------------------------------------------------------------- split B proxy

def run_eval_b():
    """Leave-one-country-out proxy for unseen France: train on one country's stack, evaluate on the other
    country's val (stage-1 features, no cross-encoder), compared with same-country val."""
    sp = load_splits()
    country = load_prep("train", "s1", ["country"]).country.values
    df = load_feats("train")
    feats = feature_columns(df)
    rep = {}
    with Pool(n_workers()) as pool:
        for tr_c, te_c in (("us", "india"), ("india", "us")):
            st = df[np.isin(df.q.values, sp["g_stack"]) & (country[df.q.values] == tr_c)].reset_index(drop=True)
            fold = fold_of(st.q.values, 3)
            oof = np.zeros(len(st))
            ms = []
            for k in range(3):
                tr = st[fold != k]
                m = train_lgb(tr[feats].values.astype(np.float32), tr.y.values, tr.q.values,
                              MODELS / f"B_{tr_c}_fold{k}.txt", 200 + k)
                oof[fold == k] = m.predict(st.loc[fold == k, feats].values.astype(np.float32))
                ms.append(m)
            iso = fit_isotonic(oof, st.y.values)
            for ev_c in (tr_c, te_c):
                qv = sp["g_val"][country[sp["g_val"]] == ev_c]
                v = df[np.isin(df.q.values, qv)]
                p = iso.predict(sum(m.predict(v[feats].values.astype(np.float32)) for m in ms) / 3)
                sel = decide(v[["q", "c"]].assign(p=p), "expf05", None, CFG["arb_max_iter"], pool)
                f, _ = macro_f05(sel, qv, sp["n_true"], sp["gt_q"], sp["gt_c"])
                rep[f"train_{tr_c}->val_{ev_c}"] = {"f05": f, **_pred_stats(sel, qv, country)[ev_c]}
                log.info(f"split B train {tr_c} -> val {ev_c}: F0.5={f:.5f}")
    save_json(rep, WORK / "reports" / "split_B.json")