#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

#include <Python.h>
#include <stdlib.h> // malloc/free
#include <numpy/arrayobject.h>
#include "numpy/ufuncobject.h"
#include "compute_bounds.h"

/* Define docstrings */
static char module_docstring[] = "Fast sigma clipping";
static char _sigma_clip_fast_docstring[] =
    "Compute sigma clipping bounds.\n\n"
    "Parameters\n"
    "----------\n"
    "data : 2D float64 array, shape (1, N) - input data\n"
    "mask : 1D bool array, shape (N,) - True means already masked\n"
    "use_median : bool - use median (True) or mean (False)\n"
    "use_mad_std : bool - use MAD-based std (True) or regular std (False)\n"
    "max_iter : int - maximum number of sigma-clipping iterations\n"
    "sigma_lower : float - lower sigma threshold\n"
    "sigma_upper : float - upper sigma threshold\n";

/* Declare the C functions here. */
static void _sigma_clip_fast(
    char **args, npy_intp const *dimensions, npy_intp const *steps, void *data
);

/* Define the methods that will be available on the module. */
static PyMethodDef module_methods[] = {{NULL, NULL, 0, NULL}};

/* This is the function that is called on import. */

#define MOD_INIT(name) PyMODINIT_FUNC PyInit_##name(void)
#define MOD_DEF(ob, name, doc, methods) \
    static struct PyModuleDef moduledef = { \
        PyModuleDef_HEAD_INIT, name, doc, -1, methods, NULL, NULL, NULL, NULL \
    }; \
    ob = PyModule_Create(&moduledef);

MOD_INIT(_fast_sigma_clip)
{
    PyObject *m, *d = NULL;
    PyUFuncObject *ufunc;
    static char types[9] = {
        NPY_DOUBLE, /* data array */
        NPY_BOOL,   /* mask array */
        NPY_BOOL,   /* use median */
        NPY_BOOL,   /* use mad_std */
        NPY_INT,    /* max iter */
        NPY_DOUBLE, /* sigma low */
        NPY_DOUBLE, /* sigma high */
        NPY_DOUBLE, /* output: lower bound */
        NPY_DOUBLE  /* output: upper bound */
    };
    /* In principle, can have multiple functions for multiple input types */
    static PyUFuncGenericFunction funcs[1] = {&_sigma_clip_fast};
    static void *data[1] = {NULL};
    MOD_DEF(m, "_fast_sigma_clip", module_docstring, module_methods);
    if (m == NULL) {
        goto fail;
    }
    d = PyModule_GetDict(m); /* borrowed ref. */
    if (d == NULL) {
        goto fail;
    }
    import_array();
    import_umath();

    ufunc = (PyUFuncObject *)PyUFunc_FromFuncAndDataAndSignature(
        funcs,
        data,
        types,
        1,
        7,
        2,
        PyUFunc_None,
        "_sigma_clip_fast",
        _sigma_clip_fast_docstring,
        0,
        "(n),(n),(),(),(),(),()->(),()"
    );
    if (ufunc == NULL) {
        goto fail;
    }
    PyDict_SetItemString(d, "_sigma_clip_fast", (PyObject *)ufunc);
    Py_DECREF(ufunc);
    return m;

fail:
    Py_XDECREF(m);
    Py_XDECREF(d);
    return NULL;
}


static void _sigma_clip_fast(
    char **args, npy_intp const *dimensions, npy_intp const *steps, void *data
)
{
    npy_intp i_o, i;
    int count;
    /* dimensions, pointers and step sizes for outer loop */
    npy_intp n_o = *dimensions++;
    char *array = *args++;
    npy_intp s_array = *steps++;
    char *mask = *args++;
    npy_intp s_mask = *steps++;
    char *use_median = *args++;
    npy_intp s_use_median = *steps++;
    char *use_mad_std = *args++;
    npy_intp s_use_mad_std = *steps++;
    char *max_iter = *args++;
    npy_intp s_max_iter = *steps++;
    char *sigma_low = *args++;
    npy_intp s_sigma_low = *steps++;
    char *sigma_high = *args++;
    npy_intp s_sigma_high = *steps++;
    char *bound_low = *args++;
    npy_intp s_bound_low = *steps++;
    char *bound_high = *args++;
    npy_intp s_bound_high = *steps++;
    /* dimension and step sizes for inner loop */
    npy_intp n_i = dimensions[0];
    char *in_array, *in_mask;
    npy_intp is_array = *steps++;
    npy_intp is_mask = *steps++;

    double *data_buffer = NULL;
    double *mad_buffer = NULL;

    // data_buffer is used to store the current values being sigma clipped
    data_buffer = (double *)malloc(n_i * sizeof(double));
    if (data_buffer == NULL) {
        PyErr_NoMemory();
        return;
    }

    for (i_o = 0; i_o < n_o; i_o++,
        array += s_array,
        mask += s_mask,
        use_median += s_use_median,
        use_mad_std += s_use_mad_std,
        max_iter += s_max_iter,
        sigma_low += s_sigma_low,
        sigma_high += s_sigma_high,
        bound_low += s_bound_low,
        bound_high += s_bound_high) {
        /* copy to buffer */
        in_array = array;
        in_mask = mask;
        count = 0;
        for (i = 0; i < n_i; i++, in_array += is_array, in_mask += is_mask) {
            if (*(uint8_t *)in_mask == 0) {
                data_buffer[count] = *(double *)in_array;
                count += 1;
            }
        }
        if (count > 0) {

            // If we are using mad_std, we need to prepare an additional buffer
            // that is used in the calculation. We just need to allocate this once
            // and can use it in any future loop iteration that needs it.
            if ((*(npy_bool *)use_mad_std) && mad_buffer == NULL) {
                mad_buffer = (double *)malloc(n_i * sizeof(double));
                if (mad_buffer == NULL) {
                    PyErr_NoMemory();
                    return;
                }
            }

            compute_sigma_clipped_bounds(
                data_buffer,
                count,
                (int)(*(npy_bool *)use_median),
                (int)(*(npy_bool *)use_mad_std),
                *(int *)max_iter,
                *(double *)sigma_low,
                *(double *)sigma_high,
                (double *)bound_low,
                (double *)bound_high,
                mad_buffer
            );
        }
        else {
            *(double *)bound_low = NPY_NAN;
            *(double *)bound_high = NPY_NAN;
        }
    }
    free((void *)data_buffer);
    if (mad_buffer != NULL) {
        free((void *)mad_buffer);
    }
}
