import numpy as np
from ase import units
from ase.spacegroup import crystal
from ase.build import bulk

np.random.seed(12345)

a = 3.52678
super_cell = bulk('C', 'diamond', a=a, cubic=True) 

"""
Explanation: On-the-fly training using ASE
Yu Xie (xiey@g.harvard.edu)
This is a quick introduction of how to set up our ASE-OTF interface to train a force field. We will train a force field model for diamond. To run the on-the-fly training, we will need to

Create a supercell with ASE Atoms object
Set up FLARE ASE calculator, including the kernel functions, hyperparameters, cutoffs for Gaussian process, and mapping parameters (if Mapped Gaussian Process is used)
Set up DFT ASE calculator. Here we will give an example of Quantum Espresso
Set up on-the-fly training with ASE MD engine

Please make sure you are using the LATEST FLARE code in our master branch.
Step 1: Set up supercell with ASE
Here we create a 2x1x1 supercell with lattice constant 3.855, and randomly perturb the positions of the atoms, so that they will start MD with non-zero forces.
End of explanation
"""


from flare.gp import GaussianProcess
from flare.utils.parameter_helper import ParameterHelper

# set up GP hyperparameters
kernels = ['twobody', 'threebody'] # use 2+3 body kernel
parameters = {'cutoff_twobody': 5.0, 
              'cutoff_threebody': 3.5}
pm = ParameterHelper(
    kernels = kernels,   
    random = True,
    parameters=parameters
)

hm = pm.as_dict()
hyps = hm['hyps']
cut = hm['cutoffs']
print('hyps', hyps)

gp_model = GaussianProcess(
    kernels = kernels,
    component = 'mc', # If you are using ASE, please set to "mc" no matter for single-component or multi-component
    hyps = hyps,
    cutoffs = cut,
    hyp_labels = ['sig2','ls2','sig3','ls3','noise'],
    opt_algorithm = 'L-BFGS-B',
    n_cpus = 1
)

"""
Explanation: Step 2: Set up FLARE calculator
Now let’s set up our Gaussian process model in the same way as introduced before
End of explanation
"""


from flare.mgp import MappedGaussianProcess

grid_params = {'twobody':   {'grid_num': [64]},
               'threebody': {'grid_num': [16, 16, 16]}}

mgp_model = MappedGaussianProcess(grid_params, 
                                  unique_species = [6], 
                                  n_cpus = 1,
                                  var_map=None)

"""
Explanation: Optional
If you want to use Mapped Gaussian Process (MGP), then set up MGP as follows
End of explanation
"""


from flare.ase.calculator import FLARE_Calculator

flare_calculator = FLARE_Calculator(gp_model, 
                                    par = True, 
                                    mgp_model = None,
                                    use_mapping = False)

super_cell.set_calculator(flare_calculator)

"""
Explanation: Now let's set up FLARE's ASE calculator. If you want to use MGP model, then set use_mapping = True and mgp_model = mgp_model below.
End of explanation
"""


from ase.calculators.lj import LennardJones
lj_calc = LennardJones() 

"""
Explanation: Step 3: Set up DFT calculator
For DFT calculator, you can use any calculator provided by ASE, e.g. Quantum Espresso (QE), VASP, etc. 
For a quick illustration of our interface, we use the Lennard-Jones (LJ) potential as an example.
End of explanation
"""


import os
from ase.calculators.espresso import Espresso

# ---------------- set up executable ---------------- 
label = 'C'
input_file = label+'.pwi'
output_file = label+'.pwo'
no_cpus = 32
npool = 32
pw_loc = 'path/to/pw.x'

# serial
os.environ['ASE_ESPRESSO_COMMAND'] = f'{pw_loc} < {input_file} > {output_file}'

## parallel qe using mpirun
# os.environ['ASE_ESPRESSO_COMMAND'] = f'mpirun -np {no_cpus} {pw_loc} -npool {npool} < {input_file} > {output_file}'

## parallel qe using srun (for slurm system)
# os.environ['ASE_ESPRESSO_COMMAND'] = 'srun -n {no_cpus} --mpi=pmi2 {pw_loc} -npool {npool} < {input_file} > {output_file}'


# -------------- set up input parameters -------------- 
input_data = {'control':   {'prefix': label, 
                            'pseudo_dir': './',
                            'outdir': './out',
                            'calculation': 'scf'},
              'system':    {'ibrav': 0, 
                            'ecutwfc': 60,
                            'ecutrho': 360},
              'electrons': {'conv_thr': 1.0e-9,
                            'electron_maxstep': 100,
                            'mixing_beta': 0.7}}

# ----------------  pseudo-potentials -----------------         
ion_pseudo = {'C': 'C.pz-rrkjus.UPF'}

# -------------- create ASE calculator ---------------- 
dft_calc = Espresso(pseudopotentials=ion_pseudo, label=label, 
                    tstress=True, tprnfor=True, nosym=True, 
                    input_data=input_data, kpts=(8, 8, 8)) 

"""
Explanation: Optional: alternatively, set up Quantum Espresso calculator
We also give the code below for setting up the ASE quantum espresso calculator, following the instruction on ASE website.
First, we need to set up our environment variable ASE_ESPRESSO_COMMAND to our QE executable, so that ASE can find this calculator. Then set up our input parameters of QE and create an ASE calculator
End of explanation
"""


from ase import units
from ase.md.velocitydistribution import (MaxwellBoltzmannDistribution,
                                         Stationary, ZeroRotation)

temperature = 500
MaxwellBoltzmannDistribution(super_cell, temperature * units.kB)
Stationary(super_cell)  # zero linear momentum
ZeroRotation(super_cell)  # zero angular momentum

md_engine = 'VelocityVerlet'
md_kwargs = {}

"""
Explanation: Step 4: Set up On-The-Fly MD engine
Finally, our OTF is compatible with  


5 MD engines that ASE supports: VelocityVerlet, NVTBerendsen, NPTBerendsen, NPT and Langevin, 


and 1 MD engine implemented by FLARE: NoseHoover.


We can choose any of them, and set up the parameters based on ASE requirements. After everything is set up, we can run the on-the-fly training.
Note 1: Currently, only VelocityVerlet is tested on real system, NPT may have issue with pressure and stress.
Set up ASE_OTF training engine:


Initialize the velocities of atoms as 500K


Set up MD arguments as a dictionary based on ASE MD parameters. For VelocityVerlet, we don't need to set up extra parameters.
E.g. for NVTBerendsen, we can set md_kwargs = {'temperature': 500, 'taut': 0.5e3 * units.fs}


Note 2: For some tricks and tips related to the on-the-fly training (e.g. how to set up temperatures, how to optimize hyperparameters), see FAQs
End of explanation
"""


from flare.ase.otf import ASE_OTF

otf_params = {'init_atoms': [0, 1, 2, 3],
              'output_name': 'otf',
              'std_tolerance_factor': 2,
              'max_atoms_added' : 4,
              'freeze_hyps': 10,
              'write_model': 3} # If you will probably resume the training, please set to 3

test_otf = ASE_OTF(super_cell, 
                   timestep = 1 * units.fs,
                   number_of_steps = 3,
                   dft_calc = lj_calc,
                   md_engine = md_engine,
                   md_kwargs = md_kwargs,
                   **otf_params)

test_otf.run()

"""
Explanation: Set up parameters for On-The-Fly (OTF) training. The descriptions of the parameters are in ASE OTF module.

Set up the ASE_OTF training engine, and run
Note: the ASE Trajectory is supported, but NOT recommended. 


Check otf.out after the training is done.
End of explanation
"""


new_otf = ASE_OTF.from_checkpoint("<output_name>_checkpt.json")
new_otf.run()

"""
Explanation: Then check the *.out file for the training log.
Step 5 (Optional): Resume Interrupted Training
At the end of each OTF training step, there will be several checkpoint files dumpped


&lt;output_name&gt;_checkpt.json: checkpoint of the current MD step of OTF. In the above example, al_otf_qe_checkpt.json. 


&lt;output_name&gt;_flare.json: If you've set write_model=3, then there will be another file saving the trained FLARE calculator, which will be loaded when restarting OTF.


&lt;output_name&gt;_atoms.json: The ASE Atoms of the current MD step in the format of json


&lt;output_name&gt;_dft.pickle: The DFT calculator saved in the format of .pickle.


Then, use ASE_OTF.from_checkpoint(&lt;output_name&gt;_checkpt.json) to load the OTF state, and resume the training by run().
End of explanation
"""


