#!/usr/bin/env python
# coding=utf8

"""
%prog <BAM> <BED> <TXT> <INT> <BAM,...> [options]

Based on the input <BAM> file DE-peaks will be simulated in the provided <BED> regions, <TXT> chromosmes to be used and their length tab separated, 
<INT> defines the number of simulated replicates of the two samples,
the input/control Bam files <BAM,...> comma-separated, will be used to create an input/control Bam file for the simulation.
@author: Thomas Eder
"""
from __future__ import print_function
from __future__ import division

import matplotlib.pyplot as plt
import numpy as np
import scipy as sp
import pysam
import os
import sys

from scipy.stats import truncnorm
from future.utils import lrange
from random import random
from random import randrange
from pybedtools import BedTool
from optparse import OptionParser
from numpy.random import dirichlet

class HelpfulOptionParser(OptionParser):
	"""An OptionParser that prints full help on errors."""
	def error(self, msg):
		self.print_help(sys.stderr)
		self.exit(2, "\n%s: error: %s\n" % (self.get_prog_name(), msg))

def sort_index(file_name):
	"""Sort and index Bam-file"""
	pysam.sort("-o",'tmp.bam',file_name)
	os.rename('tmp.bam',file_name)
	pysam.index(file_name)

def get_truncated_normal(mean=0, sd=1, low=0, upp=10):
	"""Get truncnorm object with defined mean and sd, it represents a truncated normal distribution"""
	return truncnorm((low - mean) / sd, (upp - mean) / sd, loc=mean, scale=sd)
	
def _callback_list_float(option, opt, value, parser): #get a list of float options and save it
	try:
		tmplist = list(map(lambda x: float(x), value.split(',')))
		setattr(parser.values, option.dest, tmplist)  
	except Exception as e:
		print("got exception: %r, please check input parameters" % (e,))
		exit()

def _scale_laplace(nfrags, read_frac, scale):
	"""Scale beta result based on Laplace distribution"""
	loc = 0.5 #fixed
	
	top_nfrgs_scale = laplace.pdf(0.5, loc, scale)#make sure at postion 0.5 is 100%
	scale_factor = 1 / top_nfrgs_scale
	nfrags_scale = laplace.pdf(read_frac, loc, scale) *scale_factor #the bigger the difference from beta the smaller the sample...
	
	if(nfrags_scale == 0):
		nfrags = nfrags 
	else:
		nfrags = int(nfrags *nfrags_scale) #new fragment scaling based on beta result
		
	return nfrags

def _scale_lognorm(nfrags, read_frac, sigma, loc, scale):
	"""Scale number_frags result based on Lognorm distribution"""
	#loc=scale*-1
	
	top_beta_scale = sp.stats.lognorm.pdf(1,s=sigma,loc=loc, scale=scale)#to get to 100% at size 10
	scale_factor=1/top_beta_scale #this is the factor to set beta_frac to 100% at 1
	beta_scale = sp.stats.lognorm.pdf(nfrags,s=sigma,loc=loc, scale=scale) *scale_factor
	
	if(beta_scale == 0):
		read_frac = read_frac 
	else:
		read_frac = ((read_frac-0.5)*beta_scale)+0.5 #center and scale it based on beta_scale from exponential distri, so reduce to 0.5 for high values
	return read_frac

def _scale_exp(nfrags, read_frac, loc, scale):
	"""Scale number_frags result based on Exponential distribution"""
	
	top_beta_scale = sp.stats.expon.pdf(loc,loc=loc, scale=scale)#to get to 100% at size loc
	scale_factor=1/top_beta_scale #this is the factor to set beta_frac to 100% at 1
	beta_scale = sp.stats.expon.pdf(nfrags,loc=loc, scale=scale) *scale_factor
	
	read_frac_old = read_frac
	
	if(beta_scale == 0):
		read_frac = read_frac 
	else:
		read_frac = ((read_frac-0.5)*beta_scale)+0.5 #center and scale it based on beta_scale from exponential distribution, we do this to reduce the read_frac towards 0.5 for high values
	
	return read_frac


def _reads_to_replicates(read_list, skewness, n_replicates):
	"""puts reads in a randomly selected replicate"""
	#get reads into replicates
	bins = np.cumsum([0] + list(dirichlet(np.array([skewness] * n_replicates), 1)[0])) #sample dirichlet dist, get array, make list, add 0 to beginning and sum up all values to 1, so we end up with an array of bins with probabilities
	replicate_counts = [0] * n_replicates
	replicate_list = [[] for i in lrange(n_replicates)]
	
	for readobj in read_list: #go through reads
		index = np.digitize([random()], bins)[0].item() - 1 #get index-1 of random number (between 0 or 1) fitting to the right bin
		replicate_list[index].append(readobj) #add element to replicate via choosen index
		replicate_counts[index] += 1
	
	return replicate_list, replicate_counts #return number of reads per replicate

#main function
if __name__ == '__main__':
	
	#handle input arguments
	parser = HelpfulOptionParser(usage=__doc__)
	parser.add_option("-c", "--chrom", default='chr19', dest="chromosome", type="string", help="Chromosome used for simulation [default: %default]")
	parser.add_option("-r", "--read_length", default=50, dest="read_length", type="int", help="Read length [default: %default]")
	
	parser.add_option("-b","--beta", default=[0.5, 0.5], dest="beta_values", type="string", action='callback', callback=_callback_list_float, help="Alpha and Beta of Beta-distribution [default: %default]")
	
	parser.add_option("--max_reads", default=1, dest="max_reads", type="float", help="[Non peak reagion] Maximum percentage of reads for repcilates [default: %default]")
	parser.add_option("--min_reads", default=0, dest="min_reads", type="float", help="[Non peak reagion] Minimum percentage of reads for repcilates [default: %default]")
	parser.add_option("--rep_sd", default=0.1, dest="rep_sd", type="float", help="[Non peak reagion] Standard deviation of replicate read numbers [default: %default]")
	parser.add_option("--rep_mean", default=0.9, dest="rep_mean", type="float", help="[Non peak reagion] Mean of replicate read numbers [default: %default]")
	
	parser.add_option("--frag-count-scaling", default="none", dest="frag_count_scaling", type="string", help="Scaling of fragment distribution, no scaling, scaling of beta result based on fragment counts (with exp) or scaling of fragment counts based on beta result (with laplace) : none , frag , beta [default: %default]")
	parser.add_option("--frag-count-lp-scale", default=0.1, dest="frag_count_lp_sc", type="float", help="Scale for Laplace distribution if frag-count-scaling is beta [default: %default]")
	parser.add_option("--frag-count-ex-loc", default=10, dest="frag_count_ex_lo", type="float", help="Loc for exponential distribution if frag-count-scaling is frag [default: %default]")
	parser.add_option("--frag-count-ex-scale", default=100, dest="frag_count_ex_sc", type="float", help="Scale for exponential distribution if frag-count-scaling is frag [default: %default]")
	
	parser.add_option("-d", "--dp-thres", default=0.7, dest="dp_thres", type="float", help="Threshold of reads to define a DB peak [default: %default]")
	parser.add_option("-m", "--min-counts", default=10, dest="min_counts", type="int", help="Minimum number of reads for a DB peak [default: %default]")
	
	parser.add_option("-n", "--non_peak_percetage", default=0.9, dest="noise_treshold", type="float", help="Percentage of reads outside the defined peak regions [default: %default]")
	
	parser.add_option("-s", "--skewness", default=10, dest="skewness", type="float", help="Variance between replicates (the higher, the less variance) [default: %default]")
	
	parser.add_option("--control_read_num", default=0, dest="num_control_reads", type="int", help="Number of reads in the input control Bam files, if 0 use mean of input control files [default: %default]")
	
	parser.add_option("--out_bed", default="results.bed", dest="outbed_filename", type="string", help="Name of the results bed file [default: %default]")
	parser.add_option("--out_sample", default="out", dest="output_name", type="string", help="Basename of the results Bam files [default: %default]")
	parser.add_option("--out_control", default="out", dest="output_control_name", type="string", help="Basename of the input/control Bam files [default: %default]")
	
	(options, args) = parser.parse_args()
	num_args = 5
	
	#read in required input arguments
	if len(args) != num_args:
		parser.error("At minimum %s parameters are requred to run this tool" % num_args)
	input_bam_file = args[0]
	inbed_filename = args[1]
	in_chromsizes = args[2]
	num_replicates = int(args[3])
	input_control_bam_files_string = args[4]
	input_control_bam_files = input_control_bam_files_string.split(",")

	#get sample bam file...
	input_bam = pysam.AlignmentFile(input_bam_file,'rb')
	subsample_chrom = options.chromosome 

	#set up output bams
	output_s1_bams=[]
	output_s2_bams=[]
	
	for x in lrange(1,num_replicates+1):
		output_s1_bams.append(pysam.AlignmentFile(options.output_name+"_sample1-rep"+str(x)+".bam", "wb", template=input_bam))
		output_s2_bams.append(pysam.AlignmentFile(options.output_name+"_sample2-rep"+str(x)+".bam", "wb", template=input_bam))
	
	#read bed
	input_bed = BedTool(inbed_filename)
	#open output bed
	output_bed = open(options.outbed_filename+".bed", "w")
	#header
	output_bed.write("#Chromsome\tStart\tStop\tSample1_mean_counts\tSample2_mean_counts%s%s\tStatus\tPeak_name\n" %( \
	"".join(list(map(lambda x: '\tSample1_replicate%i_counts' % x, lrange(1,num_replicates+1)))),\
	"".join(list(map(lambda x: '\tSample2_replicate%i_counts' % x, lrange(1,num_replicates+1))))))

	region_count=1
	
	#normal distributions sampling with max arround num_reads_sample to get num_reads_replicate
	get_num_reads_rep = get_truncated_normal(mean=options.rep_mean , sd=options.rep_sd, low=options.min_reads, upp=options.max_reads)
	num_reads_s1rep = list(get_num_reads_rep.rvs(num_replicates))
	get_num_reads_rep = get_truncated_normal(mean=options.rep_mean , sd=options.rep_sd, low=options.min_reads, upp=options.max_reads)
	num_reads_s2rep = list(get_num_reads_rep.rvs(num_replicates))

	#for each bed entry get reads...
	sample1_replicate_counts = [0] * num_replicates
	sample2_replicate_counts = [0] * num_replicates
	sample1_replicate_list = [[] for i in lrange(num_replicates)]
	sample2_replicate_list = [[] for i in lrange(num_replicates)]
	multi_factor = 2*num_replicates #per read multiply its use, so that we can definitely fill 2 samples with x replicates
	
	for region in input_bed:
		sample1_list = []
		sample2_list = []
		sample1_count = 0
		sample2_count = 0
		
		#beta distribution to get aprox. same size or down regulated read number for sample x, here we deside which sample to prefer
		read_frac = np.random.beta(options.beta_values[0], options.beta_values[1])
		
		#get number of reads in this region, name id nfragments like in DCSsim
		number_frags = input_bam.count(region.chrom, region.start, region.end)
		
		#now choose which way we want to select the number of fragments
		if( options.frag_count_scaling == "none"):
			#no scaling, no changes in nfrags and beta
			number_frags = number_frags
			read_frac = read_frac
		elif (options.frag_count_scaling == "beta") :#nfrags scaling via Laplace based on beta result, beta not changed
			number_frags = _scale_laplace(number_frags, read_frac, options.frag_count_lp_sc)
		elif (options.frag_count_scaling == "frag") :#nfrags scaling via exp distribution based on fragment counts, number_frags not changed
			read_frac = _scale_exp(number_frags, read_frac, options.frag_count_ex_lo, options.frag_count_ex_sc)
		else:
			print("Unknown scaling method, %s, please choose 'none','frag' or 'beta', exiting now" % (frag_count_scaling))
			exit(1)
		
		#for each region change height based on beta distribution
		read_counter=0
		read_pool=[]
		
		#first get all reads of this region
		for read in input_bam.fetch(region.chrom, region.start, region.end):
			read_pool.append(read)
		
		#then randomly select on of the reads for x-times
		for _ in lrange(number_frags):
			for _ in lrange(multi_factor): #do this x times
				if(read_counter > number_frags): #check if we meet calculated and scaled (if turned on) number of fragments/reads
					break
				read_counter+=1
				
				#get one random read out of region-pool
				random_index = randrange(len(read_pool))
				read = read_pool[random_index]
				
				if random() <= read_frac: #add to sample 1 or sample 2 and to list
					sample1_count += 1
					sample1_list.append(read)
				else:
					sample2_count += 1
					sample2_list.append(read)
			else:
				continue
			break
		
		#now decide to which replicate
		sample1_replicate_list, sample1_replicate_counts = _reads_to_replicates(sample1_list, options.skewness, num_replicates)
		sample2_replicate_list, sample2_replicate_counts = _reads_to_replicates(sample2_list, options.skewness, num_replicates)
		
		#write reads to files
		for repnum in lrange(num_replicates):
			for read in sample1_replicate_list[repnum]:
				output_s1_bams[repnum].write(read)
			for read in sample2_replicate_list[repnum]:
				output_s2_bams[repnum].write(read)
		
		#get mean
		sample1_mean_count = sample1_count / num_replicates
		sample2_mean_count = sample2_count / num_replicates
		
		#check if me meet threshold
		if sample1_mean_count > sample2_mean_count and sample1_mean_count > options.min_counts and (sample1_mean_count / (sample1_mean_count + sample2_mean_count)) >= options.dp_thres: #up in sample1
			stat=1
		elif sample2_mean_count > sample1_mean_count and sample2_mean_count > options.min_counts and (sample2_mean_count / (sample1_mean_count + sample2_mean_count)) >= options.dp_thres: #up in sample2
			stat=2
		else: #not meeting user defined DP parameters or is not up in either sample
			stat=0
		
		#id
		name='peak'+str(region_count)
		
		#now save data for this region
		output_bed.write("%s\t%s\t%s\t%.2f\t%.2f\t%s\t%s\t%s\t%s\n" %(region.chrom, region.start, region.end, sample1_mean_count, sample2_mean_count,\
		"\t".join(list(map(lambda x: str(x), sample1_replicate_counts))),  "\t".join(list(map(lambda x: str(x), sample2_replicate_counts))), stat, name))
		
		region_count+=1
	
	
	#now take all other regions and leave them as they are...
	complement_bed = input_bed.complement(g=in_chromsizes)
	
	for region in complement_bed:
		
		#check region if its to narrow
		if not (region.start+options.read_length+(2+options.read_length) >= region.end-options.read_length): #leave region out if to narrow we alread processes these reads...
			
			for read in input_bam.fetch(region.chrom, region.start+options.read_length, region.end-options.read_length): #decrease size of region by read length as reads with only 1 bp match will also be fetched for a spezific region
				#sample1
				if random() <= options.noise_treshold: #do not just take all, bring in some variance
					rep_cnt=0
					for rep_treshold in num_reads_s1rep:
						if random() <= rep_treshold:
							output_s1_bams[rep_cnt].write(read)
						rep_cnt+=1
				#sample2
				if random() <= options.noise_treshold:
					rep_cnt=0
					for rep_treshold in num_reads_s2rep:
						if random() <= rep_treshold:
							output_s2_bams[rep_cnt].write(read)
						rep_cnt+=1


	#input/control bam file(s)
	all_control_reads=[]

	#read all and store
	for input_control_bam_file in input_control_bam_files:
		input_control_bam = pysam.AlignmentFile(input_control_bam_file,'rb')
		for read in input_control_bam.fetch(options.chromosome):
			all_control_reads.append(read)
		input_control_bam.close()

	if options.num_control_reads == 0 : #if not spezified by user use mean number of reads in the input control files
		options.num_control_reads = len(all_control_reads)//len(input_control_bam_files)

	#select random and write one control file per sample 
	selected_control_reads = np.random.choice(all_control_reads, options.num_control_reads, replace=False)
	output_s1_control_bam = pysam.AlignmentFile(options.output_control_name+"_sample1-INPUT.bam", "wb", template=input_bam)
	for read in selected_control_reads:
		output_s1_control_bam.write(read)

	selected_control_reads = np.random.choice(all_control_reads, options.num_control_reads, replace=False)
	output_s2_control_bam = pysam.AlignmentFile(options.output_control_name+"_sample2-INPUT.bam", "wb", template=input_bam)
	for read in selected_control_reads:
		output_s2_control_bam.write(read)


	#close bam files and index new bam files
	input_bam.close()
	output_bed.close()

	for out_bam in output_s1_bams:
		out_bam.close()
		sort_index(out_bam.filename)

	for out_bam in output_s2_bams:
		out_bam.close()
		sort_index(out_bam.filename)

	#close, sort and index control files
	output_s1_control_bam.close()
	sort_index(output_s1_control_bam.filename)

	output_s2_control_bam.close()
	sort_index(output_s2_control_bam.filename)
	
	exit(0)
