"""
Script to call nucleosome positions-- track making, nucleosome calling, and nfr calling!

@author: Alicia Schep
"""
##### IMPORT MODULES #####
# import necessary python modules
import matplotlib as mpl
mpl.use('PS')
import multiprocessing as mp
import numpy as np
import traceback
import itertools
import pysam
from pyatac.utils import shell_command,read_chrom_sizes_from_bam
from pyatac.chunk import ChunkList
from nucleoatac.NucleosomeCalling import NucChunk, NucParameters
from pyatac.fragmentsizes import FragmentSizes
from pyatac.bias import PWM
from pyatac.VMat import VMat

def _nucHelper(arg):
    """function to get occupancy for a set of bed regions

    """
    (chunk, params) = arg
    try:
        nuc = NucChunk(chunk)
        nuc.process(params)
        out = {'nucpos' : [nuc.nuc_collection[i] for i in sorted(nuc.nonredundant)], 'nucpos.redundant' : [nuc.nuc_collection[i] for i in sorted(nuc.redundant)],
                               'nucleoatac_signal' : nuc.norm_signal, 'nucleoatac_raw' : nuc.nuc_signal, 'nucleoatac_background' : nuc.bias,
                               'nucleoatac_signal.smooth' : nuc.smoothed}
        nuc.removeData()
    except Exception as e:
        print('Caught exception when processing:\n'+  chunk.asBed()+"\n")
        traceback.print_exc()
        print()
        raise e
    return out



def _writeNucSig(track_queue, out):
    out_handle = open(out + '.nucleoatac_signal.bedgraph','a')
    try:
        for track in iter(track_queue.get, 'STOP'):
            track.write_track(out_handle)
            track_queue.task_done()
    except Exception, e:
        print('Caught exception when writing NucleoATAC signal track\n')
        traceback.print_exc()
        print()
        raise e
    out_handle.close()
    return True


def _writeBackground(track_queue, out):
    out_handle = open(out + '.nucleoatac_background.bedgraph','a')
    try:
        for track in iter(track_queue.get, 'STOP'):
            track.write_track(out_handle)
            track_queue.task_done()
    except Exception, e:
        print('Caught exception when writing NucleoATAC background track\n')
        traceback.print_exc()
        print()
        raise e
    out_handle.close()
    return True


def _writeSmooth(track_queue, out):
    out_handle = open(out + '.nucleoatac_signal.smooth.bedgraph','a')
    try:
        for track in iter(track_queue.get, 'STOP'):
            track.write_track(out_handle)
            track_queue.task_done()
    except Exception, e:
        print('Caught exception when writing smoothed NucleoATAC signal track\n')
        traceback.print_exc()
        print()
        raise e
    out_handle.close()
    return True

def _writeRaw(track_queue, out):
    out_handle = open(out + '.nucleoatac_raw.bedgraph','a')
    try:
        for track in iter(track_queue.get, 'STOP'):
            track.write_track(out_handle)
            track_queue.task_done()
    except Exception, e:
        print('Caught exception when writing un-normalized NucleoATAC signal track\n')
        traceback.print_exc()
        print()
        raise e
    out_handle.close()
    return True



def _writeNucPos(pos_queue, out):
    out_handle = open(out + '.nucpos.bed','a')
    try:
        for poslist in iter(pos_queue.get, 'STOP'):
            for pos in poslist:
                pos.write(out_handle)
            pos_queue.task_done()
    except Exception, e:
        print('Caught exception when writing nucleosome position file\n')
        traceback.print_exc()
        print()
        raise e
    out_handle.close()
    return True

def _writeNucPosRedundant(pos_queue, out):
    out_handle = open(out + '.nucpos.redundant.bed','a')
    try:
        for poslist in iter(pos_queue.get, 'STOP'):
            for pos in poslist:
                pos.write(out_handle)
            pos_queue.task_done()
    except Exception, e:
        print('Caught exception when writing redundant nucleosome position file\n')
        traceback.print_exc()
        print()
        raise e
    out_handle.close()
    return True




_writeFuncs = {'nucpos' : _writeNucPos, 'nucpos.redundant' : _writeNucPosRedundant,
               'nucleoatac_signal' : _writeNucSig, 'nucleoatac_raw' : _writeRaw, 'nucleoatac_background' : _writeBackground,
               'nucleoatac_signal.smooth' : _writeSmooth}


def run_nuc(args):
    """run occupancy calling

    """
    vmat = VMat.open(args.vmat)
    chrs = read_chrom_sizes_from_bam(args.bam)
    pwm = PWM.open(args.pwm)
    chunks = ChunkList.read(args.bed, chromDict = chrs, min_offset = vmat.mat.shape[1] + vmat.upper/2 + max(pwm.up,pwm.down) + args.nuc_sep/2, min_length = args.nuc_sep * 2)
    chunks.slop(chrs, up = args.nuc_sep/2, down = args.nuc_sep/2)
    chunks.merge()
    maxQueueSize = args.cores*10
    if args.sizes is not None:
        fragment_dist = FragmentSizes.open(args.sizes)
    else:
        fragment_dist = FragmentSizes(0, upper = vmat.upper)
        fragment_dist.calculateSizes(args.bam, chunks)
    params = NucParameters(vmat = vmat, fragmentsizes = fragment_dist, bam = args.bam, fasta = args.fasta, pwm = args.pwm,
                           occ_track = args.occ_track,
                           sd = args.sd, nonredundant_sep = args.nuc_sep, redundant_sep = args.redundant_sep,
                           min_z = args.min_z, min_lr = args.min_lr , atac = args.atac)
    sets = chunks.split(items = args.cores*5)
    pool1 = mp.Pool(processes = max(1,args.cores-1))
    if args.write_all:
        outputs = ['nucpos','nucpos.redundant','nucleoatac_signal','nucleoatac_signal.smooth',
                       'nucleoatac_background','nucleoatac_raw']
    else:
        outputs = ['nucpos','nucpos.redundant','nucleoatac_signal','nucleoatac_signal.smooth']
    handles = {}
    write_queues = {}
    write_processes = {}
    for i in outputs:
        if i not in ['nucpos','nucpos.redundant','nfrpos']:
            handles[i] = open(args.out + '.'+i+'.bedgraph','w')
        else:
            handles[i] = open(args.out + '.'+i+'.bed','w')
        handles[i].close()
        write_queues[i] = mp.JoinableQueue(maxsize = maxQueueSize)
        write_processes[i] = mp.Process(target = _writeFuncs[i], args=(write_queues[i], args.out))
        write_processes[i].start()
    for j in sets:
        tmp = pool1.map(_nucHelper, zip(j,itertools.repeat(params)))
        for result in tmp:
            for i in outputs:
                write_queues[i].put(result[i])
    pool1.close()
    pool1.join()
    for i in outputs:
        write_queues[i].put('STOP')
    for i in outputs:
        write_processes[i].join()
        if i not in ['nucpos','nucpos.redundant']:
            pysam.tabix_compress(args.out + '.' + i + '.bedgraph', args.out +  '.' + i + '.bedgraph.gz',force = True)
            shell_command('rm ' + args.out +  '.' + i + '.bedgraph')
            pysam.tabix_index(args.out +  '.' + i + '.bedgraph.gz', preset = "bed", force = True)
        else:
            pysam.tabix_compress(args.out + '.' + i + '.bed', args.out +  '.' + i + '.bed.gz',force = True)
            shell_command('rm ' + args.out +  '.' + i + '.bed')
            pysam.tabix_index(args.out +  '.' + i + '.bed.gz', preset = "bed", force = True)
 





