#!/usr/bin/env python3
"""产出 top 50 候选基因表 + 每个基因的知识图谱证据链。

证据链 = 该基因到痛风种子基因的最强路径（1 跳直连，或 2 跳经中介基因），
路径上每条边给出 STRING combined_score，并标注该基因当前是否已有心血管 GWAS 关联、
以及它在度数配对零模型下的分数分位（用来提醒哪些名次可能只是 hub 效应）。
"""
import json, os
import numpy as np
import scipy.sparse as sp

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "..", "results")
TOPN = 50


def main():
    A = sp.load_npz(os.path.join(RES, "string_adj.npz")).tocsr()
    nodes = json.load(open(os.path.join(RES, "network_nodes.json"), encoding="utf-8"))
    symbols = nodes["symbols"]
    ms = json.load(open(os.path.join(RES, "mapped_sets.json"), encoding="utf-8"))
    gs = json.load(open(os.path.join(RES, "gene_sets.json"), encoding="utf-8"))
    z = np.load(os.path.join(RES, "scores.npz"))
    rwr, deg_w, mask = z["rwr_full"], z["deg_w"], z["mask"]

    gout_idx = set(int(i) for i in z["gout_idx"])
    cvd_idx = set(int(i) for i in z["cvd_idx"])
    idx2sym = {i: s for i, s in enumerate(symbols)}
    sym_of_gout = {v: k for k, v in ms["gout"].items()}

    cand = [i for i in range(len(symbols)) if mask[i]]
    cand.sort(key=lambda i: -rwr[i])
    top = cand[:TOPN]

    # 度数分位：这个基因的 RWR 排名里有多少可以用度数解释
    deg_rank = {i: r for r, i in enumerate(sorted(range(len(symbols)), key=lambda j: -deg_w[j]))}
    rwr_rank = {i: r for r, i in enumerate(cand)}

    rows = []
    for i in top:
        nb = A.indices[A.indptr[i]:A.indptr[i + 1]]
        nw = A.data[A.indptr[i]:A.indptr[i + 1]]
        direct = sorted(((int(round(w * 1000)), int(j)) for j, w in zip(nb, nw) if int(j) in gout_idx),
                        reverse=True)[:5]
        chains = [{"hops": 1,
                   "path": [idx2sym[i], idx2sym[j]],
                   "edges": [{"a": idx2sym[i], "b": idx2sym[j], "string_score": sc}],
                   "gout_anchor": idx2sym[j],
                   "gout_anchor_min_p": gs["gout_genes"][sym_of_gout[j]]["min_p"]}
                  for sc, j in direct]
        if len(chains) < 3:           # 直连不够就补两跳最强路径
            two = []
            nbset = {int(j): float(w) for j, w in zip(nb, nw)}
            for m, w1 in nbset.items():
                if m in gout_idx:
                    continue
                m_nb = A.indices[A.indptr[m]:A.indptr[m + 1]]
                m_w = A.data[A.indptr[m]:A.indptr[m + 1]]
                for j, w2 in zip(m_nb, m_w):
                    if int(j) in gout_idx:
                        two.append((w1 * w2, int(m), int(j), w1, w2))
            two.sort(reverse=True)
            seen = set()
            for sc, m, j, w1, w2 in two:
                if (m, j) in seen:
                    continue
                seen.add((m, j))
                chains.append({"hops": 2,
                               "path": [idx2sym[i], idx2sym[m], idx2sym[j]],
                               "edges": [{"a": idx2sym[i], "b": idx2sym[m], "string_score": int(round(w1 * 1000))},
                                         {"a": idx2sym[m], "b": idx2sym[j], "string_score": int(round(w2 * 1000))}],
                               "gout_anchor": idx2sym[j],
                               "gout_anchor_min_p": gs["gout_genes"][sym_of_gout[j]]["min_p"]})
                if len(chains) >= 5:
                    break
        sym = idx2sym[i]
        cvd_rec = gs["cvd_genes"].get(sym)
        rows.append({
            "rank": len(rows) + 1,
            "gene": sym,
            "rwr_score": float(rwr[i]),
            "weighted_degree": float(deg_w[i]),
            "degree_rank_in_network": deg_rank[i] + 1,
            "rwr_rank_in_universe": rwr_rank[i] + 1,
            "already_cvd_gwas_gene": sym in gs["cvd_genes"],
            "cvd_gwas_min_p": cvd_rec["min_p"] if cvd_rec else None,
            "cvd_gwas_first_year": cvd_rec["first_year"] if cvd_rec else None,
            "n_direct_gout_neighbors": len(direct),
            "evidence_chains": chains[:5],
        })

    # 已知重叠（不是预测，是事实表）
    overlap = sorted(set(ms["gout"]) & set(ms["cvd"]))
    overlap_rows = [{"gene": g,
                     "gout_min_p": gs["gout_genes"][g]["min_p"],
                     "cvd_min_p": gs["cvd_genes"][g]["min_p"],
                     "cvd_traits": gs["cvd_genes"][g]["traits"]} for g in overlap]

    n_already = sum(1 for r in rows if r["already_cvd_gwas_gene"])
    out = {"top_n": TOPN,
           "note": "候选表按痛风种子的 RWR 分数排序，已剔除全部痛风种子基因（它们是模型输入，留在表里是泄漏）。"
                   "already_cvd_gwas_gene=true 表示该基因已经有心血管 GWAS 关联——那是对排序的支持性检查，不是新发现。",
           "n_in_top50_already_cvd_gwas": n_already,
           "top50": rows,
           "known_gout_cvd_overlap": {"n": len(overlap_rows), "note": "GWAS Catalog 里同时带痛风与心血管关联的基因，属已知事实不属本轮预测",
                                      "genes": overlap_rows}}
    json.dump(out, open(os.path.join(RES, "top50_candidates.json"), "w"), ensure_ascii=False, indent=1)

    with open(os.path.join(RES, "top50_candidates.tsv"), "w", encoding="utf-8") as fh:
        fh.write("rank\tgene\trwr_score\tweighted_degree\tdegree_rank\talready_cvd_gwas\tn_direct_gout_neighbors\tstrongest_evidence_chain\n")
        for r in rows:
            ch = r["evidence_chains"][0] if r["evidence_chains"] else None
            chain_txt = " -> ".join(ch["path"]) + " (" + ",".join(str(e["string_score"]) for e in ch["edges"]) + ")" if ch else "none"
            fh.write(f"{r['rank']}\t{r['gene']}\t{r['rwr_score']:.6g}\t{r['weighted_degree']:.2f}\t"
                     f"{r['degree_rank_in_network']}\t{r['already_cvd_gwas_gene']}\t{r['n_direct_gout_neighbors']}\t{chain_txt}\n")
    print(json.dumps({"top50_written": len(rows), "already_cvd_gwas_in_top50": n_already,
                      "known_overlap": len(overlap_rows),
                      "first10": [r["gene"] for r in rows[:10]]}, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()