#!/usr/bin/env amspython

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

# Load molecule and run initial geometry optimization
config.default_jobrunner = JobRunner(parallel=True)
settings = Settings()
settings.runscript.nproc = 1

struct_dir = os.path.join(os.getenv("AMSHOME"), "examples", "Hybrid", "PeriodicCapping", "systems")

mol = ChemicalSystem.from_xyz(os.path.join(struct_dir, "var5.xyz"))
mol.guess_bonds()

driver_geom_opt = drivers.AMS()
driver_geom_opt.Task = "GeometryOptimization"
driver_geom_opt.Engine = engines.DFTB()
driver_geom_opt.GeometryOptimization.OptimizeLattice = True
driver_geom_opt.GeometryOptimization.Method = "FIRE"
driver_geom_opt.GeometryOptimization.MaxIterations = 300

job_dftb = AMSJob(molecule=mol, settings=settings + Settings({"input": driver_geom_opt}), name="reference")
results_dftb = job_dftb.run()

# Get charges and add them to molecule
mol.enable_atom_attributes("forcefield")
charges = results_dftb.get_charges()
for atom, charge in zip(mol.atoms, charges):
    atom.forcefield.charge = charge

# Run cheap calculation using forcefield engine
driver_geom_opt.Engine = engines.ForceField()
driver_geom_opt.Engine.NonBondedCutoff = Units.convert("bohr", "angstrom", 50)
job_ff = AMSJob(molecule=mol, settings=settings + Settings({"input": driver_geom_opt}), name="cheap")
results_ff = job_ff.run()
mol = results_ff.get_main_system()

print("begin report")
print(
    """We optimize the lattice and test several distances.
We divide the system in such a way that there are two equivalent, and hence neutral regions.
Here are the distances (Angstrom) as obtained with a QM and an MM method:
"""
)

header = "{:<10} {:>10} {:>10} {:>10}".format("Distance", "QM", "MM", "Err(MM)")
print(header)
print("-" * len(header))

mol_qm = results_dftb.get_main_system()
mol_mm = results_ff.get_main_system()
for bond, i, j in [("B-H", 0, 1), ("N-H", 2, 3), ("B-N", 0, 2)]:
    dist_qm = mol_qm.get_distance(i, j)
    dist_mm = mol_mm.get_distance(i, j)
    print(f"{bond:<10} {dist_qm:>10.3f} {dist_mm:>10.3f} {dist_mm - dist_qm:>10.3f}")

# Run hybrid engine
results = {}
for embedding in ["mechanical", "electrostatic"]:
    for capping in ["fixed", "fractional"]:
        driver = copy.deepcopy(driver_geom_opt)
        driver.Properties.Gradients = True
        driver.GeometryOptimization.MaxIterations = 100

        driver.Engine = engines.Hybrid()
        driver.Engine.Capping.AllowHighBondOrders = True
        driver.Engine.Capping.Option = capping
        driver.Engine.QMMM.QMRegion = "qm"
        driver.Engine.QMMM.QMEngineID = "DFTB"
        driver.Engine.QMMM.MMEngineID = "ForceField"
        driver.Engine.QMMM.Embedding = embedding

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

        job = AMSJob(
            molecule=mol,
            settings=settings + Settings({"input": driver}),
            name=f"system.embedding={embedding}.capping={capping}.go",
        )

        results[(embedding, capping)] = job.run()

hybrid_results = {k: (v.get_energy(engine="hybrid"), v.get_main_system()) for k, v in results.items()}

print(
    """Of course the force field results do not exactly match the QM results, the error displayed in the last column.
Now we try the hybrid engine, can we improve the bonds in the QM region?
We start from the geometry calculated with the (cheap) forcefield.
In this table we show the errors in bond lengths (in the QM region) of the hybrid method with respect to the QM method:
"""
)

header = "{:<16} {:>16} {:>12} {:>8} {:>8} {:>8}".format("Embedding", "Capping", "Energy", "B-H", "N-H", "B-N")
print(header)
print("-" * len(header))

for (embedding, capping), (energy, mol_hybrid) in hybrid_results.items():
    errs = [mol_hybrid.get_distance(i, j) - mol_qm.get_distance(i, j) for i, j in [(0, 1), (2, 3), (0, 2)]]
    print(f"{embedding:<16} {capping:>16} {energy:>12.6f} {errs[0]:>8.3f} {errs[1]:>8.3f} {errs[2]:>8.3f}")

print(
    """Here are some observations
     * the B-H distance is a bit worse than with a plain forcefield, especially with fractional capping
     * the N-H distance is much better than with the plain forcefield 
     * the B-N distance is a bit better than with the plain forcefield, now too short. Fractional capping works best.
     * Electrostatic embedding is doing slightly better than mechanical embedding, the biggest improvement is on the B-N bond
"""
)
print("end report")