#!/usr/bin/env python3
"""跑后补充分析（非预注册，见 AMENDMENTS.md）：top-k 命中率，以及度数校正后的排序。

为什么必须补：需求方实际要的交付物是 top 50 基因表，而全局 AUROC 是整条排序的平均性质，
完全可能「整体略优于零模型」但「最顶端那 50 个全是 hub 基因」。只报 AUROC 会掩盖这一点。

Pass 1：3000 次度数配对零模型抽样，流式累计每个基因的零模型分数均值/标准差，
        同时记录每次抽样的 precision@k 分布。
Pass 2：换一批独立的随机种子再抽 3000 次，用 Pass 1 的统计量算 z 分数排序，
        得到「度数校正排序」的 precision@k 零模型分布。
"""
import json, os
import numpy as np
import scipy.sparse as sp
import importlib.util

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "..", "results")
spec = importlib.util.spec_from_file_location("score", os.path.join(HERE, "03_score.py"))
S = importlib.util.module_from_spec(spec); spec.loader.exec_module(S)

KS = [10, 20, 50, 100, 200, 500]
N_DRAW = 1000
PASS2_OFFSET = 100


def precision_at_k(scores, labels, ks):
    order = np.argsort(-scores, kind="mergesort")
    y = labels[order]
    cum = np.cumsum(y)
    return {k: float(cum[k - 1] / k) for k in ks if k <= len(y)}


def main():
    A = sp.load_npz(os.path.join(RES, "string_adj.npz")).tocsr()
    z = np.load(os.path.join(RES, "scores.npz"))
    rwr, deg_w, mask, labels = z["rwr_full"], z["deg_w"], z["mask"], z["labels"]
    gout_idx = z["gout_idx"]
    n = A.shape[0]
    P_T = S.row_normalize(A).T.tocsr()
    bin_of, members = S.degree_bins(deg_w)
    forbidden = set(int(i) for i in gout_idx)
    prevalence = float(labels.mean())

    obs = {"gout_rwr": precision_at_k(rwr[mask], labels, KS),
           "degree_only": precision_at_k(deg_w[mask], labels, KS),
           "prevalence": prevalence}

    # ---- Pass 1：零模型 precision@k + 每基因零模型分数统计量 ----
    ssum = np.zeros(n); ssq = np.zeros(n); ndraw = 0
    null_prec = {k: [] for k in KS}
    for s in S.SEEDS:
        rng = np.random.default_rng(s)
        for start in range(0, N_DRAW, 100):
            sets = [S.sample_degree_matched(rng, gout_idx, bin_of, members, forbidden) for _ in range(100)]
            X, _ = S.rwr_batch(P_T, S.seed_vectors(n, sets))
            ssum += X.sum(axis=1); ssq += (X ** 2).sum(axis=1); ndraw += X.shape[1]
            for c in range(X.shape[1]):
                p = precision_at_k(X[mask, c], labels, KS)
                for k in KS:
                    null_prec[k].append(p[k])
        print(f"pass1 seed={s} done", flush=True)
    nmean = ssum / ndraw
    nsd = np.sqrt(np.maximum(ssq / ndraw - nmean ** 2, 0))
    nsd[nsd == 0] = np.inf

    # ---- 度数校正排序：z = (实测 - 零模型均值) / 零模型标准差 ----
    zscore = (rwr - nmean) / nsd
    obs["degree_corrected_z"] = precision_at_k(zscore[mask], labels, KS)
    auc_z = S.auroc(zscore[mask], labels)

    # ---- Pass 2：独立抽样，给 z 排序做零模型 ----
    null_prec_z = {k: [] for k in KS}
    null_auc_z = []
    for s in S.SEEDS:
        rng = np.random.default_rng(s + PASS2_OFFSET)
        for start in range(0, N_DRAW, 100):
            sets = [S.sample_degree_matched(rng, gout_idx, bin_of, members, forbidden) for _ in range(100)]
            X, _ = S.rwr_batch(P_T, S.seed_vectors(n, sets))
            Z = (X - nmean[:, None]) / nsd[:, None]
            for c in range(Z.shape[1]):
                p = precision_at_k(Z[mask, c], labels, KS)
                for k in KS:
                    null_prec_z[k].append(p[k])
                null_auc_z.append(S.auroc(Z[mask, c], labels))
        print(f"pass2 seed={s} done", flush=True)

    def emp_p(nulls, obs_val):
        a = np.array(nulls)
        return float((1 + (a >= obs_val).sum()) / (len(a) + 1))

    out = {
        "note": "非预注册的跑后补充分析，见 AMENDMENTS.md 第 1 条。",
        "prevalence_in_eval_universe": prevalence,
        "n_eval": int(mask.sum()), "n_positive": int(labels.sum()),
        "n_null_draws_per_pass": ndraw,
        "precision_at_k": {
            "gout_rwr": {str(k): {"value": obs["gout_rwr"][k],
                                  "n_hits": int(round(obs["gout_rwr"][k] * k)),
                                  "null_mean": float(np.mean(null_prec[k])),
                                  "null_p95": float(np.quantile(null_prec[k], 0.95)),
                                  "empirical_p": emp_p(null_prec[k], obs["gout_rwr"][k]),
                                  "fold_over_prevalence": obs["gout_rwr"][k] / prevalence}
                         for k in KS},
            "degree_only": {str(k): {"value": obs["degree_only"][k],
                                     "n_hits": int(round(obs["degree_only"][k] * k)),
                                     "fold_over_prevalence": obs["degree_only"][k] / prevalence}
                            for k in KS},
            "degree_corrected_z": {str(k): {"value": obs["degree_corrected_z"][k],
                                            "n_hits": int(round(obs["degree_corrected_z"][k] * k)),
                                            "null_mean": float(np.mean(null_prec_z[k])),
                                            "null_p95": float(np.quantile(null_prec_z[k], 0.95)),
                                            "empirical_p": emp_p(null_prec_z[k], obs["degree_corrected_z"][k]),
                                            "fold_over_prevalence": obs["degree_corrected_z"][k] / prevalence}
                                   for k in KS},
        },
        "auroc_degree_corrected_z": auc_z,
        "auroc_degree_corrected_z_null_mean": float(np.mean(null_auc_z)),
        "auroc_degree_corrected_z_empirical_p": emp_p(null_auc_z, auc_z),
    }
    json.dump(out, open(os.path.join(RES, "topk_precision.json"), "w"), ensure_ascii=False, indent=1)
    np.savez_compressed(os.path.join(RES, "null_gene_stats.npz"), null_mean=nmean, null_sd=nsd, zscore=zscore)
    print(json.dumps({"prevalence": round(prevalence, 5),
                      "p@50_gout_rwr": out["precision_at_k"]["gout_rwr"]["50"],
                      "p@50_degree_corrected": out["precision_at_k"]["degree_corrected_z"]["50"],
                      "auroc_z": auc_z, "auroc_z_emp_p": out["auroc_degree_corrected_z_empirical_p"]},
                     ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()