
#  Copyright (C) 2023 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell 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.
#
#  Py-ChemShell 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 Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

import numpy as np
from numpy.typing import NDArray
from ctypes import Structure, c_long
from typing import Tuple, Union, Literal, List

from chemsh.objects.fragbase import _FragmentBase
from chemsh.objects.result import Result

from .Layer import Layer
from .HybridHelperFunctions import CorrectGradientsForLinkAtoms, ComputeFragsLinksIndices#, RecursiveGradientsResolve


class NLayerSubtractive(Structure):

    _result = Result()
    result = _result
    _gradients = False

    """
    # This is all just extras needed to get the class to work with Opt
    """
    input = None
    output = None
    _extra_outputs = None
    _kwargs = {}
    restart = False
    """
    # End of undeeded extras
    """

    def __init__(
        self,
        frag: _FragmentBase,
        layers: Tuple[Layer, ...],
        embedding: Union[Tuple[Union[Literal["mechanical"], Literal["electrostatic"]], ...], None] = None,
        coupling: Union[Tuple[Union[Literal["covalent"], Literal["ionic"]], ...], None] = None,
        strict_checks: bool = True,
    ) -> None:

        if embedding != None:
            if "electrostatic" in embedding:
                raise NotImplementedError(
                    'embedding "electrostatic" is currently not available'
                )

        if coupling != None:
            if "ionic" in coupling:
                raise NotImplementedError('coupling "ionic" is currently not available')

        # Check there are atleast 2 layers defined
        if len(layers) < 2:
            raise RuntimeError(
                "The number of layers passed into NLayerSubtractive, must be >= 2"
            )

        # Set the fragment
        self.frag: _FragmentBase = frag

        # We expect the structure of the layers to be (innermost (highest), ..., outermost (lowest))
        self.nlayers: int = len(layers)
        self.layers: Tuple[Layer, ...] = layers

        # Set the embedding and coupling parameters TODO: ensure len of these parameters == 1 or nlayers-1
        # defaults to "mechanical" and "covalent" if not supplied
        if embedding != None:
            self.embedding: Tuple[
                Union[Literal["mechanical"], Literal["electrostatic"]], ...
            ] = embedding
        else:
            self.embedding: Tuple[
                Union[Literal["mechanical"], Literal["electrostatic"]], ...
            ] = tuple(["mechanical" for _ in range(self.nlayers)])
        if coupling != None:
            self.coupling: Tuple[
                Union[Literal["covalent"], Literal["ionic"]], ...
            ] = coupling
        else:
            self.coupling: Tuple[
                Union[Literal["covalent"], Literal["ionic"]], ...
            ] = tuple(["covalent" for _ in range(self.nlayers)])

        # Get all the Frags, Links and Indices
        self.frags, self.links, self.indices = ComputeFragsLinksIndices(
            self.frag, self.layers
        )

        # Perform some sanity checks
        if strict_checks:
            # Check there are more atoms in the lower layers than in the higher layers
            prev_natoms: int = 0
            for f in self.frags:
                if f.natoms <= prev_natoms:
                    raise RuntimeError(
                        "NLayerSubtractive expects the highest layer to be listed first and have the least number of atoms!"
                    )
                prev_natoms: int = f.natoms

            # Check indices in the higher layers are also contained in the lower layers
            for i, _ in enumerate(self.indices[1:], start=1):
                if not np.all(np.isin(self.indices[i - 1], self.indices[i])):
                    raise RuntimeError(
                        f"Elements from layer {i} are not in layer {i-1}"
                    )

            # Check that the number of frags is the same as the number of layers
            if len(self.frags) != len(self.layers):
                raise RuntimeError("Number of fragments and layers must be equal")


    def run(self, **args) -> c_long:

        # TWK TODO: Consider use of SP here in comparison with other drivers which call .run
        #           Appears SP is necessary for use of finite difference gradients; we should
        #           decide what the consistent solution should be across the QM/MM schemes etc.
        from chemsh.tasks.sp import SP
        import copy

        print('\n '+'='*70)
        print(" ChemShell subtractive embedding driver")
        print(' '+'='*70)

        # Update all the Frags, Links and Indices
        self.frags, self.links, self.indices = ComputeFragsLinksIndices(
            self.frag, self.layers
        )

        try:
            self._gradients: bool = args["_gradients"]
        except:
            self._gradients: bool = True

        # This stores all the terms that are added
        addtn_energies: List[float] = []
        addtn_gradients: List[NDArray[np.float64]] = []

        # To output energy info
        eprt = ""

        # This computes all the energies and gradients and stores them
        for i in range(self.nlayers):  # Runs over methods
            print("\n >>> Layer %1d (+)" % i)
            self.layers[i].method.frag = copy.deepcopy(self.frags[i])
            if (self.layers[i].gradients_eval != "analytic"):
                self.layers[i].method.restart = False
            sp: SP = SP(theory=self.layers[i].method, gradients=self.layers[i].gradients_eval)
            sp.run(gradients=self._gradients)

            # store the energies to add and the relevant gradients
            addtn_energies: List[float] = addtn_energies + [sp.result.energy[0]]
            addtn_gradients: List[NDArray[np.float64]] = addtn_gradients + [
                np.copy(sp.result.gradients)
            ]

            # Add to output
            eprt += " Layer %1d energy (+)       = %25.12f a.u.\n" % (i, sp.result.energy[0])

        # This stores all the terms that are subtracted
        subtr_energies: List[float] = []
        subtr_gradients: List[NDArray[np.float64]] = []

        # In this case, the i-th frag ran by the j-th method
        for i in range(self.nlayers - 1):
            print("\n >>> Layer %1d (-)" % i)
            j: int = i + 1
            self.layers[j].method.frag = copy.deepcopy(self.frags[i])

            # If using MM, the user MUST provide link atom types... 
            if self.layers[j].method_type == "MM":
                nlinks: int = self.links[i].natoms
                if self.layers[j].link_atoms == None and nlinks != 0:
                    raise RuntimeError("User has defined an MM layer that requires a link atom, but provided no link atom type! ChemShell will not assume a link atom in this case")
                if self.layers[j].link_atoms != None:
                    if len(self.layers[j].link_atoms) != nlinks:
                        raise RuntimeError(f"Fragment requires {nlinks} link atoms, but user has supplied {len(self.layers[j].link_atoms)}.")
                    natoms: int = self.layers[j].method.frag.natoms
                    start_links: int = natoms - nlinks
                    for l, link_idx in enumerate(range(start_links,natoms)):
                        #current_link_symbol = self.layers[j].method.frag.symbols[link_idx]
                        new_link_symbol = self.layers[j].link_atoms[l]
                        self.layers[j].method.frag.names[link_idx] = new_link_symbol
                        new_link_charge = self.layers[j].link_atom_charges[l]
                        self.layers[j].method.frag.charges[link_idx] = new_link_charge

            # Restart flags can prevent correct finite-differences
            if (self.layers[j].gradients_eval != "analytic"):
                self.layers[j].method.restart = False

            sp: SP = SP(theory=self.layers[j].method, gradients=self.layers[j].gradients_eval)
            sp.run(gradients=self._gradients)

            # store the energies to subtract and the relevant gradients
            subtr_energies: List[float] = subtr_energies + [sp.result.energy[0]]
            subtr_gradients: List[NDArray[np.float64]] = subtr_gradients + [
                np.copy(sp.result.gradients)
            ]

            # Add to output
            eprt += " Layer %1d energy (-)       = %25.12f a.u.\n" % (i, -sp.result.energy[0])

        # Compute the total energy
        total_energy: float = np.sum(addtn_energies) - np.sum(subtr_energies)
        self._result.energy = total_energy
        eprt += " Total subtractive energy = %25.12f a.u." % self._result.energy[0]

        # Print energy output
        print("\n "+"*"*70)
        print(eprt)
        print(" "+"*"*70)
        print("")

        if self._gradients:
            # First Correct all the link atom gradients 
            total_gradients = np.copy(addtn_gradients[-1])

            for i, _ in enumerate(subtr_gradients):

                trunc_len: int = len(self.indices[i])

                dE_dr = np.zeros(addtn_gradients[-1].shape)
                dE_dr[self.indices[i],:] += addtn_gradients[i][:trunc_len,:]
                addtn_g = CorrectGradientsForLinkAtoms(self.frag, self.links[i], dE_dr, addtn_gradients[i][trunc_len:,:])

                dE_dr = np.zeros(addtn_gradients[-1].shape)
                dE_dr[self.indices[i],:] += subtr_gradients[i][:trunc_len,:]
                subtr_g = CorrectGradientsForLinkAtoms(self.frag, self.links[i], dE_dr, subtr_gradients[i][trunc_len:,:])

                total_gradients += (addtn_g - subtr_g)

            self._result.gradients = total_gradients

        return c_long(0)
