#!/usr/bin/env python3
"""按预注册顺序成对跑 RELAX：每对(候选,背景) × 每个前景分支(洞穴/地表) = 4 次。

预算用尽即停；顺序在跑前写死且与结果无关，因此截断不引入选择偏倚。
"""
import json, os, subprocess, sys, time, tempfile
from pathlib import Path


def run_one(hyphy, aln, tree, outjson, log):
    t0 = time.time()
    p = subprocess.run([hyphy, "CPU=2", "relax", "--alignment", str(aln), "--tree", str(tree),
                        "--test", "Test", "--mode", "Minimal", "--output", str(outjson)],
                       capture_output=True, text=True)
    dt = time.time() - t0
    if p.returncode != 0 or not Path(outjson).is_file():
        log.write(f"FAIL {aln.name} {tree.name} rc={p.returncode} {p.stderr[-300:]}\n")
        return None, dt
    d = json.loads(Path(outjson).read_text())
    tr = d.get("test results", {})
    fits = d.get("fits", {})
    alt = fits.get("RELAX alternative", {})
    return {"K": tr.get("relaxation or intensification parameter"),
            "p_value": tr.get("p-value"), "LRT": tr.get("LRT"),
            "logL_alt": alt.get("Log Likelihood"),
            "n_codons": d.get("input", {}).get("number of sites"),
            "runtime_sec": round(dt, 1)}, dt


def main(pkg, alndir, budget_min="115"):
    pkg, alndir = Path(pkg), Path(alndir)
    res_dir = pkg / "work" / "results"
    pairs = json.loads((res_dir / "pairs.json").read_text())
    hyphy = os.environ["HYPHY_BIN"]
    trees = {"cave": res_dir / "tree_test_cave_rosa.nwk",
             "surface": res_dir / "tree_test_surface_dalaica.nwk"}
    out = open(res_dir / "relax_raw.jsonl", "a")
    log = open(pkg / "work" / "logs" / "relax_run.log", "a")
    done = set()
    f = res_dir / "relax_raw.jsonl"
    if f.is_file():
        for line in open(f):
            try: r = json.loads(line); done.add((r["gene"], r["set"], r["branch"]))
            except Exception: pass
    budget = float(budget_min) * 60
    t_start = time.time()
    log.write(f"=== start {time.strftime('%FT%T')} budget={budget_min}min pairs={pairs['n_pairs']} ===\n")
    log.flush()
    stopped = None
    with tempfile.TemporaryDirectory() as wd:
        for rank, idx in enumerate(pairs["execution_order"]):
            pr = pairs["pairs"][idx]
            if time.time() - t_start > budget:
                stopped = {"reason": "cpu_budget_exhausted", "pairs_attempted": rank}
                break
            for setname, gene, gfile, ncod in (
                    ("candidate", pr["candidate"], pr["candidate_file"], pr["candidate_codons"]),
                    ("background", pr["background"], pr["background_file"], pr["background_codons"])):
                for branch, tree in trees.items():
                    if (gene, setname, branch) in done: continue
                    r, dt = run_one(hyphy, alndir / f"{gfile}.fna", tree,
                                    Path(wd) / "o.json", log)
                    rec = {"pair_rank": rank, "pair_index": idx, "gene": gene, "set": setname,
                           "branch": branch, "codons_prefilter": ncod,
                           "ok": r is not None, **(r or {})}
                    out.write(json.dumps(rec) + "\n"); out.flush()
                    log.write(f"{rank:3d} {setname:10s} {branch:7s} {gene:14s} "
                              f"K={r['K'] if r else 'NA'} t={dt:.0f}s\n")
                    log.flush()
    summary = {"finished_at": time.strftime("%FT%T"), "elapsed_min": round((time.time()-t_start)/60, 1),
               "stopped": stopped or {"reason": "all_pairs_completed", "pairs_attempted": pairs["n_pairs"]}}
    (res_dir / "run_stop_record.json").write_text(json.dumps(summary, indent=1))
    log.write(json.dumps(summary) + "\n"); log.close(); out.close()
    print(json.dumps(summary))


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