import numpy as np
from loguru import logger

from kliff.descriptors.bispectrum import bs  # C extension
from kliff.descriptors.descriptor import (
    Descriptor,
    generate_full_cutoff,
    generate_species_code,
)
from kliff.neighbor import NeighborList


class Bispectrum(Descriptor):
    """Bispectrum descriptor.

    Process dataset to generate fingerprints using the Bispectrum descriptor as discussed
    in [Bartok2010]_ and [Thompson2015]_.

    Parameters
    ----------
    cut_dists: dict
        Cutoff distances, with key of the form ``A-B`` where ``A`` and ``B`` are atomic
        species string, and value should be a float.

    cut_name: str
        Name of the cutoff function.

    hyperparams: dict
        A dictionary of the hyperparams of the descriptor.

    normalize: bool (optional)
        If ``True``, the fingerprints is centered and normalized according to:
        ``zeta = (zeta - mean(zeta)) / stdev(zeta)``

    dtype: np.dtype
        Data type for the generated fingerprints, such as ``np.float32`` and
        ``np.float64``.

    Example
    -------
    >>> cut_name = 'cos'
    >>> cut_dists = {'C-C': 5.0, 'C-H': 4.5, 'H-H': 4.0}
    >>> hyperparams = {'jmax': 4, 'weight': {'C':1.0, 'H':1.0}}
    >>> desc = Bispectrum(cut_dists, cut_name, hyperparams)

    References
    ----------
    .. [Bartok2010] Bartók, Albert P., Mike C. Payne, Risi Kondor, and Gábor Csányi.
       "Gaussian approximation potentials: The accuracy of quantum mechanics, without
       the electrons." Physical review letters 104, no. 13 (2010): 136403.
    .. [Thompson2015] Thompson, Aidan P., Laura P. Swiler, Christian R. Trott,
       Stephen M. Foiles, and Garritt J. Tucker. "Spectral neighbor analysis method
       for automated generation of quantum-accurate interatomic potentials." Journal
       of Computational Physics 285 (2015): 316-330.
    """

    def __init__(
        self,
        cut_dists,
        cut_name=None,
        hyperparams=None,
        normalize=True,
        dtype=np.float32,
    ):
        super(Bispectrum, self).__init__(
            cut_dists, cut_name, hyperparams, normalize, dtype
        )

        self.update_hyperparams(self.hyperparams)

        # init cdesc
        rfac0 = self.hyperparams["rfac0"]
        jmax = self.hyperparams["jmax"]
        diagonalstyle = self.hyperparams["diagonalstyle"]
        rmin0 = self.hyperparams["rmin0"]
        switch_flag = self.hyperparams["switch_flag"]
        bzero_flag = self.hyperparams["bzero_flag"]
        use_shared_arrays = 0

        self._cdesc = bs.Bispectrum(
            rfac0,
            2 * jmax,
            diagonalstyle,
            use_shared_arrays,
            rmin0,
            switch_flag,
            bzero_flag,
        )

        self._set_cutoff()

        self._set_hyperparams()

    def transform(self, conf, grad=False):
        # neighbor list
        infl_dist = max(self.cutoff.values())

        nei = NeighborList(conf, infl_dist, padding_need_neigh=False)

        coords = nei.coords
        image = nei.image
        species = np.asarray([self.species_code[i] for i in nei.species], dtype=np.intc)

        numneigh, neighlist = nei.get_numneigh_and_neighlist_1D()

        Natoms = len(coords)
        Ncontrib = conf.get_num_atoms()
        Ndesc = self.get_size()

        if grad:
            zeta, dzeta_dr = self._cdesc.compute_zeta_and_dzeta_dr(
                coords, species, neighlist, numneigh, image, Natoms, Ncontrib, Ndesc
            )
            # reshape to 4D array
            dzeta_dr = dzeta_dr.reshape(Ncontrib, Ndesc, Ncontrib, 3)
        else:
            zeta = self._cdesc.compute_zeta(
                coords, species, neighlist, numneigh, image, Natoms, Ncontrib, Ndesc
            )
            dzeta_dr = None

        logger.debug(
            "\n" + "=" * 25 + "descriptor values (no normalization)" + "=" * 25
        )

        logger.debug(f"\nconfiguration name: {conf.identifier}")
        logger.debug("\natomid    descriptor values ...")

        for i, line in enumerate(zeta):
            s = f"\n{i}    "
            for j in line:
                s += f"{j:.15g} "
            logger.debug(s)

        return zeta, dzeta_dr

    def update_hyperparams(self, params):
        """Update the hyperparameters based on the input at initialization."""
        default_hyperparams = {
            "jmax": 4,
            "rfac0": 0.99363,
            "diagonalstyle": 3,
            "rmin0": 0,
            "switch_flag": 1,
            "bzero_flag": 0,
            "weight": None,
        }

        if params is not None:
            for key, value in params.items():
                if key not in default_hyperparams:
                    raise BispectrumError(
                        'Hyperparameter "{}" not supported by this descriptor.'.format(
                            key
                        )
                    )
                else:
                    default_hyperparams[key] = value

        self.hyperparams = default_hyperparams

    def _set_cutoff(self):
        supported = ["cos"]

        if self.cut_name is None:
            self.cut_name = supported[0]

        if self.cut_name not in supported:
            spd = ['"{}", '.format(s) for s in supported]
            raise BispectrumError(
                'Cutoff "{}" not supported by this descriptor. Use {}.'.format(
                    self.cut_name, spd
                )
            )

        self.cutoff = generate_full_cutoff(self.cut_dists)

        self.species_code = generate_species_code(self.cut_dists)

        num_species = len(self.species_code)

        rcutsym = np.zeros([num_species, num_species], dtype=np.double)

        for si, i in self.species_code.items():
            for sj, j in self.species_code.items():
                rcutsym[i][j] = self.cutoff[si + "-" + sj]

        self._cdesc.set_cutoff(self.cut_name, rcutsym)

    def _set_hyperparams(self):
        weight_in = self.hyperparams["weight"]

        if weight_in is None:
            weight = np.ones(len(self.species_code), dtype=np.double)
        else:
            weight = np.zeros(len(self.species_code), dtype=np.double)

            for spec, code in self.species_code.items():
                try:
                    weight[code] = weight_in[spec]
                except KeyError:
                    raise BispectrumError(
                        '"weight" for species "{}" not provided.'.format(spec)
                    )

        self._cdesc.set_weight(weight)

    def get_size(self):
        """
        Return the size of descriptor.
        """
        diagonal = self.hyperparams["diagonalstyle"]
        twojmax = int(2 * self.hyperparams["jmax"])

        N = 0
        for j1 in range(0, twojmax + 1):
            if diagonal == 2:
                N += 1
            elif diagonal == 1:
                for j in range(0, min(twojmax, 2 * j1) + 1, 2):
                    N += 1
            elif diagonal == 0:
                for j2 in range(0, j1 + 1):
                    for j in range(j1 - j2, min(twojmax, j1 + j2) + 1, 2):
                        N += 1
            elif diagonal == 3:
                for j2 in range(0, j1 + 1):
                    for j in range(j1 - j2, min(twojmax, j1 + j2) + 1, 2):
                        if j >= j1:
                            N += 1
        return N


class BispectrumError(Exception):
    def __init__(self, msg):
        super(BispectrumError, self).__init__(msg)

        self.msg = msg
