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

# ## Initial imports

from scm.simple_active_learning import SimpleActiveLearningJob
import scm.plams as plams
from scm.external_engines.core import interface_is_installed

assert interface_is_installed("m3gnet"), "You must first install m3gnet with the AMS package manager"


# ## Initialize PLAMS

plams.init()


# ## Input system

mol = plams.from_smiles("OCC=O")
for at in mol:
    at.properties = {}
plams.plot_molecule(mol)


# ## Reference engine settings
# For time reasons we use the UFF force field as the reference method. Typically you would instead train to DFT using ADF, BAND, or Quantum ESPRESSO.

ref_s = plams.Settings()
ref_s.input.ForceField.Type = "UFF"
ref_s.runscript.nproc = 1


print(plams.AMSJob(settings=ref_s).get_input())


# ## Molecular dynamics settings
# Here, we use the convenient ``AMSNVTJob`` recipe to easily initialize sone MD settings.

md_s = plams.AMSNVTJob(temperature=300, timestep=0.5, nsteps=10000).settings


print(plams.AMSJob(settings=md_s).get_input())


# ## ParAMS ML Training settings
#
# (Technical note: When using ``SimpleActiveLearningJob`` the ParAMS settings go under ``input.ams``. When using ``ParAMSJob`` the settings instead simply go under ``input``. See the ParAMS Python tutorials.)

ml_s = plams.Settings()
ml_s.input.ams.MachineLearning.Backend = "M3GNet"
ml_s.input.ams.MachineLearning.CommitteeSize = 1
ml_s.input.ams.MachineLearning.M3GNet.Model = "UniversalPotential"
ml_s.input.ams.MachineLearning.MaxEpochs = 200
print(SimpleActiveLearningJob(settings=ml_s).get_input())


# ## Active learning settings

al_s = plams.Settings()
al_s.input.ams.ActiveLearning.Steps.Type = "Geometric"
al_s.input.ams.ActiveLearning.Steps.Geometric.Start = 10  # 10 MD frames
al_s.input.ams.ActiveLearning.Steps.Geometric.NumSteps = 5  # 10 AL steps
print(SimpleActiveLearningJob(settings=al_s).get_input())


# ## Simple Active Learning Job

settings = ref_s + md_s + ml_s + al_s
job = SimpleActiveLearningJob(settings=settings, molecule=mol, name="sal")
print(job.get_input())


# ## Run the job

job.run(watch=True)