# -*- coding: utf-8 -*-
def setAndUpdateMDCalculator(configuration):
    # Calculator
    potentialSet = ReaxFF_LiS_2015(strict_bondpairs = None)
    calculator = TremoloXCalculator(parameters=potentialSet)
    calculator.setVerletListsDelta(0.25*Angstrom)

    # Set calculator and run ATK-ForceField update
    bulk_configuration.setCalculator(calculator)
    bulk_configuration.update()

    return bulk_configuration

# -------------------------------------------------------------
# Set LiS composition
# -------------------------------------------------------------
composition = 0.40

#######################################
# DO NOT CHANGE BELOW THIS LINE
# UNLESS YOU KNOW WHAT YOU ARE DOING
#######################################

# -------------------------------------------------------------
# HDF5 data file
# -------------------------------------------------------------
datafile = 'x%.2f.hdf5' % composition

# -------------------------------------------------------------
# Initial configuration
# -------------------------------------------------------------
bulk_configuration = nlread(datafile, BulkConfiguration)[0]
bulk_configuration = setAndUpdateMDCalculator(bulk_configuration)

# -------------------------------------------------------------
# Relax the initial structure
# -------------------------------------------------------------
bulk_configuration = OptimizeGeometry(
    bulk_configuration,
    max_forces=0.5*eV/Ang,
    max_stress=0.1*GPa,
    max_steps=1000,
    max_step_length=0.2*Ang,
    trajectory_filename=datafile,
    optimizer_method=LBFGS(),
    constrain_bravais_lattice=False,
)
nlsave(datafile, bulk_configuration)

# -------------------------------------------------------------
# MD equilibration at 1600 K for 50 ps using NPT Berendsen.
# -------------------------------------------------------------
initial_velocity = MaxwellBoltzmannDistribution(
    temperature=1600.0*Kelvin,
    remove_center_of_mass_momentum=True
)

method = NPTBerendsen(
    time_step=1*femtoSecond,
    reservoir_temperature=1600*Kelvin,
    reservoir_pressure=1*bar,
    thermostat_timescale=100*femtoSecond,
    barostat_timescale=500*femtoSecond,
    initial_velocity=initial_velocity,
    compressibility=0.0001*bar**-1,
    heating_rate=0*Kelvin/picoSecond,
)

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=[],
    trajectory_filename=None,
    steps=50000,
    log_interval=1000,
    method=method
)

bulk_configuration = md_trajectory.lastImage()
nlsave(datafile, md_trajectory)

# -------------------------------------------------------------
# MD cool-down from 1600 K to 300 K in 1 ns using NPT Berendsen.
# -------------------------------------------------------------
initial_velocity = MaxwellBoltzmannDistribution(
    temperature=1600.0*Kelvin,
    remove_center_of_mass_momentum=True
)

method = NPTBerendsen(
    time_step=1*femtoSecond,
    reservoir_temperature=1600*Kelvin,
    reservoir_pressure=1*bar,
    thermostat_timescale=100*femtoSecond,
    barostat_timescale=500*femtoSecond,
    initial_velocity=initial_velocity,
    compressibility=0.0001*bar**-1,
    heating_rate=-1.3*Kelvin/picoSecond,
)

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=[],
    trajectory_filename=None,
    steps=1000000,
    log_interval=1000,
    method=method
)

bulk_configuration = md_trajectory.lastImage()
nlsave(datafile, md_trajectory)

# -------------------------------------------------------------
# MD equilibration at 300 K for 100 ps using NPT Martyna-Tobias-Klein
# -------------------------------------------------------------
initial_velocity = MaxwellBoltzmannDistribution(
    temperature=300.0*Kelvin,
    remove_center_of_mass_momentum=True
)

method = NPTMartynaTobiasKlein(
    time_step=1*femtoSecond,
    reservoir_temperature=300*Kelvin,
    reservoir_pressure=1*bar,
    thermostat_timescale=100*femtoSecond,
    barostat_timescale=500*femtoSecond,
    initial_velocity=initial_velocity,
    heating_rate=0*Kelvin/picoSecond,
)

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=[],
    trajectory_filename=None,
    steps=100000,
    log_interval=100,
    method=method
)

bulk_configuration = md_trajectory.lastImage()
nlsave(datafile, md_trajectory)

# -------------------------------------------------------------
# Minimize atomic forces in the 10 last geometries obtained
# in the 300 K MD equilibration, and compute and save the energies.
# -------------------------------------------------------------
images = range(-1, -11, -1)
for i in images:
    bulk_configuration = md_trajectory.image(i)
    bulk_configuration = setAndUpdateMDCalculator(bulk_configuration)

    # Optimize Geometry
    bulk_configuration = OptimizeGeometry(
        bulk_configuration,
        max_forces=0.01*eV/Ang,
        max_steps=1000,
        max_step_length=0.2*Ang,
        trajectory_filename=None,
        disable_stress=True,
        optimizer_method=LBFGS(),
    )
    nlsave(datafile, bulk_configuration)

    # Total Energy
    total_energy = TotalEnergy(bulk_configuration)
    nlsave(datafile, total_energy)
    nlprint(total_energy)