# Generate graph from KEGG kgml file and expression data file

import xml.etree.ElementTree as ET
import networkx as NX
import itertools

edgelabel = {'activation' : 'activate', 
             'expression' : 'activate',
             'inhibition' : 'inhibit',
             'repression' : 'inhibit',
             'indirect effect' : 'interact',
             'state change' : 'interact',
             'binding/association' : 'interact',
             'dissociation' : 'interact',
#             'missing interaction' : 'interact',
             'phosphorylation' : 'interact',
             'dephosphorylation' : 'interact',
             'glycosylation' : 'interact',
             'ubiquitination' : 'interact',
             'methylation' : 'interact' }; # keys are kegg terms for describing interactions, values are reduced representations of these terms used as edge labels in pathway graphs

class pathway:
    """
    This class reads KGML files and generates traversable graph
    """

    def getpathway( self ):
        """
        Parameters: 
		self
	Output:
		pathway as networkx object
	Description:
		reads the tree generated by xml file and renders the pathway as a graph
	""" 
        G = NX.Graph() # create graph object
        self.pathname = self.root.attrib['title'];
        for child in self.root:
             if (child.tag == "entry"):
                 entryid = child.attrib['id']
                 if (child.attrib['type'] == "gene"):
                     G.add_node(entryid) # only add "genes", not "ortholog" or "compound" or other entry types
                     for gc in child:
                         if (gc.tag == "graphics"):
                             #genelist = gc.attrib['name'].split(', ')
                             genes = gc.attrib['name'] # do not split the list of aliases, because we want to save non-redundant gene names from all pathways. This was it is faster to check redundancy.
                             self.genedict[entryid] = genes
                             #print entryid, genelist
                 # elif (child.attrib['type'] == "group"):
             elif (child.tag == "relation"):
                 node1 = child.attrib['entry1']
                 node2 = child.attrib['entry2']
                 edgetype = "None" # default value for interaction between nodes
                 for gc in child:
                     if (gc.tag == "subtype"):
                         edgetype = gc.attrib['name']
#                 print node1, node2, edgetype, edgelabel[edgetype];
                 if edgetype in edgelabel.keys(): # add edge only if interaction is defined
                     G.add_edge(node1, node2, type = edgelabel[edgetype]) 
	return G;   

    def deapFormat( self, fname ):
        counter = 1;
        for g in self.connectedpaths:
            if (len(g) > 2): # only look at connected subgraphs with > 2 nodes
                # print "subgraph\n";
                newf = fname+"_"+str(counter)+".deap";
                counter += 1;
                fh = open(newf, 'w')
                subgraph = self.pathwaygraph.subgraph(g)
                # print subgraph.nodes();
                # print subgraph.edges(data=True);
                for e in subgraph.edges(data=True):
                    if e[0] in self.genedict and e[1] in self.genedict: # if the nodes exist in genedict. I'm only saving nodes that are genes, and not compounds. But all edges are saved, including compounds'.
                        product1 = '\t'.join( self.genedict[ e[0] ] )
                        product2 = '\t'.join( self.genedict[ e[1] ] )
                        if (e[2]['type'] == "inhibition"): # only inhibition edges are deemed False
                            fh.write(product1+"\t::"+product2+"\t::False\n")
                        else:
                            fh.write(product1+"\t::"+product2+"\t::True\n")

    def getSubgraphs( self, nnodes ): # find all connected subgraphs of size nnodes
        all_subg = []
        for path in self.connectedpaths: # for all connected components in the pathway file (path only has set of nodes, no edge information)
            g = self.pathwaygraph.subgraph(path) # get the graph information for path (nodes + edges + attributes)
            if (len(g.nodes()) >= nnodes): # only those subgraphs that have at least nnodes nodes
                # print '\n', g.nodes();
                for sub_nodes in itertools.combinations(g.nodes(),nnodes): # for all combinations of nnodes 
                    subgraph = g.subgraph(sub_nodes)
                    if NX.is_connected(subgraph):
                        # print subgraph.edges()
                        all_subg.append( subgraph )
        return all_subg

    def removeNodes( self, explabelmap ): # nodes (genes) that don't have an expression value/ have multiple expression values
        nodeswithlabel = set( explabelmap.keys() );
        withoutlabel = list( set(self.pathwaygraph.nodes()) - set(nodeswithlabel) );
        self.pathwaygraph.remove_nodes_from( withoutlabel );     

    def addExpLabelNode( self, mapping ): #mapping has entry id as key, and exp label as value. Nodes have to be unique, so add expression label for each node as node attribute 

        NX.set_node_attributes(self.pathwaygraph, 'elabel', mapping); 
        #for n in self.pathwaygraph.nodes():
        #    print n, self.pathwaygraph.node[n]['elabel'];

    def __init__( self, xml_file ): # read kegg xml file and parse it
        """
        Parameters: self, KGML file name

        Returns: nothing

        Description: Initializes self.file, self.tree, self.root
        """
        tree = ET.parse(xml_file)
        root = tree.getroot()
        self.file = xml_file
        self.tree = tree
        self.root = root
        # print xml_file, self.root.attrib['title'];
        self.genedict = {} # key is id, value is gene name and its aliases
        self.pathwaygraph = self.getpathway() # convert kgml pathway file into a graph object
        self.connectedpaths = NX.connected_components( self.pathwaygraph )
#        self.all_subg_size3 = self.getSubgraphs( 3 )
#        self.all_subg_size4 = self.getSubgraphs( 4 )
#        return self.deapFormat(xml_file)
