import os
from copy import deepcopy
import numpy as np
from numpy.random import random
import sys

# FLARE imports
from flare.learners.otf import OTF

# FLARE_PP imports
from flare.bffs.sgp._C_flare import SparseGP, NormalizedDotProduct, B2,     Structure
from flare.bffs.sgp.sparse_gp import SGP_Wrapper
from flare.bffs.sgp.calculator import SGP_Calculator

# ASE imports
from ase import units
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, \
    Stationary, ZeroRotation
from ase.calculators.espresso import Espresso
from ase.calculators.vasp import Vasp
from ase.spacegroup import crystal
from ase.build import supercells, fcc111, add_adsorbate
from ase import Atoms
from ase.io import read

# Set NVT MD parameters.
timestep = 5.0e-4 # in ps
temperature = 473   # K

taut = 0.01 * 1000 * units.fs
md_engine = 'NVTBerendsen'

md_dict = {'temperature_K': temperature, 'taut': taut}

# Create initial structure.
atoms = read("POSCAR") #initial structure

jitter_factor = 0.01 # this structure already has nonzero forces, so we only perturb the positions by small amount
# increase jitter factor to 0.1 or more if you are already in a local minima (optimized structure)
for atom_pos in atoms.positions:
    for coord in range(3):
        atom_pos[coord] += (2 * random() - 1) * jitter_factor

n_atoms = len(atoms)

# Create sparse GP model.
cutoff = 4.5

sigma = 2.0
power = 2
kernel = NormalizedDotProduct(sigma, power)
cutoff_function = "quadratic"
many_body_cutoffs = [cutoff]
radial_basis = "chebyshev"
radial_hyps = [0., cutoff]
cutoff_hyps = []
descriptor_settings = [6, 8, 4] # the first index here is the number of atom types
# second and third indices are nmax = 8 for radial basis, and lmax = 4 for angular basis (Refer to the SI pdf)
# cutoffs = np.array([[cutoff, h_cutoff], [h_cutoff, h_cutoff]]) # this is how you define a cutoff matrix if needed
descriptor_calculator = \
    B2(radial_basis, cutoff_function, radial_hyps, cutoff_hyps,
       descriptor_settings)#, cutoffs)
sigma_e = 0.001 * n_atoms  # 1 meV/atom
sigma_f = 0.1
sigma_s = 0.0006  

# Wrap C++ model in Python.
rydberg_to_ev = 13.6056980659
species_map = {14: 0, 8: 1, 13: 2,27:3,1:4,7:5} # assign species labels for each type of atoms
single_atom_energies = {0: 0.0, 1: 0.0, 2: 0.0,3:0.0,4:0.0,5:0.0}
variance_type = "local"
max_iterations = 60
opt_method = "L-BFGS-B"
bounds = [(None, None), (sigma_e, None), (None, None), (sigma_s, None)]
gp_model = SGP_Wrapper([kernel], [descriptor_calculator], cutoff,
                       sigma_e, sigma_f, sigma_s,
                       species_map, single_atom_energies=single_atom_energies,
                       variance_type=variance_type,
                       opt_method=opt_method, bounds=bounds,
                       max_iterations=max_iterations)


# Load previous GP #This is how to load the previously trained model for retraining/adding data
#gp_file = "/Path to your file/your_MLP_file.json"
#prev_gp = SparseGP.from_json(gp_file)
#gp_model.sparse_gp = prev_gp

# Create ASE FLARE calculator.
flare_calculator = SGP_Calculator(gp_model)

# Create DFT calculator.

dft_calculator = Vasp(xc = 'pbe',
    ispin = 2,
    encut = 400,
    ismear = 0,
    algo='Fast',
    istart=0,
    lreal='Auto',
    prec = 'Accurate',
    sigma = 0.01,
    ivdw=12,
    ediff=1e-5,
    nelm=500,
    nelmin=6,
    nelmdl=-9,
    nelect=39) # assign the correct number of valence electrons (nelect) for the complexes in the vacuum


# Set up OTF object.
init_atoms = list(range(n_atoms))
output_name = 'output'
std_tolerance_factor = -0.1
max_atoms_added = n_atoms
freeze_hyps = (0,300)
os.mkdir('./output_files')
store_dft_output = ('OUTCAR', 'output_files')
min_steps_with_model = 5 # This forces model to do 5 steps without calling DFT, regardless of the prediction uncertanity
update_style = "threshold"
update_threshold = 0.1

otf_params = {'init_atoms': init_atoms, 'output_name': output_name,
              'std_tolerance_factor': std_tolerance_factor,
              'max_atoms_added': max_atoms_added, 'train_hyps': freeze_hyps,
              'store_dft_output': store_dft_output,
              'min_steps_with_model': min_steps_with_model,
              'update_style': update_style,
              'update_threshold': update_threshold}

# Initialize velocity.
temperature = temperature * units.kB
MaxwellBoltzmannDistribution(atoms, temperature)
Stationary(atoms)  # zero linear momentum
ZeroRotation(atoms)  # zero angular momentum

# Create OTF object.
number_of_steps = 100000 # change this to access the total time you need
test_otf = OTF(
    atoms, dt=timestep,
    number_of_steps=number_of_steps,
    dft_calc=dft_calculator, md_engine=md_engine, md_kwargs=md_dict,
    flare_calc=flare_calculator,
    **otf_params)

test_otf.run()

# Save the GP.
SparseGP.to_json("CoNH3x4_tetra.json", gp_model.sparse_gp)
