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 reproduceable results
random.seed(1234)
np.random.seed(1234)


# Set up a model object representing the null hypothesis and fit it to the data.
model0 = MyModel()
model0.fix({'p1' : 1.25})
fitter0 = mf.Fitter(model0, tries=1)
fitter0.sample(2000)
fit0 = fitter0.fit()

# Set up a model object representing the full model (with p1 free) and fit it to
# the data.
model1 = MyModel()
fitter1 = mf.Fitter(model1, tries=1)
fitter1.sample(2000)
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 LRTIntegrand object from the two Fitter objects and the observed
# Delta chi-square. `lrt` is callable and takes a dict describing a set of toy
# observables as argument. It returns an array with two elements.  The second
# element is the Delta chi-square for the set of toy observables.  The first
# element is 0 if that value is smaller than `dchisq` and 1 otherwise.
# ``nested=True`` indicates that ``fitter0.model`` was obtained from
# ``fitter1.model`` by fixing one or more parameters. This extra information
# helps to improve the stability of the fits performed by `lrt`.
lrt = mf.LRTIntegrand(fitter0, fitter1, dchisq, nested=True)

# Create a ToySimulator object which simulates the PDF of `model0` with
# parameters set to the maximum likelihood estimates found in the fit (plug-in
# method). The `parallel` keyword means that the simulation will be parallelised
# with two jobs running on the local machine.
sim = mf.ToySimulator(model0, fit0.bestFitParameters,
                      parallel={'njobs'  : 2,
                                'runner' : mf.LocalRunner})

# Compute the p-value with a VEGAS Monte Carlo integration, using 5 iterations
# with (at most) 2000 sample points each. `result` will be an array of two
# IntegrationResult objects. IntegrationResult objects hold the results from the
# individual iterations as well as their weighted average. The first element of
# the array holds the result for the p-value. ``drop=1`` means that the first
# iteration is excluded from the weighted averages. ``log=sys.stdout`` means that
# information will be printed to ``sys.stdout`` after each iteration.
result = sim(lrt, nitn=5, neval=2000, drop=1, log=sys.stdout)

# Print the value and error of the numerically computed p-value by converting
# result[0] first to a GVar object and then to a string. Also print the
# asymptotic approximation for the p-value.
print
print 'p-value:            {0:s}'.format(str(mf.GVar(result[0])))
print 'asymptotic p-value: {0:f}'.format(mf.pvalueFromChisq(dchisq, 1))