%matplotlib inline
import matplotlib
import matplotlib.cm as cm
import numpy as np
import matplotlib.pyplot as plt
import jinja2, json, sys
from math import log, fabs, pi, cos, sin
from scipy.stats import ks_2samp
try:
    try:
        import ruamel.yaml as yaml
    except ImportError:
        # anaconda packs it as a standalone module (underscore instead of dot)
        import ruamel_yaml as yaml
except ImportError:
    import yaml

number_of_replicas = 4
scale_array = np.geomspace(1, 0.1, number_of_replicas)
temper = True # run with parallel tempering or not?

# generate input files from the template file; one for each replica
with open("temper.yml") as template_file:
    template = jinja2.Template(template_file.read())
    for replica, scale in enumerate(scale_array):
        contents = template.render(scale=scale, micro=20000)
        yml = yaml.safe_load(contents) # parse contents as YAML
        if not temper: 
            del yml["moves"][1]
        file = open("mpi{}.input.json".format(replica), 'w')
        json.dump(yml, file, indent=4) # save JSON file
        file.close()

"""
Explanation: Parallel Tempering
This is a simple example of how to use parallel tempering in Faunus. The example is a particle moving in one dimension
and exposed to an osciallating potential. The example is taken from the book of Frenkel and Smit for parallel tempering.
End of explanation
"""


%%bash
if [[ -z "${FAUNUS_EXECUTABLE}" ]]; then
  mpirun -np 4 faunus -i input.json
else
  echo "Seems we're running CTest - use Faunus target from CMake"
  "${MPIEXEC}" -np 4 "${FAUNUS_EXECUTABLE}" -i input.json --nobar
fi

"""
Explanation: Run simulation
End of explanation
"""


def energy(x):
    s = 1 + sin(2*pi*x)
    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
    return 10000000

u_vec = np.vectorize(energy)
offset = 1e10 # used to offset pmf to match above energy function
for replica in range(number_of_replicas):
    if (replica==0):
        x = np.loadtxt("mpi{}.x.dat.gz".format(replica), usecols=[1])
        hist, bins = np.histogram(x, bins=150)
        x = (bins[:-1] + bins[1:]) / 2
        pmf = -np.log(hist)
        if (pmf.min() < offset):
            offset = pmf.min();
        plt.plot(x, pmf, label='{}'.format(replica))
        
plt.legend(loc=0, frameon=False, title='replica')
plt.plot(x, u_vec(x) + offset, 'k--', alpha=0.6)
plt.xlabel(r'x')
plt.ylabel(r'PMF ($k_BT$)');

"""
Explanation: Plot sampled histogram against expected result
End of explanation
"""


statistic, pvalue = ks_2samp(pmf, u_vec(x)+offset)
if (pvalue < 0.95):
    sys.exit(1)

"""
Explanation: Statistical test of output against expected result
End of explanation
"""


