"""
Featurizer implementations used in ChemCeption models.
SmilesToImage featurizer for ChemCeption models taken from https://arxiv.org/abs/1710.02238
"""
import numpy as np

from deepchem.utils.typing import RDKitMol
from deepchem.feat.base_classes import MolecularFeaturizer


class SmilesToImage(MolecularFeaturizer):
  """Convert SMILES string to an image.

  SmilesToImage Featurizer takes a SMILES string, and turns it into an image.
  Details taken from [1]_.

  The default size of for the image is 80 x 80. Two image modes are currently
  supported - std & engd. std is the gray scale specification,
  with atomic numbers as pixel values for atom positions and a constant value of
  2 for bond positions. engd is a 4-channel specification, which uses atom
  properties like hybridization, valency, charges in addition to atomic number.
  Bond type is also used for the bonds.

  The coordinates of all atoms are computed, and lines are drawn between atoms
  to indicate bonds. For the respective channels, the atom and bond positions are
  set to the property values as mentioned in the paper.

  Examples
  --------
  >>> import deepchem as dc
  >>> smiles = ['CC(=O)OC1=CC=CC=C1C(=O)O']
  >>> featurizer = dc.feat.SmilesToImage(img_size=80, img_spec='std')
  >>> images = featurizer.featurize(smiles)
  >>> type (images[0])
  <class 'numpy.ndarray'>
  >>> images[0].shape # (img_size, img_size, 1)
  (80, 80, 1)

  References
  ----------
  .. [1] Goh, Garrett B., et al. "Using rule-based labels for weak supervised
     learning: a ChemNet for transferable chemical property prediction."
     Proceedings of the 24th ACM SIGKDD International Conference on Knowledge
     Discovery & Data Mining. 2018.

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

  def __init__(self,
               img_size: int = 80,
               res: float = 0.5,
               max_len: int = 250,
               img_spec: str = "std"):
    """
    Parameters
    ----------
    img_size: int, default 80
      Size of the image tensor
    res: float, default 0.5
      Displays the resolution of each pixel in Angstrom
    max_len: int, default 250
      Maximum allowed length of SMILES string
    img_spec: str, default std
      Indicates the channel organization of the image tensor
    """
    if img_spec not in ["std", "engd"]:
      raise ValueError(
          "Image mode must be one of std or engd. {} is not supported".format(
              img_spec))

    self.img_size = img_size
    self.max_len = max_len
    self.res = res
    self.img_spec = img_spec
    self.embed = int(img_size * res / 2)

  def _featurize(self, datapoint: RDKitMol, **kwargs) -> np.ndarray:
    """Featurizes a single SMILE into an image.

    Parameters
    ----------
    datapoint: rdkit.Chem.rdchem.Mol
      RDKit Mol object

    Returns
    -------
    np.ndarray
      A 3D array of image, the shape is `(img_size, img_size, 1)`.
      If the length of SMILES is longer than `max_len`, this value is an empty array.
    """
    try:
      from rdkit import Chem
      from rdkit.Chem import AllChem
    except ModuleNotFoundError:
      raise ImportError("This class requires RDKit to be installed.")
    if 'mol' in kwargs:
      datapoint = kwargs.get("mol")
      raise DeprecationWarning(
          'Mol is being phased out as a parameter, please pass "datapoint" instead.'
      )

    smile = Chem.MolToSmiles(datapoint)
    if len(smile) > self.max_len:
      return np.array([])

    cmol = Chem.Mol(datapoint.ToBinary())
    cmol.ComputeGasteigerCharges()
    AllChem.Compute2DCoords(cmol)
    atom_coords = cmol.GetConformer(0).GetPositions()

    if self.img_spec == "std":
      # Setup image
      img = np.zeros((self.img_size, self.img_size, 1))
      # Compute bond properties
      bond_props = np.array(
          [[2.0, bond.GetBeginAtomIdx(),
            bond.GetEndAtomIdx()] for bond in datapoint.GetBonds()])
      # Compute atom properties
      atom_props = np.array([[atom.GetAtomicNum()] for atom in cmol.GetAtoms()])

      bond_props = bond_props.astype(np.float32)
      atom_props = atom_props.astype(np.float32)

    else:
      # Setup image
      img = np.zeros((self.img_size, self.img_size, 4))
      # Compute bond properties
      bond_props = np.array([[
          bond.GetBondTypeAsDouble(),
          bond.GetBeginAtomIdx(),
          bond.GetEndAtomIdx()
      ] for bond in datapoint.GetBonds()])
      # Compute atom properties
      atom_props = np.array([[
          atom.GetAtomicNum(),
          atom.GetProp("_GasteigerCharge"),
          atom.GetExplicitValence(),
          atom.GetHybridization().real,
      ] for atom in cmol.GetAtoms()])

      bond_props = bond_props.astype(np.float32)
      atom_props = atom_props.astype(np.float32)

      partial_charges = atom_props[:, 1]
      if np.any(np.isnan(partial_charges)):
        return np.array([])

    frac = np.linspace(0, 1, int(1 / self.res * 2))
    # Reshape done for proper broadcast
    frac = frac.reshape(-1, 1, 1)

    bond_begin_idxs = bond_props[:, 1].astype(int)
    bond_end_idxs = bond_props[:, 2].astype(int)

    # Reshapes, and axes manipulations to facilitate vector processing.
    begin_coords = atom_coords[bond_begin_idxs]
    begin_coords = np.expand_dims(begin_coords.T, axis=0)
    end_coords = atom_coords[bond_end_idxs]
    end_coords = np.expand_dims(end_coords.T, axis=0)

    # Draw a line between the two atoms.
    # The coordinates of this line, are indicated in line_coords
    line_coords = frac * begin_coords + (1 - frac) * end_coords
    # Turn the line coordinates into image positions
    bond_line_idxs = np.ceil(
        (line_coords[:, 0] + self.embed) / self.res).astype(int)
    bond_line_idys = np.ceil(
        (line_coords[:, 1] + self.embed) / self.res).astype(int)
    # Turn atomic coordinates into image positions
    atom_idxs = np.round(
        (atom_coords[:, 0] + self.embed) / self.res).astype(int)
    atom_idys = np.round(
        (atom_coords[:, 1] + self.embed) / self.res).astype(int)

    try:
      # Set the bond line coordinates to the bond property used.
      img[bond_line_idxs, bond_line_idys, 0] = bond_props[:, 0]

      # Set the atom positions in image to different atomic properties in channels
      img[atom_idxs, atom_idys, :] = atom_props

    except IndexError:
      # With fixed res and img_size some molecules (e.g. long chains) may not fit.
      raise IndexError(
          "The molecule does not fit into the image. Consider increasing img_size or res of the SmilesToImage featurizer."
      )
    return img
