from __future__ import division, unicode_literals, print_function
import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import os.path, os, sys, json
plt.rcParams.update({'font.size': 18, 'figure.figsize': [8.0, 6.0]})

"""
Explanation: Monte Carlo Simulation and Electric Multipoles
Mikael Lund, 2017
This is a Jupyter notebook (http://jupyter.org) for studying the interaction between multipolar particles.
The notebook can be run from a browser by first installing Jupyter through i.e. the Anaconda project
(https://www.continuum.io).

Load required packages
End of explanation
"""


%%bash
if [ ! -d "faunus" ]; then
    git clone https://github.com/mlund/faunus.git
    cd faunus
    git checkout e2d7b625516f68343460948206f16f23a98eb868
    git apply ../faunus.diff
    CXX=clang++ CC=clang cmake -DCMAKE_BUILD_TYPE=Release .
    make example_twobody
else
    cd faunus
    make example_twobody
fi

"""
Explanation: Download the MC program "Faunus" and compile a particular version
End of explanation
"""


%%writefile monopole.pqr
ATOM      1 X    MP      1       0.000    0.000    0.000 -1.000 1.000

%%writefile dipole.pqr
ATOM      1 X    MP      1      -0.250    0.000    0.000 -3.000 1.000                               
ATOM      2 X    MP      1       0.000    0.000    0.000  3.000 1.000    

%%writefile quadrupole.pqr
ATOM      1 X    MP      1       1.000    1.000    0.000  1.000 2.000
ATOM      2 X    MP      1      -1.000   -1.000    0.000  1.000 2.000
ATOM      3 X    MP      1       1.000   -1.000    0.000 -1.500 2.000
ATOM      4 X    MP      1      -1.000    1.000    0.000 -1.000 2.000

"""
Explanation: Create artificial "molecules" with multipolar moments
Each molecule consists of a number of particles given by $x$, $y$, $z$ coordinates (column 6-8), charge (column 9), and radius (last column). The PQR file format is a variant of PDB and can be read by many molecular viewers.
End of explanation
"""


plt.figure(figsize=plt.figaspect(1/3.))
plt.gray()
plt.figure(1)
for name, pos in {'monopole':131, 'dipole':132, 'quadrupole':133}.iteritems():   
    plt.subplot(pos)
    plt.title(name)
    plt.xticks([])
    plt.yticks([])
    plt.xlim((-3,3))
    plt.ylim((-3,3))
        
    d = np.genfromtxt(name+'.pqr', invalid_raise=False, usecols=(5,6,8), unpack=True)
    
    norm = plt.Normalize()
    colors = plt.cm.jet(norm(d[2]))
    plt.scatter(d[0], d[1], s=10000, c=colors)

"""
Explanation: Visualize multipoles
End of explanation
"""


len=20 # cylinder length
js = {
    "energy" : {
        "cmconstrain" : { "mol1 mol2" : { "mindist": 0, "maxdist": 0.5*len } },
        "nonbonded" : {
            "ljsimple" : { "eps":0.005  },
            "coulomb" : { "epsr" : 78.7, "ionicstrength" : 1e-20 }
        }
    },
    "atomlist" : {
        "X":  { "q":0,  "r":1.0 },
        "MP":  { "q":0,  "r":1.0 }
    },
    "moleculelist": {
        "mol1":  { "structure":"monopole.pqr", "Ninit":1, "insdir":"0 0 1" },
        "mol2":  { "structure":"dipole.pqr", "Ninit":1, "insdir":"0 0 1" }
    },
    "moves" : {
        "moltransrot" : {
            "mol1" : { "dp":20, "dprot":3, "prob":1.0, "permol":True, "dir":"0 0 1"}, 
            "mol2" : { "dp":20, "dprot":3, "prob":1.0, "permol":True, "dir":"0 0 1"}
        } 
    },
    "system" : {
        "temperature" : 298.15,
        "cylinder" : { "length" : len, "radius" : 10 },
        "mcloop"   : { "macro" : 10, "micro" : 100000 }
    }
}
with open('twobody.json', 'w+') as f:
    f.write(json.dumps(js, indent=4))

!rm -f state
!./faunus/src/examples/twobody

"""
Explanation: Run Monte Carlo Simulation (external program)
End of explanation
"""


r, u_exact, u_tot, u_ii, u_id, u_dd, u_iq = np.loadtxt('multipole.dat', unpack=True)

plt.plot(r, u_exact, 'k-', label='exact', lw=8, ms=4.0, alpha=0.3)
plt.plot(r, u_tot, 'k--', label='total', lw=3)
plt.plot(r, u_ii, label='ion-ion', lw=3)
plt.plot(r, u_id, 'ro', label='ion-dip', lw=3, ms=4.0)
plt.plot(r, u_dd, label='dip-dip', lw=3)
plt.plot(r, u_iq, label='ion-quad', lw=3)

plt.legend(loc=0,frameon=False, fontsize=14)
plt.ylabel(r'$u(R)/k_BT$')
plt.xlabel(r'$R$/Å')
plt.xlim((1,8))

"""
Explanation: Plot electrostatic interaction energies
End of explanation
"""


def func_id(r, Z, mu):
    ''' expression for the angularly averaged ion-dipole energy (approximate) '''
    lB=7.0
    return -(lB*Z*mu)**2 / (6*r**4) + (1.6/r)**12

def func_id_exact(r, Z, mu):
    ''' expression for the angularly averaged ion-dipole energy (exact) '''
    lB=7.0
    Emu = lB*Z/r**2 * mu
    return -np.log(np.sinh(Emu)/Emu) + (1.6/r)**12

r, g = np.loadtxt('rdf.dat', unpack=True)
g = g / g[r>7].mean() # w(r)->0 for large r
w = -np.log(g)        # Boltzmann inversion

plt.plot( r, w, 'ro-', label='mc', lw=2 )
plt.plot( r, func_id(r, 1, 0.75), label='perturbation', lw=3 )
plt.plot( r, func_id_exact(r, 1, 0.75), label='exact', lw=3 )

plt.xlim((1,6))
plt.ylim((-3,2))
plt.xlabel(r'$R$/Å')
plt.ylabel(r'$w(R)/k_BT$')

plt.legend(loc=0,frameon=False)

"""
Explanation: Let's compare with angularly averaged ion-dipole interaction
The angularly averaged ion-dipole interaction free energy is,
$$
\beta w(r)_{id} \approx -\frac{(l_BZ\mu)^2}{6r^4}
$$
which we can now plot against the simulated potential of mean force,
$$
\beta w(r)_{mc} = -\ln g(r) + const
$$
Questions:
0. Identify approximations in the multipole approximation
0. what is the difference between $u(r)$ and $w(r)$ ?
0. how could we fix the discrepancy at short separations ? (The analytical expression just becomes more and more attractive)
End of explanation
"""


