"""
Utilities to score protein-ligand poses using DeepChem.
"""
import numpy as np


def pairwise_distances(coords1: np.ndarray, coords2: np.ndarray) -> np.ndarray:
    """Returns matrix of pairwise Euclidean distances.

    Parameters
    ----------
    coords1: np.ndarray
        A numpy array of shape `(N, 3)`
    coords2: np.ndarray
        A numpy array of shape `(M, 3)`

    Returns
    -------
    np.ndarray
        A `(N,M)` array with pairwise distances.
    """
    return np.sum((coords1[None, :] - coords2[:, None])**2, -1)**0.5


def cutoff_filter(d: np.ndarray, x: np.ndarray, cutoff=8.0) -> np.ndarray:
    """Applies a cutoff filter on pairwise distances

    Parameters
    ----------
    d: np.ndarray
        Pairwise distances matrix. A numpy array of shape `(N, M)`
    x: np.ndarray
        Matrix of shape `(N, M)`
    cutoff: float, optional (default 8)
        Cutoff for selection in Angstroms

    Returns
    -------
    np.ndarray
        A `(N,M)` array with values where distance is too large thresholded to 0.
    """
    return np.where(d < cutoff, x, np.zeros_like(x))


def vina_nonlinearity(c: np.ndarray, w: float, Nrot: int) -> np.ndarray:
    """Computes non-linearity used in Vina.

    Parameters
    ----------
    c: np.ndarray
        A numpy array of shape `(N, M)`
    w: float
        Weighting term
    Nrot: int
        Number of rotatable bonds in this molecule

    Returns
    -------
    np.ndarray
        A `(N, M)` array with activations under a nonlinearity.
    """
    out_tensor = c / (1 + w * Nrot)
    return out_tensor


def vina_repulsion(d: np.ndarray) -> np.ndarray:
    """Computes Autodock Vina's repulsion interaction term.

    Parameters
    ----------
    d: np.ndarray
        A numpy array of shape `(N, M)`.

    Returns
    -------
    np.ndarray
        A `(N, M)` array with repulsion terms.
    """
    return np.where(d < 0, d**2, np.zeros_like(d))


def vina_hydrophobic(d: np.ndarray) -> np.ndarray:
    """Computes Autodock Vina's hydrophobic interaction term.

    Here, d is the set of surface distances as defined in [1]_

    Parameters
    ----------
    d: np.ndarray
        A numpy array of shape `(N, M)`.

    Returns
    -------
    np.ndarray
        A `(N, M)` array of hydrophoboic interactions in a piecewise linear curve.

    References
    ----------
    .. [1] Jain, Ajay N. "Scoring noncovalent protein-ligand interactions:
        a continuous differentiable function tuned to compute binding affinities."
        Journal of computer-aided molecular design 10.5 (1996): 427-440.
    """
    out_tensor = np.where(d < 0.5, np.ones_like(d),
                          np.where(d < 1.5, 1.5 - d, np.zeros_like(d)))
    return out_tensor


def vina_hbond(d: np.ndarray) -> np.ndarray:
    """Computes Autodock Vina's hydrogen bond interaction term.

    Here, d is the set of surface distances as defined in [1]_

    Parameters
    ----------
    d: np.ndarray
        A numpy array of shape `(N, M)`.

    Returns
    -------
    np.ndarray
        A `(N, M)` array of hydrophoboic interactions in a piecewise linear curve.

    References
    ----------
    .. [1] Jain, Ajay N. "Scoring noncovalent protein-ligand interactions:
        a continuous differentiable function tuned to compute binding affinities."
        Journal of computer-aided molecular design 10.5 (1996): 427-440.
    """
    out_tensor = np.where(
        d < -0.7, np.ones_like(d),
        np.where(d < 0, (1.0 / 0.7) * (0 - d), np.zeros_like(d)))
    return out_tensor


def vina_gaussian_first(d: np.ndarray) -> np.ndarray:
    """Computes Autodock Vina's first Gaussian interaction term.

    Here, d is the set of surface distances as defined in [1]_

    Parameters
    ----------
    d: np.ndarray
        A numpy array of shape `(N, M)`.

    Returns
    -------
    np.ndarray
        A `(N, M)` array of gaussian interaction terms.

    References
    ----------
    .. [1] Jain, Ajay N. "Scoring noncovalent protein-ligand interactions:
        a continuous differentiable function tuned to compute binding affinities."
        Journal of computer-aided molecular design 10.5 (1996): 427-440.
    """
    out_tensor = np.exp(-(d / 0.5)**2)
    return out_tensor


def vina_gaussian_second(d: np.ndarray) -> np.ndarray:
    """Computes Autodock Vina's second Gaussian interaction term.

    Here, d is the set of surface distances as defined in [1]_

    Parameters
    ----------
    d: np.ndarray
        A numpy array of shape `(N, M)`.

    Returns
    -------
    np.ndarray
        A `(N, M)` array of gaussian interaction terms.

    References
    ----------
    .. [1] Jain, Ajay N. "Scoring noncovalent protein-ligand interactions:
        a continuous differentiable function tuned to compute binding affinities."
        Journal of computer-aided molecular design 10.5 (1996): 427-440.
    """
    out_tensor = np.exp(-((d - 3) / 2)**2)
    return out_tensor


def weighted_linear_sum(w: np.ndarray, x: np.ndarray) -> np.ndarray:
    """Computes weighted linear sum.

    Parameters
    ----------
    w: np.ndarray
        A numpy array of shape `(N,)`
    x: np.ndarray
        A numpy array of shape `(N, M, L)`

    Returns
    -------
    np.ndarray
        A numpy array of shape `(M, L)`
    """
    return np.tensordot(w, x, axes=1)


def vina_energy_term(coords1: np.ndarray, coords2: np.ndarray,
                     weights: np.ndarray, wrot: float, Nrot: int) -> np.ndarray:
    """Computes the Vina Energy function for two molecular conformations

    Parameters
    ----------
    coords1: np.ndarray
        Molecular coordinates of shape `(N, 3)`
    coords2: np.ndarray
        Molecular coordinates of shape `(M, 3)`
    weights: np.ndarray
        A numpy array of shape `(5,)`. The 5 values are weights for repulsion interaction term,
        hydrophobic interaction term, hydrogen bond interaction term,
        first Gaussian interaction term and second Gaussian interaction term.
    wrot: float
        The scaling factor for nonlinearity
    Nrot: int
        Number of rotatable bonds in this calculation

    Returns
    -------
    np.ndarray
        A scalar value with free energy
    """
    # TODO(rbharath): The autodock vina source computes surface distances
    # which take into account the van der Waals radius of each atom type.
    dists = pairwise_distances(coords1, coords2)
    repulsion = vina_repulsion(dists)
    hydrophobic = vina_hydrophobic(dists)
    hbond = vina_hbond(dists)
    gauss_1 = vina_gaussian_first(dists)
    gauss_2 = vina_gaussian_second(dists)

    # Shape (N, M)
    interactions = weighted_linear_sum(
        weights, np.array([repulsion, hydrophobic, hbond, gauss_1, gauss_2]))

    # Shape (N, M)
    thresholded = cutoff_filter(dists, interactions)

    free_energies = vina_nonlinearity(thresholded, wrot, Nrot)
    return np.sum(free_energies)
