'''Interface to the SQC QM/MM Driver'''

#  Copyright (C) 2017 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 .Layer import Layer
from .NLayerSubtractive import NLayerSubtractive
import chemsh

def additive(qmmm, **kwargs):
    '''H(QM/MM) = H_QM(QM) + H_MM(MM) + H_QM/MM(inter)'''

    from numpy import arange, char
    from sys   import stdout

# YL 05/05/2020: not in use
    # part and counterpart
#    parts = { qmmm.qm:qmmm.mm,
#              qmmm.mm:qmmm.qm }

    # no need to recalculate
    qm = qmmm._qm
    mm = qmmm._mm

    qmmm._result.natoms = qmmm.frag.natoms
    qmmm._result.names  = qmmm.frag.names

    ### ENERGY ###
    # must be zeroed out first! because it is kept in memory and accumulated during iterations (for example in geometry optimisation)
    qmmm._result.energy = 0.0
    qmmm._result.energy += qm._result.energy
    qmmm._result.energy += mm._result.energy

    print("\n "+"-"*70)
    print(" QM energy          = %25.12f a.u."%qm._result.energy[0])
    print(" MM energy          = %25.12f a.u."%mm._result.energy[0])
    print(" QM/MM total energy = %25.12f a.u."%qmmm._result.energy[0])
    print(" "+"-"*70)
    print("")

    ### END OF ENERGY ###
    ### GRADIENTS ###

    # do gradients merging unless explicitly switched off
    if kwargs.get('_gradients', True):

        # MM gradients
        qmmm._result.gradients = mm._result.gradients
    
        # TODO: turn off when gradients are not calculated
        # YL 11/08/2019: it was wrong to print atoms as symbols as Result contains no symbols 
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n QM derivatives (a.u.)\n")
            print(qm._result.coords2str(coords=12,
                                        charges=None,
                                        atoms='names',
                                        inserts=[(0,char.mod('%6d',arange(qm._result.natoms)))]))
            print("\n MM derivatives (a.u.)")
            print(" NB: printed in the original order of QMMM fragment that can be different from MM output\n")
            print(mm._result.coords2str(coords=12,
                                        charges=None,
                                        atoms='names',
                                        inserts=[(0,char.mod('%6d',arange(mm._result.natoms)))]))
    
        # QM gradients
        # YL 29/01/2020: reckon with new derived classes FragmentBio, FragmentMat, and so on
        qmmm._result.gradients[qmmm.qm_region] += qm._result.gradients[qm.frag.stacks[qm.frag.__class__.__name__]]

        # resolve background point charges
        resolveBQGradients(qmmm)
    
        # resolve link atoms
        resolveLinkGradients(qmmm)


    if chemsh.CHEMSH_DEBUG > 4:
        print("\n >>> QM/MM total gradients (a.u.)\n")
        print(qmmm._result.coords2str(coords=12,
                                      charges=None,
                                      atoms='names',
                                      inserts=[(0,char.mod('%6d',arange(qmmm._result.natoms)))]))

    ### END OF GRADIENTS ###



def subtractive(qmmm, **kwargs):
    '''H(QM/MM) = H_QM(QM) - H_MM(QM) + H_MM(all)'''

    print("\n ChemShell QM/MM Driver ERROR: subtractive scheme not yet implemented\n")


def qmme(qmmm, **kwargs):
    '''QM/Me periodic embedding scheme by Meyer and Reuter'''

    from numpy  import arange, char, isin 
    from ..data import references
    from ..io   import format

    # no need to recalculate
    qm = qmmm._qm
    mm = qmmm._mm

    print("\n QM/Me periodic embedding scheme:\n")
    print(format.wrap(references.ref100))

    ### ENERGY ###
    # must be zeroed out first! because it is kept in memory and accumulated during iterations (for example in geometry optimisation)
    qmmm._result.energy = 0.0
    qmmm._result.energy += qm._result.energy
    qmmm._result.energy += mm._result.energy

    # do gradients merging unless explicitly switched off
    if kwargs.get('_gradients', True):

        # zero out all entries!
        qmmm._result.gradients[:] = 0.0

        # MM forces: adsorbate not included
        qmmm._result.gradients[~isin(arange(qmmm.natoms), qmmm.adsorbate)] = mm._result.gradients
        
        # debug print
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n QM derivatives (a.u.)\n")
            print(qm._result.coords2str(coords=12,
                                        charges=None,
                                        atoms='names',
                                        inserts=[(0,char.mod('%6d',arange(qm._result.natoms)))]))
            print("\n MM derivatives (a.u.)")
            print(" NB: printed in the original order of QMMM fragment that can be different from MM output\n")
            print(mm._result.coords2str(coords=12,
                                        charges=None,
                                        atoms='names',
                                        inserts=[(0,char.mod('%6d',arange(mm._result.natoms)))]))
        
        # QM gradients: adsorbate included
        qmmm._result.gradients[qmmm.qm_region] += qm._result.gradients

    print("\n "+"*"*70)
    print(" QM energy          = %25.12f a.u."%qm._result.energy[0])
    print(" MM energy          = %25.12f a.u."%mm._result.energy[0])
    print(" QM/MM total energy = %25.12f a.u."%qmmm._result.energy[0])
    print(" "+"*"*70)
    print("")


def merge(qmmm, **kwargs):
    '''Merge energies/gradients of QM and MM parts'''

    models = {
               'additive'   : additive,
               'subtractive': subtractive,
               'qm/me': qmme
             }

    models[qmmm.scheme](qmmm, **kwargs)


# bqs (see resolve_adq_gradients in src/hybrid/adq.c of Tcl-ChemShell)
#     Tcl-ChemShell loops over adqs in two directions (`shiftqs` in Py-ChemShell) and atoms_dst/src,
#     while here we produce the arrays of `shiftqs.grads_dst` and `shiftqs.grads_src` as a whole, respectively
def resolveBQGradients(qmmm):
    '''Resolve gradients due to BQs'''

    from sys   import stdout
    from numpy import argwhere

    if qmmm._qm.frag.nbqs:

        # MM part bqs gradients subtracting atoms_src (qmmm._shiftqs._mmatoms)
        # YL 17/09/2020: added support for FragmentBio type
        try:
            qmmm._result.gradients[qmmm._shiftqs._effective_bqs] += qmmm._qm.frag.bqs.gradients[qmmm._qm.frag.bqs.stacks['Fragment']]
        except:
            qmmm._result.gradients[qmmm._shiftqs._effective_bqs] += qmmm._qm.frag.bqs.gradients[qmmm._qm.frag.bqs.stacks['FragmentBio']]
    
        # save gradients on shiftqs multiplied by coefs_dst
        for i in range(3):
            qmmm._shiftqs.grads_dst[:,i] = qmmm._shiftqs.coefs_dst * qmmm._qm.frag.bqs.gradients[qmmm._qm.frag.bqs.stacks['ShiftQs']][:,i]
    
        # correct gradients on atoms_dst
        qmmm._result.gradients[qmmm._shiftqs.atoms_dst[::2]] += qmmm._shiftqs.grads_dst[0::2] + qmmm._shiftqs.grads_dst[1::2]
    
        # save gradients on shiftqs multiplied by coefs_src
        for i in range(3):
            qmmm._shiftqs.grads_src[:,i] = qmmm._shiftqs.coefs_src * qmmm._qm.frag.bqs.gradients[qmmm._qm.frag.bqs.stacks['ShiftQs']][:,i]
    
        # correct gradients on atoms_src
        for atom_src in qmmm._shiftqs.excluded:
            qmmm._result.gradients[atom_src] += sum(qmmm._shiftqs.grads_src[argwhere(qmmm._shiftqs.atoms_src == atom_src).flat])


# links (see link_force_analytical in src/hybrid/link_atom.c of Tcl-ChemShell)
#     Py-ChemShell: M-L-Q
#     Tcl-ChemShell: B-D-C
def resolveLinkGradients(qmmm):
    '''Resolve gradients due to link atoms'''

    from numpy import dot

    if qmmm.nlinks:

        qmmm._links.gradients = qmmm._qm._result.gradients[qmmm._qm.frag.stacks['Links']]

        for bpair in qmmm._links.bpairs:

            linkgradients = qmmm._links.gradients[qmmm._links.boundariesstr.index(str(bpair))]

            # resolve gradients at Q
            # note: in Tcl-ChemShell, "link atom force (gd)" means QM gradients at dummy L
            qmmm._result.gradients[bpair[0]] += dot(qmmm._links.getQMatrix(bpair),
                                                    linkgradients.T).flatten()
            # resolve gradients at M
            qmmm._result.gradients[bpair[1]] += dot(qmmm._links.getMMatrix(bpair),
                                                    linkgradients.T).flatten()



def __del__():
    pass

