import numpy as np

from deepchem.utils.typing import PymatgenComposition
from deepchem.feat import MaterialCompositionFeaturizer
from typing import Any


class ElementPropertyFingerprint(MaterialCompositionFeaturizer):
  """
  Fingerprint of elemental properties from composition.

  Based on the data source chosen, returns properties and statistics
  (min, max, range, mean, standard deviation, mode) for a compound
  based on elemental stoichiometry. E.g., the average electronegativity
  of atoms in a crystal structure. The chemical fingerprint is a
  vector of these statistics. For a full list of properties and statistics,
  see ``matminer.featurizers.composition.ElementProperty(data_source).feature_labels()``.

  This featurizer requires the optional dependencies pymatgen and
  matminer. It may be useful when only crystal compositions are available
  (and not 3D coordinates).

  See references [1]_, [2]_, [3]_, [4]_ for more details.

  References
  ----------
  .. [1] MagPie data: Ward, L. et al. npj Comput Mater 2, 16028 (2016).
     https://doi.org/10.1038/npjcompumats.2016.28
  .. [2] Deml data: Deml, A. et al. Physical Review B 93, 085142 (2016).
     10.1103/PhysRevB.93.085142
  .. [3] Matminer: Ward, L. et al. Comput. Mater. Sci. 152, 60-69 (2018).
  .. [4] Pymatgen: Ong, S.P. et al. Comput. Mater. Sci. 68, 314-319 (2013).

  Examples
  --------
  >>> import deepchem as dc
  >>> import pymatgen as mg
  >>> comp = mg.core.Composition("Fe2O3")
  >>> featurizer = dc.feat.ElementPropertyFingerprint()
  >>> features = featurizer.featurize([comp])
  >>> type(features[0])
  <class 'numpy.ndarray'>
  >>> features[0].shape
  (65,)

  Note
  ----
  This class requires matminer and Pymatgen to be installed.
  `NaN` feature values are automatically converted to 0 by this featurizer.
  """

  def __init__(self, data_source: str = 'matminer'):
    """
    Parameters
    ----------
    data_source: str of "matminer", "magpie" or "deml" (default "matminer")
      Source for element property data.
    """
    self.data_source = data_source
    self.ep_featurizer: Any = None

  def _featurize(self, datapoint: PymatgenComposition, **kwargs) -> np.ndarray:
    """
    Calculate chemical fingerprint from crystal composition.

    Parameters
    ----------
    datapoint: pymatgen.core.Composition object
      Composition object.

    Returns
    -------
    feats: np.ndarray
      Vector of properties and statistics derived from chemical
      stoichiometry. Some values may be NaN.
    """
    if 'composition' in kwargs and datapoint is None:
      datapoint = kwargs.get("composition")
      raise DeprecationWarning(
          'Composition is being phased out as a parameter, please pass "datapoint" instead.'
      )

    if self.ep_featurizer is None:
      try:
        from matminer.featurizers.composition import ElementProperty
        self.ep_featurizer = ElementProperty.from_preset(self.data_source)
      except ModuleNotFoundError:
        raise ImportError("This class requires matminer to be installed.")

    try:
      feats = self.ep_featurizer.featurize(datapoint)
    except:
      feats = []

    return np.nan_to_num(np.array(feats))
