#!/usr/bin/env python3
"""从 NCBI cds_from_genomic.fna 提取「每个基因一条最长 CDS」，并翻译成蛋白。

CDS 文件是唯一真值来源：蛋白由 CDS 翻译而来，保证密码子比对时 CDS 与蛋白严格对应，
不依赖 protein.faa 与 CDS 的外部一致性。
"""
import gzip, json, re, sys
from pathlib import Path

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


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


def read_fasta(path):
    name, buf = None, []
    with open(path) as fh:
        for line in fh:
            if line.startswith(">"):
                if name is not None:
                    yield name, "".join(buf)
                name, buf = line[1:].rstrip("\n"), []
            else:
                buf.append(line.strip())
    if name is not None:
        yield name, "".join(buf)


def parse_tags(header):
    return dict(re.findall(r"\[(\w+)=([^\]]*)\]", header))


def main(acc, cds_path, out_prefix):
    best = {}          # gene_key -> (len, protein_id, gene_symbol, nt)
    n_raw = n_kept = 0
    seen_ids = set()
    for header, seq in read_fasta(cds_path):
        # cds_from_genomic 对同一 CDS 可能重复出现，按 protein_id 去重
        tags = parse_tags(header)
        pid = tags.get("protein_id")
        if not pid:
            continue
        if pid in seen_ids:
            continue
        seen_ids.add(pid)
        n_raw += 1
        if tags.get("pseudo") == "true" or tags.get("partial") == "true":
            continue
        seq = seq.upper()
        if len(seq) < 300 or len(seq) % 3 != 0:
            continue
        if set(seq) - set("ACGT"):
            continue
        aa = translate(seq)
        if aa.endswith("*"):
            aa, seq = aa[:-1], seq[:-3]
        if "*" in aa or "X" in aa or not aa.startswith("M"):
            continue
        gene = tags.get("gene") or pid          # tibetana 无 gene 标签，退化为 per-protein
        cur = best.get(gene)
        if cur is None or len(seq) > cur[0]:
            best[gene] = (len(seq), pid, gene, seq)
    with open(out_prefix + ".cds.fna", "w") as fn, open(out_prefix + ".faa", "w") as fa:
        for gene, (_, pid, sym, nt) in sorted(best.items()):
            n_kept += 1
            fn.write(f">{pid}\n{nt}\n")
            fa.write(f">{pid}\n{translate(nt)}\n")
    meta = {"accession": acc, "cds_records_unique": n_raw, "genes_kept": n_kept,
            "gene_symbols_present": sum(1 for g in best if not g.startswith(("XP_", "KAA", "NP_"))),
            "protein_to_gene": {v[1]: v[2] for v in best.values()}}
    Path(out_prefix + ".meta.json").write_text(json.dumps(meta))
    print(f"{acc}: unique_cds={n_raw} genes_kept={n_kept} symbols={meta['gene_symbols_present']}")


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