#!/usr/bin/env amspython

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

# Here we treat H3O+ as qm and OH- as the MM region (Optimizing without regions gives two H2O molecules)
# We do this with a QUILD setup (mechanical embedding) and electrostatic embedding (QMMM)
# We obtain the charges from a DFTB calculation
# In this case the results (QUILD vs. QMMM) are quite similar as apparently the OH does not polarize the QM region much

# First we do two DFTB calculations on the two fragments
config.default_jobrunner = JobRunner(parallel=True)
settings = Settings()
settings.runscript.nproc = 1

h30 = ChemicalSystem(
    """
System
    Atoms
        O -1.527946410885647 -0.2107366711137158 -0.0008116899510243671
        H -0.8459142126057956 0.3517312394359257 0.4094504676540848
        H -1.834953147575289 0.1051014241823828 -0.8704652381864062
        H -1.328032016244278 -1.164422847242489 0.02894848344144469
    End
    Charge 1.0
    GuessBonds True
End
"""
)

oh = ChemicalSystem(
    """
System
    Atoms
        O 0.6370858511871781 -0.3378071707560572 -0.0006181020627287671
        H 1.318474396634582 0.2241299231185073 0.4092568796869673
    End
    Charge -1.0
    GuessBonds True
End
"""
)

fragments = {"H3O+": h30, "OH-": oh}

results_dftb = {}
for name, mol in fragments.items():
    driver_dftb = drivers.AMS()
    driver_dftb.Task = "SinglePoint"
    driver_dftb.Properties.Charges = True
    driver_dftb.GeometryOptimization.Convergence.Gradients = 1.0e-6
    driver_dftb.Engine = engines.DFTB()
    job = AMSJob(molecule=mol, settings=settings + Settings({"input": driver_dftb}), name=name)
    results_dftb[name] = job.run()

# Now we run it in a QUILD-like setup (mechanical embedding) and in a QMMM-like setup
h20_2 = ChemicalSystem(
    """
System
    Atoms
        O -1.527946410885647 -0.2107366711137158 -0.0008116899510243671 region=QM
        H -0.8459142126057956 0.3517312394359257 0.4094504676540848     region=QM
        H -1.834953147575289 0.1051014241823828 -0.8704652381864062     region=QM
        H -1.328032016244278 -1.164422847242489 0.02894848344144469     region=QM
        O 0.6370858511871781 -0.3378071707560572 -0.0006181020627287671 region=MM
        H 1.318474396634582 0.2241299231185073 0.4092568796869673       region=MM
    End
    GuessBonds True
End
"""
)

qm_charges = results_dftb["H3O+"].get_charges()
mm_charges = results_dftb["OH-"].get_charges()
h20_2.enable_atomic_properties("forcefield")

for i, charge in zip(h20_2.get_atoms_in_region("QM"), qm_charges):
    h20_2.atoms[i].forcefield.charge = charge

for i, charge in zip(h20_2.get_atoms_in_region("MM"), mm_charges):
    h20_2.atoms[i].forcefield.charge = charge

results = {}
for name in ["quild", "qmmm"]:
    driver = copy.deepcopy(driver_dftb)
    driver.Task = "GeometryOptimization"
    driver.Engine = engines.Hybrid()

    driver.Engine.Engine[0] = engines.DFTB()
    driver.Engine.Engine[0].Model = "GFN1-xTB"
    driver.Engine.Engine[1] = engines.ForceField()

    if name == "quild":
        driver.Engine.Energy.Term[0].Factor = 1.0
        driver.Engine.Energy.Term[0].Region = "*"
        driver.Engine.Energy.Term[0].EngineID = "ForceField"

        driver.Engine.Energy.Term[1].Factor = -1.0
        driver.Engine.Energy.Term[1].Region = "QM"
        driver.Engine.Energy.Term[1].EngineID = "ForceField"
        driver.Engine.Energy.Term[1].Charge = 1.0

        driver.Engine.Energy.Term[2].Factor = 1.0
        driver.Engine.Energy.Term[2].Region = "QM"
        driver.Engine.Energy.Term[2].EngineID = "DFTB"
        driver.Engine.Energy.Term[2].Charge = 1.0

    elif name == "qmmm":
        driver.Engine.QMMM.QMRegion = "QM"
        driver.Engine.QMMM.QMEngineID = "DFTB"
        driver.Engine.QMMM.MMEngineID = "ForceField"
        driver.Engine.QMMM.QMCharge = 1.0
        driver.Engine.QMMM.MMCharge = -1.0

    job = AMSJob(molecule=h20_2, settings=settings + Settings({"input": driver}), name=name)
    results[name] = job.run()

results = {k: (v.get_charges(), v.get_main_system()) for k, v in results.items()}

print("start of report")
print("method distance charges")
for name, (charges, system) in results.items():
    print(
        f"{name} {system.get_distance(0, 4):5.3f} {charges[0]:5.3f} {charges[1]:5.3f} {charges[2]:5.3f} {charges[3]:5.3f} {charges[4]:5.3f} {charges[5]:5.3f}"
    )
print("end of report")