#!/usr/bin/env python3
"""AML 跨队列一致性初检：TCGA-LAML 全基因组发现 → BeatAML 真实临床标签复制。

判定规则写在 ../claims.json，跑之前定死。零模型是标签置换 1000 次。
"""
import gzip, json, urllib.request, hashlib
import numpy as np
from scipy.stats import mannwhitneyu

TOP_N = 200
N_PERM = 1000
SEED = 20260826


def bh(p):
    p = np.asarray(p, float); n = len(p); o = np.argsort(p); adj = np.empty(n); prev = 1.0
    for r, i in enumerate(o[::-1]):
        prev = min(prev, p[i] * n / (n - r)); adj[i] = prev
    return adj


def load_tcga():
    lab = json.load(open("cohort_labels.json"))
    res, sen = set(lab["tcga_resistant"]), set(lab["tcga_sensitive"])
    genes, mat, cols = [], [], None
    with gzip.open("tcga_laml.tsv.gz", "rt") as fh:
        cols = fh.readline().rstrip("\n").split("\t")[1:]
        for line in fh:
            p = line.rstrip("\n").split("\t")
            genes.append(p[0]); mat.append(p[1:])
    X = np.array(mat, dtype=float)
    pat = [c[:12] for c in cols]
    ri = [i for i, p in enumerate(pat) if p in res]
    si = [i for i, p in enumerate(pat) if p in sen]
    return genes, X, np.array(ri), np.array(si)


def _post(url, body, timeout=600):
    req = urllib.request.Request(url, data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json",
                                          "Accept": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as fh:
        return json.load(fh)


def beataml_expression(genes):
    """只取候选基因，避免全基因组走 API。接口要 Entrez ID，先做符号→ID 映射。"""
    info = _post("https://www.cbioportal.org/api/genes/fetch?geneIdType=HUGO_GENE_SYMBOL",
                 genes, timeout=180)
    sym = {g["entrezGeneId"]: g["hugoGeneSymbol"] for g in info}
    print(f"[B] 候选 {len(genes)} 个基因里，cBioPortal 认得 {len(sym)} 个")
    rows = _post("https://www.cbioportal.org/api/molecular-profiles/"
                 "aml_ohsu_2018_rna_seq_mrna/molecular-data/fetch?projection=SUMMARY",
                 {"entrezGeneIds": sorted(sym), "sampleListId": "aml_ohsu_2018_all"})
    out = {}
    for r in rows:
        if r.get("value") is None:
            continue
        out.setdefault(sym[r["entrezGeneId"]], {})[r["sampleId"]] = float(r["value"])
    return out


def main():
    rng = np.random.default_rng(SEED)
    genes, X, ri, si = load_tcga()
    print(f"[A] TCGA-LAML 基因 {len(genes)}｜耐药 {len(ri)} 例｜敏感 {len(si)} 例")

    def screen(ri_, si_):
        ps, ds = [], []
        for g in range(X.shape[0]):
            a, b = X[g, ri_], X[g, si_]
            if a.std() + b.std() == 0:
                ps.append(1.0); ds.append(0.0); continue
            u, p = mannwhitneyu(a, b, alternative="two-sided")
            ps.append(p); ds.append(np.median(a) - np.median(b))
        return np.array(ps), np.array(ds)

    p_a, d_a = screen(ri, si)
    q_a = bh(p_a)
    top = np.argsort(q_a)[:TOP_N]
    cand = [genes[i] for i in top]
    print(f"[A] top{TOP_N} 候选，最小 FDR {q_a[top].min():.2e}")

    expr_b = beataml_expression(cand)
    print(f"[B] BeatAML 取回 {len(expr_b)} 个基因的表达")
    labels = json.load(open("cohort_labels.json"))["beataml_induction_response"]
    ref = [s for s, v in labels.items() if v == "Refractory"]
    cr = [s for s, v in labels.items() if v.startswith("Complete Response")]
    print(f"[B] 难治 {len(ref)} 例｜完全缓解 {len(cr)} 例")

    def replicate(perm_labels=None):
        ok = 0; detail = []
        rl, cl = (ref, cr) if perm_labels is None else perm_labels
        ps, dirs, gs = [], [], []
        for g in cand:
            m = expr_b.get(g)
            if not m: continue
            a = np.array([m[s] for s in rl if s in m]); b = np.array([m[s] for s in cl if s in m])
            if len(a) < 10 or len(b) < 10 or a.std() + b.std() == 0: continue
            u, p = mannwhitneyu(a, b, alternative="two-sided")
            ps.append(p); dirs.append(np.median(a) - np.median(b)); gs.append(g)
        if not ps: return 0, []
        qs = bh(ps)
        for g, q, dd in zip(gs, qs, dirs):
            da = d_a[genes.index(g)]
            if q < 0.10 and np.sign(dd) == np.sign(da) and dd != 0:
                ok += 1; detail.append({"gene": g, "tcga_delta": float(da),
                                        "beataml_delta": float(dd), "beataml_fdr": float(q)})
        return ok, detail

    n_ok, detail = replicate()
    print(f"[复制] 方向一致且 FDR<0.10 的基因: {n_ok}")

    alls = ref + cr
    null = []
    for k in range(N_PERM):
        sh = list(alls); rng.shuffle(sh)
        null.append(replicate((sh[:len(ref)], sh[len(ref):]))[0])
    null = np.array(null)
    p95 = float(np.percentile(null, 95))
    h1 = bool(n_ok > p95); h2 = bool(n_ok >= 5)
    print(f"[零模型] {N_PERM} 次置换：均值 {null.mean():.2f}｜95分位 {p95:.1f}｜最大 {null.max()}")
    print(f"[H1] 实际 {n_ok} > 零模型95分位 {p95:.1f} → {'✅成立' if h1 else '❌被推翻'}")
    print(f"[H2] 复制成功 >= 5 → {'✅成立' if h2 else '❌被推翻'}")

    json.dump({"seed": SEED, "top_n": TOP_N, "n_perm": N_PERM,
               "tcga_sha256": hashlib.sha256(open("tcga_laml.tsv.gz","rb").read()).hexdigest(),
               "n_tcga_res": int(len(ri)), "n_tcga_sen": int(len(si)),
               "n_beataml_ref": len(ref), "n_beataml_cr": len(cr),
               "n_candidates": len(cand), "n_replicated": int(n_ok),
               "null_mean": float(null.mean()), "null_p95": p95, "null_max": int(null.max()),
               "H1_passes": h1, "H2_passes": h2, "replicated": detail,
               "null_distribution": null.tolist()},
              open("results.json", "w"), ensure_ascii=False, indent=1)


if __name__ == "__main__":
    main()