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


import chemsh
from itertools       import combinations, product, zip_longest
from sys             import stdout
from time            import time
from .               import constants
from ..data          import atomicdata, periodicdata
from ..objects       import cell
from ..dl_py2f.utils import miscutils


# YL 03/02/2021: avoid using the term "cutoff" which will be confused with the cell lists/verlet list cutoff
# ij_pairs must a (npairs, 2) array
def getCovalentThresholds(fragObj, ij_pairs):
    '''Return array of thresholds for covalent connections according to the given pairs of I and J atoms'''

    znums  = fragObj.znums
    scale  = fragObj.connect_scale
    toler  = fragObj.connect_toler
    radii  = atomicdata.covalent_radii

    iatoms = ij_pairs[:,0]
    jatoms = ij_pairs[:,1]

    thresholds = (radii[znums[iatoms]]
                + radii[znums[jatoms]])*scale + toler

    return thresholds


# TWK 24/06/2020: same charge sign scaling value
#                 originally from the UCL cluster construct program
# YL 03/02/2021: not yet used in connectVerletList(); ij_pairs must a (npairs, 2) array
# YL 03/02/2021: avoid using the term "cutoff" which will be confused with the cell lists/verlet list cutoff
def getIonicThresholds(fragObj, ij_pairs, scale_covalent=0.5):
    '''Return array of thresholds for covalent connections according to the given pairs of I and J atoms'''

    from numpy import sign

    charges  = fragObj.effectcharges
    znums    = fragObj.znums
    scale    = fragObj.connect_scale
    toler    = fragObj.connect_toler
    radii    = atomicdata.ionic_radii

    iatoms   = ij_pairs[:,0]
    jatoms   = ij_pairs[:,1]
    icharges = charges[iatoms]
    jcharges = charges[jatoms]

    sum_radii = radii[znums[iatoms]] + radii[znums[jatoms]]
    thresholds = (sum_radii)*scale + toler

    # apply charges
    masks_same_sign = sign(icharges) == sign(jcharges)
    thresholds[masks_same_sign] -= sum_radii[masks_same_sign]*scale_covalent

    return thresholds


# this is Python rewrite of "subroutine images" in chemsh3/src/dl_poly/utility_pack.f
# YL 30/01/2021: PBC imcon magic number
#                0 no PBC apply
#                1 standard cubic (3D)
#                2 orthorhombic (3D)
#                3 parallelepiped (3D)
#                4 truncated octahedron (3D)
#                5 rhombic dodecahedron (3D)
#                6 x-y parallelogram (no periodicity in z) (2D)
#                7 hexagonal prism (3D)
# YL TODO 02/02/2021: complement the PBCs
def getMinImages(coords, cellObj):
    '''Calculate the minimum images of coords from the given cell'''

    from numpy import abs, around, array, diag_indices, dot, empty_like, linalg, sign, sqrt

    # imcon = 0
    def __calcNonPeriodic():
        '''No boundary condition (e.g. non-periodical system)'''

        # cannot images = numpy.copy(coords) or so
        images = empty_like(coords)
        images.setfield(coords, coords.dtype)

        return images

    # YL 31/01/2021: added cubic PBC
    # imcon = 1
    def __cubic():
        '''Cubic'''

        v0 = 1.0/(cellObj._vectors[diag_indices(3)].mean())
        _invvecs = array([[v0,0.0,0.0],[0.0,v0,0.0],[0.0,0.0,v0]])
        abuff  = dot(_invvecs, coords.T)
        images = dot(cellObj._vectors, abuff-around(abuff)).T

        return images

    # YL 31/01/2021: added orthorhombic PBC
    # imcon = 2
    def __orthorhombic():
        '''Orthorhombic'''

        v0 = 1.0/cellObj._vectors[0,0]
        v1 = 1.0/cellObj._vectors[1,1]
        v2 = 1.0/cellObj._vectors[2,2]
        _invvecs = array([[v0,0.0,0.0],[0.0,v1,0.0],[0.0,0.0,v2]])
        abuff  = dot(_invvecs, coords.T)
        images = dot(cellObj._vectors, abuff-around(abuff)).T

        return images

    # imcon = 3
    def __parallelepiped():
        '''Parallelepiped'''

        # YL 09/10/2017: to enhance the performance moved the cellObj.dim selection out of the function
        # YL: since Aug 2016 the performance here has been greatly improved
        # keep in mind: vec[x/y/z, a/b/c] (first dimension means axis)
        # YL 04/02/2021: matmul or @ is preferred to dot
        abuff  = coords @ cellObj._inverse_vecs.T
        images = (abuff-around(abuff)) @ cellObj._vectors.T
#        abuff  = dot(cellObj._inverse_vecs, coords.T)
#        images = dot(cellObj._vectors, abuff-around(abuff)).T

        return images

    # imcon = 4
    def __truncated_octahedron():
        ''''''

    # imcon = 5
    def __rhombic_dodecahedron():
        ''''''

    # imcon = 6
    def __xy_parallelogram():
        '''XYParallelogram'''

        abuff = dot(cellObj._inverse_vecs, coords.T)
        # cellObj._inverse_vecs contains dummy Z vector
        _vectors = linalg.inv(cellObj._inverse_vecs)
        images   = dot(_vectors, abuff-around(abuff)).T

        return images

    # imcon = 7
    def __hexagonal_prism():
        '''Hexagonal prism'''

        v = cellObj._vectors
        # mean a and b (a~=b)
        a0 = (v[0,0]+v[1,1])*0.5

        _invvecs = array([[1.0,0.0,0.0],[0.0,sqrt(3.0)/a0,0.0],[0.0,0.0,1.0/v[2,2]]])
        _vectors = array([[a0,0.0,0.0] ,[0.0,a0/sqrt(3.0),0.0],[0.0,0.0,v[2,2]]])

        abuff  = dot(_invvecs, coords.T)
        images = dot(_vectors, abuff-around(abuff)).T

        masks = abs(images[:,0]*sqrt(3.0)) + abs(images[:,1]) >= a0/sqrt(3.0)
        images[masks,0] -= a0/2.0*sign(images[masks,0])
        images[masks,1] -= (a0/sqrt(3.0)/2.0)*sign(images[masks,1])

        return images

    # imcon = ?
    def __calc1D():
        ''''''

        print("\nWARNING: Sorry we do not support PBC for 1D system\n")
        exit(999)

    selectcases = { 'no'                  :__calcNonPeriodic,
                    'cubic'               :__cubic,
                    'orthorhombic'        :__orthorhombic,
                    'parallelepiped'      :__parallelepiped,
                    'truncated_octahedron':__truncated_octahedron,
                    'rhombic_dodecahedron':__rhombic_dodecahedron,
                    'x-y_parallelogram'   :__xy_parallelogram,
                    'xy_parallelogram'    :__xy_parallelogram,
                    'hexagonal_prism'     :__hexagonal_prism,
                  }

    try:
        res = selectcases[cellObj.pbc]()
    except KeyError:
        print(" >>> ERROR: unsupported PBC")
        raise

    return res


# Python rewrite of parlink() in DL_POLY Classic
# YL 03/02/2021: rewritten using NumPy which is much faster
#                "cell lists" (or "cell linked-lists") is a correct jargon for this algorithm of O(N)
#                see: Quentrec, B. and C. Brot (1973). "New method for searching for neighbors in molecular dynamics computations." Journal of Computational Physics 13(3): 430-432.
#                also see: https://en.wikipedia.org/wiki/Cell_lists
# YL TODO 03/02/2021: not yet parallelised but the new implementation is highly vectorised already
#@miscutils.time
def connectPeriodicCellLists(fragObj, **kwargs):
    '''Compute the whole connectivity table of periodic systems using the cell lists (O(N)) algorithm'''

    from numpy import add, arange, argsort, argwhere, array, asarray, diff, dot, einsum, full, isin, linalg
    from numpy import meshgrid, minimum, multiply, repeat, searchsorted, sort, squeeze, subtract, swapaxes
    from numpy import unique, vstack, zeros

    # TODO honour fragment connect_scale and connect_toler values as we do in connectVerletList

    # YL 21/02/2020: make sure all conn entries are -1 before calculating connectivity
    fragObj._master.conn = -1
    connectIJ = fragObj.connectIJ

    # make a mock cell
    mockcell = fragObj.mockcell

    # YL FIXME 03/02/2021: in Tcl-ChemShell's connect() function there is another operation of PBC--is it redundant?
    mockcell.applyConnectPBC(version='DL_POLY')

    # minimum images
    images = getMinImages(fragObj.coords, mockcell)

    # in DL_POLY: uxx, uyy, uzz
    # YL 01/02/2021: rewritten in a much faster way
#    ucoords = asarray([ [ dot(mockcell._inverse_vecs[axis,:], images[iatom,:]) + 0.5 for axis in range(3) ]
#                                                                                     for iatom in range(fragObj.natoms) ])
# YL 04/02/2021: I forget why I was using einsum()....
#    ucoords = einsum('ijk,ij->ij', mockcell._inverse_vecs[None,:,:], images[:,:]) + 0.5
    ucoords = images @ mockcell._inverse_vecs.T + 0.5

    nlinkcells_xyz = mockcell._nlinkcells_xyz

    # YL 01/02/2021: a much faster implementation
    cell_positions = minimum(multiply(nlinkcells_xyz[None,:], ucoords[:,:]).astype(int), nlinkcells_xyz-1)
    all_cells = cell_positions[:,0] + nlinkcells_xyz[0]*(cell_positions[:,1] + cell_positions[:,2]*nlinkcells_xyz[1])

    # unique cells (repeats removed)
    unique_all_cells = unique(all_cells)

    # YL 02/02/2021: reimplemented with the much faster NumPy--DL_POLY order is honoured (for convenient comparison)
    jcell_indices = array(meshgrid(arange(nlinkcells_xyz[1]), arange(nlinkcells_xyz[0]), arange(nlinkcells_xyz[2]))).T.reshape(-1,3)[:,[1,0,2]][unique_all_cells]

    # hardcoded
    nsubcells = periodicdata.nsubcells[1]

    # construct jcells of shape (unique_all_cells.size, nsubcells)
    jcell_indices = add(jcell_indices[:,:,None], periodicdata.ni_xyz[:,:nsubcells])
    # YL 02/02/2021: in theory we don't need to reshape, but we do (NumPy doesn't create copies anyway) for purpose of comparing with Tcl-ChemShell
    jcell_indices = swapaxes(jcell_indices, 1, 2).reshape(nsubcells*unique_all_cells.size,3)

    # construct offsets of shape (unique_all_cells.size, nsubcells, 3): first offset and then cellposition since the latter will be altered
    # in DL_POLY: cx, cy, cz
    offsets = zeros(jcell_indices.shape, dtype=float)
    jcell_indices_subtracted = jcell_indices - nlinkcells_xyz
    masks_p = jcell_indices_subtracted > -1
    masks_q = jcell_indices            <  0
    offsets[masks_p] =  1.0
    offsets[masks_q] = -1.0
    offsets = offsets.reshape(unique_all_cells.size, nsubcells, 3)

    # in DL_POLY: jx, jy, jz
    jcell_indices[masks_p] =  jcell_indices_subtracted     [masks_p]
    jcell_indices[masks_q] = (jcell_indices+nlinkcells_xyz)[masks_q]
    # NB: jcells align in the order of unique_all_cells (so that is indexed using enumerate() below)
    # in DL_POLY: jc=jx+ilx*((jy-1)+ily*(jz-1))
    jcells = (jcell_indices[:,0] + nlinkcells_xyz[0]*(jcell_indices[:,1] + jcell_indices[:,2]*nlinkcells_xyz[1])).reshape(unique_all_cells.size,nsubcells)
    jcell_indices = jcell_indices.reshape(unique_all_cells.size, nsubcells, 3)

    # container arrays of final I-J pairs
    final_pairs   = array([], dtype=int).reshape(0,2)
    final_offsets = array([], dtype=float).reshape(0,3)

    # loop over cells
    for i, icell in enumerate(unique_all_cells):

        # atoms from icell
        iatoms = argwhere(all_cells==icell).ravel()

        # indexing with [1:] to exclude icell itself
        jcells_i = jcells[i][1:]
        # atoms from jcell (subcell)
        jatoms = argwhere(isin(all_cells, jcells_i)).ravel()

        # indices of the current subcells (remember we exclude icell by [1:]) used for constructing jatoms,
        # of the same size as jatoms and in the order of jatoms' appearance
        sorter = argsort(jcells_i)
        jcells_appearing = searchsorted(jcells_i, all_cells[jatoms], sorter=sorter)
        indices_subcell = arange(1,nsubcells)[sorter][jcells_appearing]

        # atom pairs from icells and jcells
        ij_atoms = array(meshgrid(iatoms, jatoms)).T.reshape(-1,2)
        # atom pairs from just icells (excluding self-self pairs)
        # iatoms and jatoms have no overlap but we still must calculate iatoms-iatoms pairs
        ii_atoms = array(list(combinations(iatoms, 2)), dtype=int)

        # pairs of I-J cells (cells and subcells)
        ij_cells = array(meshgrid(full(iatoms.shape, i), indices_subcell)).T.reshape(-1,2)
        # iatoms are obtained from icell and isubcell==0
        ii_cells = repeat(array([i,0])[None,:], int(iatoms.size*(iatoms.size-1)/2), axis=0)

        # actual offsets to use
        ij_offsets = offsets[ij_cells[:,0],ij_cells[:,1]]
        ii_offsets = offsets[ii_cells[:,0],ii_cells[:,1]]
        final_pairs   = vstack((final_pairs  , ij_atoms  , ii_atoms.reshape(ii_atoms.shape[0],2)))
        final_offsets = vstack((final_offsets, ij_offsets, ii_offsets))

    coords_diff = ucoords[final_pairs[:,1]] - ucoords[final_pairs[:,0]] + final_offsets
    vectors_diff = dot(mockcell.vectors[:,:], coords_diff.T)
    dists = linalg.norm(vectors_diff, axis=0)

    # ionic
    if fragObj.connect_mode.strip().lower() in [ "ionic" ]:
        thresholds = getIonicThresholds(fragObj, final_pairs)
    # covalent and default
    else:
        thresholds = getCovalentThresholds(fragObj, final_pairs)

    # connect
    masks = dists <= thresholds
    if chemsh.CHEMSH_DEBUG > 4:
        for i in arange(masks.size):
            print(">>> atoms {:>6d}-{:<6d} {:12.6f} threshold = {:12.6f} connect = {}".format(final_pairs[i,0], final_pairs[i,1], dists[i], thresholds[i], masks[i]))
    for i, (iatom, jatom) in enumerate(final_pairs[masks]):
        connectIJ(iatom[()], jatom[()])

    # YL 26/08/2021: return the connectivity array (do NOT use fragObj.conn which causes an indefinite recursion!)
    return fragObj._master.conn


# YL 11/03/2021: this is a naive but very elegant (so I think!) implementation of the generic cell lists algorithm
#                see: https://en.wikipedia.org/wiki/Cell_lists
# YL 12/03/2021: we haven't parallelised the current implementation
#@miscutils.time
def connectNonperiodicCellLists(fragObj, unitary_dist=0.0, update_nconns_only=False):
    '''Compute the whole connectivity table of non-periodic systems using the cell lists (O(N)) algorithm'''

    from numpy import all, isin, arange, searchsorted, unique, zeros

    # YL 16/03/2021: divide each virtual cell into subcells to reduce the search space
    nsubcells = 5

    # YL 16/03/2021: 28.35/5 a.u. = 5.670 a.u. = 3.000 Angstrom which we take as the longest possible covalent distance
    #                28.35 a.u. is close to an optimal parameter for the current implementation
#    cutoff = fragObj.cell_lists_cutoff # 10.0 a.u. by default
    # YL 25/08/2021: cutoff should be at least 3.0 A or unitary_dist+1.0 whichever is greater
    cutoff = max(5.670, unitary_dist+1.0)*nsubcells

    # again make sure everything starts from -1
    fragObj._master.conn = -1

    # YL 24/08/2021: added a unitary distance for sepecial purpose
    if unitary_dist:
        if update_nconns_only:
            func_connectII = fragObj._updateNConnsUnitaryDistIIs
            func_connectIJ = fragObj._updateNConnsUnitaryDistIJs
        else:
            func_connectII = fragObj._connectUnitaryDistIIs
            func_connectIJ = fragObj._connectUnitaryDistIJs
    else:
        # get references to functions
        if fragObj.connect_mode == "covalent":
            func_connectII = fragObj._connectCovalentIIs
            func_connectIJ = fragObj._connectCovalentIJs
        elif fragObj.connect_mode == "ionic":
            func_connectII = fragObj._connectIonicIIs
            func_connectIJ = fragObj._connectIonicIJs

    # prepare empty arrays
    tags    = zeros(fragObj.coords.shape, dtype=int)
    subtags = zeros(fragObj.coords.shape, dtype=int)
    indices_all = arange(fragObj.natoms)

    # divide into cells and subcells according to the X, Y, and Z spaces
    for i in range(3):
        nodes = arange(fragObj.coords[:,i].min()-cutoff/nsubcells,
                       fragObj.coords[:,i].max()+cutoff/nsubcells,
                       cutoff)
        # YL 15/03/2021: slice with finer tags to futher reduce the search space
        subnodes = arange(fragObj.coords[:,i].min()-cutoff/nsubcells,
                          fragObj.coords[:,i].max()+cutoff/nsubcells,
                          cutoff/nsubcells)
        # tag all atoms
        tags[:,i]    = searchsorted(nodes, fragObj.coords[:,i])
        subtags[:,i] = searchsorted(subnodes, fragObj.coords[:,i])

    # tag all cells along three dimensions
    tags_x, tags_y, tags_z = [ unique(tags[:,i]) for i in range(3) ]
    subtags_x, subtags_y, subtags_z = [ unique(subtags[:,i]) for i in range(3) ]
    subtags_0, subtags_1, subtags_2 = subtags[:,0], subtags[:,1], subtags[:,2]

    # first connect within cell A locally: this is the slowest step now because it's O(N**2)
    for x in tags_x:
        masks_x = tags[:,0]==x
        for y in tags_y:
            masks_xy = masks_x*(tags[:,1]==y)
            for z in tags_z:
                indices_A = indices_all[masks_xy*tags[:,2]==z]
                func_connectII(indices_A, unitary_dist=unitary_dist)

    # fix X- and Y-axis to scan along Z
    for x in tags_x:
        masks_x0 = tags[:,0]==x
        # must search along the diagonal direction, too, but we only include the closest subcells
        masks_x1 = masks_x0+(subtags_0==nsubcells*x+1)+(subtags_0==nsubcells*(x-1))
        for y in tags_y:
            masks_y0  = tags[:,1]==y
            masks_xy0 = masks_x0*masks_y0
            # must search along the diagonal direction, too, but we only include the closest subcells
            masks_xy1 = masks_x1*(masks_y0+(subtags_1==nsubcells*y+1)+(subtags_1==nsubcells*(y-1)))
            for z in tags_z[:-1]:
                # closest subcells in A and B only
                indices_A = indices_all[masks_xy0*(subtags_2==nsubcells*z)]
                indices_B = indices_all[masks_xy1*(subtags_2==nsubcells*z+1)]
                # connect between A and B
                func_connectIJ(indices_A, indices_B, unitary_dist=unitary_dist)

    # fix X- and Z-axis to scan along Y
    for x in tags_x:
        masks_x0 = tags[:,0]==x
        # must search along the diagonal direction, too, but we only include the closest subcells
        masks_x1 = masks_x0+(subtags_0==nsubcells*x+1)+(subtags_0==nsubcells*(x-1))
        for z in tags_z:
            masks_z0  = tags[:,2]==z
            masks_xz0 = masks_x0*masks_z0
            # must search along the diagonal direction, too, but we only include the closest subcells
            masks_xz1 = masks_x1*(masks_z0+(subtags_2==nsubcells*z+1)+(subtags_2==nsubcells*(z-1)))
            for y in tags_y[:-1]:
                # closest subcells in A and B only
                indices_A = indices_all[masks_xz0*(subtags_1==nsubcells*y)]
                indices_B = indices_all[masks_xz1*(subtags_1==nsubcells*y+1)]
                # connect between A and B
                func_connectIJ(indices_A, indices_B, unitary_dist=unitary_dist)

    # fix Y- and Z-axis to scan along X
    for y in tags_y:
        masks_y0 = tags[:,1]==y
        # must search along the diagonal direction, too, but we only include the closest subcells
        masks_y1 = masks_y0+(subtags[:,1]==nsubcells*y+1)+(subtags[:,1]==nsubcells*(y-1))
        for z in tags_z:
            masks_z0  = tags[:,2]==z
            masks_yz0 = masks_y0*masks_z0
            # must search along the diagonal direction, too, but we only include the closest subcells
            masks_yz1 = masks_y1*(masks_z0+(subtags_2==nsubcells*z+1)+(subtags_2==nsubcells*(z-1)))
            for x in tags_x[:-1]:
                # closest subcells in A and B only
                indices_A = indices_all[masks_yz0*(subtags_0==nsubcells*x)]
                indices_B = indices_all[masks_yz1*(subtags_0==nsubcells*x+1)]
                # connect between A and B
                func_connectIJ(indices_A, indices_B, unitary_dist=unitary_dist)

    # YL 26/08/2021: return the connectivity array (do NOT use fragObj.conn!)
    if unitary_dist:
        return fragObj._conn_aux
    else:
        return fragObj._master.conn


#@miscutils.time
# YL 03/02/2021: renamed as "Verlet list" which is the correct jargon for this real-space-distances based method of O(N**2)
def connectVerletList(fragObj, multiproc=False, **kwargs):
    '''Calculate connectivity using O(N**2) algorithm'''

#    from collections import deque
    from sys             import stdout
    from numpy           import add, array, array_split, arange, full, hstack, nditer, linalg, ravel, subtract, unravel_index, where
    from itertools       import combinations
    from ..parallel      import parallel
    from multiprocessing import cpu_count, Pipe, Process, set_start_method
    from functools       import partial

    from os import getpid

    # YL 06/03/2020: again make sure everything starts from -1
    fragObj._master.conn = -1

    mockcell = fragObj.mockcell

    # Get scale and tolerance for determining connectivity
    # according to the formula Rij < scale * (radius_i + radius_j) + toler
    scale = fragObj.connect_scale
    toler = fragObj.connect_toler

#    @miscutils.time
    def _setCovalentConnPairs(_childObj=None, _indices_of_wg=None, _multiproc=False):

        # YL 09/10/2017: method 3 is memory-efficient but three times slower
        # FIXME: comments needed
        fragObj._reconnect = False
        # avoid "dot" in loops to improve the efficiency
        natoms = fragObj.natoms
        coords = fragObj.coords
        znums  = fragObj.znums
        radii  = atomicdata.covalent_radii
        connectIJ = fragObj.connectIJ
        norm   = linalg.norm
        # YL TWK 09/09/2020: disable task-farming in connectiviy temporarily
        nworkgroups = parallel.nworkgroups
        workgroup = parallel.workgroup
        nworkgroups = 1
        workgroup   = 0

        iiter = nditer(arange(0,natoms-1), ['zerosize_ok'])

        if _indices_of_wg is not None:
            indices_of_wg = _indices_of_wg
        else:
            indices_of_wg = array_split(arange(0,natoms-1), nworkgroups)[workgroup]

        if mockcell.ndimensions == 3:
            for iatom in indices_of_wg:
                coords_i = coords[iatom]
                jatoms = arange(iatom+1,natoms)
                #c = subtract(coords[jatoms], coords[iatom])
                c = coords[jatoms].__rsub__(coords_i)
                _diffs = getMinImages(c, mockcell)
                _cutoffs = (radii[znums[iatom]]
                          + radii[znums[jatoms]])*scale + toler
                jatoms = jatoms[where(_cutoffs-norm(_diffs,axis=1)>=0.0)[0]]
                jiter = nditer(jatoms, ['zerosize_ok'])
                # TODO: vectorise this scalar function
                for jatom in jiter:
                    connectIJ(iatom[()], jatom[()])
        else:
            for iatom in indices_of_wg:
                jatoms = arange(iatom+1,natoms)
                _diffs = coords[jatoms] - coords[iatom]
                _cutoffs = (radii[znums[iatom]]
                          + radii[znums[jatoms]])*scale + toler
                jatoms = jatoms[where(norm(_diffs, axis=1) <= _cutoffs)[0]]
                jiter = nditer(jatoms, ['zerosize_ok'])
                for jatom in jiter:
                    connectIJ(iatom[()], jatom[()])

        fragObj._reconnect = False

        if _multiproc and _childObj:
            _childObj.send(fragObj.conn)
            _childObj.close()

        # YL TWK 08/09/2020: this line enables task-farmed conn calculations but currently
        #                    assumes that all workgroups are ready to receive conn data.
        #                    In other contexts, e.g. tasked-farmed QM/MM NEB, this is 
        #                    not the case and the MPI processes halt at this line
        #                    which calls synchronise_python_object() in parallel.f90
        #                    we now comment out the line temporarily and will fix it in future
#        parallel.syncPyObject(fragObj, 'conn', attach=True)


        
    def _setIonicConnPairs():

        from numpy import copy    as np_copy
        from numpy import nonzero, where

        natoms = fragObj.natoms

        coords     = fragObj.coords
        echarges   = fragObj.effectcharges
        znums      = fragObj.znums
        radii      = atomicdata.ionic_radii
        connectIJ     = fragObj.connectIJ
        norm       = linalg.norm
        np_copy    = np_copy

        # TWK 24/06/2020: same charge sign scaling value
        #                 originally from the UCL cluster construct program
        scale_covalent = 0.5

        iiter = nditer(arange(0,natoms-1),['zerosize_ok'])
        if mockcell.ndimensions == 3:
            for iatom in iiter:
                jatoms = arange(iatom+1,natoms)
                _diffs0 = coords[jatoms]-coords[iatom]
                _diffs = getMinImages(coords[jatoms]-coords[iatom], mockcell)
                _sumradii = radii[znums[iatom]] + radii[znums[jatoms]]
                # YL 03/02/2021: why was a np_copy here?
                _cutoffs = np_copy(_sumradii)*scale + toler
                icharges = echarges[iatom]
                jcharges = echarges[jatoms]
                # TWK 24/06/2020: For ions of the same sign the distance cut-off is scaled down
                #                 to exclude "unintentional" coordination. Only covalent pairing
                #                 is thus allowed as in O22- etc.
                abuff = (icharges < 0.0) == (jcharges < 0.0)
                _cutoffs[nonzero(abuff)] -= _sumradii[nonzero(abuff)]*scale_covalent
# YL 12/03/2021: commented out because I believe this is a redundant operation
#                _diffs2 = getMinImages(_diffs, mockcell)
                jatoms = jatoms[where(_cutoffs>=norm(_diffs,axis=1))[0]]
                jiter = nditer(jatoms, ['zerosize_ok'])
                for jatom in jiter:
                    connectIJ(iatom[()], jatom[()])
        else:
            for iatom in iiter:
                jatoms = arange(iatom+1,natoms)
                _diffs = getMinImages(coords[jatoms]-coords[iatom], mockcell)
                _sumradii = radii[znums[iatom]] + radii[znums[jatoms]]
                # YL 03/02/2021: why was a np_copy here?
                _cutoffs = np_copy(_sumradii)*scale + toler
                icharges = echarges[iatom]
                jcharges = echarges[jatoms]
                # TWK 24/06/2020: For ions of the same sign the distance cut-off is scaled down
                #                 to exclude "unintentional" coordination. Only covalent pairing
                #                 is thus allowed as in O22- etc.
                abuff = (icharges < 0.0) == (jcharges < 0.0)
                _cutoffs[nonzero(abuff)] -= _sumradii[nonzero(abuff)]*scale_covalent
# YL 12/03/2021: commented out because I believe this is a redundant operation
#                _diffs = getMinImages(_diffs, mockcell)
                jatoms = jatoms[where(_cutoffs>=norm(_diffs,axis=1))[0]]
                jiter = nditer(jatoms, ['zerosize_ok'])
                for jatom in jiter:
                    connectIJ(iatom[()], jatom[()])


    # number of available CPU cores
    nprocesses = min(cpu_count(), parallel.nprocs_per_workgroup)

    if fragObj.connect_mode == "covalent":

        if multiproc and nprocesses > 1:

            print("\n ChemShell connectivity: number of processes = {}".format(nprocesses))
            # "fork" is said to be faster than "spawn"
            try:
                set_start_method('fork', force=True)
            except:
                try:
                    set_start_method('spawn', force=True)
                except:
                    pass
            natoms       = fragObj.natoms
            # similar to task-farmed splitting
            split_ranges = array_split(arange(0,natoms-1), nprocesses)
            # processes list
            processes    = []
            # result buffer
            abuff = full((fragObj.natoms,0), -1, dtype='int8')
    
            for pid in range(nprocesses):
                # pipe communication
                parentObj, childObj = Pipe()
                indices_of_wg = split_ranges[pid]
                # create process instance
                proc = Process(target=_setCovalentConnPairs,
                               kwargs={'_multiproc'    :multiproc,
                                       '_childObj'     :childObj,
                                       '_indices_of_wg':indices_of_wg})
                processes.append(proc)
                proc.start()
                # collect results from child processes
                abuff = hstack((abuff, parentObj.recv()))
                abuff.sort(axis=1)
                abuff = abuff[:,-constants.MAXCONN:][:,::-1]

            # join
            for proc in processes:
                proc.join()
    
            # sort
            for conn in abuff:
                conn[:conn.argmin()].sort()
    
            # set values
            fragObj.conn = abuff

        else:
            _setCovalentConnPairs()


    elif fragObj.connect_mode == "ionic":
        _setIonicConnPairs()

    # YL 26/08/2021: return the connectivity array (do NOT use fragObj.conn!)
    return fragObj._master.conn


#@miscutils.time
# YL TODO 25/08/2021: unitary_dist only valid for connectNonperiodicCellLists()
def connect(fragObj, multiproc=False, unitary_dist=0.0):
    '''Calculate the connectivity table of fragObj'''

    from numpy import count_nonzero, full, int32

    # if disabled by user
    if not fragObj.connect_mode or not fragObj.natoms:
        return

    # make a mock cell
    # NOTE: while copying cell object, if cell.consts is empty and processed after cell.vectors, the latter (cell.vectors) will be wiped!!! (A bug of numpy?)
    mockcell = fragObj.mockcell

    # YL 26/08/2021: using a unitary distance may exceed MAXCONN (14) so that we must know
    #                the maximum number of connections in advance
    t0 = time()
    if unitary_dist:
        if mockcell._use_cell_lists and fragObj.cell.ndimensions != 0:
            # YL TODO 26/08/2021: connectPeriodicCellLists(..., update_nconns_only=True)
            pass
        else:
            # YL TODO 26/08/2021: connectVerletList(..., update_nconns_only=True)
            if fragObj.cell.ndimensions == 0:
                connectNonperiodicCellLists(fragObj, unitary_dist=unitary_dist, update_nconns_only=True)
            else:
                pass
        fragObj._conn_aux = full((fragObj.natoms, fragObj._nconns_aux), -1, dtype=int32)

    # only periodic fragObj
    if mockcell._use_cell_lists and fragObj.cell.ndimensions != 0:
        if mockcell._nlinkcells > mockcell._max_nlinkcells:
            print("\n >>> ERROR: connect() mockcell._nlinkcells should not be greater than mockcell._max_nlinkcells. Exiting..\n")
            exit(999)
        connectPeriodicCellLists(fragObj, unitary_dist=unitary_dist)
    # without cells: cell._max_nlinkcells = 1
    else:
        # YL 12/03/2021: use the cell lists algorithm which is way faster
        if fragObj.cell.ndimensions == 0:
            connectNonperiodicCellLists(fragObj, unitary_dist=unitary_dist)
# YL 12/03/2021: uncomment to test the old implementation
#            connectVerletList(fragObj, multiproc=multiproc)
        # YL 11/03/2021: we still use the old method for periodic systems
        else:
            connectVerletList(fragObj, multiproc=multiproc, unitary_dist=unitary_dist)

    # sort the connected atoms
    if unitary_dist:
        for conn in fragObj._conn_aux:
            conn[:conn.argmin()].sort()
    else:
        for conn in fragObj._master.conn:
            conn[:conn.argmin()].sort()


# TODO: three- and four-body forces not considered, see DL_POLY/utility/parset0
