import math
import numpy
import matplotlib.pyplot as plt

def FillHisto(x, histo, xMin, xMax, nBins):
	"""Adds one entry in a histogram (represented by an array)"""
	dx = (xMax-xMin) / nBins
	i  = int(numpy.floor((x - xMin) / dx))
	if i >= 0 and i < nBins:
		histo[i] = histo[i] + 1

def P_poisson(k,nu):
	"""function to return Poissonian probablity"""
	P = numpy.exp(-nu)
	for i in range(1,k+1):
		P *= nu / i
	return P

def P_binomial(k,n,p):
	"""function to return binomial probablity"""
	P = 1.0
	for i in range(1,n+1):
		if i <= k and i <= n-k:
			P = P / i
		elif i > k and i > n-k:
			P = P * i
	return P * pow(p,k) * pow(1-p,n-k)


nCyc  = 1000
nGen  = 100
nBins = 10
iBinRes  = int(nBins/2)
a = 0
b = 1

binEntries  = numpy.empty(nCyc)
res = numpy.zeros(nBins)

# binning and proability density of x:
dx = (b-a) / nBins
x  = numpy.empty(nBins)		# array of bin centres
fx = numpy.empty(nBins)		# probability density at bin centres
for i in range(nBins):
	x[i]  = a + dx * (i+0.5)
	fx[i] = 2 / (b-a) * (x[i] - a) / (b-a)
#	sigma2 = pow((b-a)/4,2)
#	fx[i] = 0.5 * (math.erf((a+dx*(i+1)-(a+b)/2)/numpy.sqrt(2*sigma2))\
#			- math.erf((a+dx*i-(a+b)/2)/numpy.sqrt(2*sigma2))) / dx

fig, ((fig1,fig2), (fig3,fig4)) = plt.subplots(2,2,figsize=(12, 8))

for j in range(nCyc):
	hist = numpy.zeros(nBins)
	for i in range(nGen):
		xGen = a + numpy.sqrt(numpy.random.rand()) * (b-a)
#		xGen = numpy.random.normal((a+b)/2,(b-a)/4)
		FillHisto(xGen,hist,a,b,nBins)
	if j == 0:
		fig1.bar(x,hist,dx,yerr=numpy.sqrt(hist))
		fig1.plot(x,fx*nGen*dx,color="red")
		fig1.set_xlabel("x")
	binEntries[j] = hist[iBinRes]
	for i in range(nBins):
		res[i] = res[i] + pow(hist[i] - fx[i]*nGen*dx, 2)

nExp = fx[iBinRes]*nGen*dx
nMin = max(int(nExp - 3*numpy.sqrt(nExp)), 0)
nMax = int(nExp + 3*numpy.sqrt(nExp))
fig2.hist(binEntries, nMax-nMin, range=(nMin-0.5,nMax-0.5),label="histogram")

pois  = numpy.empty(nMax-nMin)
binom = numpy.empty(nMax-nMin)

for i in range(nMax-nMin):
	pois[i]  = P_poisson(nMin+i,nExp) * nCyc
	binom[i] = P_binomial(nMin+i,nGen,fx[iBinRes]*dx) * nCyc

fig2.plot(range(nMin,nMax), pois, label="Poissonovo rozdělení")
fig2.plot(range(nMin,nMax), binom, label="binomické rozdělení")
fig2.set_xlabel("Počet případů v " + str(iBinRes) + ". binu")
fig2.legend()

res = numpy.sqrt(res / nCyc)
fig3.plot(range(nBins),numpy.sqrt(fx*dx*nGen),color='orange',\
	label="$\sqrt{pn_{gen}} = \sqrt{n_{exp}}$")
fig3.plot(range(nBins),numpy.sqrt(fx*dx*(1-fx*dx)*nGen),color='green',\
	label="$\sqrt{p(1-p)n_{gen}}$")
fig3.scatter(range(nBins),res,marker='o')
fig3.set_xlabel("bin")
fig3.set_ylabel('standardní odchylka počtu případů v binu')
fig3.legend()

plt.show()