"""High-order sparse graph propagation features on top of rich RW stack."""

from __future__ import annotations

import argparse
import importlib.util
import sys
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
from gensim.models import Word2Vec
from scipy import sparse
from sklearn.metrics import precision_recall_curve, roc_auc_score
from sklearn.model_selection import StratifiedKFold


def load_module(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


def read_txt(path: Path) -> list[list[int]]:
    return [list(map(int, line.strip().split())) for line in path.open()]


def best_f1(y: np.ndarray, s: np.ndarray):
    p, r, t = precision_recall_curve(y, s)
    f = 2 * p * r / (p + r + 1e-12)
    i = int(np.argmax(f))
    return float(f[i]), float(t[i] if i < len(t) else 0.5), float(roc_auc_score(y, s)), float(p[i]), float(r[i])


def row_norm(mat: sparse.csr_matrix) -> sparse.csr_matrix:
    mat = mat.tocsr().astype(np.float32)
    deg = np.asarray(mat.sum(axis=1)).ravel()
    inv = np.zeros_like(deg, dtype=np.float32)
    inv[deg > 0] = 1.0 / deg[deg > 0]
    return sparse.diags(inv).dot(mat).tocsr()


def topk_prune_rows(mat: sparse.csr_matrix, k: int) -> sparse.csr_matrix:
    mat = mat.tocsr()
    if k <= 0:
        return mat
    data = []
    indices = []
    indptr = [0]
    for i in range(mat.shape[0]):
        lo, hi = mat.indptr[i], mat.indptr[i + 1]
        vals = mat.data[lo:hi]
        cols = mat.indices[lo:hi]
        if len(vals) > k:
            keep = np.argpartition(vals, -k)[-k:]
            order = np.argsort(cols[keep])
            keep = keep[order]
            vals = vals[keep]
            cols = cols[keep]
        data.append(vals)
        indices.append(cols)
        indptr.append(indptr[-1] + len(vals))
    return sparse.csr_matrix((np.concatenate(data).astype(np.float32), np.concatenate(indices).astype(np.int32), np.asarray(indptr, dtype=np.int32)), shape=mat.shape)


def extract_scores(mat: sparse.csr_matrix, pairs: np.ndarray) -> np.ndarray:
    return np.asarray(mat[pairs[:, 0], pairs[:, 1]]).ravel().astype(np.float32)


def build_high_order(root: Path, train_refs: pd.DataFrame, pairs: np.ndarray, tag: str, topk: int = 1500) -> np.ndarray:
    cache = root / "validation_runs" / "feature_cache"
    cache.mkdir(parents=True, exist_ok=True)
    path = cache / f"high_order_{tag}_{len(pairs)}_{int(pairs[:,0].sum())}_{int(pairs[:,1].sum())}_k{topk}.npy"
    if path.exists():
        return np.load(path)

    train = train_refs[["source", "target"]].to_numpy(np.int64)
    ap = sparse.csr_matrix((np.ones(len(train), dtype=np.float32), (train[:, 0], train[:, 1])), shape=(6611, 79937))
    apn = row_norm(ap)

    cite = np.array(read_txt(root / "data_and_docs/paper_file_ann.txt"), dtype=np.int64)
    pp = sparse.csr_matrix((np.ones(len(cite), dtype=np.float32), (cite[:, 0], cite[:, 1])), shape=(79937, 79937))
    pp = (pp + pp.T).astype(np.float32)
    pp.data[:] = 1.0
    ppn = row_norm(pp)

    co = np.array(read_txt(root / "data_and_docs/author_file_ann.txt"), dtype=np.int64)
    aa = sparse.csr_matrix((np.ones(len(co), dtype=np.float32), (co[:, 0], co[:, 1])), shape=(6611, 6611))
    aa = (aa + aa.T).astype(np.float32)
    aa.data[:] = 1.0
    aan = row_norm(aa)

    paper_deg = np.asarray(ap.sum(axis=0)).ravel().astype(np.float32)
    cite_deg = np.asarray(pp.sum(axis=0)).ravel().astype(np.float32)
    denom = np.log1p(paper_deg[pairs[:, 1]] + cite_deg[pairs[:, 1]] + 1.0).astype(np.float32)

    cols = []
    names = []
    S = apn.copy()
    for k in range(1, 5):
        S = topk_prune_rows(S.dot(ppn).tocsr(), topk)
        s = extract_scores(S, pairs)
        cols.extend([s, s / (denom + 1e-6), np.log1p(s * 1000.0)])
        names.extend([f"ap_pp{k}", f"ap_pp{k}_popnorm", f"ap_pp{k}_log"])

    C = topk_prune_rows(aan.dot(apn).tocsr(), topk)
    for k in range(0, 4):
        if k > 0:
            C = topk_prune_rows(C.dot(ppn).tocsr(), topk)
        s = extract_scores(C, pairs)
        cols.extend([s, s / (denom + 1e-6), np.log1p(s * 1000.0)])
        names.extend([f"aa_ap_pp{k}", f"aa_ap_pp{k}_popnorm", f"aa_ap_pp{k}_log"])

    # Blend historical and coauthor propagation as a cheap agreement signal.
    H = np.column_stack(cols).astype(np.float32)
    np.save(path, H)
    (cache / f"high_order_{tag}_names.txt").write_text("\n".join(names) + "\n")
    return H


def build_high_order_directed(root: Path, train_refs: pd.DataFrame, pairs: np.ndarray, tag: str, topk: int = 1500) -> np.ndarray:
    cache = root / "validation_runs" / "feature_cache"
    cache.mkdir(parents=True, exist_ok=True)
    path = cache / f"high_order_directed_{tag}_{len(pairs)}_{int(pairs[:,0].sum())}_{int(pairs[:,1].sum())}_k{topk}.npy"
    if path.exists():
        return np.load(path)

    train = train_refs[["source", "target"]].to_numpy(np.int64)
    ap = sparse.csr_matrix((np.ones(len(train), dtype=np.float32), (train[:, 0], train[:, 1])), shape=(6611, 79937))
    apn = row_norm(ap)
    cite = np.array(read_txt(root / "data_and_docs/paper_file_ann.txt"), dtype=np.int64)
    pp_fwd = sparse.csr_matrix((np.ones(len(cite), dtype=np.float32), (cite[:, 0], cite[:, 1])), shape=(79937, 79937))
    pp_bwd = pp_fwd.T.tocsr()
    pp_undir = (pp_fwd + pp_bwd).astype(np.float32)
    pp_undir.data[:] = 1.0
    matrices = {
        "fwd": row_norm(pp_fwd),
        "bwd": row_norm(pp_bwd),
        "undir": row_norm(pp_undir),
    }

    co = np.array(read_txt(root / "data_and_docs/author_file_ann.txt"), dtype=np.int64)
    aa = sparse.csr_matrix((np.ones(len(co), dtype=np.float32), (co[:, 0], co[:, 1])), shape=(6611, 6611))
    aa = (aa + aa.T).astype(np.float32)
    aa.data[:] = 1.0
    aan = row_norm(aa)
    paper_deg = np.asarray(ap.sum(axis=0)).ravel().astype(np.float32)
    cite_deg = np.asarray(pp_undir.sum(axis=0)).ravel().astype(np.float32)
    denom = np.log1p(paper_deg[pairs[:, 1]] + cite_deg[pairs[:, 1]] + 1.0).astype(np.float32)

    cols = []
    names = []
    for label, ppn in matrices.items():
        S = apn.copy()
        prev = None
        for k in range(1, 4):
            S = topk_prune_rows(S.dot(ppn).tocsr(), topk)
            s = extract_scores(S, pairs)
            cols.extend([s, s / (denom + 1e-6), s - (prev if prev is not None else 0.0)])
            names.extend([f"ap_{label}{k}", f"ap_{label}{k}_popnorm", f"ap_{label}{k}_delta"])
            prev = s
        C = topk_prune_rows(aan.dot(apn).tocsr(), topk)
        for k in range(0, 3):
            if k > 0:
                C = topk_prune_rows(C.dot(ppn).tocsr(), topk)
            s = extract_scores(C, pairs)
            cols.extend([s, s / (denom + 1e-6)])
            names.extend([f"aa_ap_{label}{k}", f"aa_ap_{label}{k}_popnorm"])
    H = np.column_stack(cols).astype(np.float32)
    np.save(path, H)
    (cache / f"high_order_directed_{tag}_names.txt").write_text("\n".join(names) + "\n")
    return H


def fit_lgb_oof(X: np.ndarray, y: np.ndarray, seed: int, n_splits: int) -> np.ndarray:
    oof = np.zeros(len(y), dtype=np.float32)
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
    for fold, (tr, va) in enumerate(skf.split(X, y), start=1):
        clf = lgb.LGBMClassifier(
            n_estimators=1200,
            learning_rate=0.025,
            num_leaves=15,
            subsample=0.9,
            colsample_bytree=0.9,
            reg_lambda=8.0,
            min_child_samples=100,
            objective="binary",
            n_jobs=8,
            verbose=-1,
            random_state=seed + fold,
        )
        clf.fit(X[tr], y[tr])
        oof[va] = clf.predict_proba(X[va])[:, 1].astype(np.float32)
    return oof


def fit_full_predict(X: np.ndarray, y: np.ndarray, Xt: np.ndarray, seed: int) -> np.ndarray:
    clf = lgb.LGBMClassifier(
        n_estimators=1400,
        learning_rate=0.022,
        num_leaves=15,
        subsample=0.9,
        colsample_bytree=0.9,
        reg_lambda=8.0,
        min_child_samples=100,
        objective="binary",
        n_jobs=8,
        verbose=-1,
        random_state=seed,
    )
    clf.fit(X, y)
    return clf.predict_proba(Xt)[:, 1].astype(np.float32)


def write_ratio_submission(path: Path, score: np.ndarray, ratio: float, known: np.ndarray, anchor: np.ndarray) -> dict:
    pred = np.zeros(len(score), dtype=np.int8)
    pred[np.argsort(score, kind="mergesort")[-int(round(len(score) * ratio)):]] = 1
    pred[known] = 1
    pd.DataFrame({"Index": np.arange(len(pred), dtype=np.int64), "Predicted": pred}).to_csv(path, index=False)
    return {
        "path": str(path),
        "ratio": float(ratio),
        "positive_ratio": float(pred.mean()),
        "changed_vs_anchor": int((pred != anchor).sum()),
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--package-root", type=Path, default=Path(__file__).resolve().parents[1])
    ap.add_argument("--split-seed", type=int, default=202)
    ap.add_argument("--seed", type=int, default=202)
    ap.add_argument("--n-splits", type=int, default=5)
    ap.add_argument("--make-submission", action="store_true")
    args = ap.parse_args()

    root = args.package_root.resolve()
    rw = load_module("rw", root / "code/randomwalk_systematic_ablation.py")
    stack = load_module("stack", root / "code/stack_rank_calibration.py")
    rich = load_module("rich", root / "code/content_rich_ablation.py")
    ens = load_module("ens", root / "code/generate_randomwalk_ensemble_submission.py")
    main_val = root / "validation_runs/dynamic_seed202/dyn202_l2d512_bpr_bigbatch_more/scores/val_vanilla_ensemble_mean.npy"
    train_refs, pairs, y, X_base = rw.build_base_features(root, args.split_seed, main_val)
    builder = stack.ExplicitGraphFeatures(root, train_refs)
    X_rich = rich.content_rich_features(root, pairs, builder)

    versions = [
        "dw_base_d128_l40_w10_win10",
        "dw_long_d128_l80_w10_win10",
        "dw_highdim_d256_l40_w10_win10",
        "dw_d256_l80_w10_win10",
        "dw_seed3407_d128_l40_w10_win10",
        "dw_graph_ap_pp",
        "n2v_p2_q1_d128_l40_w10_win10",
    ]
    cfgs = {c.version_name: c for c in rw.small_configs() + rw.graph_configs() + rw.extra_configs()}
    sys_dir = root / "validation_runs/dynamic_seed202/randomwalk_systematic"
    blocks = []
    for version in versions:
        cfg = cfgs[version]
        model = Word2Vec.load(str(sys_dir / "models" / f"{version}.model"))
        block, _ = rw.pair_feature_block(model, pairs, cfg, root, args.split_seed, train_refs)
        blocks.append(block)
    X_high = build_high_order(root, train_refs, pairs, "val202")
    X_high_dir = build_high_order_directed(root, train_refs, pairs, "val202")
    out = root / "validation_runs/dynamic_seed202/high_order_graph_stack"
    out.mkdir(parents=True, exist_ok=True)

    rows = []
    for name, X in [
        ("rich_rw7", np.column_stack([X_base, X_rich, *blocks, ens.aggregate(blocks)]).astype(np.float32)),
        ("rich_rw7_highorder", np.column_stack([X_base, X_rich, *blocks, ens.aggregate(blocks), X_high]).astype(np.float32)),
        ("rich_rw7_highorder_directed", np.column_stack([X_base, X_rich, *blocks, ens.aggregate(blocks), X_high, X_high_dir]).astype(np.float32)),
        ("base_highorder", np.column_stack([X_base, X_high]).astype(np.float32)),
    ]:
        print("fit", name, X.shape)
        oof = fit_lgb_oof(X, y, args.seed + len(rows) * 19, args.n_splits)
        np.save(out / f"{name}_oof.npy", oof)
        f1, th, auc, p, r = best_f1(y, oof)
        rows.append({"stage": name, "validation_f1": f1, "threshold": th, "auc": auc, "precision": p, "recall": r, "n_features": X.shape[1]})
        print(rows[-1])
    pd.DataFrame(rows).sort_values("validation_f1", ascending=False).to_csv(out / "validation_summary.csv", index=False)
    print(pd.DataFrame(rows).sort_values("validation_f1", ascending=False).to_string(index=False))

    if not args.make_submission:
        return

    best_X = np.column_stack([X_base, X_rich, *blocks, ens.aggregate(blocks), X_high, X_high_dir]).astype(np.float32)
    gen = load_module("gen", root / "code/generate_post95_submission.py")
    post = load_module("post", root / "code/post95_ablation.py")
    extra = load_module("extra", root / "code/extra_score_sources_ablation.py")
    test_pairs = np.array(read_txt(root / "data_and_docs/bipartite_test_ann.txt"), dtype=np.int64)
    main_test = np.load(root / "validation_runs/dynamic_seed202/post95_test_scores/dyn202_l2d512_bpr_bigbatch_more/scores/test_vanilla_ensemble_mean.npy").astype(np.float32)
    full_refs = pd.DataFrame(read_txt(root / "data_and_docs/bipartite_train_ann.txt"), columns=["source", "target"])
    test_builder = stack.ExplicitGraphFeatures(root, full_refs)
    Xht = test_builder.transform(test_pairs)
    Xt = np.column_stack(
        [
            stack.add_rank_features(test_pairs, main_test),
            Xht,
            post.negative_evidence_features(Xht, main_test),
            gen.topk_content_similarity_fast(root, test_pairs, test_builder),
        ]
    ).astype(np.float32)
    selected = [Path(x.strip()) for x in (root / "validation_runs/dynamic_seed202/post95_submission/selected_variant_val_scores.txt").read_text().splitlines() if x.strip()]
    test_scores = []
    for p in selected:
        rel = p.resolve().relative_to(root / "validation_runs/dynamic_seed202")
        tp = root / "validation_runs/dynamic_seed202/post95_test_scores" / rel.parent / rel.name.replace("val_", "test_", 1)
        test_scores.append(np.load(tp).astype(np.float32))
    Xt = np.column_stack([Xt, gen.variant_feature_matrix(post, test_scores)]).astype(np.float32)
    content_test = extra.content_mean_score(root, test_pairs, test_builder)
    mf_test = np.load(root / "validation_runs/dynamic_seed202/extra_bprmf_submission/test_mf_bpr_dynamic_s202_d256_e220.npy").astype(np.float32)
    Xct, _ = extra.score_to_features(content_test, "content_mean_cos", test_pairs)
    Xmt, _ = extra.score_to_features(mf_test, "mf_bpr", test_pairs)
    Xt = np.column_stack([Xt, Xct, Xmt, rich.content_rich_features(root, test_pairs, test_builder)]).astype(np.float32)
    test_blocks = []
    for version in versions:
        cfg = cfgs[version]
        model = Word2Vec.load(str(sys_dir / "models" / f"{version}.model"))
        block, _ = rw.pair_feature_block(model, test_pairs, cfg, root, args.split_seed, full_refs)
        test_blocks.append(block)
    X_high_test = build_high_order(root, full_refs, test_pairs, "test_full")
    X_high_dir_test = build_high_order_directed(root, full_refs, test_pairs, "test_full")
    Xt = np.column_stack([Xt, *test_blocks, ens.aggregate(test_blocks), X_high_test, X_high_dir_test]).astype(np.float32)
    print("fit full / predict test", best_X.shape, Xt.shape)
    test_score = fit_full_predict(best_X, y, Xt, args.seed + 900)
    np.save(out / "rich_rw7_highorder_directed_test_pred.npy", test_score)
    known = np.load(root / "cached_scores/test_known_mask.npy").astype(bool)
    anchor = pd.read_csv(root / "validation_runs/dynamic_seed202/node2vec_deepwalk_submission/submission_content_mf_deepwalk_node2vec_lgb_th0.480000.csv")["Predicted"].to_numpy(np.int8)
    sub_dir = out / "submissions"
    sub_dir.mkdir(parents=True, exist_ok=True)
    sub_rows = []
    best_row = max(rows, key=lambda r: r["validation_f1"])
    oof = np.load(out / "rich_rw7_highorder_directed_oof.npy")
    ratios = [0.498, 0.499, 0.500, 0.501, 0.502, float((oof >= best_row["threshold"]).mean())]
    for ratio in ratios:
        path = sub_dir / f"submission_rich_rw7_highorder_directed_r{ratio:.6f}.csv"
        row = write_ratio_submission(path, test_score, ratio, known, anchor)
        row.update(best_row)
        sub_rows.append(row)
    pd.DataFrame(sub_rows).to_csv(out / "submission_summary.csv", index=False)
    print(pd.DataFrame(sub_rows).to_string(index=False))


if __name__ == "__main__":
    main()