"""
Computes putative binding pockets on protein.
"""
import logging
import numpy as np
from typing import Any, List, Optional, Tuple

from deepchem.models import Model
from deepchem.utils.rdkit_utils import load_molecule
from deepchem.utils.coordinate_box_utils import CoordinateBox, get_face_boxes, merge_overlapping_boxes
from deepchem.utils.fragment_utils import get_contact_atom_indices

logger = logging.getLogger(__name__)


def extract_active_site(
        protein_file: str,
        ligand_file: str,
        cutoff: float = 4.0) -> Tuple[CoordinateBox, np.ndarray]:
    """Extracts a box for the active site.

    Parameters
    ----------
    protein_file: str
        Location of protein PDB
    ligand_file: str
        Location of ligand input file
    cutoff: float, optional (default 4.0)
        The distance in angstroms from the protein pocket to
        consider for featurization.

    Returns
    -------
    Tuple[CoordinateBox, np.ndarray]
        A tuple of `(CoordinateBox, np.ndarray)` where the second entry is
        of shape `(N, 3)` with `N` the number of atoms in the active site.
    """
    protein = load_molecule(protein_file, add_hydrogens=False)
    ligand = load_molecule(ligand_file, add_hydrogens=True, calc_charges=True)
    protein_contacts, ligand_contacts = get_contact_atom_indices(
        [protein, ligand], cutoff=cutoff)
    protein_coords = protein[0]
    pocket_coords = protein_coords[protein_contacts]

    x_min = int(np.floor(np.amin(pocket_coords[:, 0])))
    x_max = int(np.ceil(np.amax(pocket_coords[:, 0])))
    y_min = int(np.floor(np.amin(pocket_coords[:, 1])))
    y_max = int(np.ceil(np.amax(pocket_coords[:, 1])))
    z_min = int(np.floor(np.amin(pocket_coords[:, 2])))
    z_max = int(np.ceil(np.amax(pocket_coords[:, 2])))
    box = CoordinateBox((x_min, x_max), (y_min, y_max), (z_min, z_max))
    return box, pocket_coords


class BindingPocketFinder(object):
    """Abstract superclass for binding pocket detectors

    Many times when working with a new protein or other macromolecule,
    it's not clear what zones of the macromolecule may be good targets
    for potential ligands or other molecules to interact with. This
    abstract class provides a template for child classes that
    algorithmically locate potential binding pockets that are good
    potential interaction sites.

    Note that potential interactions sites can be found by many
    different methods, and that this abstract class doesn't specify the
    technique to be used.
    """

    def find_pockets(self, molecule: Any):
        """Finds potential binding pockets in proteins.

        Parameters
        ----------
        molecule: object
            Some representation of a molecule.
        """
        raise NotImplementedError


class ConvexHullPocketFinder(BindingPocketFinder):
    """Implementation that uses convex hull of protein to find pockets.

    Based on https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4112621/pdf/1472-6807-14-18.pdf
    """

    def __init__(self, scoring_model: Optional[Model] = None, pad: float = 5.0):
        """Initialize the pocket finder.

        Parameters
        ----------
        scoring_model: Model, optional (default None)
            If specified, use this model to prune pockets.
        pad: float, optional (default 5.0)
            The number of angstroms to pad around a binding pocket's atoms
            to get a binding pocket box.
        """
        self.scoring_model = scoring_model
        self.pad = pad

    def find_all_pockets(self, protein_file: str) -> List[CoordinateBox]:
        """Find list of binding pockets on protein.

        Parameters
        ----------
        protein_file: str
            Protein to load in.

        Returns
        -------
        List[CoordinateBox]
            List of binding pockets on protein. Each pocket is a `CoordinateBox`
        """
        coords, _ = load_molecule(protein_file)
        return get_face_boxes(coords, self.pad)

    def find_pockets(self, macromolecule_file: str) -> List[CoordinateBox]:
        """Find list of suitable binding pockets on protein.

        This function computes putative binding pockets on this protein.
        This class uses the `ConvexHull` to compute binding pockets. Each
        face of the hull is converted into a coordinate box used for
        binding.

        Parameters
        ----------
        macromolecule_file: str
            Location of the macromolecule file to load

        Returns
        -------
        List[CoordinateBox]
            List of pockets. Each pocket is a `CoordinateBox`
        """
        coords, _ = load_molecule(macromolecule_file,
                                  add_hydrogens=False,
                                  calc_charges=False)
        boxes = get_face_boxes(coords, self.pad)
        boxes = merge_overlapping_boxes(boxes)
        return boxes
