#!/opt/anaconda1anaconda2anaconda3/bin/python
'''-------------------------------------------------------------------------------------------------
Convert BAM/SAM file into wig file. BAM file must be sorted and indexed
How to sort and index BAM file: http://genome.ucsc.edu/goldenPath/help/bam.html
If single BAM file is provided. Wiggle format file represented the 5' end of reads will be generated
If multiple BAM files (replicates) are provided. Wiggle format file represented the 
reproducibility score (for 5' end of reads) will be generated.
-------------------------------------------------------------------------------------------------'''
#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
import warnings
import string
import collections
from optparse import OptionParser
from string import maketrans
import subprocess

#import third-party modules
import numpy
from bx.bitset import *
from bx.bitset_builders import *
from bx.intervals import *
import pysam

#import my own modules
from exomodule import SAM
from exomodule import BED
from exomodule import FrameKmer
from exomodule import mystat
#changes to the paths

#changing history to this module


__author__ = "Liguo Wang"
__copyright__ = "Copyright 2011-2013. 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 revComp(seq):
	'''reverse complement DNA sequences'''
	from string import maketrans
	transtab = maketrans("ACGTNX","TGCANX")
	return seq.upper().translate(transtab)[::-1]
	
def build_wig_bak(samfile_obj,chrom,st,end):
        '''build wig files from using pileup object
        st is 0 based, end is 1 based'''

        v2p = collections.defaultdict(list)     #value to position
        range2p={}      #coorindate range to value, bedgraph
        for pileupcolumn in samfile.pileup(chrom,st,end):
                if pileupcolumn.pos +1 <= st:continue
                if pileupcolumn.pos +1 > end :continue
                v2p[pileupcolumn.n].append(pileupcolumn.pos +1)
                #p_v.append((pileupcolumn.pos +1,pileupcolumn.n))
        for v in v2p:
                for k,g in groupby(enumerate(v2p[v]), lambda (i,x):i-x):
                        for l in [map(itemgetter(1), g)]:
                                range2p[l[0]-1] = [len(l),v]
                                #print l[0]-1,l[-1],v 
        return range2p  #[start]=[len,value]

def load_chromsize(file):
	'''read chrom.size file'''
	chromSize={}
	for line in open(file,'r'):
		if line.startswith('#'):continue
		if not line.strip():continue
		fields = line.strip().split()
		chromSize[fields[0]] = int(fields[1])
	return chromSize

def position2value(samfile_obj,chrom,st,end):
	'''using samtools pileup engine, calculate wig signal for each position'''

	p2v = collections.defaultdict(int)     #position to value
	for pileupcolumn in samfile.pileup(chrom,st,end):
		if pileupcolumn.pos +1 <= st:continue
		if pileupcolumn.pos +1 > end :continue
		p2v[pileupcolumn.pos +1]=pileupcolumn.n
	return p2v  #[start]=[len,value]

def snr(lst):
	'''calculate signal-to-noise ratio'''
	lst_avg = numpy.mean(lst)
	lst_std = numpy.std(lst,ddof=1)
	if numpy.isnan(lst_std) or lst_std == 0:
		lst_std=1
	return float(lst_avg/lst_std)

def arithymetic_mean(lst):
	'''calculate signal-to-noise ratio'''
	lst_avg = numpy.mean(lst)
	return lst_avg

def geometric_mean(lst):
	'''calculate geometrix mean'''
	#lst2 = [i for i in lst if i >0]		#remove 0
	return numpy.array(lst).prod()**(1.0/len(lst))

def entropy_mean(lst,maxH):
	'''return mean weighted by entropy'''
	if maxH >0:
		return numpy.mean(lst) * (mystat.shannon_entropy(lst)/maxH)
	else:
		return numpy.mean(lst)
	
def build_wig(samfile_objs,chrom,st,end,nfs,wghts,kmer,maxH,met,qcut):
	'''
	build wig files from alignedReads object.
	nfs: normalization factor, used to correct sequencing depth
	wghts: weights for each kmer, used to correct nucleotide composition bias
	kmer: kmer length
	met: normalization method
	qcut: mapping quality cutoff
	'''
		 
	transtab = maketrans("ACGTNX","TGCANX")
	all_p_wig=collections.defaultdict(list)
	all_n_wig=collections.defaultdict(list)
	all_p_snr={}
	all_n_snr={}
	all_p_gm={}
	all_n_gm={}	
	
	if met not in ("EM","AM","GM","SNR"):
		print >>sys.stderr, "Unrecognized parameter " + met
		print >>sys.stderr, 'Must be one of "EM","AM","GM" or "SNR"'
		sys.exit(0)
	
	if len(nfs) != len(samfile_objs):
		print >>sys.stderr, "normalizing factors and number of bam file do not match"
		sys.exit(1)
	
	for (f,nf) in zip(samfile_objs,nfs):
		try:
			aligned_reads = f.fetch(chrom,st,end)
		except:
			print "error"
		Pwig=collections.defaultdict(float)
		Nwig=collections.defaultdict(float)
		for aligned_read in aligned_reads:
			if aligned_read.is_qcfail:			#skip low quanlity
				continue
			if aligned_read.is_duplicate:		#skip duplicate read
				continue
			if aligned_read.is_secondary:		#skip non primary hit
				continue
			if aligned_read.is_unmapped:		#skip unmap read
				continue	
			if aligned_read.mapq < qcut:
				continue	
			read_seq = aligned_read.seq.upper()
			if aligned_read.is_reverse:
				read_seq =  read_seq.translate(transtab)[::-1]
				read5end = aligned_read.pos + aligned_read.qlen
				try:
					if wghts is not None:
						Nwig[read5end] += wghts[read_seq[0:kmer]]
					else:  #turn off nucleotide composition correction
						Nwig[read5end] += 1.0
				except:
					Nwig[read5end] += 1.0
			else:
				read5end = aligned_read.pos + 1
				try:
					if wghts is not None:
						Pwig[read5end] += wghts[read_seq[0:kmer]]
					else:		#turn off nucleotide composition correction
						Pwig[read5end] +=1.0
				except:
					Pwig[read5end] +=1.0
		for i in xrange(st+1,end+1):
			if i in Nwig:
				all_n_wig[i].append(Nwig[i]*nf)
			else:
				all_n_wig[i].append(0)
			if i in Pwig:
				all_p_wig[i].append(Pwig[i]*nf)
			else:
				all_p_wig[i].append(0)

	for i in xrange(st+1,end+1):
		if sum(all_p_wig[i])==0:continue
		if met == "EM":
			all_p_gm[i] = entropy_mean(all_p_wig[i],maxH)
		elif met == "AM":
			all_p_gm[i] = arithymetic_mean(all_p_wig[i])
		elif met == "GM":
			all_p_gm[i] = geometric_mean(all_p_wig[i])
		elif met == "SNR":
			all_p_gm[i] = snr(all_p_wig[i])
		
	for i in xrange(st+1,end+1):
		if sum(all_n_wig[i])==0:continue
		if met == "EM":
			all_n_gm[i] = entropy_mean(all_n_wig[i],maxH)
		elif met == "AM":
			all_n_gm[i] = arithymetic_mean(all_n_wig[i])
		elif  met == "GM":
			all_n_gm[i] = geometric_mean(all_n_wig[i])
		elif met == "SNR":
			all_n_gm[i] = snr(all_n_wig[i])
	return(all_p_gm, all_n_gm)

def check_alignment(samfiles,chrom,chrom_size):
	count = 0
	for f in samfiles:
		try:
			f.fetch(chrom,0,chrom_size)
		except:
			count +=1
			continue
	if count == len(samfiles):
		return False	#all samfiles don't have alignments in this chromosome
	else:
		return True

def NucleoCompBias(seqs,word_size=6):
	'''Calculate weights to correct nucleotide composition bias'''
	from random import randint
	first_hex_frq = collections.defaultdict(int)
	rand_hex_frq = collections.defaultdict(int)
	
	all_hexamer = FrameKmer.all_possible_kmer(word_size)	#only A,C,G,T no "N
	
	weights={}
	for seq in seqs:
		first_hex_frq[seq[0:word_size]] +=1.0
		
		rand_starts=[]
		sampleSize = len(seq) - word_size*3
		if sampleSize <1:
			print >>sys.stderr, "Lenght of reads must be larger than 2 * word_size"
			sys.exit(1)
		for i in xrange(sampleSize):
			rand_starts.append(randint(word_size*2, len(seq) - word_size))
		for randstart in rand_starts:
			rand_hex_frq[seq[randstart:randstart+word_size]] += 1.0
	
	for hexmer in all_hexamer:
		#print hexmer + '\t' + str( first_hex_frq[hexmer]),
		#if not first_hex_frq.has_key(hexmer):
		#	weights[hexmer]=0.0
		#	continue
		#if not rand_hex_frq.has_key(hexmer):
		#	weights[hexmer]=0.0
		try:
			weights[hexmer] = rand_hex_frq[hexmer] / (first_hex_frq[hexmer] * sampleSize)
		except:
			weights[hexmer] = 1.0
	return weights
			
def main():
	usage="%prog [options]"
	parser = OptionParser(usage,version="%prog " + __version__)
	parser.add_option("-i","--inputFile",action="store",type="string",dest="input_file",help="Input file in BAM format. BAM file must be sorted and indexed using samTools. Replicates separated by comma(',') e.g. \"-i rep1.bam,rep2.bam,rep3.bam\"")
	parser.add_option("-r","--chromSize",action="store",type="string",dest="chromSize",help="Chromosome size file. Tab or space separated text file with 2 columns: first column is chromosome name, second column is size of the chromosome.")
	parser.add_option("-o","--outPrefix",action="store",type="string",dest="output_prefix",help="Prefix of output wig files(s). \"Prefix_Forward.wig\" and \"Prefix_Reverse.wig\" will be generated")
	parser.add_option("-w","--kmerSize",action="store",type="int",dest="word_size",default=6,help="Kmer size [6,12] to correct nucleotide composition bias. kmerSize < 0.5*read_lenght. larger KmerSize might make program slower. Set kmerSize = 0 to turn off nucleotide compsition bias correction. default=%default")
	parser.add_option("-b","--bin",action="store",type="int",dest="bin",default=100000,help="Chromosome chunk size. Each chomosome will be cut into small chunks of this size. Decrease chunk size will save more RAM. default=%default (bp)")
	parser.add_option("-d","--depth",action="store",type="int",dest="refReadN",default=10000000,help="Reference reads count (default = 10 million). Sequencing depth will be normailzed to this number, so that wig files are comparable between replicates.")
	parser.add_option("-q","--qCut",action="store",type="int",dest="qual_cut",default=30,help="phred scaled mapping quality threshhold to determine \"uniqueness\" of alignments. default=%default")
	parser.add_option("-m","--method",action="store",type="string",dest="norm_method",default="EM",help='methods ("EM", "AM", "GM", or "SNR") used to consolidate replicates and reduce noise. "EM" = Entropy weighted mean, "AM"=Arithmetic mean, "GM"=Geometric mean, "SNR"=Signal-to-noise ratio. default=%default')
	
	(options,args)=parser.parse_args()

	if not (options.output_prefix and options.chromSize and options.input_file):
		print logo
		parser.print_help()
		sys.exit(0)
	if not os.path.exists(options.chromSize):
		print >>sys.stderr, '\n\n' + str(options.chromSize) + " does NOT exists" + '\n'
		parser.print_help()
		sys.exit(0)
	if options.refReadN <10:
		print >>sys.stderr, '\n\n Reference reads count too small \n'
		parser.print_help()
		sys.exit(0)	
	if options.word_size >12 or (options.word_size < 6 and options.word_size != 0) :
		print >>sys.stderr, '\n\n Specify kmer size between 6 and 12. \n'
		parser.print_help()
		sys.exit(0)
	if options.norm_method not in ("AM","EM","GM","SNR"):
		print >>sys.stderr, '\n\n Specify "AM","EM","GM" or "SNR" to "-m" \n'
		parser.print_help()
		sys.exit(0)		
	bam_files = options.input_file.split(',')
	ReplicateNum = len(bam_files)
	maxEntropy = mystat.shannon_entropy([1]*ReplicateNum)
	for f in bam_files:
		if not os.path.exists(f):
			print >>sys.stderr, str(f) + " does not exists"
			parser.print_help()
			sys.exit(0)
		if not os.path.exists(f + '.bai'):
			print >>sys.stderr, str(f) + '.bai' + " does not exists"
			parser.print_help()
			sys.exit(0)
	chrom_sizes = load_chromsize(options.chromSize)	
	
	print >>sys.stderr, logo
	
	if options.word_size >=6 and options.word_size <= 12:
		print >>sys.stderr, "Calculate weights to correct nucleotide composition bias ... "
		a=SAM.ParseBAM(bam_files[0])	#this bias is reproducible between replicates
		b=a.batchIterBam(q_cut = options.qual_cut)
	
		my_weights = NucleoCompBias(b,word_size = options.word_size)
	elif options.word_size == 0:
		print >>sys.stderr, "Skip nucleotide composition bias correction ... "
		my_weights = None
	
	#for hex in sorted(my_weights):
	#	print  >>sys.stderr, hex + '\t' + str(my_weights[hex])
	#for read in open('Reb1_orphan.reads','r'):
	#	read = read.strip()
	#	print read + '\t' + str(my_weights[read[0:options.word_size]])
	
	print >>sys.stderr, "Normalize sequencing depth ... "
	samfile_objs = []
	norm_factors = []
	for f in bam_files:
		reads_sum = 0
		obj = SAM.ParseBAM(f)
		reads_sum = obj.countReads(chrom_sizes = chrom_sizes,q_cut=options.qual_cut)
		print >>sys.stderr, "\n\n" + f + ' ' + "has unique reads:" + str(reads_sum) + '\n'
		try:
			norm_factor = float(options.refReadN/reads_sum)
		except:
			norm_factor = 1.0
		norm_factors.append(norm_factor)
		
		samfile_objs.append(pysam.Samfile(f,'rb'))

	FWOUT = open(options.output_prefix + "_Forward.wig",'w')
	RWOUT = open(options.output_prefix + "_Reverse.wig",'w')
	
	print >>sys.stderr, "Writing wig files ... "
	for chr_name, chr_size in chrom_sizes.items():		#iterate each chrom
		if not check_alignment(samfile_objs,chr_name,chr_size):
			continue
		print >>sys.stderr, "Writing " + chr_name + " ..."
		FWOUT.write('variableStep chrom='+chr_name+'\n')
		RWOUT.write('variableStep chrom='+chr_name+'\n')
		for interval in BED.tillingBed(chrName = chr_name,chrSize = chr_size,stepSize = options.bin):	#cut chrom into bins, interval such as ('chr1', 235000000, 236000000)
			(Fwig,Rwig) = build_wig(samfile_objs,chr_name,interval[1],interval[2], nfs = norm_factors,wghts=my_weights,kmer=options.word_size,maxH = maxEntropy,met=options.norm_method, qcut=options.qual_cut)
			for i,j in Fwig.items():
				FWOUT.write("%d\t%.4f\n" % (i, j)) 
			for i,j in Rwig.items():
				RWOUT.write("%d\t%.4f\n" % (i, j))
	try:
		print >>sys.stderr, "Converting wig into bigwig format ... "
		subprocess.call("wigToBigWig " + options.output_prefix + "_Forward.wig " + options.chromSize + ' ' +  options.output_prefix + "_Forward.bw ",shell=True)
		subprocess.call("wigToBigWig " + options.output_prefix + "_Reverse.wig " + options.chromSize + ' ' +  options.output_prefix + "_Reverse.bw ",shell=True)
	except:
		 print >>sys.stderr, "Cannot call \"wigToBigWig\". Conversion need to be done by yourself."
		 pass

if __name__ == '__main__':
	main()

