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

# TODO: check charge groups
# TODO: subtract_charge_charge_interactions for Dalton and Gaussian
import chemsh
from .               import _Interfaces
from .qm             import _QM
from .mm             import _MM
from ..data          import atomicdata
from ..objects       import links, result, shiftqs
from ..dl_py2f.utils import iterutils, miscutils, objutils

# REFACTOR: Is a separate wrapper around the hybrid module necessary? theory=hybrid instead?
# REFACTOR: If keeping qmmm.py it should be in its own folder as other interfaces don't inherit from it (unlike qm, mm)


class QMMM(_Interfaces):
    '''QMMM object'''

    # REFACTOR Rename mm -> mm_theory, qm -> qm_theory for consistency with Tcl-ChemShell

    _kwargs = {
                'adaptive'     : False,
                'bond_modifiers':{},
                'coupling'     :'covalent',
                'dipole_adjust': False,
                'embedding'    :'electrostatic',
                'adsorbate'    :[],
                'links'        : links.Links(),
                'shiftqs'      : shiftqs.ShiftQs(),
                'mm'           : _MM(),
                'scheme'       :'additive',
                'surface'      :'',
                'qm'           : _QM(),
                'qm_region'    :[],
                'restart'      : False,
                'shl_maxcycles': 5,
                'shl_tolerance': 0.0001,
              }

    _internals = {
                   '__adjust_dipole': True,
                 }

    _synons = { 'max_shl_cycles':'shl_maxcycle' }

    # this is a list of priorities since kwargs is not sorted
    # we have to first deal with <frag> on which <qmmm.qm> and <qmmm.mm> objects depend
    # 'embedding' determins 'mm'
    # SW: `dipole_adjust` switch should be handled before other things 
    _priorities = 'dipole_adjust', 'coupling', 'qm_region', 'bond_modifiers', 'frag', 'scheme', 'embedding'


    # PROPERTIES #


    @property
    def frag(self):
        '''QM/MM fragment'''

        # `coupling` will be processed beforehand (see self._priorities)
        if self._coupling in [ 'ionic' ]:
            self._frag.connect_mode = None
            # two necessary conditions to switch off recalculating fragIDs
            self._frag._recalc_fragIDs = False
            self._frag.fragIDs[:] = 0

        # if the QM/MM fragment comprises QM and MM fragments (rather than qm_region from a whole)
        # this will disallow a QM/MM fragment containing only QM or MM part
        if self._qm.frag.natoms and self._mm.frag.natoms and not self.qm_region:
            self._frag.merge(self._qm.frag,
                             self._mm.frag,
                             noself=True)
            # qm always comes prior to mm
            self.qm_region = range(self._qm.frag.natoms)

        return self._frag


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

        from numpy import any, arange
        from ..io  import format

        self._frag = val

        # YL 26/10/2021: issue the warning only when new indices are different from the saved old ones
        # TWK 2/11/2021: moved check that _old_order exists to prevent crash at first assignment
        if hasattr(self.frag, '_old_order'):
            old_indices = self.frag._old_order[self.qm_region]
            new_indices = arange(self.frag.natoms)[self.qm_region]
            if any((old_indices-new_indices)!=0):
                # self.qm_region is set before self.frag (see _priorities)
                print('\n '+'*'*70)
                format.printwrap('WARNING: fragment reordered and the QM indices may NOT reflect the original PDB atoms')
                print('\n '+'*'*70)
                print('\n new index  old index  name  res.  resID  x (a.u.)  y (a.u.)  z (a.u.)  charge\n')
                strbuff = self.frag.coords2str(inserts=[(0,''),
                                                        (1,new_indices,'%7d'),
                                                        (2,''),
                                                        (3,old_indices,'%7d'),
                                                        (5,''),
                                                        (6,self.frag.residues[self.qm_region]),
                                                        (7,self.frag.resIDs[self.qm_region],'%4d')],
                                               coords=3,
                                               charges=2,
                                               select=self.qm_region)
                print(strbuff)


    @property
    def qm_region(self):
        '''Indices of QM atoms'''

        return self._qm_region


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

        from ..io import format

        self._qm_region = sorted(iterutils.getFlattenedIter(val))

        # print the indices
        if len(self._qm_region):
            print('\n Indices of QM atoms:')
            format.printwrap(' '.join(list(map(str, self._qm_region))), prepend_blanc_line=False)
            print('\n NB: Py-ChemShell atom indices count from 0 (cf. 1 in Tcl-ChemShell)\n')


    # when qm_region is set up, qmmm.frag may not exist at all. however mm_region relies on full list of qmmm atoms
    @property
    def mm_region(self):
        '''Indices of QM atoms'''

        mmatoms = list(set(range(self._frag.natoms))-set(self._qm_region))

        return mmatoms


    # set QM object (self._qm instead of self.qm!)
    # return a copy everytime (efficient?)
    @property
#    @utils.miscutils.memusage
    def qm(self):
        '''QM object'''

        from numpy  import arange, array, copy, count_nonzero, delete, isin, searchsorted, unique
        from ..base import constants

        SMALL = constants.SMALL

        # This allows part of the QM/Me algorithm to be run at the QM level
        # REFACTOR ideally we should handle the whole scheme at the QM/MM level if possible to promote modularity
        self._qm._qmmm_scheme = self.scheme

        # ZMatrix should not be calculated on the QM frag as `conn` is differenct from the whole
        self._qm.frag._calcZMatrix = False

        # do not check with self._qm.natoms here, because self._qm.frag may not have been defined
        if not self.qm_region:
            return self._qm

        self._qm._gradients = self._gradients

        if self.scheme != "qm/me":
            # avoid multiple recalculation
            links = self.links

            # SW: apply the function to adjust the dipoles accross the boundary, after `links` is generated and before charge shift takes place
            # apply the switch to make sure this is done only ONCE otherwise the number accumulates!
            if self.dipole_adjust and getattr(self, '__adjust_dipole', True):
                self.boundary_charge_adjust(self.bond_modifiers)

        # Set up the charge shift method used in electrostatic and polarised embedding
        if self.scheme == "additive" and self.embedding not in [ 'mechanical' ]:
            shiftqs = self.shiftqs

#        # FIXME: self.shiftqs.charges is not accessed anywhere and so not initialised. Very unsafe!!!
#        self.shiftqs.charges

        # QM/Me scheme (Meyer and Reuter)
        # TWK 16/03/21: QM/Me is best thought of as a scheme rather than an embedding option
        if self.scheme == "qm/me":

            # define surface
            self._qm._surface = self.surface

            # keep a reference to non-adsorbate atoms
            self._qm._non_adsorbate = arange(len(self.qm_region))[~isin(self.qm_region, self.adsorbate)]

            # construct the fragment with cell: do not include link atoms
            self._qm.frag.merge(self.frag.getSelected(self.qm_region),
                                noself=True)

            # NB: the QM/Me scheme deals with a 2D QM region (slabs) for surface reactions
            # but codes such as CP2K still need parameters like alpha/beta
            # alpha: angle b & c, beta: angle a & c, gamma: angle a & b.
            # we put the original lattice vectors for later constructing the selected part's lattice constants (see QM.run())
            self._qm.frag.cell.vectors[:] = self.frag.cell.vectors[:]

            # reset k-points sampling, otherwise same as qmmm.kpoints
#            if self._recalc_kpoints:
#                self._qm.kpoints = tuple()


        # standard methods
        else:
            # include link atoms in the QM fragment
            # merge (with link atoms) to form the QM fragment
            self._qm.frag.merge(self.frag.getSelected(self.qm_region),
                                links,
                                noself=True)

        # cast to the same class as self.frag
        if self._qm.frag.__class__ != self.frag.__class__:
            self._qm.frag = self._qm.frag.castTo(self.frag)

        # Save link atom indices in case the QM interface needs them
        # (e.g. MNDO inputs require a list of link atoms)
        self._qm._linkatoms = list(range(len(self.qm_region),self._qm.frag.natoms))

        # qm fragment inherits the total charge 
        # TWK TODO check if this necessary (when QM calc charge is decoupled from system charge)
        self._qm.frag.totalcharge = self.frag.totalcharge

        # if QM/MM connectivity is not calculated, QM or MM is not either
        self._qm.frag.connect_mode = self.frag.connect_mode

        # construct bqs
        if self.scheme == 'additive' and self.embedding not in [ 'mechanical' ]:

            # YL 18/03/2021: we want the BQs to be generated on the fly, but if on entry self._qm.frag.bqs holds gradients (for example,
            #                during shells relaxtation) and there is deletion (see below) it loses the gradients data on exit; it's caused
            #                by the merge() method but not delete()
            bq_grads = copy(self._qm.frag.bqs.gradients)

            # YL 14/08/2019: added self.frag.bqs to allow pre-defined BQs
            self._qm.frag.bqs.merge(self.frag.getSelected(shiftqs._mmatoms),
                                    self.frag.shells,
                                    shiftqs,
                                    self.frag.bqs,
                                    noself=True)

            # apply charge shifts by indexing atoms_dst in _mmatoms
            if chemsh.CHEMSH_DEBUG > 2:
                print("\n Applying charge shift to (QMMM.frag indices):\n", shiftqs.atoms_dst[::2])
            # YL 05/05/2020: find indices of nonrepeated atoms_dst to apply charge shift to _mmatoms because there might be atoms_dst needing charge shift for more than once
            shift_mmatoms_dst = arange(shiftqs._mmatoms.size)[isin(shiftqs._mmatoms, shiftqs.atoms_dst)]
            u, shift_atoms_dst = unique(shiftqs.atoms_dst[::2], return_index=True)
            self._qm.frag.bqs.charges.put(shift_mmatoms_dst,
                                          (self.frag.charges.take(shiftqs.atoms_dst)+shiftqs.delqs)[::2][shift_atoms_dst])

            # YL 30/04/2020: exclude zero charges
            zero_bqs = arange(self._qm.frag.bqs._master.shape[0])[(-SMALL<self._qm.frag.bqs.charges)*(self._qm.frag.bqs.charges<SMALL)]
            indices = delete(arange(self._qm.frag.nbqs), zero_bqs)

            # YL 18/03/2021: get an indices array of undeleted shells
            self.frag.shells._effective_shells = arange(self.frag.shells.natoms)[isin(self._qm.frag.bqs.stacks['Shells'], indices)]

            # YL 30/04/2020: get a reference to bqs excluding deleted zeros for future merging of BQs gradients
            try:
                shiftqs._effective_bqs = shiftqs._mmatoms[isin(self._qm.frag.bqs.stacks['Fragment'], indices)]
            except KeyError:
                shiftqs._effective_bqs = shiftqs._mmatoms[isin(self._qm.frag.bqs.stacks['FragmentBio'], indices)]

            # YL 30/04/2020: delete zero-charged bqs
            if chemsh.CHEMSH_DEBUG > 2 and zero_bqs.size:
                print("\n Deleting zero BQs (QMMM.qm.frag.bqs indices):\n", zero_bqs)
            self._qm.frag.bqs.delete(zero_bqs)

            # YL 18/03/2021: the gradients must be inherited if the shape of BQs, after deletion of zeros, remains unchanged as in the beginning
            if bq_grads.shape[0] == self._qm.frag.bqs.natoms:
                self._qm.frag.bqs.gradients = bq_grads

        #self._qm._result.natoms = self._qm.frag.natoms
        self._qm._result.znums = self._qm.frag.znums
        self._qm._result.names = self._qm.frag.names

        # TWK 23/6/2020 Ensure separate boundary ECPs are NOT used except in the context of the ionic embedding model
        if self.coupling not in ['ionic']:
            self._qm.separate_ecps = False

        return self._qm


    @qm.setter
    def qm(self, val):
        '''Setter of QM object'''

        self._qm = objutils.checkType(val, _QM, info="qm")


    @property
    def qmmm_pairs_by_junctions(self):
        '''All non-bonded Q-M pairs connected to junction QM atoms, e.g.:
           [ Q1-M1, Q1-M2, Q2-M1, Q2-M2 ] in
                  Q1
                  |
               Q2-J-M1
                  |
                  M2
        '''

        from numpy import append, array, isin, meshgrid, int64, unique

        pairs = array([], dtype=int64).reshape((0,2))

        # connectivity on junction atoms
        # YL 06/04/2023: remove repeated QM atoms and throw out a warning
        qm_unique, doubly_counts = unique(self.boundaries[:,0], return_counts=True)
        connected = self.frag.conn[qm_unique]
        if sum(doubly_counts > 1) == 1:
            strbuff = ''.join([str(qm_unique[i]) for i, c in enumerate(doubly_counts) if c > 1])
            print(f'\n WARNING: QM/MM boundary formed by choosing of the QM atom {strbuff} may cause problems!')
        elif sum(doubly_counts > 1) > 1:
            strbuff = ', '.join([str(qm_unique[i]) for i, c in enumerate(doubly_counts) if c > 1])
            print(f'\n WARNING: QM/MM boundary formed by choosing of the QM atoms {strbuff} may cause problems!')

        for i in connected:
            pairs = append(pairs,
                           array(meshgrid(i[isin(i, self.qm_region)],
                                          i[isin(i, self.mm_region)])).T.reshape(-1,2),
                           axis=0)

        return pairs


    @property
#    @utils.miscutils.memusage
    def mm(self):
        '''MM object'''

        from numpy import arange, array, isin, unique

        # do not check with self._qm.natoms here, because self._qm.frag may not have been defined
        if not self.mm_region:
            return self._mm

        # task flag for gradient
        self._mm._gradients = self._gradients

        # QM/Me scheme (Meyer and Reuter)
        if self.scheme == "qm/me":

            # define surface
            self._mm._surface = self.surface

            # keep a reference to non-adsorbate atoms
            self._mm._non_adsorbate = arange(self.natoms)[~isin(arange(self.natoms), self.adsorbate)]

            # construct the fragment with cell
            self._mm.frag.natoms = self.frag.natoms - len(self.adsorbate)
            self._mm.frag = self.frag.getSelected(self._mm._non_adsorbate)
            self._mm.frag.cell = self.frag.cell
            # YL TODO: False before we can determine the point group
            self._mm.periodic = True
            # not in use, but needs to be initialised before passing to PY2F
#            self._mm.frag.fragIDs

        # other methods
        else:
            # natoms
            self._mm.frag.natoms = self.frag.natoms

            # MM fragment is copy of QM/MM fragment
            # YL 10/09/2020
            self._mm.frag = self.frag.getCopy()

        # connectivity
        self._mm.frag.connect = self.frag.connect
        # YL 03/03/2020: avoid reconnecting
        self._mm.frag._reconnect = self.frag._reconnect
        self._mm.frag._recalc_fragIDs = self.frag._recalc_fragIDs

        # shells (MM shells is a ref to QMMM shells)
        self._mm.frag.shells = self.frag.shells

        # This allows part of the QM/Me algorithm to be run at the MM level
        # REFACTOR ideally we should handle the whole scheme at the QM/MM level if possible to promote modularity
        self._mm._qmmm_scheme = self.scheme

        # embedding model classification
        self._mm._qmmm_embedding = self.embedding

        # delete superfluous charges on QM atoms if electrostatic or higher embedding is in use
        if self.scheme == "additive" and self.embedding not in [ 'mechanical' ]:
            self._mm.frag._master.charges.put(self.qm_region, 0.0)

        # partitions
        self._mm._qmatoms = self.qm_region
        self._mm._mmatoms = self.mm_region

        # extra terms to exclude
        self._mm.qmmm_pairs_by_junctions = self.qmmm_pairs_by_junctions

        # shells
        # gradients on shell atoms are an image of part of self._qm.frag.bqs.gradients
        try:
            # YL NB 18/03/2021: self.frag.nshells may not match self._qm.frag.nbqs because the latter doesn't include zero-charge ones
            self.frag.shells.gradients[self.frag.shells._effective_shells] = self._qm.frag.bqs.gradients[self._qm.frag.bqs.stacks['Shells']]
        except:
            pass

        self._mm._result.znums = self._mm.frag.znums
        self._mm._result.names = self._mm.frag.names

        return self._mm


    @mm.setter
    def mm(self, val):
        '''Setter of MM object'''

        self._mm = objutils.checkType(val, _MM, info="mm")


    # This is a reimplementation of part of function construct_atom_lists() in chemsh3/src/hybrid/interface.c
    # and function ziccd() in chemsh3/src/zmatrix/zmatrix.c
    # Py-ChemShell      Tcl-ChemShell   ziccd
    # Q-L-M             i1-?-i2         c-?-b
    @property
    def links(self):
        '''Link atoms'''

        # containing an implicit recalculation of connectivity
        boundaries = self.boundaries

        # check the size
        self._links.nlinks = boundaries.shape[0]
        

        # boundaries
        self._links.bpairs = [ bpair for bpair in boundaries ]
        #print(self._links.bpairs)

        # distances between Q and L
        self._links.qldists = [ atomicdata.linkdists[self.frag.znums[bpair[0]]] for bpair in boundaries ]

        # distances between Q and M
        self._links.qmdists = [ self.frag.dist(tuple(bpair)) for bpair in boundaries ]

        # differences between Q and M
        self._links.qmdiffs = [ self.frag.diff(tuple(bpair)) for bpair in boundaries ]

        # compute link atoms' coodinates (in a.u.)
        # 'minus' since self.frag.diff returns iatom-jatom (but jatom-iatom needed)
        self._links.coords = [ self.frag.coords[boundaries[i][0]] - self._links.qmdiffs[i]/self._links.qmdists[i]*self._links.qldists[i]
        #                                  |                                  |                        |                    |
        #                        current atom's coords           coords[iatom]-coords[jatom]           |          (scaled by) link_dist
        #                                                                     (normalised by) ||coords[iatom]-coords[jatom]||
                               for i in range(len(boundaries)) ]

        return self._links


    @links.setter
    def links(self, val):
        ''''''
        self._links = val


    @property
    def nlinks(self):
        '''Number of link atoms'''

        return self._links.nlinks


    # YL 03/03/2020: improved to be 10^2 times faster using NumPy to replace the built-in Python methods
    # not to confuse with the cell boundary conditions
    @property
    def boundaries(self):
        '''QM/MM boundaries'''

        from numpy     import copy    as np_copy
        from numpy     import all, append, array, asarray, flatnonzero, meshgrid, trunc, vstack, isin
        from itertools import product
        from ..base    import constants

        # YL NB 03/03/2020: do not recalculate for now until we implement adaptive QM/MM
        try:
            return self.__boundaries
        except:
            pass

        # algorithm:
        #        QM   |  MM
        #     1 2 3 4 | 5 6 7
        #   1 \
        # Q 2   \
        # M 3     \
        #   4       \
        # -------------
        # M 5         | \
        # M 6         |   \
        #   7         |     \
        # step 1: get QM/MM pairs (crosspairs)
        # step 2: get connected pairs (connpairs)
        # step 3: get intersections (crosspairs ∩ connpairs)

        # all QM/MM atom pairs (this step is not slow at all)
        
        crosspairs = array(list(product(self.qm_region, self.mm_region)))
        # print(crosspairs) # This prints every connection from atom 0.
        if not len(self.qm_region):
            print("\n >>> ChemShell ERROR: no QM atom has been selected\n")
            exit(999)

        # TWK: TODO efficiency? do we need a flag for changes of qm_region? 
        # make a copy to avoid conn calculation every loop 
        conn = np_copy(self.frag.conn)

        # no connectivity for ionic coupling
        if self.coupling in ['ionic']:
            conn[:] = -1

        # get all connected pairs of the QM atoms
        MAXCONN = constants.MAXCONN
        iatoms = trunc(flatnonzero(conn+1)/MAXCONN).astype(int, copy=False)
        jatoms = conn[conn>=0]
        ij = vstack((iatoms, jatoms)).T
        # NB: we can't directly use FragmentBase.bonds which excludes duplcates (viz I-J and J-I)
        connpairs = ij[isin(ij[:,0], self.qm_region)]
        # pairs in both crosspairs and connpairs are QM/MM boundaries
        i = isin(connpairs[:,0], crosspairs[:,0])
        j = isin(connpairs[:,1], crosspairs[:,1])
        self.__boundaries = connpairs[i*j]

        return self.__boundaries
        

    @property
    def nboundaries(self):
        '''Number of boundaries'''
        
        return self.boundaries.shape[0]
      


    # YL: note that this method depends on self.boundaries which has already recomputed connectivity,
    #     therefore use self.frag._master.conn rather than self.frag.conn
    # YL NB 05/05/2020: this is a bit slow and will be called multiple times but it is safe to do so if we implement adaptive embedding in the future
    @property
    def shiftqs(self):
        '''Shifted charges'''

        from numpy  import arange, argwhere, array, bincount, count_nonzero, isin, intersect1d, repeat, unique, count_nonzero
        from ..base import errors
        from time import time

        # The QM/Me scheme does not use charge shifting
        if self.scheme == 'qm/me':
            return None

        # Mechanical embedding should not involve shiftqs method at any point
        if self.embedding == 'mechanical':
            return None

        # Sanity check for presence of charges to shift
        if not count_nonzero(self.frag.charges):
            msg = '''Electrostatic embedding with charge shifting is requested but there are no non-zero charges in the system. Please provide charges for your input fragment or use mechanical embedding.'''
            # YL 14/06/2022: do not check at initialising or if the automated charge is used
            if not (self._qm.frag.natoms <= 1 and all(self._qm.frag.znums==0)) and self._qm.charge != 'automatic':
                raise errors.ChemShGeneralError(msg)
                exit(999)

        try:
            # here is an implicit recalculation of connectivity
            # use unique() to remove repeats
            atoms_src = unique(self.boundaries[:,1])
        except:
            # if boundaries are:
            #     a) not calculated (for example because connectivity is switched off) or
            #     b) unavailable because QM and MM region are not connected
            self._shiftqs._mmatoms  = array(self.mm_region)
            return self._shiftqs

        delqs              = []
        coords             = []
        factors            = []
        repeated_atoms_src = []
        repeated_atoms_dst = []
        all_atoms_dst      = []

        # resize shiftqsObj
        self._shiftqs.nshiftqs = 2*count_nonzero(isin(self.frag._master.conn[atoms_src], self.mm_region))

        # find MM atoms connected to QM and other MM atoms
        atoms_to_edit = []
        for atom_src in atoms_src:
            # all boundaries[0] are QM and all boundaries[1] are MM
            atoms_dst = intersect1d(self.frag._master.conn[atom_src], self.mm_region)
            natoms_dst = atoms_dst.size

            # it is possible (especially in the solid state) an MM atom is not connected to any other MM atom
            if not natoms_dst:
                continue

            repeated_atoms_src.extend(atom_src.repeat(2*natoms_dst))
            repeated_atoms_dst.extend(atoms_dst.repeat(2).tolist())
            factors.extend(array(1.0/float(atoms_dst.size)).repeat(2*natoms_dst))

            for atom_dst in atoms_dst:
                # normalisation in Tcl-ChemShell is not necessary
                coords.extend(self.frag.coords[atom_dst] + self.frag.diff(atom_dst, atom_src)*self._shiftqs.shift0/self._shiftqs.r0)
                coords.extend(self.frag.coords[atom_dst] - self.frag.diff(atom_dst, atom_src)*self._shiftqs.shift0/self._shiftqs.r0)
                # factor = 1.0/float(atoms_dst.size) 
                all_atoms_dst.append(atom_dst)
                delqs.append(self.frag.charges[atom_src]/float(atoms_dst.size))

                # repeat
                delqs.extend(delqs[-1:])

        # YL 05/05/2020: save the nonrepeated delqs beforing checking to repeats for calculating charges
        self._shiftqs.unique_delqs = delqs

        # YL 05/05/2020: to initialise (moved here from self.qm)
        self._shiftqs.charges

        # YL 05/05/2020: charge shift should be performed repeatedly if an atom_dst appears more than once
        delqs = array(delqs[::2])
        u, indices = unique(all_atoms_dst, return_inverse=True)
        for t in unique(indices):
            w = argwhere(indices==t)
            delqs[w] = delqs[w].sum()
            
        # repeat
        delqs = repeat(delqs, 2)
 
        self._shiftqs.atoms_src = array(repeated_atoms_src)
        self._shiftqs.atoms_dst = array(repeated_atoms_dst)
        self._shiftqs.factors   = array(factors)
        self._shiftqs.delqs     = delqs
        self._shiftqs.coords    = array(coords)

        # exclude atoms from which charges are shifted but include those charges are shifted to
        intersec_src_dst = set(self._shiftqs.atoms_src) & set(self._shiftqs.atoms_dst)
        if len(intersec_src_dst):
            raise errors.ChemShGeneralError('Charge shift scheme fails due to inappropriate QM/MM division along covalent bond. Check the structure around atoms %s or coupling mode.\n'%(str(sorted(intersec_src_dst))))
        
        self._shiftqs._mmatoms = array(list(set(self.mm_region)-(set(self._shiftqs.atoms_src)-set(self._shiftqs.atoms_dst))))

        return self._shiftqs


    @shiftqs.setter
    def shiftqs(self, val):
        ''''''

        self._shiftqs = val


    @property
    def nshiftqs(self):
        '''Number of shifted charges'''

        return self.shiftqs.nshiftqs


    @property
    def coupling(self):
        '''Coupling model: covalent or ionic'''

        return self._coupling


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

        self._coupling = val.lower()



    @property
    def embedding(self):
        '''Embedding method: mechanical or electrostatic'''

        return self._embedding


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

        self._embedding = val.lower()


#    @property
#    def _embedding(self):
#        '''Emebedding model'''
#
#        selectcases = { 'shift'     : 'electrostatic',
#                        'l2'        : 'electrostatic',
#                        'rc'        : 'electrostatic',
#                        'polarised' : 'electrostatic',
#                        'polarized' : 'electrostatic',
#                        'ionic'     : 'electrostatic',
#                        'mechanical': 'mechanical',
#                      }
#
#        return selectcases[self.coupling]

    @property
    def scheme(self):
        '''QM/MM scheme: (default: additive)'''

        return self._scheme

    @scheme.setter
    def scheme(self, val):
        '''Setter of QM/MM scheme'''

        self._scheme = val.lower()


    @property
    def restart(self):
        '''Restart from a previous calculation'''

        return self._restart


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

        self._restart = val
        self._qm.restart = val
        self._mm.restart = val


    @property
    def adsorbate(self):
        '''Indices of adsorbate molecule(s)'''

        from numpy import array, count_nonzero

        if not len(self._adsorbate):
            # do nothing if there is no cell
            if not count_nonzero(self.frag.cell.vectors):
                return array([])

            adsorbate = self.frag.selectOutOfCell()
            if adsorbate.size:
                print("\n QMMM: atoms detected out of cell and taken as adsorbate\n ", adsorbate)
                self._adsorbate = adsorbate

        return self._adsorbate


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

        self._adsorbate = iterutils.getFlattenedIter(val)


    @property
    def surface(self):
        '''Surface of interest: ab|ac|bc'''

        # YL 04/08/2021: Fragment only (may be moved to FragmentMat in future)
        if not self._surface and hasattr(self.frag, 'getSurface'):
            # if has previously chosen by other methods such as Fragment.bind() and/or Fragment.selectPeriodic()
            try:
                self._surface = self.frag._surface
            except:
                # detect surface of interest
                self._surface = self.frag.getSurface(adsorbate=self.adsorbate)

        return self._surface


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

        from ..base import lattice

        self._surface = lattice.surfaces.get(val.strip().lower(), 'ab')


# END OF PROPERTIES #
# CLASSMETHODS #


    def printBanner(self):
        '''Print some information'''

        from ..io import banner

        print('\n '+'='*70)
        print(" ChemShell QM/MM Driver")
        print(' '+'='*70)
        print(" QM/MM scheme      : {}".format(self.scheme.strip().lower()))
        # TWK These options are only relevant to additive QM/MM
        if self.scheme == 'additive':
            print(' Embedding method  : {}'.format(self.embedding.strip().lower()))
            print(' Coupling method   : {}'.format(self.coupling.strip().lower()))
        print(' QM theory         : {}'.format(self.qm.__class__.__name__))
        print(' MM theory         : {}'.format(self.mm.__class__.__name__))
        print(' Number of atoms   : {}'.format(self.frag.natoms))
        if self.coupling.strip().lower().startswith('cova'):
            if self.nlinks <= 1:
                print(' Number of QM atoms: {} (incl. {} link atom)'.format(self._qm.frag.natoms, self.nlinks))
            else:
                print(' Number of QM atoms: {} (incl. {} link atoms)'.format(self._qm.frag.natoms, self.nlinks))
        else:
            print(' Number of QM atoms: {}'.format(self._qm.frag.natoms))
        print(' Number of MM atoms: {}'.format(self._mm.frag.natoms))
        # YL 15/06/2022: use self._qm._charge instead of self._qm.charge to reflect the "automatic" mode
        #                given that self._qm.charge is never accessed before the current point
        print(' QM charge         : {}'.format(self._qm._charge))
        print(' '+'-'*70)


    @miscutils.memusage
    def run(self, restart=False, dryrun=False, **kwargs):
        ''''''

        from sys       import stdout
        from numpy     import copy
        from ..hybrid  import merge, shells
        from ..io      import format
        from ..objects import fragbio

        if not self.frag.natoms:
            print("\n\n >>> ChemShell ERROR: no valid fragment has been assigned to theory {}!\n\n".format(self.__class__.__name__))
            exit(999)

        self.printBanner()

        # throw out a warning if the total charge of QM region is not near an integer
        if self._qm.frag.natoms and len(self.qm_region):
            total_charge_qm = self._frag.charges[self._qm_region].sum()
            remainder = total_charge_qm - int(total_charge_qm)
            # this is an arbitrary criterion
            if 0.2 < remainder < 0.8 and isinstance(self._frag, fragbio.FragmentBio):
                format.printwrap('WARNING: Sum of the selected QM atoms\' partial charges {} has a significant remainder'.format(total_charge_qm))

        # MUST calculate MM first to relax shells!
        # and qm later, whose bqs will be updated by EXPLICIT accessing (so a pre-defined list [self.qm,self.mm] does not work)
        # YL 06/11/2020: accessing QM.qm/QM.mm also unnecessarily initialises them again
        for partname in [ '_mm', '_qm' ]:
            part = getattr(self, partname)
            if part.frag.natoms:
                # relax shells
                # TWK: added a check for QM if qm.frag.nshells > 0
                if part.frag.nshells:
                    shells.relaxShells(self, dryrun=dryrun)
                # respect part.restart if it is set during shells relaxation
                # YL 18/04/2018: during QM/MM geometry optimisation, theory.restart is set True in callback()
                #                of dl_find/callback.py because part.restart is probably False in the beginning
                if chemsh.CHEMSH_DEBUG > 2:
                    part.frag.save('_'+partname+'.xyz')

                # YL 06/01/2021: update qm object
                if partname == '_qm' and self.scheme == "additive" and self.embedding not in [ 'mechanical' ]:
                    # REFACTOR make the update mechanism explicit as implicit updating is too opaque
                    self.qm
        
                part.run(restart=part.restart, dryrun=dryrun, **kwargs) 

                # YL 14/06/2022: synchronise the charges with the MM fragment for the latter may be obtained from FF drivers such as DL_FIELD
                if partname == '_mm' and self.scheme != "qm/me":
                    # YL 02/03/2023: preserve the QM charges (because charges on QM atoms in the MM fragment are zeros)
                    charges_qm = copy(self.frag.charges[self._qm_region])
                    self.frag.charges[:] = self._mm.frag.charges[:]
                    self.frag.charges[self._qm_region] = charges_qm[:]

            else:
                print(" \n QM/MM WARNING: the {} region contains no atom\n".format(partname[1:].upper()))


        # merge results (used to be a @property and now moved back here, because of trouble in invoking the merger)
        merge(self, **kwargs)


    # SW: Work on dipole_adjust
    #Harry's version
    def boundary_charge_adjust(self, bond_modifiers):
        """Modifies the charge of atoms on the QM/MM boundary due to ghost H.
        bond_modifiers is a dictionary of bond-type keys and modification values.
        eg. bond_modifiers = {("si", "o"):0.3, ("s1", "s2"):val, ...}"""

        print("\n Applying bond modifiers on Si and O charges...\n")

        for bond in self.boundaries:
            """Work out which species are involved in the bond."""

            qm_atom, mm_atom = bond[0], bond[1]

            qm_atom_species = self._frag.names[qm_atom].decode('UTF-8')
            mm_atom_species = self._frag.names[mm_atom].decode('UTF-8')

            bond_type = (qm_atom_species, mm_atom_species)

            """Modify each atom in the bond as appropriate."""
            strbuff = ' QM atom {} ({}) and MM atom {} ({}):'.format(qm_atom, bond_type[0], mm_atom, bond_type[1])
            strbuff += ' from {:-.6f} and {:-.6f}'.format(self._frag.charges[qm_atom], self._frag.charges[mm_atom])

            if bond_type in bond_modifiers.keys():
                self._frag.charges[qm_atom] -= bond_modifiers[bond_type]
                self._frag.charges[mm_atom] += bond_modifiers[bond_type]

            elif bond_type[::-1] in bond_modifiers.keys():  #Bond in reverse.
                self._frag.charges[qm_atom] += bond_modifiers[bond_type[::-1]]
                self._frag.charges[mm_atom] -= bond_modifiers[bond_type[::-1]]

            else:
                print(" >>> ERROR: Bond type not assigned a modification value.")

            strbuff += ' to {:-.6f} and {:-.6f}\n'.format(self._frag.charges[qm_atom], self._frag.charges[mm_atom])
            print(strbuff)


        setattr(self, '__adjust_dipole', False)
    #call example:
    #silicate_modifiers = {("si","o"):0.3}
    #qmmmobj.boundary_charge_adjust(silicate_modifiers)




# END OF CLASSMETHODS #
