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

def setNeighbours(fragObj, optionsObj):
    '''Get a periodic neighbour list'''

    from itertools import product
    from numpy     import dot, zeros, float64, int64, arange, tile, where, repeat, count_nonzero
    from ..        import CHEMSH_LIB_DIR
    from ..base    import constants
    from ..dl_py2f import py2f

    # constant
    MAXCONN = constants.MAXCONN

    # an empty list of neighbours
    optionsObj.nneighbours = zeros((fragObj.natoms), int64)
    optionsObj.neighbours = zeros((fragObj.natoms,MAXCONN), int64)

    selectcases = { 3 : 27,
                    2 : 9,
                    1 : 3,
                    0 : 1 }
    
    # super cell consists of unit cells in all 27 translations (9 for 2D etc.)
    supercell = fragObj.getSuperCell(redundant=True)

    # connectivity
    supercell.connect()

    # tiled list of unit cell's atom indices
    ntrans = selectcases[fragObj._cell.ndimensions]
    tiled = tile(arange(fragObj.natoms), ntrans)

#    # flattened connectivity list of the original unit cell (the central one in super cell)
#    centreconn1 = supercell.conn[13*fragObj.natoms:14*fragObj.natoms].ravel()
    # connectivity list of the original unit cell (the central one in super cell)
    endCell = int((ntrans+1)/2)
    startCell = endCell - 1
    centreconn = supercell.conn[startCell*fragObj.natoms:endCell*fragObj.natoms]
    # calculate the neighbours list
    # NOTE: "+ 1" to make it Fortran
#    neighbours1[where(centreconn1>0)] = tiled[centreconn1[where(centreconn1>0)]] + 1
    optionsObj.neighbours[where(centreconn>0)] = tiled[centreconn[where(centreconn>0)]] + 1
    #print(optionsObj.neighbours)
    
    # neighbour count for each atom
    for i in range(fragObj.natoms):
        optionsObj.nneighbours[i] = count_nonzero(optionsObj.neighbours[i])
    
    # coordinates of the original unit cell repeated MAXCONN times
    repeated = repeat(fragObj.coords, MAXCONN, axis=0)

    # flattened
    centreconn = centreconn.ravel()

    # neighbour vectors
    optionsObj.neighb_vectors = zeros((MAXCONN*fragObj.natoms,3), float64)
    optionsObj.neighb_vectors[where(centreconn>0)] = supercell.coords[centreconn[where(centreconn>0)]] - repeated[where(centreconn>0)]
#    optionsObj.neighb_vectors = optionsObj.neighb_vectors.ravel()

#    libcluster.set_neighbours(byref(py2f(optionsObj)))


