from orgasm.indexer._orgasm cimport *
from _asmbgraph cimport *
from orgasm.graph._graphmultiedge cimport *

from orgasm.utils._progress cimport progressBar
import math

import sys

cdef tuple findDeepestRead(Index index, bytes seed):
    cdef list frg
    cdef list fcount
    cdef int mfc
    
    frg   = [seed[x:(x+index.getReadSize())] for x in xrange(len(seed)-index.getReadSize())]
    fcount= [index.count(x) for x in frg]
    mfc = max(fcount)
    return frg[fcount.index(mfc)],mfc



cpdef cmpPoints(p1,p2):
    '''
    Compare two extensions point.
    INTERNAL functio uses for ordering extension point:
    the smallest point is point closest to the initial extension point
    with the highest coverage.
    
    :param p1:
    :type p1:
    :param p2:
    :type p2:
    '''
    
    rep = cmp(p2[0],p1[0])
    if not rep:
        rep= cmp(p1[4],p2[4])
    return -rep

def cmpRead(r1,r2):
    return cmp(r1[1],r2[1])


cdef int32_t deleteBranch(AsmbGraph graph,list path, int32_t maxlength, int32_t deleted=0):
    
    cdef int32_t   node
    cdef list      parents
    cdef list      sons
    cdef int32_t   i
    
    if len(path) > maxlength:
        return deleted
        

    node = path[-1]
    parents = [y for y in graph.parentIterator(node)]
    if parents :
        for p in parents:                
            sons    = [y for y in graph.neighbourIterator(p)]

            path+=[p]

            if len(sons) < 2:
                deleted = deleteBranch(graph,path,maxlength,deleted)
            else:
                for i in range(len(path)-1,0,-1):
                    try:
                        graph.deleteEdge(path[i],path[i-1])
                        deleted+=1
                    except KeyError:
                        pass
    else:
        for i in range(len(path)-1,0,-1):
            try:
                graph.deleteEdge(path[i],path[i-1])
                deleted+=1
            except KeyError:
                pass
        
    return deleted

cdef class Assembler:
    '''
    totolitoto
    '''

    cdef AsmbGraph _graph
    cdef Index _index
    cdef int _overlap
    cdef set extensionPoints 
    cdef int _startread
    cdef int _finalread
    cdef list _seeds
    cdef dict _annotations

    def __init__(self,Index index, int overlap=90):
        '''
        Create a new assembler object.
        
        :param index: the file name of an read index previously formated using orgasmi binary
        :type index: bytes
        :param seed: a DNA sequenced used as started for the assembly
        :type seed: bytes
        :param overlap: the minimum overlap between two reads
        :type overlap: int
        '''
        cdef bytes starter
        cdef int depth 
        cdef tuple reads
        cdef dict attr
                
        self._graph = AsmbGraph(b"assembling",index)
        
        self._index = index
        self._overlap=overlap
        self._seeds = []
        self._annotations = {}
                        
        
        
    cpdef tuple readType(self,ids):
        '''
        Internal function : Given a set of read ids, return the one that have to be used
        as standard id. 
         
        The set of ids given to this function corresponds to a set of all strictly identical reads.
          
        :param ids: an iterable elements contining read ids.
        :type ids: iterable
         
        :returns: a tuble of three elements
         
                    - the standard id to use
                    - the length of ids
                    - a set containing all unique ids given as parameter
                     
        :rtype: tuple
        '''
         
        cdef int t 
        cdef int y 
         
        if not ids:
            return None
         
        t = len(self._index)+1
         
        for y in ids:
            if abs(y) < abs(t):
                t = y
                 
        return (t,len(ids),set(ids))
    
    def isEnd(self, int node):
        cdef list s
        cdef list f
        s = [x for x in self.graph.neighbourIterator(node)]
        f = [x for x in self.graph.parentIterator(node)]
        if len(s)==0 and len(f)==1:
            return f[0]
        return 0

    def forkIterator(self):
        
        def has_several_sons(int node):
            return  len(set(self._graph.neighbourIterator(node))) > 1
        
        return self._graph.nodeIterator(has_several_sons)
    
        

           
    def endNodeSet(self, set excluded=None, bint alllink=False):
        if alllink:
            def isfinalnode(int node):
                return len(set(self._graph.neighbourIterator(node)))==0
        else:
            def isfinalnode(int node):                  # @DuplicatedSignature
                return len(set(self._graph.neighbourIterator(node)))==0
        if excluded is None:
            excluded=set()
        return set(x for x in self._graph.nodeIterator(isfinalnode) if x not in excluded)

    def startNodeSet(self, set excluded=None, bint alllink=False):
        if alllink:
            def isstartnode(int node):
                return len(set(self._graph.parentIterator(node)))==0
        else:
            def isstartnode(int node):                  # @DuplicatedSignature
                return len(set(self._graph.parentIterator(node)))==0
            
        if excluded is None:
            excluded=set()
        return set(x for x in self._graph.nodeIterator(isstartnode) if x not in excluded)

    def cleanDeadBranches(self,int32_t maxlength=10, bint alllink=False):

        cdef AsmbGraph graph
        cdef int32_t   d
        cdef int32_t   cd
        cdef int32_t   i                                # @DuplicatedSignature
        cdef set       ep
        cdef set       sp
                    
        graph=self.graph

        cd=0    
        d=1
            
        print >>sys.stderr
        
        while d > 0:
            d=0
            for i in self.endNodeSet(alllink=alllink):
                print >>sys.stderr,"Remaining edges : %d node : %d\r" % (graph.edgeCount(),len(graph)),
                d+=deleteBranch(graph,[i],maxlength)
            cd+=d
            
        
        ep = self.endNodeSet(alllink=alllink)
        sp = self.startNodeSet(alllink=alllink)
        
        for i in ep & sp:
            try:
                graph.deleteNode(i)
            except KeyError:
                pass
            
        print >>sys.stderr,"Remaining edges : %d node : %d\r" % (graph.edgeCount(),len(graph))

        return cd            
                
    def __len__(self):
        return self._graph.nodeCount()
    
    def stemIterator(self, bint alllink=False):
        cdef NodeIterator ni
        cdef int first
        cdef int n
        cdef set branches
        cdef set brothers
        cdef set fathers
        cdef list lsequence
        cdef int last
        cdef int lseq
        cdef bytes sequence
        cdef list path
        
        if (alllink==True):
            def sons(int node):
                return set(self._graph.neighbourIterator(node))
                
            def parents(int node):
                return set(self._graph.parentIterator(node))
        else:
            def sons(int node):                           # @DuplicatedSignature
                return set(self._graph.neighbourIterator(node))
                
            def parents(int node):                        # @DuplicatedSignature
                return set(self._graph.parentIterator(node))

        def is_junction(int node):
            return  len(sons(node)) > 1 or len(parents(node)) > 1 or len(parents(node))==0
        
        ni=self._graph.nodeIterator(is_junction)

        for first in ni:
            branches=sons(first)
            for n in branches:
                path = [first,n]                               
                brothers=sons(n)
                fathers=parents(n)
                lsequence=[self._graph.getEdgeAttr(first,n)[b'ext']]
                while len(brothers)==1 and len(fathers)==1:
                    son=brothers.pop()
                    lsequence.append(self._graph.getEdgeAttr(n,son)[b'ext'])
                    path.append(son)
                    n=son
                    brothers=sons(n)
                    fathers=parents(n)
    
                last = n
                sequence=b''.join(lsequence)
                lseq = len(sequence)
                
                yield first,last,lseq,sequence,path

        
    def compactAssembling(self, bint verbose=True):
        cdef DiGraph graph = DiGraphMultiEdge('compact')
        cdef int n=0
        cdef int x
        cdef int i
        cdef lcontig=0
        cdef int first 
        cdef int last                                # @DuplicatedSignature
        cdef int lseq                                # @DuplicatedSignature
        cdef int lgraph
        cdef int lpath
        cdef bytes sequence                          # @DuplicatedSignature
        cdef dict attr                               # @DuplicatedSignature
        cdef double weight
        cdef double minweight
        cdef int readSize = self._index.getReadSize()
        cdef dict edgeName={}

        def sons(int node):                          # @DuplicatedSignature
            '''
            Build a set composed of the following reads  
            '''
            return set(self._graph.neighbourIterator(node))
            
        def parents(int node):                       # @DuplicatedSignature
            '''
            Build a set composed of the previous reads  
            '''
            return set(self._graph.parentIterator(node))

        def is_junction(int node):                   # @DuplicatedSignature
            '''
            Predicate testing if a node is a fork or the beginning of the assembly
            '''
            return  len(sons(node)) > 1 or len(parents(node)) > 1 or len(parents(node))==0

        lgraph=len(self._graph)

        if verbose:
            print >>sys.stderr,b"Compacting graph :"
        else:
            progressBar(1,lgraph,True,b"Compacting graph")

        minweight=1000000.
        
        
        for first,last,lseq,sequence,path in self.stemIterator():
            cpath = tuple(-path[-i] for i in range(1,len(path)+1))

            # Define edgeName
            stemid = edgeName.get(tuple(path),0)
            if stemid==0:
                stemid = -edgeName.get(cpath,0)
                if stemid==0:
                    n+=1
                    stemid=n 
                    edgeName[tuple(path)]=stemid
                    
            lcontig+=lseq
            graph.addEdge(first,last)
            attr = graph.getEdgeAttr(first,last)
            attr[b'sequence']=sequence
            attr[b'stemid']=stemid
            attr[b'path']=path
            attr[b'length']=lseq
            weight=sum(self.graph[x][b'coverage'] for x in path if b'coverage' in self.graph[x]) / float(len(path))
            attr[b'weight']= int(weight * readSize)
            attr[b'graphics']={b'width':int(weight)+1,
                               b'arrow':b'last'}
                        
            if weight < minweight:
                minweight=weight

            if verbose:
                print >>sys.stderr,b" Stem   %9d : %6d bp (total : %6d) coverage : %6.2f" % (stemid,lseq,lcontig,weight * readSize)
            else:
                progressBar(lcontig,lgraph,False,b"Compacting graph")
                
        if lcontig + 2 * n < len(self):
            ccs = self._graph.connectedComponentIterator()
            
            for cc in ccs:
                circle = reduce(lambda a,b: a and b, (not is_junction(x) for x in cc))

                if circle:
                    minabs = min(abs(i) for i in cc)
                    if minabs in cc:
                        first = minabs
                    else:
                        first = sons(-minabs).pop()
                    nn = first
                    path=[first]
                    lsequence=[]
                    son=sons(nn).pop()
                    while first != son:
                        path.append(son)
                        lsequence.append(self._graph.getEdgeAttr(nn,son)[b'ext'])
                        nn=son
                        son=sons(nn).pop()
    
                    last = first
                    sequence=b''.join(lsequence)
                    lseq = len(sequence)
                                      
                    stemid = edgeName.get(tuple(path),0)
                    
                    if stemid==0:
                        cpath = list(path)
                        cpath.reverse()
                        cpath=tuple(-i for i in cpath)
                        stemid = -edgeName.get(cpath,0)
                        edgeName[tuple(path)]=stemid
                        if stemid==0:
                            n+=1
                            stemid=n 
                            edgeName[tuple(path)]=stemid
                            
                        lcontig+=lseq
                        graph.addEdge(first,first)
                        attr = graph.getEdgeAttr(first,last)
                        attr[b'sequence']=sequence
                        attr[b'stemid']=stemid
                        attr[b'path']=path
                        attr[b'length']=lseq
                        weight=sum(self.graph[x][b'coverage'] for x in path if b'coverage' in self.graph[x]) / float(len(path))
                        attr[b'weight']= weight * readSize
                        attr[b'graphics']={b'width':int(weight)+1,
                                           b'arrow':b'last',
                                           b'fill':"#0000FF"}
                                    
                        if weight < minweight:
                            minweight=weight
        
        
                        print >>sys.stderr,b" Circle %9d : %6d bp (total : %6d) coverage : %6.2f" % (stemid,lseq,lcontig,weight * readSize)
            
                    
        # We are normalizing the edge id
        es = [e for e in graph.edgeIterator(edgePredicate = lambda i: graph.getEdgeAttr(i[0],i[1],i[2])['stemid']>0)]
        esc = [(graph.getEdgeAttr(e[0],e[1],e[2])['stemid'],e) 
               for e in graph.edgeIterator(edgePredicate = lambda i: graph.getEdgeAttr(i[0],i[1],i[2])['stemid']<0)]
        esc = dict(esc)
        es.sort(key=lambda e:e[0]*10+e[2])
        stemid = 1
        for e in es:
            ea = graph.getEdgeAttr(e[0],e[1],e[2])
            ei = ea['stemid']
            ea['stemid']=stemid
            lseq = ea['length']
            sequence = ea['sequence']
            if ea['length'] > 10:
                ea[b'label']=b"%d : %s->(%d)->%s  [%d]" % (stemid,sequence[0:5],lseq,sequence[-5:],int(ea[b'weight']))
            else:
                ea[b'label']=b"%d : %s->(%d)  [%d]" % (stemid,sequence,lseq,int(ea[b'weight']))
            e = esc[-ei]
            ea = graph.getEdgeAttr(e[0],e[1],e[2])
            ea['stemid']=-stemid
            lseq = ea['length']
            sequence = ea['sequence']
            if ea['length'] > 10:
                ea[b'label']=b"%d : %s->(%d)->%s  [%d]" % (-stemid,sequence[0:5],lseq,sequence[-5:],int(ea[b'weight']))
            else:
                ea[b'label']=b"%d : %s->(%d)  [%d]" % (-stemid,sequence,lseq,int(ea[b'weight']))
            stemid+=1
            
             
        print >>sys.stderr,b"\nMinimum stem coverage = %d" % int(minweight * readSize)
        return graph
    
    property graph:

        "A doc string can go here."

        def __get__(self):
            return self._graph


    property seeds:

        "A doc string can go here."

        def __get__(self):                              # @DuplicatedSignature
            return self._seeds


    property index:

        "A doc string can go here."

        def __get__(self):                              # @DuplicatedSignature
            return self._index



                       
