#!/usr/bin/env python
# coding: utf-8

# # Joint fitting XRT and GBM data with XSPEC models
# 
# ### Goals
# 
# 3ML is designed to properly joint fit data from different instruments with thier instrument dependent likelihoods.
# We demostrate this with joint fitting data from GBM and XRT while simultaneously showing hwo to use the XSPEC models form **astromodels**
# 
# ### Setup
# 
# You must have you HEASARC initiated so that **astromodels** can find the XSPEC libraries. 
# 
# 
# 

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('matplotlib', 'notebook')
from threeML import *
import os


# ## Load XRT data
# 
# Make a likelihood for the XRT including all the appropriate files

# In[3]:


trigger="GRB110731A"
dec=-28.546
ra=280.52
xrt_dir='xrt'
xrt = SwiftXRTLike("XRT",pha_file=os.path.join(xrt_dir,"xrt_src.pha"),
                   bak_file=os.path.join(xrt_dir,"xrt_bkg.pha"),
                   rsp_file=os.path.join(xrt_dir,"xrt.rmf"),
                   arf_file=os.path.join(xrt_dir,"xrt.arf"))



xrt.view_count_spectrum()


# ## Load GBM data
# 
# Load all the GBM data you need and make appropriate background, source time, and energy selections. Make sure to check the light curves! 

# In[9]:


data_dir_gbm=os.path.join('gbm','bn110731A')
trigger_number = 'bn110731465'
gbm_data = download_GBM_trigger_data(trigger_number,detectors=['n3'],destination_directory=data_dir_gbm,compress_tte=True)

# Select the time interval
src_selection = "100.169342-150.169342"

nai3 = FermiGBMTTELike('NAI3',
                         os.path.join(data_dir_gbm,"glg_tte_n3_bn110731465_v00.fit.gz"),
                         "20-90,160-250", # background selection
                         src_selection,          # source interval
                         rsp_file=os.path.join(data_dir_gbm, "glg_cspec_n3_bn110731465_v00.rsp2"))


# View the light curve

# In[10]:


nai3.view_lightcurve(20,250)


# Make energy selections and check them out

# In[11]:


nai3.set_active_measurements("8-900")
nai3.view_count_spectrum()


# ## Setup the model
# 
# **astromodels** allows you to use XSPEC models if you have XSPEC installed.
# Set all the normal parameters you would in XSPEC and build a model the normal **3ML/astromodel** way!
# 

# In[6]:


xspec_abund('angr')

spectral_model =  XS_phabs()* XS_zphabs() * XS_powerlaw()


spectral_model.nh_1=0.101
spectral_model.nh_1.fix = True

spectral_model.nh_2=0.1114424
spectral_model.nh_2.fix = True

spectral_model.redshift_2 = 0.618
spectral_model.redshift_2.fix =True


# In[7]:


spectral_model.display()


# ## Setup the joint likelihood
# 
# Create a point source object and model. 
# 
# Load the data into a data list and create the joint likelihood
# 

# In[8]:


ptsrc = PointSource(trigger,ra,dec,spectral_shape=spectral_model)
model = Model(ptsrc)


# In[9]:


data = DataList(xrt,nai3)

jl = JointLikelihood(model, data, verbose=False)
model.display()


# ## Fitting

# ### Maximum Likelihood style

# In[10]:


res = jl.fit()


# In[11]:


res = jl.get_errors()


# In[12]:


res = jl.get_contours(spectral_model.phoindex_3,1.5,2.5,50)


# In[13]:


res = jl.get_contours(spectral_model.norm_3,.1,.3,25,spectral_model.phoindex_3,1.5,2.5,50)


# ### And then go Bayesian!

# In[14]:


spectral_model.phoindex_3.prior = Uniform_prior(lower_bound=-5.0, upper_bound=5.0)
spectral_model.norm_3.prior = Log_uniform_prior(lower_bound=1E-5, upper_bound=1)


# In[15]:


bayes = BayesianAnalysis(model, data)


# In[16]:


samples = bayes.sample(n_walkers=50,burn_in=100, n_samples=1000)


# In[17]:


fig = bayes.corner_plot(plot_contours=True, plot_density=False)


# In[18]:


bayes.get_highest_density_interval()


# In[12]:


cleanup_downloaded_GBM_data(gbm_data)


# In[ ]: