"""Simulate a mixture of magnetic and non-magnetic particles."""

import espressomd

required_features = ["DP3M", "WCA"]
espressomd.assert_features(required_features)

import espressomd.magnetostatics
import numpy as np


###########################################################################
# System parameters
###########################################################################

box_l = 10.
density = 0.7
wca_sig = 1.0

volume = box_l**3
N = int(volume * density)


###########################################################################
# System setup
###########################################################################

system = espressomd.System(box_l=[box_l] * 3)
np.random.seed(seed=42)

system.time_step = 0.01
system.cell_system.skin = 0.4
system.non_bonded_inter[0, 0].wca.set_params(epsilon=1.0, sigma=wca_sig)

# Random positions and dipoles
pos = box_l * np.random.random((N, 3))
dip_phi = 2. * np.pi * np.random.random((N, 1))
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])
# optional: make half the particles non-magnetic
dip[::2] = [0., 0., 0.]

system.part.add(pos=pos, rotation=N * [(True, True, True)],
                dip=dip, fix=N * [(False, False, False)])

def print_energy(step, key='total'):
    print(f"{step}: E_{key}={system.analysis.energy()[key]:+.2e}")


###########################################################################
# Energy minimization
###########################################################################

# convergence criterion (particles are separated by at least 90% sigma)
min_dist = 0.9 * wca_sig

# steepest descent
system.integrator.set_steepest_descent(f_max=0., gamma=1e-3,
                                       max_displacement=0.01)
i = 0
print_energy('minimization', key='total')
while i < 20 and system.analysis.min_dist() < min_dist:
    system.integrator.run(20)
    i += 1
    print_energy('minimization', key='total')


###########################################################################
# Thermalization
###########################################################################

# activate thermostat
system.thermostat.set_langevin(kT=1.0, gamma=1.0, seed=42)
system.integrator.set_vv()

# activate P3M actor
dp3m = espressomd.magnetostatics.DipolarP3M(prefactor=1.0, accuracy=1e-2)
system.actors.add(dp3m)

print_energy('thermalization', key='total')
for i in range(6):
    system.integrator.run(20)
    print_energy('thermalization', key='total')


###########################################################################
# Integration
###########################################################################

print_energy('integration', key='dipolar')
for i in range(6):
    system.integrator.run(200)
    print_energy('integration', key='dipolar')


###########################################################################
# Debugging
###########################################################################

print('debugging: turning off LJ potential')
system.non_bonded_inter[0, 0].wca.set_params(epsilon=0.0, sigma=wca_sig)
print('debugging: turning off thermostat')
system.thermostat.turn_off()
print('debugging: keeping magnetostatics')
system.integrator.run(0, recalc_forces=True)
print_energy('debugging', key='dipolar')
print_energy('debugging', key='kinetic')

with np.printoptions(edgeitems=6):
    print('magnetic dipoles:')
    print(dip)
    print('particle forces:')
    print(system.part[:].f)