#!/usr/bin/env python3
"""按比对长度 1:1 配对候选基因与背景基因，并固定一个随机顺序（预注册的执行顺序）。

长度配对的理由：RELAX 检出 K≠1 的统计功效随比对长度上升。若候选集比背景集系统性更长，
「候选集显著松弛的基因更多」可以纯粹由功效差异产生，与洞穴适应无关。
"""
import json, random, sys
from pathlib import Path


def aln_len(p):
    n = 0
    for line in open(p):
        if not line.startswith(">"): n += len(line.strip())
    return n // 3 // 6 if False else None


def read_len(p):
    seq, started = 0, False
    for line in open(p):
        if line.startswith(">"):
            if started: break
            started = True
        elif started: seq += len(line.strip())
    return seq // 3


def main(alndir, cand_list, bg_list, out, seed="20260908"):
    alndir = Path(alndir)
    def avail(listfile):
        got = {}
        for g in open(listfile).read().split():
            safe = "".join(c if c.isalnum() or c in "._-" else "_" for c in g)
            f = alndir / f"{safe}.fna"
            if f.is_file(): got[g] = (safe, read_len(f))
        return got
    cand, bg = avail(cand_list), avail(bg_list)
    used, pairs = set(), []
    for g, (safe, L) in sorted(cand.items(), key=lambda kv: -kv[1][1]):
        best = min((b for b in bg if b not in used), key=lambda b: abs(bg[b][1] - L), default=None)
        if best is None: break
        used.add(best)
        pairs.append({"candidate": g, "candidate_file": safe, "candidate_codons": L,
                      "background": best, "background_file": bg[best][0],
                      "background_codons": bg[best][1]})
    rng = random.Random(int(seed))
    order = list(range(len(pairs)))
    rng.shuffle(order)
    Path(out).write_text(json.dumps(
        {"seed": seed, "n_pairs": len(pairs),
         "note": "execution_order 在任何 RELAX 生产运行之前写死；若算力预算用尽提前停止，"
                 "截断发生在与结果无关的随机顺序上，因此不引入选择偏倚。",
         "execution_order": order, "pairs": pairs}, indent=1))
    import statistics
    print("pairs", len(pairs),
          "median codons cand", statistics.median(p["candidate_codons"] for p in pairs),
          "bg", statistics.median(p["background_codons"] for p in pairs))


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