#!/usr/bin/env python3
################################################################################
#
#  CDDL HEADER START
#
#  The contents of this file are subject to the terms of the Common Development
#  and Distribution License Version 1.0 (the "License").
#
#  You can obtain a copy of the license at
#  http:# www.opensource.org/licenses/CDDL-1.0.  See the License for the
#  specific language governing permissions and limitations under the License.
#
#  When distributing Covered Code, include this CDDL HEADER in each file and
#  include the License file in a prominent location with the name LICENSE.CDDL.
#  If applicable, add the following below this CDDL HEADER, with the fields
#  enclosed by brackets "[]" replaced with your own identifying information:
#
#  Portions Copyright (c) [yyyy] [name of copyright owner]. All rights reserved.
#
#  CDDL HEADER END
#
#  Copyright (c) 2019, Regents of the University of Minnesota.
#  All rights reserved.
#
#  Contributor(s):
#     Ellad B. Tadmor
#
################################################################################

"""
Do a series of calculations while valgrind is looking.
"""

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

try:
    input = raw_input
except NameError:
    pass

from ase.lattice.cubic import FaceCenteredCubic
from ase.calculators.kim import KIM, get_model_supported_species
import kim_python_utils.ase as kim_ase_utils
import random
import sys

__author__ = "Ellad Tadmor"


################################################################################
#
#   FUNCTIONS
#
################################################################################

################################################################################
def _compute_energy_and_forces(vc, atoms, heading, dashwidth, verbose):
    """
    Compute and print energy and forces of ASE object in 'atoms'
    """
    # compute analytical forces (negative gradient of cohesive energy)
    energy = atoms.get_potential_energy()
    forces = atoms.get_forces()

    if verbose:
        vc.rwrite(heading)
        vc.rwrite("-" * dashwidth)
        vc.rwrite("Energy = {}".format(energy))
        vc.rwrite("")
        vc.rwrite("Forces:")
        for i in range(len(forces)):
            vc.rwrite(
                "{:>3d}  {: .8e}  {: .8e}  {: .8e}".format(
                    i + 1, forces[i, 0], forces[i, 1], forces[i, 2]
                )
            )
        vc.rwrite("")


################################################################################
def memory_leaks_test_process(model, vc, verbose=True):
    """
    Do a series of calculations based on the numerical derivative verification
    check that will be used by valgrind to look for memory leaks.
    """
    # Max iterations allowed for some of the while loops below
    max_iters = 2000

    # Get supported species
    species = get_model_supported_species(model)
    species = kim_ase_utils.remove_species_not_supported_by_ASE(list(species))
    species.sort()

    # Basic cell parameters
    lattice_constant_orig = 3.0
    pert_amp_orig = 0.1 * lattice_constant_orig
    ncells_per_side = 2
    seed = 13
    random.seed(seed)

    # Finite domain in which to embed the finite cluster of atoms we'll translate and
    # invert
    large_cell_len = 7 * lattice_constant_orig * ncells_per_side

    # Print VC info
    dashwidth = 101
    if verbose:
        vc.rwrite("")
        vc.rwrite("-" * dashwidth)
        vc.rwrite("Results for KIM Model      : %s" % model.strip())
        vc.rwrite("Supported species          : %s" % " ".join(species))
        vc.rwrite("")
        vc.rwrite("random seed                = %d" % seed)
        vc.rwrite("lattice constant (orig)    = %0.3f" % lattice_constant_orig)
        vc.rwrite("perturbation amplitude     = %0.3f" % pert_amp_orig)
        vc.rwrite("number unit cells per side = %d" % ncells_per_side)
        vc.rwrite("-" * dashwidth)
        vc.rwrite("")

    # Heading format
    if verbose:
        headfmt = (
            "{0} STRUCTURE (pbc={1})-- Species = {2}"
            + '   (Configuration in file "{3}")'
        )

    # Perform memory leak check for monotatomic systems
    for spec in species:
        calc = KIM(model)
        lattice_constant = lattice_constant_orig
        got_initial_config = False
        while not got_initial_config:
            atoms = FaceCenteredCubic(
                size=(ncells_per_side, ncells_per_side, ncells_per_side),
                latticeconstant=lattice_constant,
                symbol=spec,
                pbc=False,
            )

            # Move our finite cluster of atoms to the center of our large cell
            atoms.set_cell([large_cell_len, large_cell_len, large_cell_len])
            trans = [0.5 * large_cell_len] * 3
            atoms.translate(trans)

            atoms.set_calculator(calc)
            try:
                kim_ase_utils.rescale_to_get_nonzero_forces(atoms, 0.01)
                got_initial_config = True
            except kim_ase_utils.KIMASEError:
                # Routine failed in on recoverable manner
                raise  # re-raise same exception
            except Exception:
                # Initial config failed. This most likely due to an evaluation
                # outside the legal model range. Increase lattice constant and
                # try again.
                lattice_constant += 0.25
                if lattice_constant > 10.0:
                    raise RuntimeError(
                        "Cannot find a working configuration within a reasonable lattice constant range."
                    )

        # Randomize positions
        save_positions = atoms.get_positions()
        pert_amp = pert_amp_orig
        got_randomized_config = False
        iters = 0
        while not got_randomized_config:
            try:
                kim_ase_utils.randomize_positions(atoms, pert_amp)
                atoms.get_forces()  # make sure forces can be computed
                got_randomized_config = True
            except:  # noqa: E722
                # Failed to compute forces; reset to original posns and retry
                atoms.set_positions(save_positions)
                pert_amp *= 0.5  # cut perturbation amplitude by half
                iters += 1
                if iters >= max_iters:
                    raise RuntimeError(
                        "Iteration limit exceeded when randomizing positions "
                        "during check for species {}".format(spec)
                    )

        kim_ase_utils.perturb_until_all_forces_sizeable(atoms, pert_amp)
        for pbc in [False, True]:
            heading = ""
            if verbose:
                aux_file = "config-" + str(pbc)[0] + "-" + spec + ".xyz"
                vc.vc_files.append(aux_file)
                vc.write_aux_ase_atoms(aux_file, atoms, "xyz")
                heading = headfmt.format("MONOATOMIC", pbc, spec, aux_file)
            if pbc:
                atoms.set_pbc(pbc)  # change periodic BCs to True
            try:
                _compute_energy_and_forces(vc, atoms, heading, dashwidth, verbose)
            except:  # noqa: E722
                pass

    # Perform memory leak check for mixed system
    lattice_constant = lattice_constant_orig
    if len(species) > 1:
        while True:
            atoms = FaceCenteredCubic(
                size=(ncells_per_side, ncells_per_side, ncells_per_side),
                latticeconstant=lattice_constant,
                symbol="H",
                pbc=False,
            )
            if len(atoms) < len(species):
                ncells_per_side += 1
            else:
                break
        kim_ase_utils.randomize_species(atoms, species)
        calc = KIM(model)
        atoms.set_calculator(calc)
        got_initial_config = False
        while not got_initial_config:
            try:
                kim_ase_utils.rescale_to_get_nonzero_forces(atoms, 0.01)
                got_initial_config = True
            except kim_ase_utils.KIMASEError:
                # Routine failed in on recoverable manner
                raise  # re-raise same exception
            except Exception:
                # Initial config failed. This most likely due to an evaluation
                # outside the legal model range. Increase lattice constant and
                # try again.
                lattice_constant += 0.25
                if lattice_constant > 10.0:
                    raise RuntimeError(
                        "Cannot find a working configuration within # a reasonable lattice constant range."
                    )
                acell = lattice_constant * ncells_per_side
                atoms.set_cell([acell, acell, acell], scale_atoms=True)

        # Randomize positions
        save_positions = atoms.get_positions()
        pert_amp = pert_amp_orig
        got_randomized_config = False
        iters = 0
        while not got_randomized_config:
            try:
                kim_ase_utils.randomize_positions(atoms, pert_amp)
                atoms.get_forces()  # make sure forces can be computed
                got_randomized_config = True
            except:  # noqa: E722
                # Failed to compute forces; reset to original posns and retry
                atoms.set_positions(save_positions)
                pert_amp *= 0.5  # cut perturbation amplitude by half
                iters += 1
                if iters >= max_iters:
                    raise RuntimeError(
                        "Iteration limit exceeded when randomizing positions "
                        "during check for species {}".format(spec)
                    )

        kim_ase_utils.perturb_until_all_forces_sizeable(atoms, pert_amp)
        for pbc in [False, True]:
            heading = ""
            if verbose:
                aux_file = "config-" + str(pbc)[0] + "-" + "".join(species) + ".xyz"
                vc.vc_files.append(aux_file)
                vc.write_aux_ase_atoms(aux_file, atoms, "xyz")
                heading = headfmt.format("MIXED", pbc, " ".join(species), aux_file)
            if pbc:
                atoms.set_pbc(pbc)  # change periodic BCs to True
            try:
                _compute_energy_and_forces(vc, atoms, heading, dashwidth, verbose)
            except:  # noqa: E722
                pass


################################################################################
#
#   MAIN PROGRAM
#
###############################################################################
if __name__ == "__main__":

    # Get KIM model from argument stream
    model = sys.argv[1]

    # Run memory leak test process in silent mode
    memory_leaks_test_process(model, "", verbose=False)