import math

import numpy as np
import torch
from ase.neighborlist import neighbor_list

__all__ = [
    "BaseEnvironmentProvider",
    "SimpleEnvironmentProvider",
    "AseEnvironmentProvider",
    "TorchEnvironmentProvider",
]


class BaseEnvironmentProvider:
    """
    Environment Providers are supposed to collect neighboring atoms within
    local, atom-centered environments. All environment providers should inherit
    from this class.

    """

    def get_environment(self, atoms):
        """
        Returns the neighbor indices and offsets

        Args:
            atoms (ase.Atoms): atomistic system

        Returns:
            neighborhood_idx (np.ndarray): indices of the neighbors with shape
                n_atoms x n_max_neighbors
            offset (np.ndarray): offset in lattice coordinates for periodic
                systems (otherwise zero matrix) of shape
                n_atoms x n_max_neighbors x 3

        """

        raise NotImplementedError


class SimpleEnvironmentProvider(BaseEnvironmentProvider):
    """
    A simple environment provider for small molecules where all atoms are each
    other's neighbors. It calculates full distance matrices and does not
    support cutoffs or periodic boundary conditions.
    """

    def get_environment(self, atoms, grid=None):
        n_atoms = atoms.get_global_number_of_atoms()

        if n_atoms == 1:
            neighborhood_idx = -np.ones((1, 1), dtype=np.float32)
            offsets = np.zeros((n_atoms, 1, 3), dtype=np.float32)
        else:
            neighborhood_idx = np.tile(
                np.arange(n_atoms, dtype=np.float32)[np.newaxis], (n_atoms, 1)
            )

            neighborhood_idx = neighborhood_idx[
                ~np.eye(n_atoms, dtype=np.bool)
            ].reshape(n_atoms, n_atoms - 1)

            if grid is not None:
                n_grid = grid.shape[0]
                neighborhood_idx = np.hstack([neighborhood_idx, -np.ones((n_atoms, 1))])
                grid_nbh = np.tile(
                    np.arange(n_atoms, dtype=np.float32)[np.newaxis], (n_grid, 1)
                )
                neighborhood_idx = np.vstack([neighborhood_idx, grid_nbh])

            offsets = np.zeros(
                (neighborhood_idx.shape[0], neighborhood_idx.shape[1], 3),
                dtype=np.float32,
            )
        return neighborhood_idx, offsets


class AseEnvironmentProvider(BaseEnvironmentProvider):
    """
    Environment provider making use of ASE neighbor lists. Supports cutoffs
    and PBCs.

    """

    def __init__(self, cutoff):
        self.cutoff = cutoff

    def get_environment(self, atoms, grid=None):
        if grid is not None:
            raise NotImplementedError

        n_atoms = atoms.get_global_number_of_atoms()
        idx_i, idx_j, idx_S = neighbor_list(
            "ijS", atoms, self.cutoff, self_interaction=False
        )
        if idx_i.shape[0] > 0:
            uidx, n_nbh = np.unique(idx_i, return_counts=True)
            n_max_nbh = np.max(n_nbh)

            n_nbh = np.tile(n_nbh[:, np.newaxis], (1, n_max_nbh))
            nbh_range = np.tile(
                np.arange(n_max_nbh, dtype=np.int)[np.newaxis], (n_nbh.shape[0], 1)
            )

            mask = np.zeros((n_atoms, np.max(n_max_nbh)), dtype=np.bool)
            mask[uidx, :] = nbh_range < n_nbh
            neighborhood_idx = -np.ones((n_atoms, np.max(n_max_nbh)), dtype=np.float32)
            neighborhood_idx[mask] = idx_j

            offset = np.zeros((n_atoms, np.max(n_max_nbh), 3), dtype=np.float32)
            offset[mask] = idx_S
        else:
            neighborhood_idx = -np.ones((n_atoms, 1), dtype=np.float32)
            offset = np.zeros((n_atoms, 1, 3), dtype=np.float32)

        return neighborhood_idx, offset


class TorchEnvironmentProvider(BaseEnvironmentProvider):
    """
    Environment provider making use of neighbor lists as implemented in TorchAni
    (https://github.com/aiqm/torchani/blob/master/torchani/aev.py).
    Supports cutoffs, PBCs and can be performed on either CPU or GPU.

    """

    def __init__(self, cutoff, device):
        """
        Args:
            cutoff (float): the cutoff inside which atoms are considered pairs
            device (:class:`torch.device`): pass torch.device('cpu') or torch.device('cuda') to
                perform the calculation on a CPU or GPU, respectively.
        """
        self.cutoff = cutoff
        self.device = device

    def get_environment(self, atoms):

        species = torch.FloatTensor(atoms.numbers).to(self.device)
        coordinates = torch.FloatTensor(atoms.positions).to(self.device)
        pbc = torch.from_numpy(atoms.pbc.astype("uint8")).to(self.device)

        if not atoms.cell.any():
            cell = torch.eye(3, dtype=species.dtype).to(self.device)
        else:
            cell = torch.Tensor(atoms.cell).to(self.device)

        shifts = compute_shifts(cell=cell, pbc=pbc, cutoff=self.cutoff)

        # The returned indices are only one directional
        idx_i, idx_j, idx_S = neighbor_pairs(
            species == -1, coordinates, cell, shifts, self.cutoff
        )

        idx_i = idx_i.cpu().detach().numpy()
        idx_j = idx_j.cpu().detach().numpy()
        idx_S = idx_S.cpu().detach().numpy()

        # Create bidirectional id arrays, similar to what the ASE neighbor_list returns
        bi_idx_i = np.hstack((idx_i, idx_j))
        bi_idx_j = np.hstack((idx_j, idx_i))
        bi_idx_S = np.vstack((-idx_S, idx_S))

        n_atoms = atoms.get_global_number_of_atoms()
        if bi_idx_i.shape[0] > 0:
            uidx, n_nbh = np.unique(bi_idx_i, return_counts=True)
            n_max_nbh = np.max(n_nbh)

            n_nbh = np.tile(n_nbh[:, np.newaxis], (1, n_max_nbh))
            nbh_range = np.tile(
                np.arange(n_max_nbh, dtype=np.int)[np.newaxis], (n_nbh.shape[0], 1)
            )

            mask = np.zeros((n_atoms, np.max(n_max_nbh)), dtype=np.bool)
            mask[uidx, :] = nbh_range < n_nbh
            neighborhood_idx = -np.ones((n_atoms, np.max(n_max_nbh)), dtype=np.float32)
            offset = np.zeros((n_atoms, np.max(n_max_nbh), 3), dtype=np.float32)

            # Assign neighbors and offsets according to the indices in bi_idx_i, since in contrast
            # to the ASE provider the bidirectional arrays are no longer sorted.
            # TODO: There might be a more efficient way of doing this than a loop
            for idx in range(n_atoms):
                neighborhood_idx[idx, mask[idx]] = bi_idx_j[bi_idx_i == idx]
                offset[idx, mask[idx]] = bi_idx_S[bi_idx_i == idx]

        else:
            neighborhood_idx = -np.ones((n_atoms, 1), dtype=np.float32)
            offset = np.zeros((n_atoms, 1, 3), dtype=np.float32)

        return neighborhood_idx, offset


def compute_shifts(cell, pbc, cutoff):
    """Compute the shifts of unit cell along the given cell vectors to make it
    large enough to contain all pairs of neighbor atoms with PBC under
    consideration.
    Copyright 2018- Xiang Gao and other ANI developers
    (https://github.com/aiqm/torchani/blob/master/torchani/aev.py)

    Arguments:
        cell (:class:`torch.Tensor`): tensor of shape (3, 3) of the three
        vectors defining unit cell:
            tensor([[x1, y1, z1], [x2, y2, z2], [x3, y3, z3]])
        cutoff (float): the cutoff inside which atoms are considered pairs
        pbc (:class:`torch.Tensor`): boolean vector of size 3 storing
            if pbc is enabled for that direction.

    Returns:
        :class:`torch.Tensor`: long tensor of shifts. the center cell and
            symmetric cells are not included.
    """
    # type: (Tensor, Tensor, float) -> Tensor
    reciprocal_cell = cell.inverse().t()
    inv_distances = reciprocal_cell.norm(2, -1)
    num_repeats = torch.ceil(cutoff * inv_distances).to(torch.long)
    num_repeats = torch.where(pbc.bool(), num_repeats, torch.zeros_like(num_repeats))

    r1 = torch.arange(1, num_repeats[0] + 1, device=cell.device)
    r2 = torch.arange(1, num_repeats[1] + 1, device=cell.device)
    r3 = torch.arange(1, num_repeats[2] + 1, device=cell.device)
    o = torch.zeros(1, dtype=torch.long, device=cell.device)

    return torch.cat(
        [
            torch.cartesian_prod(r1, r2, r3),
            torch.cartesian_prod(r1, r2, o),
            torch.cartesian_prod(r1, r2, -r3),
            torch.cartesian_prod(r1, o, r3),
            torch.cartesian_prod(r1, o, o),
            torch.cartesian_prod(r1, o, -r3),
            torch.cartesian_prod(r1, -r2, r3),
            torch.cartesian_prod(r1, -r2, o),
            torch.cartesian_prod(r1, -r2, -r3),
            torch.cartesian_prod(o, r2, r3),
            torch.cartesian_prod(o, r2, o),
            torch.cartesian_prod(o, r2, -r3),
            torch.cartesian_prod(o, o, r3),
        ]
    )


def neighbor_pairs(padding_mask, coordinates, cell, shifts, cutoff):
    """Compute pairs of atoms that are neighbors
    Copyright 2018- Xiang Gao and other ANI developers
    (https://github.com/aiqm/torchani/blob/master/torchani/aev.py)

    Arguments:
        padding_mask (:class:`torch.Tensor`): boolean tensor of shape
            (molecules, atoms) for padding mask. 1 == is padding.
        coordinates (:class:`torch.Tensor`): tensor of shape
            (molecules, atoms, 3) for atom coordinates.
        cell (:class:`torch.Tensor`): tensor of shape (3, 3) of the three vectors
            defining unit cell: tensor([[x1, y1, z1], [x2, y2, z2], [x3, y3, z3]])
        cutoff (float): the cutoff inside which atoms are considered pairs
        shifts (:class:`torch.Tensor`): tensor of shape (?, 3) storing shifts
    """
    # type: (Tensor, Tensor, Tensor, Tensor, float) -> Tuple[Tensor, Tensor, Tensor, Tensor]

    coordinates = coordinates.detach()
    cell = cell.detach()
    num_atoms = padding_mask.shape[0]
    all_atoms = torch.arange(num_atoms, device=cell.device)

    # Step 2: center cell
    p1_center, p2_center = torch.combinations(all_atoms).unbind(-1)
    shifts_center = shifts.new_zeros(p1_center.shape[0], 3)

    # Step 3: cells with shifts
    # shape convention (shift index, molecule index, atom index, 3)
    num_shifts = shifts.shape[0]
    all_shifts = torch.arange(num_shifts, device=cell.device)
    shift_index, p1, p2 = torch.cartesian_prod(all_shifts, all_atoms, all_atoms).unbind(
        -1
    )
    shifts_outside = shifts.index_select(0, shift_index)

    # Step 4: combine results for all cells
    shifts_all = torch.cat([shifts_center, shifts_outside])
    p1_all = torch.cat([p1_center, p1])
    p2_all = torch.cat([p2_center, p2])

    shift_values = torch.mm(shifts_all.to(cell.dtype), cell)

    # step 5, compute distances, and find all pairs within cutoff
    distances = (coordinates[p1_all] - coordinates[p2_all] + shift_values).norm(2, -1)

    padding_mask = (padding_mask[p1_all]) | (padding_mask[p2_all])
    distances.masked_fill_(padding_mask, math.inf)
    in_cutoff = torch.nonzero(distances < cutoff, as_tuple=False)
    pair_index = in_cutoff.squeeze()
    atom_index1 = p1_all[pair_index]
    atom_index2 = p2_all[pair_index]
    shifts = shifts_all.index_select(0, pair_index)
    return atom_index1, atom_index2, shifts


def collect_atom_triples(nbh_idx):
    """
    Collect all valid triples of atoms by rearranging neighbor indices obtained
    from an environment provider.

    Args:
        nbh_idx (numpy.ndarray): neighbor indices

    Returns:
        nbh_idx_j, nbh_idx_k (numpy.ndarray): triple indices
        offset_idx_j, offset_idx_k (numpy.ndarray): offset indices for PBC

    """
    natoms, nneigh = nbh_idx.shape

    # Construct possible permutations
    nbh_idx_j = np.tile(nbh_idx, nneigh)
    nbh_idx_k = np.repeat(nbh_idx, nneigh).reshape((natoms, -1))

    # Remove same interactions and non unique pairs
    triu_idx_row, triu_idx_col = np.triu_indices(nneigh, k=1)
    triu_idx_flat = triu_idx_row * nneigh + triu_idx_col
    nbh_idx_j = nbh_idx_j[:, triu_idx_flat]
    nbh_idx_k = nbh_idx_k[:, triu_idx_flat]

    # Keep track of periodic images
    offset_idx = np.tile(np.arange(nneigh), (natoms, 1))

    # Construct indices for pairs of offsets
    offset_idx_j = np.tile(offset_idx, nneigh)
    offset_idx_k = np.repeat(offset_idx, nneigh).reshape((natoms, -1))

    # Remove non-unique pairs and diagonal
    offset_idx_j = offset_idx_j[:, triu_idx_flat]
    offset_idx_k = offset_idx_k[:, triu_idx_flat]

    return nbh_idx_j, nbh_idx_k, offset_idx_j, offset_idx_k


def collect_atom_triples_batch(neighbors, neighbor_mask):
    """
    Batch/torch version for collecting atom triples, offset indices and the corresponding mask directly from a
    batch of neighbor indices and their corresponding mask. This is e.g. used in the
    schnetpack.md.calculators.SchnetPackCalculator class to generate extended inputs for Behler type symmetry
    functions.
    Args:
        neighbors (torch.LongTensor): (n_batch x n_atoms x n_neighbors) tensor holding the indices of all
                                      neighbor atoms (e.g. from NeighborList or EnvironmentProvider).
        neighbor_mask (torch.LongTensor): (n_batch x n_atoms x n_neighbors) binary tensor indicating non-existent
                                          atoms due to padding.
    Returns:
        torch.LongTensor: (n_batch x n_atoms x n_triples) indices of the first neighbor in all triples.
        torch.LongTensor: (n_batch x n_atoms x n_triples) indices of the second neighbor in all triples.
        torch.LongTensor: (n_batch x n_atoms x n_triples) first neighbor offset indices for PBC.
        torch.LongTensor: (n_batch x n_atoms x n_triples) second neighbor offset indices for PBC.
        torch.LongTensor: (n_batch x n_atoms x n_triples) mask indicating all invalid pairs due to padding.
    """
    B, A, N = neighbors.shape

    # Generate indices of all possible unique pairs
    idx_k, idx_j = torch.combinations(
        torch.arange(N, device=neighbors.device).long(), r=2, with_replacement=False
    ).unbind(1)

    # Generate triple indices
    nbh_idx_j = neighbors[:, :, idx_j]
    nhb_idx_k = neighbors[:, :, idx_k]

    # Generate triple offset indices
    offset_idx_j = idx_j.repeat((B, A, 1))
    offset_idx_k = idx_k.repeat((B, A, 1))

    # Generate mask for triples
    mask_j = neighbor_mask[:, :, idx_j]
    mask_k = neighbor_mask[:, :, idx_k]
    mask_triples = mask_j * mask_k

    return nbh_idx_j, nhb_idx_k, offset_idx_j, offset_idx_k, mask_triples
