from espressomd import System, electrostatics
import espressomd
import numpy
import matplotlib.pyplot as plt
plt.ion()

# Print enabled features
required_features = ["EXTERNAL_FORCES", "MASS", "ELECTROSTATICS", "LENNARD_JONES"]
espressomd.assert_features(required_features)
print(espressomd.features())

# System Parameters
n_part = 200
n_ionpairs = n_part / 2
density = 0.5
time_step = 0.01
temp = 1.0
gamma = 1.0
l_bjerrum = 7.0

num_steps_equilibration = 1000
num_configs = 500
integ_steps_per_config = 1000

# Particle Parameters
types       = {"Anion":          0, "Cation": 1}
numbers     = {"Anion": n_ionpairs, "Cation": n_ionpairs}
charges     = {"Anion":       -1.0, "Cation": 1.0}
lj_sigmas   = {"Anion":        1.0, "Cation": 1.0}
lj_epsilons = {"Anion":        1.0, "Cation": 1.0}

WCA_cut = 2.**(1. / 6.)
lj_cuts     = {"Anion":  WCA_cut * lj_sigmas["Anion"],
               "Cation": WCA_cut * lj_sigmas["Cation"]}

"""
Explanation: Tutorial 2: A Simple Charged System, Part 1
1 Introduction
This tutorial introduces some of the basic features of ESPResSo for charged systems by constructing a simulation script for a simple salt crystal. In the subsequent task, we use a more realistic force-field for a NaCl crystal. Finally, we introduce constraints and 2D-Electrostatics to simulate a molten salt in a parallel plate capacitor. We assume that the reader is familiar with the basic concepts of Python and MD simulations.  Compile ESPResSo with the following features in your <tt>myconfig.hpp</tt> to be set throughout the whole tutorial:
```c++
define EXTERNAL_FORCES
define MASS
define ELECTROSTATICS
define LENNARD_JONES
```
2 Basic Set Up
The script for the tutorial can be found in your build directory at <tt>/doc/tutorials/02-charged_system/scripts/nacl.py</tt>.
We start by importing numpy, pyplot, espressomd and setting up the simulation parameters:
End of explanation
"""


# Setup System
box_l = (n_part / density)**(1. / 3.)
system = System(box_l=[box_l, box_l, box_l])
system.seed = 42
system.periodicity = [True, True, True]
system.time_step = time_step
system.cell_system.skin = 0.3

"""
Explanation: These variables do not change anything in the simulation engine, but
are just standard Python variables. They are used to increase the
readability and flexibility of the script. The box length is not a
parameter of this simulation, it is calculated from the number of
particles and the system density. This allows to change the parameters
later easily, e.g. to simulate a bigger system.
We use dictionaries for all particle related parameters, which is less error-prone and
readable as we will see later when we actually need the values. The parameters here define a purely repulsive, 
equally sized, monovalent salt.
The simulation engine itself is modified by changing the
<tt>espressomd.System()</tt> properties. We create an instance <tt>system</tt> and
set the box length, periodicity and time step. The skin depth <tt>skin</tt> 
is a parameter for the link--cell system which tunes its
performance, but shall not be discussed here.
End of explanation
"""


for i in range(int(n_ionpairs)):
    system.part.add(id=len(system.part), type=types["Anion"],
                    pos=numpy.random.random(3) * box_l, q=charges["Anion"])
for i in range(int(n_ionpairs)):
    system.part.add(id=len(system.part), type=types["Cation"],
                    pos=numpy.random.random(3) * box_l, q=charges["Cation"])

"""
Explanation: We now fill this simulation box with particles at random positions, using type and charge from our dictionaries.
Using the length of the particle list <tt>system.part</tt> for the id, we make sure that our particles are numbered consecutively.
The particle type is used to link non-bonded interactions to a certain group of particles.
End of explanation
"""


def combination_rule_epsilon(rule, eps1, eps2):
    if rule == "Lorentz":
        return (eps1 * eps2)**0.5
    else:
        return ValueError("No combination rule defined")

def combination_rule_sigma(rule, sig1, sig2):
    if rule == "Berthelot":
        return (sig1 + sig2) * 0.5
    else:
        return ValueError("No combination rule defined")

# Lennard-Jones interactions parameters
for s in [["Anion", "Cation"], ["Anion", "Anion"], ["Cation", "Cation"]]:
    lj_sig = combination_rule_sigma("Berthelot", lj_sigmas[s[0]], lj_sigmas[s[1]])
    lj_cut = combination_rule_sigma("Berthelot", lj_cuts[s[0]], lj_cuts[s[1]])
    lj_eps = combination_rule_epsilon("Lorentz", lj_epsilons[s[0]], lj_epsilons[s[1]])

    system.non_bonded_inter[types[s[0]], types[s[1]]].lennard_jones.set_params(
        epsilon=lj_eps, sigma=lj_sig, cutoff=lj_cut, shift="auto")

"""
Explanation: Before we can really start the simulation, we have to specify the
interactions between our particles. We already defined the Lennard-Jones parameters at the beginning,
what is left is to specify the combination rule and to iterate over particle type pairs. For simplicity, 
we implement only the Lorentz-Berthelot rules. 
We pass our interaction pair to <tt>system.non_bonded_inter[*,*]</tt> and set the 
pre-calculated LJ parameters <tt>epsilon</tt>, <tt>sigma</tt> and <tt>cutoff</tt>. With <tt>shift="auto"</tt>,
we shift the interaction potential to the cutoff so that $U_\mathrm{LJ}(r_\mathrm{cutoff})=0$.
End of explanation
"""


# Lennard-Jones Equilibration
max_sigma = max(lj_sigmas.values())
min_dist = 0.0
system.minimize_energy.init(f_max=0, gamma=10.0, max_steps=10,
                            max_displacement=max_sigma * 0.01)

while min_dist < max_sigma:
    min_dist = system.analysis.min_dist([types["Anion"], types["Cation"]],
                                        [types["Anion"], types["Cation"]])
    system.minimize_energy.minimize()

# Set thermostat
system.thermostat.set_langevin(kT=temp, gamma=gamma, seed=42)

"""
Explanation: 3 Equilibration
With randomly positioned particles, we most likely have huge overlap and the strong repulsion will
cause the simulation to crash. The next step in our script therefore is a suitable LJ equilibration.
This is known to be a tricky part of a simulation and several approaches exist to reduce the particle overlap.
Here, we use the steepest descent algorithm and cap the maximal particle displacement per integration step
to 1% of sigma.
We use <tt>system.analysis.min_dist()</tt> to get the minimal distance between all particles pairs. This value
is used to stop the minimization when particles are far away enough from each other. At the end, we activate the
Langevin thermostat for our NVT ensemble with temperature <tt>temp</tt> and friction coefficient <tt>gamma</tt>.
End of explanation
"""


p3m = electrostatics.P3M(prefactor=l_bjerrum * temp,
                         accuracy=1e-3)
system.actors.add(p3m)

"""
Explanation: ESPResSo uses so-called <tt>actors</tt> for electrostatics, magnetostatics and hydrodynamics. This ensures that unphysical combinations of algorithms are
avoided, for example simultaneous usage of two electrostatic interactions.
Adding an actor to the system also activates the method and calls necessary
initialization routines. Here, we define a P$^3$M object with parameters Bjerrum
length and rms force error. This automatically starts a
tuning function which tries to find optimal parameters for P$^3$M and prints them
to the screen:
End of explanation
"""


# Temperature Equilibration
system.time = 0.0
for i in range(int(num_steps_equilibration / 50)):
    energy = system.analysis.energy()
    temp_measured = energy['kinetic'] / ((3.0 / 2.0) * n_part)
    print("t={0:.1f}, E_total={1:.2f}, E_coulomb={2:.2f},T={3:.4f}"
          .format(system.time, energy['total'], energy['coulomb'], temp_measured), end='\r')
    system.integrator.run(200)

print()

"""
Explanation: Before the production part of the simulation, we do a quick temperature 
equilibration. For the output, we gather all energies with <tt>system.analysis.energy()</tt>, calculate the "current" temperature from the ideal part and print it to the screen along with the total and Coulomb energies. Note that for the ideal gas the temperature is given via $1/2 m \sqrt{\langle v^2 \rangle}=3/2 k_BT$, where $\langle \cdot \rangle$ denotes the ensemble average. Calculating some kind of "current temperature" via $T_\text{cur}=\frac{m}{3 k_B} \sqrt{ v^2 }$ won't produce the temperature in the system. Only when averaging the squared velocities first one would obtain the temperature for the ideal gas. $T$ is a fixed quantity and does not fluctuate in the canonical ensemble.
We integrate for a certain amount of steps with <tt>system.integrator.run(100)</tt>.
End of explanation
"""


# Integration
system.time = 0.0
for i in range(num_configs):
    energy = system.analysis.energy()
    temp_measured = energy['kinetic'] / ((3.0 / 2.0) * n_part)
    print("progress={:.0f}%, t={:.1f}, E_total={:.2f}, E_coulomb={:.2f}, T={:.4f}"
          .format((i + 1) * 100. / num_configs, system.time, energy['total'],
                  energy['coulomb'], temp_measured), end='\r')
    system.integrator.run(integ_steps_per_config)

    # Internally append particle configuration
    system.analysis.append()

print()

"""
Explanation: <figure>
<img src='figures/salt.png' alt='missing' style="width: 300px;"/>
<center>
<figcaption>Figure 1: VMD Snapshot of the Salt System</figcaption>
</center>
</figure>

4 Running the Simulation
Now we can integrate the particle trajectories for a couple of time
steps. Our integration loop basically looks like the equilibration:
End of explanation
"""


# Analysis
# Calculate the averaged rdfs
rdf_bins = 100
r_min = 0.0
r_max = system.box_l[0] / 2.0
r, rdf_00 = system.analysis.rdf(rdf_type='<rdf>',
                                type_list_a=[types["Anion"]],
                                type_list_b=[types["Anion"]],
                                r_min=r_min,
                                r_max=r_max,
                                r_bins=rdf_bins)
r, rdf_01 = system.analysis.rdf(rdf_type='<rdf>',
                                type_list_a=[types["Anion"]],
                                type_list_b=[types["Cation"]],
                                r_min=r_min,
                                r_max=r_max,
                                r_bins=rdf_bins)

"""
Explanation: Additionally, we append all particle configurations in the core with <tt>system.analysis.append()</tt> for a very convenient analysis later on.
5 Analysis
Now, we want to calculate the averaged radial distribution functions
$g_{++}(r)$ and $g_{+-}(r)$ with the <tt>rdf()</tt> command from <tt>system.analysis</tt>:
End of explanation
"""


with open('rdf.data', 'w') as rdf_fp:
    for i in range(rdf_bins):
        rdf_fp.write("%1.5e %1.5e %1.5e\n" % (r[i], rdf_00[i], rdf_01[i]))

"""
Explanation: The shown <tt>rdf()</tt> commands return the radial distribution functions for
equally and oppositely charged particles for specified radii and number of bins. 
In this case, we calculate the averaged rdf of the stored
configurations, denoted by the chevrons in <tt>rdf_type='$<\mathrm{rdf}>$'</tt>. Using <tt>rdf_type='rdf'</tt> would simply calculate the rdf of the current particle
configuration. The results are two NumPy arrays containing the $r$ and $g(r)$
values. We can then write the data into a file with standard python output routines.
End of explanation
"""


# Plot the distribution functions
plt.figure(figsize=(10, 6), dpi=80)
plt.plot(r[:], rdf_00[:], label='Na$-$Na')
plt.plot(r[:], rdf_01[:], label='Na$-$Cl')
plt.xlabel('$r$', fontsize=20)
plt.ylabel('$g(r)$', fontsize=20)
plt.legend(fontsize=20)
plt.show()

"""
Explanation: Finally we can plot the two radial distribution functions using pyplot.
End of explanation
"""


