#!/usr/bin/python 

# -------- basicDistributions.py ------------------------------------------
# Description: example showing 
#              - how to define a function
#              - use matplotlib to plot histograms and function
# Author:      G. Quast   Dec. 2013
# dependencies: PYTHON v2.7, numpy, scipy.special, matplotlib.pyplot 
# last modified: 
#  
#--------------------------------------------------------------

import numpy as np, matplotlib.pyplot as plt, scipy.special as sp

#Binomial distribution
def fBinomial(x,n,p):
  k=np.around(x)
  return sp.binom(n,k) * p**k * (1.-p)**(n-k)

# Poisson distribution
def fPoisson(x,mu):
  k=np.around(x)
  return (mu**k)/np.exp(mu)/sp.gamma(k+1.)

# Gauss distribution
def fGauss(x,mu=0.,sigma=1.):
    return (np.exp(-(x-mu)**2/2./sigma**2)/np.sqrt(2.*np.pi)/sigma)

# ---------------------------------------------------------------

# plot basic distributions
mu=20.
sigma=np.sqrt(mu)
xmx=mu+3.5*sigma
xmn=mu-3.5*sigma
x = np.arange(xmn, xmx, 0.1)

plt.xkcd()
fig = plt.figure(figsize=(5.,3.) )
ax = fig.add_subplot( 1, 1, 1)
ax.set_ylim(0.,0.08)
ax.plot(x, np.sqrt(x)/8.*fGauss(x,mu,sigma), 'r-', 
  label='unbekannte Verteilung')   
#     make plot nicer:
ax.set_xlabel('t') # axis labels
ax.set_ylabel('probability density')

t = mu + sigma * np.random.randn(25)
ax.plot(t, np.ones(len(t))*0.005, 'b|', markersize=20, label='Stichprobenwerte')
plt.legend(loc='best')

# create bootstrapped samples
t_bootstrap=[]
labels=[]
nbs = 5

t_bootstrap.append(t) 
labels.append('original')
for i in range(nbs):
  t_bootstrap.append( np.random.choice(t, size=len(t)) )
  labels.append('sample %i'%(i+1))        
bins = np.linspace(min(t)-1, max(t)+1, (max(t)-min(t)+3)*10)

# 3d-plot of resamplings
from mpl_toolkits.mplot3d import Axes3D
fig2d = plt.figure(figsize=(7.5, 7.5) )
ax2d = fig2d.add_subplot(1,1,1, projection='3d')

for i in range(nbs+1):
   bc, be = np.histogram(t_bootstrap[i], bins, density=True)
   ax2d.bar(bins[:-1], bc, zs=i*20, zdir='y', alpha=0.5)
 
# summed distribution of als samples
fig2 = plt.figure(figsize=(7.5, 7.5) )
ax2 = fig2.add_subplot( 1, 1, 1)
ax2.hist(t_bootstrap, bins, density=True, stacked=True, label=labels,
          align='mid', alpha=0.8) 
#     make plot nicer:
ax2.set_xlabel('t') # axis labels
ax2.set_ylabel('frequency from bootstrapped samples')

plt.legend(loc='best')
# finally, show the plots on the screen
plt.show()