"""
This is a C compiled module to compute atomic bonds.
"""                      
from libc.math cimport sqrt, abs
import cython
cimport cython
import numpy as np
cimport numpy as np
from numpy cimport ndarray
from cython.parallel import prange
from fullrmc.Core.pairs_distances import pairs_distances_to_point

# declare types
NUMPY_FLOAT32 = np.float32
NUMPY_INT32   = np.int32
ctypedef np.float32_t C_FLOAT32
ctypedef np.int32_t   C_INT32

# declare constants
cdef C_FLOAT32 FLOAT_ZERO      = 0.0
cdef C_FLOAT32 FLOAT_ONE       = 1.0
cdef C_INT32   INT32_ZERO      = 0
cdef C_INT32   INT32_ONE       = 1



@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
cdef C_FLOAT32 _single_atom_single_shell_dists( C_FLOAT32[:] distances,
                                                C_FLOAT32    lowerShell,
                                                C_FLOAT32    upperShell,
                                                C_INT32      ncores = 1) nogil:
    # declare variables
    cdef C_INT32 i
    cdef C_FLOAT32 coordNumber = FLOAT_ZERO
    cdef C_INT32 num_threads = ncores
    # loop
    for i in prange(INT32_ZERO, <C_INT32>distances.shape[0], INT32_ONE, nogil=True, schedule="static", num_threads=num_threads):
    #for i from 0 <= i < <C_INT32>distances.shape[0]:
        if lowerShell <= distances[i] <= upperShell:      
            coordNumber += FLOAT_ONE
            #with gil:
            #    print i, lowerShell, upperShell, distances[i]
    return coordNumber



@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def single_atom_single_shell_subdists( C_FLOAT32[:] distances,
                                       C_FLOAT32    lowerShell,
                                       C_FLOAT32    upperShell,
                                       C_INT32      ncores = 1):
    return _single_atom_single_shell_dists( distances  = distances,
                                            lowerShell = lowerShell,
                                            upperShell = upperShell,
                                            ncores     = ncores)


    
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def single_atom_single_shell_totdists( ndarray[C_FLOAT32, ndim=1] distances,
                                       C_INT32[:]                 shellIndexes,
                                       C_FLOAT32                  lowerShell,
                                       C_FLOAT32                  upperShell,
                                       C_INT32                    ncores = 1):
    # declare variables
    return _single_atom_single_shell_dists( distances  = distances[shellIndexes],
                                            lowerShell = lowerShell,
                                            upperShell = upperShell,
                                            ncores     =  ncores)
    
    
    
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def single_atom_single_shell_coords( C_INT32                    coreIndex,
                                     C_INT32[:]                 shellIndexes,
                                     ndarray[C_FLOAT32, ndim=2] boxCoords not None,
                                     C_FLOAT32[:,:]             basis not None,
                                     bint                       isPBC,
                                     C_FLOAT32                  lowerShell,
                                     C_FLOAT32                  upperShell,
                                     C_INT32                    ncores = 1):
    # declare variables
    distances = pairs_distances_to_point( point  = boxCoords[ coreIndex ], 
                                          coords = boxCoords[ shellIndexes ],
                                          basis  = basis,
                                          isPBC  = isPBC,
                                          ncores = ncores)  
    # compute and return
    return single_atom_single_shell_subdists(distances  = distances,
                                             lowerShell = lowerShell,
                                             upperShell = upperShell,
                                             ncores     = ncores)    
    
    
    
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def single_atom_multi_shells_totdists( ndarray[C_FLOAT32, ndim=1] distances,
                                       list                       shellsIndexes,
                                       C_FLOAT32[:]               lowerShells,
                                       C_FLOAT32[:]               upperShells,
                                       C_INT32                    ncores = 1):
    # declare variables
    cdef C_INT32 i
    cdef ndarray[C_FLOAT32,  mode="c", ndim=1] coordNumbers = np.zeros(len(shellsIndexes), dtype=NUMPY_FLOAT32)
    # compute coordination numbers
    for i from 0 <= i < <C_INT32>len(shellsIndexes):
        coordNumbers[i] = _single_atom_single_shell_dists( distances   = distances[shellsIndexes[i]],
                                                           lowerShell  = lowerShells[i],
                                                           upperShell  = upperShells[i],
                                                           ncores      = ncores)
    return coordNumbers
    
    
    
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def single_atom_multi_shells_coords( C_INT32                    coreIndex,
                                     list                       shellsIndexes,
                                     ndarray[C_FLOAT32, ndim=2] boxCoords not None,
                                     C_FLOAT32[:,:]             basis not None,
                                     bint                       isPBC,
                                     C_FLOAT32[:]               lowerShells,
                                     C_FLOAT32[:]               upperShells,
                                     C_INT32                    ncores = 1):
    # declare variables
    cdef C_INT32 i
    cdef ndarray[C_FLOAT32,  mode="c", ndim=1] coordNumbers = np.zeros(len(shellsIndexes), dtype=NUMPY_FLOAT32)
    # compute coordination numbers
    for i from 0 <= i < <C_INT32>len(shellsIndexes):
        # declare variables
        distances = pairs_distances_to_point( point  = boxCoords[ coreIndex ], 
                                              coords = boxCoords[ shellsIndexes[i] ],
                                              basis  = basis,
                                              isPBC  = isPBC,
                                              ncores = ncores)      
        coordNumbers[i] = _single_atom_single_shell_dists( distances   = distances,
                                                           lowerShell  = lowerShells[i],
                                                           upperShell  = upperShells[i],
                                                           ncores      = ncores)
    return coordNumbers        
    
    
    
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def single_atom_coord_number_totdists( C_INT32                    atomIndex,
                                       ndarray[C_FLOAT32, ndim=1] distances,
                                       list                       coresIndexes,
                                       list                       shellsIndexes,
                                       list                       lowerShells,
                                       list                       upperShells,
                                       list                       asCoreDefIdxs,
                                       list                       inShellDefIdxs,
                                       C_FLOAT32[:]               coordNumData,
                                       C_INT32                    ncores = 1):
    # declare variables
    cdef C_INT32 defIdx
    # compute coordination numbers as core
    for defIdx in asCoreDefIdxs[atomIndex]:
        coordNumber = single_atom_single_shell_totdists(distances    = distances,
                                                        shellIndexes = shellsIndexes[defIdx],
                                                        lowerShell   = lowerShells[defIdx] ,
                                                        upperShell   = upperShells[defIdx],
                                                        ncores       = ncores)
        coordNumData[defIdx] += coordNumber
    # compute coordination numbers in shell
    for defIdx in asCoreDefIdxs[atomIndex]:
        coordNumber = single_atom_single_shell_totdists(distances    = distances,
                                                        shellIndexes = coresIndexes[defIdx],
                                                        lowerShell   = lowerShells[defIdx] ,
                                                        upperShell   = upperShells[defIdx],
                                                        ncores       = ncores)
        coordNumData[defIdx] += coordNumber
    
    
    
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def single_atom_coord_number_coords( C_INT32                    atomIndex,
                                     ndarray[C_FLOAT32, ndim=2] boxCoords not None,
                                     C_FLOAT32[:,:]             basis not None,
                                     bint                       isPBC,
                                     list                       coresIndexes,
                                     list                       shellsIndexes,
                                     list                       lowerShells,
                                     list                       upperShells,
                                     list                       asCoreDefIdxs,
                                     list                       inShellDefIdxs,
                                     C_FLOAT32[:]               coordNumData,
                                     C_INT32                    ncores = 1):
    # compute coordination numbers as core
    for defIdx in asCoreDefIdxs[atomIndex]:
        coordNumber = single_atom_single_shell_coords( coreIndex    = atomIndex,
                                                       shellIndexes = shellsIndexes[defIdx],
                                                       boxCoords    = boxCoords,
                                                       basis        = basis,
                                                       isPBC        = isPBC,
                                                       lowerShell   = lowerShells[defIdx],
                                                       upperShell   = upperShells[defIdx],
                                                       ncores       = ncores)
        coordNumData[defIdx] += coordNumber
    # compute coordination numbers in shell
    for defIdx in inShellDefIdxs[atomIndex]:
        coordNumber = single_atom_single_shell_coords( coreIndex    = atomIndex,
                                                       shellIndexes = coresIndexes[defIdx],
                                                       boxCoords    = boxCoords,
                                                       basis        = basis,
                                                       isPBC        = isPBC,
                                                       lowerShell   = lowerShells[defIdx],
                                                       upperShell   = upperShells[defIdx],
                                                       ncores       = ncores)
        coordNumData[defIdx] += coordNumber    
        
        
            
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def multi_atoms_coord_number_totdists( C_INT32[:]   indexes,
                                       list         distances,
                                       list         coresIndexes,
                                       list         shellsIndexes,
                                       list         lowerShells,
                                       list         upperShells,
                                       list         asCoreDefIdxs,
                                       list         inShellDefIdxs,
                                       C_FLOAT32[:] coordNumData,
                                       C_INT32      ncores = 1):
    # declare variables
    cdef C_INT32 i
    for i from 0 <= i < <C_INT32>len(indexes):
        single_atom_coord_number_totdists( atomIndex      = indexes[i],
                                           distances      = distances[i],
                                           coresIndexes   = coresIndexes,
                                           shellsIndexes  = shellsIndexes,
                                           lowerShells    = lowerShells,
                                           upperShells    = upperShells,
                                           asCoreDefIdxs  = asCoreDefIdxs,
                                           inShellDefIdxs = inShellDefIdxs,
                                           coordNumData   = coordNumData,
                                           ncores         = ncores)
    
    

@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def multi_atoms_coord_number_coords( C_INT32[:]                 indexes,
                                     ndarray[C_FLOAT32, ndim=2] boxCoords not None,
                                     C_FLOAT32[:,:]             basis not None,
                                     bint                       isPBC,
                                     list                       coresIndexes,
                                     list                       shellsIndexes,
                                     list                       lowerShells,
                                     list                       upperShells,
                                     list                       asCoreDefIdxs,
                                     list                       inShellDefIdxs,
                                     C_FLOAT32[:]               coordNumData,
                                     C_INT32                    ncores = 1):
    # declare variables
    cdef C_INT32 i
    for i from 0 <= i < <C_INT32>len(indexes):
        single_atom_coord_number_coords( atomIndex      = indexes[i],
                                         boxCoords      = boxCoords,
                                         basis          = basis,
                                         isPBC          = isPBC,
                                         coresIndexes   = coresIndexes,
                                         shellsIndexes  = shellsIndexes,
                                         lowerShells    = lowerShells,
                                         upperShells    = upperShells,
                                         asCoreDefIdxs  = asCoreDefIdxs,
                                         inShellDefIdxs = inShellDefIdxs,
                                         coordNumData   = coordNumData,
                                         ncores         = ncores)
    
    

    
    
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def all_atoms_coord_number_totdists( ndarray[C_FLOAT32, ndim=2] distances,
                                     list                       coresIndexes,
                                     list                       shellsIndexes,
                                     list                       lowerShells,
                                     list                       upperShells,
                                     list                       asCoreDefIdxs,
                                     list                       inShellDefIdxs,
                                     C_FLOAT32[:]               coordNumData,
                                     C_INT32                    ncores = 1):
    # declare variables
    cdef C_INT32 i
    # get indexes
    cdef ndarray[C_INT32,  mode="c", ndim=1] indexes = np.arange(<C_INT32>distances.shape[1], dtype=NUMPY_INT32)
    # run multiple atoms coordination number using distances
    multi_atoms_coord_number_totdists( indexes        = indexes,
                                       distances      = distances,
                                       coresIndexes   = coresIndexes,
                                       shellsIndexes  = shellsIndexes,
                                       lowerShells    = lowerShells,
                                       upperShells    = upperShells,
                                       asCoreDefIdxs  = asCoreDefIdxs,
                                       inShellDefIdxs = inShellDefIdxs,
                                       coordNumData   = coordNumData,
                                       ncores         = ncores)
    
    
    
@cython.nonecheck(False)
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
@cython.always_allow_keywords(False)
def all_atoms_coord_number_coords( ndarray[C_FLOAT32, ndim=2] boxCoords not None,
                                   C_FLOAT32[:,:]             basis not None,
                                   bint                       isPBC,
                                   list                       coresIndexes,
                                   list                       shellsIndexes,
                                   list                       lowerShells,
                                   list                       upperShells,
                                   list                       asCoreDefIdxs,
                                   list                       inShellDefIdxs,
                                   C_FLOAT32[:]               coordNumData,
                                   C_INT32                    ncores = 1):
    # declare variables
    cdef C_INT32 i
    # get indexes
    cdef ndarray[C_INT32,  mode="c", ndim=1] indexes = np.arange(<C_INT32>boxCoords.shape[0], dtype=NUMPY_INT32)
    # run multiple atoms coordination number using coordinates
    multi_atoms_coord_number_coords( indexes        = indexes,
                                     boxCoords      = boxCoords,
                                     basis          = basis,
                                     isPBC          = isPBC,
                                     coresIndexes   = coresIndexes,
                                     shellsIndexes  = shellsIndexes,
                                     lowerShells    = lowerShells,
                                     upperShells    = upperShells,
                                     asCoreDefIdxs  = asCoreDefIdxs,
                                     inShellDefIdxs = inShellDefIdxs,
                                     coordNumData   = coordNumData,
                                     ncores         = ncores)


            
    
    
    
    