import torch
import itertools
import math
from . import utils


def _cutoff_cosine(distances, cutoff):
    return torch.where(
        distances <= cutoff,
        0.5 * torch.cos(math.pi * distances / cutoff) + 0.5,
        torch.zeros_like(distances)
    )


class AEVComputer(torch.nn.Module):
    r"""The AEV computer that takes coordinates as input and outputs aevs.

    Arguments:
        Rcr (:class:`torch.Tensor`): The scalar tensor of :math:`R_C` in
            equation (2) when used at equation (3) in the `ANI paper`_.
        Rca (:class:`torch.Tensor`): The scalar tensor of :math:`R_C` in
            equation (2) when used at equation (4) in the `ANI paper`_.
        EtaR (:class:`torch.Tensor`): The 1D tensor of :math:`\eta` in
            equation (3) in the `ANI paper`_.
        ShfR (:class:`torch.Tensor`): The 1D tensor of :math:`R_s` in
            equation (3) in the `ANI paper`_.
        EtaA (:class:`torch.Tensor`): The 1D tensor of :math:`\eta` in
            equation (4) in the `ANI paper`_.
        Zeta (:class:`torch.Tensor`): The 1D tensor of :math:`\zeta` in
            equation (4) in the `ANI paper`_.
        ShfA (:class:`torch.Tensor`): The 1D tensor of :math:`R_s` in
            equation (4) in the `ANI paper`_.
        ShfZ (:class:`torch.Tensor`): The 1D tensor of :math:`\theta_s` in
            equation (4) in the `ANI paper`_.
        num_species (int): Number of supported atom types.

    .. _ANI paper:
        http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
    """

    def __init__(self, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ,
                 num_species):
        super(AEVComputer, self).__init__()
        self.register_buffer('Rcr', Rcr)
        self.register_buffer('Rca', Rca)
        # convert constant tensors to a ready-to-broadcast shape
        # shape convension (..., EtaR, ShfR)
        self.register_buffer('EtaR', EtaR.view(-1, 1))
        self.register_buffer('ShfR', ShfR.view(1, -1))
        # shape convension (..., EtaA, Zeta, ShfA, ShfZ)
        self.register_buffer('EtaA', EtaA.view(-1, 1, 1, 1))
        self.register_buffer('Zeta', Zeta.view(1, -1, 1, 1))
        self.register_buffer('ShfA', ShfA.view(1, 1, -1, 1))
        self.register_buffer('ShfZ', ShfZ.view(1, 1, 1, -1))

        self.num_species = num_species

    def radial_sublength(self):
        """Returns the length of radial subaev of a single species"""
        return self.EtaR.numel() * self.ShfR.numel()

    def radial_length(self):
        """Returns the length of full radial aev"""
        return self.num_species * self.radial_sublength()

    def angular_sublength(self):
        """Returns the length of angular subaev of a single species"""
        return self.EtaA.numel() * self.Zeta.numel() * self.ShfA.numel() * \
            self.ShfZ.numel()

    def angular_length(self):
        """Returns the length of full angular aev"""
        s = self.num_species
        return (s * (s + 1)) // 2 * self.angular_sublength()

    def aev_length(self):
        """Returns the length of full aev"""
        return self.radial_length() + self.angular_length()

    def _radial_subaev_terms(self, distances):
        """Compute the radial subAEV terms of the center atom given neighbors

        This correspond to equation (3) in the `ANI paper`_. This function just
        compute the terms. The sum in the equation is not computed.
        The input tensor have shape (conformations, atoms, N), where ``N``
        is the number of neighbor atoms within the cutoff radius and output
        tensor should have shape
        (conformations, atoms, ``self.radial_sublength()``)

        .. _ANI paper:
            http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
        """
        distances = distances.unsqueeze(-1).unsqueeze(-1)
        fc = _cutoff_cosine(distances, self.Rcr)
        # Note that in the equation in the paper there is no 0.25
        # coefficient, but in NeuroChem there is such a coefficient.
        # We choose to be consistent with NeuroChem instead of the paper here.
        ret = 0.25 * torch.exp(-self.EtaR * (distances - self.ShfR)**2) * fc
        # At this point, ret now have shape
        # (conformations, atoms, N, ?, ?) where ? depend on constants.
        # We then should flat the last 4 dimensions to view the subAEV as one
        # dimension vector
        return ret.flatten(start_dim=-2)

    def _angular_subaev_terms(self, vectors1, vectors2):
        """Compute the angular subAEV terms of the center atom given neighbor pairs.

        This correspond to equation (4) in the `ANI paper`_. This function just
        compute the terms. The sum in the equation is not computed.
        The input tensor have shape (conformations, atoms, N), where N
        is the number of neighbor atom pairs within the cutoff radius and
        output tensor should have shape
        (conformations, atoms, ``self.angular_sublength()``)

        .. _ANI paper:
            http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
        """
        vectors1 = vectors1.unsqueeze(
            -1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        vectors2 = vectors2.unsqueeze(
            -1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        distances1 = vectors1.norm(2, dim=-5)
        distances2 = vectors2.norm(2, dim=-5)

        # 0.95 is multiplied to the cos values to prevent acos from
        # returning NaN.
        cos_angles = 0.95 * \
            torch.nn.functional.cosine_similarity(
                vectors1, vectors2, dim=-5)
        angles = torch.acos(cos_angles)

        fcj1 = _cutoff_cosine(distances1, self.Rca)
        fcj2 = _cutoff_cosine(distances2, self.Rca)
        factor1 = ((1 + torch.cos(angles - self.ShfZ)) / 2) ** self.Zeta
        factor2 = torch.exp(-self.EtaA *
                            ((distances1 + distances2) / 2 - self.ShfA) ** 2)
        ret = 2 * factor1 * factor2 * fcj1 * fcj2
        # At this point, ret now have shape
        # (conformations, atoms, N, ?, ?, ?, ?) where ? depend on constants.
        # We then should flat the last 4 dimensions to view the subAEV as one
        # dimension vector
        return ret.flatten(start_dim=-4)

    def _terms_and_indices(self, species, coordinates):
        """Returns radial and angular subAEV terms, these terms will be sorted
        according to their distances to central atoms, and only these within
        cutoff radius are valid. The returned indices stores the source of data
        before sorting.
        """

        vec = coordinates.unsqueeze(2) - coordinates.unsqueeze(1)
        """Shape (conformations, atoms, atoms, 3) storing Rij vectors"""

        distances = vec.norm(2, -1)
        """Shape (conformations, atoms, atoms) storing Rij distances"""

        padding_mask = (species == -1).unsqueeze(1)
        distances = torch.where(
            padding_mask,
            torch.tensor(math.inf, dtype=self.EtaR.dtype,
                         device=self.EtaR.device),
            distances)

        distances, indices = distances.sort(-1)

        min_distances, _ = distances.flatten(end_dim=1).min(0)
        inRcr = (min_distances <= self.Rcr).nonzero().flatten()[1:]
        inRca = (min_distances <= self.Rca).nonzero().flatten()[1:]

        distances = distances.index_select(-1, inRcr)
        indices_r = indices.index_select(-1, inRcr)
        radial_terms = self._radial_subaev_terms(distances)

        indices_a = indices.index_select(-1, inRca)
        new_shape = list(indices_a.shape) + [3]

        # TODO: remove this workaround when gather support broadcasting
        # https://github.com/pytorch/pytorch/pull/9532
        _indices_a = indices_a.unsqueeze(-1).expand(*new_shape)
        vec = vec.gather(-2, _indices_a)

        vec = self._combinations(vec, -2)
        angular_terms = self._angular_subaev_terms(*vec)

        # Returned tensors has shape:
        # (conformations, atoms, neighbors, ``self.radial_sublength()``)
        # (conformations, atoms, pairs, ``self.angular_sublength()``)
        # (conformations, atoms, neighbors)
        # (conformations, atoms, pairs)
        return radial_terms, angular_terms, indices_r, indices_a

    def _combinations(self, tensor, dim=0):
        # TODO: remove this when combinations is merged into PyTorch
        # https://github.com/pytorch/pytorch/pull/9393
        n = tensor.shape[dim]
        r = torch.arange(n, dtype=torch.long, device=tensor.device)
        grid_x, grid_y = torch.meshgrid([r, r])
        index1 = grid_y.masked_select(
            torch.triu(torch.ones(n, n, device=tensor.device),
                       diagonal=1) == 1)
        index2 = grid_x.masked_select(
            torch.triu(torch.ones(n, n, device=tensor.device),
                       diagonal=1) == 1)
        return tensor.index_select(dim, index1), \
            tensor.index_select(dim, index2)

    def _compute_mask_r(self, species, indices_r):
        """Get mask of radial terms for each supported species from indices"""
        species_r = species.gather(-1, indices_r)
        mask_r = (species_r.unsqueeze(-1) ==
                  torch.arange(self.num_species, device=self.EtaR.device))
        return mask_r

    def _compute_mask_a(self, species, indices_a, present_species):
        """Get mask of angular terms for each supported species from indices"""
        species_a = species.gather(-1, indices_a)
        species_a1, species_a2 = self._combinations(species_a, -1)
        mask_a1 = (species_a1.unsqueeze(-1) == present_species).unsqueeze(-1)
        mask_a2 = (species_a2.unsqueeze(-1).unsqueeze(-1) == present_species)
        mask = mask_a1 * mask_a2
        mask_rev = mask.permute(0, 1, 2, 4, 3)
        mask_a = (mask + mask_rev) > 0
        return mask_a

    def _assemble(self, radial_terms, angular_terms, present_species,
                  mask_r, mask_a):
        """Returns radial and angular AEV computed from terms according
        to the given partition information.

        Arguments:
            radial_terms (:class:`torch.Tensor`): shape (conformations, atoms,
                neighbors, ``self.radial_sublength()``)
            angular_terms (:class:`torch.Tensor`): shape (conformations, atoms,
                pairs, ``self.angular_sublength()``)
            present_species (:class:`torch.Tensor`):  Long tensor for species
                of atoms present in the molecules.
            mask_r (:class:`torch.Tensor`): shape (conformations, atoms,
                neighbors, supported species)
            mask_a (:class:`torch.Tensor`): shape (conformations, atoms,
                pairs, present species, present species)
        """
        conformations = radial_terms.shape[0]
        atoms = radial_terms.shape[1]

        # assemble radial subaev
        present_radial_aevs = (
            radial_terms.unsqueeze(-2) *
            mask_r.unsqueeze(-1).type(radial_terms.dtype)
        ).sum(-3)
        """shape (conformations, atoms, present species, radial_length)"""
        radial_aevs = present_radial_aevs.flatten(start_dim=2)

        # assemble angular subaev
        rev_indices = {present_species[i].item(): i
                       for i in range(len(present_species))}
        angular_aevs = []
        zero_angular_subaev = torch.zeros(
            conformations, atoms, self.angular_sublength(),
            dtype=self.EtaR.dtype, device=self.EtaR.device)
        for s1, s2 in itertools.combinations_with_replacement(
                                        range(self.num_species), 2):
            if s1 in rev_indices and s2 in rev_indices:
                i1 = rev_indices[s1]
                i2 = rev_indices[s2]
                mask = mask_a[..., i1, i2].unsqueeze(-1).type(self.EtaR.dtype)
                subaev = (angular_terms * mask).sum(-2)
            else:
                subaev = zero_angular_subaev
            angular_aevs.append(subaev)

        return radial_aevs, torch.cat(angular_aevs, dim=2)

    def forward(self, species_coordinates):
        """Compute AEVs

        Arguments:
            species_coordinates (tuple): Two tensors: species and coordinates.
                species must have shape ``(C, A)`` and coordinates must have
                shape ``(C, A, 3)``, where ``C`` is the number of conformations
                in a chunk, and ``A`` is the number of atoms.

        Returns:
            tuple: Species and AEVs. species are the species from the input
            unchanged, and AEVs is a tensor of shape
            ``(C, A, self.aev_length())``
        """
        species, coordinates = species_coordinates

        present_species = utils.present_species(species)

        # TODO: remove this workaround after gather support broadcasting
        atoms = coordinates.shape[1]
        species_ = species.unsqueeze(1).expand(-1, atoms, -1)

        radial_terms, angular_terms, indices_r, indices_a = \
            self._terms_and_indices(species, coordinates)
        mask_r = self._compute_mask_r(species_, indices_r)
        mask_a = self._compute_mask_a(species_, indices_a, present_species)

        radial, angular = self._assemble(radial_terms, angular_terms,
                                         present_species, mask_r, mask_a)
        fullaev = torch.cat([radial, angular], dim=2)
        return species, fullaev
