# system functions that are always useful to have
import time, sys, os

# basic numeric setup
import numpy as np

# inline plotting
%matplotlib inline

# plotting
import matplotlib
from matplotlib import pyplot as plt

# seed the random number generator
rstate = np.random.default_rng(736109)

# re-defining plotting defaults
from matplotlib import rcParams
rcParams.update({'xtick.major.pad': '7.0'})
rcParams.update({'xtick.major.size': '7.5'})
rcParams.update({'xtick.major.width': '1.5'})
rcParams.update({'xtick.minor.pad': '7.0'})
rcParams.update({'xtick.minor.size': '3.5'})
rcParams.update({'xtick.minor.width': '1.0'})
rcParams.update({'ytick.major.pad': '7.0'})
rcParams.update({'ytick.major.size': '7.5'})
rcParams.update({'ytick.major.width': '1.5'})
rcParams.update({'ytick.minor.pad': '7.0'})
rcParams.update({'ytick.minor.size': '3.5'})
rcParams.update({'ytick.minor.width': '1.0'})
rcParams.update({'font.size': 30})

import dynesty

"""
Explanation: Eggbox
Setup
First, let's set up some environmental dependencies. These just make the numerics easier and adjust some of the plotting defaults to make things more legible.
End of explanation
"""


# define the eggbox log-likelihood
tmax = 5.0 * np.pi
def loglike(x):
    t = 2.0 * tmax * x - tmax
    return (2.0 + np.cos(t[0] / 2.0) * np.cos(t[1] / 2.0)) ** 5.0

# define the prior transform
def prior_transform(x):
    return x

# plot the log-likelihood surface
plt.figure(figsize=(10., 10.))
axes = plt.axes(aspect=1)
xx, yy = np.meshgrid(np.linspace(0., 1., 50),
                     np.linspace(0., 1., 50))
L = loglike(np.array([xx, yy]))
axes.contourf(xx, yy, L, 50, cmap=plt.cm.Purples_r)
plt.title('Log-Likelihood Surface')
plt.xlabel(r'$x$')
plt.ylabel(r'$y$')

# ln(evidence)
lnz_truth = 235.88

"""
Explanation: The "Eggbox" likelihood is a helpful case that demonstrate Nested Sampling's ability to properly sample/integrate over multi-modal distributions.
End of explanation
"""


dsampler = dynesty.DynamicNestedSampler(loglike, prior_transform, ndim=2,
                                        bound='multi', sample='unif', rstate=rstate)

"""
Explanation: Let's sample from this distribution using multi-ellipsoidal decomposition.
End of explanation
"""


dsampler.run_nested(dlogz_init=0.01, nlive_init=500, nlive_batch=500,
                    wt_kwargs={'pfrac': 0.0}, stop_kwargs={'pfrac': 0.0})
dres_z = dsampler.results

"""
Explanation: Let's first start by sampling with a focus on deriving the evidence.
End of explanation
"""


dsampler.reset()
dsampler.run_nested(dlogz_init=0.01, nlive_init=500, nlive_batch=500,
                    wt_kwargs={'pfrac': 1.0}, stop_kwargs={'pfrac': 1.0})
dres_p = dsampler.results

"""
Explanation: Now let's add samples with a focus on deriving the posterior.
End of explanation
"""


dsampler = dynesty.DynamicNestedSampler(loglike, prior_transform, ndim=2,
                                        bound='balls', sample='unif', rstate=rstate)
dsampler.run_nested(dlogz_init=0.01, nlive_init=500, nlive_batch=500,
                    wt_kwargs={'pfrac': 0.0}, stop_kwargs={'pfrac': 0.0})
dres_z2 = dsampler.results

"""
Explanation: Finally, let's switch to using 'balls'.
End of explanation
"""


from dynesty import plotting as dyplot

fig, axes = dyplot.runplot(dres_z, color='blue')
fig, axes = dyplot.runplot(dres_z2, color='red', lnz_truth=lnz_truth,
                           truth_color='black', fig=(fig, axes))
fig.tight_layout()

fig, axes = dyplot.cornerplot(dres_p, quantiles=None, color='darkviolet',
                              span=[[0, 1], [0, 1]],
                              fig=plt.subplots(2, 2, figsize=(10, 10)))

"""
Explanation: Let's see how we did.
End of explanation
"""


