%matplotlib inline
import os
import matplotlib.pyplot as plt
import matplotlib
import numpy as np
from fermipy.gtanalysis import GTAnalysis
from fermipy.plotting import ROIPlotter

"""
Explanation: Fitting Spatial Extension of IC443
This tutorial demonstrates how to perform a measurement of spatial extension with the extension method in the fermipy package.  This tutorial assumes that you have first gone through the PG 1553 analysis tutorial.
Get the Data and Setup the Analysis
End of explanation
"""


if os.path.isfile('../data/ic443.tar.gz'):
    !tar xzf ../data/ic443.tar.gz
else:
    !curl -OL https://raw.githubusercontent.com/fermiPy/fermipy-extras/master/data/ic443.tar.gz
    !tar xzf ic443.tar.gz

"""
Explanation: In this thread we will use a pregenerated data set which is contained in a tar archive in the data directory of the fermipy-extra repository.
End of explanation
"""


gta = GTAnalysis('ic443/config.yaml')
matplotlib.interactive(True)
gta.setup()

"""
Explanation: We first instantiate a GTAnalysis instance using the config file in the ic443 directory and the run the setup() method.  This will prepare all of the ancillary files and create the pylikelihood instance for binned analysis.  Note that in this example these files have already been generated so the routines that will normally be executed to create these files will be skipped.
End of explanation
"""


gta.print_roi()

"""
Explanation: Print the ROI model
We can print the ROI object to see a list of sources in the model along with their distance from the ROI center (offset), TS, and number of predicted counts (Npred).  Since we haven't yet fit any sources, the ts of all sources will initially be set to nan.
End of explanation
"""


gta.optimize()

gta.print_roi()

"""
Explanation: Now we will run the optimize() method.  This method refits the spectral parameters of all sources in the ROI and gives us baseline model that we can use as a starting point for fitting the spatial extension.
End of explanation
"""


resid = gta.residmap('ic443_roifit',model={'SpatialModel' : 'PointSource', 'Index' : 2.0})

o = resid
fig = plt.figure(figsize=(14,6))
ROIPlotter(o['sigma'],roi=gta.roi).plot(vmin=-5,vmax=5,levels=[-5,-3,3,5,7,9],subplot=121,cmap='RdBu_r')
plt.gca().set_title('Significance')
ROIPlotter(o['excess'],roi=gta.roi).plot(vmin=-200,vmax=200,subplot=122,cmap='RdBu_r')
plt.gca().set_title('Excess Counts')

"""
Explanation: To check the quality of the ROI model fit we can generate a residual map with the residmap method.  This will produce smoothed maps of the counts distribution and residuals (counts-model) using a given spatial kernel.  The spatial kernel can be defined with a source dictionary.  In the following example we use a PointSource with a PowerLaw index of 2.0.
End of explanation
"""


resid_noj0621 = gta.residmap('ic443_roifit_noj0621',
                model={'SpatialModel' : 'PointSource', 'Index' : 2.0},
                exclude=['3FGL J0621.0+2514'])

o = resid_noj0621
fig = plt.figure(figsize=(14,6))
ROIPlotter(o['sigma'],roi=gta.roi).plot(vmin=-5,vmax=5,levels=[-5,-3,3,5,7,9],subplot=121,cmap='RdBu_r')
plt.gca().set_title('Significance')
ROIPlotter(o['excess'],roi=gta.roi).plot(vmin=-200,vmax=200,subplot=122,cmap='RdBu_r')
plt.gca().set_title('Excess Counts')

"""
Explanation: We can see the effect of removing sources from the model by running residmap with the exclude option.  Here we generate a residual map with the source 3FGL J0621.0+2514 removed from the model.
End of explanation
"""


tsmap_noj0621 = gta.tsmap('ic443_noj0621',
                model={'SpatialModel' : 'PointSource', 'Index' : 2.0},
                exclude=['3FGL J0621.0+2514'])

o = tsmap_noj0621

fig = plt.figure(figsize=(6,6))
ROIPlotter(o['sqrt_ts'],roi=gta.roi).plot(vmin=0,vmax=5,levels=[3,5,7,9],subplot=111,cmap='magma')
plt.gca().set_title('sqrt(TS)')

"""
Explanation: We can get alternative assessment of the model by generating a TS map of the region.  Again we see a hotspot at the position of 3FGL J0621.0+2514 which we excluded from the model.
End of explanation
"""


import pprint
pprint.pprint(gta.config['extension'])

"""
Explanation: Measuring Source Extension
After optimizing the model we are ready to run an extension analysis on IC 443.  As reported in Abdo et al. 2010, this source has a spatial extension of 0.27 deg $\pm$ 0.01 (stat).  We can run an extension test of this source by calling the extension method with the source name.  The extension method has a number of options which can be changed at runtime by passing keyword arguments.  To see the default settings we can look at the extension sub-dictionary of the config property of our GTAnalysis instance.
End of explanation
"""


ext_gauss = gta.extension('3FGL J0617.2+2234e',width=np.linspace(0.25,0.30,11).tolist())
gta.write_roi('ext_gauss_fit')

"""
Explanation: By default the method will use a 2D Gaussian source template and scan the width parameter between 0.00316 and 1 degrees in 26 steps.  The width parameter can be used to provide an explicit  vector of points for the scan.  Since we know the extension of IC 443 is well localized around 0.27 deg we use a width vector centered around this point.  The analysis results are returned as an output dictionary and are also written to the internal source object of the GTAnalysis instance.
End of explanation
"""


plt.figure(figsize=(8,6))
plt.plot(ext_gauss['width'],ext_gauss['dloglike'],marker='o')
plt.gca().set_xlabel('Width [deg]')
plt.gca().set_ylabel('Delta Log-Likelihood')
plt.gca().axvline(ext_gauss['ext'])
plt.gca().axvspan(ext_gauss['ext']-ext_gauss['ext_err_lo'],ext_gauss['ext']+ext_gauss['ext_err_hi'],
                      alpha=0.2,label='This Measurement',color='b')

plt.gca().axvline(0.27,color='k')
plt.gca().axvspan(0.27-0.01,0.27+0.01,alpha=0.2,label='Abdo et al. 2010',color='k')
plt.gca().set_ylim(2030,2070)
plt.gca().set_xlim(0.20,0.34)

plt.annotate('TS$_{\mathrm{ext}}$ = %.2f\nR$_{68}$ = %.3f $\pm$ %.3f'%
            (ext_gauss['ts_ext'],ext_gauss['ext'],ext_gauss['ext_err']),xy=(0.05,0.05),xycoords='axes fraction')
plt.gca().legend(frameon=False)

"""
Explanation: To inspect the results of the analysis we can make a plot of the likelihood profile.  From this we can see that the spatial extension is in good agreement with the value from Abdo et al. 2010.
End of explanation
"""


ext_gauss_free = gta.extension('3FGL J0617.2+2234e',width=np.linspace(0.25,0.30,11).tolist(),free_radius=1.0)
print 'Fixed Sources: %f +/- %f'%(ext_gauss['ext'],ext_gauss['ext_err'])
print 'Free Sources:  %f +/- %f'%(ext_gauss_free['ext'],ext_gauss_free['ext_err'])

"""
Explanation: As an additional cross-check we can look at what happens when we free sources and rerun the extension analysis.
End of explanation
"""


