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

# ## Initial imports

import scm.plams as plams
from scm.params import ResultsImporter
from scm.plams import Settings, AMSJob, log, Molecule, packmol_on_slab
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,
    m3gnet_up_settings,
    replay_settings,
    slice_slab,
    check_installation,
)


# ## Initialize PLAMS working directory

old_ref_dir = "reference_data_2"
check_installation(old_ref_dir)
new_ref_dir = "reference_data_3"
ri = ResultsImporter.from_yaml(old_ref_dir)
plams.init()


# ## Construct the Ru(10-10)/H2(gas) interface
#
# For details about the construction of the slab, see the previous notebook.

optimized_bulk = ri.job_collection["hcp_lattopt_Ru_dft"].molecule
slab_100 = slice_slab(optimized_bulk, miller=(1, 0, 0), thickness=7.0, cell_z=15, ref_atom=0)
slab_100 = slab_100.supercell(3, 2, 1)
for at in slab_100:
    at.properties = Settings()  # remove details about supercell generation
plams.plot_molecule(slab_100, rotation=rotation)
plt.title("Ru(10-10)")


# Now use the ``packmol_slab`` function to add hydrogen molecules:

from scm.plams import packmol_on_slab

h2_mol = plams.from_smiles("[HH]")
density = 0.3  # approximate density of the gas phase in g/cm^3
slab_100_H2_gas_raw = packmol_on_slab(slab_100, h2_mol, density=density)
slab_100_H2_gas = plams.preoptimize(slab_100_H2_gas_raw, settings=m3gnet_up_settings(), maxiterations=100)


plams.plot_molecule(slab_100_H2_gas, rotation=rotation, radii=0.8)


# Now let's run some MD:

mdjob = plams.AMSNVTJob(
    settings=m3gnet_up_settings(),
    name="md_ru10-10_h2_m3gnet",
    temperature=501,
    nsteps=10000,
    molecule=slab_100_H2_gas,
    samplingfreq=100,
)
mdjob.run()


from scm.params import ResultsImporter

ri = ResultsImporter.from_yaml(old_ref_dir)


settings = dft_settings(QEKPointsConfig(3, 3, 1))
settings += replay_settings(mdjob.results.rkfpath(), frames=[10, 30, 50, 70, 90])

dft_replay_job = plams.AMSJob(
    settings=settings,
    name="snapshots_from_md_ru10-10_h2_dft",
)
dft_replay_job.run(watch=True)


ri.add_trajectory_singlepoints(dft_replay_job, properties=["energy", "forces"])
ri.store(new_ref_dir)