# This file is part of tblite.
# SPDX-Identifier: LGPL-3.0-or-later
#
# tblite is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# tblite 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 Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with tblite.  If not, see <https://www.gnu.org/licenses/>.
"""
The Python API of *tblite* natively support integration with the atomic
simulation environment (`ASE`_).
By constructing a calculator most functionality of ASE is readily available.
For details on building the Python API checkout the
:ref:`installation guide <python-build>`.

.. _ase: https://wiki.fysik.dtu.dk/ase/

Example
-------

The calculator can be used like other ASE calculators by attaching it to an atoms object.

>>> from ase.build import molecule
>>> from tblite.ase import TBLite
>>> atoms = molecule("H2O")
>>> atoms.calc = TBLite(method="GFN2-xTB")
>>> atoms.get_potential_energy()  # in eV
-137.96777594361677

The calculator will perform single point calculations as necessary, if properties for
the same atoms are requested the properties will be returned without additional single point
calculations.

>>> atoms.get_dipole_moment()  # in Debye
array([-7.83154860e-17,  1.10157136e-16, -4.76020974e-01])
>>> atoms.get_forces()  # in eV/A
array([[-3.31859412e-15,  7.83871951e-15, -7.49527168e-01],
       [ 3.08351626e-15, -1.53506496e-01,  3.74763584e-01],
       [ 2.35077857e-16,  1.53506496e-01,  3.74763584e-01]])

Changed properties of the atoms object will lead to an reevalution of the single point,
for example updating the initial charges of the atoms to set a total charge.

>>> atoms.set_initial_charges([0, 0, -1])
>>> atoms.get_potential_energy()  # in eV
-131.50015843891816

Some properties, like the total charge, can also be set via the parameters of the
calculator.

>>> atoms = molecule("H2O")
>>> atoms.calc = TBLite(method="GFN2-xTB", charge=-1)
>>> atoms.get_potential_energy()  # in eV
-131.50015843891822

Finally, the parameters of the calculator can be updated at any time using the *set*
method, which resets the calculation as necessary.

>>> atoms.calc.set(method="GFN1-xTB", charge=0)
>>> atoms.get_potential_energy()  # in eV
-156.96750590276173
>>> atoms.get_dipole_moment()  # in Debye
array([ 3.71500789e-16,  3.08439980e-16, -5.97048616e-01])
>>> atoms.get_forces()  # in eV/A
array([[ 2.88997963e-15,  5.70620869e-15, -7.84252969e-01],
       [-4.90812472e-15, -2.28835628e-01,  3.92126485e-01],
       [ 2.01814509e-15,  2.28835628e-01,  3.92126485e-01]])

The calculator also support atom partitioned properties, including energies and charges.

>>> atoms.get_charges()  # in e
array([-0.66558478,  0.33279239,  0.33279239])
>>> atoms.get_potential_energies()  # in eV
array([-132.37935247,  -12.29407672,  -12.29407672])
"""

try:
    import ase
    import ase.calculators.calculator
    from ase.atoms import Atoms
    from ase.units import Bohr, Hartree, kB
except ModuleNotFoundError as e:
    raise ModuleNotFoundError("This submodule requires ASE installed") from e

from typing import Any, Dict, List, Optional

from .interface import Calculator


class TBLite(ase.calculators.calculator.Calculator):
    r"""
    ASE calculator for using xTB Hamiltonians from the tblite library.
    Supported properties by this calculator are:

    - *energy* (*free_energy*): The total energy of the system (in eV)
    - *energies*: The atom-partitioned energy of the system (in eV)
    - *forces*: The derivative of the energy with respect to the atomic positions (in eV/Ang)
    - *stress*: The derivative of the energy with respect to cell deformations (PBC only, in eV/Ang³)
    - *dipole*: The dipole moment of the system (in Debye)
    - *charges*: Mulliken charges of the system (in e)

    Supported keywords are

    ======================== ================= ====================================================
     Keyword                  Default           Description
    ======================== ================= ====================================================
     method                   "GFN2-xTB"        Underlying method for energy and forces
     charge                   None              Total charge of the system
     multiplicity             None              Total multiplicity of the system
     accuracy                 1.0               Numerical accuracy of the calculation
     electronic_temperature   300.0             Electronic temperatur in Kelvin
     max_iterations           250               Iterations for self-consistent evaluation
     mixer                    "broyden"         Self-consistent mixer
     mixer_memory             0                 Number of mixer history vectors (0 uses max_iterations)
     initial_guess            "sad"             Initial guess for wavefunction (sad, eeq, or eeqbc)
     mixer_damping            0.4               Damping parameter for self-consistent mixer
     annealing                None              Initial temperature or (initial, hold, cycles)
     electric_field           None              Uniform electric field vector (in V/A)
     spin_polarization        None              Spin polarization (scaling factor)
     solvation                None              Solvation model to use (see below for details)
     cache_api                True              Reuse generate API objects (recommended)
     verbosity                1                 Set verbosity of printout
     xtb_config               None              Additional configuration for the API calculator (dict)
    ======================== ================= ====================================================

    Solvation models are supported by passing a tuple to the *solvation* keyword.
    The first element of the tuple is the name of the solvation model, the following
    arguments are passed to initialize the solvation model.

    ================== ====================================================
     Solvation model    Arguments
    ================== ====================================================
     "alpb"             Solvent name (str), solution state (optional)
     "gbsa"             Solvent name (str), solution state (optional)
     "cosmo"            Epsilon (float) / Solvent name (str)
     "cpcm"             Epsilon (float) / Solvent name (str)
     "pcm"              Epsilon (float) / Solvent name (str)
     "gbe"              Epsilon (float), Born kernel
     "gb"               Epsilon (float), Born kernel
    ================== ====================================================

    Example
    -------

    An ASE calculator can be constructed by using the *TBLite* class provided
    by the *tblite.ase* module.
    For example to perform a single point calculation for a CO\ :sub:`2`
    crystal use

    >>> from tblite.ase import TBLite
    >>> from ase.atoms import Atoms
    >>> import numpy as np
    >>> atoms = Atoms(
    ...     symbols="C4O8",
    ...     positions=np.array(
    ...         [
    ...             [0.9441259872, 0.9437851680, 0.9543505632],
    ...             [3.7179966528, 0.9556570368, 3.7316862240],
    ...             [3.7159517376, 3.7149292800, 0.9692330016],
    ...             [0.9529872864, 3.7220864832, 3.7296981120],
    ...             [1.6213905408, 1.6190616096, 1.6313879040],
    ...             [0.2656685664, 0.2694175776, 0.2776540416],
    ...             [4.3914553920, 1.6346256864, 3.0545920000],
    ...             [3.0440834880, 0.2764611744, 4.4080419264],
    ...             [4.3910577696, 3.0416409504, 0.2881058304],
    ...             [3.0399936576, 4.3879335936, 1.6497353376],
    ...             [0.2741322432, 4.4003734944, 3.0573754368],
    ...             [1.6312174944, 3.0434586528, 4.4023048032],
    ...         ]
    ...     ),
    ...     cell=np.array([5.68032, 5.68032, 5.68032]),
    ...     pbc=np.array([True, True, True]),
    ... )
    >>> atoms.calc = TBLite(method="GFN1-xTB", verbosity=0)
    >>> atoms.get_potential_energy()  # result in eV
    -1257.0943962462964

    The resulting calculator can be used like most ASE calculator, *e.g.* for optimizing geometries.
    """

    implemented_properties = [
        "energy",
        "energies",
        "forces",
        "charges",
        "dipole",
        "stress",
    ]

    default_parameters = {
        "method": "GFN2-xTB",
        "charge": None,
        "multiplicity": None,
        "accuracy": 1.0,
        "guess": "sad",
        "max_iterations": 250,
        "mixer": "broyden",
        "mixer_memory": 0,
        "mixer_damping": 0.4,
        "annealing": None,
        "electric_field": None,
        "spin_polarization": None,
        "solvation": None,
        "electronic_temperature": 300.0,
        "cache_api": True,
        "verbosity": 1,
        "xtb_config": None,
    }

    _res = None
    _xtb = None

    def __init__(
        self,
        atoms: Optional[Atoms] = None,
        **kwargs,
    ):
        """
        Construct the TBLite base calculator object.
        """

        ase.calculators.calculator.Calculator.__init__(self, atoms=atoms, **kwargs)

    def set(self, **kwargs) -> dict:
        """
        Set new parameters to TBLite. Will automatically reconstruct the underlying
        model in case critical parameters change.

        Example
        -------
        >>> from ase.build import molecule
        >>> from tblite.ase import TBLite
        >>> atoms = molecule("H2O")
        >>> atoms.calc = TBLite(method="GFN2-xTB")
        >>> atoms.get_potential_energy()
        -137.96777625229421
        >>> atoms.calc.set(method="GFN1-xTB")
        {'method': 'GFN1-xTB'}
        >>> atoms.get_potential_energy()
        -156.9675057724589
        """

        _update_parameters(kwargs)

        changed_parameters = ase.calculators.calculator.Calculator.set(self, **kwargs)

        # Always reset the calculation if parameters change
        if changed_parameters:
            self.reset()

        # If the method is changed, invalidate the cached calculator as well
        if (
            "method" in changed_parameters
            or "electric_field" in changed_parameters
            or "spin_polarization" in changed_parameters
            or "solvation" in changed_parameters
            or "annealing" in changed_parameters
            or "xtb_config" in changed_parameters
        ):
            self._xtb = None
            self._res = None

        # Minor changes can be updated in the API calculator directly
        if self._xtb is not None:
            if "accuracy" in changed_parameters:
                self._xtb.set("accuracy", self.parameters.accuracy)

            if "electronic_temperature" in changed_parameters:
                self._xtb.set(
                    "temperature",
                    self.parameters.electronic_temperature * kB / Hartree,
                )

            if "max_iterations" in changed_parameters:
                self._xtb.set("max-iter", self.parameters.max_iterations)

            if "mixer" in changed_parameters:
                self._xtb.set("mixer", self.parameters.mixer)

            if "mixer_memory" in changed_parameters:
                self._xtb.set("mixer-memory", self.parameters.mixer_memory)

            if "initial_guess" in changed_parameters:
                self._xtb.set(
                    "guess", {"sad": 0, "eeq": 1, "eeqbc": 2}[self.parameters.guess]
                )

            if "mixer_damping" in changed_parameters:
                self._xtb.set("mixer-damping", self.parameters.mixer_damping)

            if "annealing" in changed_parameters:
                if self.parameters.annealing is not None:
                    self._xtb.set(
                        "annealing",
                        self.parameters.annealing,
                    )

            if (
                "charge" in changed_parameters or "multiplicity" in changed_parameters
            ) and self.atoms is not None:
                self._xtb.update(
                    charge=_get_charge(self.atoms, self.parameters),
                    uhf=_get_uhf(self.atoms, self.parameters),
                )

        return changed_parameters

    def reset(self) -> None:
        """
        Clear all information from old calculation. This will only remove the cached
        API objects in case the `cache_api` is set to False.
        """
        ase.calculators.calculator.Calculator.reset(self)

        if not self.parameters.cache_api:
            self._xtb = None
            self._res = None

    def _check_api_calculator(self, system_changes: List[str]) -> None:
        """Check state of API calculator and reset if necessary"""

        # Changes in positions, cell, charges and magnetic moments can use a normal update
        _reset = system_changes.copy()
        for _change in system_changes:
            if _change in ("positions", "cell", "initial_charges", "initial_magmoms"):
                _reset.remove(_change)

        # Invalidate cached calculator and results object
        if _reset:
            self._xtb = None
            self._res = None
        else:
            if system_changes and self._xtb is not None:
                try:
                    _cell = self.atoms.cell
                    self._xtb.update(
                        positions=self.atoms.positions / Bohr,
                        lattice=_cell / Bohr,
                        charge=_get_charge(self.atoms, self.parameters),
                        uhf=_get_uhf(self.atoms, self.parameters),
                    )
                # An exception in this part means the geometry is bad,
                # still we will give a complete reset a try as well
                except RuntimeError:
                    self._xtb = None
                    self._res = None

    def calculate(
        self,
        atoms: Optional[Atoms] = None,
        properties: Optional[List[str]] = None,
        system_changes: List[str] = ase.calculators.calculator.all_changes,
    ) -> None:
        """
        Perform actual calculation by calling the TBLite API

        Example
        -------

        >>> from ase.build import molecule
        >>> from tblite.ase import TBLite
        >>> calc = TBLite(method="GFN2-xTB")
        >>> calc.calculate(molecule("H2O"))
        >>> calc.get_potential_energy()
        -137.96777625229421
        >>> calc.calculate(molecule("CH4"))
        >>> calc.get_potential_energy()
        -113.60956621093894

        Raises
        ------
        ase.calculators.calculator.InputError
            on invalid input passed to the interface module

        ase.calculators.calculator.CalculationFailed
            in case of an `RuntimeError` in the library
        """

        if not properties:
            properties = ["energy"]
        ase.calculators.calculator.Calculator.calculate(
            self, atoms, properties, system_changes
        )

        self._check_api_calculator(system_changes)

        if self._xtb is None:
            self._xtb = _create_api_calculator(self.atoms, self.parameters)

        try:
            self._res = self._xtb.singlepoint(self._res)
        except RuntimeError as e:
            raise ase.calculators.calculator.CalculationFailed(str(e)) from e

        # These properties are garanteed to exist for all implemented calculators
        self.results["energy"] = self._res["energy"] * Hartree
        self.results["energies"] = self._res["energies"] * Hartree
        self.results["free_energy"] = self.results["energy"]
        self.results["forces"] = -self._res["gradient"] * Hartree / Bohr
        self.results["charges"] = self._res["charges"]
        self.results["dipole"] = self._res["dipole"] * Bohr
        self.results["bond-orders"] = self._res["bond-orders"]
        # stress tensor is only returned for periodic systems
        if self.atoms.pbc.any():
            _stress = self._res["virial"] * Hartree / self.atoms.get_volume()
            self.results["stress"] = _stress.flat[[0, 4, 8, 5, 2, 1]]


def _create_api_calculator(
    atoms: ase.Atoms,
    parameters: ase.calculators.calculator.Parameters,
) -> Calculator:
    """Create a new API calculator object"""

    try:
        _cell = atoms.cell
        _periodic = atoms.pbc
        _charge = _get_charge(atoms, parameters)
        _uhf = _get_uhf(atoms, parameters)

        calc = Calculator(
            parameters.method,
            atoms.numbers,
            atoms.positions / Bohr,
            _charge,
            _uhf,
            _cell / Bohr,
            _periodic,
            xtb_config=parameters.xtb_config,
        )
        calc.set("accuracy", parameters.accuracy)
        calc.set(
            "temperature",
            parameters.electronic_temperature * kB / Hartree,
        )
        calc.set("max-iter", parameters.max_iterations)
        calc.set("mixer", parameters.mixer)
        calc.set("mixer-memory", parameters.mixer_memory)
        calc.set("guess", {"sad": 0, "eeq": 1, "eeqbc": 2}[parameters.guess])
        calc.set("mixer-damping", parameters.mixer_damping)
        if parameters.annealing is not None:
            calc.set("annealing", parameters.annealing)
        calc.set("verbosity", parameters.verbosity)
        if parameters.electric_field is not None:
            calc.add(
                "electric-field",
                parameters.electric_field * Bohr / Hartree,
            )
        if parameters.spin_polarization is not None:
            calc.add("spin-polarization", parameters.spin_polarization)
        if parameters.solvation is not None:
            solvation_model, *solvation_args = parameters.solvation
            if isinstance(solvation_args, (tuple, list)):
                calc.add(f"{solvation_model}-solvation", *solvation_args)
            else:
                calc.add(f"{solvation_model}-solvation", solvation_args)

    except RuntimeError as e:
        raise ase.calculators.calculator.InputError(str(e)) from e

    return calc


def _get_charge(
    atoms: ase.Atoms, parameters: ase.calculators.calculator.Parameters
) -> int:
    """
    Get the total charge of the system.
    If no charge is provided, the total charge of the system is calculated
    by summing the initial charges of all atoms.
    """
    return (
        atoms.get_initial_charges().sum()
        if parameters.charge is None
        else parameters.charge
    )


def _get_uhf(
    atoms: ase.Atoms, parameters: ase.calculators.calculator.Parameters
) -> int:
    """
    Get the number of unpaired electrons.
    If no multiplicity is provided, the number of unpaired electrons
    is calculated by summing the initial magnetic moments of all atoms.
    """
    return (
        int(atoms.get_initial_magnetic_moments().sum().round())
        if parameters.multiplicity is None
        else parameters.multiplicity - 1
    )


def _update_parameters(kwargs: Dict[str, Any]) -> None:
    """
    Update the parameters of the calculator with the provided keyword arguments.
    This function is used to update the parameters of the calculator
    when new values are provided.
    """

    ddx_models = ("cosmo", "cpcm", "pcm")

    for key in (
        "alpb",
        "gbsa",
        "cosmo",
        "cpcm",
        "pcm",
        "gbe",
        "gb",
    ):
        if f"{key}_solvation" in kwargs:
            value = kwargs.pop(f"{key}_solvation")

            if key in ddx_models:
                value = ("ddX", value, key)
            elif isinstance(value, (tuple, list)):
                value = (key, *value)
            else:
                value = (key, value)

            if "solvation" in kwargs:
                raise ase.calculators.calculator.InputError(
                    "Multiple solvation models provided, can only use one"
                )

            kwargs["solvation"] = value


if "tblite" not in ase.calculators.calculator.external_calculators:
    ase.calculators.calculator.register_calculator_class("tblite", TBLite)
