#!/usr/bin/env amspython
from scm.plams import *
from scm.params import *
import numpy as np
import os
import matplotlib.pyplot as plt


def main():
    init()

    inputfile = os.path.expandvars("$AMSHOME/scripting/scm/params/examples/LJ_Ar_restart/params.in")

    job = ParAMSJob.from_inputfile(inputfile, name="initial")
    job.run()

    earliest_checkpoint_file = job.results.get_checkpoints()[0]

    restart_job = ParAMSJob.from_inputfile(inputfile, name="resume_from_checkpoint")
    restart_job.resume_checkpoint = earliest_checkpoint_file
    restart_job.settings.input.ExitCondition = []
    restart_job.add_exit_condition("MaxTotalFunctionCalls", 400)

    restart_job.run()

    evaluation, loss = job.results.get_running_loss()
    restart_evaluation, restart_loss = restart_job.results.get_running_loss()

    plt.plot(evaluation, np.log10(loss), restart_evaluation, np.log10(restart_loss))
    plt.legend(["Initial", "Resume from Checkpoint"])
    plt.xlabel("Evaluation id")
    plt.ylabel("log10(loss)")
    plt.show()

    finish()


if __name__ == "__main__":
    main()