#!/usr/bin/env python
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.
"""
Script to do final null hypothesis test that the events in the 20 MeV - 10 GeV
range are consistent with atmospheric neutrino events. To run it just run:

    $ ./chi2 [list of data fit results] --mc [list of atmospheric MC files] --muon-mc [list of muon MC files] --steps [steps]

After running you will get a plot showing the best fit results of the MC and
data along with p-values for each of the possible particle combinations (single
electron, single muon, double electron, etc.).
"""
from __future__ import print_function, division
import numpy as np
from scipy.stats import iqr, poisson
from matplotlib.lines import Line2D
from scipy.stats import iqr, norm, beta, percentileofscore
from scipy.special import spence
from itertools import izip_longest
from sddm.stats import *
from sddm.dc import estimate_errors, EPSILON, truncnorm_scaled
import emcee
from sddm import printoptions
from sddm.utils import fast_cdf, correct_energy_bias

# Likelihood Fit Parameters
# 0 - Atmospheric Neutrino Flux Scale
# 1 - Electron energy bias
# 2 - Electron energy resolution
# 3 - Muon energy bias
# 4 - Muon energy resolution
# 5 - External Muon scale

FIT_PARS = [
    'Atmospheric Neutrino Flux Scale',
    'Electron energy bias',
    'Electron energy resolution',
    'Muon energy bias',
    'Muon energy resolution',
    'External Muon scale']

# Uncertainty on the energy scale
#
# - the muon energy scale and resolution terms come directly from measurements
#   on stopping muons, so those are known well.
# - for electrons, we only have Michel electrons at the low end of our energy
#   range, and therefore we don't really have any good way of constraining the
#   energy scale or resolution. However, if we assume that the ~7% energy bias
#   in the muons is from the single PE distribution (it seems likely to me that
#   that is a major part of the bias), then the energy scale should be roughly
#   the same. Since the Michel electron distributions are consistent, we leave
#   the mean value at 0, but to be conservative, we set the error to 10%.
# - The energy resolution for muons was pretty much spot on, and so we expect
#   the same from electrons. In addition, the Michel spectrum is consistent so
#   at that energy level we don't see anything which leads us to expect a major
#   difference. To be conservative, and because I don't think it really affects
#   the analysis at all, I'll leave the uncertainty here at 10% anyways.
# - For the electron + muon energy resolution, I don't have any real good way
#   to estimate this, so we are conservative and set the error to 10%.
PRIORS = [
    1.0,   # Atmospheric Neutrino Scale
    0.0,   # Electron energy scale
    0.0,   # Electron energy resolution
    0.066, # Muon energy scale
    0.0,   # Muon energy resolution
    0.0,   # Muon scale
]

PRIOR_UNCERTAINTIES = [
    0.2  , # Atmospheric Neutrino Scale
    0.1,   # Electron energy scale
    0.1,   # Electron energy resolution
    0.011, # Muon energy scale
    0.014, # Muon energy resolution
    10.0,  # Muon scale
]

# Lower bounds for the fit parameters
PRIORS_LOW = [
    EPSILON,
    -10,
    EPSILON,
    -10,
    EPSILON,
    EPSILON
]

# Upper bounds for the fit parameters
PRIORS_HIGH = [
    10,
    10,
    10,
    10,
    10,
    10
]

particle_id = {20: 'e', 22: r'\mu'}

def plot_hist2(hists, bins, color=None):
    for id in (20,22,2020,2022,2222):
        if id == 20:
            plt.subplot(2,3,1)
        elif id == 22:
            plt.subplot(2,3,2)
        elif id == 2020:
            plt.subplot(2,3,4)
        elif id == 2022:
            plt.subplot(2,3,5)
        elif id == 2222:
            plt.subplot(2,3,6)

        bincenters = (bins[1:] + bins[:-1])/2
        plt.hist(bincenters, bins=bins, histtype='step', weights=hists[id],color=color)
        plt.gca().set_xscale("log")
        plt.xlabel("Energy (MeV)")
        plt.title('$' + ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + '$')

    if len(hists):
        plt.tight_layout()

def get_mc_hists(data,x,bins,scale=1.0,reweight=False):
    """
    Returns the expected Monte Carlo histograms for the atmospheric neutrino
    background.

    Args:
        - data: pandas dataframe of the Monte Carlo events
        - x: fit parameters
        - bins: histogram bins
        - scale: multiply histograms by an overall scale factor

    This function does two basic things:

        1. apply the energy bias and resolution corrections
        2. histogram the results

    Returns a dictionary mapping particle id combo -> histogram.
    """
    df_dict = {}
    for id in (20,22,2020,2022,2222):
        df_dict[id] = data[data.id == id]

    return get_mc_hists_fast(df_dict,x,bins,scale,reweight)

def get_mc_hists_fast(df_dict,x,bins,scale=1.0,reweight=False):
    """
    Same as get_mc_hists() but the first argument is a dictionary mapping
    particle id -> dataframe. This is much faster than selecting the events
    from the dataframe every time.
    """
    mc_hists = {}

    for id in (20,22,2020,2022,2222):
        df = df_dict[id]

        if id == 20:
            ke = df.energy1.values*(1+x[1])
            resolution = df.energy1.values*max(EPSILON,x[2])
        elif id == 2020:
            ke = df.energy1.values*(1+x[1]) + df.energy2.values*(1+x[1])
            resolution = np.sqrt((df.energy1.values*max(EPSILON,x[2]))**2 + (df.energy2.values*max(EPSILON,x[2]))**2)
        elif id == 22:
            ke = df.energy1.values*(1+x[3])
            resolution = df.energy1.values*max(EPSILON,x[4])
        elif id == 2222:
            ke = df.energy1.values*(1+x[3]) + df.energy2.values*(1+x[3])
            resolution = np.sqrt((df.energy1.values*max(EPSILON,x[4]))**2 + (df.energy2.values*max(EPSILON,x[4]))**2)
        elif id == 2022:
            ke = df.energy1.values*(1+x[1]) + df.energy2.values*(1+x[3])
            resolution = np.sqrt((df.energy1.values*max(EPSILON,x[2]))**2 + (df.energy2.values*max(EPSILON,x[4]))**2)

        if reweight:
            cdf = fast_cdf(bins[:,np.newaxis],ke,resolution)*df.weight.values
        else:
            cdf = fast_cdf(bins[:,np.newaxis],ke,resolution)

        mc_hists[id] = np.sum(cdf[1:] - cdf[:-1],axis=-1)
        mc_hists[id] *= scale
    return mc_hists

def get_data_hists(data,bins,scale=1.0):
    """
    Returns the data histogrammed into `bins`.
    """
    data_hists = {}
    for id in (20,22,2020,2022,2222):
        data_hists[id] = np.histogram(data[data.id == id].ke.values,bins=bins)[0]*scale
    return data_hists

def make_nll(data, muons, mc, atmo_scale_factor, muon_scale_factor, bins, print_nll=False):
    df_dict = {}
    for id in (20,22,2020,2022,2222):
        df_dict[id] = mc[mc.id == id]

    df_dict_muon = {}
    for id in (20,22,2020,2022,2222):
        df_dict_muon[id] = muons[muons.id == id]

    data_hists = get_data_hists(data,bins)

    def nll(x, grad=None):
        if any(x[i] < 0 for i in (0,2,4,5)):
            return np.inf

        # Get the Monte Carlo histograms. We need to do this within the
        # likelihood function since we apply the energy resolution parameters
        # to the Monte Carlo.
        mc_hists = get_mc_hists_fast(df_dict,x,bins,scale=1/atmo_scale_factor)
        muon_hists = get_mc_hists_fast(df_dict_muon,x,bins,scale=1/muon_scale_factor)

        # Calculate the negative log of the likelihood of observing the data
        # given the fit parameters

        nll = 0
        for id in data_hists:
            oi = data_hists[id]
            ei = mc_hists[id]*x[0] + muon_hists[id]*x[5] + EPSILON
            N = ei.sum()
            nll -= -N - np.sum(gammaln(oi+1)) + np.sum(oi*np.log(ei))

        # Add the priors
        nll -= norm.logpdf(x,PRIORS,PRIOR_UNCERTAINTIES).sum()

        if print_nll:
            # Print the result
            print("nll = %.2f" % nll)

        return nll
    return nll

def get_mc_hists_posterior(data_mc,muon_hists,data_hists,atmo_scale_factor,muon_scale_factor,x,bins):
    """
    Returns the posterior on the Monte Carlo histograms.

    Basically this function just histograms the Monte Carlo data. However,
    there is one extra thing it does. In general when doing a fit, the Monte
    Carlo histograms have some uncertainty since you can never simulate an
    infinite number of statistics. I don't think I've ever really seen anyone
    properly treat this. Since the uncertainty on the central value in each bin
    is just given by the Dirichlet distribution, we treat the problem of
    finding the best value of the posterior as a problem in which you're prior
    is equal to the expected number of events from the Monte Carlo, and then
    you actually see the data.  Since the likelihood on the true mean in each
    bin is a multinomial, the posterior is also a dirichlet where the alpha
    parameters are given by a sum of the prior and observed counts.

    All that is a long way of saying we calculate the posterior as the sum of
    the Monte Carlo events (unscaled) and the observed events. In the limit of
    infinite statistics, this is just equal to the Monte Carlo predicted
    histogram, but deals with the fact that we don't have infinite statistics,
    and so a single outlier event isn't necessarily a problem with the model.

    Returns a dictionary mapping particle id combo -> histogram.
    """
    mc_hists = get_mc_hists(data_mc,x,bins,reweight=True)
    for id in (20,22,2020,2022,2222):
        mc_hists[id] = get_mc_hist_posterior(mc_hists[id],data_hists[id],norm=x[0]/atmo_scale_factor)
        # FIXME: does the orering of when we add the muons matter here?
        mc_hists[id] += muon_hists[id]*x[5]/muon_scale_factor
    return mc_hists

def get_multinomial_prob(data, data_muon, data_mc, weights, atmo_scale_factor, muon_scale_factor, id, x_samples, bins, percentile=50.0, size=10000):
    """
    Returns the p-value that the histogram of the data is drawn from the MC
    histogram.

    The p-value is calculated by first sampling the posterior of the fit
    parameters `size` times. For each iteration we calculate a p-value. We then
    return the `percentile` percentile of all the p-values. This approach is
    similar to both the supremum and posterior predictive methods of
    calculating a p-value. For more information on different methods of
    calculating p-values see https://cds.cern.ch/record/1099967/files/p23.pdf.

    Arguments:

        data: 1D array of KE values
        data_mc: 1D array of MC KE values
        x_samples: MCMC samples of the floated parameters in the fit
        bins: bins used to bin the mc histogram
        size: number of values to compute
    """
    df_dict_muon = {}
    for _id in (20,22,2020,2022,2222):
        df_dict_muon[_id] = data_muon[data_muon.id == _id]

    data_hists = get_data_hists(data,bins)

    # Get the total number of "universes" simulated in the GENIE reweight tool
    if len(data_mc):
        nuniverses = weights['universe'].max()+1
    else:
        nuniverses = 0

    ps = []
    for i in range(size):
        x = x_samples[np.random.randint(x_samples.shape[0])]

        muon_hists = get_mc_hists_fast(df_dict_muon,x,bins)

        if nuniverses > 0:
            universe = np.random.randint(nuniverses)
            data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
            data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
        else:
            data_mc_with_weights = data_mc.copy()
            data_mc['weight'] = 1.0
        mc = get_mc_hists_posterior(data_mc_with_weights,muon_hists,data_hists,atmo_scale_factor,muon_scale_factor,x,bins)[id]
        N = mc.sum()
        # Fix a bug in scipy(). See https://github.com/scipy/scipy/issues/8235 (I think).
        mc = mc + 1e-10
        p = mc/mc.sum()
        chi2_data = nllr(data_hists[id],mc)
        # To draw the multinomial samples we first draw the expected number of
        # events from a Poisson distribution and then loop over the counts and
        # unique values. The reason we do this is that you can't call
        # multinomial.rvs with a multidimensional `n` array, and looping over every
        # single entry takes forever
        ns = np.random.poisson(N,size=1000)
        samples = []
        for n, count in zip(*np.unique(ns,return_counts=True)):
            samples.append(multinomial.rvs(n,p,size=count))
        samples = np.concatenate(samples)
        # Calculate the negative log likelihood ratio for the data simulated under
        # the null hypothesis
        chi2_samples = nllr(samples,mc)
        ps.append(np.count_nonzero(chi2_samples >= chi2_data)/len(chi2_samples))
    return np.percentile(ps,percentile)

def get_prob(data,muon,mc,weights,atmo_scale_factor,muon_scale_factor,samples,bins,size):
    prob = {}
    for id in (20,22,2020,2022,2222):
        prob[id] = get_multinomial_prob(data,muon,mc,weights,atmo_scale_factor,muon_scale_factor,id,samples,bins,size=size)
        print(id, prob[id])
    return prob

def do_fit(data,muon,data_mc,atmo_scale_factor,muon_scale_factor,bins,steps,print_nll=False,walkers=100,thin=10):
    """
    Run the fit and return the minimum along with samples from running an MCMC
    starting near the minimum.

    Args:
        - data: pandas dataframe representing the data to fit
        - muon: pandas dataframe representing the expected background from
                external muons
        - data_mc: pandas dataframe representing the expected background from
                   atmospheric neutrino events
        - bins: an array of bins to use for the fit
        - steps: the number of MCMC steps to run

    Returns a tuple (xopt, samples) where samples is an array of shape (steps,
    number of parameters).
    """
    nll = make_nll(data,muon,data_mc,atmo_scale_factor,muon_scale_factor,bins,print_nll)

    pos = np.empty((walkers, len(PRIORS)),dtype=np.double)
    for i in range(pos.shape[0]):
        pos[i] = truncnorm_scaled(PRIORS_LOW,PRIORS_HIGH,PRIORS,PRIOR_UNCERTAINTIES)

    nwalkers, ndim = pos.shape

    # We use the KDEMove here because I think it should sample the likelihood
    # better. Because we have energy scale parameters and we are doing a binned
    # likelihood, the likelihood is discontinuous. There can also be several
    # local minima. The author of emcee recommends using the KDEMove with a lot
    # of workers to try and properly sample a multimodal distribution. In
    # addition, I've found that the autocorrelation time for the KDEMove is
    # much better than the other moves.
    sampler = emcee.EnsembleSampler(nwalkers, ndim, lambda x: -nll(x), moves=emcee.moves.KDEMove())
    with np.errstate(invalid='ignore'):
        sampler.run_mcmc(pos, steps)

    print("Mean acceptance fraction: {0:.3f}".format(np.mean(sampler.acceptance_fraction)))

    try:
        print("autocorrelation time: ", sampler.get_autocorr_time(quiet=True))
    except Exception as e:
        print(e)

    samples = sampler.get_chain(flat=True,thin=thin)

    return sampler.get_chain(flat=True)[sampler.get_log_prob(flat=True).argmax()], samples

if __name__ == '__main__':
    import argparse
    import numpy as np
    import pandas as pd
    import sys
    import h5py
    from sddm.plot_energy import *
    from sddm.plot import *
    from sddm import setup_matplotlib
    import nlopt

    parser = argparse.ArgumentParser("plot fit results")
    parser.add_argument("filenames", nargs='+', help="input files")
    parser.add_argument("--save", action='store_true', default=False, help="save corner plots for backgrounds")
    parser.add_argument("--mc", nargs='+', required=True, help="atmospheric MC files")
    parser.add_argument("--muon-mc", nargs='+', required=True, help="muon MC files")
    parser.add_argument("--nhit-thresh", type=int, default=None, help="nhit threshold to apply to events before processing (should only be used for testing to speed things up)")
    parser.add_argument("--steps", type=int, default=1000, help="number of steps in the MCMC chain")
    parser.add_argument("--multinomial-prob-size", type=int, default=10000, help="number of p values to compute")
    parser.add_argument("--coverage", type=int, default=0, help="plot p value coverage")
    parser.add_argument("--pull", type=int, default=0, help="plot pull plots")
    parser.add_argument("--weights", nargs='+', required=True, help="GENIE reweight HDF5 files")
    parser.add_argument("--print-nll", action='store_true', default=False, help="print nll values")
    parser.add_argument("--walkers", type=int, default=100, help="number of walkers")
    parser.add_argument("--thin", type=int, default=10, help="number of steps to thin")
    args = parser.parse_args()

    setup_matplotlib(args.save)

    import matplotlib.pyplot as plt

    # Loop over runs to prevent using too much memory
    evs = []
    rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in args.filenames],ignore_index=True)
    for run, df in rhdr.groupby('run'):
        evs.append(get_events(df.filename.values, merge_fits=True, nhit_thresh=args.nhit_thresh))
    ev = pd.concat(evs)

    ev = correct_energy_bias(ev)

    ev_mc = get_events(args.mc, merge_fits=True, nhit_thresh=args.nhit_thresh)
    muon_mc = get_events(args.muon_mc, merge_fits=True, nhit_thresh=args.nhit_thresh)
    weights = pd.concat([read_hdf(filename, "weights") for filename in args.weights],ignore_index=True)

    ev_mc = correct_energy_bias(ev_mc)
    muon_mc = correct_energy_bias(muon_mc)

    # Set all prompt events in the MC to be muons
    muon_mc.loc[muon_mc.prompt & muon_mc.filename.str.contains("cosmic"),'muon'] = True

    ev = ev.reset_index()
    ev_mc = ev_mc.reset_index()
    muon_mc = muon_mc.reset_index()

    # 00-orphan cut
    ev = ev[(ev.gtid & 0xff) != 0]
    ev_mc = ev_mc[(ev_mc.gtid & 0xff) != 0]
    muon_mc = muon_mc[(muon_mc.gtid & 0xff) != 0]

    # remove events 200 microseconds after a muon
    ev = ev.groupby('run',group_keys=False).apply(muon_follower_cut)

    # Get rid of events which don't have a successful fit
    ev = ev[~np.isnan(ev.fmin)]
    ev_mc = ev_mc[~np.isnan(ev_mc.fmin)]
    muon_mc = muon_mc[~np.isnan(muon_mc.fmin)]

    # require (r/r_psup)^3 < 0.9
    ev = ev[ev.r_psup < 0.9]
    ev_mc = ev_mc[ev_mc.r_psup < 0.9]
    muon_mc = muon_mc[muon_mc.r_psup < 0.9]

    # require psi < 6
    ev = ev[ev.psi < 6]
    ev_mc = ev_mc[ev_mc.psi < 6]
    muon_mc = muon_mc[muon_mc.psi < 6]

    data = ev[ev.signal & ev.prompt & ~ev.atm]
    data_atm = ev[ev.signal & ev.prompt & ev.atm]

    # Right now we use the muon Monte Carlo in the fit. If you want to use the
    # actual data, you can comment the next two lines and then uncomment the
    # two after that.
    muon = muon_mc[muon_mc.muon & muon_mc.prompt & ~muon_mc.atm]
    muon_atm = muon_mc[muon_mc.muon & muon_mc.prompt & muon_mc.atm]
    #muon = ev[ev.muon & ev.prompt & ~ev.atm]
    #muon_atm = ev[ev.muon & ev.prompt & ev.atm]

    data_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ~ev_mc.atm]
    data_atm_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ev_mc.atm]

    bins = np.logspace(np.log10(20),np.log10(10e3),21)

    atmo_scale_factor = 100.0
    muon_scale_factor = len(muon) + len(muon_atm)

    if args.coverage:
        p_values = {id: [] for id in (20,22,2020,2022,2222)}
        p_values_atm = {id: [] for id in (20,22,2020,2022,2222)}

        # Set the random seed so we get reproducible results here
        np.random.seed(0)

        xtrue = truncnorm_scaled(PRIORS_LOW,PRIORS_HIGH,PRIORS,PRIOR_UNCERTAINTIES)

        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == 0],how='left',on=['run','evn'])
        data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == 0],how='left',on=['run','evn'])

        for i in range(args.coverage):
            # Calculate expected number of events
            N = len(data_mc)*xtrue[0]/atmo_scale_factor
            N_atm = len(data_atm_mc)*xtrue[0]/atmo_scale_factor
            N_muon = len(muon)*xtrue[5]/muon_scale_factor
            N_muon_atm = len(muon_atm)*xtrue[5]/muon_scale_factor

            # Calculate observed number of events
            n = np.random.poisson(N)
            n_atm = np.random.poisson(N_atm)
            n_muon = np.random.poisson(N_muon)
            n_muon_atm = np.random.poisson(N_muon_atm)

            # Sample data from Monte Carlo
            data = pd.concat((data_mc_with_weights.sample(n=n,replace=True,weights='weight'), muon.sample(n=n_muon,replace=True)))
            data_atm = pd.concat((data_atm_mc_with_weights.sample(n=n_atm,replace=True,weights='weight'), muon_atm.sample(n=n_muon_atm,replace=True)))

            # Smear the energies by the additional energy resolution
            data.loc[data.id1 == 20,'energy1'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data.id1 == 20))*xtrue[2])
            data.loc[data.id1 == 22,'energy1'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data.id1 == 22))*xtrue[4])
            data.loc[data.id2 == 20,'energy2'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data.id2 == 20))*xtrue[2])
            data.loc[data.id2 == 22,'energy2'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data.id2 == 22))*xtrue[4])
            data['ke'] = data['energy1'].fillna(0) + data['energy2'].fillna(0) + data['energy3'].fillna(0)

            data_atm.loc[data_atm.id1 == 20,'energy1'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data_atm.id1 == 20))*xtrue[2])
            data_atm.loc[data_atm.id1 == 22,'energy1'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data_atm.id1 == 22))*xtrue[4])
            data_atm.loc[data_atm.id2 == 20,'energy2'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data_atm.id2 == 20))*xtrue[2])
            data_atm.loc[data_atm.id2 == 22,'energy2'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data_atm.id2 == 22))*xtrue[4])
            data_atm['ke'] = data_atm['energy1'].fillna(0) + data_atm['energy2'].fillna(0) + data_atm['energy3'].fillna(0)

            xopt, samples = do_fit(data,muon,data_mc,atmo_scale_factor,muon_scale_factor,bins,args.steps,args.print_nll,args.walkers,args.thin)

            prob = get_prob(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)
            prob_atm = get_prob(data_atm,muon_atm,data_atm_mc,weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)

            for id in (20,22,2020,2022,2222):
                p_values[id].append(prob[id])
                p_values_atm[id].append(prob_atm[id])

        fig = plt.figure()
        for id in (20,22,2020,2022,2222):
            if id == 20:
                plt.subplot(2,3,1)
            elif id == 22:
                plt.subplot(2,3,2)
            elif id == 2020:
                plt.subplot(2,3,4)
            elif id == 2022:
                plt.subplot(2,3,5)
            elif id == 2222:
                plt.subplot(2,3,6)
            plt.hist(p_values[id],bins=np.linspace(0,1,101),histtype='step')
            plt.title('$' + ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + '$')
        despine(fig,trim=True)
        plt.tight_layout()

        if args.save:
            fig.savefig("chi2_p_value_coverage_plot.pdf")
            fig.savefig("chi2_p_value_coverage_plot.eps")
        else:
            plt.suptitle("P-value Coverage without Neutron Follower")

        fig = plt.figure()
        for id in (20,22,2020,2022,2222):
            if id == 20:
                plt.subplot(2,3,1)
            elif id == 22:
                plt.subplot(2,3,2)
            elif id == 2020:
                plt.subplot(2,3,4)
            elif id == 2022:
                plt.subplot(2,3,5)
            elif id == 2222:
                plt.subplot(2,3,6)
            plt.hist(p_values_atm[id],bins=np.linspace(0,1,101),histtype='step')
            plt.title('$' + ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + '$')
        despine(fig,trim=True)
        plt.tight_layout()

        if args.save:
            fig.savefig("chi2_p_value_coverage_plot_atm.pdf")
            fig.savefig("chi2_p_value_coverage_plot_atm.eps")
        else:
            plt.suptitle("P-value Coverage with Neutron Follower")

        sys.exit(0)

    if args.pull:
        pull = [[] for i in range(len(FIT_PARS))]

        # Set the random seed so we get reproducible results here
        np.random.seed(0)

        for i in range(args.pull):
            xtrue = truncnorm_scaled(PRIORS_LOW,PRIORS_HIGH,PRIORS,PRIOR_UNCERTAINTIES)

            # Calculate expected number of events
            N = len(data_mc)*xtrue[0]/atmo_scale_factor
            N_atm = len(data_atm_mc)*xtrue[0]/atmo_scale_factor
            N_muon = len(muon)*xtrue[5]/muon_scale_factor
            N_muon_atm = len(muon_atm)*xtrue[5]/muon_scale_factor

            # Calculate observed number of events
            n = np.random.poisson(N)
            n_atm = np.random.poisson(N_atm)
            n_muon = np.random.poisson(N_muon)
            n_muon_atm = np.random.poisson(N_muon_atm)

            # Sample data from Monte Carlo
            data = pd.concat((data_mc.sample(n=n,replace=True), muon.sample(n=n_muon,replace=True)))
            data_atm = pd.concat((data_atm_mc.sample(n=n_atm,replace=True), muon_atm.sample(n=n_muon_atm,replace=True)))

            # Smear the energies by the additional energy resolution
            data.loc[data.id1 == 20,'energy1'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data.id1 == 20))*xtrue[2])
            data.loc[data.id1 == 22,'energy1'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data.id1 == 22))*xtrue[4])
            data.loc[data.id2 == 20,'energy2'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data.id2 == 20))*xtrue[2])
            data.loc[data.id2 == 22,'energy2'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data.id2 == 22))*xtrue[4])
            data['ke'] = data['energy1'].fillna(0) + data['energy2'].fillna(0) + data['energy3'].fillna(0)

            data_atm.loc[data_atm.id1 == 20,'energy1'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data_atm.id1 == 20))*xtrue[2])
            data_atm.loc[data_atm.id1 == 22,'energy1'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data_atm.id1 == 22))*xtrue[4])
            data_atm.loc[data_atm.id2 == 20,'energy2'] *= (1+xtrue[1]+np.random.randn(np.count_nonzero(data_atm.id2 == 20))*xtrue[2])
            data_atm.loc[data_atm.id2 == 22,'energy2'] *= (1+xtrue[3]+np.random.randn(np.count_nonzero(data_atm.id2 == 22))*xtrue[4])
            data_atm['ke'] = data_atm['energy1'].fillna(0) + data_atm['energy2'].fillna(0) + data_atm['energy3'].fillna(0)

            xopt, samples = do_fit(data,muon,data_mc,atmo_scale_factor,muon_scale_factor,bins,args.steps,args.print_nll,args.walkers,args.thin)

            for i in range(len(FIT_PARS)):
                # The "pull plots" we make here are actually produced via a
                # procedure called "Simulation Based Calibration".
                #
                # See https://arxiv.org/abs/1804.06788.
                pull[i].append(percentileofscore(samples[:,i],xtrue[i]))

        fig = plt.figure()
        axes = []
        for i, name in enumerate(FIT_PARS):
            axes.append(plt.subplot(3,3,i+1))
            plt.hist(pull[i],bins=np.linspace(0,100,101),histtype='step',normed=True)
            plt.title(name)
        for ax in axes:
            despine(ax=ax,left=True,trim=True)
            ax.get_yaxis().set_visible(False)
        plt.tight_layout()

        if args.save:
            fig.savefig("chi2_pull_plot.pdf")
            fig.savefig("chi2_pull_plot.eps")
        else:
            plt.show()

        sys.exit(0)

    xopt, samples = do_fit(data,muon,data_mc,atmo_scale_factor,muon_scale_factor,bins,args.steps,args.print_nll,args.walkers,args.thin)

    prob = get_prob(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)
    prob_atm = get_prob(data_atm,muon_atm,data_atm_mc,weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)

    plt.figure()
    for i in range(len(FIT_PARS)):
        plt.subplot(3,2,i+1)
        plt.hist(samples[:,i],bins=100,histtype='step')
        plt.xlabel(FIT_PARS[i].title())
        despine(ax=plt.gca(),left=True,trim=True)
        plt.gca().get_yaxis().set_visible(False)
    plt.tight_layout()

    if args.save:
        plt.savefig("chi2_fit_posterior.pdf")
        plt.savefig("chi2_fit_posterior.eps")
    else:
        plt.suptitle("Fit Posteriors")

    handles = [Line2D([0], [0], color='C0'),
               Line2D([0], [0], color='C1'),
               Line2D([0], [0], color='C2')]
    labels = ('Data','Monte Carlo','External Muons')

    fig = plt.figure()
    hists = get_data_hists(data,bins)
    hists_muon = get_mc_hists(muon,xopt,bins,scale=xopt[5]/muon_scale_factor)
    hists_mc = get_mc_hists(data_mc,xopt,bins,scale=xopt[0]/atmo_scale_factor)
    plot_hist2(hists,bins=bins,color='C0')
    plot_hist2(hists_mc,bins=bins,color='C1')
    plot_hist2(hists_muon,bins=bins,color='C2')
    for id in (20,22,2020,2022,2222):
        if id == 20:
            plt.subplot(2,3,1)
        elif id == 22:
            plt.subplot(2,3,2)
        elif id == 2020:
            plt.subplot(2,3,4)
        elif id == 2022:
            plt.subplot(2,3,5)
        elif id == 2222:
            plt.subplot(2,3,6)
        plt.text(0.95,0.95,"p = %.2f" % prob[id],horizontalalignment='right',verticalalignment='top',transform=plt.gca().transAxes)
    fig.legend(handles,labels,loc='upper right')

    despine(fig,trim=True)
    if args.save:
        plt.savefig("chi2_prompt.pdf")
        plt.savefig("chi2_prompt.eps")
    else:
        plt.suptitle("Without Neutron Follower")
    fig = plt.figure()
    hists = get_data_hists(data_atm,bins)
    hists_muon = get_mc_hists(muon_atm,xopt,bins,scale=xopt[5]/muon_scale_factor)
    hists_mc = get_mc_hists(data_atm_mc,xopt,bins,scale=xopt[0]/atmo_scale_factor)
    plot_hist2(hists,bins=bins,color='C0')
    plot_hist2(hists_mc,bins=bins,color='C1')
    plot_hist2(hists_muon,bins=bins,color='C2')
    for id in (20,22,2020,2022,2222):
        if id == 20:
            plt.subplot(2,3,1)
        elif id == 22:
            plt.subplot(2,3,2)
        elif id == 2020:
            plt.subplot(2,3,4)
        elif id == 2022:
            plt.subplot(2,3,5)
        elif id == 2222:
            plt.subplot(2,3,6)
        plt.text(0.95,0.95,"p = %.2f" % prob_atm[id],horizontalalignment='right',verticalalignment='top',transform=plt.gca().transAxes)
    fig.legend(handles,labels,loc='upper right')

    despine(fig,trim=True)
    if args.save:
        plt.savefig("chi2_atm.pdf")
        plt.savefig("chi2_atm.eps")
    else:
        plt.suptitle("With Neutron Follower")
        plt.show()