#!/usr/bin/env amspython

# In this example we use the hybrid engine in a molecular gun MD application, shooting HF molecules at a BN surface

# The BN slab represents the MM region and the "bullets" are the QM region
# The regions are defined in the xyz files using end of line strings (atom attributes)

from scm.libbase import ChemicalSystem, Units
from scm.input_classes import drivers, engines
from scm.plams import AMSJob, JobRunner, config, Settings
import os

# Load chemical systems from xyz files
struct_dir = os.path.join(os.getenv("AMSHOME"), "examples", "Hybrid", "HybridGun", "molecules")
bn_slab = ChemicalSystem.from_xyz(os.path.join(struct_dir, "BNSlab.xyz"))
hf = ChemicalSystem.from_xyz(os.path.join(struct_dir, "HF.xyz"))
systems = {"bn_slab": bn_slab, "hf": hf}

# First we do two dftb calculations to get a guess of the charges to be used by the force field.
# Setup jobs and run in parallel
config.default_jobrunner = JobRunner(parallel=True)
settings = Settings()
settings.runscript.nproc = 1

driver_dftb = drivers.AMS()
driver_dftb.Task = "SinglePoint"
driver_dftb.Engine = engines.DFTB()

results_dftb = {}
for name, system in systems.items():
    job_dftb = AMSJob(molecule=system, settings=settings + Settings({"input": driver_dftb}), name=name)
    results_dftb[name] = job_dftb.run()

# Load forcefield charges from the dftb calculations onto the chemical systems
for name, system in systems.items():
    results = results_dftb[name]
    charges = results.get_charges()
    system.enable_atom_attributes("forcefield")
    for atom, charge in zip(system.atoms, charges):
        atom.forcefield.charge = charge
    system.guess_bonds()
    print(f"Total Energy (hartree) {results.get_energy('hartree', 'dftb'):.8f}")
    print(f"Total Energy (eV) {results.get_energy('eV', 'dftb'):.8f}")

# Now we can run our MD simulation using both mechanical and electrostatic embedding
# (because electrostatic embedding is more expensive we limit here the number of steps)
results_md = {}
for embedding, steps in [("mechanical", 1400), ("electrostatic", 300)]:

    driver_md = drivers.AMS()
    driver_md.Task = "MolecularDynamics"
    driver_md.RNGSeed = [-1341016088, 83513668, 1764626453, -87803069, -1149690266, 1963370818, -1393571175, 1985130742]
    driver_md.Constraints.Atom = 1

    driver_md.MolecularDynamics.NSteps = steps
    driver_md.MolecularDynamics.Trajectory.SamplingFreq = 20
    driver_md.MolecularDynamics.InitialVelocities.Temperature = 300
    driver_md.MolecularDynamics.Preserve.Momentum = "No"
    driver_md.MolecularDynamics.Preserve.AngularMomentum = "No"
    driver_md.MolecularDynamics.AddMolecules.System = "HF"
    driver_md.MolecularDynamics.AddMolecules.Frequency = 159
    driver_md.MolecularDynamics.AddMolecules.CoordsBox = [0, 3, 0, 8.57, 6, 7]
    driver_md.MolecularDynamics.AddMolecules.VelocityDirection = [0.45752820, 0, -0.5540656]
    driver_md.MolecularDynamics.AddMolecules.Velocity = 0.07
    driver_md.MolecularDynamics.AddMolecules.Rotate = "Yes"
    driver_md.MolecularDynamics.AddMolecules.MinDistance = 3.0
    driver_md.MolecularDynamics.RemoveMolecules.Formula = "*"
    driver_md.MolecularDynamics.RemoveMolecules.Frequency = 101
    driver_md.MolecularDynamics.RemoveMolecules.SinkBox.FractionalCoordsBox = [0, 1, 0, 1, 8, 1000]

    driver_md.Engine = engines.Hybrid()
    driver_md.Engine.QMMM.QMRegion = "QM"
    driver_md.Engine.QMMM.QMEngineID = "DFTB"
    driver_md.Engine.QMMM.MMEngineID = "ForceField"
    driver_md.Engine.QMMM.Embedding = embedding

    driver_md.Engine.Engine[0] = engines.DFTB()
    driver_md.Engine.Engine[1] = engines.ForceField()
    driver_md.Engine.Engine[1].NonBondedCutoff = Units.convert("bohr", "angstrom", 50)

    job_md = AMSJob(
        molecule={"": systems["bn_slab"], "HF": systems["hf"]},
        settings=settings + Settings({"input": driver_md}),
        name=f"SinkBox.embedding={embedding}"
    )
    results_md[embedding] = job_md.run()

for result in results_md.values():
    res = result.read_rkf_section("MDResults")
    header = "{:<20} {:>16} {:>16} {:>16} {:>16}".format("Measure", "Min", "Max", "Mean", "StdDev")
    print(header)
    print("-" * len(header))
    print(f"{'Potential Energy':<20} {res['MinPotentialEnergy']:>16.8f} {res['MaxPotentialEnergy']:>16.8f} {res['MeanPotentialEnergy']:>16.8f} {res['StdDevPotentialEnergy']:>16.8f}")
    print(f"{'Kinetic Energy':<20} {res['MinKineticEnergy']:>16.8f} {res['MaxKineticEnergy']:>16.8f} {res['MeanKineticEnergy']:>16.8f} {res['StdDevKineticEnergy']:>16.8f}")
    print(f"{'Total Energy':<20} {res['MinTotalEnergy']:>16.8f} {res['MaxTotalEnergy']:>16.8f} {res['MeanTotalEnergy']:>16.8f} {res['StdDevTotalEnergy']:>16.8f}")
    print(f"{'Temperature':<20} {res['MinTemperature']:>16.8f} {res['MaxTemperature']:>16.8f} {res['MeanTemperature']:>16.8f} {res['StdDevTemperature']:>16.8f}")