import string
from Bio import Alphabet
from Bio.Alphabet import IUPAC
from Bio.Data import IUPACData

unambiguous_dna_by_name = {}
unambiguous_dna_by_id = {}
unambiguous_rna_by_name = {}
unambiguous_rna_by_id = {}

# standard IUPAC unambiguous codons
standard_dna_table = None
standard_rna_table = None

# In the future, the back_table could return a statistically
# appropriate distribution of codons, so do not cache the results of
# back_table lookups!

class TranslationError(Exception):
    pass

class CodonTable:
    nucleotide_alphabet = Alphabet.generic_nucleotide
    protein_alphabet = Alphabet.generic_protein
    
    forward_table = {}    # only includes codons which actually code
    back_table = {}       # for back translations
    start_codons = []
    stop_codons = []
    # Not always called from derived classes!
    def __init__(self, nucleotide_alphabet = nucleotide_alphabet,
                 protein_alphabet = protein_alphabet,
                 forward_table = forward_table, back_table = back_table,
                 start_codons = start_codons, stop_codons = stop_codons):
        self.nucleotide_alphabet = nucleotide_alphabet
        self.protein_alphabet = protein_alphabet
        self.forward_table = forward_table
        self.back_table = back_table
        self.start_codons = start_codons
        self.stop_codons = stop_codons

def make_back_table(table, default_stop_codon):
    #  ONLY RETURNS A SINGLE CODON
    # Do the sort so changes in the hash implementation won't affect
    # the result when one amino acid is coded by more than one codon.
    back_table = {}
    keys = table.keys() ; keys.sort()
    for key in keys:
        back_table[table[key]] = key
    back_table[None] = default_stop_codon
    return back_table


class NCBICodonTable(CodonTable):
    nucleotide_alphabet = Alphabet.generic_nucleotide
    protein_alphabet = IUPAC.protein
    
    def __init__(self, id, names, table, start_codons, stop_codons):
        self.id = id
        self.names = names
        self.forward_table = table
        self.back_table = make_back_table(table, stop_codons[0])
        self.start_codons = start_codons
        self.stop_codons = stop_codons


class NCBICodonTableDNA(NCBICodonTable):
    nucleotide_alphabet = IUPAC.unambiguous_dna

class NCBICodonTableRNA(NCBICodonTable):
    nucleotide_alphabet = IUPAC.unambiguous_rna



def register_ncbi_table(name, alt_name, id,
                        table, start_codons, stop_codons):
    names = string.split(name, "; ")
    
    dna = NCBICodonTableDNA(id, names + [alt_name], table, start_codons,
                            stop_codons)
    # replace all T's with U's for the RNA tables
    rna_table = {}
    for codon, val in table.items():
        rna_table[string.replace(codon, "T", "U")] = val
    rna_start_codons = []
    for codon in start_codons:
        rna_start_codons.append(string.replace(codon, "T", "U"))
    rna_stop_codons = []
    for codon in stop_codons:
        rna_stop_codons.append(string.replace(codon, "T", "U"))
    
    rna = NCBICodonTableRNA(id, names + [alt_name], rna_table,
                            rna_start_codons, rna_stop_codons)

    if id == 1:
        global standard_dna_table, standard_rna_table
        standard_dna_table = dna
        standard_rna_table = rna

    unambiguous_dna_by_id[id] = dna
    unambiguous_rna_by_id[id] = rna

    if alt_name is not None:
        names.append(alt_name)

    for name in names:
        unambiguous_dna_by_name[name] = dna
        unambiguous_rna_by_name[name] = rna

### These tables created from the data file
###  ftp://ncbi.nlm.nih.gov/entrez/misc/data/gc.prt
### using the following:
##import re
##for line in open("gc.prt").readlines():
##    if line[:2] == " {":
##        names = []
##        id = None
##        aa = None
##        start = None
##        bases = []
##    elif line[:6] == "  name":
##        names.append(re.search('"([^"]*)"', line).group(1))
##    elif line[:8] == "    name":
##        names.append(re.search('"(.*)$', line).group(1))
##    elif line == ' Mitochondrial; Mycoplasma; Spiroplasma" ,\n':
##        names[-1] = names[-1] + " Mitochondrial; Mycoplasma; Spiroplasma"
##    elif line[:4] == "  id":
##        id = int(re.search('(\d+)', line).group(1))
##    elif line[:10] == "  ncbieaa ":
##        aa = line[12:12+64]
##    elif line[:10] == "  sncbieaa":
##        start = line[12:12+64]
##    elif line[:9] == "  -- Base":
##        bases.append(line[12:12+64])
##    elif line[:2] == " }":
##        assert names != [] and id is not None and aa is not None
##        assert start is not None and bases != []
##        if len(names) == 1:
##            names.append(None)
##        print "register_ncbi_table(name = %s," % repr(names[0])
##        print "                    alt_name = %s, id = %d", % \
##              (repr(names[1]), id)
##        print "                    table = {"
##        s = "    "
##        for i in range(64):
##            if aa[i] != "*":
##                t = " '%s%s%s': '%s'," % (bases[0][i], bases[1][i],
##                                          bases[2][i], aa[i])
##                if len(s) + len(t) > 75:
##                    print s
##                    s = "    " + t
##                else:
##                    s = s + t
##        print s, "},"

##        s = "                    stop_codons = ["
##        for i in range(64):
##            if aa[i] == "*":
##                t = " '%s%s%s'," % (bases[0][i], bases[1][i], bases[2][i])
##                if len(s) + len(t) > 75:
##                    print s
##                    s = "                                    " + t
##                else:
##                    s = s + t
##        print s, "],"

##        s = "                    start_codons = ["
##        for i in range(64):
##            if start[i] == "M":
##                t = " '%s%s%s'," % (bases[0][i], bases[1][i], bases[2][i])
##                if len(s) + len(t) > 75:
##                    print s
##                    s = "                                    " + t
##                else:
##                    s = s + t
##        print s, "]"
##        print "                    )"
##    elif line[:2] == "--" or line == "\n" or line == "}\n" or \
##         line == 'Genetic-code-table ::= {\n':
##        pass
##    else:
##        raise "Unparsed", repr(line)

register_ncbi_table(name = 'Standard',
                    alt_name = 'SGC0', id = 1,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGG': 'W', 'CTT': 'L', 'CTC': 'L',
     'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P', 'CCA': 'P',
     'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q', 'CAG': 'Q',
     'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R', 'ATT': 'I',
     'ATC': 'I', 'ATA': 'I', 'ATG': 'M', 'ACT': 'T', 'ACC': 'T',
     'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N', 'AAA': 'K',
     'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'R', 'AGG': 'R',
     'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V', 'GCT': 'A',
     'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D', 'GAC': 'D',
     'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G', 'GGA': 'G',
     'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', 'TGA', ],
                    start_codons = [ 'TTG', 'CTG', 'ATG', ]
                    )
register_ncbi_table(name = 'Vertebrate Mitochondrial',
                    alt_name = 'SGC1', id = 2,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGA': 'W', 'TGG': 'W', 'CTT': 'L',
     'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P',
     'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q',
     'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
     'ATT': 'I', 'ATC': 'I', 'ATA': 'M', 'ATG': 'M', 'ACT': 'T',
     'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N',
     'AAA': 'K', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'GTT': 'V',
     'GTC': 'V', 'GTA': 'V', 'GTG': 'V', 'GCT': 'A', 'GCC': 'A',
     'GCA': 'A', 'GCG': 'A', 'GAT': 'D', 'GAC': 'D', 'GAA': 'E',
     'GAG': 'E', 'GGT': 'G', 'GGC': 'G', 'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', 'AGA', 'AGG', ],
                    start_codons = [ 'ATT', 'ATC', 'ATA', 'ATG', 'GTG', ]
                    )
register_ncbi_table(name = 'Yeast Mitochondrial',
                    alt_name = 'SGC2', id = 3,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGA': 'W', 'TGG': 'W', 'CTT': 'T',
     'CTC': 'T', 'CTA': 'T', 'CTG': 'T', 'CCT': 'P', 'CCC': 'P',
     'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q',
     'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
     'ATT': 'I', 'ATC': 'I', 'ATA': 'M', 'ATG': 'M', 'ACT': 'T',
     'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N',
     'AAA': 'K', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'R',
     'AGG': 'R', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
     'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D',
     'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G',
     'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', ],
                    start_codons = [ 'ATG', ]
                    )
register_ncbi_table(name = 'Mold Mitochondrial; Protozoan Mitochondrial; Coelenterate Mitochondrial; Mycoplasma; Spiroplasma',
                    alt_name = 'SGC3', id = 4,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGA': 'W', 'TGG': 'W', 'CTT': 'L',
     'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P',
     'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q',
     'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
     'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M', 'ACT': 'T',
     'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N',
     'AAA': 'K', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'R',
     'AGG': 'R', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
     'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D',
     'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G',
     'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', ],
                    start_codons = [ 'TTA', 'TTG', 'CTG', 'ATT', 'ATC',
                                     'ATA', 'ATG', 'GTG', ]
                    )
register_ncbi_table(name = 'Invertebrate Mitochondrial',
                    alt_name = 'SGC4', id = 5,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGA': 'W', 'TGG': 'W', 'CTT': 'L',
     'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P',
     'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q',
     'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
     'ATT': 'I', 'ATC': 'I', 'ATA': 'M', 'ATG': 'M', 'ACT': 'T',
     'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N',
     'AAA': 'K', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'S',
     'AGG': 'S', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
     'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D',
     'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G',
     'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', ],
                    start_codons = [ 'TTG', 'ATT', 'ATC', 'ATA', 'ATG',
                                     'GTG', ]
                    )
register_ncbi_table(name = 'Ciliate Nuclear; Dasycladacean Nuclear; Hexamita Nuclear',
                    alt_name = 'SGC5', id = 6,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TAA': 'Q', 'TAG': 'Q', 'TGT': 'C', 'TGC': 'C', 'TGG': 'W',
     'CTT': 'L', 'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P',
     'CCC': 'P', 'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H',
     'CAA': 'Q', 'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R',
     'CGG': 'R', 'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M',
     'ACT': 'T', 'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N',
     'AAC': 'N', 'AAA': 'K', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S',
     'AGA': 'R', 'AGG': 'R', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V',
     'GTG': 'V', 'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A',
     'GAT': 'D', 'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G',
     'GGC': 'G', 'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TGA', ],
                    start_codons = [ 'ATG', ]
                    )
register_ncbi_table(name = 'Echinoderm Mitochondrial',
                    alt_name = 'SGC8', id = 9,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGA': 'W', 'TGG': 'W', 'CTT': 'L',
     'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P',
     'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q',
     'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
     'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M', 'ACT': 'T',
     'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N',
     'AAA': 'N', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'S',
     'AGG': 'S', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
     'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D',
     'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G',
     'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', ],
                    start_codons = [ 'ATG', ]
                    )
register_ncbi_table(name = 'Euplotid Nuclear',
                    alt_name = 'SGC9', id = 10,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGA': 'C', 'TGG': 'W', 'CTT': 'L',
     'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P',
     'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q',
     'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
     'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M', 'ACT': 'T',
     'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N',
     'AAA': 'K', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'R',
     'AGG': 'R', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
     'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D',
     'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G',
     'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', ],
                    start_codons = [ 'ATG', ]
                    )
register_ncbi_table(name = 'Bacterial',
                    alt_name = None, id = 11,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGG': 'W', 'CTT': 'L', 'CTC': 'L',
     'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P', 'CCA': 'P',
     'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q', 'CAG': 'Q',
     'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R', 'ATT': 'I',
     'ATC': 'I', 'ATA': 'I', 'ATG': 'M', 'ACT': 'T', 'ACC': 'T',
     'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N', 'AAA': 'K',
     'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'R', 'AGG': 'R',
     'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V', 'GCT': 'A',
     'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D', 'GAC': 'D',
     'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G', 'GGA': 'G',
     'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', 'TGA', ],
                    start_codons = [ 'TTG', 'CTG', 'ATT', 'ATC', 'ATA',
                                     'ATG', 'GTG', ]
                    )
register_ncbi_table(name = 'Alternative Yeast Nuclear',
                    alt_name = None, id = 12,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGG': 'W', 'CTT': 'L', 'CTC': 'L',
     'CTA': 'L', 'CTG': 'S', 'CCT': 'P', 'CCC': 'P', 'CCA': 'P',
     'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q', 'CAG': 'Q',
     'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R', 'ATT': 'I',
     'ATC': 'I', 'ATA': 'I', 'ATG': 'M', 'ACT': 'T', 'ACC': 'T',
     'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N', 'AAA': 'K',
     'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'R', 'AGG': 'R',
     'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V', 'GCT': 'A',
     'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D', 'GAC': 'D',
     'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G', 'GGA': 'G',
     'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', 'TGA', ],
                    start_codons = [ 'CTG', 'ATG', ]
                    )
register_ncbi_table(name = 'Ascidian Mitochondrial',
                    alt_name = None, id = 13,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TGT': 'C', 'TGC': 'C', 'TGA': 'W', 'TGG': 'W', 'CTT': 'L',
     'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P',
     'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q',
     'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
     'ATT': 'I', 'ATC': 'I', 'ATA': 'M', 'ATG': 'M', 'ACT': 'T',
     'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N',
     'AAA': 'K', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'G',
     'AGG': 'G', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
     'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D',
     'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G',
     'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TAG', ],
                    start_codons = [ 'ATG', ]
                    )
register_ncbi_table(name = 'Flatworm Mitochondrial',
                    alt_name = None, id = 14,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TAA': 'Y', 'TGT': 'C', 'TGC': 'C', 'TGA': 'W', 'TGG': 'W',
     'CTT': 'L', 'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P',
     'CCC': 'P', 'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H',
     'CAA': 'Q', 'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R',
     'CGG': 'R', 'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M',
     'ACT': 'T', 'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N',
     'AAC': 'N', 'AAA': 'N', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S',
     'AGA': 'S', 'AGG': 'S', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V',
     'GTG': 'V', 'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A',
     'GAT': 'D', 'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G',
     'GGC': 'G', 'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAG', ],
                    start_codons = [ 'ATG', ]
                    )
register_ncbi_table(name = 'Blepharisma Macronuclear',
                    alt_name = None, id = 15,
                    table = {
     'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L', 'TCT': 'S',
     'TCC': 'S', 'TCA': 'S', 'TCG': 'S', 'TAT': 'Y', 'TAC': 'Y',
     'TAG': 'Q', 'TGT': 'C', 'TGC': 'C', 'TGG': 'W', 'CTT': 'L',
     'CTC': 'L', 'CTA': 'L', 'CTG': 'L', 'CCT': 'P', 'CCC': 'P',
     'CCA': 'P', 'CCG': 'P', 'CAT': 'H', 'CAC': 'H', 'CAA': 'Q',
     'CAG': 'Q', 'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
     'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M', 'ACT': 'T',
     'ACC': 'T', 'ACA': 'T', 'ACG': 'T', 'AAT': 'N', 'AAC': 'N',
     'AAA': 'K', 'AAG': 'K', 'AGT': 'S', 'AGC': 'S', 'AGA': 'R',
     'AGG': 'R', 'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
     'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A', 'GAT': 'D',
     'GAC': 'D', 'GAA': 'E', 'GAG': 'E', 'GGT': 'G', 'GGC': 'G',
     'GGA': 'G', 'GGG': 'G', },
                    stop_codons = [ 'TAA', 'TGA', ],
                    start_codons = [ 'ATG', ]
                    )

#########  Deal with ambiguous forward translations

class AmbiguousCodonTable(CodonTable):
    def __init__(self, codon_table,
                 ambiguous_nucleotide_alphabet,
                 ambiguous_nucleotide_values,
                 ambiguous_protein_alphabet,
                 ambiguous_protein_values):
        CodonTable.__init__(self,
                            ambiguous_nucleotide_alphabet,
                            ambiguous_protein_alphabet,
                            AmbiguousForwardTable(codon_table.forward_table,
                                                  ambiguous_nucleotide_values,
                                                  ambiguous_protein_values),
                            codon_table.back_table,

                            # These two are WRONG!  I need to get the
                            # list of ambiguous codons which code for
                            # the stop codons  XXX
                            list_ambiguous_codons(codon_table.start_codons),
                            list_ambiguous_codons(codon_table.stop_codons)
                            )
        self._codon_table = codon_table

    # Be sneaky and forward attribute lookups to the original table.
    # This lets us get the names, if the original table is an NCBI
    # table.
    def __getattr__(self, name):
        return getattr(self._codon_table, name)

def list_possible_proteins(codon, forward_table, ambiguous_nucleotide_values):
        c1, c2, c3 = codon
        x1 = ambiguous_nucleotide_values[c1]
        x2 = ambiguous_nucleotide_values[c2]
        x3 = ambiguous_nucleotide_values[c3]
        possible = {}
        stops = []
        for y1 in x1:
            for y2 in x2:
                for y3 in x3:
                    try:
                        possible[forward_table[y1+y2+y3]] = 1
                    except KeyError:
                        # If tripping over a stop codon
                        stops.append(y1+y2+y3)
        if stops:
            if possible.keys():
                raise TranslationError, ("ambiguous codon codes for both proteins and stop codons", codon)
            # This is a true stop codon - tell the caller about it
            raise KeyError, codon
        return possible.keys()

def list_ambiguous_codons(codons):
    # XXX not implemented!
    return codons

# Forward translation is "onto", that is, any given codon always maps
# to the same protein, or it doesn't map at all.  Thus, I can build
# off of an existing table to produce the ambiguous mappings.
#
# This handles the general case.  Perhaps it's overkill?
#  >>> t = CodonTable.ambiguous_dna_by_id[1]
#  >>> t.forward_table["AAT"]
#  'N'
#  >>> t.forward_table["GAT"]
#  'D'
#  >>> t.forward_table["RAT"]
#  'B'
#  >>> t.forward_table["YTA"]
#  'L'

class AmbiguousForwardTable:
    def __init__(self, forward_table, ambiguous_nucleotide, ambiguous_protein):
        self.forward_table = forward_table

        self.ambiguous_nucleotide = ambiguous_nucleotide
        self.ambiguous_protein = ambiguous_protein

        inverted = {}
        for name, val in ambiguous_protein.items():
            for c in val:
                x = inverted.get(c, {})
                x[name] = 1
                inverted[c] = x
        for name, val in inverted.items():
            inverted[name] = val.keys()
        self._inverted = inverted
        
        self._cache = {}

    def get(self, codon, failobj = None):
        try:
            return self.__getitem__(codon)
        except KeyError:
            return failobj
        
    def __getitem__(self, codon):
        try:
            x = self._cache[codon]
        except KeyError:
            pass
        else:
            if x is TranslationError:
                raise TranslationError, codon   # no unique translation
            if x is KeyError:
                raise KeyError, codon  # it's a stop codon
            return x
        try:
            x = self.forward_table[codon]
            self._cache[codon] = x
            return x
        except KeyError:
            pass

        # XXX Need to make part of this into a method which returns
        # a list of all possible encodings for a codon!
        try:
            possible = list_possible_proteins(codon,
                                              self.forward_table,
                                              self.ambiguous_nucleotide)
        except KeyError:
            self._cache[codon] = KeyError
            raise KeyError, codon  # stop codon
        except TranslationError:
            self._cache[codon] = TranslationError
            raise TranslationError, codon  # does not code
        assert len(possible) > 0, "unambiguous codons must code"

        # Hah!  Only one possible protein, so use it
        if len(possible) == 1:
            self._cache[codon] = possible[0]
            return possible[0]

        # See if there's an ambiguous protein encoding for the multiples.
        # Find residues which exist in every coding set.
        ambiguous_possible = {}
        for amino in possible:
            for term in self._inverted[amino]:
                ambiguous_possible[term] = ambiguous_possible.get(term, 0) + 1

        n = len(possible)
        possible = []
        for amino, val in ambiguous_possible.items():
            if val == n:
                possible.append(amino)

        # No amino acid encoding for the results
        if len(possible) == 0:
            self._cache[codon] = TranslationError
            raise TranslationError, codon   # no valid translation

        # All of these are valid, so choose one
        # To be unique, sort by smallet ambiguity then alphabetically
        # Can get this if "X" encodes for everything.
        def _sort(x, y, table = self.ambiguous_protein):
            a = cmp(len(table[x]), len(table[y]))
            if a == 0:
                return cmp(x, y)
            return a
        possible.sort(_sort)
                          
        x = possible[0]
        self._cache[codon] = x
        return x


ambiguous_dna_by_name = {}
for key, val in unambiguous_dna_by_name.items():
    ambiguous_dna_by_name[key] = AmbiguousCodonTable(val,
                                     IUPAC.ambiguous_dna,
                                     IUPACData.ambiguous_dna_values,
                                     IUPAC.extended_protein,
                                     IUPACData.extended_protein_values)
ambiguous_dna_by_id = {}
for key, val in unambiguous_dna_by_id.items():
    ambiguous_dna_by_id[key] = AmbiguousCodonTable(val,
                                     IUPAC.ambiguous_dna,
                                     IUPACData.ambiguous_dna_values,
                                     IUPAC.extended_protein,
                                     IUPACData.extended_protein_values)

ambiguous_rna_by_name = {}
for key, val in unambiguous_rna_by_name.items():
    ambiguous_rna_by_name[key] = AmbiguousCodonTable(val,
                                     IUPAC.ambiguous_rna,
                                     IUPACData.ambiguous_rna_values,
                                     IUPAC.extended_protein,
                                     IUPACData.extended_protein_values)
ambiguous_rna_by_id = {}
for key, val in unambiguous_rna_by_id.items():
    ambiguous_rna_by_id[key] = AmbiguousCodonTable(val,
                                     IUPAC.ambiguous_rna,
                                     IUPACData.ambiguous_rna_values,
                                     IUPAC.extended_protein,
                                     IUPACData.extended_protein_values)
del key, val

