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

# # Bayesian Plotting
# 
# 3ML's Bayesian capabilites include the ability to examine marginal distributions and compare model parameters between fits. 
# 
# This tutorial demonstrates some of these capabilities. 
# 
# Simple corner plots require the corner package. 
# 
# ChainConsumer is required for more advanced plots. https://github.com/Samreay/ChainConsumer
# 
# 

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('matplotlib', 'notebook')

from threeML import *


# ## Data setup and sampling
# 
# 
# We will use Fermi GBM data for this example. The following code sets up the data and Bayesian fits. 
# 
# We use emcee to sample, but any of the included samplers will work.
# 

# In[6]:


data_dir = os.path.join('gbm','bn080916009')
trigger_number = 'bn080916009'
src_selection = '0-71'
# Download the data

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




nai3 = FermiGBMTTELike('NAI3',
                         os.path.join(data_dir, "glg_tte_n3_bn080916009_v01.fit.gz"),
                         "-10-0, 100-200",
                         src_selection,
                         rsp_file=os.path.join(data_dir, "glg_cspec_n3_bn080916009_v00.rsp2"))

bgo0 = FermiGBMTTELike('BGO0',
                         os.path.join(data_dir, "glg_tte_b0_bn080916009_v01.fit.gz"),
                         "-10-0,100-200",
                         src_selection,
                         rsp_file=os.path.join(data_dir, "glg_cspec_b0_bn080916009_v00.rsp2"))


nai3.set_active_measurements("10.0-30.0", "40.0-900.0")

bgo0.set_active_measurements("250-43000")



triggerName = 'bn080916009'
ra = 121.8
dec = -61.3


data_list = DataList(nai3,bgo0 )

band = Band()


GRB = PointSource( triggerName, ra, dec, spectral_shape=band )

model = Model( GRB )


band.K.prior = Log_uniform_prior(lower_bound=1e-4, upper_bound=3)
band.xp.prior = Log_uniform_prior(lower_bound=10, upper_bound=1e5)

# or use the set_uninformative_prior method, which will use as lower_bound
# and upper_bound the current boundaries for the parameter. Such boundaries
# must exists and be finite

band.alpha.set_uninformative_prior(Uniform_prior)
band.beta.set_uninformative_prior(Uniform_prior)

bayes = BayesianAnalysis(model, data_list)

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


# ### Simple corner plot
# 
# We can look at the marginal distributions for the parameters with corner.
# 
# 
# 
# 

# In[7]:


_=bayes.corner_plot()


# ### Advanced corner plot
# 
# Using chain consumer, we can have more options for our plots.
# 
# Here is the default plot
# 
# 

# In[8]:


_=bayes.corner_plot_cc()


# We can look at the cloud  or sampled points. Additionally, we can change the sigma level of the contours. Colors must be provided as HTML codes to facilitate the contour levels.

# In[11]:


_=bayes.corner_plot_cc(sigmas=[0,1,3,5],color_params='alpha',shade_alpha=.4)


# It is also possible to change parameter names. A dictionary of {old:new} names can be provided. 

# In[13]:


renamed ={'xp':'$E_p$',
          'alpha':'$\\alpha$',
          'beta':'$\\beta$'}

_=bayes.corner_plot_cc(renamed_parameters=renamed,shade=False)


# ### Comparing distributions
# 
# Perhaps we want to look at two fits that are related. Here we use the example of Band+ Blackbody and will examine how the band parameters change when we included the extra component.
# 
# 

# In[14]:


bandpl = Band() + Powerlaw()


GRB2 = PointSource( triggerName, ra, dec, spectral_shape=bandpl )

model2 = Model( GRB2 )


bandpl.K_1.prior = Log_uniform_prior(lower_bound=1e-4, upper_bound=3)
bandpl.xp_1.prior = Log_uniform_prior(lower_bound=10, upper_bound=1e5)


bandpl.alpha_1.set_uninformative_prior(Uniform_prior)
bandpl.beta_1.set_uninformative_prior(Uniform_prior)

bandpl.K_2.prior = Log_uniform_prior(lower_bound=1e-4, upper_bound=3)
bandpl.index_2.set_uninformative_prior(Uniform_prior)

bayes2 = BayesianAnalysis(model2, data_list)

samples2 = bayes2.sample(n_walkers=50,burn_in=200, n_samples=1000)


# We can examine the marginals of all parameters first

# In[15]:


renamed ={'xp_1':'$E_p$',
          'alpha_1':'$\\alpha$',
          'beta_1':'$\\beta$'}

_ = bayes2.corner_plot_cc(renamed_parameters=renamed)


# #### Comparing parameters
# 
# Now let's compare to the original Band fit. We will have to do some manipulation because composite models from astromodels append numbers to clearly identify components.
# 
# We need to rename parameters so that those we wish to compare have the same names:
# 
# 
# 

# In[19]:


renamed ={'xp_1':'xp',
          'alpha_1':'alpha',
          'beta_1':'beta',
          'K_1':'K'}

_ = bayes2.compare_posterior(bayes,renamed_parameters=renamed,color_params=['index 2','xp'])


# We can see that there is a shift in some parameters of the Band function due to the additional components.
# 
# Perhaps we want to focus on a subset of the parameters to draw more attention to them. This is easily accomplished by passing the parameters we are interested in. **NOTE:** We must use the original parameter names so that they are picked up!

# In[20]:


_=bayes2.compare_posterior(bayes,renamed_parameters=renamed,shade_alpha=.7,parameters=['xp','alpha','index 2'],colors=["#740595","#FCA816"])


# In[21]:


mc = ModelComparison(bayes,bayes2)


# In[22]:


mc.report()


# In[ ]:


cleanup_downloaded_GBM_data(gbm_data)