import numpy as np

cimport cython
cimport numpy as np
cimport openmp
from cython.parallel cimport prange
from libc.math cimport atan, pow
from libc.stdlib cimport free, malloc

from .. import config


cdef extern from "geometry_selection.hpp" nogil:
    void wrapfn[T](T & x, T & y, T & z, const T & wrap, const T & wrap_by_two)
    void sphere_selection[T](T* position_array, char* output_array,
                              T x0, T y0, T z0, T max_radius,
                              T wrap,
                              Py_ssize_t num_particles, int num_threads)
    void cube_selection[T](T* position_array, char* output_array,
                            T x0, T y0, T z0, T x1, T y1, T z1,
                            T wrap,
                            Py_ssize_t num_particles, int num_threads)

ctypedef fused fused_float:
    np.float32_t
    np.float64_t

def selection(np.ndarray[fused_float, ndim=2] pos_ar, region, parameters, fused_float wrap):
    """OpenMP selection algorithm for spheres and cubes, optionally wrapped around the box.

    Parameters
    ----------
    pos_ar : array_like
        The positions of the particles, with shape (N,3). Must be C-contiguous.
        The datatype can either be float32 or float64, represented by fused_float.
    region : str
        The region to select particles from. Can be 'sphere' or 'cube'.
    parameters : array_like
        The parameters of the region. For a sphere, it is (x0, y0, z0, r_max), and for a cube,
        it is (x0, y0, z0, x1, y1, z1). Each parameter must be a float (will be cast to fused_float).
    wrap : fused_float
        The size of the box to wrap around. If <0, no wrapping is done.
    """

    cdef long i
    cdef long N=len(pos_ar)
    cdef fused_float cx,cy,cz,x,y,z,r2
    cdef fused_float r_max_2
    cdef np.ndarray[np.uint8_t, ndim=1] output = np.empty(len(pos_ar),dtype=np.uint8)

    cdef char* output_data = <char*> np.PyArray_DATA(output)

    if len(pos_ar) == 0:
        # The check for C-contiguous 3-col fails when len is zero, but there is nothing to do anyway, so just return
        return output

    if (np.PyArray_NDIM(pos_ar)!=2 or np.PyArray_DIMS(pos_ar)[1]!=3 or
            np.PyArray_STRIDES(pos_ar)[0] != sizeof(fused_float)*3 or
            np.PyArray_STRIDES(pos_ar)[1] != sizeof(fused_float)):
        raise ValueError("Input array must be C-contiguous and have 3 columns")

    cdef fused_float* pos_ar_data = <fused_float*> np.PyArray_DATA(pos_ar)

    cdef fused_float wrap_by_two = wrap/2
    cdef int num_threads = config['number_of_threads']

    if region == 'sphere':
        if len(parameters) != 4:
            raise ValueError("Sphere selection requires 4 parameters: (x0, y0, z0, r_max)")
        # (ugly syntax below is because starred parameters not allowed)
        sphere_selection(pos_ar_data, output_data,
                         <fused_float> parameters[0],
                         <fused_float> parameters[1],
                         <fused_float> parameters[2],
                         <fused_float> parameters[3], wrap,
                         len(pos_ar), num_threads)
    elif region == 'cube':
        if len(parameters) != 6:
            raise ValueError("Cube selection requires 6 parameters: (x0, y0, z0, x1, y1, z1)")
        cube_selection(pos_ar_data, output_data,
                       <fused_float> parameters[0],
                       <fused_float> parameters[1],
                       <fused_float> parameters[2],
                       <fused_float> parameters[3],
                       <fused_float> parameters[4],
                       <fused_float> parameters[5], wrap, len(pos_ar), num_threads)
    else:
        raise ValueError("Region must be either 'sphere' or 'cube'")

    return output
