#!/usr/bin/env python3
"""取两个队列的分组标签：BeatAML 诱导治疗反应 / TCGA-LAML 生存分位数代理。

只用 cBioPortal 与 UCSC Xena 的公开匿名接口。
"""
import json, urllib.request, urllib.parse, gzip
import numpy as np

API = "https://www.cbioportal.org/api"


def clinical(study, attr, level="PATIENT"):
    """注意 key：SAMPLE 层要用 sampleId，用 patientId 会和表达矩阵对不上（2026-08-26 踩过，
    结果是队列 B 直接空掉、伪造出"复制不了"的假结论）。"""
    url = f"{API}/studies/{study}/clinical-data?clinicalDataType={level}&attributeId={attr}&pageSize=5000"
    key = "sampleId" if level == "SAMPLE" else "patientId"
    with urllib.request.urlopen(url, timeout=180) as fh:
        return {r[key]: r["value"] for r in json.load(fh)}


# --- 队列 A：BeatAML 诱导治疗反应 ---
resp = clinical("aml_ohsu_2018", "INDUCTION_RESPONSE", level="SAMPLE")
from collections import Counter
print("  BeatAML INDUCTION_RESPONSE 取值:", dict(Counter(resp.values()).most_common(8)))

# --- 队列 B：TCGA-LAML 生存 ---
os_m = clinical("laml_tcga", "OS_MONTHS")
os_s = clinical("laml_tcga", "OS_STATUS")
vals = {p: float(v) for p, v in os_m.items() if v not in (None, "", "NA")}
print(f"  TCGA-LAML 有生存月数的病例: {len(vals)}")
arr = np.array(list(vals.values()))
q1, q3 = np.percentile(arr, [25, 75])
print(f"  生存下四分位 {q1:.1f} 月 / 上四分位 {q3:.1f} 月")
resistant = {p for p, v in vals.items() if v <= q1 and os_s.get(p, "").startswith("1")}
sensitive = {p for p, v in vals.items() if v >= q3}
print(f"  TCGA 耐药代理 {len(resistant)} 例 / 敏感代理 {len(sensitive)} 例")

json.dump({"beataml_induction_response": resp,
           "tcga_resistant": sorted(resistant), "tcga_sensitive": sorted(sensitive),
           "tcga_os_q1": float(q1), "tcga_os_q3": float(q3)},
          open("cohort_labels.json", "w"), ensure_ascii=False, indent=1)
print("  → cohort_labels.json")