#!/usr/bin/env python3
"""用背景基因串联比对估计 6 物种拓扑（NJ，蛋白 p-distance）。RELAX 需要固定物种树。"""
import json, sys, itertools
from pathlib import Path

SP = ["Triplophysa_rosa", "Triplophysa_dalaica", "Triplophysa_tibetana",
      "Barbatula_barbatula", "Misgurnus_anguillicaudatus", "Danio_rerio"]


def read_fasta(p):
    d, n, b = {}, None, []
    for line in open(p):
        if line.startswith(">"):
            if n: d[n] = "".join(b)
            n, b = line[1:].strip(), []
        else: b.append(line.strip())
    if n: d[n] = "".join(b)
    return d


def main(alndir, genelist, out):
    want = {g.strip() for g in open(genelist) if g.strip()}
    diff = {p: 0 for p in itertools.combinations(SP, 2)}
    tot = {p: 0 for p in diff}
    nfile = 0
    for f in sorted(Path(alndir).glob("*.fna")):
        if f.stem not in want: continue
        a = read_fasta(f)
        if set(a) != set(SP): continue
        nfile += 1
        for x, y in diff:
            sx, sy = a[x], a[y]
            for i in range(0, len(sx), 3):
                cx, cy = sx[i:i+3], sy[i:i+3]
                if "-" in cx or "-" in cy: continue
                tot[(x, y)] += 1
                if cx != cy: diff[(x, y)] += 1
    D = {p: diff[p] / tot[p] for p in diff}
    # 邻接法（NJ）
    nodes = list(SP)
    dist = {}
    for (x, y), v in D.items(): dist[(x, y)] = dist[(y, x)] = v
    for n in nodes: dist[(n, n)] = 0.0
    while len(nodes) > 3:
        N = len(nodes)
        r = {i: sum(dist[(i, j)] for j in nodes) for i in nodes}
        best, bp = None, None
        for i, j in itertools.combinations(nodes, 2):
            q = (N - 2) * dist[(i, j)] - r[i] - r[j]
            if best is None or q < best: best, bp = q, (i, j)
        i, j = bp
        new = f"({i},{j})"
        for k in nodes:
            if k in bp: continue
            dist[(new, k)] = dist[(k, new)] = (dist[(i, k)] + dist[(j, k)] - dist[(i, j)]) / 2
        dist[(new, new)] = 0.0
        nodes = [n for n in nodes if n not in bp] + [new]
    tree = "(" + ",".join(nodes) + ");"
    res = {"n_genes_used": nfile,
           "codon_p_distance": {f"{x}|{y}": round(v, 5) for (x, y), v in sorted(D.items())},
           "nj_tree": tree}
    Path(out).write_text(json.dumps(res, indent=1))
    print(json.dumps(res, indent=1))


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