"""
Generate coulomb matrices for molecules.

See Montavon et al., _New Journal of Physics_ __15__ (2013) 095003.

"""
import numpy as np
from typing import Any, List, Optional

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


class CoulombMatrix(MolecularFeaturizer):
    """Calculate Coulomb matrices for molecules.

    Coulomb matrices provide a representation of the electronic structure of
    a molecule. For a molecule with `N` atoms, the Coulomb matrix is a
    `N X N` matrix where each element gives the strength of the
    electrostatic interaction between two atoms. The method is described
    in more detail in [1]_.

    Examples
    --------
    >>> import deepchem as dc
    >>> featurizers = dc.feat.CoulombMatrix(max_atoms=23)
    >>> input_file = 'deepchem/feat/tests/data/water.sdf' # really backed by water.sdf.csv
    >>> tasks = ["atomization_energy"]
    >>> loader = dc.data.SDFLoader(tasks, featurizer=featurizers)
    >>> dataset = loader.create_dataset(input_file)

    References
    ----------
    .. [1] Montavon, Grégoire, et al. "Learning invariant representations of
        molecules for atomization energy prediction." Advances in neural information
        processing systems. 2012.

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

    """

    def __init__(self,
                 max_atoms: int,
                 remove_hydrogens: bool = False,
                 randomize: bool = False,
                 upper_tri: bool = False,
                 n_samples: int = 1,
                 seed: Optional[int] = None):
        """Initialize this featurizer.

        Parameters
        ----------
        max_atoms: int
            The maximum number of atoms expected for molecules this featurizer will
            process.
        remove_hydrogens: bool, optional (default False)
            If True, remove hydrogens before processing them.
        randomize: bool, optional (default False)
            If True, use method `randomize_coulomb_matrices` to randomize Coulomb matrices.
        upper_tri: bool, optional (default False)
            Generate only upper triangle part of Coulomb matrices.
        n_samples: int, optional (default 1)
            If `randomize` is set to True, the number of random samples to draw.
        seed: int, optional (default None)
            Random seed to use.

        """
        self.max_atoms = int(max_atoms)
        self.remove_hydrogens = remove_hydrogens
        self.randomize = randomize
        self.upper_tri = upper_tri
        self.n_samples = n_samples
        if seed is not None:
            seed = int(seed)
        self.seed = seed

    def _featurize(self, datapoint: RDKitMol, **kwargs) -> np.ndarray:
        """
        Calculate Coulomb matrices for molecules. If extra randomized
        matrices are generated, they are treated as if they are features
        for additional conformers.

        Since Coulomb matrices are symmetric, only the (flattened) upper
        triangular portion is returned.

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

        Returns
        -------
        np.ndarray
            The coulomb matrices of the given molecule.
            The default shape is `(num_confs, max_atoms, max_atoms)`.
            If num_confs == 1, the shape is `(max_atoms, max_atoms)`.

        """
        if 'mol' in kwargs:
            datapoint = kwargs.get("mol")
            raise DeprecationWarning(
                'Mol is being phased out as a parameter, please pass "datapoint" instead.'
            )

        features = self.coulomb_matrix(datapoint)
        if self.upper_tri:
            features = np.asarray(
                [f[np.triu_indices_from(f)] for f in features])
        if features.shape[0] == 1:
            # `(1, max_atoms, max_atoms)` -> `(max_atoms, max_atoms)`
            features = np.squeeze(features, axis=0)
        return features

    def coulomb_matrix(self, mol: RDKitMol) -> np.ndarray:
        """
        Generate Coulomb matrices for each conformer of the given molecule.

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

        Returns
        -------
        np.ndarray
            The coulomb matrices of the given molecule

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

        # Check whether num_confs >=1 or not
        num_confs = len(mol.GetConformers())
        if num_confs == 0:
            mol = Chem.AddHs(mol)
            AllChem.EmbedMolecule(mol, AllChem.ETKDG())

        if self.remove_hydrogens:
            mol = Chem.RemoveHs(mol)
        n_atoms = mol.GetNumAtoms()
        z = [atom.GetAtomicNum() for atom in mol.GetAtoms()]
        rval = []
        for conf in mol.GetConformers():
            d = self.get_interatomic_distances(conf)
            m = np.outer(z, z) / d
            m[range(n_atoms), range(n_atoms)] = 0.5 * np.array(z)**2.4
            if self.randomize:
                for random_m in self.randomize_coulomb_matrix(m):
                    random_m = pad_array(random_m, self.max_atoms)
                    rval.append(random_m)
            else:
                m = pad_array(m, self.max_atoms)
                rval.append(m)
        return np.asarray(rval)

    def randomize_coulomb_matrix(self, m: np.ndarray) -> List[np.ndarray]:
        """Randomize a Coulomb matrix as decribed in [1]_:

        1. Compute row norms for M in a vector row_norms.
        2. Sample a zero-mean unit-variance noise vector e with dimension
            equal to row_norms.
        3. Permute the rows and columns of M with the permutation that
            sorts row_norms + e.

        Parameters
        ----------
        m: np.ndarray
            Coulomb matrix.

        Returns
        -------
        List[np.ndarray]
            List of the random coulomb matrix

        References
        ----------
        .. [1] Montavon et al., New Journal of Physics, 15, (2013), 095003

        """
        rval = []
        row_norms = np.asarray([np.linalg.norm(row) for row in m], dtype=float)
        rng = np.random.RandomState(self.seed)
        for i in range(self.n_samples):
            e = rng.normal(size=row_norms.size)
            p = np.argsort(row_norms + e)
            new = m[p][:, p]  # permute rows first, then columns
            rval.append(new)
        return rval

    @staticmethod
    def get_interatomic_distances(conf: Any) -> np.ndarray:
        """
        Get interatomic distances for atoms in a molecular conformer.

        Parameters
        ----------
        conf: rdkit.Chem.rdchem.Conformer
            Molecule conformer.

        Returns
        -------
        np.ndarray
            The distances matrix for all atoms in a molecule

        """
        n_atoms = conf.GetNumAtoms()
        coords = [
            # Convert AtomPositions from Angstrom to bohr (atomic units)
            conf.GetAtomPosition(i).__idiv__(0.52917721092)
            for i in range(n_atoms)
        ]
        d = np.zeros((n_atoms, n_atoms), dtype=float)
        for i in range(n_atoms):
            for j in range(i):
                d[i, j] = coords[i].Distance(coords[j])
                d[j, i] = d[i, j]
        return d


class CoulombMatrixEig(CoulombMatrix):
    """Calculate the eigenvalues of Coulomb matrices for molecules.

    This featurizer computes the eigenvalues of the Coulomb matrices for provided
    molecules. Coulomb matrices are described in [1]_.

    Examples
    --------
    >>> import deepchem as dc
    >>> featurizers = dc.feat.CoulombMatrixEig(max_atoms=23)
    >>> input_file = 'deepchem/feat/tests/data/water.sdf' # really backed by water.sdf.csv
    >>> tasks = ["atomization_energy"]
    >>> loader = dc.data.SDFLoader(tasks, featurizer=featurizers)
    >>> dataset = loader.create_dataset(input_file)

    References
    ----------
    .. [1] Montavon, Grégoire, et al. "Learning invariant representations of
        molecules for atomization energy prediction." Advances in neural information
        processing systems. 2012.

    """

    def __init__(self,
                 max_atoms: int,
                 remove_hydrogens: bool = False,
                 randomize: bool = False,
                 n_samples: int = 1,
                 seed: Optional[int] = None):
        """Initialize this featurizer.

        Parameters
        ----------
        max_atoms: int
            The maximum number of atoms expected for molecules this featurizer will
            process.
        remove_hydrogens: bool, optional (default False)
            If True, remove hydrogens before processing them.
        randomize: bool, optional (default False)
            If True, use method `randomize_coulomb_matrices` to randomize Coulomb matrices.
        n_samples: int, optional (default 1)
            If `randomize` is set to True, the number of random samples to draw.
        seed: int, optional (default None)
            Random seed to use.

        """
        self.max_atoms = int(max_atoms)
        self.remove_hydrogens = remove_hydrogens
        self.randomize = randomize
        self.n_samples = n_samples
        if seed is not None:
            seed = int(seed)
        self.seed = seed

    def _featurize(self, datapoint: RDKitMol, **kwargs) -> np.ndarray:
        """
        Calculate eigenvalues of Coulomb matrix for molecules. Eigenvalues
        are returned sorted by absolute value in descending order and padded
        by max_atoms.

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

        Returns
        -------
        np.ndarray
            The eigenvalues of Coulomb matrix for molecules.
            The default shape is `(num_confs, max_atoms)`.
            If num_confs == 1, the shape is `(max_atoms,)`.

        """
        if 'mol' in kwargs:
            datapoint = kwargs.get("mol")
            raise DeprecationWarning(
                'Mol is being phased out as a parameter, please pass "datapoint" instead.'
            )

        cmat = self.coulomb_matrix(datapoint)
        features_list = []
        for f in cmat:
            w, v = np.linalg.eig(f)
            w_abs = np.abs(w)
            sortidx = np.argsort(w_abs)
            sortidx = sortidx[::-1]
            w = w[sortidx]
            f = pad_array(w, self.max_atoms)
            features_list.append(f)
        features = np.asarray(features_list)
        if features.shape[0] == 1:
            # `(1, max_atoms)` -> `(max_atoms,)`
            features = np.squeeze(features, axis=0)
        return features
