%matplotlib inline
import matplotlib
import matplotlib.cm as cm
import numpy as np
import matplotlib.pyplot as plt
from math import log, fabs, pi, cos, sin

"""
Explanation: Penalty Function
This is a simple example of how to use the flat histogram method in in Faunus
End of explanation
"""


%%bash
rm -fR penalty.dat
if [[ -z "${FAUNUS_EXECUTABLE}" ]]; then
  yason.py penalty.yml | faunus --nobar
else
  echo "Seems we're running CTest - use Faunus target from CMake"
  "${YASON_EXECUTABLE}" penalty.yml | "${FAUNUS_EXECUTABLE}" --nobar
fi

"""
Explanation: Run penalty example
End of explanation
"""


PMF = -np.loadtxt('penalty.dat', skiprows=1).T
HIST = np.loadtxt('penalty-histogram.dat').T
HIST = np.log(HIST/HIST.min())
print('Penalty maximum barrier    =', fabs(PMF.max()-PMF.min()), 'kT')
print('Maximum simulation barrier =', HIST.max(), 'kT')

fig, (ax1, ax2) = plt.subplots(1,2, sharex=True, sharey=True)
fig.set_size_inches(10,4)
im1 = ax1.imshow(PMF, interpolation='bilinear', cmap=cm.afmhot, extent=[-2,2,-2,2],
               origin='lower', vmin=PMF.min(), vmax=PMF.max())
im2 = ax2.imshow(HIST, interpolation='bilinear', cmap=cm.afmhot, extent=[-2,2,-2,2],
               origin='lower', vmin=HIST.min(), vmax=HIST.max())
ax1.set_title('Potential of Mean Force')
ax2.set_title('Barriers during simulation')
ax1.set_xlabel(r'$x$')
ax2.set_xlabel(r'$x$')
ax1.set_ylabel(r'$y$')
ax2.set_ylabel(r'$y$')
cbar1 = fig.colorbar(im1, ax=ax1)
cbar2 = fig.colorbar(im2, ax=ax2)
cbar1.set_label('$k_BT$')
cbar2.set_label('$k_BT$')
plt.show()

"""
Explanation: Load penalty function
This loads the penalty function from disk which is simply the negative of the final potential of mean force. Also load the final histogram, which ideally should be flat over the reaction coordinates.
End of explanation
"""


def energy(x,y):
    ''' Energy function used during simulation '''
    s=1+sin(2*pi*x) + cos(2*pi*y)
    if x>=-2.00 and x<=-1.25: return 1*s
    if x>=-1.25 and x<=-0.25: return 2*s
    if x>=-0.25 and x<= 0.75: return 3*s
    if x>= 0.75 and x<= 1.75: return 4*s
    if x>= 1.75 and x<= 2.00: return 5*s
    raise Exception

x = y = np.arange(-2.0, 2.0, 0.1)
X, Y = np.meshgrid(x,y) 
Z = np.vectorize(energy)(X,Y)
Z = Z-Z.max()
fig, ax = plt.subplots()
im = ax.imshow(Z, interpolation='bilinear', cmap=cm.afmhot,
               origin='lower', extent=[-2, 2, -2, 2],
               vmax=Z.max(), vmin=Z.min())
ax.set_title('External potential')
fig.colorbar(im, ax=ax)

"""
Explanation: External Potential
Now let's plot the external potential applied during simulation. The PMF should tend towards this.
End of explanation
"""


