# -*- coding: utf-8 -*-
# HORTON-GRID: GRID for Helpful Open-source Research TOol for N-fermion systems.
# Copyright (C) 2011-2023 The HORTON-GRID Development Team
#
# This file is part of HORTON-GRID
#
# HORTON-GRID is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# HORTON-GRID is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>
#
# --
"""Molecular integration grids"""


import numpy as np

from horton_grid.grid.base import IntGrid
from horton_grid.grid.atgrid import AtomicGrid, AtomicGridSpec
from horton_grid.grid.cext import becke_helper_atom
from horton_grid.log import log, timer, biblio
from horton_grid.periodic import periodic
from horton_grid.utils import typecheck_geo, doc_inherit


__all__ = ["BeckeMolGrid"]


class BeckeMolGrid(IntGrid):
    """Molecular integration grid using Becke weights"""

    @timer.with_section("Becke-Lebedev")
    def __init__(
        self,
        centers,
        numbers,
        pseudo_numbers=None,
        agspec="medium",
        k=3,
        random_rotate=True,
        mode="discard",
    ):
        """
        **Arguments:**

        centers
             An array (N, 3) with centers for the atom-centered grids.

        numbers
             An array (N,) with atomic numbers.

        **Optional arguments:**

        pseudo_numbers
             An array (N,) with effective core charges. When not given, this
             defaults to ``numbers``.

        agspec
             A specifications of the atomic grid. This can either be an
             instance of the AtomicGridSpec object, or the first argument
             of its constructor.

        k
             The order of the switching function in Becke's weighting scheme.

        random_rotate
             Flag to control random rotation of spherical grids.

        mode
             Select one of the following options regarding atomic subgrids:

             * ``'discard'`` (the default) means that all information about
               subgrids gets discarded.

             * ``'keep'`` means that a list of subgrids is kept, including
               the integration weights of the local grids.

             * ``'only'`` means that only the subgrids are constructed and
               that the computation of the molecular integration weights
               (based on the Becke partitioning) is skipped.
        """
        natom, centers, numbers, pseudo_numbers = typecheck_geo(centers, numbers, pseudo_numbers)
        self._centers = centers
        self._numbers = numbers
        self._pseudo_numbers = pseudo_numbers

        # check if the mode argument is valid
        if mode not in ["discard", "keep", "only"]:
            raise ValueError("The mode argument must be 'discard', 'keep' or 'only'.")

        # transform agspec into a usable format
        if not isinstance(agspec, AtomicGridSpec):
            agspec = AtomicGridSpec(agspec)
        self._agspec = agspec

        # assign attributes
        self._k = k
        self._random_rotate = random_rotate
        self._mode = mode

        # allocate memory for the grid
        size = sum(agspec.get_size(self.numbers[i], self.pseudo_numbers[i]) for i in range(natom))
        points = np.zeros((size, 3), float)
        weights = np.zeros(size, float)
        self._becke_weights = np.ones(size, float)

        # construct the atomic grids
        if mode != "discard":
            atgrids = []
        else:
            atgrids = None
        offset = 0

        if mode != "only":
            # More recent covalent radii are used than in the original work of Becke.
            # No covalent radius is defined for elements heavier than Curium and a
            # default value of 3.0 Bohr is used for heavier elements.
            cov_radii = np.array([(periodic[n].cov_radius or 3.0) for n in self.numbers])

        # The actual work:
        if log.do_medium:
            log("Preparing Becke-Lebedev molecular integration grid.")
        pb = log.progress(natom)
        for i in range(natom):
            atsize = agspec.get_size(self.numbers[i], self.pseudo_numbers[i])
            atgrid = AtomicGrid(
                self.numbers[i],
                self.pseudo_numbers[i],
                self.centers[i],
                agspec,
                random_rotate,
                points[offset : offset + atsize],
            )
            if mode != "only":
                atbecke_weights = self._becke_weights[offset : offset + atsize]
                becke_helper_atom(
                    points[offset : offset + atsize],
                    atbecke_weights,
                    cov_radii,
                    self.centers,
                    i,
                    self._k,
                )
                weights[offset : offset + atsize] = atgrid.weights * atbecke_weights
            if mode != "discard":
                atgrids.append(atgrid)
            offset += atsize
            pb()

        # finish
        IntGrid.__init__(self, points, weights, atgrids)

        # Some screen info
        self._log_init()

    @classmethod
    def from_hdf5(cls, grp):
        return BeckeMolGrid(
            grp["centers"][:],
            grp["numbers"][:],
            grp["psuedo_numbers"][:],
            AtomicGridSpec.from_hdf5(grp["agspec"]),
            grp["k"][()],
            grp["random_rotate"][()],
            grp.attrs["mode"],
        )

    def to_hdf5(self, grp):
        grp.attrs["class"] = self.__class__.__name__
        grp["centers"] = self._centers
        grp["numbers"] = self._numbers
        grp["psuedo_numbers"] = self._pseudo_numbers
        grp_agspec = grp.require_group("agspec")
        self.agspec.to_hdf5(grp_agspec, (self._numbers, self._pseudo_numbers))
        grp["random_rotate"] = self._random_rotate
        grp["k"] = self._k
        grp.attrs["mode"] = self._mode

    def _get_centers(self):
        """The positions of the nuclei"""
        return self._centers.view()

    centers = property(_get_centers)

    def _get_numbers(self):
        """The element numbers"""
        return self._numbers.view()

    numbers = property(_get_numbers)

    def _get_pseudo_numbers(self):
        """The effective core charges"""
        return self._pseudo_numbers.view()

    pseudo_numbers = property(_get_pseudo_numbers)

    def _get_agspec(self):
        """The specifications of the atomic grids."""
        return self._agspec

    agspec = property(_get_agspec)

    def _get_k(self):
        """The order of the Becke switching function."""
        return self._k

    k = property(_get_k)

    def _get_random_rotate(self):
        """The random rotation flag."""
        return self._random_rotate

    random_rotate = property(_get_random_rotate)

    def _get_mode(self):
        """The MO of this molecular grid"""
        return self._mode

    mode = property(_get_mode)

    def _get_becke_weights(self):
        """The becke weights of the grid points"""
        return self._becke_weights

    becke_weights = property(_get_becke_weights)

    def _log_init(self):
        if log.do_medium:
            log("Initialized: %s" % self)
            log.deflist(
                [
                    ("Size", self.size),
                    ("Switching function", "k=%i" % self._k),
                ]
            )
            log.blank()
        # Cite reference
        biblio.cite(
            "becke1988_multicenter",
            "the multicenter integration scheme used for the molecular integration grid",
        )
        biblio.cite(
            "cordero2008",
            "the covalent radii used for the Becke-Lebedev molecular integration grid",
        )

    @doc_inherit(IntGrid)
    def integrate(self, *args, **kwargs):
        if self.mode == "only":
            raise NotImplementedError(
                "When mode=='only', only the subgrids can be used for integration."
            )
        return IntGrid.integrate(self, *args, **kwargs)
