#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include "Python.h"
#include "numpy/arrayobject.h"
#include <complex.h>


static PyObject *rotate_grid3d(PyObject *dummy, PyObject *args)
{
    // Argument variables
    PyObject *arg1=NULL, *arg2=NULL, *arg4=NULL;
    int radius, nearest;
    PyArrayObject *py_grid=NULL, *py_rotmat=NULL, *py_out=NULL;

    // Pointers
    float *grid, *rotmat, *out;
    npy_intp *out_shape, *grid_shape;

    npy_intp out_size, out_slice, grid_size, grid_slice;
    int z, y, x, radius2, x0, y0, z0, dist2_z, dist2_zy, dist2_zyx, 
        out_z, out_zy, out_zyx, grid_zyx;
    float xcoor_z, ycoor_z, zcoor_z, xcoor_zy, ycoor_zy, zcoor_zy,
           xcoor_zyx, ycoor_zyx, zcoor_zyx;
    // Trilinear interpolation
    int x1, y1, z1, offset0, offset1;
    float dx, dy, dz, dx1, dy1, dz1, c00, c10, c01, c11, c0, c1, c;

    // Parse arguments
    if (!PyArg_ParseTuple(args, "OOiOi", &arg1, &arg2, &radius, &arg4, &nearest))
        return NULL;

    py_grid = (PyArrayObject *) PyArray_FROM_OTF(arg1, NPY_FLOAT32, NPY_ARRAY_IN_ARRAY);
    if (py_grid == NULL)
        goto fail;
    py_rotmat = (PyArrayObject *) PyArray_FROM_OTF(arg2, NPY_FLOAT32, NPY_ARRAY_IN_ARRAY);
    if (py_rotmat == NULL)
        goto fail;
    py_out = (PyArrayObject *) PyArray_FROM_OTF(arg4, NPY_FLOAT32, NPY_ARRAY_OUT_ARRAY);
    if (py_out == NULL)
        goto fail;

    // Get pointers to arrays and shape info
    grid = (float *) PyArray_DATA(py_grid);
    rotmat = (float *) PyArray_DATA(py_rotmat);
    out = (float *) PyArray_DATA(py_out);

    out_shape = PyArray_DIMS(py_out);
    out_slice = out_shape[2] * out_shape[1];
    out_size = PyArray_SIZE(py_out);
    grid_shape = PyArray_DIMS(py_grid);
    grid_slice = grid_shape[2] * grid_shape[1];
    grid_size = PyArray_SIZE(py_grid);

    radius2 = radius * radius;

    for (z = -radius; z <= radius; z++) {
        dist2_z = z * z;
        if (dist2_z > radius2)
            continue;
        // Rotate points
        xcoor_z = rotmat[6] * z;
        ycoor_z = rotmat[7] * z;
        zcoor_z = rotmat[8] * z;

        // Indices
        out_z = z * out_slice;
        if (z < 0)
            out_z += out_size;

        for (y = -radius; y <= radius; y++) {
            dist2_zy = dist2_z + y * y;
            if (dist2_zy > radius2)
                continue;

            xcoor_zy = xcoor_z + rotmat[3] * y;
            ycoor_zy = ycoor_z + rotmat[4] * y;
            zcoor_zy = zcoor_z + rotmat[5] * y;

            out_zy = out_z + y * out_shape[2];
            if (y < 0)
                out_zy += out_slice;

            for (x = -radius; x <= radius; x++) {
                dist2_zyx = dist2_zy + x * x;
                if (dist2_zyx > radius2)
                    continue;

                xcoor_zyx = xcoor_zy + rotmat[0] * x;
                ycoor_zyx = ycoor_zy + rotmat[1] * x;
                zcoor_zyx = zcoor_zy + rotmat[2] * x;
                // Indices
                out_zyx = out_zy + x;
                if (x < 0)
                    out_zyx += out_shape[2];

                if (nearest > 0) {
                    // Nearest interpolation
                    x0 = (int) round(xcoor_zyx);
                    y0 = (int) round(ycoor_zyx);
                    z0 = (int) round(zcoor_zyx);

                    // Grid index
                    grid_zyx = z0 * grid_slice + y0 * grid_shape[2] + x0;
                    if (x0 < 0)
                        grid_zyx += grid_shape[2];
                    if (y0 < 0)
                        grid_zyx += grid_slice;
                    if (z0 < 0)
                        grid_zyx += grid_size;

                    out[out_zyx] = grid[grid_zyx];
                } else {
                    // Tri-linear interpolation
                    //
                    x0 = (int) floor(xcoor_zyx);
                    y0 = (int) floor(ycoor_zyx);
                    z0 = (int) floor(zcoor_zyx);
                    x1 = x0 + 1;
                    y1 = y0 + 1;
                    z1 = z0 + 1;

                    // Grid index
                    grid_zyx = z0 * grid_slice + y0 * grid_shape[2] + x0;
                    if (x0 < 0)
                        grid_zyx += grid_shape[2];
                    if (y0 < 0)
                        grid_zyx += grid_slice;
                    if (z0 < 0)
                        grid_zyx += grid_size;

                    dx = xcoor_zyx - x0;
                    dy = ycoor_zyx - y0;
                    dz = zcoor_zyx - z0;
                    dx1 = 1 - dx;
                    dy1 = 1 - dy;
                    dz1 = 1 - dz;

                    offset1 = 1;
                    if (x1 == 0)
                        offset1 -= grid_shape[2];
                    c00 = grid[grid_zyx] * dx1 + 
                          grid[grid_zyx + offset1] * dx;

                    offset0 = grid_shape[2];
                    if (y1 == 0)
                        offset0 -= grid_slice;
                    offset1 = offset0 + 1;
                    if (x1 == 0)
                        offset1 -= grid_shape[2];

                    c10 = grid[grid_zyx + offset0] * dx1 + 
                          grid[grid_zyx + offset1] * dx;

                    offset0 = grid_slice;
                    if (z1 == 0)
                        offset0 -= grid_size;
                    offset1 = offset0 + 1;
                    if (x1 == 0)
                        offset1 -= grid_shape[2];
                    c01 = grid[grid_zyx + offset0] * dx1 + 
                          grid[grid_zyx + offset1] * dx;

                    offset0 = grid_slice + grid_shape[2];
                    if (z1 == 0)
                        offset0 -= grid_size;
                    if (y1 == 0)
                        offset0 -= grid_slice;
                    offset1 = offset0 + 1;
                    if (x1 == 0)
                        offset1 -= grid_shape[2];
                    c11 = grid[grid_zyx + offset0] * dx1 + 
                          grid[grid_zyx + offset1] * dx;

                    c0 = c00 * dy1 + c10 * dy;
                    c1 = c01 * dy1 + c11 * dy;

                    c = c0 * dz1 + c1 * dz;

                    out[out_zyx] = c;
                }
            }
        }
    }
    // Clean up objects
    Py_DECREF(py_grid);
    Py_DECREF(py_rotmat);
    Py_DECREF(py_out);
    Py_INCREF(Py_None);
    return Py_None;

fail:
    // Clean up objects
    Py_XDECREF(py_grid);
    Py_XDECREF(py_rotmat);
    PyArray_XDECREF(py_out);
    return NULL;
}


static PyObject *conj_multiply(PyObject *dummy, PyObject *args)
{
    //Calculate conj(a) * (b)
    //

    PyObject *arg1=NULL, *arg2=NULL, *arg3=NULL;
    PyArrayObject *py_in1=NULL, *py_in2=NULL, *py_out=NULL;

    float complex *in1, *in2, *out;
    npy_intp out_size, i;

    // Parse arguments
    if (!PyArg_ParseTuple(args, "OOO", &arg1, &arg2, &arg3))
        return NULL;

    py_in1 = (PyArrayObject *) PyArray_FROM_OTF(arg1, NPY_COMPLEX64, NPY_ARRAY_IN_ARRAY);
    if (py_in1 == NULL)
        goto fail;
    py_in2 = (PyArrayObject *) PyArray_FROM_OTF(arg2, NPY_COMPLEX64, NPY_ARRAY_IN_ARRAY);
    if (py_in2 == NULL)
        goto fail;
    py_out = (PyArrayObject *) PyArray_FROM_OTF(arg3, NPY_COMPLEX64, NPY_ARRAY_OUT_ARRAY);
    if (py_out == NULL)
        goto fail;

    // Get pointers to arrays and shape info
    in1 = (float complex *) PyArray_DATA(py_in1);
    in2 = (float complex *) PyArray_DATA(py_in2);
    out = (float complex *) PyArray_DATA(py_out);
    out_size = PyArray_SIZE(py_out);

    for (i = 0; i < out_size; i++)
        out[i] = conj(in1[i]) * in2[i];

    // Clean up objects
    Py_DECREF(py_in1);
    Py_DECREF(py_in2);
    Py_DECREF(py_out);
    Py_INCREF(Py_None);
    return Py_None;

fail:
    // Clean up objects
    Py_XDECREF(py_in1);
    Py_XDECREF(py_in2);
    PyArray_XDECREF(py_out);
    return NULL;
}


static PyMethodDef mymethods[] = {
    {"conj_multiply", conj_multiply, METH_VARARGS, ""},
    {"rotate_grid3d", rotate_grid3d, METH_VARARGS, ""},
    {NULL, NULL, 0, NULL}
};


// NEW: Module definition struct for Python 3
static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
    "_extensions",        // name of module
    NULL,                 // docstring (or NULL)
    -1,                   // size of per-interpreter module state (-1 = global)
    mymethods             // your function table
};

// NEW: Module init function for Python 3
PyMODINIT_FUNC PyInit__extensions(void) {
    import_array();       // numpy setup
    return PyModule_Create(&moduledef);
}

// PyMODINIT_FUNC
// init_extensions(void)
// {
//     (void) Py_InitModule("_extensions", mymethods);
//     import_array();
// };
//
