# fmt: off
import numpy as np
from pytest import mark

from ase import Atoms


def squeeze_dimer(atoms, d):
    """Squeeze the atoms together by the absolute distance ``d`` (Angstroms)
    """
    pos = atoms.get_positions()
    pos[0] += np.asarray([d, 0, 0])
    atoms.set_positions(pos)


def set_positions_to_orig(atoms, box_len, dimer_separation):
    pos0 = 0.5 * np.full(3, 0.5 * box_len)
    displacement = np.array([0.5 * dimer_separation, 0, 0])

    pos1 = pos0 - displacement
    pos2 = pos0 + displacement
    atoms.set_positions([pos1, pos2])


@mark.calculator_lite
def test_update_coords(KIM, testdir):
    """
    Check that the coordinates registered with the KIM API are updated
    appropriately when the atomic positions are updated.  This can go awry
    if the 'coords' attribute of the relevant NeighborList subclass is
    reassigned to a new memory location -- a problem which was indeed
    occurring at one point
    (see https://gitlab.com/ase/ase/merge_requests/1442)!
    """

    # We know that ex_model_Ar_P_Morse_07C has a cutoff of 8.15 Angstroms
    model = "ex_model_Ar_P_Morse_07C"
    model_cutoff = 8.15  # Angstroms

    # Create a dimer centered in a small box that's small enough so
    # that there will be padding atoms when we make it periodic
    box_len = 0.5 * model_cutoff
    dimer_separation = model_cutoff * 0.3

    atoms = Atoms("Ar" * 2, cell=[[box_len, 0, 0],
                  [0, box_len, 0], [0, 0, box_len]])

    # Create calculator.  Either the kimpy neighbor list library or
    # ASE's native neighbor lists should suffice to check this since
    # update_kim_coords() belongs to their parent class, NeighborList.
    # Here, we'll use the default mode (kimpy neighbor list).
    neigh_skin_ratio = 0.2
    skin = neigh_skin_ratio * model_cutoff
    calc = KIM(model, options={"neigh_skin_ratio": neigh_skin_ratio})
    atoms.calc = calc

    squeezed_energies_ref = {
        False: 5.784620078721877,  # finite
        True: 6.766293119162073,  # periodic
    }

    for pbc in [False, True]:

        # Reset dimer positions to original configuration
        set_positions_to_orig(atoms, box_len, dimer_separation)

        atoms.set_pbc(pbc)

        # Get potential energy so that it will get rid of "pbc" being
        # in the system_changes.  The only way that update_kim_coords
        # is called is when system_changes only contains "positions",
        # as otherwise a neighbor list rebuild is triggered.
        atoms.get_potential_energy()

        # First squeeze the dimer together by a distance less than the
        # skin and compute the energy.  This doesn't trigger a
        # neighbor list rebuild (which we avoid here because we're
        # only trying to test update_kim_coords, which is not called
        # in the event that a neighbor list rebuild is necessary).
        #
        # This will update the coordinate values in the ``coords``
        # attribute of the relevant NeighborList subclass instance
        # (see ase/calculators/kim/neighborlist.py) and, since
        # ``coords`` is pointing at the same memory the KIM API is
        # reading the coordinates from, the KIM Model will see the
        # updated coordinates with no problem.  *However*, it's
        # possible that after updating these values, the ``coords``
        # attribute is bound to a *new* object in memory if
        # update_kim_coords is broken somehow!
        squeeze_dimer(atoms, 0.2 * skin)
        atoms.get_potential_energy()

        # Now squeeze the dimer together by the same amount again,
        # where the sum of this squeeze and the last is still less
        # than the skin distance so that we still avoid a neighbor
        # list rebuild.  If all is well, the `coords` attribute of the
        # relevant NeighborList subclass still points at the same
        # location as before we did the previous squeeze.  Thus, when
        # we update the coordinates again, it should still be updating
        # the same coordinates being read by the KIM API, giving us
        # the expected value of the energy.  If update_kim_coords is
        # broken, ``coords`` will already be bound to a new object in
        # memory while the KIM API is still looking at its previous
        # address to read coordinates.  This means that when we update
        # the coordinates in the ``coords`` value, the KIM API never
        # sees them.  In fact, it's not even guaranteed that the KIM
        # API will read the same coordinates as before since the
        # memory where it's looking may have since been overwritten by
        # some other python objects.
        squeeze_dimer(atoms, 0.2 * skin)
        squeezed_energy = atoms.get_potential_energy()
        assert np.isclose(squeezed_energy, squeezed_energies_ref[pbc])
