#!/usr/bin/env python3
"""Reframe the NA alignments to be codon-in-frame for hyphaeon meme.

The GISAID NA alignments come with a 2-column leading offset (aln len 1435, L%3==1;
frame 2 is the clean ORF starting MNPNQK... — canonical influenza NA N-terminus).
This is a uniform whole-alignment offset, NOT ragged ends, so we trim columns
(preserving the authors' homology columns) rather than re-align.

Trim = drop OFFSET leading cols, then drop trailing cols so length % 3 == 0.
Verified against the longest ungapped sequence: frame-2 translation has 1 internal stop.
Writes <name>.inframe.fasta next to each input and prints a per-file report.
"""
import sys, glob, os

OFFSET = 2  # leading columns to drop (frame-2 ORF)

BASES="TCAG"; AAS="FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG"
CODON={}
for i,a in enumerate(BASES):
    for j,b in enumerate(BASES):
        for k,c in enumerate(BASES):
            CODON[a+b+c]=AAS[i*16+j*4+k]
def transl(seq):
    return ''.join(CODON.get(seq[i:i+3].upper().replace('-','N'),'X') for i in range(0,len(seq)-2,3))

def read(f):
    seqs={};name=None
    for line in open(f):
        line=line.rstrip("\n")
        if line.startswith('>'): name=line[1:]; seqs[name]=[]
        elif name is not None: seqs[name].append(line)
    return {k:''.join(v) for k,v in seqs.items()}

def main(na_dir):
    for f in sorted(glob.glob(os.path.join(na_dir,"*.fasta"))):
        if f.endswith(".inframe.fasta"): continue
        seqs=read(f)
        L=len(next(iter(seqs.values())))
        trimmed=L-OFFSET; trimmed-=trimmed%3
        out=f.replace(".fasta",".inframe.fasta")
        worst=0
        with open(out,"w") as o:
            for name,s in seqs.items():
                assert len(s)==L, f"ragged aln in {f}: {name} len {len(s)}!={L}"
                sub=s[OFFSET:OFFSET+trimmed]
                worst=max(worst, transl(sub)[:-1].count('*'))
                o.write(f">{name}\n{sub}\n")
        print(f"{os.path.basename(f):24s} {L}->{trimmed} cols ({trimmed//3} codons) "
              f"max_internal_stops_any_seq={worst}  -> {os.path.basename(out)}")

if __name__=="__main__":
    main(sys.argv[1] if len(sys.argv)>1 else ".")