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

# <form action="index.ipynb">
#     <input type="submit" value="Return to Index" style="background-color: green; color: white; width: 150px; height: 35px; float: right"/>
# </form>
# 
# # Radiation Source Localization
# Author(s): Paul Miles, Jason Hite | Date Created: July 19, 2019
# 
# This example demonstrates how MCMC methods can be used to locate unknown sources of gamma radiation in an urban environment.  For this demonstration we consider a simulated 250m x 178m block of downtown Washington D.C.
# 
# We utilize a highly simplified radiation transport model which ignores scattering.  The model consists of the following components.  There are $N_d$ gamma detectors located at positions $\mathbf{r}_i, i = 1, ..., N_d$.  A ray tracing algorithm is implemented for a source located at $\mathbf{r}_0 = (x_0, y_0)$ with intensity $I_0$.  Buildings affect the readings of the gamma detectors, so we have modeled their impact using a exponential decay function that depends on the building dimensions and average cross-section. For the $i^{th}$ detector, we let $N_i$ denote the number of buildings between $\mathbf{r}_0$ and $\mathbf{r}_i$.  For $n_i = 1,...,N_i$, we let $s_{n_i}$ and $\sigma_{n_i}$ denote the length of the building segment and the average cross-section of the $n_i^{th}$ building, respectively.  The surface area, efficiency, and measurement duration for the $i^{th}$ detector are denoted by $A_i, \varepsilon_i,$ and $\Delta t_i$, respectively.  Finally, $B_i$ represents the expected background count rate at position $\mathbf{r}_i$. 
# $$
# \Gamma_i = \frac{\Delta t_i\varepsilon_iA_iI_0}{4\pi|\mathbf{r}_i - \mathbf{r}_0|^2}\exp\Big(-\sum_{n_i=1}^{N_i}\sigma_{n_i}s_{n_i}\Big) + E[B_i\Delta t_i]
# $$
# 
# This model is implemented in the open source Python package [gefry3](https://github.com/jasonmhite/gefry3).  The input deck for this example can be found [here](https://github.com/jasonmhite/gefry3/blob/master/examples/g3_deck.yml).  For additional reading on this area of research, please refer to the following articles:
# 
# - Hite, J., & Mattingly, J. (2018). Bayesian Metropolis Methods for Source Localization in an Urban Environment. Radiation Physics and Chemistry. [https://doi.org/10.1016/j.radphyschem.2018.06.024](https://doi.org/10.1016/j.radphyschem.2018.06.024)
# - Ştefănescu, R., Schmidt, K., Hite, J., Smith, R. C., & Mattingly, J. (2017). Hybrid optimization and Bayesian inference techniques for a non‐smooth radiation detection problem. International Journal for Numerical Methods in Engineering, 111(10), 955-982. [https://doi.org/10.1002/nme.5491](https://doi.org/10.1002/nme.5491)
# - Hite, J. M., Mattingly, J. K., Schmidt, K. L., Ştefănescu, R., & Smith, R. (2016, September). Bayesian metropolis methods applied to sensor networks for radiation source localization. In Multisensor Fusion and Integration for Intelligent Systems (MFI), 2016 IEEE International Conference on (pp. 389-393). IEEE. [https://doi.org/10.1109/MFI.2016.7849519](https://doi.org/10.1109/MFI.2016.7849519)
# 
# __Acknowledgment:__
# Funding for this research provided by the [Consortium for Nonproliferation Enabling Capabilities (CNEC)](https://cnec.ncsu.edu/).
# 
# __Package Installation:__
# All packages can be install via
# ```python
# pip install package_name
# ```
# with the exception of [gefry3](https://github.com/jasonmhite/gefry3), which requires
# ```python
# pip install git+https://github.com/jasonmhite/gefry3
# ```
# Advanced statistical plots are generated using the [mcmcplot](https://prmiles.wordpress.ncsu.edu/codes/python-packages/mcmcplot/) package.
# 
# To run this particular notebook, please execute the two cells below.  The first cell will install [gefry3](https://github.com/jasonmhite/gefry3) and the second [descartes](https://pypi.org/project/descartes/).

# In[ ]:


pip install git+https://github.com/jasonmhite/gefry3


# In[ ]:


pip install descartes


# In[7]:


# import required packages
import os
import gefry3
import numpy as np
from pymcmcstat.MCMC import MCMC
from descartes.patch import PolygonPatch
import seaborn as sns
import matplotlib.pyplot as plt
np.seterr(over = 'ignore')
NOBS = 1 # number of observations for each detector
NS = int(1e3) # number of MCMC simulations


# # Define Radiation Model, Problem Domain, and Background Radiation

# In[8]:


# Setup Radiation Model - Downtown Urban Environment Simulation
P = gefry3.read_input_problem('data_files' + os.sep + 'g3_deck.yml', 'Simple_Problem')
# Setup background
nominal_background = 300
dwell = np.array([detector.dwell for detector in P.detectors])
B0 = nominal_background * dwell
# True source and intensity
S0 = P.source.R  # m
I0 = P.source.I0  # Bq
nominal = P(S0, I0)
ndet = len(P.detectors)
# Define observations
observations = np.random.poisson(
    nominal + nominal_background * dwell,
    size=(
        NOBS,
        ndet,
    ),
)
observations = observations.astype(np.float64)
# Setup parameter bounds
XMIN, YMIN, XMAX, YMAX = P.domain.all.bounds
IMIN, IMAX = I0 * 0.1, I0 * 10
XMIN, YMIN, XMAX, YMAX = P.domain.all.bounds
print(observations)
print('S0 = {} m, I0 = {} Bq'.format(S0, I0))


# # Define Sum-of-Squares Error Function.

# In[9]:


# Radiation Sum of Squares Function
def radiation_ssfun(theta, data):
    x, y, I = theta
    model, background = data.user_defined_object[0]
    output = model((x, y), I) + background

    res = data.ydata[0] - output
    ss = (res ** 2).sum(axis = 0)
    return ss


# # Initialize MCMC Object and Setup Simulation
# Note, we pass the radiation model into our residual function via the data structure.  Also note, the background is included as a parameter, but it is not sampled during the simulation.

# In[10]:


# Initialize MCMC object
mcstat = MCMC()
# setup data structure for dram
mcstat.data.add_data_set(
    x=np.zeros(observations.shape),
    y=observations,
    user_defined_object=[
        P,
        B0,
    ],
)
mcstat.parameters.add_model_parameter(
    name='$x_0$',
    theta0=140.,
    minimum=XMIN,
    maximum=XMAX,
)
mcstat.parameters.add_model_parameter(
    name='$y_0$',
    theta0=85.,
    minimum=YMIN,
    maximum=YMAX,
)
mcstat.parameters.add_model_parameter(
    name='$I_0$',
    theta0=3.0e9,
    minimum=IMIN,
    maximum=IMAX
)

mcstat.simulation_options.define_simulation_options(
    nsimu=NS,
    updatesigma=False,
    method='dram',
    adaptint=100,
    verbosity=1,
    waitbar=1,
    save_to_json=False,
    save_to_bin=False,
)

mcstat.model_settings.define_model_settings(
    sos_function=radiation_ssfun,
    sigma2=observations.sum(axis=0),
)


# # Run Simulation

# In[11]:


# Run mcmcrun
mcstat.run_simulation()


# # Process Results and Display Statistics

# In[14]:


from pymcmcstat import mcmcplot as mcp
import seaborn as sns

results = mcstat.simulation_results.results
# specify burnin period
burnin = int(results['nsimu']/2)
# display chain statistics
chain = results['chain'][burnin:, :]
sschain = results['sschain'][burnin:, :]
names = results['names']
mcstat.chainstats(chain, results)
print('Acceptance rate: {:6.4}%'.format(100*(1 - results['total_rejected'])))
print('Model Evaluations: {}'.format(results['nsimu']
                                     - results['iacce'][0]
                                     + results['nsimu']))
# generate mcmc plots
sns.set_style('white')
mcp.plot_density_panel(chain, names)
mcp.plot_chain_panel(chain, names)
mcp.plot_pairwise_correlation_panel(chain, names)


# # Generate Plot of Urban Environment With Source Location Distributions
# First we define a function to allow us to add axes within a figure.

# In[15]:


def add_subplot_axes(ax,rect,**kwargs):
    fig = plt.gcf()
    box = ax.get_position()
    width = box.width
    height = box.height
    inax_position  = ax.transAxes.transform(rect[0:2])
    transFigure = fig.transFigure.inverted()
    infig_position = transFigure.transform(inax_position)    
    x = infig_position[0]
    y = infig_position[1]
    width *= rect[2]
    height *= rect[3]  # <= Typo was here
    subax = fig.add_axes([x,y,width,height],**kwargs)
    x_labelsize = subax.get_xticklabels()[0].get_size()
    y_labelsize = subax.get_yticklabels()[0].get_size()
    x_labelsize *= rect[2]**0.5
    y_labelsize *= rect[3]**0.5
    subax.xaxis.set_tick_params(labelsize=x_labelsize)
    subax.yaxis.set_tick_params(labelsize=y_labelsize)
    return subax


# We set the context of our displays to "talk" which will improve their visibility for presentation.  To begin, we will plot the joint distribution relative to the true source location.

# In[18]:


sns.set_context('talk')
settings = dict(jointplot=dict(kind='kde', color='b'),
                marginal_kws=dict(alpha=1))
fig0 = mcp.plot_joint_distributions(
    chain[:, 0:2],
    names,
    settings=settings
);
# add source
fig0.ax_joint.plot(P.source.R[0], P.source.R[1], 'or',
                      linewidth=3, markerfacecolor='r',
                      markersize=10, label='True Source');


# Next, we plot the joint distributions relative to our entire simulated urban environment.  We have added an insert to highlight the relative positioning of the posterior distributions and the true source location.

# In[25]:


settings = dict(jointplot=dict(kind='kde', color='b'),
                marginal_kws=dict(alpha=1))
fig = mcp.plot_joint_distributions(
    chain[:, 0:2],
    names,
    settings=settings);
fig.ax_marg_x.set_xlim([XMIN, XMAX])
fig.ax_marg_y.set_ylim([YMIN, YMAX])
ax = fig.ax_joint
# add source
ax.plot(P.source.R[0], P.source.R[1],
        linewidth=3,
        markerfacecolor='r',
        markersize=8,
        marker='*')
# add buildings to plot
for building in P.domain.solids:
    patch = PolygonPatch(building.geom, alpha=0.3,
                         facecolor='g', edgecolor='k')
    ax.add_patch(patch)
# Create insert
rect = [0.1, 0.1, 0.3, 0.3]
ax1 = add_subplot_axes(ax, rect, facecolor='w')
sns.kdeplot(chain[:, 0], chain[:, 1], ax=ax1, color='b')
mu = results['mean']
xlimits = [mu[0] - 7.5, mu[0] + 7.5]
ylimits = [mu[1] - 7.5, mu[1] + 7.5]
ax1.set_xlim(xlimits)
ax1.set_ylim(ylimits)
# add source
ax1.plot(P.source.R[0], P.source.R[1],
         linewidth=3,
         markerfacecolor='r',
         markersize=20,
         marker='*')
# add buildings to plot
for building in P.domain.solids:
    patch = PolygonPatch(building.geom, alpha=0.3,
                         facecolor='g', edgecolor='k')
    ax1.add_patch(patch)
# Create lines outlining insert location
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
mark_inset(ax, ax1, loc1=2, loc2=4,
           facecolor="none", edgecolor="k", linewidth=2)
ax1.xaxis.set_visible(False)
ax1.yaxis.set_visible(False)


# We observe that the poster distributions for the source location are within a couple meters of the true source location (red star - $R_{true} = [158, 98]$ meters).  Furthermore, the true location is within the uncertainty bounds of our posteriors.