import logging
import numpy as np
from deepchem.utils.typing import RDKitBond, RDKitMol, List
from deepchem.feat.base_classes import MolecularFeaturizer
from deepchem.utils.typing import OneOrMany

from typing import Optional

logger = logging.getLogger(__name__)


class GraphMatrix:
  """
  This is class used to store data for MolGAN neural networks.

  Parameters
  ----------
  node_features: np.ndarray
    Node feature matrix with shape [num_nodes, num_node_features]
  edge_features: np.ndarray,
    Edge feature matrix with shape [num_nodes, num_nodes]

  Returns
  -------
  graph: GraphMatrix
    A molecule graph with some features.
  """

  def __init__(self, adjacency_matrix: np.ndarray, node_features: np.ndarray):
    self.adjacency_matrix = adjacency_matrix
    self.node_features = node_features


class MolGanFeaturizer(MolecularFeaturizer):
  """
  Featurizer for MolGAN de-novo molecular generation [1]_.
  The default representation is in form of GraphMatrix object.
  It is wrapper for two matrices containing atom and bond type information.
  The class also provides reverse capabilities.

  Examples
  --------
  >>> import deepchem as dc
  >>> from rdkit import Chem
  >>> rdkit_mol, smiles_mol = Chem.MolFromSmiles('CCC'), 'C1=CC=CC=C1'
  >>> molecules = [rdkit_mol, smiles_mol]
  >>> featurizer = dc.feat.MolGanFeaturizer()
  >>> features = featurizer.featurize(molecules)
  >>> len(features) # 2 molecules
  2
  >>> type(features[0])
  <class 'deepchem.feat.molecule_featurizers.molgan_featurizer.GraphMatrix'>
  >>> molecules = featurizer.defeaturize(features) # defeaturization
  >>> type(molecules[0])
  <class 'rdkit.Chem.rdchem.Mol'>

  """

  def __init__(
      self,
      max_atom_count: int = 9,
      kekulize: bool = True,
      bond_labels: List[RDKitBond] = None,
      atom_labels: List[int] = None,
  ):
    """
    Parameters
    ----------
    max_atom_count: int, default 9
      Maximum number of atoms used for creation of adjacency matrix.
      Molecules cannot have more atoms than this number
      Implicit hydrogens do not count.
    kekulize: bool, default True
      Should molecules be kekulized.
      Solves number of issues with defeaturization when used.
    bond_labels: List[RDKitBond]
      List of types of bond used for generation of adjacency matrix
    atom_labels: List[int]
      List of atomic numbers used for generation of node features

    References
    ---------
    .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for
       small molecular graphs" (2018), https://arxiv.org/abs/1805.11973
    """

    self.max_atom_count = max_atom_count
    self.kekulize = kekulize

    try:
      from rdkit import Chem
    except ModuleNotFoundError:
      raise ImportError("This class requires RDKit to be installed.")

    # bond labels
    if bond_labels is None:
      self.bond_labels = [
          Chem.rdchem.BondType.ZERO,
          Chem.rdchem.BondType.SINGLE,
          Chem.rdchem.BondType.DOUBLE,
          Chem.rdchem.BondType.TRIPLE,
          Chem.rdchem.BondType.AROMATIC,
      ]
    else:
      self.bond_labels = bond_labels

    # atom labels
    if atom_labels is None:
      self.atom_labels = [0, 6, 7, 8, 9]  # C,N,O,F
    else:
      self.atom_labels = atom_labels

    # create bond encoders and decoders
    self.bond_encoder = {l: i for i, l in enumerate(self.bond_labels)}
    self.bond_decoder = {i: l for i, l in enumerate(self.bond_labels)}
    # create atom encoders and decoders
    self.atom_encoder = {l: i for i, l in enumerate(self.atom_labels)}
    self.atom_decoder = {i: l for i, l in enumerate(self.atom_labels)}

  def _featurize(self, datapoint: RDKitMol, **kwargs) -> Optional[GraphMatrix]:
    """
    Calculate adjacency matrix and nodes features for RDKitMol.
    It strips any chirality and charges

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

    Returns
    -------
    graph: GraphMatrix
      A molecule graph with some features.
    """

    try:
      from rdkit import Chem
    except ModuleNotFoundError:
      raise ImportError("This method 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.'
      )

    if self.kekulize:
      Chem.Kekulize(datapoint)

    A = np.zeros(
        shape=(self.max_atom_count, self.max_atom_count), dtype=np.float32)
    bonds = datapoint.GetBonds()

    begin, end = [b.GetBeginAtomIdx() for b in bonds], [
        b.GetEndAtomIdx() for b in bonds
    ]
    bond_type = [self.bond_encoder[b.GetBondType()] for b in bonds]

    A[begin, end] = bond_type
    A[end, begin] = bond_type

    degree = np.sum(
        A[:datapoint.GetNumAtoms(), :datapoint.GetNumAtoms()], axis=-1)
    X = np.array(
        [
            self.atom_encoder[atom.GetAtomicNum()]
            for atom in datapoint.GetAtoms()
        ] + [0] * (self.max_atom_count - datapoint.GetNumAtoms()),
        dtype=np.int32,
    )
    graph = GraphMatrix(A, X)

    return graph if (degree > 0).all() else None

  def _defeaturize(self,
                   graph_matrix: GraphMatrix,
                   sanitize: bool = True,
                   cleanup: bool = True) -> RDKitMol:
    """
    Recreate RDKitMol from GraphMatrix object.
    Same featurizer need to be used for featurization and defeaturization.
    It only recreates bond and atom types, any kind of additional features
    like chirality or charge are not included.
    Therefore, any checks of type: original_smiles == defeaturized_smiles
    will fail on chiral or charged compounds.

    Parameters
    ----------
    graph_matrix: GraphMatrix
      GraphMatrix object.
    sanitize: bool, default True
      Should RDKit sanitization be included in the process.
    cleanup: bool, default True
      Splits salts and removes compounds with "*" atom types

    Returns
    -------
    mol: RDKitMol object
      RDKitMol object representing molecule.
    """

    try:
      from rdkit import Chem
    except ModuleNotFoundError:
      raise ImportError("This method requires RDKit to be installed.")

    if not isinstance(graph_matrix, GraphMatrix):
      return None

    node_labels = graph_matrix.node_features
    edge_labels = graph_matrix.adjacency_matrix

    mol = Chem.RWMol()

    for node_label in node_labels:
      mol.AddAtom(Chem.Atom(self.atom_decoder[node_label]))

    for start, end in zip(*np.nonzero(edge_labels)):
      if start > end:
        mol.AddBond(
            int(start), int(end), self.bond_decoder[edge_labels[start, end]])

    if sanitize:
      try:
        Chem.SanitizeMol(mol)
      except Exception:
        mol = None

    if cleanup:
      try:
        smiles = Chem.MolToSmiles(mol)
        smiles = max(smiles.split("."), key=len)
        if "*" not in smiles:
          mol = Chem.MolFromSmiles(smiles)
        else:
          mol = None
      except Exception:
        mol = None

    return mol

  def defeaturize(self, graphs: OneOrMany[GraphMatrix],
                  log_every_n: int = 1000) -> np.ndarray:
    """
    Calculates molecules from corresponding GraphMatrix objects.

    Parameters
    ----------
    graphs: GraphMatrix / iterable
      GraphMatrix object or corresponding iterable
    log_every_n: int, default 1000
      Logging messages reported every `log_every_n` samples.

    Returns
    -------
    features: np.ndarray
      A numpy array containing RDKitMol objext.
    """

    # Special case handling of single molecule
    if isinstance(graphs, GraphMatrix):
      graphs = [graphs]
    else:
      # Convert iterables to list
      graphs = list(graphs)

    molecules = []
    for i, gr in enumerate(graphs):
      if i % log_every_n == 0:
        logger.info("Featurizing datapoint %i" % i)

      try:
        molecules.append(self._defeaturize(gr))
      except Exception as e:
        logger.warning(
            "Failed to defeaturize datapoint %d, %s. Appending empty array",
            i,
            gr,
        )
        logger.warning("Exception message: {}".format(e))
        molecules.append(np.array([]))

    return np.asarray(molecules)
