import numpy as np
from flare.gp import GaussianProcess

# make gp model
hyps = np.array([0.1, 1, 0.01])
hyp_labels = ['Signal Std', 'Length Scale', 'Noise Std']
cutoffs = {'threebody':3.9}

gp = \
    GaussianProcess(kernels=['threebody'],
                    hyps=hyps,
                    cutoffs=cutoffs,
                    hyp_labels=hyp_labels,
                    maxiter=50)

"""
Explanation: On-the-fly aluminum potential
For the scheme/workflow of the on-the-fly training, please refer to Vandermause et al.. 
In our FLARE python package, we provide two ways for the on-the-fly training:
1. Use OTF module and our MD engine (only supports NVE) $\to$ this tutorial
2. Use ASE Atoms with our ASE interface with OTF $\to$ see ase-tutorial
Here we give an example of running OTF (on-the-fly) training with QE (Quantum Espresso) and NVE ensemble.
Step 1: Set up a GP Model
Let’s start up with the GP model with three-body kernel function. (See kernels.py (single component) or mc_simple.py (multi-component) for more options.)
End of explanation
"""


import os
# set up DFT calculator
qe_input = './pwscf.in' # quantum espresso input file
dft_loc = os.environ.get('PWSCF_COMMAND')

"""
Explanation: Some Explanation about the parameters


kernels: set to be the name list of kernel functions to use

Currently we have the choices of twobody, threebody and manybody kernel functions.
If multiple kernels are listed, the resulted kernel is simply the summation of all listed kernels,



hyps: the array of hyperparameters, whose names are shown in hyp_labels.

For two-body kernel function, an array of length 3 is needed, hyps=[sigma_2, ls_2, sigma_n];
For three-body, hyps=[sigma_3, ls_3, sigma_n];
For twobody plus threebody, hyps=[sigma_2, ls_2, sigma_3, ls_3, sigma_n].
For twobody, threebody plus manybody, hyps=[sigma_2, ls_2, sigma_3, ls_3, sigma_m, ls_m, sigma_n].



cutoffs: a dictionary consists of corresponding cutoff values for each kernel. Usually we will set a larger one for two-body, and smaller one for threebody and manybody


maxiter: set to constrain the number of steps in training hyperparameters.


Note

See GaussianProcess for complete description of arguments of GaussianProcess class.
See AdvancedHyperparametersSetUp for more complicated hyper-parameters set up.

Step 2: Set up DFT Calculator
The next step is to set up DFT calculator, here we use QE (quantum espresso). Suppose we’ve prepared a QE input file in current directory ./pwscf.in, and have set the environment variable PWSCF_COMMAND to the location of our QE’s executable pw.x. Then we specify the input file and executable by qe_input and dft_loc.
End of explanation
"""


# set up OTF parameters
dt = 0.001                  # timestep (ps)
number_of_steps = 100       # number of steps
std_tolerance_factor = 1
max_atoms_added = 2
freeze_hyps = 3

otf = OTF(qe_input, dt, number_of_steps, gp, dft_loc,
          std_tolerance_factor, 
          init_atoms=[0],
          calculate_energy=True, 
          output_name='al_otf_qe',
          freeze_hyps=freeze_hyps, 
          skip=5,
          max_atoms_added=max_atoms_added,
          write_model=3)

"""
Explanation: Step 3: Set up OTF MD Training Engine
Then we can set up our On-The-Fly (OTF) MD engine for training and simulation.
End of explanation
"""


# run OTF MD
otf.run()

"""
Explanation: Some Explanation about the parameters


dt: the time step in unit of ps


number_of_steps: the number of steps that the MD is run


std_tolerance_factor: the uncertainty threshold = std_tolerance_factor x hyps[-1]. In OTF training, when GP predicts uncertainty above the uncertainty threshold, it will call DFT


max_atoms_added: constrain the number of atoms added to the training set after each DFT call


freeze_hyps: stop training hyperparameters and fix them from the freeze_hyps th step. Usually set to a small number, because for large dataset the training will take long.


init_atoms: list of atoms to be added in the first DFT call. Because there’s no uncertainty predicted in the initial DFT call, so there’s no selection rule to pick up “maximully uncertain” atoms into the training set, we have to specify which atoms to pick up by this variable.


calculate_energy: if True, the local energy on each atom will be calculated


output_name: the name of the logfile


skip: record/dump the information every skip steps.


write_model: identify the frequency of dumpping GP model during the training. 


1: default, dump GP model only when the OTF is complete


2: dump GP model every time after training


3: dump GP model at every step


(We recommend using 3 if you are possibly gonna resume the OTF training, as instructed by the next section.)


Step 4: Launch the OTF Training
Finally, let’s run it!
End of explanation
"""


new_otf = OTF.from_checkpoint(log_name + "_checkpt.json")
new_otf.run()

"""
Explanation: After OTF training is finished, we can check log file al_otf_qe.out for all the information dumped. This output file can be parsed using our otf_parser.py module, which will be introduced in the after-training-tutorial
Step 5 (Optional): Resume Interrupted Training
At the end of each OTF training step, there will be a checkpoint file dumpped of name &lt;output_name&gt;_checkpt.json, in the above example, al_otf_qe_checkpt.json. This file saves the state at the current OTF MD step. 
If you've set write_model=3, then there will be another file saving the trained GP model, &lt;output_name&gt;_gp.json, which will be loaded when restarting OTF.
Then, use OTF.from_checkpoint(&lt;output_name&gt;_checkpt.json) to load the OTF state, and resume the training by run().
End of explanation
"""


