#  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/>.

from typing import Literal, Union, List, Iterable
import numpy as np
from numpy.typing import NDArray
from ..interfaces.mm import _MM
from ..interfaces.qm import _QM

from ..interfaces import _Interfaces


class Layer:
    """
    The layer dataclass is used to store the method of a given layer
    and the atoms on which it applies
    """

    # Stores the method to use for this layer
    method: _Interfaces
    # Defines the atom indices in the layer,
    # if "[-1]" use remaining leftover atoms
    atom_indices: NDArray[np.int32]
    # Define how the force will be computed
    gradients_eval: Union[Literal["analytic"], Literal["finitediff"]]
    method_type: Union[Literal["MM"], Literal["QM"]]
    link_atoms: Union[List[str], None]
    link_atom_charges: Union[List[float], None]

    def __init__(
        self,
        method: _Interfaces,
        atom_indices: Union[NDArray[np.int32], Iterable[int]] = np.array([-1]),
        gradients_eval: Union[Literal["analytic"], Literal["finitediff"]] = "analytic",
        link_atoms: Union[List[str], None] = None,
        link_atom_charges: Union[List[float], None] = None,
    ) -> None:

        self.method: _Interfaces = method
        self.atom_indices: NDArray[np.int32] = np.array(atom_indices, dtype=np.int32)
        self.gradients_eval = gradients_eval
        self.link_atoms = link_atoms
        self.link_atom_charges = link_atom_charges
        self.method_type = "MM"
        if issubclass(self.method.__class__, _MM):
            self.method_type = "MM"
        elif issubclass(method.__class__, _QM):
            self.method_type = "QM"
        else:
            raise RuntimeError("layer.method must inherit _QM or _MM")

        if link_atoms != None and link_atom_charges == None:
            raise ValueError(
                "link_atoms have been defined, therefore the link_atom_charges need also be defined"
            )

        if link_atoms == None and link_atom_charges != None:
            raise ValueError(
                "link_atoms has not been defined, but link_atom_charges has, please define link_atoms"
            )

        if link_atoms != None and link_atom_charges != None:
            if len(link_atoms) != len(link_atom_charges):
                raise ValueError(
                    "The number of link_atoms != the number of link_atom_charges"
                )

        if np.all(self.atom_indices != -1):
            for i, idx in enumerate(self.atom_indices):
                if idx < 0:
                    raise ValueError(
                        f"Index {i} of atom_indices out of range "
                        + "must be postive (currently {idx})"
                    )
