"""
Utilities for constructing node features or bond features.
Some functions are based on chainer-chemistry or dgl-lifesci.

Repositories:
- https://github.com/chainer/chainer-chemistry
- https://github.com/awslabs/dgl-lifesci
"""

import os
import logging
from typing import List, Union, Tuple

import numpy as np

from deepchem.utils.typing import RDKitAtom, RDKitBond, RDKitMol

logger = logging.getLogger(__name__)

DEFAULT_ATOM_TYPE_SET = [
    "C",
    "N",
    "O",
    "F",
    "P",
    "S",
    "Cl",
    "Br",
    "I",
]
DEFAULT_HYBRIDIZATION_SET = ["SP", "SP2", "SP3"]
DEFAULT_TOTAL_NUM_Hs_SET = [0, 1, 2, 3, 4]
DEFAULT_FORMAL_CHARGE_SET = [-2, -1, 0, 1, 2]
DEFAULT_TOTAL_DEGREE_SET = [0, 1, 2, 3, 4, 5]
DEFAULT_RING_SIZE_SET = [3, 4, 5, 6, 7, 8]
DEFAULT_BOND_TYPE_SET = ["SINGLE", "DOUBLE", "TRIPLE", "AROMATIC"]
DEFAULT_BOND_STEREO_SET = ["STEREONONE", "STEREOANY", "STEREOZ", "STEREOE"]
DEFAULT_GRAPH_DISTANCE_SET = [1, 2, 3, 4, 5, 6, 7]
DEFAULT_ATOM_IMPLICIT_VALENCE_SET = [0, 1, 2, 3, 4, 5, 6]
DEFAULT_ATOM_EXPLICIT_VALENCE_SET = [1, 2, 3, 4, 5, 6]
ALLEN_ELECTRONEGATIVTY = {  # Allen scale electronegativity
    'H': 2.3,
    'He': 4.160,
    'Li': 0.912,
    'Be': 1.576,
    'B': 2.051,
    'C': 2.544,
    'N': 3.066,
    'O': 3.610,
    'F': 4.193,
    'Ne': 4.787,
    'Na': 0.869,
    'Mg': 1.293,
    'Al': 1.613,
    'Si': 1.916,
    'P': 2.253,
    'S': 2.589,
    'Cl': 2.869,
    'Ar': 3.242,
    'K': 0.734,
    'Ca': 1.034,
    'Sc': 1.19,
    'Ti': 1.38,
    'V': 1.53,
    'Cr': 1.65,
    'Mn': 1.75,
    'Fe': 1.80,
    'Co': 1.84,
    'Ni': 1.88,
    'Cu': 1.85,
    'Zn': 1.588,
    'Ga': 1.756,
    'Ge': 1.994,
    'As': 2.211,
    'Se': 2.424,
    'Br': 2.685,
    'Kr': 2.966,
    'Rb': 0.706,
    'Sr': 0.963,
    'Y': 1.12,
    'Zr': 1.32,
    'Nb': 1.41,
    'Mo': 1.47,
    'Tc': 1.51,
    'Ru': 1.54,
    'Rh': 1.56,
    'Pd': 1.58,
    'Ag': 1.87,
    'Cd': 1.521,
    'In': 1.656,
    'Sn': 1.824,
    'Sb': 1.984,
    'Te': 2.158,
    'I': 2.359,
    'Xe': 2.582,
    'Cs': 0.659,
    'Ba': 0.881,
    'Lu': 1.09,
    'Hf': 1.16,
    'Ta': 1.34,
    'W': 1.47,
    'Re': 1.60,
    'Os': 1.65,
    'Ir': 1.68,
    'Pt': 1.72,
    'Au': 1.92,
    'Hg': 1.765,
    'Tl': 1.789,
    'Pb': 1.854,
    'Bi': 2.01,
    'Po': 2.19,
    'At': 2.39,
    'Rn': 2.60,
    'Fr': 0.67,
    'Ra': 0.89
}


class _ChemicalFeaturesFactory:
  """This is a singleton class for RDKit base features."""
  _instance = None

  @classmethod
  def get_instance(cls):
    try:
      from rdkit import RDConfig
      from rdkit.Chem import ChemicalFeatures
    except ModuleNotFoundError:
      raise ImportError("This class requires RDKit to be installed.")

    if not cls._instance:
      fdefName = os.path.join(RDConfig.RDDataDir, 'BaseFeatures.fdef')
      cls._instance = ChemicalFeatures.BuildFeatureFactory(fdefName)
    return cls._instance


def one_hot_encode(val: Union[int, str],
                   allowable_set: Union[List[str], List[int]],
                   include_unknown_set: bool = False) -> List[float]:
  """One hot encoder for elements of a provided set.

  Examples
  --------
  >>> one_hot_encode("a", ["a", "b", "c"])
  [1.0, 0.0, 0.0]
  >>> one_hot_encode(2, [0, 1, 2])
  [0.0, 0.0, 1.0]
  >>> one_hot_encode(3, [0, 1, 2])
  [0.0, 0.0, 0.0]
  >>> one_hot_encode(3, [0, 1, 2], True)
  [0.0, 0.0, 0.0, 1.0]

  Parameters
  ----------
  val: int or str
    The value must be present in `allowable_set`.
  allowable_set: List[int] or List[str]
    List of allowable quantities.
  include_unknown_set: bool, default False
    If true, the index of all values not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    An one-hot vector of val.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.

  Raises
  ------
  ValueError
    If include_unknown_set is False and `val` is not in `allowable_set`.
  """
  if include_unknown_set is False:
    if val not in allowable_set:
      logger.info("input {0} not in allowable set {1}:".format(
          val, allowable_set))

  # init an one-hot vector
  if include_unknown_set is False:
    one_hot_legnth = len(allowable_set)
  else:
    one_hot_legnth = len(allowable_set) + 1
  one_hot = [0.0 for _ in range(one_hot_legnth)]

  try:
    one_hot[allowable_set.index(val)] = 1.0  # type: ignore
  except:
    if include_unknown_set:
      # If include_unknown_set is True, set the last index is 1.
      one_hot[-1] = 1.0
    else:
      pass
  return one_hot


#################################################################
# atom (node) featurization
#################################################################


def get_atom_type_one_hot(atom: RDKitAtom,
                          allowable_set: List[str] = DEFAULT_ATOM_TYPE_SET,
                          include_unknown_set: bool = True) -> List[float]:
  """Get an one-hot feature of an atom type.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object
  allowable_set: List[str]
    The atom types to consider. The default set is
    `["C", "N", "O", "F", "P", "S", "Cl", "Br", "I"]`.
  include_unknown_set: bool, default True
    If true, the index of all atom not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    An one-hot vector of atom types.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.
  """
  return one_hot_encode(atom.GetSymbol(), allowable_set, include_unknown_set)


def construct_hydrogen_bonding_info(mol: RDKitMol) -> List[Tuple[int, str]]:
  """Construct hydrogen bonding infos about a molecule.

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

  Returns
  -------
  List[Tuple[int, str]]
    A list of tuple `(atom_index, hydrogen_bonding_type)`.
    The `hydrogen_bonding_type` value is "Acceptor" or "Donor".
  """
  factory = _ChemicalFeaturesFactory.get_instance()
  feats = factory.GetFeaturesForMol(mol)
  hydrogen_bonding = []
  for f in feats:
    hydrogen_bonding.append((f.GetAtomIds()[0], f.GetFamily()))
  return hydrogen_bonding


def get_atom_hydrogen_bonding_one_hot(
    atom: RDKitAtom, hydrogen_bonding: List[Tuple[int, str]]) -> List[float]:
  """Get an one-hot feat about whether an atom accepts electrons or donates electrons.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object
  hydrogen_bonding: List[Tuple[int, str]]
    The return value of `construct_hydrogen_bonding_info`.
    The value is a list of tuple `(atom_index, hydrogen_bonding)` like (1, "Acceptor").

  Returns
  -------
  List[float]
    A one-hot vector of the ring size type. The first element
    indicates "Donor", and the second element indicates "Acceptor".
  """
  one_hot = [0.0, 0.0]
  atom_idx = atom.GetIdx()
  for hydrogen_bonding_tuple in hydrogen_bonding:
    if hydrogen_bonding_tuple[0] == atom_idx:
      if hydrogen_bonding_tuple[1] == "Donor":
        one_hot[0] = 1.0
      elif hydrogen_bonding_tuple[1] == "Acceptor":
        one_hot[1] = 1.0
  return one_hot


def get_atom_is_in_aromatic_one_hot(atom: RDKitAtom) -> List[float]:
  """Get ans one-hot feature about whether an atom is in aromatic system or not.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object

  Returns
  -------
  List[float]
    A vector of whether an atom is in aromatic system or not.
  """
  return [float(atom.GetIsAromatic())]


def get_atom_hybridization_one_hot(
    atom: RDKitAtom,
    allowable_set: List[str] = DEFAULT_HYBRIDIZATION_SET,
    include_unknown_set: bool = False) -> List[float]:
  """Get an one-hot feature of hybridization type.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object
  allowable_set: List[str]
    The hybridization types to consider. The default set is `["SP", "SP2", "SP3"]`
  include_unknown_set: bool, default False
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    An one-hot vector of the hybridization type.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.
  """
  return one_hot_encode(str(atom.GetHybridization()), allowable_set,
                        include_unknown_set)


def get_atom_total_num_Hs_one_hot(
    atom: RDKitAtom,
    allowable_set: List[int] = DEFAULT_TOTAL_NUM_Hs_SET,
    include_unknown_set: bool = True) -> List[float]:
  """Get an one-hot feature of the number of hydrogens which an atom has.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object
  allowable_set: List[int]
    The number of hydrogens to consider. The default set is `[0, 1, ..., 4]`
  include_unknown_set: bool, default True
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    A one-hot vector of the number of hydrogens which an atom has.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.
  """
  return one_hot_encode(atom.GetTotalNumHs(), allowable_set,
                        include_unknown_set)


def get_atom_chirality_one_hot(atom: RDKitAtom) -> List[float]:
  """Get an one-hot feature about an atom chirality type.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object

  Returns
  -------
  List[float]
    A one-hot vector of the chirality type. The first element
    indicates "R", and the second element indicates "S".
  """
  one_hot = [0.0, 0.0]
  try:
    chiral_type = atom.GetProp('_CIPCode')
    if chiral_type == "R":
      one_hot[0] = 1.0
    elif chiral_type == "S":
      one_hot[1] = 1.0
  except:
    pass
  return one_hot


def get_atom_formal_charge(atom: RDKitAtom) -> List[float]:
  """Get a formal charge of an atom.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object

  Returns
  -------
  List[float]
    A vector of the formal charge.
  """
  return [float(atom.GetFormalCharge())]


def get_atom_formal_charge_one_hot(
    atom: RDKitAtom,
    allowable_set: List[int] = DEFAULT_FORMAL_CHARGE_SET,
    include_unknown_set: bool = True) -> List[float]:
  """Get one hot encoding of formal charge of an atom.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object
  allowable_set: List[int]
    The degree to consider. The default set is `[-2, -1, ..., 2]`
  include_unknown_set: bool, default True
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.


  Returns
  -------
  List[float]
    A vector of the formal charge.
  """
  return one_hot_encode(atom.GetFormalCharge(), allowable_set,
                        include_unknown_set)


def get_atom_partial_charge(atom: RDKitAtom) -> List[float]:
  """Get a partial charge of an atom.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object

  Returns
  -------
  List[float]
    A vector of the parital charge.

  Notes
  -----
  Before using this function, you must calculate `GasteigerCharge`
  like `AllChem.ComputeGasteigerCharges(mol)`.
  """
  gasteiger_charge = atom.GetProp('_GasteigerCharge')
  if gasteiger_charge in ['-nan', 'nan', '-inf', 'inf']:
    gasteiger_charge = 0.0
  return [float(gasteiger_charge)]


def get_atom_total_degree_one_hot(
    atom: RDKitAtom,
    allowable_set: List[int] = DEFAULT_TOTAL_DEGREE_SET,
    include_unknown_set: bool = True) -> List[float]:
  """Get an one-hot feature of the degree which an atom has.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object
  allowable_set: List[int]
    The degree to consider. The default set is `[0, 1, ..., 5]`
  include_unknown_set: bool, default True
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    A one-hot vector of the degree which an atom has.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.
  """
  return one_hot_encode(atom.GetTotalDegree(), allowable_set,
                        include_unknown_set)


def get_atom_implicit_valence_one_hot(
    atom: RDKitAtom,
    allowable_set: List[int] = DEFAULT_ATOM_IMPLICIT_VALENCE_SET,
    include_unknown_set: bool = True) -> List[float]:
  """Get an one-hot feature of implicit valence of an atom.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object
  allowable_set: List[int]
    Atom implicit valence to consider. The default set is `[0, 1, ..., 6]`
  include_unknown_set: bool, default True
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    A one-hot vector of implicit valence an atom has.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.

  """
  return one_hot_encode(atom.GetImplicitValence(), allowable_set,
                        include_unknown_set)


def get_atom_explicit_valence_one_hot(
    atom: RDKitAtom,
    allowable_set: List[int] = DEFAULT_ATOM_EXPLICIT_VALENCE_SET,
    include_unknown_set: bool = True) -> List[float]:
  """Get an one-hot feature of explicit valence of an atom.

  Parameters
  ---------
  atom: rdkit.Chem.rdchem.Atom
    RDKit atom object
  allowable_set: List[int]
    Atom explicit valence to consider. The default set is `[1, ..., 6]`
  include_unknown_set: bool, default True
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    A one-hot vector of explicit valence an atom has.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.

  """
  return one_hot_encode(atom.GetExplicitValence(), allowable_set,
                        include_unknown_set)


#################################################################
# bond (edge) featurization
#################################################################


def get_bond_type_one_hot(bond: RDKitBond,
                          allowable_set: List[str] = DEFAULT_BOND_TYPE_SET,
                          include_unknown_set: bool = False) -> List[float]:
  """Get an one-hot feature of bond type.

  Parameters
  ---------
  bond: rdkit.Chem.rdchem.Bond
    RDKit bond object
  allowable_set: List[str]
    The bond types to consider. The default set is `["SINGLE", "DOUBLE", "TRIPLE", "AROMATIC"]`.
  include_unknown_set: bool, default False
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    A one-hot vector of the bond type.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.
  """
  return one_hot_encode(str(bond.GetBondType()), allowable_set,
                        include_unknown_set)


def get_bond_is_in_same_ring_one_hot(bond: RDKitBond) -> List[float]:
  """Get an one-hot feature about whether atoms of a bond is in the same ring or not.

  Parameters
  ---------
  bond: rdkit.Chem.rdchem.Bond
    RDKit bond object

  Returns
  -------
  List[float]
    A one-hot vector of whether a bond is in the same ring or not.
  """
  return [int(bond.IsInRing())]


def get_bond_is_conjugated_one_hot(bond: RDKitBond) -> List[float]:
  """Get an one-hot feature about whether a bond is conjugated or not.

  Parameters
  ---------
  bond: rdkit.Chem.rdchem.Bond
    RDKit bond object

  Returns
  -------
  List[float]
    A one-hot vector of whether a bond is conjugated or not.
  """
  return [int(bond.GetIsConjugated())]


def get_bond_stereo_one_hot(bond: RDKitBond,
                            allowable_set: List[str] = DEFAULT_BOND_STEREO_SET,
                            include_unknown_set: bool = True) -> List[float]:
  """Get an one-hot feature of the stereo configuration of a bond.

  Parameters
  ---------
  bond: rdkit.Chem.rdchem.Bond
    RDKit bond object
  allowable_set: List[str]
    The stereo configuration types to consider.
    The default set is `["STEREONONE", "STEREOANY", "STEREOZ", "STEREOE"]`.
  include_unknown_set: bool, default True
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    A one-hot vector of the stereo configuration of a bond.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.
  """
  return one_hot_encode(str(bond.GetStereo()), allowable_set,
                        include_unknown_set)


def get_bond_graph_distance_one_hot(
    bond: RDKitBond,
    graph_dist_matrix: np.ndarray,
    allowable_set: List[int] = DEFAULT_GRAPH_DISTANCE_SET,
    include_unknown_set: bool = True) -> List[float]:
  """Get an one-hot feature of graph distance.

  Parameters
  ---------
  bond: rdkit.Chem.rdchem.Bond
    RDKit bond object
  graph_dist_matrix: np.ndarray
    The return value of `Chem.GetDistanceMatrix(mol)`. The shape is `(num_atoms, num_atoms)`.
  allowable_set: List[int]
    The graph distance types to consider. The default set is `[1, 2, ..., 7]`.
  include_unknown_set: bool, default False
    If true, the index of all types not in `allowable_set` is `len(allowable_set)`.

  Returns
  -------
  List[float]
    A one-hot vector of the graph distance.
    If `include_unknown_set` is False, the length is `len(allowable_set)`.
    If `include_unknown_set` is True, the length is `len(allowable_set) + 1`.
  """
  graph_dist = graph_dist_matrix[bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()]
  return one_hot_encode(graph_dist, allowable_set, include_unknown_set)
