#!/usr/bin/env amspython
"""
This script runs several MD simulations.

The final production simulation ("prod") has the BinLog option enabled for the
PressureTensor. This means that the pressure tensor components are written to
ams.rkf for every time step. This is needed to be able to calculate the
viscosity accurately with the Green-Kubo relation.

You can run this script either with
* APPLE&P for an ionic liquid (engine = 'applenp')
* GAFF for benzene (engine = 'gaff')

Change the temperature, timestep, and number of MD steps for each individual simulation to your liking.

"""
import os

from scm.plams import *

engine = "applenp"
# engine = 'gaff'

temperature = 333
timestep = 0.5

tutorial_mdsteps = {
    "scan_density": 10000,
    "NVT_before_NPT": 10000,
    "NPT_eq": 150000,
    "NVT_eq": 20000,
    "prod": 4000000,
}

recommended_mdsteps = {
    "scan_density": 20000,
    "NVT_before_NPT": 10000,
    "NPT_eq": 1000000,  # or even longer!
    "NVT_eq": 10000,
    "prod": 10000000,  # or even longer!
}


def main():
    # change the below to recommended_mdsteps for a more realistic simulation!
    mdsteps = tutorial_mdsteps

    if engine == "applenp":
        forcefield_file = os.path.abspath("PYR14_TFSI.ff")
        s = applenp_engine_settings(forcefield_file)
        box = applenp_initial_system(forcefield_file)  # ionic liquid
        scan_density_upper = 1.3  # reasonable high density for the ionic liquid
    else:
        s = gaff_engine_settings()
        box = gaff_initial_system()  # benzene
        scan_density_upper = 1.0  # reasonable high density for benzene

    init()

    job = AMSJob(settings=s + preoptimization_settings(), name="preoptimization", molecule=box)
    job.run(watch=True)
    if not check_job(job):
        return

    mol = job.results.get_main_molecule()

    job = AMSMDScanDensityJob(
        scan_density_upper=scan_density_upper,
        molecule=mol,
        settings=s,
        name="scan_density",
        nsteps=mdsteps["scan_density"],
        timestep=timestep,
        temperature=temperature,
        writevelocities=False,
        writemolecules=False,
        writebonds=True,
    )
    job.run(watch=True)
    if not check_job(job):
        return

    job = AMSNVTJob(
        molecule=job.results.get_lowest_energy_molecule(),
        settings=s,
        name="NVT_before_NPT",
        nsteps=mdsteps["NVT_before_NPT"],
        timestep=timestep,
        temperature=temperature,
        writevelocities=False,
        writemolecules=False,
        writebonds=False,
    )
    job.run(watch=True)
    if not check_job(job):
        return

    job = AMSNPTJob.restart_from(
        job,
        name="NPT_eq",
        nsteps=mdsteps["NPT_eq"],
        timestep=timestep,
        samplingfreq=500,
        thermostat="NHC",
        temperature=temperature,
        tau=100,
        barostat="MTK",
        barostat_tau=1000,
        equal="XYZ",
        writebonds=True,
        writevelocities=False,
        writemolecules=False,
    )
    job.run(watch=True)
    if not check_job(job):
        return

    mol = job.results.get_equilibrated_molecule()
    job = AMSNVTJob(
        molecule=mol,
        settings=s,
        name="NVT_eq",
        nsteps=mdsteps["NVT_eq"],
        timestep=timestep,  # femtoseconds
        thermostat="NHC",
        temperature=temperature,
        tau=100,
    )
    job.run(watch=True)
    if not check_job(job):
        return

    job = AMSNVTJob(
        molecule=job.results.get_main_molecule(),
        settings=s,
        name="prod",
        temperature=temperature,
        thermostat="NHC",
        tau=100,
        nsteps=mdsteps["prod"],
        samplingfreq=2000,
        timestep=timestep,
        calcpressure=True,
        binlog_time=True,
        binlog_pressuretensor=True,
        writevelocities=False,
        writemolecules=False,
        writebonds=False,
    )
    job.run(watch=True)
    if not check_job(job):
        return

    finish()


def gaff_engine_settings():
    s = Settings()
    s.input.ForceField.Type = "GAFF"
    s.input.ForceField.AnteChamberIntegration = "Yes"
    return s


def gaff_initial_system():
    box = packmol(from_smiles("c1ccccc1"), n_atoms=600, density=0.6)
    return box


def applenp_engine_settings(forcefield_file):
    s = Settings()
    s.input.ForceField.Type = "APPLE&P"
    s.input.ForceField.ForceFieldFile = forcefield_file
    return s


def applenp_initial_system(forcefield_file):
    from scm.appleandp import appleandp_packmol

    PYR14 = from_smiles("CCCC[N+]1(CCCC1)C")
    TFSI = from_smiles("C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F")
    box = appleandp_packmol(
        molecules=[PYR14, TFSI],
        n_molecules=[30, 30],
        density=0.7,
        region_names=["cation", "anion"],
        forcefield_file=forcefield_file,
    )

    return box


def preoptimization_settings():
    geo_opt_s = Settings()
    geo_opt_s.input.ams.task = "GeometryOptimization"
    geo_opt_s.input.ams.GeometryOptimization.MaxIterations = 10
    geo_opt_s.input.ams.GeometryOptimization.PretendConverged = "Yes"
    return geo_opt_s


def check_job(job):
    if job.ok():
        return True

    log(f"ERRORS for job {job.results.rkfpath()}")
    for line in job.results.get_file_chunk("$JN.err"):
        log(line)

    return False


if __name__ == "__main__":
    main()