from _orgasm cimport *
from cpython.bytes cimport PyBytes_Size
from cpython.dict cimport PyDict_Merge
from orgasm.backtranslate._ahocorasick cimport *
from orgasm.utils.dna cimport isDNA

import sys

cdef extern from "stdlib.h":
    ctypedef unsigned int size_t
    int32_t abs(int32_t i)

cdef extern from "math.h":
    double log10(double x)
    double ceil(double x)

cdef extern from "orgasm.h":
    void freeBuffer(buffer_t *buffer)

    buffer_t *loadIndexedReads(char *indexname)

    int32_t lookForString(buffer_t *buffer, char *key, size_t length, int32_t* endoflist)
    int32_t nextRead(buffer_t *buffer, int32_t current, size_t length, int32_t* endoflist)
    int32_t lookForReadsIds(buffer_t *buffer, int32_t key, int32_t* endoflist)
    int32_t nextReadIds(buffer_t *buffer, int32_t current, int32_t* endoflist)
    char *getRead(buffer_t *buffer, int32_t recordid, uint32_t begin, int32_t length, char *dest)

cdef bint lowcomplexity(bytes s):
    """
    Returns True if the word s is an homopolymer, an homedimer
    or an homotrimer

    @param s: the string to test
    @type s: bytes

    @return True if s is an homopolymer, an homedimer
            or an homotrimer, False otherwise
    @rtype: bool
    """
    return s[0:-2] == s[2:] or s[0:-3] == s[3:]

cdef class Index:

    def __init__(self,bytes name):
        cdef int minfake

        self._index=loadIndexedReads(name)
        minfake = <int>(10**ceil(log10(self._index.readCount))) + 1
        self._fakes=FakeReads(minfake,self._index.readSize)

    cpdef bint contains(self, bytes word):
        cdef int32_t endoflist
        cdef size_t  lword = len(word)

        lookForString(self._index,word,lword,&endoflist)

        return endoflist==0

    cpdef int count(self,word):
        cdef size_t n=0
        cdef int32_t readid

        for readid in self[word]:
            n+=1

        return n

    cpdef int len(self):
        return self._index.readCount

    cpdef bytes getRead(self, int32_t readid, int32_t begin, int32_t length):
        """
        getRead methods allows to access to the the sequence of the reads stored in the
        index. The read id is include in the interval [-n,-1] U [1,n] with n the count
        of read in the index. the read 0 does not exist.
        """
        cdef char buffer[500]
        cdef bytes seq

        assert begin < self._index.readSize,"Begin must be smaller than read size"
        assert begin+length <= self._index.readSize,"Begin+Length must be smaller than read size"

        if self._fakes.isFake(readid):
            seq = self._fakes.getRead(readid,begin,length)
        elif readid==0 or abs(readid)>self._index.readCount:
            raise IndexError,readid
        else:
            seq=getRead(self._index, readid, begin, length, buffer)

        return seq

    cpdef dict getExtensions(self, bytes word):
        cdef dict rep
        cdef int32_t readid                                 # @DuplicatedSignature
        cdef size_t  lword = len(word)                      # @DuplicatedSignature
        cdef size_t  lext  = self._index.readSize - lword
        cdef bytes e
        cdef list  data

        rep = {}

        for readid in self[word]:
            # assert abs(readid) < self._index.readCount
            e = self.getRead(readid,lword, lext)
            data = rep.get(e,[0,[]])
            data[0]+=1
            data[1].append(readid)
            rep[e]=data

        return rep

    cpdef dict checkedExtensions(self, bytes probe,
                                       int minread=20,
                                       double minratio=0.1,
                                       int mincov=1,
                                       int minoverlap=40,
                                       int extlength=1,
                                       bint lowfilter=True,
                                       object restrict=None,
                                       bint exact=True):
        '''
        :param probe: The sequence used as prefix in the selection
        :type probe:  bytes
        :param minread: the minimum count of read to consider
        :type minread:  int
        :param minratio: minimum ratio between occurrences of an extension
                         and the occurrences of the most frequent extension
                         to keep it.
        :type minratio:  float [0,1]
        :param mincov:   minimum occurrences of an extension to keep it
        :type mincov:    int
        :param minoverlap: minimum length of the probe usable
        :type minoverlap: int
        '''

        cdef dict  ext       = {}
        cdef dict  lext
        cdef int   covmax    = 0
        cdef int   readcount = 0
        cdef list  e,rs,ks,rk
        cdef list  allexts
        cdef dict  extcount = {}
        cdef bytes k
        cdef int   x
        cdef int   r,rn,rsize
        cdef bytes oriprobe

        oriprobe = probe
        
        while not (lowcomplexity(probe) and lowfilter) and \
              PyBytes_Size(probe) >= minoverlap   and \
              (readcount < minread or exact):
            
            trueoverlap = len(probe)
            lext = self.getExtensions(probe)

            if restrict is not None:
                ks = lext.keys()
                for k in ks:
                    rs=lext[k][1]
                    rk=[]
                    for r in rs:
                        rn = self.getPairedRead(r)
                        if rn:
                            rn = self.getIds(rn)[0]
                            if rn in restrict:
                                rk.append(r)
                    if rk:
                        lext[k][1]=rk
                    else:
                        del lext[k]

            if lext:
                for e in lext.itervalues():
                    readcount+=e[0]
                PyDict_Merge(ext,lext,1)
                
            probe=probe[1:]
            
            
        ##-------------------END OF THE WHILE LOOP ----------------------##

        if readcount < minread:
            ext={}
        elif ext:
            allexts = [(k[0:extlength],ext[k][0]) for k in ext.iterkeys()]
            extcount= {}
            for k,x in allexts:
                extcount[k]=extcount.get(k,0)+x
                
            rsize = self.getReadSize()
            
            # Evaluate the true coverage according the actual overlap
            for k,x in extcount.iteritems():
                extcount[k]=x * rsize / (rsize - trueoverlap)

            covmax = max(extcount.values())

            for k,x in extcount.items():
                if (   x < mincov 
                    or <double>x/<double>covmax < minratio):
                    del extcount[k]
                else:
                    data = ext.get(k,(0,[]))[1]
                    if not data:
                        data=list(self.getReadIds(oriprobe+k)[2])

                    extcount[k]=(extcount[k],data)


        return extcount


    cpdef int getReadSize(self):
        return self._index.readSize

    cpdef tuple getReadIds(self, bytes seq):

        cdef list x
        cdef int  y
        cdef int  t

        assert len(seq)==self._index.readSize,"seq must have the same length than reads"

        x = [y for y in self[seq]]

        if x:
            t = len(self)+1

            for y in x:
                if abs(y) < abs(t):
                    t = y
            y = len(x)
        else:
            f = self._fakes.getReadIds(seq)
            return f 

        return (t,y,set(x))

    cpdef int32_t getPairedRead(self, int32_t id):
        cdef int32_t iid = abs(id)
        cdef int32_t pid

        if iid > self._index.readCount:
            pid = 0
        else:
            pid = self._index.order2[iid-1] + 1

        if id > 0:
            pid = - pid

        return pid

    cpdef tuple normalizedPairedEndsReads(self, int32_t id):
        '''
        Returns for a read id the normalized set of paired reads corresponting to it.

        :param id: the read id
        @return: a tuple of two elements
                    - an integer value : the normalized id
                    - a list of integers : the normalized ids of the paired reads
        '''
        cdef int32_t nid
        cdef int32_t n
        cdef set synomymes
        cdef int32_t i
        cdef list paired
        cdef tuple k

        nid,n,synomymes = self.getIds(id)
        if id not in synomymes:
            nid=-nid
            synomymes=set([-i for i in synomymes])
        paired = [self.getPairedRead(i) for i in synomymes]
        paired = [(j,self.getIds(j)) for j in paired]
        paired = [k[0] if j in k[2] else -k[0] for j,k in paired]
        return nid,paired


    cpdef lookForSeeds(self, dict sequences, int kup=-1, int mincov=1):

        cdef AhoCorasick patterns
        cdef dict matches
        cdef bytes k                                                # @DuplicatedSignature

        nuc = all([isDNA(sequences[k]) for k in sequences])
            
        if nuc:
            patterns = NucAhoCorasick()
            kup = 12 if kup < 0 else kup
        else:
            patterns = ProtAhoCorasick()
            kup = 4 if kup < 0 else kup

        for k in sequences:
            patterns.addSequence(sequences[k],k,kup)

        patterns.finalize()
        
        minmatch = 50 if nuc else 15

        matches = patterns.scanIndex(self,15,-1,mincov)

        return matches



    cpdef tuple getIds(self, int32_t id):
        cdef int32_t endoflist                                      # @DuplicatedSignature
        cdef int32_t readid                                         # @DuplicatedSignature
        cdef list    reads = []
        cdef int  y                                                 # @DuplicatedSignature
        cdef int  x
        cdef int  t                                                 # @DuplicatedSignature

        if id==0:
            raise IndexError,readid

        if self._fakes.isFake(id):
            return self._fakes.getIds(id)

        y = abs(id)
        if y > self._index.readCount:
                return (0,0,set())

        readid = lookForReadsIds(self._index,id,&endoflist)

        while(endoflist==0):
            reads.append(readid)
            readid = nextReadIds(self._index,readid,&endoflist)

        if reads:
            t = len(self)+1

            for y in reads:
                if abs(y) < abs(t):
                    t = y

            if t < 0:
                reads = [-x for x in reads]
                t = -t

        return (t,len(reads),set(reads))


    def __dealloc__(self):
        freeBuffer(self._index)

    def __contains__(self, bytes word):
        return self.contains(word)

    def __getitem__(self, bytes word):
        cdef int32_t endoflist                                  # @DuplicatedSignature
        cdef int32_t readid                                     # @DuplicatedSignature
        cdef size_t  lword = len(word)                          # @DuplicatedSignature

        readid = lookForString(self._index,word,lword,&endoflist)

        while(endoflist==0):
            yield readid
            readid = nextRead(self._index,readid,lword,&endoflist)

    def __len__(self):
        return self.len()

    def checkReverseOrder(self):
        conv = {'A':'0','C':'1','T':'2','G':'3'}
        from orgasm.utils.dna import reverseComplement
        olds=" "
        for i in range(len(self)):
            j = self._index.order1[i] + 1
            s = "".join(conv[x] for x in reverseComplement(self.getRead(j,0,self.getReadSize())))
            if olds > s:
                print i,olds,s
                raise AssertionError
            else:
                olds=s



    property fakes:

        "An direct access to the fake reads."

        def __get__(self):
            return self._fakes









