# Buys's statsistics doodles
import scipy.stats
import numpy as np
import collections
import sklearn.metrics
import matplotlib.pyplot as plt

def sampleFromCounter(counter, n, replace=True, outputCounter=False):
	counterSum = sum(counter.values())
	if replace:
		samples = np.random.choice(list(counter.keys()), n, p=[ float(counter[className])/float(counterSum) for className in counter ])
		if outputCounter:
			#@todo optimize
			collections.Counter( samples )
		else:
			return(  samples )
	if counterSum<n:
		raise ValueError('Cannot pick %s samples from a list of %s samples' % (n, counterSum))

	samples = []

	if outputCounter:
		samples=collections.Counter()
	for i in range(0,n):
		counterSum = sum(counter.values())
		classProbs = [ float(counter[className])/float(counterSum) for className in counter ]
		s = list( np.random.choice( list(counter.keys()), 1, p=classProbs) )[0]
		if outputCounter:
			samples[s]+=1
		else:
			samples.append( s )
		counter[s]-=1
	return(samples)


class IterativeHistogram(object):

	def __init__(self, minV, maxV, bins):
		self.binCount = bins
		self.min = minV
		self.max = maxV
		self.counts = np.zeros( self.binCount )

		self.binStarts = np.linspace(self.min, self.max, self.binCount,endpoint=False)
		self.binSize = self.binStarts[1] - self.binStarts[0]

	def add(self, vector):
		try:
			locs = np.searchsorted(self.binStarts, vector,'right')-1
			#print('input: %s' %vector)
			#print('point: %s'%locs)
			#print('starts:%s' % self.binStarts)

			np.add.at( self.counts, locs, 1)
			print(self.counts)
		except:

			exit('Failed adding vector to histogram')
		return(self)

	def plot(self,path=None):
		fig, ax = plt.subplots() #figsize=(120, 10))
		for i in range(0,self.binCount):
			plt.bar(self.binStarts[i],self.counts[i], float(self.max-self.min)/float(self.binCount))

		plt.ylabel("Density")
		plt.xlabel("PIC score")
		plt.title('PIC score distribution')
		ax.legend(loc='upper right')
		plt.yscale('log', nonposy='clip')
		if path==None:
			plt.show()
		else:
			plt.savefig(path)

	def __add__(self, histogramToAdd):
		n = IterativeHistogram(self.min, self.max, self.binCount)
		n.counts += self.counts
		n.counts += histogramToAdd.counts
		return(n)

	def p(self, value):
		return( sum( self.counts[ np.searchsorted(self.binStarts, value,'right'): ] )/sum(self.counts) )


	def getScipyRandomVariable(self):

		values = self.binStarts
		probabilities = self.counts/np.sum(self.counts)
		X = scipy.stats.rv_discrete(values=(values, probabilities))
		return(X)


	def __iter__(self):
		self.iterIndex = 0
		return(self)

	def __next__(self):
		value = self.min
		index = 0
		for i in range(0,self.binCount):
			index+=self.counts[i]

			if index> self.iterIndex:
				self.iterIndex+=1
				return( self.binStarts[i]  ) #+ 0.5*self.binSize
		raise StopIteration

def calc_MI(x, y, bins):
    c_xy = np.histogram2d(x, y, bins)[0]
    mi = sklearn.metrics.mutual_info_score(None, None, contingency=c_xy)
    return mi

## Example of Chi-squared test between two distributions:
#
#
# observedValues = [16, 18, 16, 14, 12, 12]
# expectedValues = [16, 16, 16, 16, 16, 8]
# binCount = 20
#
#
# expectedDistribution = list( IterativeHistogram(0,20, binCount).add(expectedValues) )
# print(expectedDistribution)
#
# observedDistribution = list( IterativeHistogram(0,20, binCount).add(observedValues) )
#
# print( scipy.stats.chisquare( observedDistribution, f_exp=expectedDistribution).pvalue )
#
# print( scipy.stats.ks_2samp( observedDistribution, expectedDistribution ))
# expectedDistribution = IterativeHistogram(0,20, binCount).add(expectedValues).getScipyRandomVariable()
# observedDistribution = IterativeHistogram(0,20, binCount).add(observedValues).getScipyRandomVariable()
#
# print( scipy.stats.kstest( observedDistribution, expectedDistribution ))