"""Score-level meta stack over all cached high-value predictors.

This is intentionally cheap: it reuses OOF validation scores and cached test
scores instead of retraining GNN / random-walk embeddings.
"""

from __future__ import annotations

import argparse
import importlib.util
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
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
    spec.loader.exec_module(module)
    return module


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))
    th = float(t[i]) if i < len(t) else 0.5
    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 prf(y: np.ndarray, pred: np.ndarray):
    tp = int(((pred == 1) & (y == 1)).sum())
    fp = int(((pred == 1) & (y == 0)).sum())
    fn = int(((pred == 0) & (y == 1)).sum())
    p = tp / (tp + fp + 1e-12)
    r = tp / (tp + fn + 1e-12)
    f = 2 * p * r / (p + r + 1e-12)
    return p, r, f, tp, fp, fn


def rank01(x: np.ndarray) -> np.ndarray:
    order = np.argsort(x, kind="mergesort")
    out = np.empty(len(x), dtype=np.float32)
    out[order] = np.linspace(0, 1, len(x), dtype=np.float32)
    return out


def zscore(x: np.ndarray) -> np.ndarray:
    return ((x - x.mean()) / (x.std() + 1e-8)).astype(np.float32)


def author_rank01(pairs: np.ndarray, score: np.ndarray) -> np.ndarray:
    out = np.zeros(len(score), dtype=np.float32)
    df = pd.DataFrame({"idx": np.arange(len(score)), "author": pairs[:, 0], "score": score})
    for _, g in df.groupby("author", sort=False):
        idx = g["idx"].to_numpy()
        order = np.argsort(g["score"].to_numpy(), kind="mergesort")
        vals = np.linspace(0, 1, len(idx), dtype=np.float32) if len(idx) > 1 else np.array([1.0], dtype=np.float32)
        out[idx[order]] = vals
    return out


def add_score_block(pairs: np.ndarray, score: np.ndarray) -> np.ndarray:
    return np.column_stack([score.astype(np.float32), zscore(score), rank01(score), author_rank01(pairs, score)]).astype(np.float32)


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


def load_sources(root: Path, split_seed: int, n_val: int, n_test: int) -> list[tuple[str, np.ndarray, np.ndarray]]:
    base = root / "validation_runs" / f"dynamic_seed{split_seed}"
    candidates: list[tuple[str, Path, Path]] = [
        ("post95_lgb", base / "post95_ablation/ensemble_lgcn_oof.npy", base / "post95_submission/test_post95_ens_pred.npy"),
        ("post95_xgb", base / "post95_ablation/xgboost_76feat_oof.npy", base / "post95_xgboost_submission/test_post95_xgb_pred.npy"),
        ("content_mf", base / "extra_score_sources/bpr_mf_stack_oof.npy", base / "extra_bprmf_submission/test_post95_content_mf_lgb_pred.npy"),
        ("content_rich", base / "content_rich/rich_content_stack_oof.npy", base / "content_rich_submission/test_content_rich_mf_lgb_pred.npy"),
        ("n2v_dw_anchor", base / "node2vec_deepwalk/node2vec_stack_oof.npy", base / "node2vec_deepwalk_submission/test_content_mf_deepwalk_node2vec_lgb_pred.npy"),
    ]
    rw_dir = base / "randomwalk_systematic"
    rw_test_dir = base / "randomwalk_ensemble_submission"
    for p in sorted(rw_dir.glob("*_oof.npy")):
        stem = p.stem
        test = rw_test_dir / f"test_{stem[:-4] if stem.endswith('_oof') else stem}_pred.npy"
        if not test.exists():
            test = rw_test_dir / f"test_{stem.replace('_oof', '')}_pred.npy"
        if test.exists():
            candidates.append((stem.replace("_oof", ""), p, test))
    out = []
    seen = set()
    for name, vp, tp in candidates:
        if name in seen or not vp.exists() or not tp.exists():
            continue
        v = np.load(vp).astype(np.float32)
        t = np.load(tp).astype(np.float32)
        if len(v) == n_val and len(t) == n_test and np.std(v) > 1e-8:
            out.append((name, v, t))
            seen.add(name)
    return out


def fit_oof_predict(X: np.ndarray, y: np.ndarray, X_test: np.ndarray, kind: str, seed: int, n_splits: int):
    oof = np.zeros(len(y), dtype=np.float32)
    test = np.zeros(X_test.shape[0], dtype=np.float32)
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
    models = []
    for fold, (tr, va) in enumerate(skf.split(X, y), start=1):
        if kind == "logreg":
            clf = LogisticRegression(C=0.5, max_iter=1000, solver="lbfgs")
        else:
            clf = lgb.LGBMClassifier(
                n_estimators=900,
                learning_rate=0.025,
                num_leaves=15 if kind == "lgb_small" else 31,
                subsample=0.9,
                colsample_bytree=0.85,
                reg_lambda=10.0,
                min_child_samples=120,
                objective="binary",
                verbose=-1,
                random_state=seed + fold,
            )
        clf.fit(X[tr], y[tr])
        oof[va] = clf.predict_proba(X[va])[:, 1].astype(np.float32)
        test += clf.predict_proba(X_test)[:, 1].astype(np.float32) / n_splits
        models.append(clf)
    return oof, test, models


def write_sub(path: Path, score: np.ndarray, known: np.ndarray, anchor: np.ndarray, *, th: float | None = None, ratio: float | None = None):
    if th is not None:
        pred = (score >= th).astype(np.int8)
    else:
        pred = np.zeros(len(score), dtype=np.int8)
        pred[np.argsort(score, kind="mergesort")[-int(round(len(score) * float(ratio))):]] = 1
    pred[known] = 1
    pd.DataFrame({"Index": np.arange(len(pred), dtype=np.int64), "Predicted": pred}).to_csv(path, index=False)
    return float(pred.mean()), 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)
    args = ap.parse_args()

    root = args.package_root.resolve()
    lgcn = load_module("lgcn", root / "code/train_val_lgcn_ensemble.py")
    train_refs, val_pairs = lgcn.make_notebook_style_split(root, args.split_seed, 0.9)
    pairs = val_pairs[["source", "target"]].to_numpy(np.int64)
    y = val_pairs["label"].to_numpy(np.int8)
    test_pairs = np.array(read_txt(root / "data_and_docs/bipartite_test_ann.txt"), dtype=np.int64)
    known = np.load(root / "cached_scores/test_known_mask.npy").astype(bool)
    anchor_path = root / "validation_runs/dynamic_seed202/node2vec_deepwalk_submission/submission_content_mf_deepwalk_node2vec_lgb_th0.480000.csv"
    anchor = pd.read_csv(anchor_path)["Predicted"].to_numpy(np.int8)

    sources = load_sources(root, args.split_seed, len(y), len(test_pairs))
    print("sources", len(sources))
    for name, v, _ in sources:
        f, th, auc, p, r = best_f1(y, v)
        print(f"{name:90s} f1={f:.6f} th={th:.6f} auc={auc:.6f}")

    blocks = []
    test_blocks = []
    names = []
    raw_val = []
    raw_test = []
    for name, v, t in sources:
        blocks.append(add_score_block(pairs, v))
        test_blocks.append(add_score_block(test_pairs, t))
        names.extend([name, name + "_z", name + "_rank", name + "_author_rank"])
        raw_val.append(rank01(v))
        raw_test.append(rank01(t))
    R = np.vstack(raw_val)
    Rt = np.vstack(raw_test)
    summary = np.column_stack([R.mean(axis=0), R.std(axis=0), R.max(axis=0), R.min(axis=0), (R >= 0.5).sum(axis=0), (R >= 0.9).sum(axis=0)]).astype(np.float32)
    summary_t = np.column_stack([Rt.mean(axis=0), Rt.std(axis=0), Rt.max(axis=0), Rt.min(axis=0), (Rt >= 0.5).sum(axis=0), (Rt >= 0.9).sum(axis=0)]).astype(np.float32)
    X = np.column_stack([*blocks, summary]).astype(np.float32)
    Xt = np.column_stack([*test_blocks, summary_t]).astype(np.float32)
    print("X", X.shape, "Xt", Xt.shape)

    out = root / "validation_runs" / f"dynamic_seed{args.split_seed}" / "score_level_meta_stack"
    sub_dir = out / "submissions"
    out.mkdir(parents=True, exist_ok=True)
    sub_dir.mkdir(parents=True, exist_ok=True)
    pd.Series(names + ["rank_mean", "rank_std", "rank_max", "rank_min", "vote_ge_05", "vote_ge_09"]).to_csv(out / "feature_names.csv", index=False)

    rows = []
    for kind in ["logreg", "lgb_small", "lgb"]:
        oof, test_score, _ = fit_oof_predict(X, y, Xt, kind, args.seed + len(rows) * 31, args.n_splits)
        np.save(out / f"{kind}_oof.npy", oof)
        np.save(out / f"{kind}_test_pred.npy", test_score)
        f, th, auc, p, r = best_f1(y, oof)
        pred = (oof >= th).astype(np.int8)
        _, _, _, tp, fp, fn = prf(y, pred)
        for rule_name, kwargs in [
            (f"{kind}_valbest_th", {"th": th}),
            (f"{kind}_r_valratio", {"ratio": float(pred.mean())}),
            (f"{kind}_r0500", {"ratio": 0.500}),
            (f"{kind}_r0499", {"ratio": 0.499}),
            (f"{kind}_r0501", {"ratio": 0.501}),
        ]:
            path = sub_dir / f"submission_{rule_name}.csv"
            pos_ratio, changed = write_sub(path, test_score, known, anchor, **kwargs)
            rows.append(
                {
                    "experiment": rule_name,
                    "model": kind,
                    "validation_f1": f,
                    "threshold": th,
                    "auc": auc,
                    "precision": p,
                    "recall": r,
                    "val_pred_ratio": float(pred.mean()),
                    "tp": tp,
                    "fp": fp,
                    "fn": fn,
                    "test_positive_ratio": pos_ratio,
                    "changed_vs_anchor": changed,
                    "public_submission_path": str(path),
                }
            )
    pd.DataFrame(rows).sort_values(["validation_f1", "changed_vs_anchor"], ascending=[False, True]).to_csv(out / "summary.csv", index=False)
    print(pd.DataFrame(rows).sort_values(["validation_f1", "changed_vs_anchor"], ascending=[False, True]).to_string(index=False))


if __name__ == "__main__":
    main()