#!/usr/bin/env python3
"""主分析：按预注册判据汇总 RELAX 结果。计数式判据一律并列连续量。"""
import json, sys, statistics
from collections import defaultdict
from pathlib import Path
from scipy.stats import fisher_exact, mannwhitneyu


def load(p):
    recs = [json.loads(l) for l in open(p) if l.strip()]
    return [r for r in recs if r.get("ok") and r.get("K") is not None]


def main(pkg):
    pkg = Path(pkg); res = pkg / "work" / "results"
    recs = load(res / "relax_raw.jsonl")
    idx = {(r["gene"], r["set"], r["branch"]): r for r in recs}
    pairs = json.loads((res / "pairs.json").read_text())["pairs"]
    # 只用四格全部跑完的成对基因，保证 2x2 严格配对
    complete = []
    for pr in pairs:
        keys = [(pr["candidate"], "candidate", "cave"), (pr["candidate"], "candidate", "surface"),
                (pr["background"], "background", "cave"), (pr["background"], "background", "surface")]
        if all(k in idx for k in keys):
            complete.append(pr)
    out = {"n_pairs_complete": len(complete), "n_pairs_planned": len(pairs)}

    def cell(setname, branch):
        return [idx[(p["candidate" if setname == "candidate" else "background"], setname, branch)]
                for p in complete]

    groups = {(s, b): cell(s, b) for s in ("candidate", "background") for b in ("cave", "surface")}
    summ = {}
    for (s, b), rs in groups.items():
        Ks = [r["K"] for r in rs]
        rel = [r for r in rs if r["K"] < 1 and r["p_value"] < 0.05]
        inten = [r for r in rs if r["K"] > 1 and r["p_value"] < 0.05]
        summ[f"{s}|{b}"] = {
            "n": len(rs), "n_relaxed": len(rel), "n_intensified": len(inten),
            "frac_relaxed": round(len(rel) / len(rs), 4) if rs else None,
            "median_K": round(statistics.median(Ks), 4) if Ks else None,
            "mean_K": round(statistics.mean(Ks), 4) if Ks else None,
            "relaxed_genes": sorted(r["gene"] for r in rel),
            "intensified_genes": sorted(r["gene"] for r in inten)}
    out["groups"] = summ

    tests = {}
    for b in ("cave", "surface"):
        c, g = groups[("candidate", b)], groups[("background", b)]
        a1 = sum(1 for r in c if r["K"] < 1 and r["p_value"] < 0.05); a0 = len(c) - a1
        b1 = sum(1 for r in g if r["K"] < 1 and r["p_value"] < 0.05); b0 = len(g) - b1
        orr, pf = fisher_exact([[a1, a0], [b1, b0]], alternative="greater")
        u, pm = mannwhitneyu([r["K"] for r in c], [r["K"] for r in g], alternative="less")
        tests[b] = {"table_[[cand_relaxed,cand_not],[bg_relaxed,bg_not]]": [[a1, a0], [b1, b0]],
                    "fisher_odds_ratio": None if orr != orr else round(float(orr), 4),
                    "fisher_p_one_sided_greater": round(float(pf), 6),
                    "median_K_candidate": summ[f"candidate|{b}"]["median_K"],
                    "median_K_background": summ[f"background|{b}"]["median_K"],
                    "mannwhitney_U": float(u),
                    "mannwhitney_p_one_sided_candidate_lower": round(float(pm), 6)}
    out["tests"] = tests
    # 配对内的 K 差（洞穴分支 vs 地表分支，同一基因）——尺子的分支特异性
    for s in ("candidate", "background"):
        cav = [r["K"] for r in groups[(s, "cave")]]
        sur = [r["K"] for r in groups[(s, "surface")]]
        u, p = mannwhitneyu(cav, sur, alternative="two-sided")
        out.setdefault("branch_contrast", {})[s] = {
            "median_K_cave": round(statistics.median(cav), 4),
            "median_K_surface": round(statistics.median(sur), 4),
            "mannwhitney_p_two_sided": round(float(p), 6)}
    verdict = {}
    tc = tests["cave"]; ts = tests["surface"]
    verdict["H1"] = ("成立" if tc["fisher_p_one_sided_greater"] < 0.05 and (tc["fisher_odds_ratio"] or 0) > 1
                     else "被推翻")
    verdict["H2"] = ("成立" if tc["mannwhitney_p_one_sided_candidate_lower"] < 0.05
                     and tc["median_K_candidate"] < tc["median_K_background"] else "被推翻")
    verdict["H3_negative_control"] = ("成立（地表分支未出现同等富集）"
                                      if not (ts["fisher_p_one_sided_greater"] < 0.05) else
                                      "被推翻（地表分支也富集，H1 的因果解读不成立）")
    out["prereg_verdict"] = verdict
    (res / "analysis.json").write_text(json.dumps(out, ensure_ascii=False, indent=1))
    print(json.dumps({k: out[k] for k in ("n_pairs_complete", "tests", "prereg_verdict", "branch_contrast")},
                     ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main(*sys.argv[1:])