#!/usr/bin/env python
# Programmer : zhuxp
# Date: 
# Last-modified: 02-13-2014, 14:53:10 EST
VERSION="0.0.9"
import os,sys,argparse
from bam2x.Annotation import BED6 as Bed
from bam2x.Annotation import BED3
from bam2x.Annotation import BED12 as Bed12
from bam2x import TableIO,Tools,DBI,IO
import signal
signal.signal(signal.SIGPIPE,signal.SIG_DFL)
import gzip
import time
import multiprocessing as mp
import array,tempfile,heapq
import bam2x.prob as prob
import itertools
assert array.array('i').itemsize==4

MIN_GAP_IN_SAME_EXON=10
'''
V3:
DONE: mv sorting array to Turing Module
V4:
DONE: correct the coverage calculation
DONE: report link exon and intron

V5:
DONE: same memory using numarray instead of bed?
      only using array?
      rm bed class?
      list of list ?
      bedgraph sorting array?
DONE: report if peak has intron and the possible cDNA length and gene length.
V6:
    change format to Bed12
    get rid of terminal introns
V7:
    don't change state because of intron just link all exon will do.
    get rid of consective small exons.
V8:
    DONE: add complexity filter out [ LZW complexity / length > 0.4 ]
    DONE: mv to bam2x subcommand
V9:
    DONE: rm xplib?
    DONE: mv to namedtuple turing

TODO: compare with known gene
TODO: trim the last intron or extend the exon? ( KEY PROBLEM. how to define the end )
'''
EXON_GROUP_CODE=1
INTRON_GROUP_CODE=0
HAS_INTRON=20
NOT_HAS_INTRON=10

POSITIVE_STRAND=1
NEGATIVE_STRAND=-1

START_INDEX=0
STOP_INDEX=1
SCORE_INDEX=2
STRAND_INDEX=3
GROUP_INDEX=4
OTHER_INDEX=5
ID_INDEX=6
EXONSTARTS_INDEX=7
EXONSIZES_INDEX=8
CPLX_CUTOFF=0.4
def help():
    return "bam2peak , call exon peaks and link exons"
def set_parser(p):
    ''' This Function Parse the Argument '''
    p.add_argument('-s','--strand',dest="strand",type=str,choices=["read1","read2"],default="read2",help="read1 or read2 is the positive strand , default: %(default)s")
    p.add_argument('--prefix',dest="prefix",type=str,default="R",help="prefix for bed name default: %(default)s")
    p.add_argument('--cutoff',dest="cutoff",type=int,default=1,help="only report region covrage >= cutoff : %(default)i")
    p.add_argument('--pvalue',dest="pvalue",type=float,default=1e-05,help="cutoff for calling peak : %(default)f")
    p.add_argument('-g','--genome',dest="genome",type=str,help="genome sequence file in twobit format")
def run(args_local):
    '''
    IO TEMPLATE
    '''
    global args,out,exon_cutoff,intron_cutoff,hasGenome
    args=args_local
    dbi=DBI.init(args.input,"bam")
    out=IO.fopen(args.output,"w")
    hasGenome=False
    if args.genome:
        hasGenome=True
    
    '''
    END OF IO TEMPLATE 
    '''
    print >>out,"# This positive_data was generated by program ",sys.argv[1]," (version: %s)"%VERSION,
    print >>out,"in bam2x ( https://github.com/nimezhu/bam2x )"
    print >>out,"# Date: ",time.asctime()
    print >>out,"# The command line is :"
    print >>out,"#\t"," ".join(sys.argv)
    chrs=[]
    lengths=[]
    for i in dbi.bamfiles[0].references:
        chrs.append(i)
    for i in dbi.bamfiles[0].lengths:
        lengths.append(i)
    p=mp.Pool(processes=args.num_cpus)    
    coverage_bedgraphs=p.map(process_chrom,chrs)
    bedgraphs=[]
    coverages=[]
    for i in range(len(chrs)):
        bedgraphs.append(coverage_bedgraphs[i][1])
        coverages.append(coverage_bedgraphs[i][0])
    s=0.0  # 1000.0
    l=long(0)
    for i in range(len(chrs)):
        s+=coverages[i]
        l+=lengths[i]
    l=l*2 # Double Strand
    coverage=s/l*1000.0
    threshold=1
    while 1:
        if prob.poisson_cdf(threshold,coverage,False) < args.pvalue: break
        threshold+=1
    exon_cutoff=threshold
    intron_cutoff=2 #TODO revise it
    print >>out,"# MEAN COVERAGE:",coverage
    print >>out,"# EXON COVERAGE CUTOFF:",exon_cutoff
    
    #call_peaks(bedgraphs[0],1) #debug
    peaks=p.map(call_peaks_star,itertools.izip(chrs,bedgraphs,itertools.repeat(exon_cutoff)))
    output(chrs,peaks)

    #process_chrom("chr1")
def output(chrs,s):
    for chrom,i in itertools.izip(chrs,s):
        for j in i:
            print >>out,nice_format(chrom,j)

def nice_format(chrom,a):
    s=chrom+"\t"
    s+=str(a[START_INDEX])+"\t"
    s+=str(a[STOP_INDEX])+"\t"
    s+=chrom+"_"+str(a[ID_INDEX])+"\t"
    s+=str(a[SCORE_INDEX])+"\t"
    if a[STRAND_INDEX]==1:
        s+="+\t"
    elif a[STRAND_INDEX]==-1:
        s+="-\t"
    else:
        s+=".\t"
    s+=str(a[START_INDEX])+"\t" #cds_start
    s+=str(a[START_INDEX])+"\t" #cds_stop
    s+="0,0,0\t" #Rgb
    s+=str(len(a[EXONSIZES_INDEX]))+"\t"
    for i in a[EXONSIZES_INDEX]:
        s+=str(i)+","
    s+="\t"
    for i in a[EXONSTARTS_INDEX]: 
        s+=str(i)+","
    
    return s

    #return "\t".join("%s"%item for item in a)
from bam2x.Turing import TuringCode
from bam2x.Turing import TuringCodeBook as cb
from bam2x.Turing import TuringSortingArray
from operator import itemgetter
from bam2x.Tools.LZW import complexity 
def process_chrom(chrom):
    local_dbi=DBI.init(args.input,"bam")
    retv=list()
    intron_retv=list()
    a=[]
    #positive_data=TuringSortingArray(None,500)
    positive_data=TuringSortingArray()
    negative_data=TuringSortingArray()
    
    positive_intron_data=TuringSortingArray()
    negative_intron_data=TuringSortingArray()

    for i in local_dbi.query(chrom,method="bam1",strand=args.strand):
        if i.strand=="+" or i.strand==".": 
            positive_data.append(TuringCode(i.start,cb.ON))
            positive_data.append(TuringCode(i.stop,cb.OFF))
            for j in i.Exons():
                positive_data.append(TuringCode(j.start,cb.BLOCKON))
                positive_data.append(TuringCode(j.stop,cb.BLOCKOFF))
            for j in i.Introns():
                positive_intron_data.append(TuringCode(j.start,cb.BLOCKON))
                positive_intron_data.append(TuringCode(j.stop,cb.BLOCKOFF))
        else:
            negative_data.append(TuringCode(i.start,cb.ON))
            negative_data.append(TuringCode(i.stop,cb.OFF))
            for j in i.Exons():
                negative_data.append(TuringCode(j.start,cb.BLOCKON))
                negative_data.append(TuringCode(j.stop,cb.BLOCKOFF))
            for j in i.Introns():
                negative_intron_data.append(TuringCode(j.start,cb.BLOCKON))
                negative_intron_data.append(TuringCode(j.stop,cb.BLOCKOFF))
    cutoff=args.cutoff
    coverage=0.0
    for i,x in enumerate(codesToBedGraph(positive_data.iter())):
        if x[SCORE_INDEX] >= cutoff:
            retv.append((x[0],x[1],x[2],POSITIVE_STRAND,EXON_GROUP_CODE))
        coverage+=float(x[1]-x[0])*x[2]/1000.0
    for i,x in enumerate(codesToBedGraph(negative_data.iter())):
        if x[SCORE_INDEX] >= cutoff:
            retv.append((x[0],x[1],x[2],NEGATIVE_STRAND,EXON_GROUP_CODE))
        coverage+=float(x[1]-x[0])*x[2]/1000.0
    INTRON_CUTOFF=1.0
    for i,x in enumerate(codesToBedGraph(positive_intron_data.iter())):
        if x[SCORE_INDEX] >= INTRON_CUTOFF:
            retv.append((x[0],x[1],x[2],POSITIVE_STRAND,INTRON_GROUP_CODE))
    for i,x in enumerate(codesToBedGraph(negative_intron_data.iter())):
        if x[SCORE_INDEX] >= INTRON_CUTOFF:
            retv.append((x[0],x[1],x[2],NEGATIVE_STRAND,INTRON_GROUP_CODE))
    retv.sort(key=itemgetter(0,1,2))
    #TODO how to sort!
    local_dbi.close()
    return coverage,retv
def codesToBedGraph(iter):
    a=iter.next()
    last_pos=a[0]
    counter=0
    for i in iter:
        if i.pos!=last_pos:
            #TODO using more efficient data structure
            yield (last_pos,i.pos,counter)
            last_pos=i.pos
        if i.code==cb.BLOCKON:
            counter+=1
        if i.code==cb.BLOCKOFF:
            counter-=1
    raise StopIteration


def call_peaks_star(a_b_c):
    return call_peaks(*a_b_c)
    
def call_peaks(chrom,bedgraph,exon_cutoff):
    #TODO
    
    if hasGenome:
        genome=DBI.init(args.genome,"genome")
    def filter_low_complexity(peak):
        new_exon_starts=[]
        new_exon_sizes=[]
        local_peak=list(peak)
        offset=local_peak[START_INDEX]
        for exon_start,exon_size in itertools.izip(local_peak[EXONSTARTS_INDEX],local_peak[EXONSIZES_INDEX]):
            exon_s=offset+exon_start
            exon_e=offset+exon_start+exon_size
            seq=genome.query(BED3(chrom,exon_s,exon_e),method="seq")
            if len(seq) > 100:
                cplx=float(complexity(seq[0:100]))/100
            else:
                cplx=float(complexity(seq))/len(seq)
            if cplx > CPLX_CUTOFF:
                new_exon_starts.append(exon_start)
                new_exon_sizes.append(exon_size)
        if len(new_exon_sizes) > 0:
            local_peak[EXONSIZES_INDEX]=tuple(new_exon_sizes)
            local_peak[START_INDEX]=new_exon_starts[0]+offset
            local_peak[STOP_INDEX]=new_exon_starts[-1]+offset+new_exon_sizes[-1]
            shift_start=new_exon_starts[0]
            for i0 in range(len(new_exon_starts)):
                new_exon_starts[i0]-=shift_start
            local_peak[EXONSTARTS_INDEX]=tuple(new_exon_starts)
            return tuple(local_peak)
        else:
            return None
    gap=10
    pos_beds=[]
    neg_beds=[]
    peaks=[]
    i_p=0
    i_n=0
    last_pos_stop=0
    last_neg_stop=0
    for i in bedgraph:
        if i[STRAND_INDEX]==POSITIVE_STRAND:
            if i[GROUP_INDEX]==EXON_GROUP_CODE:
                if i[SCORE_INDEX] >= exon_cutoff:
                    if len(pos_beds)>0:
                        if  i[START_INDEX]-last_pos_stop < gap or last_pos_stop==0:
                            pos_beds.append(i)
                            if last_pos_stop < i[STOP_INDEX]:
                                last_pos_stop=i[STOP_INDEX]
                        else:
                            
                            peak=bedsToPeak(pos_beds,"p_"+str(i_p))
                            if peak is not None:
                                if hasGenome:
                                    peak=filter_low_complexity(peak)
                                    if peak is not None:
                                        peaks.append(peak)
                                        i_p+=1
                                else:
                                    peaks.append(peak)
                                    i_p+=1
                            pos_beds=[i]
                            last_pos_stop=i[STOP_INDEX]
                    else:
                        last_pos_stop=i[STOP_INDEX]
                        pos_beds.append(i)
            elif i[GROUP_INDEX]==INTRON_GROUP_CODE:
                if last_pos_stop < i[STOP_INDEX]:
                    last_pos_stop=i[STOP_INDEX]
                pos_beds.append(i)
        else:
            if i[GROUP_INDEX]==EXON_GROUP_CODE:
                if i[SCORE_INDEX] >= exon_cutoff:
                    if len(neg_beds)>0:
                        if  i[START_INDEX]-last_neg_stop < gap or last_neg_stop==0:
                            neg_beds.append(i)
                            if last_neg_stop < i[STOP_INDEX]:
                                last_neg_stop=i[STOP_INDEX]
                        else:
                            peak=bedsToPeak(neg_beds,"n_"+str(i_n))
                            if peak is not None:
                                if hasGenome:
                                    peak=filter_low_complexity(peak)
                                    if peak is not None:
                                        peaks.append(peak)
                                        i_n+=1
                                else:
                                    peaks.append(peak)
                                    i_n+=1
                            neg_beds=[i]
                            last_neg_stop=i[STOP_INDEX]
                    else:
                        neg_beds.append(i)
                        last_neg_stop=i[STOP_INDEX]
            elif i[GROUP_INDEX]==INTRON_GROUP_CODE:
                if last_neg_stop < i[STOP_INDEX]:
                    last_neg_stop=i[STOP_INDEX]
                neg_beds.append(i)

    if len(pos_beds)>0:
        peak=bedsToPeak(pos_beds,"p_"+str(i_p))
        if peak is not None:
            if hasGenome:
                peak=filter_low_complexity(peak)
                if peak is not None:        
                    peaks.append(peak)
            else:
                peaks.append(peak)
    if len(neg_beds)>0:
        peak=bedsToPeak(neg_beds,"n_"+str(i_n))
        if peak is not None:
            if hasGenome:
                peak=filter_low_complexity(peak)
                if peak is not None:        
                    peaks.append(peak)
            
            else:
                peaks.append(peak)
    peaks.sort()
    return peaks
def length(x):
    return x[STOP_INDEX]-x[START_INDEX]
def bedsToPeak(ibeds,id):
    try:
        while(ibeds[0][GROUP_INDEX]==INTRON_GROUP_CODE):
            try:
                ibeds=ibeds[1:]
            except:
                return None
        peak=[ibeds[0][START_INDEX],ibeds[0][STOP_INDEX],float(ibeds[0][SCORE_INDEX]),ibeds[0][STRAND_INDEX],NOT_HAS_INTRON,0,id]
        cdna_length=length(peak)
        exonstarts=[]
        exonsizes=[]
        exon_start_point=ibeds[0][START_INDEX]
        exon_stop_point=ibeds[0][STOP_INDEX]
        exon_signal=1
        for i in ibeds[1:]:
            if i[GROUP_INDEX]==EXON_GROUP_CODE:
                if exon_signal==0:
                    exon_start_point=i[START_INDEX]
                    exon_signal=1
                elif i[START_INDEX]-exon_stop_point > MIN_GAP_IN_SAME_EXON:
                    exonstarts.append(exon_start_point-peak[0])
                    exonsizes.append(exon_stop_point-exon_start_point)
                    #print "debug in intron"
                    exon_start_point=i[START_INDEX]
                    exon_stop_point=i[STOP_INDEX]

                peak[SCORE_INDEX]=float(peak[SCORE_INDEX]*cdna_length+i[SCORE_INDEX]*length(i))/(length(peak)+length(i))
                peak[STOP_INDEX]=i[STOP_INDEX]
                cdna_length+=length(i)
                exon_stop_point=i[STOP_INDEX]
            else:
                peak[GROUP_INDEX]=HAS_INTRON
                '''
                if exon_signal==1:
                    exonstarts.append(exon_start_point-peak[0])
                    exonsizes.append(exon_stop_point-exon_start_point)
                    #print "debug in intron"
                    exon_signal=0
                '''
        if exon_signal==1:
            exonstarts.append(exon_start_point-peak[0])
            exonsizes.append(exon_stop_point-exon_start_point)
            exon_signal=0
        peak[OTHER_INDEX]=cdna_length
        '''
        print "DEBUG META",peak
        for i in ibeds:
            print "DEBUG IN",i
        print "DEBUG",exonstarts
        print "DEBUG",exonsizes
        '''
        peak.append(tuple(exonstarts))
        peak.append(tuple(exonsizes))
        #print "debug",tuple(peak)
        return tuple(peak)
    except:
        # ibeds is zero.
        return None



        

    
if __name__=="__main__":
    run()








