#!/usr/bin/env python3
"""从 GWAS Catalog（匿名可直接下载）抽取痛风与心血管疾病的基因集。

规则在跑分析之前写死：
  - 只取全基因组显著关联 P < 5e-8
  - 只取落在基因内部的 SNP（INTERGENIC != 1 且 MAPPED_GENE 不含 " - "），
    基因间 SNP 的基因归属有歧义，一律丢弃
  - 痛风种子 = 疾病本体项 MONDO_0005393 (gout) / EFO_0021525 (renal overload-type gout)
  - 心血管标签 = 预注册的疾病本体项清单（只要疾病，不要测量指标）
"""
import csv, glob, json, os, re, sys
from collections import defaultdict

DATA = os.path.join(os.path.dirname(__file__), "..", "data")
RES = os.path.join(os.path.dirname(__file__), "..", "results")

GOUT_TERMS = {"MONDO_0005393": "gout", "EFO_0021525": "renal overload-type gout"}
CVD_TERMS = {
    "MONDO_0005010": "coronary artery disorder",
    "MONDO_0005068": "myocardial infarction",
    "MONDO_0005252": "heart failure",
    "MONDO_0005009": "congestive heart failure",
    "HP_0002140": "ischemic stroke",
    "MONDO_0005098": "stroke disorder",
    "MONDO_0004995": "cardiovascular disorder",
    "MONDO_0021661": "coronary atherosclerosis",
    "MONDO_0005311": "atherosclerosis",
    "MONDO_0005386": "peripheral arterial disease",
    "EFO_0003913": "angina pectoris",
}
PCUT = 5e-8
TERM_RE = re.compile(r"(MONDO_\d+|EFO_\d+|HP_\d+|OBA_\d+)")


def parse_p(raw):
    try:
        v = float(raw)
    except (TypeError, ValueError):
        return None
    return v if v > 0 else None


def split_genes(mapped_gene):
    """MAPPED_GENE 里 ' - ' 表示基因间(上游-下游)，整条丢弃；',' / 'x' 表示多个映射基因。"""
    if not mapped_gene or mapped_gene.strip() in ("", "NR"):
        return []
    if " - " in mapped_gene:
        return []
    out = []
    for part in re.split(r"[,;]| x ", mapped_gene):
        g = part.strip()
        if g and g != "NR":
            out.append(g)
    return out


def main():
    files = sorted(glob.glob(os.path.join(DATA, "gwas-catalog-download-associations-alt.*.tsv")))
    assert files, "GWAS Catalog 文件不存在"
    gout = defaultdict(lambda: {"n_assoc": 0, "min_p": 1.0, "pubmed": set(), "first_year": 9999, "traits": set()})
    cvd = defaultdict(lambda: {"n_assoc": 0, "min_p": 1.0, "pubmed": set(), "first_year": 9999, "traits": set()})
    stats = {"rows": 0, "rows_pass_p": 0, "gout_rows": 0, "cvd_rows": 0, "files": [os.path.basename(f) for f in files]}

    for path in files:
        with open(path, newline="", encoding="utf-8", errors="replace") as fh:
            rdr = csv.DictReader(fh, delimiter="\t")
            for row in rdr:
                stats["rows"] += 1
                p = parse_p(row.get("P-VALUE"))
                if p is None or p >= PCUT:
                    continue
                stats["rows_pass_p"] += 1
                if (row.get("INTERGENIC") or "").strip() == "1":
                    continue
                terms = set(TERM_RE.findall(row.get("MAPPED_TRAIT_URI") or ""))
                hit_gout = terms & set(GOUT_TERMS)
                hit_cvd = terms & set(CVD_TERMS)
                if not hit_gout and not hit_cvd:
                    continue
                genes = split_genes(row.get("MAPPED_GENE"))
                if not genes:
                    continue
                year = (row.get("DATE") or "")[:4]
                year = int(year) if year.isdigit() else 9999
                pmid = (row.get("PUBMEDID") or "").strip()
                for tag, hits, bucket in (("gout", hit_gout, gout), ("cvd", hit_cvd, cvd)):
                    if not hits:
                        continue
                    stats[f"{tag}_rows"] += 1
                    for g in genes:
                        rec = bucket[g]
                        rec["n_assoc"] += 1
                        rec["min_p"] = min(rec["min_p"], p)
                        rec["first_year"] = min(rec["first_year"], year)
                        if pmid:
                            rec["pubmed"].add(pmid)
                        rec["traits"].update(hits)

    def dump(bucket):
        return {g: {"n_assoc": v["n_assoc"], "min_p": v["min_p"], "first_year": v["first_year"],
                    "n_pubmed": len(v["pubmed"]), "traits": sorted(v["traits"])}
                for g, v in sorted(bucket.items())}

    out = {
        "generated_at_utc": __import__("datetime").datetime.now(__import__("datetime").timezone.utc).isoformat(timespec="seconds"),
        "p_threshold": PCUT,
        "gout_terms": GOUT_TERMS,
        "cvd_terms": CVD_TERMS,
        "stats": stats,
        "gout_genes": dump(gout),
        "cvd_genes": dump(cvd),
    }
    os.makedirs(RES, exist_ok=True)
    with open(os.path.join(RES, "gene_sets.json"), "w", encoding="utf-8") as fh:
        json.dump(out, fh, ensure_ascii=False, indent=1)
    print(json.dumps({"gout_genes": len(gout), "cvd_genes": len(cvd),
                      "overlap": len(set(gout) & set(cvd)), **stats}, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()