#!/usr/bin/env python3
"""构建 STRING v12.0 人类高置信互作网络，并把 GWAS 基因符号映射进网络。

预注册参数（跑分析前写死）：combined_score >= 700（STRING 官方 "high confidence" 档），
无向、去重、权重 = combined_score / 1000。符号映射只做精确匹配 STRING preferred_name，
不做模糊匹配、不做别名猜测，映射失败的基因如实计数并丢弃。
"""
import gzip, json, os
import numpy as np
import scipy.sparse as sp

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "..", "data")
RES = os.path.join(HERE, "..", "results")
SCORE_MIN = 700


def main():
    ensp2sym, sym2ensp = {}, {}
    with gzip.open(os.path.join(DATA, "string_info.txt.gz"), "rt", encoding="utf-8") as fh:
        next(fh)
        for line in fh:
            pid, name = line.split("\t")[0], line.split("\t")[1]
            ensp2sym[pid] = name
            sym2ensp.setdefault(name, pid)

    edges, kept, total = [], 0, 0
    with gzip.open(os.path.join(DATA, "string_links.txt.gz"), "rt", encoding="utf-8") as fh:
        next(fh)
        for line in fh:
            a, b, s = line.split()
            total += 1
            s = int(s)
            if s < SCORE_MIN:
                continue
            kept += 1
            edges.append((a, b, s))

    nodes = sorted({a for a, b, s in edges} | {b for a, b, s in edges})
    idx = {n: i for i, n in enumerate(nodes)}
    n = len(nodes)
    rows = np.fromiter((idx[a] for a, b, s in edges), dtype=np.int32, count=len(edges))
    cols = np.fromiter((idx[b] for a, b, s in edges), dtype=np.int32, count=len(edges))
    vals = np.fromiter((s / 1000.0 for a, b, s in edges), dtype=np.float64, count=len(edges))
    A = sp.coo_matrix((vals, (rows, cols)), shape=(n, n)).tocsr()
    A = A.maximum(A.T)                       # STRING 文件本身对称，取 max 去重保险
    A.setdiag(0); A.eliminate_zeros()

    symbols = [ensp2sym[p] for p in nodes]
    sp.save_npz(os.path.join(RES, "string_adj.npz"), A)
    deg_w = np.asarray(A.sum(axis=1)).ravel()
    deg_c = np.diff(A.indptr)

    gs = json.load(open(os.path.join(RES, "gene_sets.json"), encoding="utf-8"))
    sym_index = {}
    for i, s in enumerate(symbols):
        sym_index.setdefault(s, i)

    def map_set(genes):
        hit, miss = {}, []
        for g in genes:
            if g in sym_index:
                hit[g] = sym_index[g]
            else:
                miss.append(g)
        return hit, sorted(miss)

    gout_hit, gout_miss = map_set(gs["gout_genes"])
    cvd_hit, cvd_miss = map_set(gs["cvd_genes"])

    meta = {
        "string_version": "v12.0 (9606)",
        "score_min": SCORE_MIN,
        "edges_total_in_file": total,
        "edges_kept_directed_rows": kept,
        "nodes": n,
        "undirected_edges": int(A.nnz // 2),
        "median_weighted_degree": float(np.median(deg_w)),
        "median_degree": int(np.median(deg_c)),
        "gout_mapped": len(gout_hit), "gout_unmapped": len(gout_miss),
        "cvd_mapped": len(cvd_hit), "cvd_unmapped": len(cvd_miss),
        "overlap_mapped": len(set(gout_hit) & set(cvd_hit)),
        "gout_unmapped_symbols": gout_miss[:50],
        "cvd_unmapped_symbols": cvd_miss[:50],
    }
    json.dump({"symbols": symbols, "nodes_ensp": nodes}, open(os.path.join(RES, "network_nodes.json"), "w"), ensure_ascii=False)
    json.dump({"gout": gout_hit, "cvd": cvd_hit, "meta": meta},
              open(os.path.join(RES, "mapped_sets.json"), "w"), ensure_ascii=False, indent=1)
    print(json.dumps(meta, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()