import pints
import pints.toy
import numpy as np
import matplotlib.pyplot as plt

# Create log pdf
log_pdf = pints.toy.TwistedGaussianLogPDF(dimension=2)

# Contour plot of pdf
levels = np.linspace(-50, -1, 20)
x = np.linspace(-50, 50, 250)
y = np.linspace(-100, 20, 250)
X, Y = np.meshgrid(x, y)
Z = [[log_pdf([i, j]) for i in x] for j in y]
plt.contour(X, Y, Z, levels = levels)
plt.show()

"""
Explanation: Twisted Gaussian (banana) toy LogPDF
This distribution (described here) has a curved "banana" shape. The problem can be more more or less difficult by changing the "bananicity" parameter b.
End of explanation
"""


direct = log_pdf.sample(15000)

plt.contour(X, Y, Z, levels=levels, colors='k', alpha=0.2)
plt.scatter(direct[:, 0], direct[:, 1], alpha=0.2)
plt.xlim(-50, 50)
plt.ylim(-100, 20)
plt.show()

"""
Explanation: We can also sample independently from this toy LogPDF, and add that to the visualisation:
End of explanation
"""


# Create an adaptive covariance MCMC routine
x0 = np.random.uniform(-25, 25, size=(3, 2))
mcmc = pints.MCMCController(log_pdf, 3, x0, method=pints.HaarioBardenetACMC)

# Stop after 10000 iterations
mcmc.set_max_iterations(3000)

# Disable logging
mcmc.set_log_to_screen(False)

# Run!
print('Running...')
chains = mcmc.run()
print('Done!')

# Discard warm-up
chains = [chain[1000:] for chain in chains]

stacked = np.vstack(chains)
plt.contour(X, Y, Z, levels=levels, colors='k', alpha=0.2)
plt.scatter(stacked[:, 0], stacked[:, 1], alpha=0.2)
plt.xlim(-50, 50)
plt.ylim(-100, 20)
plt.show()

"""
Explanation: We now try to sample from the distribution with MCMC:
End of explanation
"""


print(log_pdf.kl_divergence(stacked))
print(log_pdf.kl_divergence(direct))

"""
Explanation: Now check how close the result is to the expected result, using the Kullback-Leibler divergence, and compare this to the result from sampling directly.
End of explanation
"""


# Create an adaptive covariance MCMC routine
x0 = np.random.uniform(-25, 25, size=(3, 2))
sigma0 = [5, 5]
mcmc = pints.MCMCController(log_pdf, 3, x0, method=pints.HamiltonianMCMC, sigma0=sigma0)

# Stop after 10000 iterations
mcmc.set_max_iterations(3000)

# Disable logging
mcmc.set_log_to_screen(False)

# Run!
print('Running...')
chains = mcmc.run()
print('Done!')

chains1 = [chain[1000:] for chain in chains]
stacked = np.vstack(chains1)

print(log_pdf.kl_divergence(stacked))
print(log_pdf.kl_divergence(direct))

"""
Explanation: Hamiltonian Monte Carlo fares much better on this curved density.
End of explanation
"""


divergent_transitions = mcmc.samplers()[0].divergent_iterations()
plt.contour(X, Y, Z, levels=levels, colors='k', alpha=0.2)
plt.plot(chains[0][:, 0], chains[0][:, 1], alpha=0.5)
plt.scatter(chains[0][divergent_transitions, 0], chains[0][divergent_transitions, 1], color='red')
plt.xlim(-50, 50)
plt.ylim(-100, 20)
plt.show()

"""
Explanation: Visualising the path of a single HMC chain, we see that it moves naturally along contours although does occassionally suffer from divergent iterations (red dots) in the neck of the banana due to the varying posterior curvature throughout the domain.
End of explanation
"""


