#!/opt/anaconda1anaconda2anaconda3/bin/python
'''---------------------------------------------------------------------------------------
MACE: Model based Analysis of ChIP-EXo. Workflow:
1) Sequencing depth normalization and nucleotide composition bias correction
2) Signal consolidation and noise reduction (Shannon's Entropy weighted mean)
3) Detect optimal border pair size (Gaussian Mixture, EM) 
4) Border detection (Chebyshen's Inequality)
5) Border matching (Gale-Sharpley)
------------------------------------------------------------------------------------------'''
#import built-in modules
import os,sys
if sys.version_info[0] != 2 or sys.version_info[1] != 7:
        print >>sys.stderr, "\nYou are using python" + str(sys.version_info[0]) + '.' \
        + str(sys.version_info[1]) + "MACE needs python2.7!\n"
        sys.exit()
import string
from optparse import OptionParser
import string
import collections
import  numpy
from time import strftime
import signal
import math
import operator

#import third-party modules
from bx.bitset import *
from bx.bitset_builders import *
from bx.intervals import *
from bx.bbi.bigwig_file import BigWigFile
from pymix import mixture

#import MACE own modules
from exomodule import SAM
from exomodule import BED
from exomodule import peakdet
from exomodule import GaleShapley
from exomodule import mystat

#changing history to this module

__author__ = "Liguo Wang"
__copyright__ = "Copyright 2011-2014, Liguo Wang. All rights reserved."
__credits__ = []
__license__ = "GPL"
__version__="1.2"
__maintainer__ = "Liguo Wang"
__email__ = "wang.liguo@mayo.edu"
__status__ = "Production"

logo="""
   ____    ____        _          ______   ________
  |_   \  /   _|      / \       .' ___  | |_   __  |
    |   \/   |       / _ \     / .'   \_|   | |_ \_|
    | |\  /| |      / ___ \    | |          |  _| _ 
   _| |_\/_| |_   _/ /   \ \_  \ `.___.'\  _| |__/ |
  |_____||_____| |____| |____|  `.____ .' |________|
   Model based    Analysis of     ChIP        Exo
"""                                          

                                                                                        
def printlog (mesg_lst):
	'''print progress into stderr and log file'''
	if len(mesg_lst)==1:
		msg = "@ " + strftime("%Y-%m-%d %H:%M:%S") + ": " +  mesg_lst[0]
	else:
		msg = "@ " + strftime("%Y-%m-%d %H:%M:%S") + ": " + ' '.join(mesg_lst)
	print >>sys.stderr,msg

def signal_handler(signal, frame):
	print >>sys.stderr, '\nYou pressed Ctrl+C. MACE was terminated!'
	os._exit() 

def load_chromsize(file):
	'''read chromosome size file into dict'''
	chromSize={}
	for line in open(file,'r'):
		if line.startswith('#'):continue
		if not line.strip():continue
		fields = line.strip().split()
		if len(fields)<2:
			print >>sys.stderr, "skip " + line
			continue
		try:
			chromSize[fields[0]] = int(fields[1])
		except:
			print >>sys.stderr, "skip " + line
			continue
	if len(chromSize)==0:
		print >>sys.stderr, "Chromosome size file is not in correct format!"
		sys.exit(1)
	return chromSize
		

def pos2sig(bw_obj, chrom,st,end):
	'''for given bigwig file and genome coordinate range. return two lists, first list  
	contains coordinates, second list contains the corresponding signals. coordinates 
	are 1-based. '''
	p=[]
	v=[]
	try:
		coord = st +1
		sig_list = bw_obj.get_as_array(chrom,st,end)
		if numpy.isnan(numpy.nansum(sig_list)):
			return (p,v)
		for i,j in enumerate( sig_list ):
			if numpy.isnan(j): continue
			v.append(j)
			p.append(coord + i)
		return (p,v)
	except:
		return(p,v)

def get_signals_from_bw (bw,chrom,st,end,anchor,u,s):
	'''return coordinates and signal for given bigwig file'''
	p2v={}	#coordinate to signal
	try:
		local_mean = numpy.ma.masked_invalid(bw.get_as_array(chrom,st,end)).mean()
		for i,j in enumerate(numpy.nan_to_num(bw.get_as_array(chrom,st,end))):
			if j <= local_mean:continue
			l = abs(st + i - anchor) - u
			scale = math.exp(-(l**2)/(s**2))
			p2v[st + i] = j * scale
		if len(p2v)==0:
			return None
	except:
		return None
	tmp = max(p2v.iteritems(), key=operator.itemgetter(1))
	if tmp[1]>0:
		return tmp[0]
	else:
		return None
	
def get_signals_from_bb (ibit,chrom,st,end):
	'''get signal from intervalTree. return ([pos1, snr1, pval1],[pos2, snr2, pval2]...), 
	sorted by snr value. coordinates are 1-based'''
	sig_list=[]
	try:
		for i in ibit[chrom].find(st,end):
			#p2v[i.end] = i.value
			sig_list.append([i.end,i.value[0],i.value[1]])
	except:
		pass
	sorted_p2v = sorted(sig_list, key=operator.itemgetter(1),reverse=True)
	return sorted_p2v
		
def build_interval(infile):
	'''build interval tree for bed file'''
	range_FP={}	#range for forward peaks
	range_RP={}	#range for reverse peaks
	
	#chr10   956486  956487  1.0     0.00660803557231        +
	for line in open(infile,'r'):
		if line.startswith(("#","track","browser")):continue
		if not line.strip():continue
		fields=line.rstrip('\n ').split()
		if len(fields) <6: continue
		
		chrom = fields[0]
		start = int(fields[1])
		end = int(fields[2])
		snr = float(fields[3])
		pval = float(fields[4])
		strand = fields[5]
		
		if strand == '+':
			if chrom not in range_FP:
				range_FP[chrom] = IntervalTree()
			range_FP[chrom].insert_interval( Interval( start, end, value=[snr,pval]) )
		elif strand == '-':
			if chrom not in range_RP:
				range_RP[chrom] = IntervalTree()
			range_RP[chrom].insert_interval( Interval( start, end, value=[snr,pval]) )
		else:
			print >>sys.stderr, "unknown strand" + line,
			continue
	return (range_FP,range_RP)

def rank_pos_val_r(lst,pos,u,s):
	'''rank tuple list (reverse strand) by position then by signal value. k is kernel width,
	u,s is mean and std of prior normal distribution. k,u,s are used to weight snr signal
	by distance.'''
	
	new_list=[]
	for r_pos, r_val,r_pval in lst:
		dist = r_pos - pos
		if dist <=0:flag=-1
		if dist >0:flag=1
		
		l = abs(dist) - u
		scale_factor = math.exp(-(l**2)/(s**2))	
		weighted_val = r_val * scale_factor	
		new_list.append((flag,r_pos,weighted_val))

	tmp = sorted(new_list, key=operator.itemgetter(0,2),reverse=True)
	return [i[1] for i in tmp]

def rank_pos_val_f(lst,pos,u,s):
	'''rank tuple list (forward strand) by position then signal value. k is kernel width,
	u,s is mean and std of prior normal distribution. k,u,s are used to weight snr signal
	by distance.'''
	
	new_list=[]
	for f_pos, f_val,f_pval in lst:
		dist = f_pos - pos
		if dist <=0:flag=1
		if dist >0:flag=-1
		
		l = abs(abs(dist) - u)
		scale_factor = math.exp(-(l**2)/(s**2))		
		weighted_val = f_val * scale_factor	
		new_list.append((flag,f_pos,weighted_val))

	tmp = sorted(new_list, key=operator.itemgetter(0,2),reverse=True)
	return [i[1] for i in tmp]	#return position list (sorted by weighted coverage signal)
	
def fisher_combp(p1,p2):
	'''combine two pvalues into single pvalue using fisher's method'''
	if p1==0 or p2==0:
		return 0
	else:
		return p1*p2*(1.0-math.log(p1*p2))

def bimodal_mixture(lst,m1=10,m2=60,sd1=10,sd2=30,maxIter=500,delta=0.01):
	'''estimate two Gaussian distribution from mixture distribution'''
	from pymix import mixture
	data = mixture.DataSet()
	data.fromArray(numpy.array(lst))
	n1 = mixture.NormalDistribution(m1,sd1)
	n2 = mixture.NormalDistribution(m2,sd2)
	m = mixture.MixtureModel(2,[0.5,0.5], [n1,n2])
	#print >>sys.stderr, m
	
	m.EM(data,maxIter,delta)
	print >>sys.stderr, m
	
	if m.components[0][0].mu > m.components[1][0].mu:
		return (m.components[0][0].mu, m.components[0][0].sigma)
	else:
		return (m.components[1][0].mu, m.components[1][0].sigma)
	
def main():
	usage="%prog [options]"
	parser = OptionParser(usage,version="%prog " + __version__)
	parser.add_option("-f","--forward",action="store",type="string",dest="forward_bw",help="BigWig format file containing coverage calcualted from reads mapped to *forward* strand.")
	parser.add_option("-r","--reverse",action="store",type="string",dest="reverse_bw",help="BigWig format file containing coverage calcualted from reads mapped to *reverse* strand.")
	parser.add_option("-s","--chromSize",action="store",type="string",dest="chromSize",help="Chromosome size file. Tab or space separated text file with 2 columns: first column contains chromosome name, second column contains chromosome size. Example:chr1	249250621 <NewLine> chr2	243199373 <NewLine> chr3	198022430 <NewLine> ... ")
	parser.add_option("-o","--out-prefix",action="store",type="string",dest="output_prefix",help="Prefix of output files. NOTE: if 'prefix.border.bed' exists and was non-empty, peak calling step will be skipped! So if you want to rerun mace.py from scratch, use different 'prefix' or delete old 'prefix.border.bed' before starting.")
	parser.add_option("-p","--pvalue",action="store",type="float",dest="pvalue_cutoff",default=0.05,help="Pvalue cutoff for border detection and subsequent border pairing. default=%default")
	parser.add_option("-m","--max-dist",action="store",type="int",dest="max_distance",default=100,help="Maximum distance allowed for border pairing. default=%default")
	parser.add_option("-e","--fz-window",action="store",type="int",dest="fuzzy_size",default=5,help="Peaks located closely within this window will be merged. default=%default (bp)")	
	parser.add_option("-w","--bg-window",action="store",type="int",dest="window_size",default=100,help="Background window size used to determine background signal level. default=%default (bp)")
	parser.add_option("-n","--fold",action="store",type="float",dest="signal_fold",default=2.0,help="Minmum coverage signal used to build model (i.e. estimate optimal peak pair size). default=%default")
		
	(options,args)=parser.parse_args()

	if not (options.output_prefix and options.chromSize and options.forward_bw and options.reverse_bw):
		print '\n' + logo + '\n'
		parser.print_help()
		sys.exit(0)
	for file in (options.chromSize,options.forward_bw,options.reverse_bw):
		if not os.path.exists(file):
			print >>sys.stderr, '\n\n' + file + " does NOT exists" + '\n'
			sys.exit(0)
	print '\n' + logo + '\n'
	printlog(['Running MACE ...'])
	chrom_sizes = load_chromsize(options.chromSize)
	fw_bw_obj = BigWigFile( file = open(options.forward_bw))
	rv_bw_obj = BigWigFile( file = open(options.reverse_bw))
	signal.signal(signal.SIGINT, signal_handler)

	#=====================================================================================
	#  peak calling if this hasn't been done already.
	#=====================================================================================
	
	outpeakFile = options.output_prefix + ".border.bed"
	peak_table=[]
	if os.path.exists(outpeakFile) and os.path.getsize(outpeakFile) >0:
		printlog (["Using", outpeakFile +'.', "Skip border detection step! Change \"-o\" to start from scratch."]) 
		pass
	else:
		printlog(["Border detection ..."])
		OUT = open(outpeakFile,'w')
		print >>OUT, '\t'.join(["#chrom","st","end","Signal","pvalue","strand"])
		for chr_name, chr_size in chrom_sizes.items():		#iterate each chrom
			#if chr_name != "chr1" :continue
			progress = 0
			printlog(["Processing", chr_name])
			#for each chunk
			for interval in BED.tillingBed(chrName = chr_name,chrSize = chr_size,stepSize = 100000):	
				finish_part = int(interval[2]*100/chr_size)
				if finish_part > progress:
					print >>sys.stderr, " %d%% finished\r" % (finish_part),
					progress = finish_part	
			
				#detect peak on forward strand
				(x_axis,y_axis) = pos2sig(fw_bw_obj,interval[0],interval[1],interval[2])
				if len(y_axis)==0:continue
				for (coord,score,pvalue) in peakdet.Chebyshev(x_axis,y_axis,h=(1/options.pvalue_cutoff)**0.5,bgwindow=options.window_size,fzwindow=options.fuzzy_size):
					if pvalue ==0:
						peak_table.append([chr_name,coord-1,coord,score,3000,'+'])
					elif pvalue >0 and pvalue <= options.pvalue_cutoff:
						peak_table.append([chr_name,coord-1,coord,score,pvalue,'+'])
					elif pvalue > options.pvalue_cutoff:
						continue
				#detect peak on reverse strand
				(x_axis,y_axis) = pos2sig(rv_bw_obj,interval[0],interval[1],interval[2])
				if len(y_axis)==0:continue
				for (coord,score,pvalue) in peakdet.Chebyshev(x_axis,y_axis,h=(1/options.pvalue_cutoff)**0.5,bgwindow=options.window_size,fzwindow=options.fuzzy_size):
					if pvalue ==0:
						peak_table.append([chr_name,coord-1,coord,score,3000,'+'])
					elif pvalue >0 and pvalue <= options.pvalue_cutoff:
						peak_table.append([chr_name,coord-1,coord,score,pvalue,'-'])
					elif pvalue > options.pvalue_cutoff:
						continue
		
		peak_table = sorted(peak_table, key=operator.itemgetter(0,1))	#sort by chrom then start
		
		for i in peak_table:
			print >>OUT, '\t'.join(str(j) for j in i)	
		OUT.close()

	#=====================================================================================
	#  border clustering
	#=====================================================================================
	printlog(["Detect border clusters ..."])
	peakClusterFile = options.output_prefix + ".border_cluster.bed"
	tmp = BED.ParseBED(outpeakFile)
	tmp.segmentBED(maxdist=options.max_distance, outfile=peakClusterFile)
	
	peakPairFile = options.output_prefix + ".border_pair.bed"
	peakPairEliteFile = options.output_prefix + ".border_pair_elite.bed"
	OUT = open(peakPairFile,'w')
	OUTE = open(peakPairEliteFile,'w')
	
	printlog(["Building interval tree ... "])
	(f_range,r_range) = build_interval(outpeakFile)
	
	#=====================================================================================
	#  estimate optimal border pair size
	#=====================================================================================
	min_sample_size=30
	reasonable_sample_size=200
	max_sample_size=10000
	pp_mean=0
	pp_std=0
	pp_elite_list=[]
	printlog(["Estimate optimal border pair size ..."])
	for l in open(peakClusterFile,'r'):		#iterate each peak cluster
		l = l.strip()
		if l.startswith('#'):continue
		
		f2r = collections.defaultdict(list)	#key is forward peak position, value is list of reverse peak positoin
		r2f = collections.defaultdict(list)
		chr,st,end = l.split()
		st = int(st)
		end = int(end)
		f_peaks = get_signals_from_bb(f_range, chr, st, end)	#(sorted by snr value) tuple list [('pos1', 2.3, 0.014), ('pos2', 12), ('pos3', 4)]
		r_peaks = get_signals_from_bb(r_range, chr, st, end)
		
		if len(f_peaks) !=1 or len(r_peaks) !=1: continue
		if f_peaks[0][1] < options.signal_fold: continue
		if r_peaks[0][1] < options.signal_fold: continue
		pp_elite_list.append(abs(f_peaks[0][0]- r_peaks[0][0] ))
		
		if f_peaks[0][0] < r_peaks[0][0]:
			print >>OUTE, chr + '\t' + str(f_peaks[0][0]-1) + '\t' +  str(r_peaks[0][0])
		else:
			print >>OUTE, chr + '\t' + str(r_peaks[0][0]-1) + '\t' +  str(f_peaks[0][0])
		
		if len(pp_elite_list) > max_sample_size: 
			break
			
	OUTE.close()	

	if len(pp_elite_list) < min_sample_size:
		print >>sys.stderr, "@ " + strftime("%Y-%m-%d %H:%M:%S") +  " Warning: Only " + str(len(pp_elite_list)) + ' border pairs found! MACE will be terminated. Rerun mace.py with lowered \"--fold (-n)\" value to increase sample size.'
		sys.exit(1)
	elif len(pp_elite_list) >=min_sample_size and len(pp_elite_list) < reasonable_sample_size:
		print >>sys.stderr, "@ " + strftime("%Y-%m-%d %H:%M:%S") + " Warning: Only " + str(len(pp_elite_list)) + ' border pairs found! Please terminate mace.py and lower the \"--fold (-n)\" value to increase sample size. Otherwise, program will continue to use this sample to estiamte optimal border pair size.'
		(pp_mean, pp_std) = bimodal_mixture(pp_elite_list)
	else:
		(pp_mean, pp_std) = bimodal_mixture(pp_elite_list)
	pp_mean = mystat.singleMode([i for i in pp_elite_list if (i > pp_mean - pp_std) and (i < pp_mean + pp_std)])
	print >>sys.stderr, "@ " + strftime("%Y-%m-%d %H:%M:%S") + ":" + " Optimal border Pari size (bp): Mean = " + str(pp_mean) + ' Std = ' + str(pp_std)
	
	#=====================================================================================
	#  peak pairing within peak clusters
	#=====================================================================================
	pp_list=[]
	printlog (["Border pairing ... "])
	for l in open(peakClusterFile,'r'):		#iterate each peak cluster
		l = l.strip()
		if l.startswith('#'):continue
		
		f2r = collections.defaultdict(list)	#key is forward peak position, value is list of reverse peak positoin
		r2f = collections.defaultdict(list)
		chr,st,end = l.split()
		st = int(st)
		end = int(end)
		f_peaks = get_signals_from_bb(f_range, chr, st, end)	#(sorted by snr value) tuple list [('pos1', 2.3, 0.014), ('pos2', 12), ('pos3', 4)]
		r_peaks = get_signals_from_bb(r_range, chr, st, end)
		
		#single forward border
		if len(f_peaks) >=1 and len(r_peaks)==0:
			(f_peak_highest_pos, f_peak_highest_sig, f_peak_highest_pval) = max(f_peaks,key=operator.itemgetter(1))		#('pos1', 2.3, 0.014)	
			
			try:
				tmp_sig = get_signals_from_bw (rv_bw_obj, chr, f_peak_highest_pos , f_peak_highest_pos + options.max_distance, f_peak_highest_pos,pp_mean,pp_std)	#the max signal although not significant
			except:
				continue
			if tmp_sig >0:
				pp_end = tmp_sig
				#if f_peak_highest_pos-1 == 23442377:print tmp_sig
				pp_list.append((chr,f_peak_highest_pos-1,pp_end+1,'SFB',f_peak_highest_pval))
			else:		#cannot find significant value from bigwig file
				pp_list.append((chr,f_peak_highest_pos-1,f_peak_highest_pos-1+pp_mean,'SFB',f_peak_highest_pval))
			
		
		#single reverse border
		elif len(f_peaks) ==0 and len(r_peaks)>=1:
			(r_peak_highest_pos, r_peak_highest_sig, r_peak_highest_pval) = max(r_peaks,key=operator.itemgetter(1))		#('pos1', 2.3, 0.014)	
			try:
				tmp_sig = get_signals_from_bw (fw_bw_obj, chr, r_peak_highest_pos - options.max_distance,  r_peak_highest_pos, r_peak_highest_pos, pp_mean,pp_std)	#the max signal although not significant
			except:
				continue
			if tmp_sig >0:
				pp_start = tmp_sig
				pp_list.append((chr,pp_start,r_peak_highest_pos,'SRB',r_peak_highest_pval))
			else:		#cannot find significant value from bigwig file
				if r_peak_highest_pos - pp_mean >= 0:
					pp_list.append((chr,r_peak_highest_pos - pp_mean,r_peak_highest_pos,'SRB',r_peak_highest_pval))
				else:
					pp_list.append((chr,0,r_peak_highest_pos,'SRB',r_peak_highest_pval))
			
		
		#both two border defined
		elif len(f_peaks) >=1 and len(r_peaks)>=1:
			f2r = {}	#key is f_position, value is list of reverse positions
			r2f = {}
			f_pos2pval={}
			r_pos2pval={}
			for pos,snr,pval in f_peaks:
				f_pos2pval[pos]=pval
			for pos,snr,pval in r_peaks:
				r_pos2pval[pos]=pval

			for f_pos, f_val,f_pval in f_peaks:
				f2r[f_pos] = rank_pos_val_r(r_peaks,f_pos,pp_mean,pp_std)
			for r_pos, r_val,r_pval in r_peaks:
				r2f[r_pos] = rank_pos_val_f(f_peaks,r_pos,pp_mean,pp_std)

			if len(f_peaks) <= len(r_peaks):
				for start, end in GaleShapley.matchmaker(f2r,r2f).items():	#small group proposed to large group
					if start < end:
						pp_list.append((chr,start-1,end,'GSB',fisher_combp(f_pos2pval[end],r_pos2pval[start])))
					else:
						pp_list.append((chr,end-1,start,'GSB',fisher_combp(f_pos2pval[end],r_pos2pval[start])))
			elif len(f_peaks) > len(r_peaks):
				for start, end in GaleShapley.matchmaker(r2f,f2r).items():	#small group proposed to large group
					if start < end:
						pp_list.append((chr,start-1,end,'GSB',fisher_combp(f_pos2pval[start],r_pos2pval[end])))
					else:
						pp_list.append((chr,end-1,start,'GSB',fisher_combp(f_pos2pval[start],r_pos2pval[end])))
			
	#=============================================================
	#sort peak pairs by chrom, then start
	#=============================================================
	pp_list = sorted(pp_list, key=operator.itemgetter(0,1))
	for k in pp_list:
		print >>OUT, '\t'.join([str(i) for i in k])	
				
if __name__ == '__main__':
	main()
