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

# ## Overview
# 
# Import reference data from a short trajectory of aqueous sodium chloride, split it into training and
# validation sets, fine-tune the packaged ``QET-PBE-2025`` foundation model with ParAMS,
# and reuse the deployed model in a production MD job.
# 
# QET combines an equivariant neural network with charge equilibration, so it can describe charge redistribution without requiring atomic-charge labels during fine-tuning.

from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import scm.plams as plams
from scm.base import ChemicalSystem
from scm.params import ParAMSJob, ResultsImporter


# ## Run a quick reference MD job for aqueous NaCl
# 
# In this example we generate some simple reference data.
# 
# If you 
# 
# * already have reference data in the ParAMS .yaml format, then you can skip this step
# * already have reference data in the ASE .xyz or .db format, you can convert it to the ParAMS .yaml format. See :ref:`convert_params_ase`.

sodium = ChemicalSystem.from_smiles("[Na+]")
chloride = ChemicalSystem.from_smiles("[Cl-]")
water = ChemicalSystem.from_smiles("O")

for system in (sodium, chloride, water):
    system.enable_atom_attributes("forcefield")

sodium.atoms[0].forcefield.type = "IP"
sodium.atoms[0].forcefield.charge = 1.0
chloride.atoms[0].forcefield.type = "IM"
chloride.atoms[0].forcefield.charge = -1.0
water.atoms[0].forcefield.type = "OW"
water.atoms[0].forcefield.charge = -0.834
water.atoms[1].forcefield.type = "HW"
water.atoms[1].forcefield.charge = 0.417
water.atoms[2].forcefield.type = "HW"
water.atoms[2].forcefield.charge = 0.417

reference_settings = plams.Settings()
reference_settings.runscript.nproc = 1
reference_settings.input.ForceField.Type = "Amber95"

box = plams.packmol(
    [sodium, chloride, water], n_molecules=[1, 1, 24], density=1.029
)
box = plams.preoptimize(box, settings=reference_settings)
plams.view(box, direction="tilt_x", show_lattice_vectors=True, picture_path="picture1.png")


reference_job = plams.AMSNVTJob(
    settings=reference_settings,
    molecule=box,
    name="amber95_md",
    nsteps=1000,
    timestep=0.5,
    temperature=300,
    thermostat="Berendsen",
    tau=100,
    writeenginegradients=True,
    samplingfreq=100,
)
reference_job.run();


# ## Import reference results with ParAMS ResultsImporter
# 
# Here we use the ``add_trajectory_singlepoints`` results importer. For more details about usage of the results importers, see the corresponding tutorials.
# 
# QET fine-tuning uses the imported energy and force labels; explicit atomic charges are not required.

importer = ResultsImporter(settings={"units": {"energy": "eV", "forces": "eV/angstrom"}})
importer.add_trajectory_singlepoints(reference_job.results.rkfpath(), properties=["energy", "forces"])


# ## Optional: split into training/validation sets
# 
# Machine learning potentials in ParAMS can only be trained if there is both a training set and a validation set.
# 
# If you do not specify a validation set, the training set will automatically be split into a training and validation set when the parametrization starts.
# 
# Here, we will manually split the data set ourselves.
# 
# Let's first print the information in the current ResultsImporter training set:

def print_data_set_summary(data_set: Any, title: str) -> None:
    print(f"{title}:")
    print(f"    number of entries: {len(data_set)}")
    print(f"    number of job IDs: {len(data_set.jobids)}")
    print(f"    job IDs: {data_set.jobids}")


print_data_set_summary(importer.data_sets["training_set"], "Original training set")


# Above, the number of entries is twice the number of jobids because the ``energy`` and ``forces`` extractors are separate entries.
# 
# The energy and force extractors for a given structure (e.g. frame006) must belong to the same data set. For this reason, when doing the split, we call ``split_by_jobid``

training_set, validation_set = importer.data_sets["training_set"].split_by_jobids(0.8, 0.2, seed=314)
importer.data_sets["training_set"] = training_set
importer.data_sets["validation_set"] = validation_set

print_data_set_summary(training_set, "Training set")
print_data_set_summary(validation_set, "Validation set")


# ## Store the reference results in ParAMS yaml format
# 
# Use ``ResultsImporter.store()`` to store all the data in the results importer in the ParAMS .yaml format:

yaml_dir = Path("yaml_ref_data")
importer.store(str(yaml_dir), backup=False)

for path in sorted(yaml_dir.iterdir()):
    print(path)


# ## Set up and run a ParAMSJob for training ML Potentials
# 
# See the ParAMS MachineLearning documentation for all available input options.
# 
# Select the MatGL backend and the packaged ``QET-PBE-2025`` model to fine-tune QET.
# 
# Training the model may take a few minutes.

job = ParAMSJob.from_yaml(str(yaml_dir))
job.name = "params_finetune_qet"
job.settings.runscript.nproc = 1
job.settings.input.Task = "MachineLearning"
job.settings.input.MachineLearning.CommitteeSize = 1
job.settings.input.MachineLearning.MaxEpochs = 250
# job.settings.input.MachineLearning.LossCoeffs.Energy = 10
job.settings.input.MachineLearning.Backend = "MatGL"
job.settings.input.MachineLearning.MatGL.Model = "QET-PBE-2025"
job.settings.input.MachineLearning.MatGL.LearningRate = 0.001
job.settings.input.MachineLearning.MatGL.LearningRateSchedule.Type = "Cosine"
job.settings.input.MachineLearning.MatGL.LearningRateSchedule.FinalFactor = 0.01
job.settings.input.MachineLearning.MatGL.TrainableLayers = "All"
job.settings.input.MachineLearning.Target.Forces.Enabled = "Yes"
job.settings.input.MachineLearning.Target.Forces.MAE = 0.05
job.settings.input.MachineLearning.RunAMSAtEnd = "Yes"


job.run(watch=True);


# ## Results of the ML potential training
# 
# Use ``job.results.get_running_loss()`` to get the loss value as a function of epoch:

fig, ax = plt.subplots()
for data_set in ("training_set", "validation_set"):
    epoch, loss = job.results.get_running_loss(data_set=data_set)
    ax.semilogy(epoch, loss, label=data_set)
ax.set_xlabel("Epoch")
ax.set_ylabel("Loss")
ax.legend()
ax;
ax.figure.savefig("picture2.png")


# The MatGL results also contain the learning rate used for every epoch. Plot it on a logarithmic scale to verify the cosine schedule:

fit_results = job.results.get_machine_learning_results()[0]
epoch, learning_rate = fit_results.get_learning_rate_history()
fig, ax = plt.subplots()
ax.semilogy(epoch, learning_rate)
ax.set_xlabel("Epoch")
ax.set_ylabel("Learning rate")
ax;
ax.figure.savefig("picture3.png")


# If you set ``MachineLearning%RunAMSAtEnd`` (it is on by default), this will run the ML potential through AMS at the end of the fitting procedure, similar to the ParAMS SinglePoint task.
# 
# This will give you access to more results, for example the predicted-vs-reference energy and forces for all entries in the training and validation set. Plot them in a scatter plot like this:

fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(8, 6))

for row, data_set in enumerate(("training_set", "validation_set")):
    evaluator = job.results.get_data_set_evaluator(data_set=data_set, source="best")
    for column, key in enumerate(("energy", "forces")):
        data = evaluator.results[key]
        ax = axes[row, column]
        ax.plot(data.reference_values, data.predictions, ".")
        ax.set_xlabel(f"Reference {key} ({data.unit})")
        ax.set_ylabel(f"Predicted {key} ({data.unit})")
        ax.set_title(f"{data_set}\n{key} MAE: {data.mae:.3f} {data.unit}")
        lower = min(min(data.reference_values), min(data.predictions))
        upper = max(max(data.reference_values), max(data.predictions))
        ax.plot([lower, upper], [lower, upper], linewidth=3, alpha=0.3, color="red")

fig.subplots_adjust(hspace=0.6, wspace=0.4)
axes[-1, -1];


# ## Get the engine settings for production jobs
# 
# First, let's find the path to where the fine-tuned QET model resides using ``get_deployed_model_paths()``. This function returns a list of paths to the trained models. In this case we only trained one model, so we access the first element of the list with ``[0]``.
# 
# The returned path is the path we need to give as the ``ParameterDir`` input option in the AMS MLPotential engine. For other backends it might instead be the ``ParameterFile`` option.
# 
# To get the complete engine settings as a PLAMS Settings object, use the method ``get_production_engine_settings()``:

print(job.results.get_deployed_model_paths()[0])

production_settings = job.results.get_production_engine_settings()
print(plams.AMSJob(settings=production_settings).get_input())


# ## Run a production MD simulation with the trained potential
# 
# The production simulation uses a Nose--Hoover chain thermostat at 300 K with a 100 fs damping constant.

production_settings.runscript.nproc = 1
production_job = plams.AMSNVTJob(
    settings=production_settings,
    molecule=box,
    nsteps=1000,
    temperature=300,
    thermostat="NHC",
    tau=100,
    samplingfreq=100,
    name="production_md",
    timestep=1.0,
)
production_job.run(watch=True);


charges = production_job.results.get_charges(engine="MDStep1000")
print(f"{'Atom':>4}  {'Element':>7}  {'Charge (e)':>12}")
for index, (atom, charge) in enumerate(zip(box, charges), start=1):
    print(f"{index:4d}  {atom.symbol:>7}  {charge:12.6f}")


# ## Open trajectory file in AMSmovie
# 
# With the production trajectory you can run analysis tools in AMSmovie, or access them from Python. See the AMS manual for details.

trajectory_file = production_job.results.rkfpath()
print(trajectory_file)
# !amsmovie "{trajectory_file}"