#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include "Python.h"
#include "math.h"
#include "numpy/ndarraytypes.h"
#include "numpy/ufuncobject.h"
#include "numpy/npy_3kcompat.h"

static PyMethodDef PopcountMethods[] = {
        {NULL, NULL, 0, NULL}
};

#define CONCAT2(A, B) A##B
#define CONCAT(A, B) CONCAT2(A, B)

#define MAKE_POPCOUNT(prefix, ty) \
    static void CONCAT(prefix,_popcount)(char** args, const npy_intp* dimensions, \
            const npy_intp* steps, void* data) { \
        npy_intp i; \
        npy_intp n = dimensions[0]; \
        char* in = args[0]; \
        char* out = args[1]; \
        npy_intp in_step = steps[0], out_step = steps[1]; \
        ty tmp; \
        for (i = 0; i < n; i++) { \
            tmp = *(ty*)in; \
            *((ty*)out) = __builtin_popcountl((unsigned long) tmp); \
            in += in_step; \
            out += out_step; \
        } \
    }

MAKE_POPCOUNT(u8, uint8_t)
MAKE_POPCOUNT(u16, uint16_t)
MAKE_POPCOUNT(u32, uint32_t)
MAKE_POPCOUNT(u64, uint64_t)

PyUFuncGenericFunction funcs[4] = {
    &u8_popcount,
    &u16_popcount,
    &u32_popcount,
    &u64_popcount
};

static char types[8] = {
    NPY_UINT8, NPY_UINT8,
    NPY_UINT16, NPY_UINT16,
    NPY_UINT32, NPY_UINT32,
    NPY_UINT64, NPY_UINT64,
};

static void* data[4] = {NULL, NULL, NULL, NULL};

static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
    "numpy_popcount",
    NULL,
    -1,
    PopcountMethods,
    NULL,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC PyInit_numpy_popcount(void) {
    PyObject* m;
    PyObject* numpy_popcount;
    PyObject* d;
    m = PyModule_Create(&moduledef);
    if (!m) {
        return NULL;
    }

    import_array();
    import_umath();

    numpy_popcount = PyUFunc_FromFuncAndData(funcs, data, types, 4, 1, 1,
            PyUFunc_None, "popcount",
            "numpy_popcount doc", 0);

    d = PyModule_GetDict(m);

    PyDict_SetItemString(d, "popcount", numpy_popcount);
    Py_DECREF(numpy_popcount);

    return m;
}