
#  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
import copy
from numpy.typing import NDArray
from typing import Tuple, Union, List
from chemsh.objects.links import Links
from chemsh.objects.fragbase import _FragmentBase
from .Layer import Layer

def GetBoundaries(
    frag: _FragmentBase, high_indices: NDArray[np.int32], low_indices: NDArray[np.int32]
) -> NDArray[np.int32]:

    # Get the connectivity of the high layer with the entire system
    high_conn: NDArray[np.int32] = frag.conn[high_indices, :]

    # Find connections between the high layer region and low region
    high_idx, low_idx = np.isin(high_conn, low_indices).nonzero()

    # make a list of all connections as atomic indices
    boundaries: NDArray[np.int32] = np.vstack(
        (high_indices[high_idx], high_conn[high_idx, low_idx])
    ).T

    return boundaries


def GetLinks(frag: _FragmentBase, boundaries: NDArray[np.int32]) -> Links:

    from chemsh.data import atomicdata

    links: Links = Links()

    links.nlinks = boundaries.shape[0]

    # Do not continue if there are no links
    links.bpairs = [[]]
    if links.nlinks == 0:
        return links

    # Set the atom pairs for which a bond is broken
    links.bpairs = [bpair for bpair in boundaries]

    # Define the indices of atom set A
    idx_a: NDArray[np.int32] = np.array(links.bpairs)[:, 0]
    # Get the coordinates in set A
    xyz_a: NDArray[np.float64] = frag.coords[idx_a, :]

    # Define the indices of atom set B
    idx_b: NDArray[np.int32] = np.array(links.bpairs)[:, 1]
    # Get the coordinates in set A
    xyz_b: NDArray[np.float64] = frag.coords[idx_b, :]

    # Get the Z values of atom set A
    znums_a: NDArray[np.int32] = frag.znums[idx_a]

    # Get the distance (r_HA) between atoms A and the link atoms
    links.qldists = atomicdata.linkdists[znums_a]

    # r_AB: compute the distance between A_i and B_i in sets A and B
    links.qmdists = np.linalg.norm(xyz_b - xyz_a, axis=1)

    # R_B - R_A
    links.qmdiffs = xyz_b - xyz_a

    # Set a vector to contain the values of g (= r_AH / r_AB)
    g: NDArray[np.float64] = np.array(links.qldists) / np.array(links.qmdists)

    # R_H = R_A + g*(R_B - R_A)
    links.coords = xyz_a + g[:, None] * links.qmdiffs

    # Set all the symbols to be hydrogen atoms
    links.symbols = np.array(["H" for _ in range(links.nlinks)], dtype=str)

    return links

def CorrectGradientsForLinkAtoms(
    frag: _FragmentBase,
    links: Links,
    dE_dr: NDArray[np.float64],
    dE_drL: NDArray[np.float64],
) -> NDArray[np.float64]:
    """
    Parameters
    ----------
    frag : _FragmentBase
        fragment object
    links : Links
        The properties of the linker atoms
    dE_dr : NDArray[np.float64] of size (natoms x 3)
        The gradients of the total system (excluding the linker atoms)
    dE_drL : NDArray[np.float64] of size (nlinks x 3)
        The gradients of the link atoms
    Returns
    -------
    NDArray[np.float64]:
        Gradients of the total system, with corrections for the linker atoms
    Notes
    -----
    The equations in this function are solved as in [Journal of Computational Chemistry, Vol 18, No 4, 463-477 (1996)]
    See equations: (21) and (22)
    """


    if links.natoms == 0:
        return dE_dr


    # First copy all the "original gradients" over for the actual system (not including link atoms)
    outgradients: NDArray[np.float64] = np.copy(dE_dr)

    # Get the indices for the inner region atom links (set A)
    idx_a: NDArray[np.int32] = np.array(links.bpairs)[:, 0]

    # Get the indices for the outer region atom links (set B)
    idx_b: NDArray[np.int32] = np.array(links.bpairs)[:, 1]

    if idx_b.size != idx_a.size:
        raise RuntimeError(
            "There should be the same number of atoms that defined cleaved bonds in the high layer and the low layer"
        )

    if idx_b.size != dE_drL.shape[0]:
        raise RuntimeError(
            "The number of link atom derivatives does not equal the number of cleaved atom bonds"
        )

    # Define the unit vectors: A->B
    n: NDArray[np.float64] = frag.coords[idx_b, :] - frag.coords[idx_a, :]
    n /= np.linalg.norm(n, axis=1)[:,None]

    # Set a vector to contain the values of g (= r_AH / r_AB), g should be of length nlinks
    g: NDArray[np.float64] = (links.qldists / links.qmdists)[:,None]

    # Precompute the required dot product
    dot_prod: NDArray[np.float64] = np.matmul( np.matmul(dE_drL, n.T), n )

    # Compute the gradients of the inner atoms, corrected for link atoms
    outgradients[idx_a, :] += g * dot_prod + (1.0 - g) * dE_drL

    # Compute the gradients of the outer atoms, corrected for link atoms
    outgradients[idx_b, :] += -g * dot_prod + g * dE_drL

    return outgradients


def GetModelFragmentWithLinkAtoms(
        frag: _FragmentBase, high_indices: NDArray[np.int32], low_indices: NDArray[np.int32]
) -> Tuple[_FragmentBase, Links]:

    model_frag: _FragmentBase = frag.getSelected(high_indices)

    boundaries: NDArray[np.int32] = GetBoundaries(frag, high_indices, low_indices)
    links: Links = GetLinks(frag, boundaries)

    model_frag_with_links: _FragmentBase = copy.deepcopy(model_frag).merge(links)

    return model_frag_with_links, links


def GetAllIndices(
    I: int, layers: Union[List[Layer], Tuple[Layer, ...]]
) -> NDArray[np.int32]:
    if I == 0:
        return layers[I].atom_indices
    else:
        return np.sort(np.append(layers[I].atom_indices, GetAllIndices(I - 1, layers)))


def ComputeFragsLinksIndices(
    frag: _FragmentBase, layers: Tuple[Layer, ...]
) -> Tuple[List[_FragmentBase], List[Links], List[NDArray[np.int32]]]:

    # Define a list of indices, the first contains all the lowest layer atoms
    indices: List[NDArray[np.int32]] = []
    link_atoms: List[Union[List[str], None]] = []
    for i in range(len(layers) - 1, -1, -1):
        indices = indices + [GetAllIndices(i, layers)]
        link_atoms = link_atoms + [layers[i].link_atoms]

    # Define the fragments, adding in link atoms where required
    frags: List[_FragmentBase] = [frag.getSelected(indices[0])]
    links: List[Links] = []
    for i, ind in enumerate(indices[1:]):
        temp_indices = np.setdiff1d(indices[0], ind)
        temp_frag, temp_links = GetModelFragmentWithLinkAtoms(frag, ind, temp_indices)
        frags = frags + [temp_frag]
        links = links + [temp_links]

    # we have worked from lowest layer to highest so far... need to reverse
    frags: List[_FragmentBase] = list(reversed(frags))
    links: List[Links] = list(reversed(links))
    indices: List[NDArray[np.int32]] = list(reversed(indices))

    return frags, links, indices


def RecursiveGradientsResolve(
    low_layer: int,
    global_indices: List[NDArray[np.int32]],
    addtn_gradients: List[NDArray[np.float64]],
    subtr_gradients: List[NDArray[np.float64]],
    links_list: List[Links],
    frag_list: List[_FragmentBase],
):

    # Define the high layer index
    high_layer: int = low_layer - 1

    # Define high/low layer indices
    high_indices: NDArray[np.int32] = global_indices[high_layer]
    low_indices: NDArray[np.int32] = global_indices[low_layer]
    # Find the overlap of indices
    _, high_in_low_indices, _ = np.intersect1d(
        low_indices, high_indices, return_indices=True
    )

    # high indices (not including link atom)
    high_len: int = len(high_indices)
    low_len: int = len(low_indices)

    if high_layer == 0:  # Stop recusion at the highest layer
        high_high_gradients: NDArray[np.float64] = addtn_gradients[0][:high_len, :]
        high_link_gradients: NDArray[np.float64] = addtn_gradients[0][high_len:, :]
    else:  # Continue recursion
        high_high_gradients, high_link_gradients = RecursiveGradientsResolve(
            low_layer - 1,
            global_indices,
            addtn_gradients,
            subtr_gradients,
            links_list,
            frag_list,
        )

    # Get the low layer gradients (this is the total system for this layer)
    gradients: NDArray[np.float64] = np.copy(addtn_gradients[low_layer])

    # Add the high_layer gradients (with no link atoms, these gradients should have been resolved from the previous recursions)
    gradients[high_in_low_indices, :] += high_high_gradients

    # Remove the current low LoT/high layer gradients
    gradients[high_in_low_indices, :] -= subtr_gradients[high_layer][:high_len, :]

    # Define the link atom gradients
    link_atom_gradients: NDArray[np.float64] = (
        high_link_gradients - subtr_gradients[high_layer][high_len:, :]
    )

    # Correct the link atoms and return
    gradients: NDArray[np.float64] = CorrectGradientsForLinkAtoms(
        frag_list[low_layer], links_list[high_layer], gradients, link_atom_gradients
    )

    return gradients[:low_len, :], gradients[low_len:, :]
