'''
Created on 28 sept. 2014

@author: coissac
'''

import orgasm.samples

from orgasm import getIndex, getProbes, getSeeds
from orgasm.tango import matchtoseed, tango, cutLowCoverage, cutSNPs,\
    estimateDeadBrancheLength, coverageEstimate, estimateFragmentLength,\
    genesincontig, scaffold, fillGaps, dumpGraph
from orgasm.assembler import Assembler
import sys


__title__="Build the initial assembling graph"


default_config = {   'minread'       : None,
                     'coverage'      : None,
                     'minratio'      : None,
                     'mincov'        : 1,
                     'minoverlap'    : 50,
                     'smallbranches' : None,
                     'lowcomplexity' : False,
                     'snp'           : False,
                 }

def addOptions(parser):
    parser.add_argument(dest='orgasm:indexfilename',  metavar='index', 
                        help='index root filename (produced by the orgasmi command)')
    
    parser.add_argument(dest='orgasm:outputfilename',     metavar='output', 
                                                          nargs='?', 
                                                          default=None,
                        help='output prefix' )
    
    
    
    parser.add_argument('--minread',          dest='buildgraph:minread', 
                                              type=int, 
                                              action='store', 
                                              default=None, 
                        help='the minimum count of read to consider [default: <estimated>]')
    
    parser.add_argument('--coverage',         dest='buildgraph:coverage', 
                                              type=int, 
                                              action='store', 
                                              default=None, 
                        help='the expected sequencing coverage [default: <estimated>]')
    
    parser.add_argument('--minratio',         dest='buildgraph:minratio', 
                                              type=float, action='store', 
                                              default=None, 
                        help='minimum ratio between occurrences of an extension'
                             ' and the occurrences of the most frequent extension '
                             'to keep it. [default: <estimated>]')
    
    parser.add_argument('--mincov',           dest='buildgraph:mincov', 
                                              type=int, 
                                              action='store', 
                                              default=1, 
                        help='minimum occurrences of an extension to '
                             'keep it. [default: %(default)d]')
    
    parser.add_argument('--minoverlap',       dest='buildgraph:minoverlap', 
                                              type=int, 
                                              action='store', 
                                              default=None, 
                        help='minimum length of the overlap between '
                             'the sequence and reads to participate in '
                             'the extension. [default: <estimated>]')
    
    parser.add_argument('--smallbranches',    dest='buildgraph:smallbranches', 
                                              type=int, 
                                              action='store', 
                                              default=None, 
                        help='maximum length of the branches to cut during '
                             'the cleaning process [default: <estimated>]')
    
    parser.add_argument('--lowcomplexity',    dest='buildgraph:lowcomplexity', 
                                              action='store_true', 
                                              default=False, 
                        help='Use also low complexity probes')
    
    parser.add_argument('--back',             dest='orgasm:back', 
                                              type=int, 
                                              action='store', 
                                              default=None, 
                        help='the number of bases taken at the end of '
                             'contigs to jump with pared-ends [default: <estimated>]')
    
    parser.add_argument('--snp',              dest='buildgraph:snp', 
                                              action='store_true', 
                                              default=False, 
                        help='desactivate the SNP clearing mode')
    
    parser.add_argument('--seeds',            dest ='orgasm:seeds', 
                                              metavar='seeds', 
                                              default=None, 
                                              type=str, 
                                              required=True,
                        help='protein seeds; either a file containing '
                        'seeds proteic sequences or internal set of seeds '
                        'among %s' % str(filter(lambda s: s.startswith('prot') or 
                                                s.startswith('nuc'),dir(orgasm.samples))) )


def estimateMinRead(index,minoverlap,coverage):
    MINREAD=10
    MINREADCOR=3
    MINOVERLAP=50
    minread =  (index.getReadSize() - minoverlap) * coverage / index.getReadSize()  / MINREADCOR
    if minread < MINREAD:
        minoverlap = index.getReadSize() - (MINREAD * MINREADCOR * index.getReadSize() / coverage)
        minread = MINREAD
    if  minoverlap< MINOVERLAP:
        minread =  MINREAD
        minoverlap = MINOVERLAP
    return minread,minoverlap


def run(config):
    
    logger=config['orgasm']['logger']
    output = config['orgasm']['outputfilename'] 
    lowfilter=not config['buildgraph']['lowcomplexity']
    coverageset=config['buildgraph']['coverage'] is not None
    snp=config['buildgraph']['snp']

    r = getIndex(config)
    p = getProbes(config)
    x = getSeeds(r,p,config)            

    # First estimation of the coverage based on protein matches
    
    common    = set(p.keys()) & set(x.keys())
    pcoverage = [len(x[i]) * r.getReadSize()  / len(p[i]) / 3 for i in common]
    pcoverage.sort()

    logger.info("Probe coverage : %s" % str(pcoverage))
    
    if coverageset:
        coverage = config['buildgraph']['coverage']
    else:
        coverage = pcoverage[-1]
    
    # according to the minread option estimate it from coverage or use the specified value
   
    if config['buildgraph']['minread'] is None:
        minread,minoverlap = estimateMinRead(r,
                                             config['buildgraph']['minoverlap'],
                                             coverage)
    else:
        minread = config['buildgraph']['minread']
        minoverlap = config['buildgraph']['minoverlap']

    logger.info("\nestimated coverage  : %d based on probe match (minread: %d)" % (coverage,minread))
    
        
        
    # Convert matches in sedd list    
    s = matchtoseed(x,r)
    
    # Create the assembler object
    asm = Assembler(r)

    ##########################
    #
    # If minread is not specified we initiate the assembling
    # based on the coverage estimated from protein match
    # to obtain a better coverage estimation
    #
    ##########################
    if config['buildgraph']['minread'] is None and not coverageset:

        # Run the first assembling pass
        a = tango(asm,s,mincov=minread,       # @UnusedVariable
                        minread=minread,
                        minoverlap=minoverlap,
                        lowfilter=lowfilter,
                        maxjump=0, 
                        cycle=1,
                        nodeLimit=30000)
    
        # Clean small unsuccessful extensions
        asm.cleanDeadBranches(maxlength=10)
        
        # and too low covered assembling
        if coverageset:
            cutLowCoverage(asm,int(coverage),terminal=True)
        else:
            cutLowCoverage(asm,int(coverage/4),terminal=True)
            
    
        if snp:
            cutSNPs(asm)
        
        if config['buildgraph']['smallbranches'] is not None:
            smallbranches = config['buildgraph']['smallbranches']
        else:
            smallbranches = estimateDeadBrancheLength(asm)
            logger.info("Dead branch length setup to : %d bp" % smallbranches)
        
        asm.cleanDeadBranches(maxlength=smallbranches)
        
        if len(asm) > 0:
            score,length,ecoverage = coverageEstimate(asm,x,r)  # @UnusedVariable
            if not coverageset:
                coverage = ecoverage    
            minread,minoverlap = estimateMinRead(r,config['buildgraph']['minoverlap'],coverage) 
            logger.info("coverage estimated : %d based on %d bp (minread: %d)" %(coverage,length,minread/2))
        
            
        # Reinitiate the assembling for running with the estimated parameters
        # Convert matches in sedd list    
        s = matchtoseed(x,r)
    
        # Create the assembler object
        asm = Assembler(r)
    
    #############################################
    #
    # We now run the main assembling process
    #
    #############################################

    # Run the first assembling pass
    a = tango(asm,s,mincov=coverage/4,       #@UnusedVariable
                    minread=minread,
                    minoverlap=minoverlap,
                    lowfilter=lowfilter,
                    maxjump=0, 
                    cycle=1)
    
    # Clean small unsuccessful extensions
    asm.cleanDeadBranches(maxlength=10)
    
    # and too low covered assembling
    if coverageset:
        cutLowCoverage(asm,int(coverage),terminal=True)
    else:
        cutLowCoverage(asm,int(coverage/4),terminal=True)
        
    
    # cleanup snp bubble in the graph    
    if snp:
        cutSNPs(asm)
    
    if config['buildgraph']['smallbranches'] is not None:
        smallbranches = config['buildgraph']['smallbranches']
    else:
        smallbranches = estimateDeadBrancheLength(asm)
        logger.info("     Dead branch length setup to : %d bp" % smallbranches)

    asm.cleanDeadBranches(maxlength=smallbranches)


    # reestimate coverage
    score,length,ecoverage = coverageEstimate(asm,x,r)  # @UnusedVariable
    
    if not coverageset:
        coverage = ecoverage  
    
    if coverage < 30:
        sys.exit()

    if config['buildgraph']['minread'] is None:
        minread,minoverlap = estimateMinRead(r,config['buildgraph']['minoverlap'],coverage)
        minread/=4
        if minread<5:
            minread=5

    logger.info("coverage estimated : %d based on %d bp (minread: %d)" %(coverage,length,minread))
        
    meanlength,sdlength = estimateFragmentLength(asm)
    
    if config['orgasm']['back'] is not None:
        back = config['orgasm']['back']
    elif config['orgasm']['back'] is None and meanlength is not None:
        back = int(meanlength + 4 * sdlength)
        if back > 500:
            back=500
    else:
        back = 300
        
    if meanlength is not None:
        logger.info("Fragment length estimated : %f pb (sd: %f)" % (meanlength,sdlength))

    cg = asm.compactAssembling(verbose=False)
    genesincontig(cg,r,x)
    scaffold(asm,cg,minlink=5,back=back,addConnectedLink=False)
    print >>open(output+'.intermediate.gml','w'),cg.gml()
    
    ###################################################
    #
    # We now fill the gaps between the contigs
    #
    ###################################################
    
    delta = 1
       
    # Run the fill gap procedure    
    while  delta > 0 or delta < -2000 :
                   
        delta = fillGaps(asm,back=back,
                       minread=minread,
                       maxjump=0,
                       minoverlap=minoverlap,
                       cmincov=2,
                       emincov=int(coverage/4),
                       gmincov=int(coverage/4),
                       lowfilter=lowfilter,
                       snp=snp)

        print >>sys.stderr
        print >>sys.stderr
        print >>sys.stderr
        print >>sys.stderr,"=================================================================="      
        print >>sys.stderr 
        
        cg = asm.compactAssembling(verbose=False)
        genesincontig(cg,r,x)
        scaffold(asm,cg,minlink=5,back=back,addConnectedLink=False)
        print >>open(output+'.intermediate.gml','w'),cg.gml()
        
        if meanlength is None:
            meanlength,sdlength = estimateFragmentLength(asm)
            if config['orgasm']['back'] is None and meanlength is not None:
                logger.info("Fragment length estimated : %f pb (sd: %f)" % (meanlength,sdlength))
                back = int(meanlength + 4 * sdlength)  
                if back > 500:
                    back=500
                         
        print >>sys.stderr
        print >>sys.stderr,"=================================================================="      
        print >>sys.stderr
        
    ###################################################
    #
    # Finishing of the assembling
    #
    ###################################################

    if snp:
        logger.info("Clean polymorphisms")
        cutSNPs(asm)
        
    asi = len(asm)+1
    
    logger.info("Clean dead branches")
    while (asi>len(asm)):
        asi=len(asm)
        smallbranches = estimateDeadBrancheLength(asm)
        logger.info("     Dead branch length setup to : %d bp" % smallbranches)
        asm.cleanDeadBranches(maxlength=smallbranches)
        
    cg = asm.compactAssembling(verbose=False)
    score,length,ecoverage = coverageEstimate(asm,x,r)  # @UnusedVariable
    if not coverageset:
        coverage=ecoverage

    if snp:
        logger.info("Clean polymorphisms phase 2")
        cutSNPs(asm)
        
    logger.info("Clean low coverage terminal branches")
    if coverageset:
        cutLowCoverage(asm,int(coverage),terminal=False)
    else:
        cutLowCoverage(asm,int(coverage/2),terminal=True)
        logger.info("Clean low coverage internal branches")
        cutLowCoverage(asm,int(coverage/3),terminal=False)
        
        
    cg = asm.compactAssembling(verbose=False)     
        
    logger.info("Scaffold the assembly")
    scaffold(asm,cg,minlink=5,back=int(back),addConnectedLink=False)
    genesincontig(cg,r,x)
    print >>open(output+'.gml','w'),cg.gml()
    dumpGraph(output+'.oax',asm)
