%matplotlib inline
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
plt.style.use('ggplot')

"""
Explanation: Fitting the H.E.S.S. Crab spectrum with iminuit and emcee
As an example of a chi^2 fit, we use the flux points from the Crab nebula
as measured by H.E.S.S. in 2006:  http://adsabs.harvard.edu/abs/2006A%26A...457..899A
End of explanation
"""


data = pd.read_table('spectrum_crab_hess_2006.txt',
                    comment='#', sep='\s*', engine='python')
data

"""
Explanation: The data
We start by loading the flux points from a text file.
It is of course possible to load this data using just Python
or Numpy, but we'll use the pandas.read_table function
because it's very flexible, i.e. by setting a few arguments you'll
be able to load most ascii tables.
End of explanation
"""


def flux_ecpl(energy, flux1, gamma, energy_cut):
    return flux1 * energy ** (-gamma) * np.exp(-energy / energy_cut)

"""
Explanation: The model
In the paper they fit a power-law with an exponential cutoff and find the following parameters (see row "all" in table 6):
* gamma = 2.39 +- 0.03
* energy_cut = 14.3 +- 2.1 TeV
* flux1 = (3.76 +- 0.07) x 1e-11 cm^-2 s^-1 TeV^-1
The flux1 is the differential flux at 1 TeV.
Let's code up that model ...
TODO: extend this tutorial to also consider a power-law model and compare the two models via chi2 / ndf.
End of explanation
"""


energy = np.logspace(-0.5, 1.6, 100)
flux = flux_ecpl(energy, flux1=3.76e-11, gamma=2.39, energy_cut=14.3)
plt.plot(energy, flux)
plt.errorbar(x=data['energy'],
             y=data['flux'],
             yerr=data['flux_err'],
             fmt='.'
            )
plt.loglog();

"""
Explanation: Plot data and model
Let's plot the data and model and compare to Figure 18b
from the paper ...
End of explanation
"""


def chi2(flux1, gamma, energy_cut):
    energy = data['energy']
    flux_model = flux_ecpl(energy, flux1, gamma, energy_cut)
    chi = (data['flux'] - flux_model) / data['flux_err']
    return np.sum(chi ** 2)

# TODO: visualise the likelihood as a 1D profile or
# 2D contour to check that the implementation is OK
# before fitting. E.g. reproduce Fig 19 from the paper?
# Maybe talk about how chi2 differences relate to
# confidence levels here?

"""
Explanation: The likelihood
In this case we'll use a chi^2 likelihood function to
fit the model to the data.
Note that the likelihood function combines the data and model, and just depends on the model parameters that
shall be estimated (whereas the model function flux_ecpl has an extra parameter energy).
Also note that we're accessing data and model flux_ecpl from the global scope instead of passing them in explicitly as parameters. Modeling and fitting frameworks like e.g. Sherpa have more elaborate ways to combine data and models and likelihood functions, but for simple, small code examples like we do here, using the global scope to tie things together works just fine.
End of explanation
"""


from iminuit import Minuit

# TODO

"""
Explanation: ML fit with Minuit
Let's use Minuit to do a maximum likelihood (ML) analysis.
Note that this is not what they did in the paper (TODO: check), so it's not surprising if best-fit results
are a little different.
End of explanation
"""


import emcee

# TODO

"""
Explanation: Analysis with emcee

Should we only do Bayesian analysis? (posterior sampling)
Or should we start with a ML analysis (likelihood sampling and compare with iminuit)
End of explanation
"""


