"""
Generates protein-ligand docked poses.
"""
import platform
import logging
import os
import tempfile
import numpy as np
from subprocess import Popen, PIPE
from typing import List, Optional, Tuple, Union

from deepchem.dock.binding_pocket import BindingPocketFinder
from deepchem.utils.data_utils import download_url, get_data_dir
from deepchem.utils.typing import RDKitMol
from deepchem.utils.geometry_utils import compute_centroid, compute_protein_range
from deepchem.utils.rdkit_utils import load_molecule, write_molecule
from deepchem.utils.docking_utils import load_docked_ligands, write_vina_conf, write_gnina_conf, read_gnina_log

logger = logging.getLogger(__name__)
DOCKED_POSES = List[Tuple[RDKitMol, RDKitMol]]


class PoseGenerator(object):
  """A Pose Generator computes low energy conformations for molecular complexes.

  Many questions in structural biophysics reduce to that of computing
  the binding free energy of molecular complexes. A key step towards
  computing the binding free energy of two complexes is to find low
  energy "poses", that is energetically favorable conformations of
  molecules with respect to each other. One application of this
  technique is to find low energy poses for protein-ligand
  interactions.
  """

  def generate_poses(self,
                     molecular_complex: Tuple[str, str],
                     centroid: Optional[np.ndarray] = None,
                     box_dims: Optional[np.ndarray] = None,
                     exhaustiveness: int = 10,
                     num_modes: int = 9,
                     num_pockets: Optional[int] = None,
                     out_dir: Optional[str] = None,
                     generate_scores: bool = False):
    """Generates a list of low energy poses for molecular complex

    Parameters
    ----------
    molecular_complexes: Tuple[str, str]
      A representation of a molecular complex. This tuple is
      (protein_file, ligand_file).
    centroid: np.ndarray, optional (default None)
      The centroid to dock against. Is computed if not specified.
    box_dims: np.ndarray, optional (default None)
      A numpy array of shape `(3,)` holding the size of the box to dock.
      If not specified is set to size of molecular complex plus 5 angstroms.
    exhaustiveness: int, optional (default 10)
      Tells pose generator how exhaustive it should be with pose
      generation.
    num_modes: int, optional (default 9)
      Tells pose generator how many binding modes it should generate at
      each invocation.
    num_pockets: int, optional (default None)
      If specified, `self.pocket_finder` must be set. Will only
      generate poses for the first `num_pockets` returned by
      `self.pocket_finder`.
    out_dir: str, optional (default None)
      If specified, write generated poses to this directory.
    generate_score: bool, optional (default False)
      If `True`, the pose generator will return scores for complexes.
      This is used typically when invoking external docking programs
      that compute scores.

    Returns
    -------
    A list of molecular complexes in energetically favorable poses.
    """
    raise NotImplementedError


class GninaPoseGenerator(PoseGenerator):
  """Use GNINA to generate binding poses.

  This class uses GNINA (a deep learning framework for molecular
  docking) to generate binding poses. It downloads the GNINA
  executable to DEEPCHEM_DATA_DIR (an environment variable you set)
  and invokes the executable to perform pose generation.

  GNINA uses pre-trained convolutional neural network (CNN) scoring
  functions to rank binding poses based on learned representations of
  3D protein-ligand interactions. It has been shown to outperform
  AutoDock Vina in virtual screening applications [1]_.

  If you use the GNINA molecular docking engine, please cite the relevant
  papers: https://github.com/gnina/gnina#citation
  The primary citation for GNINA is [1]_.

  References
  ----------
  .. [1] M Ragoza, J Hochuli, E Idrobo, J Sunseri, DR Koes.
  "Protein–Ligand Scoring with Convolutional Neural Networks."
  Journal of chemical information and modeling (2017).

  Note
  ----
  * GNINA currently only works on Linux operating systems.
  * GNINA requires CUDA >= 10.1 for fast CNN scoring.
  * Almost all dependencies are included in the most compatible way
    possible, which reduces performance. Build GNINA from source
    for production use.

  """

  def __init__(self):
    """Initialize GNINA pose generator."""

    data_dir = get_data_dir()
    if platform.system() == 'Linux':
      url = "https://github.com/gnina/gnina/releases/download/v1.0/gnina"
      filename = 'gnina'
      self.gnina_dir = data_dir
      self.gnina_cmd = os.path.join(self.gnina_dir, filename)
    else:
      raise ValueError(
          "GNINA currently only runs on Linux. Try using a cloud platform to run this code instead."
      )

    if not os.path.exists(self.gnina_cmd):
      logger.info("GNINA not available. Downloading...")
      download_url(url, data_dir)
      downloaded_file = os.path.join(data_dir, filename)
      os.chmod(downloaded_file, 755)
      logger.info("Downloaded GNINA.")

  def generate_poses(
      self,
      molecular_complex: Tuple[str, str],
      centroid: Optional[np.ndarray] = None,
      box_dims: Optional[np.ndarray] = None,
      exhaustiveness: int = 10,
      num_modes: int = 9,
      num_pockets: Optional[int] = None,
      out_dir: Optional[str] = None,
      generate_scores: bool = True,
      **kwargs) -> Union[Tuple[DOCKED_POSES, np.ndarray], DOCKED_POSES]:
    """Generates the docked complex and outputs files for docked complex.

    Parameters
    ----------
    molecular_complexes: Tuple[str, str]
      A representation of a molecular complex. This tuple is
      (protein_file, ligand_file).
    centroid: np.ndarray, optional (default None)
      The centroid to dock against. Is computed if not specified.
    box_dims: np.ndarray, optional (default None)
      A numpy array of shape `(3,)` holding the size of the box to dock.
      If not specified is set to size of molecular complex plus 4 angstroms.
    exhaustiveness: int (default 8)
      Tells GNINA how exhaustive it should be with pose
      generation.
    num_modes: int (default 9)
      Tells GNINA how many binding modes it should generate at
      each invocation.
    out_dir: str, optional
      If specified, write generated poses to this directory.
    generate_scores: bool, optional (default True)
      If `True`, the pose generator will return scores for complexes.
      This is used typically when invoking external docking programs
      that compute scores.
    kwargs:
      Any args supported by GNINA as documented
      https://github.com/gnina/gnina#usage

    Returns
    -------
    Tuple[`docked_poses`, `scores`] or `docked_poses`
      Tuple of `(docked_poses, scores)` or `docked_poses`. `docked_poses`
      is a list of docked molecular complexes. Each entry in this list
      contains a `(protein_mol, ligand_mol)` pair of RDKit molecules.
      `scores` is an array of binding affinities (kcal/mol),
      CNN pose scores, and CNN affinities predicted by GNINA.

    """

    if out_dir is None:
      out_dir = tempfile.mkdtemp()
    if not os.path.exists(out_dir):
      os.makedirs(out_dir)

    # Parse complex
    if len(molecular_complex) > 2:
      raise ValueError(
          "GNINA can only dock protein-ligand complexes and not more general molecular complexes."
      )

    (protein_file, ligand_file) = molecular_complex

    # check filetypes
    if not protein_file.endswith('.pdb'):
      raise ValueError('Protein file must be in .pdb format.')
    if not ligand_file.endswith('.sdf'):
      raise ValueError('Ligand file must be in .sdf format.')

    protein_mol = load_molecule(
        protein_file, calc_charges=True, add_hydrogens=True)
    ligand_name = os.path.basename(ligand_file).split(".")[0]

    # Define locations of log and output files
    log_file = os.path.join(out_dir, "%s_log.txt" % ligand_name)
    out_file = os.path.join(out_dir, "%s_docked.pdbqt" % ligand_name)
    logger.info("About to call GNINA.")

    # Write GNINA conf file
    conf_file = os.path.join(out_dir, "conf.txt")
    write_gnina_conf(
        protein_filename=protein_file,
        ligand_filename=ligand_file,
        conf_filename=conf_file,
        num_modes=num_modes,
        exhaustiveness=exhaustiveness,
        **kwargs)

    # Run GNINA
    args = [
        self.gnina_cmd, "--config", conf_file, "--log", log_file, "--out",
        out_file
    ]
    process = Popen(args, stdout=PIPE, stderr=PIPE)
    stdout, stderr = process.communicate()

    # read output and log
    ligands, _ = load_docked_ligands(out_file)
    docked_complexes = [(protein_mol[1], ligand) for ligand in ligands]
    scores = read_gnina_log(log_file)

    if generate_scores:
      return docked_complexes, scores
    else:
      return docked_complexes


class VinaPoseGenerator(PoseGenerator):
  """Uses Autodock Vina to generate binding poses.

  This class uses Autodock Vina to make make predictions of
  binding poses.

  Example
  -------
  >> import deepchem as dc
  >> vpg = dc.dock.VinaPoseGenerator(pocket_finder=None)
  >> protein_file = '1jld_protein.pdb'
  >> ligand_file = '1jld_ligand.sdf'
  >> poses, scores = vpg.generate_poses(
  ..        (protein_file, ligand_file),
  ..        exhaustiveness=1,
  ..        num_modes=1,
  ..        out_dir=tmp,
  ..        generate_scores=True)

  Note
  ----
  This class requires RDKit and vina to be installed.
  """

  def __init__(self, pocket_finder: Optional[BindingPocketFinder] = None):
    """Initializes Vina Pose Generator

    Parameters
    ----------
    pocket_finder: BindingPocketFinder, optional (default None)
      If specified should be an instance of
      `dc.dock.BindingPocketFinder`.
    """
    self.pocket_finder = pocket_finder

  def generate_poses(
      self,
      molecular_complex: Tuple[str, str],
      centroid: Optional[np.ndarray] = None,
      box_dims: Optional[np.ndarray] = None,
      exhaustiveness: int = 10,
      num_modes: int = 9,
      num_pockets: Optional[int] = None,
      out_dir: Optional[str] = None,
      generate_scores: Optional[bool] = False,
      **kwargs) -> Union[Tuple[DOCKED_POSES, List[float]], DOCKED_POSES]:
    """Generates the docked complex and outputs files for docked complex.

    Parameters
    ----------
    molecular_complexes: Tuple[str, str]
      A representation of a molecular complex. This tuple is
      (protein_file, ligand_file). The protein should be a pdb file
      and the ligand should be an sdf file.
    centroid: np.ndarray, optional
      The centroid to dock against. Is computed if not specified.
    box_dims: np.ndarray, optional
      A numpy array of shape `(3,)` holding the size of the box to dock. If not
      specified is set to size of molecular complex plus 5 angstroms.
    exhaustiveness: int, optional (default 10)
      Tells Autodock Vina how exhaustive it should be with pose generation. A
      higher value of exhaustiveness implies more computation effort for the
      docking experiment.
    num_modes: int, optional (default 9)
      Tells Autodock Vina how many binding modes it should generate at
      each invocation.
    num_pockets: int, optional (default None)
      If specified, `self.pocket_finder` must be set. Will only
      generate poses for the first `num_pockets` returned by
      `self.pocket_finder`.
    out_dir: str, optional
      If specified, write generated poses to this directory.
    generate_score: bool, optional (default False)
      If `True`, the pose generator will return scores for complexes.
      This is used typically when invoking external docking programs
      that compute scores.
    kwargs:
      The kwargs - cpu, min_rmsd, max_evals, energy_range supported by VINA
      are as documented in https://autodock-vina.readthedocs.io/en/latest/vina.html

    Returns
    -------
    Tuple[`docked_poses`, `scores`] or `docked_poses`
      Tuple of `(docked_poses, scores)` or `docked_poses`. `docked_poses`
      is a list of docked molecular complexes. Each entry in this list
      contains a `(protein_mol, ligand_mol)` pair of RDKit molecules.
      `scores` is a list of binding free energies predicted by Vina.

    Raises
    ------
    `ValueError` if `num_pockets` is set but `self.pocket_finder is None`.
    """
    if "cpu" in kwargs:
      cpu = kwargs["cpu"]
    else:
      cpu = 0
    if "min_rmsd" in kwargs:
      min_rmsd = kwargs["min_rmsd"]
    else:
      min_rmsd = 1.0
    if "max_evals" in kwargs:
      max_evals = kwargs["max_evals"]
    else:
      max_evals = 0
    if "energy_range" in kwargs:
      energy_range = kwargs["energy_range"]
    else:
      energy_range = 3.0

    try:
      from vina import Vina
    except ModuleNotFoundError:
      raise ImportError("This function requires vina to be installed")

    if out_dir is None:
      out_dir = tempfile.mkdtemp()

    if num_pockets is not None and self.pocket_finder is None:
      raise ValueError(
          "If num_pockets is specified, pocket_finder must have been provided at construction time."
      )

    # Parse complex
    if len(molecular_complex) > 2:
      raise ValueError(
          "Autodock Vina can only dock protein-ligand complexes and not more general molecular complexes."
      )

    (protein_file, ligand_file) = molecular_complex

    # Prepare protein
    protein_name = os.path.basename(protein_file).split(".")[0]
    protein_hyd = os.path.join(out_dir, "%s_hyd.pdb" % protein_name)
    protein_pdbqt = os.path.join(out_dir, "%s.pdbqt" % protein_name)
    protein_mol = load_molecule(
        protein_file, calc_charges=True, add_hydrogens=True)
    write_molecule(protein_mol[1], protein_hyd, is_protein=True)
    write_molecule(protein_mol[1], protein_pdbqt, is_protein=True)

    # Get protein centroid and range
    if centroid is not None and box_dims is not None:
      centroids = [centroid]
      dimensions = [box_dims]
    else:
      if self.pocket_finder is None:
        logger.info("Pockets not specified. Will use whole protein to dock")
        centroids = [compute_centroid(protein_mol[0])]
        dimensions = [compute_protein_range(protein_mol[0]) + 5.0]
      else:
        logger.info("About to find putative binding pockets")
        pockets = self.pocket_finder.find_pockets(protein_file)
        logger.info("%d pockets found in total" % len(pockets))
        logger.info("Computing centroid and size of proposed pockets.")
        centroids, dimensions = [], []
        for pocket in pockets:
          (x_min, x_max), (y_min, y_max), (
              z_min, z_max) = pocket.x_range, pocket.y_range, pocket.z_range
          # TODO(rbharath: Does vina divide box dimensions by 2?
          x_box = (x_max - x_min) / 2.
          y_box = (y_max - y_min) / 2.
          z_box = (z_max - z_min) / 2.
          centroids.append(pocket.center())
          dimensions.append(np.array((x_box, y_box, z_box)))

    if num_pockets is not None:
      logger.info("num_pockets = %d so selecting this many pockets for docking."
                  % num_pockets)
      centroids = centroids[:num_pockets]
      dimensions = dimensions[:num_pockets]

    # Prepare ligand
    ligand_name = os.path.basename(ligand_file).split(".")[0]
    ligand_pdbqt = os.path.join(out_dir, "%s.pdbqt" % ligand_name)

    ligand_mol = load_molecule(
        ligand_file, calc_charges=True, add_hydrogens=True)
    write_molecule(ligand_mol[1], ligand_pdbqt)

    docked_complexes = []
    all_scores = []
    vpg = Vina(sf_name='vina', cpu=cpu, seed=0, no_refine=False, verbosity=1)
    for i, (protein_centroid, box_dims) in enumerate(
        zip(centroids, dimensions)):
      logger.info("Docking in pocket %d/%d" % (i + 1, len(centroids)))
      logger.info("Docking with center: %s" % str(protein_centroid))
      logger.info("Box dimensions: %s" % str(box_dims))
      # Write Vina conf file
      conf_file = os.path.join(out_dir, "conf.txt")
      write_vina_conf(
          protein_pdbqt,
          ligand_pdbqt,
          protein_centroid,
          box_dims,
          conf_file,
          num_modes=num_modes,
          exhaustiveness=exhaustiveness)

      # Define locations of output files
      out_pdbqt = os.path.join(out_dir, "%s_docked.pdbqt" % ligand_name)
      logger.info("About to call Vina")

      vpg.set_receptor(protein_pdbqt)
      vpg.set_ligand_from_file(ligand_pdbqt)

      vpg.compute_vina_maps(center=protein_centroid, box_size=box_dims)
      vpg.dock(
          exhaustiveness=exhaustiveness,
          n_poses=num_modes,
          min_rmsd=min_rmsd,
          max_evals=max_evals)
      vpg.write_poses(
          out_pdbqt,
          n_poses=num_modes,
          energy_range=energy_range,
          overwrite=True)

      ligands, scores = load_docked_ligands(out_pdbqt)
      docked_complexes += [(protein_mol[1], ligand) for ligand in ligands]
      all_scores += scores

    if generate_scores:
      return docked_complexes, all_scores
    else:
      return docked_complexes
