class Codes(object):

    def __init__(self):
        codesfile = open('coding.txt')
        self.codes = {}
        self.abbrv = {}
        for line in codesfile:
            (name, abbreviation, codon) = line.split()
            self.codes[codon] = abbreviation
            self.abbrv[abbreviation] = name

    def translate(self, codon):
        """
        Translates from codon (e.g. UUG) to the 1-letter abbreviation
        (e.g. 'L')
        """
        if self.codes.has_key(codon):
            return self.codes[codon]
        else:
            return "?"

    def translateName(self, codon):
        """
        Translates from codon (e.g. UUG) to the 3-letter abbreviation
        (e.g. 'Leu')
        """
        if self.codes.has_key(codon):
            return self.abbrv[self.codes[codon]]
        else:
            return "???"

class DistanceMatrix(object):
    """
    Stores a distance matrix such as Blosum80
    """

    def __init__(self, filename):
        self.readMatrix(filename)


    def readMatrix(self, filename):
        """
        Read the matrix file.

        Does not check to make sure the file exists.
        """
        self.matrix = {}
        infile = open(filename)
        line = infile.readline()
        columns = line.split()
        for rowIdx in range(len(columns)-1):
            line = infile.readline().split()
            aaFrom = line[0]
            for colIdx in range(1,len(line)):
                aaTo = columns[colIdx]
                self.matrix[(aaFrom, aaTo)] = int(line[colIdx])

    def distance(self, x, y):
        """
        Return the distance from x to y.
        """
        try:
            return self.matrix[(x, y)]
        except KeyError:
            return None


"""
For translation
"""
codes = Codes()

def translate(codon):
    return codes.translate(codon)

def translateName(codon):
    return codes.translateName(codon)


"""
For graphing population changes
"""
def populationGraph(recessive, dominant, mixed):
    import pylab
    numOrganisms = recessive[0] + dominant[0] + mixed[0]
    pylab.plot(range(len(recessive)), recessive, label="tt")
    pylab.plot(range(len(dominant)), dominant, label="TT")
    pylab.plot(range(len(mixed)), mixed, label="Tt and tT")
    pylab.legend(loc='upper left')
    pylab.axis([0,len(recessive),0,numOrganisms])
    pylab.xlabel('Number of Generations')
    pylab.ylabel('Number of Organisms')
    pylab.show()
                            

"""
For alignment
"""
BLOSUM80 = DistanceMatrix("blosum80.txt")

def blosum80(r1,r2):
    """
    Returns the Blosum80 score for substituting amino acid r1 for amino
    acid r2.
    """
    distance = BLOSUM80.distance(r1,r2)
    if distance == None:
        raise RuntimeError("Parameters to blosum80 must be amino acids")
    else:
        return distance

def drawTree(pairwiseScores, rowLabels):
    """
    Cluster the pairwise scores using hcluster and display using the
    provided labels in pylab.
    """
    from hcluster import squareform, linkage, dendrogram
    from numpy import array
    import pylab

    #Create numpy squareform for hcluster: Set it up so that the most
    #similar organisms are 0 and the least similar are 1
    sqfrm = array(squareform(pairwiseScores))
    maxval = max(sqfrm)
    sqfrm = maxval - sqfrm + 5

    #clustmethod can be one of ['complete', 'single', 'average', 'weighted']
    clustmethod = 'complete'
    clustering = linkage(sqfrm, method=clustmethod)

    dendrogram(clustering, labels=rowLabels,orientation='left',
               leaf_font_size=6)

    pylab.subplots_adjust(left=0.3)
    #pylab.savefig(clustmethod+'.png', dpi=300, format='png')
    pylab.show()
    pylab.close()