from collections import defaultdict

import numpy as np
import kimpy
from kimpy import neighlist
from ase.neighborlist import neighbor_list
from ase import Atom

from .kimpy_wrappers import check_call_wrapper


class NeighborList:

    kimpy_arrays = {
        "num_particles": np.intc,
        "coords": np.double,
        "particle_contributing": np.intc,
        "species_code": np.intc,
        "cutoffs": np.double,
        "padding_image_of": np.intc,
        "need_neigh": np.intc,
    }

    def __setattr__(self, name, value):
        """
        Override assignment to any of the attributes listed in
        kimpy_arrays to automatically cast the object to a numpy array.
        This is done to avoid a ton of explicit numpy.array() calls (and
        the possibility that we forget to do the cast).  It is important
        to use np.asarray() here instead of np.array() because using the
        latter will mean that incrementation (+=) will create a new
        object that the reference is bound to, which becomes a problem
        if update_compute_args isn't called to reregister the
        corresponding address with the KIM API.
        """
        if name in self.kimpy_arrays and value is not None:
            value = np.asarray(value, dtype=self.kimpy_arrays[name])
        self.__dict__[name] = value

    def __init__(
        self,
        neigh_skin_ratio,
        model_influence_dist,
        model_cutoffs,
        padding_not_require_neigh,
        debug,
    ):

        self.skin = neigh_skin_ratio * model_influence_dist
        self.influence_dist = model_influence_dist + self.skin
        self.cutoffs = model_cutoffs + self.skin
        self.padding_need_neigh = not padding_not_require_neigh.all()
        self.debug = debug

        if self.debug:
            print()
            print("Calculator skin: {}".format(self.skin))
            print(f"Model influence distance: {model_influence_dist}")
            print(
                "Calculator influence distance (including skin distance): {}"
                "".format(self.influence_dist)
            )
            print("Number of cutoffs: {}".format(model_cutoffs.size))
            print("Model cutoffs: {}".format(model_cutoffs))
            print(
                "Calculator cutoffs (including skin distance): {}"
                "".format(self.cutoffs)
            )
            print(
                "Model needs neighbors of padding atoms: {}"
                "".format(self.padding_need_neigh)
            )
            print()

        # Attributes to be set by subclasses
        self.neigh = None
        self.num_contributing_particles = None
        self.padding_image_of = None
        self.num_particles = None
        self.coords = None
        self.particle_contributing = None
        self.species_code = None
        self.need_neigh = None
        self.last_update_positions = None

    def update_kim_coords(self, atoms):
        """Update atomic positions in self.coords, which is where the KIM
        API will look to find them in order to pass them to the model.
        """
        if self.padding_image_of.size != 0:
            disp_contrib = atoms.positions - self.coords[: len(atoms)]
            disp_pad = disp_contrib[self.padding_image_of]
            self.coords += np.concatenate((disp_contrib, disp_pad))
        else:
            np.copyto(self.coords, atoms.positions)

        if self.debug:
            print("Debug: called update_kim_coords")
            print()

    def need_neigh_update(self, atoms, system_changes):
        need_neigh_update = True
        if len(system_changes) == 1 and "positions" in system_changes:
            # Only position changes
            if self.last_update_positions is not None:
                a = self.last_update_positions
                b = atoms.positions
                if a.shape == b.shape:
                    delta = np.linalg.norm(a - b, axis=1)
                    # Indices of the two largest elements
                    ind = np.argpartition(delta, -2)[-2:]
                    if sum(delta[ind]) <= self.skin:
                        need_neigh_update = False

        return need_neigh_update

    def clean(self):
        pass


class ASENeighborList(NeighborList):
    def __init__(
        self,
        compute_args,
        neigh_skin_ratio,
        model_influence_dist,
        model_cutoffs,
        padding_not_require_neigh,
        debug,
    ):
        super().__init__(
            neigh_skin_ratio,
            model_influence_dist,
            model_cutoffs,
            padding_not_require_neigh,
            debug,
        )

        self.neigh = {}
        compute_args.set_callback(
            kimpy.compute_callback_name.GetNeighborList,
            self.get_neigh,
            self.neigh)

    @staticmethod
    def get_neigh(data, cutoffs, neighbor_list_index, particle_number):
        """Retrieves the neighbors of each atom using ASE's native neighbor
        list library
        """
        # We can only return neighbors of particles that were stored
        number_of_particles = data["num_particles"]
        if particle_number >= number_of_particles or particle_number < 0:
            return (np.array([]), 1)

        neighbors = data["neighbors"][neighbor_list_index][particle_number]
        return (neighbors, 0)

    def build(self, orig_atoms):
        """Build the ASE neighbor list and return an Atoms object with
        all of the neighbors added.  First a neighbor list is created
        from ase.neighbor_list, having only information about the
        neighbors of the original atoms.  If neighbors of padding atoms
        are required, they are calculated using information from the
        first neighbor list.
        """
        syms = orig_atoms.get_chemical_symbols()
        orig_num_atoms = len(orig_atoms)
        orig_pos = orig_atoms.get_positions()

        # New atoms object that will contain the contributing atoms plus
        # the padding atoms
        new_atoms = orig_atoms.copy()

        neigh_list = defaultdict(list)
        neigh_dists = defaultdict(list)

        # Information for padding atoms
        padding_image_of = []
        padding_shifts = []

        # Ask ASE to build the neighbor list using the influence distance.
        # This is not a neighbor list for each atom, but rather a listing
        # of all neighbor pairs that exist.  Atoms with no neighbors will
        # not show up.
        (neigh_indices_i,
         neigh_indices_j,
         relative_pos,
         neigh_cell_offsets,
         dists) = neighbor_list(
            "ijDSd", orig_atoms, self.influence_dist
        )

        # Loop over all neighbor pairs. Because this loop will generally
        # include image atoms (for periodic systems), we keep track of
        # which atoms/images we've accounted for in the `used` dictionary.
        used = dict()
        for neigh_i, neigh_j, rel_pos, offset, dist in zip(neigh_indices_i,
                                                           neigh_indices_j,
                                                           relative_pos,
                                                           neigh_cell_offsets,
                                                           dists):
            # Get neighbor position of neighbor
            # (mapped back into unit cell, so this may overlap with other atoms)
            wrapped_pos = orig_pos[neigh_i] + rel_pos

            shift = tuple(offset)
            uniq_index = (neigh_j,) + shift
            if shift == (0, 0, 0):
                # This atom is in the unit cell, i.e. it is contributing
                neigh_list[neigh_i].append(neigh_j)
                neigh_dists[neigh_i].append(dist)
                if uniq_index not in used:
                    used[uniq_index] = neigh_j
            else:
                # This atom is not in the unit cell, i.e. it is padding
                if uniq_index not in used:
                    # Add the neighbor as a padding atom
                    used[uniq_index] = len(new_atoms)
                    new_atoms.append(Atom(syms[neigh_j], position=wrapped_pos))
                    padding_image_of.append(neigh_j)
                    padding_shifts.append(offset)
                neigh_list[neigh_i].append(used[uniq_index])
                neigh_dists[neigh_i].append(dist)
        neighbor_list_size = orig_num_atoms

        # Add neighbors of padding atoms if the potential requires them
        if self.padding_need_neigh:
            neighbor_list_size = len(new_atoms)
            inv_used = dict((v, k) for k, v in used.items())
            # Loop over all the neighbors (k) and the image of that neighbor
            # in the cell (neigh)
            for k, neigh in enumerate(padding_image_of):
                # Shift from original atom in cell to neighbor
                shift = padding_shifts[k]
                for orig_neigh, orig_dist in zip(neigh_list[neigh],
                                                 neigh_dists[neigh]):
                    # Get the shift of the neighbor of the original atom
                    orig_shift = inv_used[orig_neigh][1:]

                    # Apply sum of original shift and current shift to
                    # neighbors of original atom
                    total_shift = orig_shift + shift

                    # Get the image in the cell of the original neighbor
                    if orig_neigh <= orig_num_atoms - 1:
                        orig_neigh_image = orig_neigh
                    else:
                        orig_neigh_image = padding_image_of[orig_neigh -
                                                            orig_num_atoms]

                    # If the original image with the total shift has been
                    # used before then it is also a neighbor of this atom
                    uniq_index = (orig_neigh_image,) + tuple(total_shift)
                    if uniq_index in used:
                        neigh_list[k + orig_num_atoms].append(used[uniq_index])
                        neigh_dists[k + orig_num_atoms].append(orig_dist)

        # If the model has multiple cutoffs, we need to return neighbor lists
        # corresponding to each of them
        neigh_lists = []
        for cut in self.cutoffs:
            neigh_list = [
                np.array(neigh_list[k], dtype=np.intc)[neigh_dists[k] <= cut]
                for k in range(neighbor_list_size)
            ]
            neigh_lists.append(neigh_list)

        self.padding_image_of = padding_image_of

        self.neigh["neighbors"] = neigh_lists
        self.neigh["num_particles"] = neighbor_list_size

        return new_atoms

    def update(self, orig_atoms, species_map):
        """Create the neighbor list along with the other required
        parameters (which are stored as instance attributes). The
        required parameters are:

            - num_particles
            - coords
            - particle_contributing
            - species_code

        Note that the KIM API requires a neighbor list that has indices
        corresponding to each atom.
        """

        # Information of original atoms
        self.num_contributing_particles = len(orig_atoms)

        new_atoms = self.build(orig_atoms)

        # Save the number of atoms and all their neighbors and positions
        num_atoms = len(new_atoms)
        num_padding = num_atoms - self.num_contributing_particles
        self.num_particles = [num_atoms]
        self.coords = new_atoms.get_positions()

        # Save which coordinates are from original atoms and which are from
        # neighbors using a mask
        indices_mask = [1] * self.num_contributing_particles + [0] * num_padding
        self.particle_contributing = indices_mask

        # Species support and code
        try:
            self.species_code = [
                species_map[s] for s in new_atoms.get_chemical_symbols()
            ]
        except KeyError as e:
            raise RuntimeError("Species not supported by KIM model; {}"
                               .format(str(e)))

        self.last_update_positions = orig_atoms.get_positions()

        if self.debug:
            print("Debug: called update_ase_neigh")
            print()


class KimpyNeighborList(NeighborList):
    def __init__(
        self,
        compute_args,
        neigh_skin_ratio,
        model_influence_dist,
        model_cutoffs,
        padding_not_require_neigh,
        debug,
    ):
        super().__init__(
            neigh_skin_ratio,
            model_influence_dist,
            model_cutoffs,
            padding_not_require_neigh,
            debug,
        )

        self.neigh = neighlist.initialize()
        compute_args.set_callback_pointer(
            kimpy.compute_callback_name.GetNeighborList,
            neighlist.get_neigh_kim(),
            self.neigh,
        )

    @check_call_wrapper
    def build(self):
        return neighlist.build(
            self.neigh, self.coords, self.influence_dist,
            self.cutoffs, self.need_neigh
        )

    @check_call_wrapper
    def create_paddings(
        self, cell, pbc, contributing_coords, contributing_species_code
    ):
        # Cast things passed through kimpy to numpy arrays
        cell = np.asarray(cell, dtype=np.double)
        pbc = np.asarray(pbc, dtype=np.intc)
        contributing_coords = np.asarray(contributing_coords, dtype=np.double)

        return neighlist.create_paddings(
            self.influence_dist,
            cell,
            pbc,
            contributing_coords,
            contributing_species_code,
        )

    def update(self, atoms, species_map):
        """Create the neighbor list along with the other required
        parameters (which are stored as instance attributes). The
        required parameters are:

            - num_particles
            - coords
            - particle_contributing
            - species_code

        Note that the KIM API requires a neighbor list that has indices
        corresponding to each atom.
        """

        # Get info from Atoms object
        cell = np.asarray(atoms.get_cell(), dtype=np.double)
        pbc = np.asarray(atoms.get_pbc(), dtype=np.intc)
        contributing_coords = np.asarray(atoms.get_positions(), dtype=np.double)
        self.num_contributing_particles = atoms.get_global_number_of_atoms()
        num_contributing = self.num_contributing_particles

        # Species support and code
        try:
            contributing_species_code = np.array(
                [species_map[s] for s in atoms.get_chemical_symbols()],
                dtype=np.intc)
        except KeyError as e:
            raise RuntimeError("Species not supported by KIM model; {}"
                               .format(str(e)))

        if pbc.any():  # Need padding atoms
            # Create padding atoms

            padding_coords, padding_species_code, self.padding_image_of = \
                self.create_paddings(
                    cell, pbc, contributing_coords, contributing_species_code)
            num_padding = padding_species_code.size

            self.num_particles = [num_contributing + num_padding]
            self.coords = np.concatenate((contributing_coords, padding_coords))
            self.species_code = np.concatenate(
                (contributing_species_code, padding_species_code)
            )
            self.particle_contributing = ([1] * num_contributing +
                                          [0] * num_padding)
            self.need_neigh = [1] * self.num_particles[0]
            if not self.padding_need_neigh:
                self.need_neigh[num_contributing:] = 0

        else:  # Do not need padding atoms
            self.padding_image_of = []
            self.num_particles = [num_contributing]
            self.coords = contributing_coords
            self.species_code = contributing_species_code
            self.particle_contributing = [1] * num_contributing
            self.need_neigh = self.particle_contributing

        # Create neighborlist
        self.build()

        self.last_update_positions = atoms.get_positions()

        if self.debug:
            print("Debug: called update_kimpy_neigh")
            print()

    def clean(self):
        neighlist.clean(self.neigh)