#!/usr/bin/env python3
"""
Lattice constant Test Driver

Computes the lattice constant for any material and any cubic crystal structure
by minimizing the energy using a simplex minimization

Date: 2012/09/03
Author: Alex Alemi & Matt Bierbaum

Updates: 2016/02/26 Junhao Li

         2017/08/21 Ellad Tadmor

         2018/02/17 Dan Karls

         2019/02/24 Dan Karls (added logic to check if lattice is unstable,
                               handle very small lattice constants for Toy
                               Models)

         2019/06/13 Dan Karls (convert to python3, check if species has non-trivial
                               energy interaction, add maxiter constraint to
                               minimizations, calculate cutoff of model and use it to
                               disallow small lattice spacings that would lead to a
                               massive number of neighbors/atom)
"""

# Python 2-3 compatible code issues
from __future__ import print_function

try:
    input = raw_input
except NameError:
    pass

import os
import json

import jinja2
import scipy.optimize as opt

from ase.lattice.cubic import BodyCenteredCubic, FaceCenteredCubic, Diamond, SimpleCubic
from ase.calculators.kim.kim import KIM
import kim_python_utils.ase as kim_ase_utils

def create_cubic_lattice(symbol, lattice, a, ncells_per_side):
    """Note that we use ase.lattice instead of ase.build.bulk so that we can construct
       finite clusters of unit cells if necessary."""
    dimensions = (ncells_per_side, ncells_per_side, ncells_per_side)
    if lattice == "bcc":
        lattice = BodyCenteredCubic(symbol, latticeconstant=a, size=dimensions)
    elif lattice == "fcc":
        lattice = FaceCenteredCubic(symbol, latticeconstant=a, size=dimensions)
    elif lattice == "diamond":
        lattice = Diamond(symbol, latticeconstant=a, size=dimensions)
    elif lattice == "sc":
        lattice = SimpleCubic(symbol, latticeconstant=a, size=dimensions)

    return lattice


def energy(a, atoms, cell, positions, isolated_energy_per_atom, min_allowed_latconst):
    """ Given the lattice constant, compute the total energy of a lattice supercell
    relative to the isolated energy of the same number of atoms.
    """
    if a < min_allowed_latconst:
        raise RuntimeError("Attempted to evaluate energy at lattice spacing ({}) below "
                "the minimum allowed value ({}).  This may mean that the model does "
                "not possess close-range repulsive forces to prevent system collapse."
                "".format(a, min_allowed_latconst))
    atoms.set_positions(positions * a)
    atoms.set_cell(cell * a)
    energy = atoms.get_potential_energy() - len(atoms) * isolated_energy_per_atom
    return energy


def get_min_allowed_latconst(lattice, cutoff):
    """
    Heuristic used to prevent neighbor list overflow.  Based on the restriction that no
    atom must be allowed to have more than 5000 neighbors using the provided cutoff.
    """
    min_prefactors = {"bcc": 0.143, "diamond": 0.227, "fcc": 0.18, "sc": 0.113}
    return min_prefactors[lattice]*cutoff


def get_lattice_constant(model, symbol, lattice, maxiter):
    """ Here, we start with an initial guess of 2.5 Angstroms for the lattice
    constant. We then attempt to minimize the potential energy using this
    initial value and, if the optimization fails for any reason, we increment
    the initial guess by 0.5 Angstroms and repeat the calculation. This
    process repeats until either the calculation succeeds or until a maximum
    initial guess is exceeded.
    """
    repeat = 0  # misc. vc info (not currently used)

    # Check if atoms have an energy interaction for this model
    atoms_interacting = kim_ase_utils.check_if_atoms_interacting(
        model, symbols=[symbol, symbol], check_force=False, etol=1e-6
    )
    if not atoms_interacting:
        raise RuntimeError(
            "The model provided, {}, does not possess a non-trivial "
            "energy interaction for species {} as required by this Test.  Aborting."
            "".format(model, symbol)
        )

    # Get model cutoff
    cutoff = kim_ase_utils.get_model_energy_cutoff(model, symbols=[symbol, symbol])

    # Get isolated energy per atom
    isolated_energy_per_atom = {}
    isolated_energy_per_atom[symbol] = kim_ase_utils.get_isolated_energy_per_atom(
        model, symbol
    )

    # Create lattice and calculator
    atoms = create_cubic_lattice(symbol, lattice, a=1, ncells_per_side=1)
    num_atoms = len(atoms)
    cell = atoms.get_cell()
    positions = atoms.get_positions()
    calc = KIM(model)
    atoms.set_calculator(calc)

    init_guess = 2.5
    max_init_guess = 10.0

    # Check if the initial guess would be outside of the cutoff.  This is unlikely, but
    # can occur for Toy Models.  On the other hand, many non-toy Models may produce an
    # error because init_guess is too small for them to handle and, in this case, we
    # simply proceed to the loop below, which will catch the associated exceptions as it
    # tries larger initial guesses.
    init_guess_too_large = init_guess >= cutoff

    # Get minimum allowed lattice constant that can be used to evaluate the energy
    # function.  This is intended to prevent a neighbor list overflow.
    min_allowed_latconst = get_min_allowed_latconst(lattice, cutoff)

    if not init_guess_too_large:
        # Start at init_guess and increment up to max_init_guess before giving up
        done = False
        while not done:
            print(
                "\nAttempting to perform relaxation using initial lattice constant "
                "guess of {} Angstroms\n".format(init_guess, ".1f")
            )
            try:
                aopt_arr, eopt, iterations, funccalls, warnflag = opt.fmin(
                    energy,
                    init_guess,
                    args=(atoms, cell, positions, isolated_energy_per_atom[symbol],
                        min_allowed_latconst),
                    maxiter=maxiter,
                    full_output=True,
                    xtol=1e-8,
                )
                done = True
            except Exception as e:
                print(
                    "\nFailed to perform relaxation using initial lattice constant "
                    "guess of {} Angstroms".format(init_guess, ".1f")
                )
                print("Exception message:\n{}\n".format(str(e)))
                init_guess += 0.5
                if init_guess > max_init_guess:
                    done = True
                    raise RuntimeError(
                        "Exceeded maximum of {} for initial lattice constant guess for "
                        "relaxation".format(max_init_guess, ".1f")
                    )

    else:
        print(
            "\nDetected that energy at a lattice constant of {} Angstroms was equal to "
            "the isolated energy (i.e. no interactions). Shrinking box until energy "
            "becomes non-trivial...\n".format(init_guess, ".1f")
        )
        # In this case, we construct a 2x2x2 cell, turn off periodic boundary
        # and start shrinking the box until we get non-zero energy (relative to the
        # isolated energy). The resulting lattice constant is then used as the initial
        # guess.
        ncells_per_side = 2
        dummylattice = create_cubic_lattice(
            symbol, lattice, a=init_guess, ncells_per_side=ncells_per_side
        )
        dummylattice.set_calculator(calc)
        kim_ase_utils.rescale_to_get_nonzero_energy(
            dummylattice, isolated_energy_per_atom, etol=1e-6
        )
        new_init_guess = dummylattice.get_cell()[0, 0] / ncells_per_side
        min_lattice_constant = 0.5 * new_init_guess
        decrement_value = 0.1 * new_init_guess

        # Try new_init_guess, 0.9*new_init_guess, 0.8*new_init_guess, ...,
        # min_lattice_constant before giving up
        done = False
        while not done:
            print(
                "\nAttempting to perform relaxation using initial lattice constant "
                "guess of {} Angstroms\n".format(new_init_guess, ".4f")
            )
            try:
                aopt_arr, eopt, iterations, funccalls, warnflag = opt.fmin(
                    energy,
                    new_init_guess,
                    args=(atoms, cell, positions, isolated_energy_per_atom[symbol],
                        min_allowed_latconst),
                    maxiter=maxiter,
                    full_output=True,
                    xtol=1e-8,
                )
                done = True
            except:  # noqa: E722
                print(
                    "\nFailed to perform relaxation using initial lattice constant "
                    "guess of {} Angstroms\n".format(init_guess, ".4f")
                )
                new_init_guess -= decrement_value
                if new_init_guess < min_lattice_constant:
                    done = True
                    raise RuntimeError(
                        "Fell below minimum of {} for initial lattice constant guess "
                        "for relaxation.".format(min_lattice_constant, ".1f")
                    )

    # Check to see if the lattice constant we converged to is actually just
    # the one large enough so that atoms are not interacting, i.e. this
    # lattice is unstable
    if abs(eopt) < 1e-8:
        raise RuntimeError(
            "Calculation converged to lattice constant ({}) for which atoms have the "
            "same energy as in isolation. This may indicate the lattice is unstable."
            "".format(aopt_arr[0])
        )

    info = {
        "iterations": iterations,
        "func_calls": funccalls,
        "warnflag": warnflag,
        "repeat": repeat,
    }
    return aopt_arr[0], -eopt / num_atoms, info


if __name__ == "__main__":
    # grab from stdin (or a file)
    symbol = input("element:\n")
    print(symbol)
    lattice = input("lattice type:\n")
    print(lattice)
    model = input("modelname:\n")
    print(model)

    aopt, eopt, info = get_lattice_constant(
        model=model, symbol=symbol, lattice=lattice, maxiter=500
    )

    space_groups = {"fcc": "Fm-3m", "bcc": "Im-3m", "sc": "Pm-3m", "diamond": "Fd-3m"}
    wyckoff_codes = {"fcc": "4a", "bcc": "2a", "sc": "1a", "diamond": "8a"}

    dummylattice = create_cubic_lattice(symbol, lattice, a=1, ncells_per_side=1)
    normed_basis = {
        lattice: json.dumps(dummylattice.positions.tolist(), separators=(" ", " "))
        for lattice in space_groups.keys()
    }
    count = dummylattice.positions.shape[0]  # count atoms in a unit cell
    # pack the results in a dictionary
    results = {
        "lattice_constant": aopt,
        "cohesive_energy": eopt,
        "element": symbol,
        "species": '" "'.join([symbol] * count),  # repeat symbols to match normed_basis
        "crystal_structure": lattice,
        "space_group": space_groups[lattice],
        "wyckoff_code": wyckoff_codes[lattice],
        "basis_atoms": normed_basis[lattice],
    }
    results.update(info)
    print(results)
    template_environment = jinja2.Environment(
        loader=jinja2.FileSystemLoader("/"),
        block_start_string="@[",
        block_end_string="]@",
        variable_start_string="@<",
        variable_end_string=">@",
        comment_start_string="@#",
        comment_end_string="#@",
        undefined=jinja2.StrictUndefined,
    )

    # template the EDN output
    with open(os.path.abspath("output/results.edn"), "w") as f:
        template = template_environment.get_template(os.path.abspath("results.edn.tpl"))
        f.write(template.render(**results))