# imports
from scipy import optimize, math
import numpy
import ROOT
from array import array


ROOT.gStyle.SetOptStat(1100)
ROOT.gStyle.SetOptTitle(0)

###############
# DEFINITIONS #
###############

# Poisson prob
def prob_poisson(n,mu):
    if mu<=0 or n<0: 
        return 0
    else:
        p = 1.0
        for i in range(0,len(n)):
            nn = int(n[i])
            p = p * (mu**nn)*math.exp(-mu)/math.factorial(nn) 
        return p

# Poisson LH
def lh_poisson(mu, *args):
    n = args[0]
    lh = prob_poisson(n,mu)
    if lh>0:
        return -2.0*math.log(lh)
    else:
        return 1e99
    

##############
# PARAMETERS #
##############

mu0=3.5
ntrials=1000



################
# MAIN PROGRAM #
################


# derived constants
xmin=max(0,int(mu0-5.*math.sqrt(mu0)))
xmax=int(mu0+5.*math.sqrt(mu0))
nbins=xmax-xmin

# lists to be used to plot the graphs
ntoys = []
dntoys = []
crbounds = []

# list of methods
means = {}
variances = {}
methods = ["moments","MLE"]

for meth in methods:
    means[meth] = 0
    variances[meth] = 0

rnd = ROOT.TRandom3()

nt = 1
ntoys.append(nt)
dntoys.append(0)
crbounds.append(1./(nt*(1./mu0)))

# trials
estimates = {}
for meth in methods:
    estimates[meth] = []
    
for i in range(0,ntrials):
    # toy generation
    toys = []
    for it in range(0,nt):
        toy = rnd.Poisson(mu0)
        toys.append(toy)

    for meth in methods:
        result=0
        if meth=="MLE":    
            # maximum likelihood estimate
            result = optimize.fmin(lh_poisson,mu0,args=(toys,),disp=False)
        elif meth=="moments":
            result = numpy.mean(toys)
            
        estimates[meth].append(result)


# check the mean and variance of the estimator for the different methods            
for meth in methods:
    means[meth] = numpy.mean(estimates[meth])
    variances[meth] = numpy.var(estimates[meth])
    print ("mean for method %s: %f" % (meth , means[meth]))
    print ("variance for method %s: %f" % (meth, variances[meth]))


# draw the distribition of estimator values
c = ROOT.TCanvas("c0","estimators")
c.Divide(len(methods))
i=0
h = []
for meth in methods:
    i=i+1
    c.cd(i)
    h.append(ROOT.TH1F("h%d" % i, "",nbins,xmin-0.5,xmax-0.5))
    for j in range(0,len(estimates[meth])):
        h[-1].Fill(estimates[meth][j])
    h[-1].Draw()
c.Draw()


raw_input("Press Enter to continue ...")