#!/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
#
################################################################################

# The docstring below is vc_description
"""Determines whether a model has a continuous energy and first
derivative, i.e. belongs to the C^1 continuity class, for all possible dimers.
For a model supporting N species, there are N + N!/(2(N-2)! distinct dimers for
all possible species combinations.  (For example if N=3, there are 3+3!/2=6
dimers.  If the species are A,B,C, the 6 dimers are AA, BB, CC, AB, AC, BC.)
For each dimer, the equilibrium separation and cutoff are determined.
The continuity across the cutoff is assessed.  Then an analysis is performed
to detect any discontinuities from half the equilibrium distance to the cutoff.
Although the verification check only requires C^1 continuity to pass, continuity
up to 3rd order is checked and reported."""

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

try:
    input = raw_input
except NameError:
    pass

import itertools
import math

from ase.calculators.kim import KIM, get_model_supported_species
import kim_python_utils.ase as kim_ase_utils
import kim_python_utils.vc as kim_vc_utils
from ase import Atoms
import scipy.optimize as opt
import numpy as np
import numdifftools as nd
from numdifftools.step_generators import MaxStepGenerator

__version__ = "003"
__author__ = "Ellad Tadmor"


################################################################################
#
#   FUNCTIONS
#
################################################################################
def create_species_pairs_list(species):
    """
    Given a list of species, create all the monoatomic and biatomic dimer pairs.
    """
    species_pairs = []
    for sp in species:
        species_pairs.append((sp, sp))
    if len(species) > 1:
        species_pairs += list(itertools.combinations(species, 2))
    return species_pairs


################################################################################
def energy(a, dimer, large_cell_len, einf):
    """
    Returns the energy of a dimer with a separation 'a', subtracting off the
    total energy of the dimer at very large separation.
    """
    dimer.set_positions(get_dimer_positions(a, large_cell_len))
    return dimer.get_potential_energy() - einf


################################################################################
def energy_cheat(a, dimer, large_cell_len, offset, einf):
    """
    Returns the energy of a dimer with a separation 'a' plus a small 'offset'
    to ensure that the root bisection has a positive number past the cutoff.
    """
    dimer.set_positions(get_dimer_positions(a, large_cell_len))
    return (dimer.get_potential_energy() - einf) + offset


################################################################################
def get_dimer_positions(a, large_cell_len):
    """
    Generate positions for a dimer of length 'a' centered in a finite simulation box
    with side length 'large_cell_len'
    """
    half_cell = 0.5 * large_cell_len
    positions = [
        [half_cell - 0.5 * a, half_cell, half_cell],
        [half_cell + 0.5 * a, half_cell, half_cell],
    ]
    return positions


################################################################################
def get_equilibrium_separation(dimer, large_cell_len, init_guess, einf):
    """
    Compute the equilibrium separation of a dimer (2-atom chain).
    """
    aopt_arr, eopt, iterations, funccalls, warnflag = opt.fmin(
        energy,
        init_guess,
        args=(dimer, large_cell_len, einf),
        full_output=True,
        xtol=1e-8,
        disp=False,
        maxiter=1000,
    )
    info = {"iterations": iterations, "func_calls": funccalls, "warnflag": warnflag}
    return aopt_arr[0], eopt, info


################################################################################
def get_cutoff_radius(dimer, a, b, offset, large_cell_len, einf):
    """
    Compute the cutoff radius (distance at which energy of dimer is zero)
    """
    x0, results = opt.bisect(
        energy_cheat,
        a,
        b,
        args=(dimer, large_cell_len, offset, einf),
        full_output=True,
        xtol=1e-8,
        maxiter=1000,
    )
    return x0, results


################################################################################
def do_vc(model, vc):
    """
    Perform Dimer Continuity C1 Verification Check
    """
    # Initialize
    max_deriv = 3  # largest derivative to be investigated
    # (NOTE: If this is increased, increase nth below)

    continuous = [True] * (1 + max_deriv)  # assume function and all
    # derivs are continuous

    etol_fine = 1e-15  # This tolerance is used to:
    # (1) Determine if a given species pair has energy interactions by checking
    #     if the bond length of a dimer can be reduced so that the total energy
    #     differs from the sum of the corresponding isolated atomic energies for
    #     each species by more than this value.
    # (2) Determine an upper bracket of the search range when attempting to
    #     calculate the cutoff radius.  In order to ensure that a given dimer
    #     separation is above the cutoff, its total energy is compared to the
    #     sum of the corresponding isolated energies; if the absolute value of
    #     the difference is less than this amount, we assert that the atoms are
    #     no longer interacting.

    etol_coarse = 1e-6  # This tolerance is used to:
    # (1) Determine the lower bracket of the search range when attempting to
    #     calculate the cutoff radius by checking if the total energy differs
    #     from the sum of the corresponding isolated atomic energies by more
    #     than this value.
    # (2) Judge whether the energy and its derivatives are truly continuous at
    #     the cutoff.  The energy must be within this value of the sum of the
    #     corresponding isolated atomic energies; the absolute value of the
    #     derivatives must not exceed this value.

    led_tolerance = 1.0  # value considered an indicator of a
    # discontinuity for the local edge detection
    # algorithm

    # Side length of large finite simulation domain
    large_cell_len = 50

    got_atleast_one_set_of_results = False

    nth = {0: "Energy", 1: "1st-Deriv", 2: "2nd-Deriv", 3: "3rd-Deriv"}
    yesno = {True: "yes", False: "** NO **"}

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

    # Print information specific to this VC
    dashwidth = 101
    vc.rwrite("")
    vc.rwrite("-" * dashwidth)
    vc.rwrite("Results for KIM Model      : %s" % model.strip())
    vc.rwrite("Supported species          : %s" % " ".join(species))

    # Get isolated energy of each species for this model
    isolated_energy_per_atom = {}
    for spec in species:
        isolated_energy_per_atom[spec] = kim_ase_utils.get_isolated_energy_per_atom(
            model, spec
        )

    # Create list of dimers species
    species_pairs = create_species_pairs_list(species)

    # Loop over all dimers and determine the smoothness of each
    for i in range(0, len(species_pairs)):

        atoms_interacting_energy = kim_ase_utils.check_if_atoms_interacting(
            model, symbols=species_pairs[i], check_force=False, etol=etol_fine
        )
        if not atoms_interacting_energy:
            vc.rwrite("")
            vc.rwrite(
                "WARNING: The model provided, {}, does not possess a non-trivial energy "
                "interaction for species {} as required by this Verification "
                "Check. Skipping...".format(model, species_pairs[i])
            )
            vc.rwrite("")
            continue

        # Infinite separation energy for this species pair
        einf = (
            isolated_energy_per_atom[species_pairs[i][0]]
            + isolated_energy_per_atom[species_pairs[i][1]]
        )

        # Print header for this dimer
        vc.rwrite("")
        vc.rwrite("-" * dashwidth)
        vc.rwrite("DIMER {0}--{1}".format(*species_pairs[i]))
        vc.rwrite("-" * dashwidth)

        # Build ASE Atoms object for the dimer
        a = 1.0
        dimer = Atoms(
            "".join(species_pairs[i]),
            positions=get_dimer_positions(a, large_cell_len),
            cell=(large_cell_len, large_cell_len, large_cell_len),
            pbc=(False, False, False),
        )

        # Create calculator
        calc = KIM(model)
        dimer.set_calculator(calc)

        vc.rwrite("")
        vc.rwrite("E n e r g y  E x t r e m a")

        # Find equilibrium separation for this dimer
        max_equilibrium_separation = 10.0
        got_equil_sep = False
        while not got_equil_sep:
            try:
                aopt, eopt, info = get_equilibrium_separation(
                    dimer, large_cell_len, init_guess=1.0, einf=einf
                )
                got_equil_sep = True
            except:  # noqa: E722
                a += 0.25
                if a > max_equilibrium_separation:
                    raise RuntimeError(
                        "Cannot find a working "
                        "configuration within a reasonable dimer "
                        "separation range."
                    )
                dimer[1].x = a
        warn = "none"
        if info["warnflag"]:
            warn = "** DID NOT CONVERGE **"
        vc.rwrite(
            "{0:25}   {1:>15}   {2:>15}   {3:>9}   {4:>9}    {5}".format(
                "", "distance", "energy", "#iter", "#fn-calls", "warnings"
            )
        )
        vc.rwrite(
            "{0:25}   {1: 11.8e}   {2: 11.8e}   {3:9d}   {4:9d}    {5}".format(
                "Equilibrium separation:",
                aopt,
                eopt,
                info["iterations"],
                info["func_calls"],
                warn,
            )
        )
        if info["warnflag"]:
            vc.rwrite("")
            vc.rwrite(
                "WARNING: NOT CHECKING CONTINUITY FOR "
                + "-".join(species_pairs[i])
                + " DIMER."
            )
            vc.rwrite("")

            # Explicitly close calculator to ensure any allocated memory
            # is freed (relevant for SMs)
            if hasattr(calc, "__del__"):
                calc.__del__()

            continue  # failed to converge, so skip this dimer

        # Check if this dimer collapsed while trying to find the equilibrium distance
        print_aopt_warning = False
        if aopt < np.finfo(float).eps:
            aopt = abs(aopt)  # In case aopt is a very small negative number
            print_aopt_warning = True

        # Find dimer cutoff radius
        #
        # Step 1: Find upper end of bracket delimiting cutoff
        #         Start from equilibrium separation, increase separation until
        #         energy reaches infinite separation energy.
        warn = "none"
        failed = False
        max_upper_cutoff_bracket = 2 * max_equilibrium_separation
        b = aopt
        db = max(aopt, 0.5)
        still_interacting = True
        while still_interacting and not failed:
            b += db
            if b > max_upper_cutoff_bracket:
                failed = True
                warn = "** Unable to obtain upper bracket for cutoff radius **"
            else:
                eb = energy(b, dimer, large_cell_len, einf)
                if abs(eb) < etol_fine:
                    still_interacting = False

        #
        # Step 2: Find lower end of bracket delimiting cutoff
        #         Start from upper end of bracket, and reduce separation until
        #         energy changes by a discernable amount.
        if not failed:
            a = b
            da = 0.01
            not_interacting = True
            while not_interacting:
                a -= da
                if a < 0:
                    failed = True
                    warn = "** Unable to obtain lower bracket for cutoff radius **"
                else:
                    ea = energy(a, dimer, large_cell_len, einf)
                    if abs(ea) > etol_coarse:
                        not_interacting = False

        #
        # Step 3: Use bisection algorithm to get cutoff radius
        #
        #         NOTE: Some Simulator Models have a history dependence due to
        #         them maintaining charges from the previous energy evaluation
        #         to use as an initial guess for the next charge equilibration.
        #         We therefore have to treat them not as single-valued
        #         functions but as distributions, i.e. for a given
        #         configuration you might get any of a range of energy values
        #         depending on the history of your previous energy evaluations.
        #         This is particularly problematic for this step, where we set
        #         up a bisection problem in order to determine the cutoff
        #         radius of the model.  Our solution for this specific case is
        #         to make a very crude estimate of the variance of that
        #         distribution with a 10% factor of safety on it.
        ran_bisection = False
        eb_new = energy(b, dimer, large_cell_len, einf)
        eb_error = abs(eb_new - eb)
        if not failed:
            # compute offset to ensure that energy before and after cutoff have
            # different signs
            if ea < eb:
                offset = -eb + 1.1 * eb_error + np.finfo(float).eps
            else:
                offset = -eb - 1.1 * eb_error - np.finfo(float).eps

            rcut, results = get_cutoff_radius(dimer, a, b, offset, large_cell_len, einf)
            ran_bisection = True
            if not results.converged:
                failed = True
                warn = "** DID NOT CONVERGE **"

        #
        # Step 4: Report to user
        if ran_bisection:
            vc.rwrite(
                "{0:25}   {1: 11.8e}   {2: 11.8e}   {3:9d}   {4:9d}    {5}".format(
                    "Cutoff separation:",
                    rcut,
                    energy(rcut, dimer, large_cell_len, einf),
                    results.iterations,
                    results.function_calls,
                    warn,
                )
            )
        else:
            vc.rwrite(warn)

        if print_aopt_warning:
            vc.rwrite("")
            vc.rwrite(
                "WARNING: Equilibrium distance ({} Angstroms) is smaller than "
                "machine epsilon ({}).".format(aopt, np.finfo(float).eps)
            )

        if failed:
            vc.rwrite("")
            vc.rwrite(
                "WARNING: NOT CHECKING CONTINUITY FOR "
                + "-".join(species_pairs[i])
                + " DIMER."
            )
            vc.rwrite("")

            # Explicitly close calculator to ensure any allocated memory
            # (particularly for Simulator Models) is freed
            if hasattr(calc, "__del__"):
                calc.__del__()
            continue  # failed to converge, so skip this dimer

        vc.rwrite("")
        vc.rwrite("C u t o f f  S m o o t h n e s s")
        vc.rwrite("")

        vc.rwrite("{0:10}   {1:>15}   {2:10}".format("", "value", "continuous"))

        # Check smoothness at cutoff
        sg = MaxStepGenerator(
            base_step=0.5 * (rcut - aopt),
            num_steps=14,
            use_exact_steps=True,
            step_ratio=1.6,
            offset=0,
        )
        for n in range(0, max_deriv + 1):
            Denergy = nd.Derivative(
                energy, step=sg, full_output=True, method="backward", n=n
            )
            val, info = Denergy(rcut, dimer, large_cell_len, einf)
            is_continuous = abs(val) <= etol_coarse
            if not is_continuous:
                continuous[n] = False
            vc.rwrite(
                "{0:10}   {1: 11.8e}   {2:>10}".format(
                    nth[n], np.asscalar(val), yesno[is_continuous]
                )
            )

        vc.rwrite("")
        vc.rwrite("C o n t i n u i t y")
        vc.rwrite("")

        # Set the range and increment for exploring internal discontinuities
        if aopt < rcut / 1000.0:
            # Start with amin=rcut/5 and keep increasing until the energy can be
            # evaluated until amin=rcut/2
            afrac = 0.6
            max_lower_limit = rcut / 2

        else:
            # Start with amin=aopt/2 and keep increasing until the energy can be
            # evaluated until amin=aopt
            afrac = 0.5
            max_lower_limit = aopt

        got_amin = False
        while not got_amin:
            try:
                amin = (1.0 - afrac) * max_lower_limit
                dimer[1].x = amin
                dimer.get_potential_energy()  # See that energy can be computed without error.
                got_amin = True
            except:  # noqa: E722
                afrac *= 0.5

        amax = rcut
        del_a = 0.01
        na = int(math.ceil((amax - amin) / del_a))
        dimer[1].x = amin
        vc.rwrite(
            "Checking continuity for r = [{0:.5f},{1:.5f}] at {2:d} points (Delta r = {3:.5f})".format(
                amin, amax, na, del_a
            )
        )

        vc.rwrite("")
        vc.rwrite(
            "Local edge detection based on a normalized 5th-order local difference formula T^5"
        )
        vc.rwrite(
            "is used to determine the presence of discontinuities. The tolerance is |T^5|>{0:.5f}.".format(
                led_tolerance
            )
        )
        vc.rwrite(
            "(For details see Anne Gelb and Eitan Tadmor, J. Sci. Comp., 28:279-306, 2006.)"
        )
        vc.rwrite("")

        for n in range(0, max_deriv + 1):
            # set numerical derivative to constant step to prevent the algorithm
            # going haywire in some cases.
            Denergy = nd.Derivative(energy, full_output=True, n=n, step=0.1 * del_a)
            if n == 0:
                aux_file = "dimer-energy-" + "".join(species_pairs[i]) + ".dat"
            else:
                aux_file = (
                    "dimer-energy-deriv-"
                    + str(n)
                    + "-"
                    + "".join(species_pairs[i])
                    + ".dat"
                )
            vc.vc_files.append(aux_file)
            vc.rwrite("")
            vc.rwrite(
                'Checking {0} ({1} vs. distance in file "{2}")'.format(
                    nth[n], nth[n].lower(), aux_file
                )
            )

            # Generate energy curve
            r = []
            e = []
            for j in range(0, na + 1):
                a = amin + j * del_a
                val, info = Denergy(a, dimer, large_cell_len, einf)
                r.append(a)
                e.append(np.asscalar(val))
            vc.write_aux_x_y(aux_file, r, e)

            # Apply local edge detection algorithm to identify discontinuities
            # Based on: A. Gelb and E. Tadmor, J. Sci. Comp., 28:279-306, 2006.
            fact = 1.0 / 6.0
            is_continuous = True
            for j in range(2, na - 3):
                # use 5-th order local difference formula
                led = fact * (
                    -e[j - 2]
                    + 5 * e[j - 1]
                    - 10 * e[j]
                    + 10 * e[j + 1]
                    - 5 * e[j + 2]
                    + e[j + 3]
                )
                if abs(led) > led_tolerance:
                    continuous[n] = False
                    is_continuous = False
                    vc.rwrite(
                        "==> Suspected discontinuity encountered at r={:11.8e} (|T^5| = {:11.8e})".format(
                            r[j], abs(led)
                        )
                    )
            if is_continuous:
                vc.rwrite("... No discontinuities found.")

        # Explicitly close calculator to ensure any allocated memory
        # (particularly for leaky Simulator Models) is freed
        if hasattr(calc, "__del__"):
            calc.__del__()

        got_atleast_one_set_of_results = True

    if got_atleast_one_set_of_results:
        # Summary of results
        vc.rwrite("")
        vc.rwrite("=" * dashwidth)
        vc.rwrite("")
        vc.rwrite("SUMMARY of Model Continuity Results Across All Dimers:")
        vc.rwrite("")

        vc.rwrite("{0:10}   {1:10}".format("", "continuous"))
        vc.rwrite("-" * 29)
        for n in range(0, max_deriv + 1):
            vc.rwrite("{0:10}   {1:>10}".format(nth[n], yesno[continuous[n]]))

        # Determine continuity class and write out properties
        k = -1
        n = 0
        while n <= max_deriv and continuous[n]:
            k = n
            n += 1
        vc_comment = "The model is C^{0:d} continuous. ".format(k)
        if k == -1:
            vc_comment += "This means that the model has discontinuous energy."
        if k == 0:
            vc_comment += "This means that the model has continuous energy, but a discontinuous first derivative."
        if k == 1:
            vc_comment += "This means that the model has continuous energy and continuous first derivative."
        if k == 2:
            vc_comment += "This means that the model has continuous energy and continuous derivatives up to order 2."
        if k == 3:
            vc_comment += "This means that the model has continuous energy and continuous derivatives at least up to order 3. (Derivatives beyond this order were not tested.)"

        vc.rwrite("")
        vc.rwrite("=" * dashwidth)
        vc.rwrite("Continuity must be C^1 or higher to pass this verification check.")
        vc.rwrite("")
        if k >= 1:
            vc_grade = "P"
        else:
            vc_grade = "F"

        return vc_grade, vc_comment

    else:
        msg = (
            "ERROR: Failed to find any pairs in the species set supported "
            "by the model that have a non-zero energy interaction. This "
            "likely means that the model only computes forces and does "
            "not compute energy; this Verification Check currently "
            "requires that models report energy."
        )
        vc.rwrite("")
        vc.rwrite(msg)
        vc.rwrite("")
        raise RuntimeError(msg)


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

    vcargs = {
        "vc_name": "vc-dimer-continuity-c1",
        "vc_author": __author__,
        "vc_description": kim_vc_utils.vc_stripall(__doc__),
        "vc_category": "informational",
        "vc_grade_basis": "passfail",
        "vc_files": [],
        "vc_debug": False,  # Set to True to get exception traceback info
    }

    # Get the model extended KIM ID:
    model = input("Model Extended KIM ID = ")

    # Execute VC
    kim_vc_utils.setup_and_run_vc(do_vc, model, **vcargs)