#!/usr/bin/env python

# ----------------------------------------------------------------------------------- #
#
#  Python/ROOT macro for illustrating the Binomial, Poisson and Gaussian distributions.
#
#  Reference:
#    Barlow 
#
#  Run this pyROOT macro by:
#    > ./BinomialPoissonGaussian.py
#
#  Author: Troels C. Petersen (NBI)
#  Email:  petersen@nbi.dk
#  Date:   7th of September 2014
#
# ----------------------------------------------------------------------------------- #

from ROOT import *
import math

r = TRandom3()                        # Random generator
r.SetSeed(1)

SavePlots = False
verbose = True
Nverbose = 5

Nexperiments = 1000


#----------------------------------------------------------------------------------
# Loop over process:
#----------------------------------------------------------------------------------

# Define process, i.e. number of trials and probability of success:
Ntrials = 10
pSuccess = 1.0/6.0
Lambda = Ntrials * pSuccess     # This is the mean and also how the Poisson parameter is defined!

Hist_nSuccess = TH1F("Hist_nSuccess", ";Number of successes;Frequency", Ntrials+1, -0.5, Ntrials+0.5)

for iexp in range( Nexperiments ) : 
    
    # Simulating process defined:
    nSuccess = 0
    for i in range( Ntrials ) : 
        x = r.Uniform()
        if (x < pSuccess) : nSuccess = nSuccess + 1

    # Record result:
    if (verbose and iexp < Nverbose) : print "  nSuccess: %4d"%(nSuccess)
    Hist_nSuccess.Fill(nSuccess)




#---------------------------------------------------------------------------------- 
# Plot result:
#---------------------------------------------------------------------------------- 

canvas = TCanvas( "canvas", "", 50, 50, 1200, 600 )
# canvas.SetLogy()          # Set the y-axis to be logarithmic

# Draw the histogram:
# Hist_nSuccess.GetXaxis().SetRangeUser(-0.5, 10.5)
Hist_nSuccess.SetLineWidth(3)
Hist_nSuccess.SetLineColor(kBlack)
Hist_nSuccess.Draw("e")


# Draw Binomial:
# --------------
def func_Binomial(x, p) :
    if (x[0] > -0.5) :
        return p[0] * TMath.Binomial(int(p[1]), int(x[0]+0.5)) * TMath.Power(p[2],int(x[0]+0.5)) * TMath.Power(1-p[2], int(p[1])-int(x[0]+0.5))
    else : return 0.0

fBinomial = TF1("fBinomial", func_Binomial, -0.5, Ntrials+0.5, 3)
fBinomial.SetParameters(Nexperiments, Ntrials, pSuccess)
fBinomial.SetNpx(1000)        # Set the number of intervals used when drawing function!
fBinomial.SetLineColor(kBlue)
fBinomial.Draw("same")

# Draw Poisson:
# -------------
def func_Poisson(x, p) :
    if (x[0] > -0.5) :
        return p[0] * TMath.PoissonI(x[0]+0.5, p[1])
    else : return 0.0

fPoisson = TF1("fPoisson", func_Poisson, -0.5, Ntrials+0.5, 2)
fPoisson.SetParameters(Nexperiments, Lambda)
fPoisson.SetNpx(1000)        # Set the number of intervals used when drawing function!
fPoisson.SetLineColor(kMagenta)
fPoisson.Draw("same")

# Draw Gaussian:
# -------------
def func_Gaussian(x, p) :
    z = ((x[0]) - p[1]) / p[2]
    return p[0] / sqrt(2*TMath.Pi()) / p[2] * exp(-0.5 *z*z)

fGaussian = TF1("fGaussian", func_Gaussian, -0.5, Ntrials+0.5, 3)
fGaussian.SetParameters(Nexperiments, Lambda, sqrt(Lambda))
fGaussian.SetNpx(1000)        # Set the number of intervals used when drawing function!
fGaussian.SetLineColor(kRed)
fGaussian.Draw("same")


# Legend:
leg = TLegend( 0.60, 0.60, 0.89, 0.75 )
leg.SetFillColor(kWhite)
leg.SetLineColor(kWhite)
leg.AddEntry(Hist_nSuccess, " Distribution of nSuccess", "L")
leg.AddEntry(fBinomial, " Binomial distribution", "L")
leg.AddEntry(fPoisson,  " Poisson distribution", "L")
leg.AddEntry(fGaussian, " Gaussian distribution", "L")
leg.Draw()

if (SavePlots) : canvas.SaveAs("BinomialPoissonGaussian.pdf")


raw_input('Press Enter to exit')


#---------------------------------------------------------------------------------- 
# 
# Make sure you understand what process yields a Binomial, a Poisson and a Gaussian
# distribution.
# 
# The program repeats an experiment, where one performs N trials, each with p chance
# of success, and then considers the distribution of successes.
# 
# 
# Questions:
# ----------
#  1) Which distribution has the longest tail towards high values? And which one has
#     the longest tail the other way? Possibly use a logarithmic plot.
# 
#  2) Using Nexperiments=1000, Ntrials=20 and pSuccess=1/6, calculate "by hand" (i.e.
#     using a loop) the ChiSquare between the data and each of the three distributions.
#     Do they give a reasonable Chi2 value?
#     Retry with Ntrials 200 and pSuccess=1/60. Does the degree of fit improve for any
#     of the three?
#     Retry again with Ntrials 200 and pSuccess=1/12. Again, does the degree of fit 
#     improve further for any of the three?
# 
#  3) Using Nexperiments=1000, Ntrials=200 and pSuccess=1/12, is the skewness consistent
#     with zero (as the Gaussian should have)? Does that contradict the conclusion from
#     above?
# 
#  4) In the above case, what is the mean, and what is the width? Does the Poisson
#     relation between mean and width hold?
# 
#---------------------------------------------------------------------------------- 
#
# Key questions:
# --------------
#   https://docs.google.com/forms/d/1yY9gfH2H5IotEXXUHXFRuMoNlMBqTOr8UZ50ijROHEI/viewform?usp=send_form
#   * Which distribution has the longest tail towards high values? And low values?
#   * Using Nexperiments=1000, Ntrials=20 and pSuccess=1/6, what is the ChiSquare
#     probability of the Binomial and Poisson based on the first seven bins?
#
#----------------------------------------------------------------------------------