#!/usr/bin/env python

# Dénes Türei EMBL 2017
# turei.denes@gmail.com

from future.utils import iteritems
from past.builtins import xrange, range
import collections
import copy

import pypath

# you can add more patterns here
# this was the one on your sketch:
setMotifPatterns = set([
    (-1, 1, -1, 1)
])

fnOut = 'loops.tab'

depth = 4

pa = pypath.PyPath()
pa.load_omnipath()

adjDict = collections.defaultdict(lambda: set([]))

for e in pa.graph.es:
    
    d = e['dirs']
    if d.is_stimulation(d.straight):
        adjDict[e.source].add((pa.p(d.straight[1]).index, 1))
    if d.is_inhibition(d.straight):
        adjDict[e.source].add((pa.p(d.straight[1]).index, -1))
    if d.is_stimulation(d.reverse):
        adjDict[e.target].add((pa.p(d.straight[0]).index, 1))
    if d.is_inhibition(d.reverse):
        adjDict[e.target].add((pa.p(d.straight[0]).index, -1))

def get_pattern(motif):
    
    return tuple(i[1] for i in motif[1:])

def motif_generator(adjDict, depth = 4):
    
    def step(adjDict, motif, depth):
        
        if depth == 0:
            
            yield motif
        
        elif motif[-1][0] in adjDict:
            
            for nxt in adjDict[motif[-1][0]]:
                
                motif1 = copy.copy(motif)
                motif1.append(nxt)
                
                for motif0 in step(adjDict, motif1, depth - 1):
                    
                    yield motif0
        
        else:
            
            return
    
    prg = pypath.progress.Progress(len(adjDict), 'Looking up paths', 1)
    
    for start in adjDict.keys():
        
        prg.step()
        
        for motif in step(adjDict, [(start, None)], depth):
            
            if (
                motif[0][0] == motif[depth][0] and
                not any(i[0] == motif[0][0] for i in motif[1:depth]) and
                get_pattern(motif) in setMotifPatterns
            ):
                
                yield motif
    
    prg.terminate()

def to_file(fnOut, pa, adjDict, depth = 4):
    
    with open(fnOut, 'w') as fp:
        
        for motif in motif_generator(adjDict, depth):
            
            fp.write(
                '%s\t%s\n' % (
                    pa.nodNam[motif[0][0]],
                    '\t'.join('%s\t%s' % (
                            '+' if motif[i][1] > 0 else '-',
                            pa.nodNam[motif[i][0]]
                        )
                        for i in xrange(1, len(motif))
                    )
                )
            )

if __name__ == '__main__':
    
    to_file(fnOut, pa, adjDict, depth)