'''
Created on 28 sept. 2014

@author: coissac
'''

from subprocess import Popen
from tempfile   import mkdtemp
from tempfile   import mktemp
from contextlib import nested

import atexit
import os
import sys
import re

from orgasm.utils.dna import reverseComplement  # @UnresolvedImport
import zlib
from collections import Counter


__title__="Index a set of reads"

default_config = { 'single'   : False,
                   'forward'  : None,
                   'reverse'  : None,
                   'maxread'  : None,
                   'length'   : None,
                   'estimate' : None,
                   'fasta'    : False,
                   'ffasta'   : False,
                   'rfasta'   : False
                 }

def addOptions(parser):
    parser.add_argument(dest='orgasm:indexfilename',  metavar='index', 
                        help='name of the produced index')

    parser.add_argument(dest='index:forward',  metavar='forward', 
                        help='Filename of the forward reads')
    
    parser.add_argument(dest='index:reverse',     metavar='reverse', 
                                                          nargs='?', 
                                                          default=None,
                        help='Filename of the reverse reads' )
    
    
    parser.add_argument('--single',           dest='index:single', 
                                              action='store_true', 
                                              default=None, 
                        help='Single read mode')

    parser.add_argument('--max-read',         dest='index:maxread', 
                                              type=int, 
                                              action='store', 
                                              default=None, 
                        help='the count of million of reads to '
                             'index (default the full file)')
    
    parser.add_argument('--length',           dest='index:length', 
                                              type=int, 
                                              action='store', 
                                              default=None, 
                        help='The length of the read to index '
                             ' (default indexed length is estimated from the first read)')
    
    parser.add_argument('--estimate-length',  dest='index:estimate', 
                                              metavar='FRACTION',
                                              type=float, 
                                              action='store', 
                                              default=None, 
                        help='Estimate the length to index for conserving FRACTION '
                             'of the overall dataset')
    
    parser.add_argument('--fasta',            dest='index:fasta', 
                                              action='store_true', 
                                              default=None, 
                        help='sequence files are fasta formated')

    parser.add_argument('--forward-fasta',            dest='index:fasta', 
                                              action='store_true', 
                                              default=None, 
                        help='forward file is a fasta file')

    parser.add_argument('--reverse-fasta',            dest='index:fasta', 
                                              action='store_true', 
                                              default=None, 
                        help='forward file is a fasta file')
tmpdir = []

# @atexit.register
# def cleanup():
#     try:
#         os.unlink(FIFO)
#     except:
#         pass


def postponerm(filename):
    def trigger():
        print >>sys.stderr,"Deleting tmp file : %s" % filename
        try:
            os.unlink(filename)
        except:
            pass
        
    return trigger

def getTmpDir():
    global tmpdir

    if not tmpdir:
        tmpdir.append(mkdtemp())
        atexit.register(postponerm(tmpdir[0]))
        
    return tmpdir[0]


def ungzip(filename):
    tmpdir = getTmpDir()
    fifo = mktemp(prefix="unziped-", dir=tmpdir)
    os.mkfifo(fifo)
    atexit.register(postponerm(fifo))
    command="gzip -d -c %s > %s" % (filename,fifo)
    process = Popen(command,                                            # @UnusedVariable
                    shell=True,
                    stderr=open('/dev/null','w')) 
    
    return fifo

def unbzip(filename):
    tmpdir = getTmpDir()
    fifo = mktemp(prefix="unziped-", dir=tmpdir)
    os.mkfifo(fifo)
    atexit.register(postponerm(fifo))
    command="bzip2 -d -c %s > %s" % (filename,fifo)
    process = Popen(command,                                            # @UnusedVariable
                    shell=True,
                    stderr=open('/dev/null','w')) 
    
    return fifo
    

    
def readFasta(filename):
    
    if isinstance(filename, str):
        filename = open(filename)
    
    seq = []    
    seqid=0
    for line in filename:
        if line[0]==">":
            if seqid > 0:
                seq = ''.join(seq)
                yield (seqid,seq)
            seqid+=1
            seq=[]
        else:
            seq.append(line.strip().upper())
            
    seq = ''.join(seq)
    yield (seqid,seq)
           
def readFastq(filename):
    
    if isinstance(filename, str):
        filename = open(filename)
        
    seq   = []
    seqid = 0
    lseq  = 0
    
    for line in filename:
        if line[0]=="@":
            if seqid > 0:
                seq = ''.join(seq)
                yield (seqid,seq)
            seqid+=1
            lseq = 0
            seq=[]
        elif line[0]!='+':
            seq.append(line.strip().upper())
            lseq+=1
        else:
            for i in range(lseq):  # @UnusedVariable
                filename.next()
        
    seq = ''.join(seq)
    yield (seqid,seq)
    

def acgtCut(seqs):
    acgt = re.compile('[^ACGT]')
    
    for s in seqs:
        yield (s[0],acgt.split(s[1],1)[0])
        
def lengthStats(seqs):
    
    store=[]
    
    def seqIterator():
        for s in store:
            yield (s[0],zlib.decompress(s[1]))
            
    def reader():
        for s in seqs:
            store.append((s[0],zlib.compress(s[1])))
            yield len(s[1])
            
    stats = Counter(reader())
    
    return stats,seqIterator()
            
def formatFastq(seq):
    
    fastq = b'@s{id:0>7}\n{seq}\n+\n{qual}'
    return fastq.format(id=seq[0],
                        seq=seq[1],
                        qual='0'*len(seq[1]))
               
        
def singleToPairedEnd(seqs,length):
    
    for s in seqs:
        f = s[1][0:length]
        r = reverseComplement(s[1][-length:])
        yield ((s[0],f),(s[0],r))
        
def doubleToPairedEnd(forward,reverse):

    for f in forward:
        r = reverse.next()
        assert f[0]==r[0]
        yield ((f[0],f[1]),(f[0],r[1]))
        
def quantiles(stats,quantile):
    ls = stats.keys()
    ls.sort(reverse=True)    
    sums=[]
    s=0
    for l in ls:
        s+=stats[l]
        sums.append(s)
    qs = [float(x)/sums[-1] for x in sums]
    
    q=0
    
    try:
        while qs[q] < quantile:
            q+=1
    except IndexError:
        q = len(ls)-1
                
    return ls[q]
        
        
def pairedEndLengthFilter(pairs,length):
    for p in pairs:
        if len(p[0][1]) >= length and len(p[1][1]) >= length:
            yield p        
    
        
def pairedEndCounter(pairs,n=-1):
    
    i=0
    for p in pairs:
        yield p
        i+=1
        if i == n:
            break
        
def writeToFifo(pairs,forward,reverse):
    with nested(open(forward,'a'),
                open(reverse,'a')) as (f,r):
        for p in pairs:
            print >>f,formatFastq(p[0])
            print >>r,formatFastq(p[1])
            
        

def run(config):

    logger  = config['orgasm']['logger']
    output  = config['orgasm']['indexfilename']
    forward = config['index']['forward']
    reverse = config['index']['reverse']
    fasta   = config['index']['fasta']
    
    convert = False
    
    if forward[-3:]=='.gz':
        logger.info('Forward file compressed by gzip')
        forward = ungzip(forward)
        
    if forward[-4:]=='.bz2':
        logger.info('Forward file compressed by bzip2')
        forward = unbzip(forward)

        
    if reverse is not None and reverse[-3:]=='.gz':
        logger.info('Reverse file compressed by gzip')
        reverse = ungzip(reverse)
 
    if reverse is not None and reverse[-4:]=='.bz2':
        logger.info('Reverse file compressed by bzip2')
        reverse = unbzip(reverse)
        
       
    cforward = readFastq(forward)
    creverse = readFastq(reverse) if reverse is not None else None
    pairs    = None
    
    if fasta or config['index']['ffasta']:
        logger.info('Forward file format is fasta')
        convert=True
        cforward = readFasta(forward)
        
    if fasta or config['index']['rfasta']:
        if creverse is not None:
            convert=True
            logger.info('Reverse file format is fasta')
            creverse = readFasta(reverse)
            
    if config['index']['estimate'] is not None:
        logger.info('Computing read length statistics for the forward file')
        stats,cforward = lengthStats(acgtCut(cforward))
        quantile = quantiles(stats, config['index']['estimate'])
        if creverse is not None:
            logger.info('Computing read length statistics for the reverse file')
            stats,creverse = lengthStats(acgtCut(creverse))
            quantile = min(quantiles(stats, config['index']['estimate']),quantile)
            
        logger.info('Read size considered for a quantile of %3.2f : %d',
                    config['index']['estimate'],
                    quantile)
        config['index']['length']=quantile

    if config['index']['single']:
        convert=True
        length = config['index']['length']
        logger.info('Simulate paired reads of %dbp' % length)
        if length is None:
            length=100
        pairs = singleToPairedEnd(cforward,length)
    else:
        if convert:
            pairs = doubleToPairedEnd(cforward,creverse)
            
    if config['index']['maxread'] is not None:
        logger.info('limit indexation to the first %d millions of reads' % config['index']['maxread'])
        
        if convert : 
            pairs = pairedEndCounter(pairs,n=config['index']['maxread'])
        
    command = ['orgasmi','-o',output]
    
    if config['index']['maxread'] is not None:
        command.append('-M')
        command.append(str(config['index']['maxread']))
        
    if config['index']['length'] is not None:
        command.append('-l')
        command.append(str(config['index']['length']))
        
    if convert:
        forward = mktemp(prefix="forward-", dir=getTmpDir())
        os.mkfifo(forward)
        atexit.register(postponerm(forward))
        
        reverse = mktemp(prefix="reverse-", dir=getTmpDir())
        os.mkfifo(reverse)
        atexit.register(postponerm(reverse))
        
    command.append(forward)
    command.append(reverse)
    
    process=Popen(command)
    
    if convert:
        writeToFifo(pairs,forward,reverse)

    process.wait()
    
    
            
        

              
