"""
Various utilities around voxel grids.
"""
import logging
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import numpy as np

from deepchem.utils.noncovalent_utils import compute_pi_stack

logger = logging.getLogger(__name__)


def convert_atom_to_voxel(coordinates: np.ndarray, atom_index: int,
                          box_width: float, voxel_width: float) -> np.ndarray:
    """Converts atom coordinates to an i,j,k grid index.

    This function offsets molecular atom coordinates by
    (box_width/2, box_width/2, box_width/2) and then divides by
    voxel_width to compute the voxel indices.

    Parameters
    -----------
    coordinates: np.ndarray
        Array with coordinates of all atoms in the molecule, shape (N, 3).
    atom_index: int
        Index of an atom in the molecule.
    box_width: float
        Size of the box in Angstroms.
    voxel_width: float
        Size of a voxel in Angstroms

    Returns
    -------
    indices: np.ndarray
        A 1D numpy array of length 3 with `[i, j, k]`, the voxel coordinates
        of specified atom.
    """

    indices = np.floor(
        (coordinates[atom_index] + box_width / 2.0) / voxel_width).astype(int)

    return indices


def convert_atom_pair_to_voxel(coordinates_tuple: Tuple[np.ndarray, np.ndarray],
                               atom_index_pair: Tuple[int,
                                                      int], box_width: float,
                               voxel_width: float) -> np.ndarray:
    """Converts a pair of atoms to i,j,k grid indexes.


    Parameters
    ----------
    coordinates_tuple: Tuple[np.ndarray, np.ndarray]
        A tuple containing two molecular coordinate arrays of shapes `(N, 3)` and `(M, 3)`.
    atom_index_pair: Tuple[int, int]
        A tuple of indices for the atoms in the two molecules.
    box_width: float
        Size of the box in Angstroms.
    voxel_width: float
        Size of a voxel in Angstroms

    Returns
    -------
    indices_list: np.ndarray
        A numpy array of shape `(2, 3)`, where `3` is `[i, j, k]` of the
        voxel coordinates of specified atom.
    """

    indices_list = []
    for coordinates, atom_index in zip(coordinates_tuple, atom_index_pair):
        indices_list.append(
            convert_atom_to_voxel(coordinates, atom_index, box_width,
                                  voxel_width))
    return np.array(indices_list)


def voxelize(get_voxels: Callable[..., Any],
             coordinates: Any,
             box_width: float = 16.0,
             voxel_width: float = 1.0,
             hash_function: Optional[Callable[..., Any]] = None,
             feature_dict: Optional[Dict[Any, Any]] = None,
             feature_list: Optional[List[Union[int, Tuple[int]]]] = None,
             nb_channel: int = 16,
             dtype: str = 'int') -> np.ndarray:
    """Helper function to voxelize inputs.

    This helper function helps convert a hash function which
    specifies spatial features of a molecular complex into a voxel
    tensor. This utility is used by various featurizers that generate
    voxel grids.

    Parameters
    ----------
    get_voxels: Function
        Function that voxelizes inputs
    coordinates: Any
        Contains the 3D coordinates of a molecular system.  This should have
        whatever type get_voxels() expects as its first argument.
    box_width: float, optional (default 16.0)
        Size of a box in which voxel features are calculated. Box
        is centered on a ligand centroid.
    voxel_width: float, optional (default 1.0)
        Size of a 3D voxel in a grid in Angstroms.
    hash_function: Function
        Used to map feature choices to voxel channels.
    feature_dict: Dict, optional (default None)
        Keys are atom indices or tuples of atom indices, the values are
        computed features. If `hash_function is not None`, then the values
        are hashed using the hash function into `[0, nb_channels)` and
        this channel at the voxel for the given key is incremented by `1`
        for each dictionary entry. If `hash_function is None`, then the
        value must be a vector of size `(n_channels,)` which is added to
        the existing channel values at that voxel grid.
    feature_list: List, optional (default None)
        List of atom indices or tuples of atom indices. This can only be
        used if `nb_channel==1`. Increments the voxels corresponding to
        these indices by `1` for each entry.
    nb_channel: int, , optional (default 16)
        The number of feature channels computed per voxel. Should
        be a power of 2.
    dtype: str ('int' or 'float'), optional (default 'int')
        The type of the numpy ndarray created to hold features.

    Returns
    -------
    feature_tensor: np.ndarray
        The voxel of the input with the shape
        `(voxels_per_edge, voxels_per_edge, voxels_per_edge, nb_channel)`.
    """
    # Number of voxels per one edge of box to voxelize.
    voxels_per_edge = int(box_width / voxel_width)
    if dtype == "int":
        feature_tensor = np.zeros(
            (voxels_per_edge, voxels_per_edge, voxels_per_edge, nb_channel),
            dtype=np.int8)
    else:
        feature_tensor = np.zeros(
            (voxels_per_edge, voxels_per_edge, voxels_per_edge, nb_channel),
            dtype=np.float16)
    if feature_dict is not None:
        for key, features in feature_dict.items():
            voxels = get_voxels(coordinates, key, box_width, voxel_width)
            if len(voxels.shape) == 1:
                voxels = np.expand_dims(voxels, axis=0)
            for voxel in voxels:
                if ((voxel >= 0) & (voxel < voxels_per_edge)).all():
                    if hash_function is not None:
                        feature_tensor[
                            voxel[0], voxel[1], voxel[2],
                            hash_function(features, nb_channel)] += 1.0
                    else:
                        feature_tensor[voxel[0], voxel[1], voxel[2],
                                       0] += features
    elif feature_list is not None:
        for key in feature_list:
            voxels = get_voxels(coordinates, key, box_width, voxel_width)
            for voxel in voxels:
                if ((voxel >= 0) & (voxel < voxels_per_edge)).all():
                    feature_tensor[voxel[0], voxel[1], voxel[2], 0] += 1.0

    return feature_tensor


def voxelize_pi_stack(prot_xyz, prot_rdk, lig_xyz, lig_rdk, distances,
                      pi_stack_dist_cutoff, pi_stack_angle_cutoff, box_width,
                      voxel_width):
    protein_pi_t, protein_pi_parallel, ligand_pi_t, ligand_pi_parallel = (
        compute_pi_stack(prot_rdk,
                         lig_rdk,
                         distances,
                         dist_cutoff=pi_stack_dist_cutoff,
                         angle_cutoff=pi_stack_angle_cutoff))
    pi_parallel_tensor = voxelize(
        convert_atom_to_voxel,
        prot_xyz,
        box_width=box_width,
        voxel_width=voxel_width,
        feature_dict=protein_pi_parallel,
        nb_channel=1,
    )

    pi_parallel_tensor += voxelize(
        convert_atom_to_voxel,
        lig_xyz,
        box_width=box_width,
        voxel_width=voxel_width,
        feature_dict=ligand_pi_parallel,
        nb_channel=1,
    )

    pi_t_tensor = voxelize(
        convert_atom_to_voxel,
        prot_xyz,
        box_width=box_width,
        voxel_width=voxel_width,
        feature_dict=protein_pi_t,
        nb_channel=1,
    )

    pi_t_tensor += voxelize(
        convert_atom_to_voxel,
        lig_xyz,
        box_width=box_width,
        voxel_width=voxel_width,
        feature_dict=ligand_pi_t,
        nb_channel=1,
    )

    return [pi_parallel_tensor, pi_t_tensor]
