import matplotlib

matplotlib.use("Agg")

import numpy as np
from ase.build import fcc100

from agox import AGOX
from agox.acquisitors import LowerConfidenceBoundAcquisitor
from agox.collectors import ParallelCollector
from agox.databases import Database
from agox.environments import Environment
from agox.evaluators import LocalOptimizationEvaluator
from agox.generators import RandomGenerator, RattleGenerator
from agox.models.descriptors.fingerprint import Fingerprint
from agox.models.GPR import GPR
from agox.models.GPR.kernels import RBF, Noise
from agox.models.GPR.kernels import Constant as C
from agox.models.GPR.priors import Repulsive
from agox.postprocessors import ParallelRelaxPostprocess
from agox.samplers import KMeansSampler

# Manually set seed and database-index
seed = 41
database_index = 0

##############################################################################
# Calculator
##############################################################################

from ase.calculators.vasp import Vasp

nslots = 8  # Number of cores
calc = Vasp(
    command="mpirun -np {} vasp".format(nslots),
    prec="Low",
    lreal=False,
    lwave=False,
    xc="PBE",
    nwrite=2,
    istart=1,
    icharg=1,
    ispin=1,
    voskown=1,
    lmaxmix=6,
    nelm=500,
    nelmin=4,
    ibrion=-1,
    potim=0.3,
    iwavpr=1,
    nsw=0,
    ialgo=38,
    ismear=0,
    sigma=0.1,
    lorbit=11,
)

##############################################################################
# System & general settings:
##############################################################################

# Define a template surface using ASE functions.
template = fcc100("Au", size=(4, 4, 2), vacuum=5.0)
template.pbc = [True, True, False]
template.positions[:, 2] -= template.positions[:, 2].min()

# Confinement cell matches the template cell in x & y but is smaller in z.
confinement_cell = template.cell.copy()
confinement_cell[2, 2] = 4.0

# Confinement corner is at the origin of the template cell, except in z where
# it is shifted up.
z0 = template.positions[:, 2].max()
confinement_corner = np.array([0, 0, z0])  # Confinement corner is at cell origin

environment = Environment(
    template=template,
    symbols="O4",
    confinement_cell=confinement_cell,
    confinement_corner=confinement_corner,
    box_constraint_pbc=[True, True, False],  # Confinement is not periodic in z
)

# Database
db_path = "db{}.db".format(database_index)  # From input argument!
database = Database(filename=db_path, order=5)

##############################################################################
# Search Settings:
##############################################################################

# Setup a ML model.
descriptor = Fingerprint(environment=environment)
beta = 0.01
k0 = C(beta, (beta, beta)) * RBF()
k1 = C(1 - beta, (1 - beta, 1 - beta)) * RBF()
kernel = C(5000, (1, 1e5)) * (k0 + k1) + Noise(0.01, (0.01, 0.01))
model = GPR(descriptor=descriptor, kernel=kernel, database=database, prior=Repulsive())

# Sampler to choose candidates to modify using the generators.
sample_size = 10
sampler = KMeansSampler(descriptor=descriptor, database=database, sample_size=sample_size)

# Generators to produce candidates structures
rattle_generator = RattleGenerator(**environment.get_confinement())
random_generator = RandomGenerator(**environment.get_confinement(), contiguous=False)

# Dict specificies how many candidates are created with and the dict-keys are iterations.
generators = [random_generator, rattle_generator]
num_candidates = {0: [10, 0], 5: [3, 7]}

# Collector creates a number of structures in each iteration.
collector = ParallelCollector(
    generators=generators,
    sampler=sampler,
    environment=environment,
    num_candidates=num_candidates,
    order=1,
)

# Acquisitor to choose a candidate to evaluate in the real potential.
acquisitor = LowerConfidenceBoundAcquisitor(model=model, kappa=2, order=3)

# Number of steps is very low - should be set higher for a real search!
relaxer = ParallelRelaxPostprocess(
    model=acquisitor.get_acquisition_calculator(),
    constraints=environment.get_constraints(),
    optimizer_run_kwargs={"steps": 5},
    start_relax=8,
    order=2,
)


# Evaluator to evaluate the candidates in the real potential.
from ase.calculators.calculator import CalculationFailed
from ase.calculators.vasp import Vasp

def callback(candidate):
    if isinstance(candidate.calc, Vasp) and not candidate.calc.read_convergence():
        raise CalculationFailed('VASP: SCF not converged')

evaluator = LocalOptimizationEvaluator(
    calc,
    gets={"get_key": "prioritized_candidates"},
    optimizer_kwargs={"logfile": None},
    optimizer_run_kwargs={"fmax": 0.05, "steps": 1},
    constraints=environment.get_constraints(),
    store_trajectory=True,
    order=4,
    check_callback=callback,
)

##############################################################################
# Let get the show running!
##############################################################################

# The oder of things here does not matter. But it can be simpler to understand
# what the expected behaviour is if they are put in order.
agox = AGOX(collector, relaxer, acquisitor, evaluator, database, seed=seed)

agox.run(N_iterations=10)
