#!/usr/bin/env python3
"""导出需求方要的「松弛/强化选择基因列表 + 显著性矩阵」。"""
import json, sys
from pathlib import Path


def main(pkg):
    pkg = Path(pkg); res = pkg / "work" / "results"
    rows = [json.loads(l) for l in open(res / "relax_raw.jsonl") if l.strip()]
    idx = {(r["gene"], r["branch"]): r for r in rows if r.get("ok")}
    meta = {r["gene"]: r for r in rows}
    genes = sorted({r["gene"] for r in rows if r.get("ok")})
    hdr = ["gene", "gene_set", "codons_analyzed",
           "K_cave_Trosa", "p_cave_Trosa", "call_cave",
           "K_surface_Tdalaica", "p_surface_Tdalaica", "call_surface"]
    def call(r):
        if r is None: return "NA"
        if r["p_value"] >= 0.05: return "ns"
        return "relaxed" if r["K"] < 1 else "intensified"
    n = 0
    with open(res / "relax_matrix.tsv", "w") as fh:
        fh.write("\t".join(hdr) + "\n")
        for g in genes:
            c, s = idx.get((g, "cave")), idx.get((g, "surface"))
            if c is None and s is None: continue
            n += 1
            ref = c or s
            fh.write("\t".join([
                g, ref["set"], str(ref.get("n_codons", "")),
                f"{c['K']:.5f}" if c else "NA", f"{c['p_value']:.5g}" if c else "NA", call(c),
                f"{s['K']:.5f}" if s else "NA", f"{s['p_value']:.5g}" if s else "NA", call(s),
            ]) + "\n")
    print("matrix rows:", n)


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