#  Copyright (C) 2019 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 chemsh
from .               import _Tasks
from ..dl_py2f.utils import miscutils, iterutils


class ChargeFitting(_Tasks):
    '''Single-point calculation is the basis of almost all tasks'''

    _kwargs = {
                'alpha'      : 1.0,
                'a'          : 0.0005,       # (Hartree/e)**2
                'b'          : 0.2,          # e
                'conserve'   : 10,
                'constraints': [],
                'frozen'     : [],
                'initial'    : [],
                'maxcycles'  : 10000,
                'maxinners'  : 10,
                'method'     :'resp',
                'nlayers'    : 4,            # number of concentric layers
                'npoints'    : 0,
                'tol_inner'  : 1.0E-5,
                'tolerance'  : 1.0E-12,
                'type'       :'shell',
                'vdw_scale'  : 1.2,
                'vdw_incr'   : 0.2,
              }


    @property
    def frozen(self):
        '''A list of frozen atoms'''

        return self._frozen


    @frozen.setter
    def frozen(self, val):
        '''Setter of frozen'''

        self._frozen = iterutils.getFlattenedIter(val)


    @property
    def active(self):
        '''A list of active atoms (inverse of frozen)'''

        from numpy import arange, isin

        all = arange(self.theory.frag.natoms)

        return all[~isin(all, self.frozen)]


    def __getCalculatedEsp(self, charges):
        '''Calculate the electrostatic potential due to point charges at the grid of points '''

        # TD 11/04/2024: This calculates the model electrostatic potential at the grid points
        # The aim of ESP is to find the value of charges that minimises the difference between
        # this and the electrostatic potential from DFT
        # More complex when there are point charges present in the DFT calc

        from numpy import dot

        # Distances from the fit charges to the esp points
        dists = self.theory.frag.grid.dists

        # Calculate contribution of charges being fitted
        # dot() is faster than inner(): rrms = inner(_charges, 1.0/dists)
        calc_esp = dot(charges, 1.0/dists.T)

        # Add the contribution of bqs charges, if present
        if self.theory.frag.nbqs > 0:
            calc_esp += dot(self.theory.frag.bqs.charges, 1.0/self.theory.frag.grid.bqs_dists.T)

        return calc_esp


    def __getRRMS(self, charges):
        '''Calculate and return the Relative Root Mean Square (RRMS) of the potential fitting error'''
    
        # TD 09/04/2024 This mimics TCL CHARGE_RRMS in esp_f.f90
        # (This is as measure of the agreement of the electrostatic potential with
        #  the point charge model. So this is minimised by ESP, but doesn't know
        #  about the penalty function included in RESP) 
       
        from numpy import sqrt, sum

        # Electrostatic grid points (esp) from QM
        qm_esp = self.theory.frag.grid.potentials
        # Electrostatic grid points (esp) from point charge model
        calc_esp = self.__getCalculatedEsp(charges)
        # Calculate potential fitting error
        delta = qm_esp - calc_esp
        # Calculate Relative Root Mean Square (RRMS) of potential fitting error
        rrms = sqrt(sum(delta**2)/sum(qm_esp**2))

        return rrms


    def __getGrads(self, charges):
        '''Gradients updater'''

        # Follows J. Phys. Chem. 1993, 97, 10269
        # ESP : sum((qm_esp-calc_esp)**2)
        # RESP: sum((qm_esp-calc_esp)**2)+A*sum(sqrt(q**2+B**2)-b)

        from numpy import average, dot, unique, sqrt
        
        # Distances from the fit charges to the esp points
        dists = self.theory.frag.grid.dists
        # Electrostatic grid points (esp) from QM
        qm_esp = self.theory.frag.grid.potentials
        # Electrostatic grid points (esp) from point charge model
        calc_esp = self.__getCalculatedEsp(charges)
        # Calculate potential fitting error
        delta =  qm_esp - calc_esp

        # Compute the gradients (Both methods)
        gradients = -2.0*dot(delta, 1.0/dists)

        # Add the extra restraint term for RESP
        if self.method == 'resp':
            gradients += self.a*charges/sqrt(charges**2+self.b**2)

        # see Tcl-ChemShell: src/chemsh/esp_f.f90
        natoms = self.theory.frag.natoms
        if self.conserve:
            gradients -= gradients.sum()/float(natoms)

        # frozen atoms
        gradients[self.frozen] = 0.0

        # use the averaged gradients if there are degenerate atoms (i.e., constrained)
        for constraint in self.constraints:
            arraybuff = unique(constraint)
#            self.theory.frag.grid.dists[:,arraybuff] = average(self.theory.frag.grid.dists[:,arraybuff], axis=1)[:,None]
            gradients[arraybuff] = average(gradients[arraybuff])

        return gradients


    def __resp(self):
        '''Restrained electrostatic potential method'''

        # TD 21/03/2024 Will also work for original ESP method
        # Note that in ESP, the gradients are calculated differently

        from time         import time
        from numpy.linalg import norm
        from numpy        import arange, average, diff, dot, einsum, int32, sqrt, zeros, hstack, unique, setdiff1d, sort

        if self.theory.frag.cell.dim != '0D':
            print("\n ChemShell ERROR: RESP charge fitting only supported for 0D systems (This is {})".format(self.theory.frag.cell.dim))
            exit(112)
            

        def __getConserved(_charges, _totalcharge):
            '''Compensate the deviations during fitting to conserve the original total charge'''

            _charges[self.active] -= (_charges.sum()-_totalcharge)/float(len(self.active))
            if(chemsh.CHEMSH_DEBUG >= 5):
                print("\n >>> The total charge has been reset to = %15.9f\n"%_charges.sum())

            return _charges

        # conservation is unstable if there are frozen atoms
        if len(self.frozen):
            print("\n >>> WARNING: Frozen atom(s) included. Total charge conservation has been switched off.\n")
            self.conserve = 0

        natoms      = self.theory.frag.natoms
        totalcharge = self.theory.frag.totalcharge

        # check for repeated indices
        if self.constraints:
            indices_all = sort(hstack(self.constraints))
            diff = indices_all[1:][diff(indices_all) == 0]
            if diff.size:
                print('\n >>> WARNING: There are repeated indices in the constrants list: ' + ', '.join(['%i'%i for i in diff]))
                print('              Please check your input or this may result in unexpected outcomes!\n')

        # get the indices of redundant variables from the constraints
        redundants = zeros((0,), dtype=int32)
        if self.constraints:
            strbuff = '\n Charge fitting constraints: '
            for constraint in self.constraints:
                arraybuff = unique(constraint)
                redundants = hstack((redundants, arraybuff[1:]))
                strbuff += '(' + ', '.join(['%i'%i for i in arraybuff]) + ') '
            print(strbuff)

        # single-point calculation
        self.theory.run(_gradients=False, _potential=False, _graphics=False, restart=False)

        # single-point calculation for grid potentials
        self.theory.run(_gradients=False, _potential=True,  _graphics=True , restart=True)

        # create an array holding scratch charges
        charges = zeros((natoms,))

        # initial values
        if self.initial == []:
            charges[:] = self.theory.frag.charges
        else:
            charges[:] = self.initial

        # apply constraints
        for constraint in self.constraints:
            arraybuff = unique(constraint)
            charges[arraybuff] = average(charges[arraybuff])

#        charges -= (charges.sum()-totalcharge)/float(natoms)
        charges = __getConserved(charges, totalcharge)

        gradients = self.__getGrads(charges)

        alpha = self.alpha

        # initial gradients and directions
        gradients  =  self.__getGrads(charges)
        directions = -gradients


        # outer iteration
        for outer in range(self.maxcycles):

            # conserve the total charge every self.conserve cycles
            if self.conserve and not outer%self.conserve:
                charges = __getConserved(charges, totalcharge)
           
            # Header for loop progress
            if chemsh.CHEMSH_DEBUG >= 5 and outer==0:
                print("   Outer   Inner        RRMS   |Gradient|")

            #linesearch.lineSearch(charges, grads, _updateGrads)
            totalnorm_direct  = norm(directions)
            dot_direct_grads = dot(directions, -gradients)/totalnorm_direct
        
       
            # inner iteration (line search) to update alpha
            for inner in range(self.maxinners):
                charges2   = charges + alpha*directions
                gradients2 = self.__getGrads(charges2)
                rrms2 = self.__getRRMS(charges2)
                dot_direct_grads2 = dot(directions, -gradients2)/totalnorm_direct
                if dot_direct_grads > dot_direct_grads2:
                    alpha *= dot_direct_grads/(dot_direct_grads-dot_direct_grads2)
                else:
                    alpha = 0.0
    
                # health check
                if alpha < 0.0:
                    print(" >>> ERROR: alpha < 0.0 at cycle %d"%outer)
                    break

                if chemsh.CHEMSH_DEBUG >= 5:
                    # Loop progress (Inner)
                    print(f"     {outer+1:>3d}     {inner+1:>3d}   {rrms2:.3e}    {norm(gradients2):.3e}")

    	        # convergence check
                if abs(dot_direct_grads2) < self.tol_inner:
                    break
                    if chemsh.CHEMSH_DEBUG >= 5:
                        print(" >>> Inner iteration has converged: dot(directs, grads) =", abs(dot_direct_grads2))

                else:
                    if inner == self.maxinners-1:
                        print(" >>> ERROR: Inner iteration did not converge within %d cycles! dot(directs, grads) = %f"%(self.maxinners, abs(dot_direct_grads2)))
                        quit()

            
            # health check
            if alpha < 0.0:
                break

            charges += alpha*directions
            totalnormOLD  = norm(gradients)
            gradients  = self.__getGrads(charges)
            rrms = self.__getRRMS(charges)
            totalnorm  = norm(gradients)
            directions = -gradients + directions*(totalnorm/totalnormOLD)**2

            if chemsh.CHEMSH_DEBUG >= 5:
                # Loop progress (Outer)
                print(f"     {outer+1:>3d}     {inner+1:>3d}   {rrms:.3e}    {totalnorm:.3e}")

            if totalnorm < self.tolerance:
                print("\n >>> Outer iteration has converged after %d cycles:"%(outer+1))
                print("\n     total norm = %15.9e"%totalnorm)
                print("\n     tolerance  = %15.9e\n"%self.tolerance)
                break

        # force conservation at the final step
        charges = __getConserved(charges, totalcharge)

        # finally we assign the fitted values to the fragment
        self.theory.frag.charges[:] = charges


    @miscutils.time
    def run(self, **kwargs):
        '''Run a single-point calculation'''

        from ..objects import grid

        if not self.theory:
            print(" ChemShell >>> Please select \"theory\" for charge-fitting calculation. Exiting...")
            quit(1)

        # default npoints
        if not self.npoints:
            self.npoints = self.theory.frag.natoms*200

        # set up the grid
        self.theory.frag.grid = grid.Grid(npoints=self.npoints,
                                          nlayers=self.nlayers,
                                          natoms=self.theory.frag.natoms,
                                          type=self.type,
                                          vdw_scale=self.vdw_scale,
                                          vdw_incr=self.vdw_incr,
                                          _frag=self.theory.frag)

        self.theory.frag.grid.buildGrid()

        # Calculate ESP/RESP charges
        self.__resp()

        return self


# aliases
QFitting = ChargeFitting
