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)

# 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 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, and that log messages about the
# running and pending jobs will be written to ``sys.stdout``.
sim = mf.ToySimulator(model0, fit0.bestFitParameters,
                      parallel={'njobs'  : 2,
                                'runner' : mf.LocalRunner,
                                'log'    : sys.stdout})

# Create a LRTIntegrand object from the two Fitter objects and the observed
# Delta chi-square. `gh` 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 always 1 since we did not specify the `dchisq` argument (see
# toy_example_2.py).  ``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`.
gh = mf.LRTIntegrand(fitter0, fitter1, nested=True)

# Generate at most 10000 sample points and fill the histogram `dchisqhist`.  For
# each toy observable `x`, the element with index 1 of the array returned by
# ``gh(x)`` will be used to select the `dchisqhist` bin.
sim.fill(gh, 10000, (1, dchisqhist))

# Convert `dchisqhist` to a density by dividing all data values by the
# corresponding bin width.
dchisqhist = mf.densityHistogram(dchisqhist)
# Plot the histogram with error bars.
dchisqhist.plot(drawstyle='steps', color='b')
dchisqhist.errorbars(color='b')
# 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()