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

__author__ = "Joseph Thacker <joseph.thacker@stfc.ac.uk>"

import pyscf
import pyscf.qmmm
from pyscf.qmmm import mm_charge, mm_charge_grad
import numpy as np


def BlockedForcesComp(blk_i, block_size, mol, qmmm_grad, dm):

    # Get the point charge coordinates and charges
    coords = qmmm_grad.base.mm_mol.atom_coords()
    charges = qmmm_grad.base.mm_mol.atom_charges()

    # Compute the forces on the point charges induced by the electrons

    # let H = 1/|r-R| and d = d/dR
    # Compute <di|H|j>
    AOInts = mol.intor("int1e_grids_ip", grids=coords[blk_i:block_size, :])

    # Get derivatives in (unoptimised) AO basis
    # d<i|H|j> = <di|H|j> + <dj|H|i>
    AOInts += AOInts.transpose((0, 1, 3, 2))

    # Transform from (unoptimised) AO basis (ij) to (optimised) MO basis (IJ)
    # d<I|H|J> = D . d<i|H|j>
    # Also scale by the charge and flip sign to account for wanting d/dr not d/dR
    point_charge_forces = np.einsum(
        "...ij,xkij,k->kx", dm, AOInts, -charges[blk_i:block_size], optimize=True
    )

    # Compute the forces on the point charges from the qm nuclear centres
    qm_coords = qmmm_grad.mol.atom_coords()
    qm_charges = qmmm_grad.mol.atom_charges()

    # Compute distance matrix (xyz components)
    dr = qm_coords[:, None, :] - coords[blk_i:block_size, :]

    # distance matrix
    r = np.linalg.norm(dr, axis=2)

    # Compute forces between QM (i) and MM (j) point charges using:
    # Fjx = Qi * Qj * dRijx / Rij^3
    point_charge_forces += np.einsum(
        "i,j,ijx,ij->jx", qm_charges, charges[blk_i:block_size], dr, np.power(r, -3)
    )

    return point_charge_forces


def ComputeElectrostaticEmbedding(
    mol: pyscf.gto.Mole, method, xyz, zcharge, compute_gradients=True
):
    #if method != pyscf.scf.RHF or method != pyscf.dft.RKS:
    #    raise NotImplementedError("Only RHF and RKS Implemented")

    # Add the point charges to the QMMM model
    method.mol = mol
    mf = mm_charge(method, xyz, zcharge)
    # Run the QMMM model
    mf.run()
    energy = mf.e_tot

    if compute_gradients:
        # Compute the gradients of the QM region
        qmmm_grad = mf.nuc_grad_method()
        qmmm_grad.run()
        qm_forces = qmmm_grad.de

        # Get the number of point charges
        n_charges = qmmm_grad.base.mm_mol.atom_coords().shape[0]

        # Compute the density matrix
        dm = mf.make_rdm1()

        # Block this bit of the computation as the matrices could get very large
        block_size = 512  # Max point charge forces to compute at a time
        remainder = n_charges % block_size
        blocked_length = n_charges - remainder
        point_charge_forces = np.zeros((n_charges, 3))
        for blk_i in range(0, blocked_length, block_size):
            point_charge_forces[blk_i:block_size, :] = BlockedForcesComp(
                blk_i, block_size, mol, qmmm_grad, dm
            )

        point_charge_forces[blocked_length:, :] = BlockedForcesComp(
            blocked_length, n_charges, mol, qmmm_grad, dm
        )

        return (energy, qm_forces, point_charge_forces)

    return energy
