# distutils: extra_compile_args = -fopenmp
# distutils: extra_link_args = -fopenmp
from cython cimport wraparound, boundscheck, cdivision

from libc.math cimport pow, sqrt, fabs
from libc.stdio cimport printf

from cython.parallel cimport prange, threadid
from openmp cimport omp_get_num_threads

import numpy as np
cimport numpy as np


@wraparound(False)
@boundscheck(False)
@cdivision(True)
def cvar(int ia, int ja, int npt,
         double[:] array,
         double degree,
         int[:, :] indices,
         double[:] dist, double[:] varg):

    cdef long n, n1, n2
    cdef int ii, jj, i, j
    cdef double idist, jdist

    n = 0
    for n1 in range(npt - 1):
        for n2 in range(n1 + 1, npt):
            i = indices[0, n1]
            j = indices[1, n1]
            ii = indices[0, n2]
            jj = indices[1, n2]
            idist = ii - i
            jdist = jj - j
            dist[n] = sqrt(idist*idist + jdist*jdist)
            varg[n] = 0.5 * pow(array[n1] - array[n2], degree)
            n += 1


@wraparound(False)
@boundscheck(False)
@cdivision(True)
def bin_cvar(int ia, int ja, int npt,
             double[:] array,
             double degree,
             int[:, :] indices,
             double[:] bin_def,
             double[:, :] meand, double[:, :] varg, long[:, :] count,
             bint absolute, int ncpu):

    cdef long n1, n2

    for n1 in prange(npt - 1, nogil=True,
                     schedule="guided", chunksize=50,
                     num_threads=ncpu):
            for n2 in range(n1 + 1, npt):
                compute_sf(indices[0, n1],
                           indices[1, n1],
                           indices[0, n2],
                           indices[1, n2],
                           array[n1], array[n2], degree, absolute,
                           bin_def,
                           count, meand, varg)


@wraparound(False)
@boundscheck(False)
@cdivision(True)
cdef void compute_sf(int i, int j,
                     int ii, int jj,
                     double a1, double a2,
                     double degree,
                     bint absolute,
                     double[:] bin_def,
                     long[:, :] count,
                     double[:, :] meand,
                     double[:, :] varg) nogil:

    cdef int ni, nt
    cdef double idist, jdist, dist

    nt = threadid()

    idist = ii - i
    jdist = jj - j
    dist = sqrt(idist*idist + jdist*jdist)

    if (dist < bin_def[0]) or (dist >= bin_def[1]):
        pass
    else:
        ni = int((dist - bin_def[0]) / bin_def[2])
        count[ni, nt] += 1
        meand[ni, nt] += dist
        if absolute:
            varg[ni, nt] += 0.5 * pow(fabs(a1 - a2), degree)
        else:
            varg[ni, nt] += 0.5 * pow(a1 - a2, degree)