##### READ BELOW FIRST #####
# Requirements:
# PyMOL, Argparse
# Running: pymol <thisfile> -- <DNA HERE> --visual
############################
import argparse
import pymol
from pymol import _cmd

RNA3 = {
 "AAA": "Lys",
 "AAC": "Asn",
 "AAG": "Lys",
 "AAU": "Asn",
 "ACA": "Thr",
 "ACC": "Thr",
 "ACG": "Thr",
 "ACU": "Thr",
 "AGA": "Arg",
 "AGC": "Ser",
 "AGG": "Arg",
 "AGU": "Ser",
 "AUA": "Lle",
 "AUC": "Lle",
 "AUG": "Mel",
 "AUU": "Lle",
 "CAA": "Gln",
 "CAC": "His",
 "CAG": "Gln",
 "CAU": "His",
 "CCA": "Pro",
 "CCC": "Pro",
 "CCG": "Pro",
 "CCU": "Pro",
 "CGA": "Arg",
 "CGC": "Arg",
 "CGG": "Arg",
 "CGU": "Arg",
 "CUA": "Leu",
 "CUC": "Leu",
 "CUG": "Leu",
 "CUU": "Leu",
 "GAA": "Glu",
 "GAC": "Asp",
 "GAG": "Glu",
 "GAU": "Asp",
 "GCA": "Ala",
 "GCC": "Ala",
 "GCG": "Ala",
 "GCU": "Ala",
 "GGA": "Gly",
 "GGC": "Gly",
 "GGG": "Gly",
 "GGU": "Gly",
 "GUA": "Val",
 "GUC": "Val",
 "GUG": "Val",
 "GUU": "Val",
 "UAA": "Stop",
 "UAC": "Tyr",
 "UAG": "Stop",
 "UAU": "Tyr",
 "UCA": "Ser",
 "UCC": "Ser",
 "UCG": "Ser",
 "UCU": "Ser",
 "UGA": "Stop",
 "UGC": "Cys",
 "UGG": "Trp",
 "UGU": "Cys",
 "UUA": "Leu",
 "UUC": "Phe",
 "UUG": "Leu",
 "UUU": "Phe",
}
AMINOACID_CHEM = {
    "Ala": "CH3-CH(NH2)-COOH",
    "Arg": "HN=C(NH2)-NH-(CH2)3-CH(NH2)-COOH",
    "Asn": "H2N-CO-CH2-CH(NH2)-COOH",
    "Asp": "HOOC-CH2-CH(NH2)-COOH",
    "Cys": "HS-CH2-CH(NH2)-COOH",
    "Gln": "H2N-CO-(CH2)2-CH(NH2)-COOH",
    "Glu": "HOOC-(CH2)2-CH(NH2)-COOH",
    "Gly": "NH2-CH2-COOH",
    "His": "NH-CH=N-CH=C-CH2-CH(NH2)-COOH",
    "Lle": "CH3-CH2-CH(CH3)-CH(NH2)-COOH",
    "Leu": "(CH3)2-CH-CH2-CH(NH2)-COOH",
    "Lys": "H2N-(CH2)4-CH(NH2)-COOH",
    "Met": "CH3-S-(CH2)2-CH(NH2)-COOH",
    "Phe": "Ph-CH2-CH(NH2)-COOH",
    "Pro": "NH-(CH2)3-CH-COOH",
    "Ser": "HO-CH2-CH(NH2)-COOH",
    "Thr": "CH3-CH(OH)-CH(NH2)-COOH",
    "Trp": "Ph-NH-CH=C-CH2-CH(NH2)-COOH",
    "Tyr": "HO-Ph-CH2-CH(NH2)-COOH",
    "Val": "(CH3)2-CH-CH(NH2)-COOH"
}
AMINOACID_PEPTIDE = {
    "Ala": "A",
    "Arg": "R",
    "Asn": "N",
    "Asp": "D",
    "Cys": "C",
    "Gln": "Q",
    "Glu": "E",
    "Gly": "G",
    "His": "H",
    "Lle": "I",
    "Leu": "L",
    "Lys": "K",
    "Met": "M",
    "Phe": "F",
    "Pro": "P",
    "Ser": "S",
    "Thr": "T",
    "Trp": "W",
    "Tyr": "Y",
    "Val": "V"
}
RNA_CONVERSION = {
    "A": "A",
    "C": "C",
    "G": "G",
    "T": "U"
}

def DNAtoRNA(DNA):
    RNA = ""
    for c in DNA:
        RNA += RNA_CONVERSION[c]
    return RNA

def RNAtoAminoAcid(RNA):
    amino = ""
    i = 0;
    while (i <= len(RNA)-1):
        next3 = RNA[i:3+i];
        if (len(next3) != 3):
            break;
        if (RNA3[next3] == "Stop"):
            break;
        amino += RNA3[next3];
        i+=3;
    return amino;

def AminoAcidtoChem(AMINO):
    chem = ""
    i = 0;
    while (i <= len(AMINO)-1):
        next3 = AMINO[i:3+i];
        if (len(next3) != 3):
            break;
        if (AMINOACID_CHEM[next3] == "Stop"):
            break;
        chem += AMINOACID_CHEM[next3] + " + ";
        i+=3;
    return chem[:-3];        

def AminoAcidtoPeptide(AMINO):
    peptide = ""
    i = 0;
    while (i <= len(AMINO)-1):
        next3 = AMINO[i:3+i];
        if (len(next3) != 3):
            break;
        if (AMINOACID_PEPTIDE[next3] == "Stop"):
            break;
        peptide += AMINOACID_PEPTIDE[next3]
        i+=3;
    return peptide

def ChemtoAminoAcid(CHEM):
    CHEM = CHEM.split(" + ")
    AMINO = ""
    for eq in CHEM:
        for amino in AMINOACID_CHEM:
            if (eq == AMINOACID_CHEM[amino]):
                AMINO += amino;
    return AMINO;

def AminoAcidtoRNA(AMINO):
    RNA = ""
    i = 0;
    while (i <= len(AMINO)):
        for rnaset in RNA3:
            if AMINO[i:i+3] == RNA3[rnaset]:
                RNA += rnaset;
                break;
        i += 3;
    return RNA;

def RNAtoDNA(RNA):
    DNA = ""
    for c in RNA:
        if (c == "U"):
            DNA += "T"
        else:
            DNA += c
    return DNA
    

parser = argparse.ArgumentParser(description='Process DNA/RNA Structures')

parser.add_argument('structure', help='DNA/RNA Structure (Example: TCAGTTCGAATC)')
parser.add_argument('--rna', action='store_true', help='Use RNA as structure');
parser.add_argument('--chem', action='store_true', help='Generate chemical structure of protein produced by RNA/DNA')
parser.add_argument('--chemtoamino', action='store_true', help='Convert CHEM to Amino acid')
parser.add_argument('--aminotorna', action='store_true', help='Convert Amino acids to RNA')
parser.add_argument('--rnatodna', action='store_true', help='Convert RNA to DNA')
parser.add_argument('--peptide', action='store_true', help='Output Peptide form')
parser.add_argument('--visual', action='store_true', help='Visualize the protein')

args = parser.parse_args()

if (args.rna and not args.chemtoamino and not args.aminotorna and not args.rnatodna):
    try:
        if (args.chem):
            print(AminoAcidtoChem(RNAtoAminoAcid(args.structure)))
        elif (args.peptide):
            print(AminoAcidtoPeptide(RNAtoAminoAcid(args.structure)))
        elif (args.visual):
            for aa in AminoAcidtoPeptide(RNAtoAminoAcid(args.structure)): cmd._alt(string.lower(aa))
        else:
            print(RNAtoAminoAcid(args.structure))
    except Exception as e:
        print("Error: Invalid RNA Structure!")
elif (not args.rna and not args.chemtoamino and not args.aminotorna and not args.rnatodna):
    try:
        if (args.chem):
            print(AminoAcidtoChem(RNAtoAminoAcid(DNAtoRNA(args.structure))))
        elif (args.peptide):
            print(AminoAcidtoPeptide(RNAtoAminoAcid(DNAtoRNA(args.structure))))
        elif (args.visual):
            for aa in AminoAcidtoPeptide(RNAtoAminoAcid(DNAtoRNA(args.structure))): cmd._alt(string.lower(aa))
        else:
            print(RNAtoAminoAcid(DNAtoRNA(args.structure)))
    except Exception as e:
        print("Error: Invalid DNA Structure!")
elif (args.chemtoamino):
    print(ChemtoAminoAcid(args.structure))
elif (args.aminotorna):
    print(AminoAcidtoRNA(args.structure))
elif (args.rnatodna):
    print(RNAtoDNA(args.structure))