'''Cluster partitioning routines'''

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

# TWK written
# YL: refactored 08/2019, this is a pure Python method
# TWK: 02/2020 revised new routine to preserve original TCS atom ordering

import chemsh
from numpy      import (any, append, array, array2string, arange, argsort, char,
                        einsum, isin, inf, ndenumerate, set_printoptions, subtract,
                        sort, sqrt, sum, unique, zeros, where)
from sys             import modules, stdout
from ..base          import constants, errors
from ..objects       import options, fragment, shells, bqs
from ..io            import keyword, file
from ..dl_py2f.utils import miscutils, objutils

# create an instance of <class Options>
optionsObj = options.Options()

# YL 15/08/2019: all items in this list are automatically made properties (with setters)
_properties = 'origin'

# partition options
_kwargs = {
            'bq_symbol'        : 'F',
            'cutoff_boundary'  :  10.0,
            'interface_exclude':  [],
            'interface_rigid'  :  [],
            'origin'           :  zeros(3),
#            'origin_point'     :  zeros(3), # YL will remove this in future
            'qmmm_interface'   : 'implicit',
            'qm_region'        :  [],
            'qm_termination'   : '',
            'qm_save_to'       : '',
            'radius_active'    :  20.0,
            'label_qm'         : '1',
            'label_interface'  : '2',
            'label_active'     : '3',
            'label_inactive'   : '4',
            'label_bq'         : '5',
            'save_to'          : ''
          }

_internals = {
             }

# YL 13/07/2018: done improvement of allowing returning result with the decorator @objutils.init()
# YL 14/08/2019: we need to think about future flexible multiple qm regions rather than the current hardlabelled 1-5 names
@miscutils.time
@objutils.init(optionsObj, inheritfrom=modules[__name__])
def partition(clusterObj, *args, **kwargs):
    '''Divide a cluster into 5 regions

    Region 1 = QM atoms
    Region 2 = QM/MM boundary atoms
    Region 3 = Active MM atoms
    Region 4 = Inactive MM atoms
    Region 5 = Point charges (approximating bulk electrostatics)
    '''

    from numpy  import any, float64, isin, ones
    from ..base import connectivity, forcefield

    print("\n >>> ChemShell CLUSTER: partitioning cluster...\n")

    # print option
    ITEMS_PER_LINE = 5

    # Read in user options
    keyword.parseKeys(optionsObj, **kwargs)    

    # Initialisation
    natoms = clusterObj.natoms
    charge_qm = 0.0

    # Initialise fragment which will hold the partitioned cluster
    # Build the output cluster object
    # - intended to resemble the output of construct,
    #   therefore no connectivity, fractional coords, etc.

    # YL 13/07/2018: renamed resultObj -> partitioned not to be confused with the <class Result> instance
    partitioned = clusterObj.getCopy()

    # Check options
    interface_methods = ['implicit', 'explicit']
    try:
        setattr(optionsObj, 'qmmm_interface', getattr(optionsObj, 'qmmm_interface').lower())
    except:
        pass
    if getattr(optionsObj, 'qmmm_interface') not in interface_methods:
        raise errors.ChemShValError('qmmm_interface', interface_methods)
    
    # TWK 02/2020
    # New order of atoms in output cluster (matching Tcl-ChemShell)
    # NB: In general, we group atoms of each region together, and retain ordering by distance
    # from the centre within each region. However, excluded atoms are grouped together
    # with the region they are excluded from.
    # The aim of all of this is to minimise changes in ordering when options are changed
    # (e.g. the overall size of the cluster, which atoms are excluded, etc. etc.)
    # for ease of comparison.    
    atom_order = zeros(natoms, dtype=int)
    nmapped = 0
    # Atom exclusions
    iflag = zeros(natoms, dtype=int)

    # List of species to be excluded or have shells removed from interface
    INCLUDED, RIGID, EXCLUDED = range(3)    
    
    # Rigid ions - remove shells
    # NB: regions are not defined yet, so all shells with a given symbol are flagged as RIGID
    # but shells must only be removed in the interface region
    for el in optionsObj.interface_rigid:
        iflag[char.startswith(char.lower(char.decode(clusterObj.names)),el.lower())] = RIGID

    # Exclusions - note this overrides the rigid option
    for el in optionsObj.interface_exclude:
        iflag[char.startswith(char.lower(char.decode(clusterObj.names)),el.lower())] = EXCLUDED

    if chemsh.CHEMSH_DEBUG:
        print("\n >>> iflags:")
        print(iflag)

    # get origin
    if array(optionsObj.origin).shape == (3,):
        r_centre = optionsObj.origin
    # YL 21/11/2019: from now io.keyword.parseKeys() takes [val] instead of val if default type is list/array
    elif array(optionsObj.origin).shape == (1,):
        if optionsObj.origin[0]+1:
            origin, r_centre = optionsObj.origin[0], zeros((3,),dtype=float64)
            r_centre[:] = clusterObj.coords[origin]
        else:
            r_centre = zeros((3,), dtype=float64)
    print("\n Origin (a.u.): %20.10f %20.10f %20.10f"%(r_centre[0], r_centre[1], r_centre[2]))

    # YL 14/08/2019: the meaning of Fragment.shellatoms was changed and caused point_shell to result wrong shells

    # YL 14/08/2019: our aim is to support unlimited number of qm and other regions with any label
    #                the current method is temporary

    # ==========================================================================
    # Inner QM region (region 1)
    # ==========================================================================
    # - user marks inner QM region with '1' in atom label
    # - point charges are marked with atom label = bq_symbol */
        # don't assign atom order yet, region 5 should appear last
                # Include charge contribution to QM atom from shell
    # YL 10/10/2019: check if there are existing QM region labels
    # qm_masks is a boolean array - True if QM label found, False otherwise
    qm_masks  = char.endswith(clusterObj.names, optionsObj.label_qm.encode())
    if qm_masks.sum():
        print("\n Pre-defined QM region label \"{}\" found in the original fragment:".format(optionsObj.label_qm), end='\n')
        prelabelled = ['{:>4} {:<8}'.format(i,clusterObj.names[i].decode()) for i in arange(clusterObj.natoms)[qm_masks]]
        print('\n'.join(' '.join(prelabelled[i:i+ITEMS_PER_LINE]) for i in range(0,len(prelabelled),ITEMS_PER_LINE)))
        # QM fragment object based on pre-labelled region
        qmFrag = clusterObj.getSelected(arange(clusterObj.natoms)[qm_masks])

    # YL 10/10/2019: check if there is a user defined `qm_region`
    qm_natoms = array(optionsObj.qm_region).size
    if qm_natoms:
        if qm_masks.sum():
            print("\n Keyword option `qm_region` has overridden the existing labels selecting {} centres:".format(qm_natoms), end='\n')
        else:
            print("\n Keyword option `qm_region` defines {} centres:".format(qm_natoms), end='\n')
        userselected = ['{:>4} {:<8}'.format(i,clusterObj.names[i].decode()) for i in array(optionsObj.qm_region)]
        print('\n'.join(' '.join(userselected[i:i+ITEMS_PER_LINE]) for i in range(0,len(userselected),ITEMS_PER_LINE)))
        # QM fragment object based on given `qm_region`
        qmFrag = clusterObj.getSelected(optionsObj.qm_region)
        # reset qm_masks (shape: (partitioned.natoms,))
        qm_masks = zeros(clusterObj.natoms, dtype=bool)
        qm_masks[optionsObj.qm_region] = True

    # array of QM indices
    qm_region = arange(clusterObj.natoms)[qm_masks]

    # MM indices
    mm_masks = ones(clusterObj.natoms, dtype=bool)
    mm_masks[optionsObj.qm_region] = False
    mm_indices = arange(clusterObj.natoms)[mm_masks]

    # if `termination` defined
    if optionsObj.qm_termination:
        # get conn
        connectivity.connect(clusterObj)
        qmFrag.conn[:] = clusterObj.conn[qm_masks]
        # masks of atoms to remove from the selected QM region: QM conn containing MM atoms and not defined as termination
        nont_masks = qmFrag.select(types=optionsObj.qm_termination, return_masks=True, invert=True)
        toremove = any(isin(qmFrag.conn,mm_indices).reshape(qmFrag.conn.shape), axis=1)*nont_masks
        # array of QM indices
        qm_region = qm_region[~toremove]
        # reset qm_masks (shape: (partitioned.natoms,))
        qm_masks = zeros(clusterObj.natoms, dtype=bool)
        qm_masks[qm_region] = True
        # QM fragment object based on updated `qm_region`
        qmFrag = clusterObj.getSelected(qm_region)

    if not qm_region.size:
        print(" >>> ChemShell ERROR: No QM atom defined. Please specify via keyword `qm_region`")
        print("{}or attach region label to atom names in structure file".format(" "*22))
        print("{}\"{}\".".format(" "*22,clusterObj._filename))
        exit(999)

    # relabel QM region by resetting names first
    partitioned.regions[qm_masks] = optionsObj.label_qm
    partitioned.names[qm_region] = char.add(forcefield.znums2names(forcefield.names2znums(clusterObj.names[qm_region])),
                                            partitioned.regions[qm_masks])

    if optionsObj.qm_termination:
        print("\n QM termination made \"{}\" only. Final QM region contains {} centres with new labels:".format(optionsObj.qm_termination, qm_region.size), end='\n')
        userselected = ['{:>4} {:<8}'.format(i,partitioned.names[i].decode()) for i in array(qm_region)]
        # print 10 items per line
        print('\n'.join(' '.join(userselected[i:i+ITEMS_PER_LINE]) for i in range(0,len(userselected),ITEMS_PER_LINE)))

    # Save to a new structure file if requested
    qmFrag.conn[:] = -1
    if optionsObj.qm_save_to:
        qmFrag.save(optionsObj.qm_save_to)
    else:
        if chemsh.CHEMSH_DEBUG:
            qmFrag.save('_qm_region_{}.pun'.format(qmFrag.formula))

    # New atom ordering for region 1
    nmapped = len(qm_region)
    atom_order[arange(nmapped)] = qm_region

    # =========================================================================
    # Fitted BQs (region 5)
    # =========================================================================
    bq_masks = char.find(clusterObj.names, optionsObj.bq_symbol.encode(), start=0, end=len(optionsObj.bq_symbol)) + 1 > 0
    partitioned.regions[bq_masks] = optionsObj.label_bq

    charge_qm += clusterObj.charges[qm_masks].sum()
    charge_qm += clusterObj.shells.charges[clusterObj.shellatoms[qm_masks*(clusterObj.shellatoms>-1)]].sum()

    if sum(partitioned.regions!=optionsObj.label_qm.encode()) == 0:
        print(" >>> ERROR cluster partition: no QM region (label: '{}') atoms specified!".format(optionsObj.label_qm))
        return False


    # =========================================================================
    # Region 2
    # =========================================================================
    # - cutoff distance is measured to region 1 atoms, not the origin
    # - character of region 2 depends on the QMMM interface options:
    #   - in IMPLICIT mode, region 2 is included in the QM region
    #     and shells are removed from all region 2 species
    #   - in EXPLICIT mode, region 2 is not included in the QM region
    #     and shells are not removed from region 2 species unless
    #     specified using the interface_rigid option.
    # - anions are no longer excluded from region 2 by default. Individual species
    #   can be excluded using the interface_exclude option.

    # YL 14/08/2019: our aim is to support unlimited number of qm and other regions with any label
    
    # this array of masks contains atoms that should be excluded
    interface_masks = (partitioned.regions!=optionsObj.label_qm.encode()) * (partitioned.regions!=optionsObj.label_bq.encode())
    excluded_masks = iflag == EXCLUDED

    vectors = subtract(clusterObj.coords[interface_masks], clusterObj.coords[partitioned.regions==optionsObj.label_qm.encode()][:,None])
    # YL 14/08/2019: faster than linalg.norm, see
    #                https://stackoverflow.com/questions/1401712/how-can-the-euclidean-distance-be-calculated-with-numpy
    #                i: region 1, j: region 2, k: xyz
    dists = sqrt(einsum("ijk,ijk->ij", vectors, vectors))

    # shape: (sum(interface_masks==True),)
    within_cutoff = any(dists<optionsObj.cutoff_boundary, axis=0)
    # YL 14/08/2019: `partitioned.regions[interface_masks][within_cutoff] = optionsObj.label_interface` strangely doesn't work
    #                so let us map within_cutoff back to interface_masks
    interface_masks[interface_masks==True] = within_cutoff

    # New atom ordering for interface atoms (whether excluded from region 2 or not)
    ninterface = sum(interface_masks)
    atom_order[arange(nmapped, nmapped+ninterface)] = where(interface_masks)
    nmapped += ninterface
    
    # label unexcluded atoms within cutoff boundary as "interface"
    partitioned.regions[interface_masks*~excluded_masks] = optionsObj.label_interface

    # label excluded atoms within cutoff boundary as "active"
    partitioned.regions[interface_masks*excluded_masks] = optionsObj.label_active

    # TWK 04/2020: In IMPLICIT mode, overall QM charge must include charge contributions from region 2 atoms and shells
    # The RIGID flag is not relevant here as shells are always deleted in IMPLICIT mode
    if optionsObj.qmmm_interface == 'implicit':
        condition_at = interface_masks*~excluded_masks
        condition_sh = condition_at*(clusterObj.shellatoms>-1)
        charge_qm += clusterObj.charges[condition_at].sum()
        charge_qm += clusterObj.shells.charges[clusterObj.shellatoms[condition_sh]].sum()

    # =========================================================================
    # Region 3
    # =========================================================================
    # - cutoff distance measured as radius around the origin
    # YL 14/08/2019:
    active_masks = clusterObj.selectByRadius(optionsObj.radius_active, centre=r_centre, return_masks=True)
    # unassigned regions
    active_masks *= partitioned.regions==b'0'
    partitioned.regions[active_masks] = optionsObj.label_active

    # New atom ordering for region 3
    nactive = sum(active_masks)
    atom_order[arange(nmapped, nmapped+nactive)] = where(active_masks)
    nmapped += nactive

    # =========================================================================
    # Region 4
    # =========================================================================
    # - all remaining atoms
    inactive_masks = (partitioned.regions!=optionsObj.label_qm.encode())        \
                    * (partitioned.regions!=optionsObj.label_interface.encode()) \
                    * (partitioned.regions!=optionsObj.label_active.encode())    \
                    * (partitioned.regions!=optionsObj.label_bq.encode())
    partitioned.regions[inactive_masks] = optionsObj.label_inactive

    # New atom ordering for region 4
    ninactive = sum(inactive_masks)
    atom_order[arange(nmapped, nmapped+ninactive)] = where(inactive_masks)
    nmapped += ninactive

    # =========================================================================
    # Region 5 again
    # =========================================================================
    # - just to assign the atom ordering
    n5 = sum(bq_masks)
    atom_order[arange(nmapped, nmapped+n5)] = where(bq_masks)
    nmapped += n5

    # All active regions have been assigned
    # Sanity checks
    if nmapped != clusterObj.natoms:
        print(" >>> ChemShell internal ERROR: nmapped != natoms")
        raise errors.ChemShError()

    if chemsh.CHEMSH_DEBUG > 4:
        print("\n Atom regions:")
        print(char.decode(partitioned.regions))

    # Atoms (reordered, charge set to zero for region 1 atoms
    # and also region 2 if IMPLICIT interface used)

    # YL 15/08/2019: Vacancy (ghost atoms) should be able to enter as bqs
    if clusterObj.nbqs:
        partitioned.bqs = clusterObj.bqs.getCopy()

    # YL 14/08/2019: effective shells
    # a: region 3 and beyond
    condition_a = (partitioned.regions==optionsObj.label_active.encode()) + (partitioned.regions==optionsObj.label_inactive.encode())
    # b: explicit interface and not excluded from region 2
    condition_b1 = optionsObj.qmmm_interface == 'explicit'
    condition_b2 = partitioned.regions == optionsObj.label_interface.encode()
    condition_b3 = iflag != RIGID
    condition_b = condition_b1*condition_b2*condition_b3
    condition_b4 = optionsObj.qmmm_interface == 'implicit'
    # c: having shells
    condition_c = clusterObj.shellatoms > -1

    # shells meeting the conditions (already excluding RIGID)
    condition_shells = (condition_a+condition_b)*condition_c
    nshells_effective = sum(condition_shells)
    
    # YL 22/11/2019: add back charge contributions from RIGID shells (what are already excluded from partitioned.charges)
    # TWK 02/2020: RIGID must only be applied in the interface region (condition_b2 must also be met)
    # TWK 04/2020: Also make sure RIGID only applies to explicit interface calculations (condition_b1)
    condition_rigid = (~condition_b3)*condition_c*condition_b2*condition_b1
    partitioned.charges[condition_rigid] += clusterObj.shells.charges[partitioned.shellatoms[condition_rigid]]
    
    # Sort atoms before dealing with shells
    partitioned._master = partitioned._master[atom_order]
    
    # YL 14/08/2019: initialise shells
    partitioned.shells.natoms = nshells_effective

    # valid shells of partitioned fragment mapped in clusterObj.shells
    # TWK 02/2020: as we know already the new atom order in partitioned, we can map the shells to new positions directly
    map_cluster_shells = []
    pcoreatoms = []
    for i in range(clusterObj.natoms):
        if condition_shells[atom_order[i]]:
            map_cluster_shells.append(clusterObj.shellatoms[atom_order[i]])
            pcoreatoms.append(i)

    if nshells_effective:
        partitioned._shells.coreatoms = pcoreatoms
        # correct partitioned.shellatoms because nshells changed
        pshellatoms = zeros(natoms, dtype=int)
        pshellatoms -= 1
        pshellatoms[partitioned._shells.coreatoms] = arange(nshells_effective)
        partitioned.shellatoms = pshellatoms
        # assign values
        partitioned.shells.names   = clusterObj.shells.names[map_cluster_shells]
        partitioned.shells.coords  = clusterObj.shells.coords[map_cluster_shells]
        partitioned.shells.charges = clusterObj.shells.charges[map_cluster_shells]

    # remove qm charges (including interface region if implicit)
    # NB: need a new array of masks after reordering!
    condition_b2 = partitioned.regions == optionsObj.label_interface.encode()
    qm_region = (partitioned.regions==optionsObj.label_qm.encode()) + condition_b2*condition_b4
    partitioned.charges[qm_region] = 0.0

    # relabel regions
    to_relabel_masks = partitioned.regions!=optionsObj.label_qm.encode()
    partitioned.names[to_relabel_masks] = char.add(char.strip(partitioned.names[to_relabel_masks]),
                                                   char.strip(partitioned.regions[to_relabel_masks]))

    # charge of QM region
    partitioned.totalcharge = charge_qm

    # title
    partitioned.title = clusterObj.title

    # Summary information
    set_printoptions(threshold=inf)
    print("\n Assignment of regions:")
    for region in unique(partitioned.regions):
        print(" No. atoms in region '{}': {}".format(region.decode(), sum(partitioned.regions==region)))
    print("\n QM charge of a neutral cluster: %12.6f" % charge_qm)
    print("\n QM region ", end="")
    if optionsObj.qmmm_interface == 'implicit':
        print("(regions {} and {}):".format(optionsObj.label_qm, optionsObj.label_interface))
    else:
        print("(region {} only):".format(optionsObj.label_qm))
    print(array2string(arange(natoms)[qm_region],
                       separator=' '))
 
    print("\n Active atoms (QM and interface not printed):")
    print(array2string(arange(natoms)[partitioned.regions==optionsObj.label_active.encode()],
                       separator=' '), end='\n\n')

    print("\n Centres possessing shells:")
    if optionsObj.qmmm_interface == 'implicit':
        print(" (Regions {} and {} shells already removed):".format(optionsObj.label_qm, optionsObj.label_interface))
    else:
        print(" (Region {} shells already removed):".format(optionsObj.label_qm))
    print(sort(partitioned.shells.coreatoms))

    # TWK 04/02/2020 in line with usual PCS philosophy, don't automatically save to disk unless explicitly requested
    if optionsObj.save_to:
        save_to = optionsObj.save_to
        print("\n Saving partitioned cluster to \"{}\"...".format(save_to))
        partitioned.save(save_to)

    return partitioned


# YL 13/07/2018: done improvement of allowing returning result with the decorator @objutils.init()
# hack as applying this decorator directly to partition seems to prevent
# partition returning an output object
# IMPROVE!
#@objutils.init2(optionsObj, inheritfrom=modules[__name__])
#def initOptions():
#    pass
@objutils.init(optionsObj, inheritfrom=modules[__name__])
def partitionOLD(clusterObj, *args, **kwargs):
    '''Divide a cluster into 5 regions

    Region 1 = QM atoms
    Region 2 = QM/MM boundary atoms
    Region 3 = Active MM atoms
    Region 4 = Inactive MM atoms
    Region 5 = Point charges (approximating bulk electrostatics)
    '''

    # is setting _kwargs twice necessary? REFACTOR
    _kwargs = {
                'cutoff_boundary' : 10.0,
                'radius_active' : 20.0,
                'qmmm_interface' : 'implicit',
                'interface_exclude' : [],
                'interface_rigid' : [],
                'bq_symbol' : 'F',
                'origin_point' : numpy.zeros(3)
              }

    _internals = {
                 }

    optionsObj._kwargs    = _kwargs
    optionsObj._internals = _internals
    
#    initOptions()

    debug = False
    
    # Initialisation
    assigned = 0
    natoms = clusterObj.natoms
    nshells = clusterObj.shells.natoms
    charge_qm = 0.0
    # Region specified for each core
    region = numpy.zeros(natoms, dtype=numpy.int)
    # Count for each region. 0 = unassigned
    nregion = numpy.zeros(6, dtype=numpy.int)
    # New order of atoms in output cluster
    atom_order = numpy.zeros(natoms, dtype=numpy.int)
    # Atom exclusions
    iflag = numpy.zeros(natoms, dtype=numpy.int)
    # Mapping from shells to cores
    point_core = numpy.zeros(nshells, dtype=numpy.int)
    # Mapping from cores to shells
    point_shell = numpy.zeros(natoms, dtype=numpy.int)
    
    # Read in user options
    keyword.parseKeys(optionsObj, **kwargs)    

    # Check options
    if getattr(optionsObj, 'qmmm_interface') not in ['implicit', 'explicit']:
        raise errors.ChemShValError(key, val)

    # List of species to be excluded or have shells removed from interface
    INCLUDED, RIGID, EXCLUDED = range(3)
    # Rigid ions - remove shells
    for el in optionsObj.interface_rigid:
        for i in range(natoms):
            if el in clusterObj.names[i].decode():
                iflag[i] = RIGID
    # Exclusions - note this overrides the rigid option
    for el in optionsObj.interface_exclude:
        for i in range(natoms):
            if el in clusterObj.names[i].decode():
                iflag[i] = EXCLUDED
    if debug:
        print("iflag")
        for i in range(natoms):
            print(i, iflag[i])
                
    # Get origin
    r_centre = optionsObj.origin_point
    if debug:
        print("r_centre: %20.10f %20.10f %20.10f" % (r_centre[0], r_centre[1], r_centre[2]))

    # Get pointers to/from shells
    # REFACTOR after shell pointers are refactored elsewhere
    for i in range(natoms):
        point_shell[i] = -1
    for i in range(nshells):
        pnt = clusterObj.shellatoms[i]
        point_core[i] = pnt
        point_shell[pnt] = i

    # TODO check that charges have been set on atoms

    # Find existing regions (1 and 5)
    # - user marks inner QM region with '1' in atom label
    # - point charges are marked with atom label = bq_sym */
    for i in range(natoms):
        if clusterObj.names[i].decode() == optionsObj.bq_symbol:
            region[i] = 5
            nregion[5] += 1
            # don't assign atom order yet, region 5 should appear last
        elif "1" in clusterObj.names[i].decode():
            region[i] = 1
            nregion[1] += 1
            atom_order[assigned] = i
            assigned += 1
            charge_qm += clusterObj.charges[i]
            if point_shell[i] != -1:
                # Include charge contribution to QM atom from shell
                charge_qm += clusterObj.shells.charges[point_shell[i]]
                nshells -= 1

    if nregion[1] == 0:
        print("cluster partition: no region 1 atoms specified!")
        return False

    # Region 2
    # - cutoff distance is measured to region 1 atoms, not the origin
    # - character of region 2 depends on the QMMM interface options:
    #   - in IMPLICIT mode, region 2 is included in the QM region
    #     and shells are removed from all region 2 species
    #   - in EXPLICIT mode, region 2 is not included in the QM region
    #     and shells are not removed from region 2 species unless
    #     specified using the interface_rigid option.
    # - anions are no longer excluded from region 2. Individual species
    #   can be excluded using the interface_exclude option.
    for i in range(natoms):
        if region[i] == 0:
            for j in range(natoms):
                if region[j] == 1:
                    disp = clusterObj.dist(i, j)
                    if disp < optionsObj.cutoff_boundary:
                        if iflag[i] == EXCLUDED:
                            # User has excluded species from region 2
		            # - put straight into region 3 instead
		            # - do not alter shells
                            region[i] = 3
                            nregion[3] += 1
                            atom_order[assigned] = i
                            assigned += 1
                            break
                        else:
                            # Place in region 2
                            region[i] = 2
                            nregion[2] += 1
                            chg = clusterObj.charges[i]
                            if point_shell[i] != -1:
	                        # Polarisable ions 
                                chg += clusterObj.shells.charges[point_shell[i]]
                                if optionsObj.qmmm_interface == 'implicit' or iflag[i] == RIGID:
                                    # shells will be removed later 
                                    nshells -= 1
                            if optionsObj.qmmm_interface == 'implicit':
	                        # Region 2 included in QM region
	                        # Include charge contributions from core and any shell 
                                charge_qm += chg
                            atom_order[assigned] = i
                            assigned += 1	    
                            break

    # Region 3
    # - cutoff distance measured as radius around the origin
    for i in range(natoms):
        if region[i] == 0:
            disp = numpy.linalg.norm(clusterObj.coords[i] - r_centre)
            if disp < optionsObj.radius_active:
                region[i] = 3
                nregion[3] += 1
                atom_order[assigned] = i
                assigned += 1
                
    # All active regions have been assigned
    nactive = assigned

    # Region 4
    # - all remaining atoms
    for i in range(natoms):
        if region[i] == 0:
            region[i] = 4
            nregion[4] += 1
            atom_order[assigned] = i
            assigned += 1

    # Region 5 again
    # - just to assign the atom ordering
    for i in range(natoms):
        if region[i] == 5:
            atom_order[assigned] = i
            assigned += 1

    # Sanity checks
    count = 0
    for i in range(1,6):
        count += nregion[i]
    if count != natoms:
        printf("cluster partition: region allocation count failed.")
        return False
    if assigned != natoms:
        printf("cluster partition: assigned atoms check failed.")
        return False

    if debug:
        print("\nAtom        region")
        for i in range(natoms):
            print("%6d %6d" % (i, region[i]))
        print("\nReordered   region")
        for i in range(natoms):
            print("%6d %6d" % (atom_order[i], region[atom_order[i]]))

    # Build the output cluster object
    # - intended to resemble the output of construct,
    #   therefore no connectivity, fractional coords, etc.
    
    # Initialise fragment which will hold the partitioned cluster
    # YL 13/07/2018: renamed resultObj -> partitioned not to be confused with the <class Result> instance
    partitioned = fragment.Fragment(connect_mode=clusterObj.connect_mode)

    # Atoms (reordered, charge set to zero for region 1 atoms
    # and also region 2 if IMPLICIT interface used)
    partitioned.natoms = natoms
    for i in range(natoms):
        j = atom_order[i]
        if region[j] == 1:
            partitioned.names[i] = clusterObj.names[j]
        else:
            partitioned.names[i] = clusterObj.names[j] + str(region[j]).encode()
        partitioned.coords[i] = clusterObj.coords[j]
        if region[j] > 2:
            partitioned.charges[i] = clusterObj.charges[j]
        elif region[j] == 2 and optionsObj.qmmm_interface == 'explicit':
            if point_shell[j] != -1 and iflag[j] == RIGID:
                # We are going to remove the shell so add shell charge to the core
                chg = clusterObj.charges[j] + clusterObj.shells.charges[point_shell[j]]
                partitioned.charges[i] = chg
            else:
                # Leave charge alone
                partitioned.charges[i] = clusterObj.charges[j]
        else:
            # Region 1 and region 2 for IMPLICIT interface
            partitioned.charges[i] = 0.0

    # No BQs are expected
    if clusterObj.nbqs:
        print("cluster partition: input cluster BQs array not empty.")
        return False

    # Shells (reordered and region 1 shells removed,
    # and region 2 as well if IMPLICIT interface is used or any
    # atoms have been specified as RIGID) 
    count = 0    
    partitioned.shells.natoms = nshells
    shellatoms = []

    # YL TODO: improve the performance
    from time import time
    st = time()
    for i in range(natoms):
        j = atom_order[i]
        if ((region[j] > 2
            or (region[j] == 2 and optionsObj.qmmm_interface == 'explicit' and iflag[j] != RIGID))
            and point_shell[j] != -1):
            partitioned.shells.names[count] = partitioned.names[i]
            partitioned.shells.coords[count] = clusterObj.shells.coords[point_shell[j]]
            # have to set displace to zero otherwise coords get messed up! 
            # TODO improve handling of coords/displacements in the shell object
            partitioned.shells.displace[count] = 0.0
            partitioned.shells.charges[count] = clusterObj.shells.charges[point_shell[j]]
            shellatoms.append(i)
            partitioned.shells.coreatoms[count] = i
            count += 1

    # appending directly to shellatoms in the loop causes a resize of the shell object(?)
    # YL 13/07/2018: as per change in 04/2018 Fragment.shellatoms is the full list of shell indices
#    partitioned.shellatoms = shellatoms

    # Charge of QM region
    partitioned.totalcharge = charge_qm

    # Title
    partitioned.title = clusterObj.title

    # Summary information
    if debug:
        print("")
    print("Assignment of regions\n")
    for i in range(1,6):
        print("No. atoms in region %d: %6d" % (i, nregion[i]))
    print("QM charge of a neutral cluster: %12.6f" % charge_qm)
    print("QM region:")
    if optionsObj.qmmm_interface == 'implicit':
        print("(regions 1 and 2)", end="")
    else:
        print("(region 1 only)", end="")
    count = 0
    for i in range(natoms):
        j = atom_order[i]
        if region[j] == 1 or (region[j] == 2 and optionsObj.qmmm_interface == 'implicit'):
            if (count % 10 == 0):
                print("")
            print("%6d " % i, end="")
            count += 1
    print("\nActive atoms: 0 - %-6d" % (nactive-1))

    return partitioned
