################################################################################
#
#       This file is part of the General Hidden Markov Model Library,
#       GHMM version __VERSION__, see http://ghmm.org
#
#       file:    modhmmer.py
#       authors: Utz Pape, Benjamin Georgi
#
#       Copyright (C) 1998-2004 Alexander Schliep
#       Copyright (C) 1998-2001 ZAIK/ZPR, Universitaet zu Koeln
#       Copyright (C) 2002-2004 Max-Planck-Institut fuer Molekulare Genetik,
#                               Berlin
#
#       Contact: schliep@ghmm.org
#
#       This library is free software; you can redistribute it and/or
#       modify it under the terms of the GNU Library General Public
#       License as published by the Free Software Foundation; either
#       version 2 of the License, or (at your option) any later version.
#
#       This library is distributed in the hope that it will be useful,
#       but WITHOUT ANY WARRANTY; without even the implied warranty of
#       MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#       Library General Public License for more details.
#
#       You should have received a copy of the GNU Library General Public
#       License along with this library; if not, write to the Free
#       Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
#
#
#       This file is version $Revision: 1057 $
#                       from $Date: 2004-11-11 10:29:18 -0500 (Thu, 11 Nov 2004) $
#             last change by $Author: cic99 $.
#
################################################################################
import sys,re,string,StringIO
from xml.dom import minidom

def gotoLine(f,res):
    rematch = None
    r = " "
    while (rematch is None) and (r != ""):
       
        r = f.readline()
        rematch = res.search(r)
    if ( rematch == None ):
        # print "Fehler in gotoLine"
        pass
    else:
        return rematch

def build_matrix(n,m):
    "builds n x m matrix with lists"
    matrix = range(n)
    for i in range(n):
        matrix[i] = range(m)
        for j in range(m):
            matrix[i][j] = 0
    return matrix

def sum_mrows(mat):
    "sums the rows of a matrix"
    mout = range(len(mat))
    for i in range(len(mat)):
        s = 0
        for j in range(len(mat[i])):
            s=s+mat[i][j]
        mout[i] = s
    return mout

def norm_mat(mat):
    #print mat
    for i in range(len(mat)):
        s = 0.0
        #print "i: " + str(i)
        for j in range(len(mat[i])):
            s = s+mat[i][j]
        for j in range(len(mat[i])):
            #print "j: " + str(j) + ",  "+ str(mat[i][j])
            mat[i][j] = mat[i][j]/s

def red_mat_end(mat,r):
    "delete <r> rows and columns from the end of the matrix"
    for i in range(len(mat)-r):
        del mat[i][-1*r:]
    del mat[-1*r:]

def del_mat(mat,r):
    "deletes the <r>th columns and row from the matrix"
    for i in range(len(mat)):
        del mat[i][r]
    del mat[r]

def toint(i):
    "return integer if i is value or 0"
    try:
        iout = int(i)
    except ValueError:
        iout = 0
    return iout

def map_entries(dic,lis):
    "translates the letters to the number of the columns"
    dicout = {}
    for k in dic.keys():
        dicout[k] = []
        for i in range(len(dic[k])):
            dicout[k].append((lis.index(dic[k][i][0]),lis.index(dic[k][i][1])))
    return dicout

def xml_newdatanode(doc,nodename,attributename,attribute,text):
    "returns node with attribute and text"
    nod = doc.createElement(nodename)
    nod.setAttribute(attributename,attribute)
    nod.appendChild(doc.createTextNode(text))
    return nod

def write_file(strf,strcontent):
    "writes <strcontent> in file <strf>"
    try:
        f = open(strf,"w")
    except IOError,info:
        sys.stderr.write(str(info) + "\n")
        sys.exit(1)
    try:
        f.write(strcontent)
    finally:
        f.close()

def remove_state(matrix, index):
	for i in range(len(matrix)):
		del(matrix[i][index])
	del(matrix[index])	
	return matrix


class hmmer:
    "reads hmmer file and converts it to xml"

    #a few constants
    intscale = 1000.0 #hmmer format specific (but it has to be a float!!!)
    
    #positions of the entries
    lisHead = ["N","B","E","J","C","T","M","I","D"]
    lisMID = ["M","I","D"]

    #map table for HMMER file format
    dicTEntries = { "XT": [("N","B"),("N","N"),("E","C"),("E","J"),
                           ("C","T"),("C","C"),("J","B"),("J","J")],
                    "BD": [("B","D")],
                    "HMM":[("M","M"),("M","I"),("M","D"),("I","M"),
                            ("I","I"),("D","M"),("D","D"),("B","M"),("M","E")],
                    }

    #dictionary for the letters
    dicLetters = { 4:  ["A","C","G","T"],
                   20: ["A","C","D","E","F","G","H","I","K","L","M","N","P","Q",
                        "R","S","T","V","W","Y"]}
    
    #map table translated with positions
    dicTE = map_entries(dicTEntries,lisHead)

    dicType = {"Amino":20,"Nucleotide":4}

    #coordinate offset for graphical presentation
    off_distx = 75
    off_disty = 75

    def __init__(self,strf):
        #print "start __init__"
        
        try:
		    # f = StringIO.StringIO(strf)
            if isinstance(strf,str): 
                f = open(strf,"r")
            else:
                f = strf    
        except IOError,info:
                 sys.stderr.write(str(info) + "\n")
            #     sys.exit(1)
        
        try:
            r = f.readline()
            
            self.acc = gotoLine(f,re.compile("^NAME\s+([\w+\s*]+)" ) ).group(1)  # NAME as unique model identifier
            # self.acc  = gotoLine(f,re.compile("^ACC\s+(\w+)" ) ).group(1)      # ACC as unique model identifier
            
            if self.acc[-1] == '\n':  # remove newline at end of string
                self.acc = self.acc[:-1]
                    	
            #get number of match states
            n = int(gotoLine(f,re.compile("^LENG\s*(\d+)")).group(1))
            self.n = n
            #get type of profile hmm: amino/nucleotide
            m = self.dicType[gotoLine(f,re.compile("^ALPH\s*(\S+)")).group(1)];
            self.m = m

            #build matrix for transition: N B E J C T M1 I1 D1 M2 I2 D2 ... Mn In Dn
            self.matrans = build_matrix(3*n+6,3*n+6)
            
			#emission matrix: match state, insert state, null_model
            self.maems = [build_matrix(n,m),build_matrix(n,m),build_matrix(1,m)]

			#get line "XT" transitions
            trans = string.split(gotoLine(f,re.compile("^XT([\s\d\S]*)")).group(1))
            self.set_matrix("XT",trans)

            #null model
            trans = string.split(gotoLine(f,re.compile("^NULT([\s\d\S]*)")).group(1))
            self.manull = [self.H2P(trans[0],1.0),self.H2P(trans[1],1.0)] #[G->G,G->F]
            trans = string.split(gotoLine(f,re.compile("^NULE([\s\d\S]*)")).group(1))
            for k in range(len(trans)):
                self.maems[2][0][k] = self.H2P(trans[k],1/float(m))
            
 
            # ***  Main section ***
            gotoLine(f,re.compile("^HMM"))
            f.readline()
            #print "33333"
            #get B->D transition probability
            self.set_matrix("BD",[string.split(f.readline())[2]])
            #recurse over number of match states
            for i in range(n):
                for j in range(3):
                    lis = string.split(f.readline())
                    del lis[0]
                    if j==2:
                        #state transition probs
                        self.set_matrix("HMM",lis,i*3)
                    else:
                        #emmission probs: [match=0/insert=1][state][emission-letter]
                        for k in range(self.m):
                            #print "state: "+ str(i) +" symbol: " +str(k) +" - score = " + str(lis[k]) + ", NP = " + str(self.maems[2][0][k]) +" -> " + str(self.H2P(lis[k],self.maems[2][0][k]))							
                            self.maems[j][i][k] = self.H2P(lis[k],self.maems[2][0][k])
                            #print "null prob: " + str(self.maems[2][0][k])
            #print "44444"
            del self.maems[1][-1] #delete last (not existing) insertion state
            self.matrans[self.lisHead.index("T")][self.lisHead.index("T")] = 1.0 #set end state as loop to prevent normalization issues
            del_mat(self.matrans,-2)
            self.matrans[-1][self.lisHead.index("E")] = 1.0 # set last D->E=1
        finally:
            #pass
            f.close()
		
        #print "55555"	
        #normalize matrices:
        for i in range(len(self.maems)):
            #print "index1: " +str(i)
            #print self.maems[i]
            norm_mat(self.maems[i])
        #print "trans:"
        norm_mat(self.matrans)
		
        #print "66666"
        
        self.matrans[self.lisHead.index("T")][self.lisHead.index("T")] = 0.0
        #print "parsen fertig"
		
    def __str__(self):
        print "oben"
        hmm_str = "N= " + str(self.n)  +", M= " + str(self.m) + "\n"
        hmm_str += "Transitions: \n"
        for row in self.matrans:
            hmm_str += str(row) + "\n" 
        hmm_str += "Emissions: \n"		
        for row in self.maems:
            hmm_str += str(row) + "\n" 
        print hmm_str	
        return hmm_str


    def set_matrix(self,type,lis,offset=0):
        "fills matrix with values from line <type> and adds optional offset"
        for k in range(len(lis)):
            if lis[k]!="*":
                x_c = self.dicTEntries[type][k][0]
                y_c = self.dicTEntries[type][k][1]
                x = self.dicTE[type][k][0]+offset*(x_c in self.lisMID)
                y = self.dicTE[type][k][1]
                if y_c in self.lisMID:
                    y = y+offset
                    if ((x_c in ["D","I"]) or (x_c==y_c=="M") or (x_c=="M" and y_c=="D")) and not (x_c==y_c=="I"):
                        y=y+3
                self.matrans[x][y] = self.H2P(lis[k],1.0)
               # print (x_c + str(offset),y_c + str(offset)),self.matrans[x][y]

    def H2P(self,score,null_prob):
        "returns the probability"
        if score == "*":
            return 0
        else:
            return null_prob * 2**(float(score)/self.intscale)

    def globalConfig(self):
        """ Configures HMM for global Needelman-Wunsch alignments. """
        #build matrix for transition: N B E J C T M1 I1 D1 M2 I2 D2 ... Mn In Dn
        p1 = self.manull[0]

        print "globalConfig:  self_p = ", p1
        
        self.matrans[0][0] = 1.0 - p1  #  N->B
        self.matrans[0][1] = p1       # N->N
        self.matrans[2][4] = 1.0      # E->C
        self.matrans[2][3] = 0        # E->J
        self.matrans[4][4] = p1        # C->C
        self.matrans[4][5] = 1- p1        # C->T        
        
    
    def getGHMMmatrices(self):
        """
            Converts the MODHMMER matrix format into matrices to 
            be used in ghmm.HMMFromMatrices.
        
        """
    
        emiss_mat = []

        # intitializing pi vector, always starting in B state (index 0)
        pi = [1] + [0] * ((self.n * 3) + 4)
    
        silent = [0] * (self.m)  # emissions for silent states
        equal = [1.0/(self.m)] * self.m  # uniform distribution over the number of symbols

        # conversion of the HMMER emission matrices into ghmm format
        # emmission probs in HMMER: [match=0/insert=1][state][emission-letter]
        # order of states in HMMER transition matrix: N B E J C T M1 I1 D1 M2 I2 D2 ... Mn In Dn
    
        emiss_mat.append(equal)    # N state (equal)
        emiss_mat.append(silent)     # B state (silent)
        emiss_mat.append(silent)       # E state (silent)
        emiss_mat.append(equal)     # J state  (equal)
        emiss_mat.append(equal)       # C state  (equal)
        emiss_mat.append(silent)       # T (silent)
    
        for ind1 in range(self.n): # n = number of M/I/D blocks
            emiss_mat.append(self.maems[0][ind1])  # M state
            if (ind1 != self.n-1):
                emiss_mat.append(self.maems[1][ind1]) # I state
            emiss_mat.append(silent) #(silent)     # D state (silent)
            
        return [self.matrans,emiss_mat,pi,self.acc]

    
    def get_dom(self):
        "returns DOM object"
        doc = minidom.Document()
        nodgraphml = doc.createElement("graphml")

        nodkey = doc.createElement("key")
        nodkey.setAttribute("for","node")
        nodkey.setAttribute("gd:type","HigherDiscreteProbDist")
        nodkey.setAttribute("id","emissions")
        nodgraphml.appendChild(nodkey)

        #declare dummy class
        nodclass = doc.createElement("hmm:class")
        nodclass.setAttribute("hmm:high","0")
        nodclass.setAttribute("hmm:low","0")
        nodmap = doc.createElement("map")
        nodsym = doc.createElement("symbol")
        nodsym.setAttribute("code","0")
        nodsym.setAttribute("desc","Simple")
        nodsym.appendChild(doc.createTextNode("N"))
        nodmap.appendChild(nodsym)
        nodclass.appendChild(nodmap)
        nodgraphml.appendChild(nodclass)

        #declare alphabet
        nodalpha = doc.createElement("hmm:alphabet")
        nodalpha.setAttribute("hmm:high",str(self.m-1))
        nodalpha.setAttribute("hmm:low","0")
        nodalpha.setAttribute("hmm:type","discrete")
        nodmap = doc.createElement("map")
        for k in self.dicLetters[self.m]:
            nodmap.appendChild(xml_newdatanode(doc,"symbol","code",str(self.dicLetters[self.m].index(k)),str(k)))
        nodalpha.appendChild(nodmap)
        nodgraphml.appendChild(nodalpha)

        dicHash = {}
        #nodes/states
        nodgraph = doc.createElement("graph")
        #add match/insert/delete states
        for k in self.lisHead:
            if k in self.lisMID:
                n = self.n
                if k == "I":
                    n=n-1
            else:
                n = 1
            for i in range(n):
                pos = self.lisHead.index(k)
                nodnode = doc.createElement("node")
                #node:data:label
                strlabel = k
                if (n>1) or (k in self.lisMID):
                    strlabel = strlabel + str(i)
                nodnode.setAttribute("id",strlabel)
                h = pos + (i*3)
                if (k=="D") and (i==n-1):
                    h=h-1 #because last I column is deleted and therefore delete column index is old one minus one
                dicHash[h] = strlabel
                nodnode.appendChild(xml_newdatanode(doc,"data","key","label",strlabel))
                #node:data:class
                nodnode.appendChild(xml_newdatanode(doc,"data","key","class","0"))
                #node:data:initial
                nodnode.appendChild(xml_newdatanode(doc,"data","key","initial",str(k=="N")))
                #node:data:ngeom
                noddata = doc.createElement("data")
                noddata.setAttribute("key","ngeom")
                nodpos = doc.createElement("pos")
                if k == "N":
                    strx = self.off_distx
                    stry = 4*self.off_disty
                elif k=="B":
                    strx = 2*self.off_distx
                    stry = 4*self.off_disty
                elif k=="J":
                    strx = (self.n+4)/2*self.off_distx
                    stry = 7*self.off_disty
                elif k=="E":
                    strx = (self.n + 3)*self.off_distx
                    stry = 2*self.off_disty
                elif k=="C":
                    strx = (self.n + 4)*self.off_distx
                    stry = 2*self.off_disty
                elif k=="T":
                    strx = (self.n + 5)*self.off_distx
                    stry = 2*self.off_disty
                elif k=="M":
                    strx = (3+i)*self.off_distx
                    stry = 3*self.off_disty
                elif k=="I":
                    strx = (3+i)*self.off_distx
                    stry = 1*self.off_disty
                elif k=="D":
                    strx = (3+i)*self.off_distx
                    stry = 5*self.off_disty
                nodpos.setAttribute("x",str(strx))
                nodpos.setAttribute("y",str(stry))
                noddata.appendChild(nodpos)
                nodnode.appendChild(noddata)
                #node:data:emissions
                if k in ["M","I"]:
                    strem = string.join(map(str,self.maems[self.lisMID.index(k)][i]),", ")
                elif k in ["N","C"]:
                    strem = string.join(map(str,self.maems[2][0]),", ")
                else:
                    strem = (self.m-1)*"0.0, " + "0.0"
                nodnode.appendChild(xml_newdatanode(doc,"data","key","emissions",strem))
                nodgraph.appendChild(nodnode)
                
        #edges/transitions
        for i in range(len(self.matrans)):
            for j in range(len(self.matrans[i])):
                if self.matrans[i][j]!=0:
                    nodedge = doc.createElement("edge")
                    nodedge.setAttribute("source",dicHash[i])
                    nodedge.setAttribute("target",dicHash[j])
                    nodedge.appendChild(xml_newdatanode(doc,"data","key","prob",str(self.matrans[i][j])))
                    nodgraph.appendChild(nodedge)

        nodgraphml.appendChild(nodgraph)

        doc.appendChild(nodgraphml)
        return doc

    def write_prettyxml(self,strf):
        "writes hmm in pretty xml format to file"
        write_file(strf,self.get_dom().toprettyxml())

    def write_xml(self,strf):
        "writes hmm in xml format to file"
        write_file(strf,self.get_dom().toxml())
