import espressomd
espressomd.assert_features('DIPOLES', 'LENNARD_JONES')

from espressomd.magnetostatics import DipolarP3M
from espressomd.magnetostatic_extensions import DLC

import numpy as np

"""
Explanation: Ferrofluid - Part II
Table of Contents

Applying an external magnetic field
Magnetization curve

Remark: The equilibration and sampling times used in this tutorial would be not sufficient for scientific  purposes, but they are long enough to get at least a qualitative insight of the behaviour of ferrofluids. They have been shortened so we achieve reasonable computation times for the purpose of a tutorial.
Applying an external magnetic field
In this part we want to investigate the influence of a homogeneous external magnetic field exposed to a ferrofluid system.
We import all necessary packages and check for the required ESPResSo features
End of explanation
"""


# Lennard-Jones parameters
lj_sigma = 1.
lj_epsilon = 1.
lj_cut = 2**(1./6.) * lj_sigma

# Particles
N = 700

# Area fraction of the mono-layer 
phi = 0.06

# Dipolar interaction parameter lambda = mu_0 m^2 /(4 pi sigma^3 kT)
dip_lambda = 4.

# Temperature
kT = 1.0

# Friction coefficient
gamma = 1.0

# Time step
dt = 0.01

# Langevin parameter alpha = mu_0 m H / kT
alpha = 10.

# vacuum permeability
mu_0 = 1.

"""
Explanation: and set up the simulation parameters where we introduce a new dimensionless parameter 
\begin{equation}
    \alpha = \frac{\mu B}{k_BT} = \frac{\mu \mu_0 H}{k_BT}
\end{equation}
which is called Langevin parameter. We intentionally choose a relatively high volume fraction $\phi$ and dipolar interaction parameter $\lambda$ to clearly see the influence of the dipole-dipole interaction
End of explanation
"""


# System setup
box_size = (N * np.pi * (lj_sigma/2.)**2. /phi)**0.5

print("Box size",box_size)
# Note that the dipolar P3M and dipolar layer correction need a cubic
# simulation box for technical reasons.
system = espressomd.System(box_l=(box_size,box_size,box_size)) 
system.time_step = dt
system.thermostat.set_langevin(kT=kT,gamma=gamma, seed=1)

# Lennard-Jones interaction
system.non_bonded_inter[0,0].lennard_jones.set_params(epsilon=lj_epsilon,sigma=lj_sigma,cutoff=lj_cut, shift="auto")

# Random dipole moments
np.random.seed(seed = 1)
dip_phi = np.random.random((N,1)) *2. * np.pi
dip_cos_theta = 2*np.random.random((N,1)) -1
dip_sin_theta = np.sin(np.arccos(dip_cos_theta))
dip = np.hstack((
   dip_sin_theta *np.sin(dip_phi),
   dip_sin_theta *np.cos(dip_phi),
   dip_cos_theta))

# Random positions in the monolayer
pos = box_size* np.hstack((np.random.random((N,2)), np.zeros((N,1))))

# Add particles
system.part.add(pos=pos,rotation=N*[(1,1,1)],dip=dip,fix=N*[(0,0,1)])

# Remove overlap between particles by means of the steepest descent method
system.integrator.set_steepest_descent(
    f_max=0,gamma=0.1,max_displacement=0.05)

while system.analysis.energy()["total"] > 5*kT*N:
    system.integrator.run(20)

# Switch to velocity Verlet integrator
system.integrator.set_vv()

# tune verlet list skin
system.cell_system.tune_skin(min_skin=0.4, max_skin=2., tol=0.2, int_steps=100)

# Setup dipolar P3M and dipolar layer correction (DLC)
dp3m = DipolarP3M(accuracy=5E-4,prefactor=dip_lambda*lj_sigma**3*kT)
dlc = DLC(maxPWerror=1E-4, gap_size=box_size-lj_sigma)
system.actors.add(dp3m)
system.actors.add(dlc)

# tune verlet list skin again
system.cell_system.tune_skin(min_skin=0.4, max_skin=2., tol=0.2, int_steps=100)

# print skin value
print('tuned skin = {}'.format(system.cell_system.skin))

"""
Explanation: Now we set up the system, where we, as we did in part I, only commit the orientation of the dipole moment to the particles and take the magnitude into account in the prefactor of Dipolar P3M (for more details see part I). 
Hint:
It should be noted that we seed both the Langevin thermostat and the random number generator of numpy. The latter means that the initial configuration of our system is the same every time this script will be executed. As the time evolution of the system depends not solely on the Langevin thermostat but also on the numeric accuracy and DP3M as well as DLC (the tuned parameters are slightly different every time) it is only partly predefined. You can change the seeds to simulate with a different initial configuration and a guaranteed different time evolution.
End of explanation
"""


# magnetic field times dipole moment
H_dipm = (alpha*kT)
H_field = [H_dipm,0,0]
print("Set magnetic field constraint...")
H_constraint = espressomd.constraints.HomogeneousMagneticField(H=H_field)
system.constraints.add(H_constraint)

"""
Explanation: We now apply the external magnetic field which is
\begin{equation}
    B = \mu_0 H = \frac{\alpha~k_BT}{\mu}
\end{equation}
As only the current orientation of the dipole moments, i.e. the unit vector of the dipole moments, is saved in the particle list but not their strength we have to commit $B\cdot \mu$ as the external magnetic field to ESPResSo.
We applying the field in x-direction using the class <tt>constraints</tt> of ESPResSo
End of explanation
"""


# Equilibrate
print("Equilibration...")
equil_rounds = 10
equil_steps = 200
for i in range(equil_rounds):
  system.integrator.run(equil_steps)
  print("progress: {:3.0f}%, dipolar energy: {:9.2f}".format(
        (i + 1) * 100. / equil_rounds, system.analysis.energy()["dipolar"]), end="\r")
print("\nEquilibration done")

"""
Explanation: and equilibrate the system for a while
End of explanation
"""


import matplotlib.pyplot as plt

plt.figure(figsize=(10,10))
plt.xlim(0, box_size)
plt.ylim(0, box_size)
plt.xlabel('x-position', fontsize=20)
plt.ylabel('y-position', fontsize=20)
plt.plot(system.part[:].pos_folded[:,0], system.part[:].pos_folded[:,1], 'o')
plt.show()

"""
Explanation: Now we can visualize the current state and see that the particles mostly create chains oriented in the direction of the external magnetic field. Also some monomers should be present.
End of explanation
"""


import matplotlib.pyplot as plt
import matplotlib.animation as animation
from tempfile import NamedTemporaryFile
import base64

VIDEO_TAG = """<video controls>
 <source src="data:video/x-m4v;base64,{0}" type="video/mp4">
 Your browser does not support the video tag.
</video>"""

def anim_to_html(anim):
    if not hasattr(anim, '_encoded_video'):
        with NamedTemporaryFile(suffix='.mp4') as f:
            anim.save(f.name, fps=20, extra_args=['-vcodec', 'libx264'])
            with open(f.name, "rb") as g:
                video = g.read()
        anim._encoded_video = base64.b64encode(video).decode('ascii')
        plt.close(anim._fig)
    return VIDEO_TAG.format(anim._encoded_video)

animation.Animation._repr_html_ = anim_to_html

def init():
    # Set x and y range
    ax.set_ylim(0, box_size)
    ax.set_xlim(0, box_size)
    xdata, ydata = [], []
    part.set_data(xdata, ydata)
    return part,

def run(i):
    system.integrator.run(50)

    # Save current system state as a plot
    xdata, ydata = system.part[:].pos_folded[:,0], system.part[:].pos_folded[:,1]
    ax.figure.canvas.draw()
    part.set_data(xdata, ydata)
    print("progress: {:3.0f}%".format(i + 1), end="\r")
    return part,

"""
Explanation: Video of the development of the system
You may want to get an insight of how the system develops in time. Thus we now create a function which will save a video and embed it in an html string to create a video of the systems development
End of explanation
"""


fig, ax = plt.subplots(figsize=(10,10))
part, = ax.plot([],[], 'o')

animation.FuncAnimation(fig, run, frames=100, blit=True, interval=0, repeat=False, init_func=init)

"""
Explanation: We now can start the sampling over the <tt>animation</tt> class of <tt>matplotlib</tt>
End of explanation
"""


# Dipolar interaction parameter lambda = mu_0 m^2 /(4 pi sigma^3 kT)
dip_lambda = 1.

# increase time step
dt = 0.02

# dipole moment
dipm = np.sqrt(dip_lambda*4*np.pi*lj_sigma**3.*kT / mu_0)

# remove all particles
system.part[:].remove()


# Random dipole moments
dip_phi=np.random.random((N,1)) *2. * np.pi
dip_cos_theta = 2*np.random.random((N,1)) -1
dip_sin_theta = np.sin(np.arccos(dip_cos_theta))
dip = np.hstack((
   dip_sin_theta *np.sin(dip_phi),
   dip_sin_theta *np.cos(dip_phi),
   dip_cos_theta))

# Random positions in the monolayer
pos = box_size* np.hstack((np.random.random((N,2)), np.zeros((N,1))))

# Add particles
system.part.add(pos=pos,rotation=N*[(1,1,1)],dip=dip,fix=N*[(0,0,1)])

# Remove overlap between particles by means of the steepest descent method
system.integrator.set_steepest_descent(f_max=0,gamma=0.1,max_displacement=0.05)

while system.analysis.energy()["total"] > 5*kT*N:
    system.integrator.run(20)

# Switch to velocity Verlet integrator
system.integrator.set_vv()

# tune verlet list skin
system.cell_system.tune_skin(min_skin=0.4, max_skin=2., tol=0.2, int_steps=100)

# Setup dipolar P3M and dipolar layer correction
system.actors.remove(dp3m)
system.actors.remove(dlc)

dp3m = DipolarP3M(accuracy=5E-4,prefactor=dip_lambda*lj_sigma**3*kT)
dlc = DLC(maxPWerror=1E-4, gap_size=box_size-lj_sigma)

system.actors.add(dp3m)
system.actors.add(dlc)

# tune verlet list skin again
system.cell_system.tune_skin(min_skin=0.4, max_skin=2., tol=0.2, int_steps=100)

"""
Explanation: In the visualization video we can see that the single chains break and connect to each other during time. Also some monomers are present which break from and connect to chains. If you want to have some more frames, i.e. a longer video, just adjust the <tt>frames</tt> parameter in <tt>FuncAnimation</tt>.
Magnetization curve
An important observable of a ferrofluid system is the magnetization $M$ of the system in direction of an external magnetic field $H$
\begin{equation}
    M = \frac{\sum_i \mu_i^H}{V}
\end{equation}
where the index $H$ means the component of $\mu_i$ in direction of the external magnetic field $H$ and the sum runs over all particles.
The magnetization plotted over the external field $H$ is called magnetization curve. For particles with non-interacting dipole moments there is an analytical solution
\begin{equation}
    M = M_{sat}\cdot L(\alpha)
\end{equation}
with $L(\alpha)$ the Langevin function
\begin{equation}
    L(\alpha) = \coth(\alpha)-\frac{1}{\alpha}
\end{equation}
and $\alpha$ the Langevin parameter
\begin{equation}
    \alpha=\frac{\mu_0\mu}{k_BT}H = \frac{\mu}{k_BT}B
\end{equation}
$M_{sat}$ is the so called saturation magnetization which is the magnetization of a system where all dipole moments are aligned to each other. Thus it is the maximum of the magnetization. In our case all dipole moments are equal, thus
\begin{equation}
    M_{sat} = \frac{N\cdot\mu}{V}
\end{equation}
For better comparability we now introduce a dimensionless magnetization
\begin{equation}
    M^* = \frac{M}{M_{sat}} = \frac{\sum_i \mu_i^H}{N\cdot \mu}
\end{equation}
Thus the analytical solution for non-interacting dipole moments $M^*$ is simply the Langevin function.
For interacting dipole moments there are only approximations for the magnetization curve available. 
Here we want to use the approximation of Ref. <a href='#[1]'>[1]</a> for a quasi two dimensional system, which reads with adjusted coefficients (Ref. <a href='#[1]'>[1]</a> used a different dipole-dipole interaction prefactor $\gamma = 1$)
\begin{equation}
    M_{||}^{q2D} = M_{sat} L(\alpha) \left( 1 + \mu_0\frac{1}{8} M_{sat} \frac{d L(\alpha)}{dB} \right)
\end{equation}
and
\begin{equation}
    M_{\perp}^{q2D} = M_{sat} L(\alpha) \left( 1 - \mu_0\frac{1}{4} M_{sat} \frac{d L(\alpha)}{dB} \right)
\end{equation}
for the magnetization with an external magnetic field parallel and perpendicular to the monolayer plane, respectively. Here the dipole-dipole interaction is approximated as a small perturbation and
\begin{equation}
    \frac{d L(\alpha)}{dB} = \left( \frac{1}{\alpha^2} - \frac{1}{\sinh^2(\alpha)} \right) \cdot \frac{\mu}{k_BT}
\end{equation}
By comparing the magnetization curve parallel $M_{||}^{q2D}$ and perpendicular $M_{\perp}^{q2D}$ to the monolayer plane we can see that the magnetization is increased in the case of an external field parallel to the monolayer plane and decreased in the case of an external field perpendicular to the monolayer plane. The latter can be explained by the fact that an orientation of all single dipole moments perpendicular to the monolayer plane results in a configuration with a repulsive dipole-dipole interaction as the particles have no freedom of movement in the direction perpendicular to the monolayer plane. This counteracts the magnetization perpendicular to the monolayer plane.
We now want to use ESPResSo to get an estimation of how the magnetization curve is affected by the dipole-dipole interaction parallel and perpendicular to the monolayer plane and compare the results with the Langevin curve and the magnetization curves of Ref. <a href='#[1]'>[1]</a>.
For the sampling of the magnetization curve we set up a new system, where we decrease the dipolar interaction parameter $\lambda$ drastically. We do this as we want to compare our results with the approximation of Ref. <a href='#[1]'>[1]</a> which is only valid for small dipole-dipole interaction between the particles (decreasing the volume fraction $\phi$ would also be an appropriate choice). For smaller dipolar interaction parameters it is possible to increase the time step. We do this to get more uncorrelated measurements.
End of explanation
"""


from espressomd.observables import MagneticDipoleMoment
dipm_tot = MagneticDipoleMoment(ids=system.part[:].id)

"""
Explanation: To increase the performance we use the built-in function <tt>MagneticDipoleMoment</tt> to calculate the dipole  moment of the whole system. In our case this is only the orientation as we never committed the strength of the dipole moments to our particles.
End of explanation
"""


alphas = [0, 0.25, 0.5, 1, 2, 3, 4, 8]

"""
Explanation: We use the dimensionless Langevin parameter $\alpha$ as the parameter for the external magnetic field. As the interesting part of the magnetization curve is the one for small external magnetic field strengths - for large external magnetic fields the magnetization goes into saturation in all cases - we increase the spacing between the Langevin parameters $\alpha$ up to higher values and write them into a list
End of explanation
"""


# sampling with magnetic field perpendicular to monolayer plane (in z-direction)

# remove all constraints
system.constraints.clear()

# list of magnetization in field direction
magnetization_perp = []

# number of loops for sampling
loops = 500

for alpha in alphas:
    print("Sampling for alpha = {}".format(alpha))
    H_dipm = (alpha*kT)
    H_field = [0,0,H_dipm]
    print("Set magnetic field constraint...")
    H_constraint = espressomd.constraints.HomogeneousMagneticField(H=H_field)
    system.constraints.add(H_constraint)
    print("done\n")
    
    # Equilibration
    print("Equilibration...")
    for i in range(equil_rounds):
        system.integrator.run(equil_steps)
        print("progress: {:3.0f}%, dipolar energy: {:9.2f}".format(
              (i+1)*100./equil_rounds, system.analysis.energy()["dipolar"]), end="\r")
    print("\nEquilibration done\n")
    
    # Sampling
    print("Sampling...")
    magn_temp = 0
    for i in range(loops):
        system.integrator.run(20)
        magn_temp += dipm_tot.calculate()[2]
        print("progress: {:3.0f}%".format((i+1)*100./loops), end="\r")
    print("\n")
    
    # save average magnetization
    magnetization_perp.append(magn_temp / loops)
    print("Sampling for alpha = {} done\n".format(alpha))
    print("magnetizations = {}".format(magnetization_perp))
    print("total progress: {:5.1f}%\n".format((alphas.index(alpha)+1)*100./len(alphas)))
    
    # remove constraint
    system.constraints.clear()
print("Magnetization curve sampling done")

"""
Explanation: For both the magnetization perpendicular and parallel to the monolayer plane we use the same system for every value of the Langevin parameter $\alpha$. Thus we use that the system is already more or less equilibrated from the previous run so we save some equilibration time. For scientific purposes one would use a new system for every value for the Langevin parameter to ensure that the systems are independent and no correlation effects are measured. Also one would perform more than just one simulation for each value of $\alpha$ to increase the precision of the results.
Now we sample the magnetization for increasing $\alpha$'s or increasing magnetic field in direction perpendicular to the monolayer plane
End of explanation
"""


# sampling with magnetic field parallel to monolayer plane (in x-direction)

# remove all constraints
system.constraints.clear()

# list of magnetization in field direction
magnetization_para = []

# number of loops for sampling
loops = 500

for alpha in alphas:
    print("Sample for alpha = {}".format(alpha))
    H_dipm = (alpha*kT)
    H_field = [H_dipm,0,0]
    print("Set magnetic field constraint...")
    H_constraint = espressomd.constraints.HomogeneousMagneticField(H=H_field)
    system.constraints.add(H_constraint)
    print("done\n")
    
    # Equilibration
    print("Equilibration...")
    for i in range(equil_rounds):
        system.integrator.run(equil_steps)
        print("progress: {:3.0f}%, dipolar energy: {:9.2f}".format(
              (i+1)*100./equil_rounds, system.analysis.energy()["dipolar"]), end="\r")
    print("\nEquilibration done\n")
    
    # Sampling
    print("Sampling...")
    magn_temp = 0
    for i in range(loops):
        system.integrator.run(20)
        magn_temp += dipm_tot.calculate()[0]
        print("progress: {:3.0f}%".format((i+1)*100./loops), end="\r")
    print("\n")
    
    # save average magnetization
    magnetization_para.append(magn_temp / loops)
    print("Sampling for alpha = {} done\n".format(alpha))
    print("magnetizations = {}".format(magnetization_para))
    print("total progress: {:5.1f}%\n".format((alphas.index(alpha)+1)*100./len(alphas)))
    
    # remove constraint
    system.constraints.clear()
print("Magnetization curve sampling done")

"""
Explanation: and now we sample the magnetization for increasing $\alpha$'s or increasing magnetic field in direction parallel to the monolayer plane
End of explanation
"""


import matplotlib.pyplot as plt

"""
Explanation: Now we can compare the resulting magnetization curves with the Langevin curve and the more advanced ones of Ref. <a href='#[1]'>[1]</a> by plotting all of them in one figure. Thus first we import matplotlib if not already done
End of explanation
"""


# dipole moment
dipm = np.sqrt(dip_lambda*4*np.pi*lj_sigma**3.*kT / mu_0)
print('dipole moment = {}'.format(dipm))

"""
Explanation: For the approximations of $ M_{||}^{q2D}$ and $ M_{\perp}^{q2D}$ of Ref. <a href='#[1]'>[1]</a> we need the dipole moment of a single particle. Thus we calculate it from our dipolar interaction parameter $\lambda$
End of explanation
"""


M_sat = phi * 4./np.pi * 1./(lj_sigma**2.) * dipm

"""
Explanation: and the saturation magnetization by using
\begin{equation}
    M_{sat} = \rho \mu = \phi \frac{4}{\pi \sigma^2} \mu
\end{equation}
thus
End of explanation
"""


def dL_dB(alpha):
    return (1./(alpha**2.) - 1./((np.sinh(alpha))**2.)) * dipm / (kT)

"""
Explanation: Further we need the derivation of the Langevin function after the external field B thus we define the function
End of explanation
"""


# approximated magnetization curve for a field parallel to the monolayer plane
def magnetization_approx_para(phi, alpha):
    return L(alpha) * ( 1. + mu_0/8. * M_sat * dL_dB(alpha) )

# approximated magnetization curve for a field perpendicular to the monolayer plane
def magnetization_approx_perp(phi, alpha):
    return L(alpha) * ( 1. - mu_0/4. * M_sat * dL_dB(alpha) )

"""
Explanation: Now we define the approximated magnetization curves parallel and perpendicular to the monolayer plane
End of explanation
"""


# Langevin function
def L(x):
    return (np.cosh(x)/np.sinh(x)) - 1./x

"""
Explanation: Now we define the Langevin function
End of explanation
"""


# list of the values for alpha (x-axis)
x = np.arange(0.01,9, 0.1, dtype=float).tolist()


L_func = []
L_perp = []
L_para = []
for i in x:
    L_func.append(L(i))
    L_perp.append(magnetization_approx_perp(phi, i))
    L_para.append(magnetization_approx_para(phi, i))
    

# divide all entries in the magnetization list by N to get the dimensionless magnetization
magnetization_perp_star = []
magnetization_para_star = []
for i in range(len(magnetization_para)):
    magnetization_perp_star.append(magnetization_perp[i] / N)
    magnetization_para_star.append(magnetization_para[i] / N)



plt.figure(figsize=(10,10))
plt.xlabel(r'$\alpha$', fontsize=20)
plt.ylabel(r'$M^*$', fontsize=20)
plt.plot(x, L_func, label='Langevin function')
plt.plot(x, L_perp, label='q2D approximation $\perp$')
plt.plot(x, L_para, label='q2D approximation ||')
plt.plot(alphas, magnetization_perp_star, 'o', label='simulation results $\perp$')
plt.plot(alphas, magnetization_para_star, 'o', label='simulation results ||')
plt.legend(fontsize=20)
plt.show()

"""
Explanation: and plot the three theoretical curves together with our simulation results
End of explanation
"""


