import espressomd
espressomd.assert_features(["ELECTROSTATICS", "ROTATION", "ROTATIONAL_INERTIA", "EXTERNAL_FORCES",
                            "MASS", "VIRTUAL_SITES_RELATIVE", "CUDA", "LENNARD_JONES"])
from espressomd import interactions
from espressomd import electrostatics
from espressomd import lb
from espressomd.virtual_sites import VirtualSitesRelative

import numpy as np

# System parameters
#############################################################
box_l = 40.  # size of the simulation box

skin = 0.3  # Skin parameter for the Verlet lists
time_step = 0.01
eq_tstep = 0.001

n_cycle = 1000
integ_steps = 150

# Interaction parameters (Lennard-Jones for raspberry)
#############################################################
radius_col = 3.
harmonic_radius = 3.0

# the subscript c is for colloid and s is for salt (also used for the surface beads)
eps_ss = 1.   # LJ epsilon between the colloid's surface particles.
sig_ss = 1.   # LJ sigma between the colloid's surface particles.
eps_cs = 48.  # LJ epsilon between the colloid's central particle and surface particles.
sig_cs = radius_col  # LJ sigma between the colloid's central particle and surface particles (colloid's radius).
a_eff = 0.32  # effective hydrodynamic radius of a bead due to the discreteness of LB.

# System setup
#############################################################
system = espressomd.System(box_l=[box_l] * 3)
system.time_step = time_step

"""
Explanation: Raspberry Electrophoresis
1 Tutorial Outline
Welcome to the raspberry electrophoresis ESPResSo tutorial! This tutorial assumes some basic knowledge of ESPResSo.
The first step is compiling ESPResSo with the appropriate flags, as listed in Sec. 2.
The tutorial starts by discussing how to build a colloid out of a series of MD beads. These particles typically
resemble a raspberry as can be seen in Fig. 1. After covering the construction of a raspberry colloid, we then
briefly discuss the inclusion of hydrodynamic interactions via a lattice-Boltzmann fluid. Finally we will cover
including ions via the restrictive primitive model (hard sphere ions) and the addition of an electric field
to measure the electrokinetic properties. This script will run a raspberry electrophoresis simulation and write the time and position of the colloid out to a file named <tt>posVsTime.dat</tt> in the same directory.
A sample set of data is included in the file <tt>posVsTime_sample.dat</tt>.
2 Compiling ESPResSo for this Tutorial
To run this tutorial, you will need to enable the following features in the myconfig.hpp file when compiling ESPResSo:
```c++
define ELECTROSTATICS
define ROTATION
define ROTATIONAL_INERTIA
define EXTERNAL_FORCES
define MASS
define VIRTUAL_SITES_RELATIVE
define LENNARD_JONES
```
3 Global MD Variables
The first thing to do in any ESPResSo simulation is to import our espressomd features and set a few global simulation parameters:
End of explanation
"""


system.cell_system.skin = skin

"""
Explanation: The parameter <tt>box_l</tt> sets the size of the simulation box. In general, one should check for finite
size effects which can be surprisingly large in simulations using hydrodynamic interactions. They
also generally scale as <tt>box_l</tt>$^{-1}$ or <tt>box_l</tt>$^{-3}$ depending on the transport mechanism
which sometimes allows for the infinite box limit to be extrapolated to, instead of using an
excessively large simulation box. As a rule of thumb, the box size should be five times greater than the characteristic
length scale of the object. Note that this example uses a small box 
to provide a shorter simulation time.
End of explanation
"""


system.periodicity = [True, True, True]

"""
Explanation: The skin is used for constructing
the Verlet lists and is purely an optimization parameter. Whatever value provides the fastest
integration speed should be used. For the type of simulations covered in this tutorial, this value turns out
to be <tt>skin</tt>$\ \approx 0.3$.
End of explanation
"""


# the LJ potential with the central bead keeps all the beads from simply collapsing into the center
system.non_bonded_inter[1, 0].wca.set_params(epsilon=eps_cs, sigma=sig_cs)
# the LJ potential (WCA potential) between surface beads causes them to be roughly equidistant on the
# colloid surface
system.non_bonded_inter[1, 1].wca.set_params(epsilon=eps_ss, sigma=sig_ss)

# the harmonic potential pulls surface beads towards the central colloid bead
col_center_surface_bond = interactions.HarmonicBond(k=3000., r_0=harmonic_radius)
system.bonded_inter.add(col_center_surface_bond)

"""
Explanation: The <tt>periodicity</tt> parameter indicates that the system is periodic in all three
dimensions. Note that the lattice-Boltzmann algorithm requires periodicity in all three directions (although
this can be modified using boundaries, a topic not covered in this tutorial). 
4 Setting up the Raspberry
Setting up the raspberry is a non-trivial task. The main problem lies in creating a relatively
uniform distribution of beads on the surface of the colloid. In general one should take about 1 bead per lattice-Boltzmann grid
point on the surface to ensure that there are no holes in the surface. The behavior of the colloid can be further improved by placing
beads inside the colloid, though this is not done in this example script. In our example
we first define a harmonic interaction causing the surface beads to be attracted
to the center, and a Lennard-Jones interaction preventing the beads from entering the colloid. There is also a Lennard-Jones
potential between the surface beads to get them to distribute evenly on the surface.
End of explanation
"""


# for the warmup we use a Langevin thermostat with an extremely low temperature and high friction coefficient
# such that the trajectories roughly follow the gradient of the potential while not accelerating too much
system.thermostat.set_langevin(kT=0.00001, gamma=40., seed=42)

print("# Creating raspberry")
center = system.box_l / 2
colPos = center

# Charge of the colloid
q_col = -40
# Number of particles making up the raspberry (surface particles + the central particle).
n_col_part = int(4 * np.pi * np.power(radius_col, 2) + 1)

# Place the central particle
system.part.add(id=0, pos=colPos, type=0, q=q_col, fix=(True, True, True),
                rotation=(1, 1, 1))  # Create central particle

# Create surface beads uniformly distributed over the surface of the central particle
for i in range(1, n_col_part):
    colSurfPos = np.random.randn(3)
    colSurfPos = colSurfPos / np.linalg.norm(colSurfPos) * radius_col + colPos
    system.part.add(id=i, pos=colSurfPos, type=1)
    system.part[i].add_bond((col_center_surface_bond, 0))
print("# Number of colloid beads = {}".format(n_col_part))

# Relax bead positions. The LJ potential with the central bead combined with the
# harmonic bond keep the monomers roughly radius_col away from the central bead. The LJ
# between the surface beads cause them to distribute more or less evenly on the surface.
system.force_cap = 1000
system.time_step = eq_tstep

print("Relaxation of the raspberry surface particles")
for i in range(n_cycle):
    system.integrator.run(integ_steps)

# Restore time step
system.time_step = time_step

"""
Explanation: We set up the central bead and the other beads are initialized at random positions on the surface of the colloid. The beads are then allowed to relax using
an integration loop where the forces between the beads are capped.
End of explanation
"""


# this loop moves the surface beads such that they are once again exactly radius_col away from the center
# For the scalar distance, we use system.distance() which considers periodic boundaries
# and the minimum image convention
colPos = system.part[0].pos
for p in system.part[1:]:
    p.pos = (p.pos - colPos) / np.linalg.norm(system.distance(p, system.part[0])) * radius_col + colPos
    p.pos = (p.pos - colPos) / np.linalg.norm(p.pos - colPos) * radius_col + colPos

"""
Explanation: The best way to ensure a relatively uniform distribution
of the beads on the surface is to simply take a look at a VMD snapshot of the system after this integration. Such a snapshot is shown in Fig. 1.
<figure>
    <img src='figures/raspberry_snapshot.png' alt='missing' style="width: 600px;"/>
    <center>
    <figcaption>Figure 1: A snapshot of the simulation consisting of positive salt ions (yellow spheres), negative salt ions (grey spheres) and surface beads (blue spheres). There is also a central bead in the middle of the colloid bearing a large negative  charge.</figcaption>
    </center>
</figure>

In order to make the colloid perfectly round, we now adjust the bead's positions to be exactly <tt>radius_col</tt> away
from the central bead.
End of explanation
"""


# Select the desired implementation for virtual sites
system.virtual_sites = VirtualSitesRelative()
# Setting min_global_cut is necessary when there is no interaction defined with a range larger than
# the colloid such that the virtual particles are able to communicate their forces to the real particle
# at the center of the colloid
system.min_global_cut = radius_col

# Calculate the center of mass position (com) and the moment of inertia (momI) of the colloid
com = np.average(system.part[1:].pos, 0)  # system.part[:].pos returns an n-by-3 array
momI = 0
for i in range(n_col_part):
    momI += np.power(np.linalg.norm(com - system.part[i].pos), 2)

# note that the real particle must be at the center of mass of the colloid because of the integrator
print("\n# moving central particle from {} to {}".format(system.part[0].pos, com))
system.part[0].fix = [False, False, False]
system.part[0].pos = com
system.part[0].mass = n_col_part
system.part[0].rinertia = np.ones(3) * momI

# Convert the surface particles to virtual sites related to the central particle
# The id of the central particles is 0, the ids of the surface particles start at 1.
for p in system.part[1:]:
    p.vs_auto_relate_to(0)

"""
Explanation: Now that the beads are arranged in the shape of a raspberry, the surface beads are made virtual particles
using the VirtualSitesRelative scheme. This converts the raspberry to a rigid body
in which the surface particles follow the translation and rotation of the central particle.
Newton's equations of motion are only integrated for the central particle.
It is given an appropriate mass and moment of inertia tensor (note that the inertia tensor
is given in the frame in which it is diagonal.)
End of explanation
"""


print("# Adding the positive ions")
salt_rho = 0.001  # Number density of ions
volume = system.volume()
N_counter_ions = int(round((volume * salt_rho) + abs(q_col)))

i = 0
while i < N_counter_ions:
    pos = np.random.random(3) * system.box_l
    # make sure the ion is placed outside of the colloid
    if (np.power(np.linalg.norm(pos - center), 2) > np.power(radius_col, 2) + 1):
        system.part.add(pos=pos, type=2, q=1)
        i += 1

print("# Added {} positive ions".format(N_counter_ions))

print("\n# Adding the negative ions")

N_co_ions = N_counter_ions - abs(q_col)
i = 0
while i < N_co_ions:
    pos = np.random.random(3) * system.box_l
    # make sure the ion is placed outside of the colloid
    if (np.power(np.linalg.norm(pos - center), 2) > np.power(radius_col, 2) + 1):
        system.part.add(pos=pos, type=3, q=-1)
        i += 1

print("# Added {} negative ions".format(N_co_ions))

"""
Explanation: 5 Inserting Counterions and Salt Ions
Next we insert enough ions at random positions (outside the radius of the colloid) with opposite charge to the colloid such that the system is electro-neutral. In addition, ions
of both signs are added to represent the salt in the solution.
End of explanation
"""


# Check charge neutrality
assert np.abs(np.sum(system.part[:].q)) < 1E-10

"""
Explanation: We then check that charge neutrality is maintained
End of explanation
"""


# WCA interactions for the ions, essentially giving them a finite volume
system.non_bonded_inter[0, 2].lennard_jones.set_params(
    epsilon=eps_ss, sigma=sig_ss,
    cutoff=sig_ss * pow(2., 1. / 6.), shift="auto", offset=sig_cs - 1 + a_eff)
system.non_bonded_inter[0, 3].lennard_jones.set_params(
    epsilon=eps_ss, sigma=sig_ss,
    cutoff=sig_ss * pow(2., 1. / 6.), shift="auto", offset=sig_cs - 1 + a_eff)
system.non_bonded_inter[2, 2].wca.set_params(epsilon=eps_ss, sigma=sig_ss)
system.non_bonded_inter[2, 3].wca.set_params(epsilon=eps_ss, sigma=sig_ss)
system.non_bonded_inter[3, 3].wca.set_params(epsilon=eps_ss, sigma=sig_ss)

"""
Explanation: A WCA potential acts between all of the ions. This potential represents a purely repulsive
version of the Lennard-Jones potential, which approximates hard spheres of diameter $\sigma$. The ions also interact through a WCA potential
with the central bead of the colloid, using an offset of around $\mathrm{radius_col}-\sigma +a_\mathrm{grid}/2$. This makes
the colloid appear as a hard sphere of radius roughly $\mathrm{radius_col}+a_\mathrm{grid}/2$ to the ions, which is approximately equal to the
hydrodynamic radius of the colloid
End of explanation
"""


print("\n# Equilibrating the ions (without electrostatics):")
# Langevin thermostat for warmup before turning on the LB.
temperature = 1.0
system.thermostat.set_langevin(kT=temperature, gamma=1.)

print("Removing overlap between ions")
ljcap = 100
CapSteps = 100
for i in range(CapSteps):
    system.force_cap = ljcap
    system.integrator.run(integ_steps)
    ljcap += 5

system.force_cap = 0

"""
Explanation: After inserting the ions, again a short integration is performed with a force cap to
prevent strong overlaps between the ions.
End of explanation
"""


# Turning on the electrostatics
# Note: Production runs would typically use a target accuracy of 10^-4
print("\n# Tuning P3M parameters...")
bjerrum = 2.
p3m = electrostatics.P3M(prefactor=bjerrum * temperature, accuracy=0.001)
system.actors.add(p3m)
print("# Tuning complete")

"""
Explanation: 6 Electrostatics
Electrostatics are simulated using the Particle-Particle Particle-Mesh (P3M) algorithm. In ESPResSo this can be added to the simulation rather trivially:
End of explanation
"""


E = 0.1  # an electric field of 0.1 is the upper limit of the linear response regime for this model
Efield = np.array([E, 0, 0])
for p in system.part:
    p.ext_force = p.q * Efield

"""
Explanation: Generally a Bjerrum length of $2$ is appropriate when using WCA interactions with $\sigma=1$, since a typical ion has a radius of $0.35\ \mathrm{nm}$, while the Bjerrum
length in water is around $0.7\ \mathrm{nm}$.
The external electric field is simulated by simply adding a constant force equal to the simulated field times the particle charge. Generally the electric field is set to $0.1$ in MD units,
which is the maximum field before the response becomes nonlinear. Smaller fields are also possible, but the required simulation time is considerably larger. Sometimes, Green-Kubo methods
are also used, but these are generally only feasible in cases where there is either no salt or a very low salt concentration.
End of explanation
"""


system.part[:].v = (0, 0, 0)

"""
Explanation: 7 Lattice-Boltzmann
Before creating the LB fluid it is a good idea to set all of the particle velocities to zero.
This is necessary to set the total momentum of the system to zero. Failing to do so will lead to an unphysical drift of the system, which
will change the values of the measured velocities.
End of explanation
"""


lb = espressomd.lb.LBFluidGPU(kT=temperature, seed=42, dens=1., visc=3., agrid=1., tau=system.time_step)
system.actors.add(lb)

"""
Explanation: The important parameters for the LB fluid are the density, the viscosity, the time step,
and the friction coefficient used to couple the particle motion to the fluid.
The time step should generally be comparable to the MD time step. While
large time steps are possible, a time step of $0.01$ turns out to provide more reasonable values for the root mean squared particle velocities. Both density and viscosity
should be around $1$, while the friction should be set around $20.$ The grid spacing should be comparable to the ions' size.
End of explanation
"""


system.thermostat.turn_off()
system.thermostat.set_lb(LB_fluid=lb, seed=123, gamma=20.0)

"""
Explanation: A logical way of picking a specific set of parameters is to choose them such that the hydrodynamic radius of an ion roughly matches its physical radius determined by the
WCA potential ($R=0.5\sigma$). Using the following equation:
\begin{equation}
\frac{1}{\Gamma}=\frac{1}{6\pi \eta R_{\mathrm{H0}}}=\frac{1}{\Gamma_0} 
+\frac{1}{g\eta a} 
 \label{effectiveGammaEq}
\end{equation}
one can see that the set of parameters grid spacing $a=1\sigma$, fluid density $\rho=1$, a 
kinematic viscosity of $\nu=3 $ and a friction of $\Gamma_0=50$ leads to a hydrodynamic radius
of approximately $0.5\sigma$.
The last step is to first turn off all other thermostats, followed by turning on the LB thermostat. The temperature is typically set to 1, which is equivalent to setting
$k_\mathrm{B}T=1$ in molecular dynamics units.
End of explanation
"""


# Reset the simulation clock
system.time = 0
initial_pos = system.part[0].pos
num_iterations = 1000
num_steps_per_iteration = 1000
with open('posVsTime.dat', 'w') as f:  # file where the raspberry trajectory will be written to
    for i in range(num_iterations):
        system.integrator.run(num_steps_per_iteration)
        pos = system.part[0].pos - initial_pos
        f.write("%.2f %.4f %.4f %.4f\n" % (system.time, pos[0], pos[1], pos[2]))
        print("# time: {:.0f} ({:.0f}%), col_pos: {}".format(
            system.time, (i + 1) * 100. / num_iterations, np.around(pos, 1), end='\r'))

print("\n# Finished")

"""
Explanation: 8 Simulating Electrophoresis
Now the main simulation can begin! The only important thing is to make sure the system has enough time to equilibrate. There are two separate equilibration times: 1) the time for the ion distribution to stabilize, and 2) the time
needed for the fluid flow profile to equilibrate. In general, the ion distribution equilibrates fast, so the needed warmup time is largely determined by the fluid relaxation time, which can be calculated via $\tau_\mathrm{relax} = \mathrm{box_length}^2/\nu$. This means for a box of size 40 with a kinematic viscosity of 3 as in our example script, the relaxation time is $\tau_\mathrm{relax} = 40^2/3 = 533 \tau_\mathrm{MD}$, or 53300 integration steps. In general it is a good idea to run for many relaxation times before starting to use the simulation results for averaging observables. To be on the safe side $10^6$ integration steps is a reasonable equilibration time. Please feel free to modify the provided script and try and get some interesting results!
End of explanation
"""


import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
%matplotlib notebook

trajectory_file = 'posVsTime.dat'
trajectory = np.loadtxt(trajectory_file)[:, 1:4]
# optional: trajectory smoothing with a running average
N = 6
trajectory = np.array(
    [np.convolve(trajectory[:, i], np.ones((N,)) / N, mode='valid') for i in range(3)])
# calculate bounding box (cubic box to preserve scaling)
trajectory_range = np.max(trajectory, axis=1) - np.min(trajectory, axis=1)
mid_range = np.median(trajectory, axis=1)
max_range = 1.01 * np.max(np.abs(trajectory_range))
bbox = np.array([mid_range - max_range / 2, mid_range + max_range / 2])
# 3D plot
fig = plt.figure(figsize=(9, 6))
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')
ax.set_zlabel('Z axis')
ax.set_xlim(*bbox[:, 0])
ax.set_ylim(*bbox[:, 1])
ax.set_zlim(*bbox[:, 2])
ax.text(*trajectory[:, 0], '\u2190 start', 'y')
ax.scatter(*trajectory[:, 0])
ax.plot(*trajectory)
plt.tight_layout()
plt.rcParams.update({'font.size': 14})

"""
Explanation: Plot the raspberry trajectory with <tt>matplotlib</tt>:
End of explanation
"""


