#!/usr/bin/env python3
"""头条数字的独立复算：不依赖 scipy，也不读 analysis.json，直接从 relax_raw.jsonl 重算。

与 08_analyze.py 的实现路径不同：Fisher 精确检验用 math.comb 直接累加超几何尾概率，
中位数用自写排序取值。两条路径给出的数必须一致，否则闸门判失败。
"""
import json, math, sys
from pathlib import Path


def fisher_greater(a, b, c, d):
    """[[a,b],[c,d]] 的单侧(greater)精确 p，以及样本 odds ratio a*d/(b*c)。"""
    n = a + b + c + d
    r1, c1 = a + b, a + c
    p = 0.0
    lo, hi = max(0, c1 - b), min(r1, c1)
    denom = math.comb(n, c1)
    for k in range(a, hi + 1):
        p += math.comb(r1, k) * math.comb(n - r1, c1 - k) / denom
    orr = float("inf") if b * c == 0 else (a * d) / (b * c)
    return p, orr


def median(xs):
    s = sorted(xs); n = len(s)
    return s[n // 2] if n % 2 else (s[n // 2 - 1] + s[n // 2]) / 2


def main(pkg):
    pkg = Path(pkg); res = pkg / "work" / "results"
    rows = []
    with open(res / "relax_raw.jsonl") as fh:
        for line in fh:
            if line.strip():
                r = json.loads(line)
                if r.get("ok") and r.get("K") is not None:
                    rows.append(r)
    have = {(r["gene"], r["set"], r["branch"]) for r in rows}
    pairs = json.loads((res / "pairs.json").read_text())["pairs"]
    keep_cand, keep_bg = set(), set()
    for p in pairs:
        need = [(p["candidate"], "candidate", "cave"), (p["candidate"], "candidate", "surface"),
                (p["background"], "background", "cave"), (p["background"], "background", "surface")]
        if all(k in have for k in need):
            keep_cand.add(p["candidate"]); keep_bg.add(p["background"])
    def sel(s, b):
        keep = keep_cand if s == "candidate" else keep_bg
        return [r for r in rows if r["set"] == s and r["branch"] == b and r["gene"] in keep]
    out = {"n_pairs_complete": len(keep_cand)}
    for b in ("cave", "surface"):
        C, B = sel("candidate", b), sel("background", b)
        a = sum(1 for r in C if r["K"] < 1 and r["p_value"] < 0.05)
        c = sum(1 for r in B if r["K"] < 1 and r["p_value"] < 0.05)
        pv, orr = fisher_greater(a, len(C) - a, c, len(B) - c)
        out[b] = {"n_candidate": len(C), "n_background": len(B),
                  "n_relaxed_candidate": a, "n_relaxed_background": c,
                  "median_K_candidate": median([r["K"] for r in C]),
                  "median_K_background": median([r["K"] for r in B]),
                  "fisher_p_one_sided_greater": pv, "fisher_odds_ratio": orr}
    (res / "recheck.json").write_text(json.dumps(out, ensure_ascii=False, indent=1))
    print(json.dumps(out, ensure_ascii=False, indent=1))


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