import torch

from .open import _PairIndexer

# Deprecated?
class StaticImagePeriodicPairIndexer(_PairIndexer):
    """Finds Pairs within a given number of images"""
    def __init__(self, *args, n_images=1, **kwargs):
        super().__init__(*args, **kwargs)
        import itertools

        # shape is (n_images+2)**3 by 3
        combinator = torch.tensor(list(itertools.product(range(-n_images, n_images + 1), repeat=3)))
        self.n_images = n_images
        self.register_buffer("combinator", combinator.to(torch.float))

    def forward(self, coordinates, nonblank, real_atoms, inv_real_atoms, cells):

        original_coordinates = coordinates

        with torch.no_grad():
            cur_device = coordinates.device

            # Figure out how many molecules and atoms we have
            n_molecules, n_atoms, _ = coordinates.shape
            n_atoms_max_batch = int(nonblank.sum(axis=1).max())
            # Construct a unique index for each atom (including blanks here)
            atom_index = torch.reshape(
                torch.arange(n_molecules * n_atoms, dtype=torch.long, device=cur_device), (n_molecules, n_atoms)
            )

            nonblank = nonblank[:, :n_atoms_max_batch]
            coordinates = coordinates[:, :n_atoms_max_batch]
            atom_index = atom_index[:, :n_atoms_max_batch]

            # Eliminate blank atom pairs, construct indices for them
            nonblank_pair = nonblank.unsqueeze(1) * nonblank.unsqueeze(2)
            nb_b, nb_a1, nb_a2 = torch.nonzero(nonblank_pair, as_tuple=True)

            # Get the indices of non-blank atoms
            pair_first = atom_index[nb_b, nb_a1]
            pair_second = atom_index[nb_b, nb_a2]
            # shape pair

            pair_diffcoords = coordinates[nb_b, nb_a1] - coordinates[nb_b, nb_a2]

            # Begin periodic part: find offsets for all images of cells.
            # Shape: n_shifts, n_batch, 3
            cell_shifts = torch.tensordot(self.combinator, cells, dims=((1,), (1,)))
            # Shape n_shifts,n_atoms,3
            offsets = cell_shifts[:, nb_b]

            # Find pairs/images close enough to be within the cutoff
            # shape n_shifts, pair
            pair_dist = (pair_diffcoords.unsqueeze(0) + offsets).norm(dim=2)
            close_pairs = (pair_dist < self.hard_dist_cutoff) & (pair_dist > 1e-5)  # get rid of self-pairs
            cp_offset, cp_pair = torch.nonzero(close_pairs, as_tuple=True)
            # cp_pair holds the indices of pairs
            # cp_offset holds the indices of the offsets

            # Extract the pairs within the cutoff
            pair_first = pair_first[cp_pair]
            pair_second = pair_second[cp_pair]
            pair_offsets = offsets[cp_offset, cp_pair]
            cell_offsets = self.combinator[cp_offset]

            # This converts the atom index from an index that counts blank atoms to one that doesn't;
            # Afterwards it indexes the flat array of atoms in the batch.
            pair_first = inv_real_atoms[pair_first]
            pair_second = inv_real_atoms[pair_second]

            del nb_b, nb_a1, nb_a2, cp_pair, pair_dist, offsets, pair_diffcoords, nonblank_pair, atom_index

        coordflat = original_coordinates.reshape(n_molecules * n_atoms, 3)[real_atoms]
        paircoord = coordflat[pair_first] - coordflat[pair_second] + pair_offsets
        distflat2 = paircoord.norm(dim=1)

        return distflat2, pair_first, pair_second, paircoord, cell_offsets, cp_offset


@torch.jit.script
def tracebatch(mat):
    """
    Trace of  of batch of matrices.
    Defaults to last two dimensions.
    """
    return torch.diagonal(mat, dim1=-2, dim2=-1).sum(dim=-1)

# For some reason torch.linalg on GPU tends to
# spend a lot of time allocating memory, especialyl
# when it is given a large batch of matrices.
# So we use the Cayley-Hamilton version of a matrix inverse
# Without calling any linalg functions.
@torch.jit.script
def cayley_inv(A):
    """
    Inverse of a batch of 3x3 matrices.
    """
    # Uses A^-1 = 1/det(A) ([tr(A)**2-tr(A**2)]I/2 - Atr(A) + A**2)
    assert A.shape[-1] == 3
    assert A.shape[-2] == 3

    trA = tracebatch(A)
    Asq = torch.bmm(A, A)
    trAsq = tracebatch(Asq)
    eye = torch.eye(3, dtype=A.dtype, device=A.device).broadcast_to(A.shape)
    eyeterm = (torch.pow(trA, 2) - trAsq).unsqueeze(-1).unsqueeze(-1) * eye / 2
    Ainvunscale = eyeterm - A * (trA.unsqueeze(-1).unsqueeze(-1)) + Asq
    det = tracebatch(torch.bmm(A, Ainvunscale)) / 3.

    return Ainvunscale / det.unsqueeze(-1).unsqueeze(-1)

@torch.jit.script
def wrap_systems_torch(coords, cell, cutoff: float):
    """
    Wraps coordinates to within the unit cell.

    Returns Inverse cells of shape (n_systems, 3 [cartesian], 3 [basis])
    Wrapped coords in same shape of coords
    Wrapped offset in shape (n_systems,n_atoms,3 [basis])
    and n_bounds in shape (n_systems,3) in terms of the number of image cells
    to search for each basis vector (radius).

    :param coords: coordinates of shape (n_systems,n_atoms,3)
    :param cell: cells of shape (n_systems, 3 [basis], 3 [cartesian])
    :param cutoff: cutoff radius to search.
    :return: inv_cells, wrapped_coords, wrapped_offset, n_bounds
    """
    # cell is (sys,basis,cartesian)

    # This is faster than calling the torch.linalg.inv directly as of torch 1.10
    inv_cell = cayley_inv(cell)
    # inv is (sys,cartesian,basis)
    projections = torch.bmm(coords, inv_cell)
    fractional_coords = torch.remainder(projections, 1)
    wrapped_offset = torch.div(projections, 1, rounding_mode="floor")
    wrapped_coords = torch.bmm(fractional_coords, cell)
    # wrapped_coords is (system,atoms,cartesian)

    # Compute number images to search in each basis direction
    eps = 1e-5
    inv_lengths = torch.linalg.norm(inv_cell, dim=1)
    n_bounds = (torch.floor(cutoff * inv_lengths + eps) + 1).to(torch.int64)

    return inv_cell, wrapped_coords, wrapped_offset.to(torch.int64), n_bounds


class PeriodicPairIndexer(_PairIndexer):
    """
    Finds pairs in general periodic conditions.
    """
    def forward(self, coordinates, nonblank, real_atoms, inv_real_atoms, cells):

        original_coordinates = coordinates

        with torch.no_grad():
            cur_device = coordinates.device

            # Figure out how many molecules and atoms we have
            n_molecules, n_atoms, _ = coordinates.shape
            n_atoms_max_batch = nonblank.sum(dim=1).max()
            # Construct a unique index for each atom (including blanks here)
            atom_index = torch.arange(n_molecules * n_atoms, dtype=torch.long, device=cur_device)
            atom_index = atom_index.reshape(n_molecules, n_atoms)

            # Trim padding in this batch to minimal level
            nonblank = nonblank[:, :n_atoms_max_batch]
            coordinates = coordinates[:, :n_atoms_max_batch]
            atom_index = atom_index[:, :n_atoms_max_batch]

            inv_cell, coordinates, wrapped_offsets, nbounds = (
                wrap_systems_torch(coordinates, cells, self.hard_dist_cutoff)
            )
            # inv cell: n_sys, n_cartesian, n_basis
            # wraped_coords: nsys, n_atoms, 3 (coords)
            # wrapped_offsets: nsys, n_atoms, 3 (basis)
            # nbounds: (nsys, 3) : number of cells-lengths to search in each direction.

            # ## make combinator of images ## #
            a1, a2, a3 = nbounds.max(dim=0)[0].unbind()
            a1 = torch.arange(-a1, a1+1, device=cur_device)
            a2 = torch.arange(-a2, a2+1, device=cur_device)
            a3 = torch.arange(-a3, a3+1, device=cur_device)
            combinator = torch.cartesian_prod(a1, a2, a3)
            # ## end make combinator ## #

            # Find images valid for each system
            inrange_images = (combinator.unsqueeze(1) >= -nbounds) & (combinator.unsqueeze(1) <= nbounds)
            inrange_images = inrange_images.all(dim=2).transpose(0, 1)
            # shape (n_sys, n_perm)

            # Eliminate blank atom pairs, construct indices for them
            nonblank_pair = nonblank.unsqueeze(1) * nonblank.unsqueeze(2)

            # ## Begin expensive part. ## #
            # In this zone we aggressively `del` things we don't need anymore
            # to allow pytorch to free memory. Delete anything that looks proportional
            # to the number of pairs. Many variables are proportional to the number of
            # 'possible' pairs, and all_sparse is proportional to an
            # even larger number of (# sys)*(# atoms)^2*(# possible images)

            # Start search space here.
            # Compute indices of all images for systems and nonblank atoms of those systems
            all_sparse = (inrange_images.unsqueeze(1).unsqueeze(2) * nonblank_pair.unsqueeze(3))
            nb_sys, nb_p1, nb_p2, nb_image = all_sparse.nonzero(as_tuple=True)
            del all_sparse

            # pair displacements without shifts
            pair_diffcoords = coordinates[nb_sys, nb_p1] - coordinates[nb_sys, nb_p2]

            # Find displacements offsets due to images.
            cell_shifts = torch.matmul(combinator.to(cells.dtype).unsqueeze(0), cells)
            # cell_shifts is  n_sys, n_perm, 3 (cartesian)
            dist_offsets = cell_shifts[nb_sys, nb_image]
            # Find pairs/images close enough to be within the cutoff
            pair_dist = (pair_diffcoords + dist_offsets).norm(dim=1)
            del dist_offsets, pair_diffcoords

            # Calculate mask of pairs that map to same atom.
            int_offsets = combinator[nb_image]
            nonself_pairs = torch.logical_or(int_offsets.to(torch.bool).any(dim=1), torch.ne(nb_p1, nb_p2))

            # Compute indices of considered pairs relative to MolAtom format.
            pair_first = atom_index[nb_sys, nb_p1]
            del nb_p1
            pair_second = atom_index[nb_sys, nb_p2]
            del nb_p2

            # Compute close pairs.
            close_pairs = (pair_dist < self.hard_dist_cutoff) & nonself_pairs  # get rid of self-pairs
            del nonself_pairs, pair_dist
            close_pairs = torch.nonzero(close_pairs)[:, 0]
            # close_pairs holds the indices of pairs

            # Extract the needed data within the cutoff.
            pair_mol = nb_sys[close_pairs]
            del nb_sys
            pair_first = pair_first[close_pairs]
            pair_second = pair_second[close_pairs]
            cell_only_offsets = int_offsets[close_pairs]  # Offset due to cell images only.
            offset_num = nb_image[close_pairs]  # index of where to store cell_offset values if caching pairs.
            del int_offsets, close_pairs
            # ## End expensive part ## #

            # This converts the atom index from an index that counts blank atoms (MolAtom) to one that doesn't;
            # Afterwards it indexes the flat array of atoms in the batch. (Atoms format)
            pair_first = inv_real_atoms[pair_first]
            pair_second = inv_real_atoms[pair_second]

            # Add back in offsets due to atom wrapping.
            atom_wrap_offsets = wrapped_offsets[nonblank]
            pf_shift = atom_wrap_offsets[pair_first]  # Offset of wrapping of first atom
            ps_shift = atom_wrap_offsets[pair_second]  # Offset of wrapping of second atom

            cell_offsets = cell_only_offsets - (pf_shift - ps_shift)  # Total image offset of pairs.
            del cell_only_offsets, pf_shift, ps_shift

        # Compute distance differentiably from total image offsets and pair indices
        pair_shifts = torch.matmul(cell_offsets.unsqueeze(1).to(cells.dtype), cells[pair_mol]).squeeze(1)
        coordflat = original_coordinates.reshape(n_molecules * n_atoms, 3)[real_atoms]
        paircoord = coordflat[pair_first] - coordflat[pair_second] + pair_shifts
        distflat2 = paircoord.norm(dim=1)

        return distflat2, pair_first, pair_second, paircoord, cell_offsets, offset_num
