import myFitter as mf
import matplotlib.pyplot as plt
import scipy.stats as stats
import sys
import numpy as np
import random

# Import our example model.
from mymodel import MyModel

# Set random seed to get reproducable results
random.seed(1234)
np.random.seed(1234)

# `model0` represents the null hypothesis ``p1=1.25``. It is a MyModel object
# with the parameter `p1` fixed to 1.25.
model0 = MyModel()
model0.fix({'p1' : 1.25})
# Create a Fitter object `fitter0` which fits `model0`.
fitter0 = mf.Fitter(model0, tries=1)
fitter0.sample(2000)
# Fit `model0` to the observed data and store the result in `fit0`.
fit0 = fitter0.fit()

# `model1` represents the full model. It is a MyModel object where no parameters
# are fixed.
model1 = MyModel()
# Create a Fitter object `fitter1` which fits `model1`.
fitter1 = mf.Fitter(model1, tries=1)
fitter1.sample(2000)
# Fit `model1` to the observed data and store the result in `fit1`.
fit1 = fitter1.fit()

# The observed Delta chi-square `dchisq` is the difference of the
# minimum chi-square values found in the two fits.
dchisq = fit0.minChiSquare - fit1.minChiSquare

# Create a Histogram1D object which will hold the distribution of the
# Delta chi-square.
dchisqhist = mf.Histogram1D(xbins=np.linspace(0.0, 3.0, 21))

# Create a ToySimulator object which generates observables distributed according
# to the PDF of `model0` with the parameters set to their maximum likelihood
# estimates found in the previous fit (plug-in method).
sim = mf.ToySimulator(model0, fit0.bestFitParameters)

# Initialise all data values of the histogram to zero.
dchisqhist.set(0.0)
# Initialise `pvalue` to zero. This variable will hold the p-value.
pvalue = 0.0

# Loop over (at most) 2000 toy observables. `obs` is a dict describing the point
# in observable space and `weight` is the associated weight. The sum of the
# weights of all generated points is guaranteed to be 1.
for obs, weight in sim.random(neval=2000):
    # Fit both models to the generated toy observables.
    toy_fit0 = fitter0.fit(obs=obs)
    toy_fit1 = fitter1.fit(obs=obs)
    # Compute the Delta chi-square for the toy observables.
    dchisq_toy = toy_fit0.minChiSquare - toy_fit1.minChiSquare
    # Add `weight` to the correct bin of the histogram. KeyError is raised if
    # `dchisq_toy` is outside the range of the histogram `dchisqhist`.
    try:
        dchisqhist[dchisq_toy] += weight
    except KeyError:
        pass
    # Add `weight` to `pvalue` if `dchisq_toy` is larger than the observed
    # Delta chi-square.
    if dchisq_toy > dchisq:
        pvalue += weight

# Print the p-value and compare with the asymptotic approximation.
asymptotic_pvalue = mf.pvalueFromChisq(dchisq, 1)
print 'p-value:            {0:f}'.format(pvalue)
print 'asymptotic p-value: {0:f}'.format(asymptotic_pvalue)

# Convert `dchisqhist` to a density by dividing all data values by the
# corresponding bin width.
dchisqhist = mf.densityHistogram(dchisqhist)
# Plot the histogram.
dchisqhist.plot(drawstyle='steps')
# Overlay a chi^2 distribution with 1 degree of freedom.
plt.plot(dchisqhist.xvals, [stats.chi2.pdf(x, 1) for x in dchisqhist.xvals])
plt.show()