import sys
import numpy as np
import numpy.random as npr
from sklearn.neighbors.kde import KernelDensity
from scipy.special import gammaln
import matplotlib.pyplot as plt
from calculate_phist import read_counts
from calculate_phist import normalize_haplotypes

def log_factorial(n):
	return gammaln(n+1)

def log_multinomial(xs, ps):
	n = np.sum(xs)
	log_prob = log_factorial(n) - np.sum(log_factorial(xs)) + np.sum(xs * np.log(ps + 0.0000000000001))
	return log_prob

class KDE_MCMC_Sampler(object):
	def __init__(self, observed_counts):
		"""
		Observed counts is 3D matrix of pop, locus, haplotype
		"""
		self.observed_counts = observed_counts
		self.individual_counts = observed_counts.sum(axis=2)
		self.observed_frequencies = normalize_haplotypes(observed_counts)

		self.n_loci, self.n_pop, self.n_haplotypes = self.observed_counts.shape
		
		# from bamova
		self.DWEIGHT = 1.0
		self.DADD = 0.00001
		self.SMALL_NUM = 0.0000000000001

		print "initializing frequencies"
		self.freq = np.zeros((self.n_loci, self.n_haplotypes))
		for l in xrange(self.n_loci):
			self.freq[l, :] = self.sample_locus_freq(self.observed_frequencies[l, 0, :])

	def sample_locus_freq(self, freq):
		alphas = self.DWEIGHT * freq + self.DADD + self.SMALL_NUM

		return npr.dirichlet(alphas)

	def locus_prob(self, locus_obs_counts, locus_freq):
		log_prob_sum = 0.0
		for p in xrange(self.n_pop):
			log_prob_sum += log_multinomial(locus_obs_counts[p], locus_freq)
		return log_prob_sum


	def step(self):
		total_log_prob = 0.0
		for l in xrange(self.n_loci):
			locus_indiv_counts = self.individual_counts[l, :]
			locus_obs_counts = self.observed_counts[l, :, :]

			log_prob = self.locus_prob(locus_obs_counts, self.freq[l, :])

			proposed_locus_freq = self.sample_locus_freq(self.freq[l, :])

			proposed_log_prob = self.locus_prob(locus_obs_counts, proposed_locus_freq)
				
			log_prob_ratio = proposed_log_prob - log_prob
			log_r = np.log(npr.random())

			if proposed_log_prob >= log_prob or log_r <= log_prob_ratio:
				self.freq[l, :] = proposed_locus_freq
				log_prob = proposed_log_prob

			total_log_prob += log_prob

		locus_prob = []
		for l in xrange(self.n_loci):
			log_prob = self.locus_prob(locus_obs_counts, self.freq[l, :])
			locus_prob.append(log_prob)

		return self.freq, total_log_prob, locus_prob

def plot_log_prob(flname, log_probs):
	plt.clf()
	plt.hold(True)
	plt.hist(log_probs, bins=30)
	plt.xlabel("Log Probability", fontsize=16)
	plt.xlim([min(log_probs), 0.0])
	plt.ylabel("Occurrences (Loci)", fontsize=16)
	plt.savefig(flname, DPI=200)

def simulate(occur_fl, n_steps, plot_flname, prob_flname):
	print "reading occurrences"
	observed_counts = read_counts(occur_fl)
	individual_counts = observed_counts.sum(axis=2)
	observed_frequencies = normalize_haplotypes(observed_counts)

	sampler = KDE_MCMC_Sampler(observed_counts)

	fl = open(prob_flname, "w")

	locus_log_prob = []
	for i in xrange(n_steps):
		freq, log_prob, locus_log_prob = sampler.step()
		print "step", i, "log prob", log_prob

		if i % 100 == 0:
			for j, prob in enumerate(locus_log_prob):
				fl.write("%s %s %s\n" % (i, j, prob))

	fl.close()


	plot_log_prob(plot_flname, locus_log_prob)

if __name__ == "__main__":
	occur_fl = sys.argv[1]
	n_steps = int(sys.argv[2])
	plot_flname = sys.argv[3]
	prob_flname = sys.argv[4]

	simulate(occur_fl, n_steps, plot_flname, prob_flname)

	