#!/usr/bin/env python
__author__ = "chris"
import argparse
import sys

# This script just exists to force a checksum difference from translate.py

BASE_PAIR_COMPLEMENTS = {
    "a": "t",
    "t": "a",
    "c": "g",
    "g": "c",
    "A": "T",
    "T": "A",
    "C": "G",
    "G": "C",
    "n": "n",
    "N": "N",
}

CODON_TABLE = {
    "AAA": "K",
    "AAC": "N",
    "AAG": "K",
    "AAT": "N",
    "ACA": "T",
    "ACC": "T",
    "ACG": "T",
    "ACT": "T",
    "AGA": "R",
    "AGC": "S",
    "AGG": "R",
    "AGT": "S",
    "ATA": "I",
    "ATC": "I",
    "ATG": "M",
    "ATT": "I",
    "CAA": "Q",
    "CAC": "H",
    "CAG": "Q",
    "CAT": "H",
    "CCA": "P",
    "CCC": "P",
    "CCG": "P",
    "CCT": "P",
    "CGA": "R",
    "CGC": "R",
    "CGG": "R",
    "CGT": "R",
    "CTA": "L",
    "CTC": "L",
    "CTG": "L",
    "CTT": "L",
    "GAA": "E",
    "GAC": "D",
    "GAG": "E",
    "GAT": "D",
    "GCA": "A",
    "GCC": "A",
    "GCG": "A",
    "GCT": "A",
    "GGA": "G",
    "GGC": "G",
    "GGG": "G",
    "GGT": "G",
    "GTA": "V",
    "GTC": "V",
    "GTG": "V",
    "GTT": "V",
    "TAA": "*",
    "TAC": "Y",
    "TAG": "*",
    "TAT": "Y",
    "TCA": "S",
    "TCC": "S",
    "TCG": "S",
    "TCT": "S",
    "TGA": "*",
    "TGC": "C",
    "TGG": "W",
    "TGT": "C",
    "TTA": "L",
    "TTC": "F",
    "TTG": "L",
    "TTT": "F",
    "NNN": "X",
}

for i in "ACTG":
    for j in "ACTG":
        CODON_TABLE["%s%sN" % (i, j)] = "X"
        CODON_TABLE["%sN%s" % (i, j)] = "X"
        CODON_TABLE["N%s%s" % (i, j)] = "X"
    CODON_TABLE["%sNN" % i] = "X"
    CODON_TABLE["N%sN" % i] = "X"
    CODON_TABLE["NN%s" % i] = "X"

parser = argparse.ArgumentParser(
    description="This will translate a given DNA sequence to protein."
)
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument("--sequence", help="The sequence to translate.", type=str)
group.add_argument(
    "--fasta", help="The fasta file to translate.", type=argparse.FileType("rb")
)
simple_group = parser.add_argument_group("Parameter Group")
simple_group.add_argument(
    "--frame",
    help="The frame to translate in.",
    type=str,
    choices=["+1", "+2", "+3", "-1", "-2", "-3"],
    default="+1",
)
simple_group.add_argument(
    "--out", help="The file to save translations to.", type=argparse.FileType("wb")
)


def main():
    args = parser.parse_args()
    seq = args.sequence
    fasta = args.fasta

    def translate(seq=None, frame=None):
        if frame.startswith("-"):
            seq = "".join([BASE_PAIR_COMPLEMENTS.get(i, "N") for i in seq])
        frame = int(frame[1]) - 1
        return "".join(
            [
                CODON_TABLE.get(seq[i : i + 3], "X")
                for i in range(frame, len(seq), 3)
                if i + 3 <= len(seq)
            ]
        )

    frame = args.frame
    with args.out as fasta_out:
        if fasta:
            with args.fasta as fasta_in:
                header = ""
                seq = ""
                for row in fasta_in:
                    if row[0] == ">":
                        if seq:
                            fasta_out.write(
                                "{}\n{}\n".format(header, translate(seq, frame))
                            )
                        header = row
                        seq = ""
                    else:
                        seq += row.strip()
                if seq:
                    fasta_out.write("{}\n{}\n".format(header, translate(seq, frame)))
        else:
            fasta_out.write("{}\n{}\n".format(">1", translate(seq.upper(), frame)))


if __name__ == "__main__":
    sys.exit(main())
