#!/usr/bin/env python3
"""对每个直系同源组做「蛋白比对 → 回译密码子比对」，并做轻度列过滤。"""
import json, re, subprocess, sys, tempfile, os
from pathlib import Path

SPECIES = ["GCF_024868665.1", "GCF_015846415.1", "GCA_008369825.1",
           "GCF_947034865.1", "GCF_027580225.2", "GCF_052040795.1"]
LABEL = {"GCF_024868665.1": "Triplophysa_rosa", "GCF_015846415.1": "Triplophysa_dalaica",
         "GCA_008369825.1": "Triplophysa_tibetana", "GCF_947034865.1": "Barbatula_barbatula",
         "GCF_027580225.2": "Misgurnus_anguillicaudatus", "GCF_052040795.1": "Danio_rerio"}
MAX_GAP_FRAC = 0.5      # 一列里超过一半物种是 gap 就删掉这个密码子列
MIN_CODONS = 100        # 过滤后不足 100 个密码子的基因不进入检验


def read_fasta(path):
    d, name, buf = {}, None, []
    with open(path) as fh:
        for line in fh:
            if line.startswith(">"):
                if name: d[name] = "".join(buf)
                name, buf = line[1:].split()[0], []
            else: buf.append(line.strip())
    if name: d[name] = "".join(buf)
    return d


def load_all(S):
    return {acc: read_fasta(Path(S) / "prep" / f"{acc}.cds.fna") for acc in SPECIES}


def build(cds_by_sp, row, mafft, workdir):
    nt = {}
    for acc in SPECIES:
        s = cds_by_sp[acc].get(row[acc])
        if not s: return None, "missing_cds"
        nt[acc] = s
    tmp_in = os.path.join(workdir, "in.faa")
    with open(tmp_in, "w") as fh:
        for acc in SPECIES:
            aa = translate(nt[acc])
            fh.write(f">{LABEL[acc]}\n{aa}\n")
    p = subprocess.run([mafft, "--auto", "--quiet", "--thread", "1", tmp_in],
                       capture_output=True, text=True)
    if p.returncode != 0: return None, "mafft_failed"
    tmp_out = os.path.join(workdir, "out.faa")
    Path(tmp_out).write_text(p.stdout)
    aln_aa = read_fasta(tmp_out)
    # 回译
    codon = {}
    for acc in SPECIES:
        lab, s, out, k = LABEL[acc], nt[acc], [], 0
        for ch in aln_aa[lab]:
            if ch == "-": out.append("---")
            else:
                out.append(s[k * 3:k * 3 + 3]); k += 1
        codon[lab] = out
    ncol = len(next(iter(codon.values())))
    keep = [i for i in range(ncol)
            if sum(1 for lab in codon if codon[lab][i] == "---") / len(codon) <= MAX_GAP_FRAC]
    if len(keep) < MIN_CODONS: return None, f"too_short:{len(keep)}"
    return {lab: "".join(codon[lab][i] for i in keep) for lab in codon}, f"ok:{len(keep)}"


TABLE = {}
_B, _AA = "TCAG", "FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG"
for i, b1 in enumerate(_B):
    for j, b2 in enumerate(_B):
        for k2, b3 in enumerate(_B):
            TABLE[b1 + b2 + b3] = _AA[i * 16 + j * 4 + k2]


def translate(nt):
    return "".join(TABLE.get(nt[i:i + 3], "X") for i in range(0, len(nt), 3))


def main(S, ortho_json, outdir, mafft, id_list=None):
    rows = json.loads(Path(ortho_json).read_text())["orthologs"]
    if id_list:
        want = set(Path(id_list).read_text().split())
        rows = [r for r in rows if r["gene_symbol"] in want]
    cds = load_all(S)
    Path(outdir).mkdir(parents=True, exist_ok=True)
    log = []
    with tempfile.TemporaryDirectory() as wd:
        for r in rows:
            aln, status = build(cds, r, mafft, wd)
            safe = re.sub(r"[^A-Za-z0-9._-]", "_", r["gene_symbol"])
            if aln:
                with open(Path(outdir) / f"{safe}.fna", "w") as fh:
                    for lab, s in aln.items(): fh.write(f">{lab}\n{s}\n")
            log.append({"gene": r["gene_symbol"], "file": safe, "status": status})
    Path(outdir).joinpath("_align_log.json").write_text(json.dumps(log))
    ok = sum(1 for x in log if x["status"].startswith("ok"))
    print(f"aligned {ok}/{len(log)}")


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