#!/usr/bin/env amspython
# coding: utf-8

# ## Initial imports

import scm.plams as plams
from scm.params import ResultsImporter, ParAMSJob
from scm.plams import Settings, AMSJob, log, Molecule
from scm.simple_active_learning import SimpleActiveLearningJob
from pathlib import Path
import matplotlib.pyplot as plt

# common_ru_h.py must exist in the current working directory
from common_ru_h import (
    rotation,
    dft_settings,
    QEKPointsConfig,
    slice_slab,
    check_installation,
)


# ## Initialize PLAMS working directory

load_model_dir = "initial_training_results"
check_installation(load_model_dir)
plams.init()


# ## Load the optimized bulk Ru structure from the job collection
#
# The lattice was optimized in the previous notebook, and the structure was stored in the job collection.
#
# Let's retrieve it from the job collection and use it to construct Ru surface slabs.

job_collection = ParAMSJob.load_external(load_model_dir).results.get_job_collection()
optimized_bulk = job_collection["hcp_lattopt_Ru_dft"].molecule
slab = slice_slab(optimized_bulk, miller=(1, 0, 0), thickness=7.0, cell_z=15, ref_atom=0)
min_z = min(at.coords[2] for at in slab)
slab.translate((0, 0, -min_z + 2.0))
slab = slab.supercell(3, 2, 1)
plams.plot_molecule(slab, rotation=rotation)
plt.title("Ru(10-10)")


min_z = min(at.coords[2] for at in slab)
for i, at in enumerate(slab, 1):
    at.properties = Settings()  # remove details about supercell generation
    if at.coords[2] == min_z:
        at.properties.region = "very_cold"
    else:
        at.properties.region = "thermostatted"


h_atom = plams.Molecule()
h_atom.add_atom(plams.Atom(symbol="H", coords=(0.0, 0.0, 0.0)))
h_atom.atoms[0].properties.region = "hydrogen"


main_system_name = ""  # must be empty string
projectile_name = "projectile"
molecules_dict = {main_system_name: slab, projectile_name: h_atom}


# ## Set up the MD settings
#
# Before starting the active learning, let's set up a molecule gun simulation using the initially trained potential.
#
# This is just to see that that simulation settings are somewhat reasonable.

s = Settings()
s.input.ams.Task = "MolecularDynamics"
md_s = s.input.ams.MolecularDynamics
md_s.NSteps = 5000  # will be increased later for active learning
md_s.Trajectory.SamplingFreq = 10  # for testing purposes to check the trajectory
md_s.InitialVelocities.Temperature = 100
md_s.Thermostat = [
    Settings(
        Region="thermostatted",
        Type="NHC",
        Temperature=[300],
        Tau=100.0,
    ),
    Settings(
        Region="very_cold",
        Type="NHC",
        Temperature=[2.0],
        Tau=10.0,
    ),
]
md_s.RemoveMolecules.Frequency = 1
md_s.RemoveMolecules.Formula = "*"
md_s.RemoveMolecules.SinkBox.FractionalCoordsBox = "0 1 0 1 0.90 0.99"
md_s.AddMolecules.System = projectile_name
md_s.AddMolecules.Frequency = 1000
md_s.AddMolecules.StartStep = 100
# insert H atoms 4.5 angstrom above the surface
max_z = max(at.coords[2] for at in slab)
projectile_insertion_z = (4.5 + max_z) / slab.lattice[2][2]
md_s.AddMolecules.FractionalCoordsBox = f"0 1 0 1 {projectile_insertion_z} {projectile_insertion_z+0.01}"
md_s.AddMolecules.VelocityDirection = "0 0 -1"  # shoot down towards slab (decrease z coordinate)
md_s.AddMolecules.DeviationAngle = 0.0
md_s.AddMolecules.Velocity = 0.03


test_md_job = AMSJob(
    settings=s + ParAMSJob.load_external(load_model_dir).results.get_production_engine_settings(),
    molecule=molecules_dict,
    name="test_molecule_gun",
)


test_md_job.run()


# Open the trajectory in AMSmovie to check if it is reasonable. We'd expect some combination of the following events:
#
# * H atoms adsorbing on the Ru surface
# * H atoms diffusing into the subsurface
# * H atoms desorbing from the Ru surface
# * H atoms combining into H2 molecules and desorbing from the surface
#
# The simulation seems reasonable, so let's couple it to the active learning with on-the-fly retraining.
#
# ## Active Learning for Ru/H molecule gun simulation

plams.config.jobmanager.hashing = None
al_s = plams.Settings()
al_s.input.ams.ActiveLearning.Steps.Type = "Linear"
al_s.input.ams.ActiveLearning.Steps.Linear.Start = 1000
al_s.input.ams.ActiveLearning.Steps.Linear.StepSize = 5000
# H atoms at high temperature, so let's decrease the minimum allowed distance a bit
al_s.input.ams.ActiveLearning.ReasonableSimulationCriteria.Distance.MinValue = 0.50
al_s.input.ams.ActiveLearning.MaxReferenceCalculationsPerAttempt = 1
al_s.input.ams.ActiveLearning.SuccessCriteria.Forces.MaxDeviationForZeroForce = 0.65

ml_s = plams.Settings()
ml_s.input.ams.MachineLearning.Backend = "M3GNet"
ml_s.input.ams.MachineLearning.LoadModel = Path(load_model_dir).resolve()
ml_s.input.ams.MachineLearning.MaxEpochs = 50
ml_s.input.ams.MachineLearning.RunAMSAtEnd = "No"

ref_s = dft_settings(QEKPointsConfig(3, 3, 1), conv_thr=1e-4)

new_md_s = s.copy()
new_md_s.input.ams.MolecularDynamics.NSteps = 100000
new_md_s.input.ams.MolecularDynamics.Trajectory.SamplingFreq = 100

al_job = SimpleActiveLearningJob(name="sal", settings=al_s + ml_s + ref_s + new_md_s, molecule=molecules_dict)


al_job.run()