#!/usr/bin/env amspython
# Usage: "$AMSBIN/amspython" optimization.py
from scm.plams import *
from scm.params import *
import os
import matplotlib.pyplot as plt


def main():
    init()  # initialize PLAMS workdir
    examples_dir = os.path.expandvars("$AMSHOME/scripting/scm/params/examples/ReaxFF_water")

    # load the optimization job from examples directory  and run
    optimization_job = ParAMSJob.from_inputfile(examples_dir + "/params.in", name="optimization")
    optimization_job.run()
    # or load a finished previously run job:
    # optimization_job = ParAMSJob.load_external('plams_workdir/optimization/results')

    print("Optimized force field stored in " + optimization_job.results.get_ffield())

    # loop over all PES scans and show plots
    # in this case there is only 1 PES Scan
    all_pesscans = optimization_job.results.get_all_pes_prediction_names()
    for scan in all_pesscans:
        x_headers, x_values, ref_values, pred_values, unit = optimization_job.results.get_pes_prediction(scan)
        plt.plot(x_values[0], ref_values, "-")  # ref. vs distance
        plt.plot(x_values[0], pred_values, "-")  # pred. vs distance
        plt.legend(["Reference", "Prediction"])
        plt.xlabel(x_headers[0])
        plt.ylabel(f"Relative energy ({unit})")
        plt.title(scan)
        plt.savefig(f"plot_{scan}_from_python.png")
        plt.show()

    finish()


if __name__ == "__main__":
    main()