#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#define PY_SSIZE_T_CLEAN

#include <Python.h>
#include <fits_hcompress.h>
#include <fits_hdecompress.h>
#include <pliocomp.h>
#include <quantize.h>
#include <unquantize.h>
#include <ricecomp.h>

/* The following is needed in order to avoid compilation warnings related to calloc() */
#include <stdlib.h>

/* Define docstrings */
static char module_docstring[] = "Core compression/decompression functions wrapped from cfitsio.";
static char compress_plio_1_c_docstring[] = "Compress data using PLIO_1";
static char decompress_plio_1_c_docstring[] = "Decompress data using PLIO_1";
static char compress_rice_1_c_docstring[] = "Compress data using RICE_1";
static char decompress_rice_1_c_docstring[] = "Decompress data using RICE_1";
static char compress_hcompress_1_c_docstring[] = "Compress data using HCOMPRESS_1";
static char decompress_hcompress_1_c_docstring[] = "Decompress data using HCOMPRESS_1";
static char quantize_float_c_docstring[] = "Quantize float data";
static char quantize_double_c_docstring[] = "Quantize float data";
static char unquantize_float_c_docstring[] = "Unquantize data to float";
static char unquantize_double_c_docstring[] = "Unquantize data to double";

/* Declare the C functions here. */
static PyObject *compress_plio_1_c(PyObject *self, PyObject *args);
static PyObject *decompress_plio_1_c(PyObject *self, PyObject *args);
static PyObject *compress_rice_1_c(PyObject *self, PyObject *args);
static PyObject *decompress_rice_1_c(PyObject *self, PyObject *args);
static PyObject *compress_hcompress_1_c(PyObject *self, PyObject *args);
static PyObject *decompress_hcompress_1_c(PyObject *self, PyObject *args);
static PyObject *quantize_float_c(PyObject *self, PyObject *args);
static PyObject *quantize_double_c(PyObject *self, PyObject *args);
static PyObject *unquantize_float_c(PyObject *self, PyObject *args);
static PyObject *unquantize_double_c(PyObject *self, PyObject *args);
static PyObject *CfitsioException = NULL;

/* Define the methods that will be available on the module. */
static PyMethodDef module_methods[] = {
    {"compress_plio_1_c", compress_plio_1_c, METH_VARARGS, compress_plio_1_c_docstring},
    {"decompress_plio_1_c", decompress_plio_1_c, METH_VARARGS, decompress_plio_1_c_docstring},
    {"compress_rice_1_c", compress_rice_1_c, METH_VARARGS, compress_rice_1_c_docstring},
    {"decompress_rice_1_c", decompress_rice_1_c, METH_VARARGS, decompress_rice_1_c_docstring},
    {"compress_hcompress_1_c",
     compress_hcompress_1_c,
     METH_VARARGS,
     compress_hcompress_1_c_docstring},
    {"decompress_hcompress_1_c",
     decompress_hcompress_1_c,
     METH_VARARGS,
     decompress_hcompress_1_c_docstring},
    {"quantize_float_c", quantize_float_c, METH_VARARGS, quantize_float_c_docstring},
    {"quantize_double_c", quantize_double_c, METH_VARARGS, quantize_double_c_docstring},
    {"unquantize_float_c", unquantize_float_c, METH_VARARGS, unquantize_float_c_docstring},
    {"unquantize_double_c", unquantize_double_c, METH_VARARGS, unquantize_double_c_docstring},
    {NULL, NULL, 0, NULL}
};

/* This is the function that is called on import. */
static PyModuleDef compression = {
    PyModuleDef_HEAD_INIT,
    "_compression",
    module_docstring,
    -1,
    module_methods,
};

PyMODINIT_FUNC PyInit__compression(void)
{
    PyObject *m;
    m = PyModule_Create(&compression);

    /* Initialize new exception object */
    CfitsioException = PyErr_NewException("_compression.CfitsioException", NULL, NULL);

    /* Add exception object to your module */
    PyModule_AddObject(m, "CfitsioException", CfitsioException);

    return m;
};

// Some of the cfitsio compression functions use this function to put an error
// message on the stack.  We provide an implementation which sets the Python
// error state with our custom exception type, and the message provided by the
// cfitsio call.  In our wrapper functions we can then check if the Python error
// state is set and then return NULL to raise the error.
void ffpmsg(const char *err_message)
{
    PyGILState_STATE gstate;
    gstate = PyGILState_Ensure();
    PyErr_SetString(CfitsioException, err_message);
    PyGILState_Release(gstate);
}

/* PLIO/IRAF compression */

static PyObject *compress_plio_1_c(PyObject *self, PyObject *args)
{

    const char *str;
    char *buf;
    Py_ssize_t count;
    PyObject *result;

    int maxelem;
    int tilesize;
    short *compressed_values;
    int compressed_length;
    int *decompressed_values;

    if (!PyArg_ParseTuple(args, "y#i", &str, &count, &tilesize)) {
        return NULL;
    }

    decompressed_values = (int *)str;

    for (int ii = 0; ii < tilesize; ii++) {
        if (decompressed_values[ii] < 0 || decompressed_values[ii] > 16777215) {
            /* plio algorithm only supports positive 24 bit ints */
            PyErr_SetString(PyExc_ValueError, "data out of range for PLIO compression (0 - 2**24)");
            return (PyObject *)NULL;
        }
    }

    // For PLIO imcomp_calc_max_elem in cfitsio does this to calculate max memory:
    maxelem = tilesize;
    // However, when compressing small numbers of random integers you can end up
    // using more memory for the compressed bytes.  In the worst case scenario we
    // tested, compressing a single 4 byte integer will compress to 16 bytes.  We
    // therefore allocate a buffer 4 ints larger than we need here to give that
    // margin of error.
    compressed_values = (short *)calloc(maxelem + 4, sizeof(int));

    decompressed_values = (int *)str;

    compressed_length = pl_p2li(decompressed_values, 1, compressed_values, tilesize);

    if (PyErr_Occurred() != NULL) {
        // If an error condition inside the cfitsio function, the call inside
        // cfitsio should have called the ffpmsg function which sets the Python
        // exception, so we just return here to raise an error.
        return (PyObject *)NULL;
    }

    buf = (char *)compressed_values;

    result = Py_BuildValue("y#", buf, compressed_length * 2);
    free(buf);
    return result;
}

static PyObject *decompress_plio_1_c(PyObject *self, PyObject *args)
{

    const char *str;
    char *buf;
    Py_ssize_t count;
    PyObject *result;

    int tilesize;

    short *compressed_values;
    int *decompressed_values;

    if (!PyArg_ParseTuple(args, "y#i", &str, &count, &tilesize)) {
        return NULL;
    }

    compressed_values = (short *)str;

    decompressed_values = (int *)calloc(tilesize, sizeof(int));

    pl_l2pi(compressed_values, 1, decompressed_values, tilesize);

    if (PyErr_Occurred() != NULL) {
        // If an error condition inside the cfitsio function, the call inside
        // cfitsio should have called the ffpmsg function which sets the Python
        // exception, so we just return here to raise an error.
        return (PyObject *)NULL;
    }

    buf = (char *)decompressed_values;

    result = Py_BuildValue("y#", buf, tilesize * sizeof(int));
    free(buf);
    return result;
}

/* RICE compression */

static PyObject *compress_rice_1_c(PyObject *self, PyObject *args)
{

    const char *str;
    Py_ssize_t count;
    PyObject *result;

    int blocksize, bytepix;

    int maxelem;
    unsigned char *compressed_values;
    int compressed_length;
    signed char *decompressed_values_byte;
    short *decompressed_values_short;
    int *decompressed_values_int;

    if (!PyArg_ParseTuple(args, "y#ii", &str, &count, &blocksize, &bytepix)) {
        return NULL;
    }

    Py_BEGIN_ALLOW_THREADS

        // maxelem adapted from cfitsio's imcomp_calc_max_elem function
        maxelem = count + count / bytepix / blocksize + 2 + 4;

    compressed_values = (unsigned char *)malloc(maxelem);

    if (bytepix == 1) {
        decompressed_values_byte = (signed char *)str;
        compressed_length = fits_rcomp_byte(
            decompressed_values_byte, (int)count, compressed_values, count * 16, blocksize
        );
    }
    else if (bytepix == 2) {
        decompressed_values_short = (short *)str;
        compressed_length = fits_rcomp_short(
            decompressed_values_short, (int)count / 2, compressed_values, count * 16, blocksize
        );
    }
    else {
        decompressed_values_int = (int *)str;
        compressed_length = fits_rcomp(
            decompressed_values_int, (int)count / 4, compressed_values, count * 16, blocksize
        );
    }

    Py_END_ALLOW_THREADS

        if (PyErr_Occurred() != NULL)
    {
        // If an error condition inside the cfitsio function, the call inside
        // cfitsio should have called the ffpmsg function which sets the Python
        // exception, so we just return here to raise an error.
        return (PyObject *)NULL;
    }

    result = Py_BuildValue("y#", compressed_values, compressed_length);
    free(compressed_values);
    return result;
}

static PyObject *decompress_rice_1_c(PyObject *self, PyObject *args)
{

    const char *str;
    char *dbytes;
    Py_ssize_t count;
    PyObject *result;

    int blocksize, bytepix, tilesize;

    unsigned char *compressed_values;
    unsigned char *decompressed_values_byte;
    unsigned short *decompressed_values_short;
    unsigned int *decompressed_values_int;

    if (!PyArg_ParseTuple(args, "y#iii", &str, &count, &blocksize, &bytepix, &tilesize)) {
        return NULL;
    }

    Py_BEGIN_ALLOW_THREADS

        compressed_values = (unsigned char *)str;

    if (bytepix == 1) {
        decompressed_values_byte = (unsigned char *)malloc(tilesize);
        fits_rdecomp_byte(
            compressed_values, (int)count, decompressed_values_byte, tilesize, blocksize
        );
        dbytes = (char *)decompressed_values_byte;
    }
    else if (bytepix == 2) {
        decompressed_values_short = (unsigned short *)malloc(tilesize * 2);
        fits_rdecomp_short(
            compressed_values, (int)count, decompressed_values_short, tilesize, blocksize
        );
        dbytes = (char *)decompressed_values_short;
    }
    else {
        decompressed_values_int = (unsigned int *)malloc(tilesize * 4);
        fits_rdecomp(compressed_values, (int)count, decompressed_values_int, tilesize, blocksize);
        dbytes = (char *)decompressed_values_int;
    }

    Py_END_ALLOW_THREADS

        if (PyErr_Occurred() != NULL)
    {
        // If an error condition inside the cfitsio function, the call inside
        // cfitsio should have called the ffpmsg function which sets the Python
        // exception, so we just return here to raise an error.
        return (PyObject *)NULL;
    }

    result = Py_BuildValue("y#", dbytes, tilesize * bytepix);
    free(dbytes);
    return result;
}

/* HCompress compression */

static PyObject *compress_hcompress_1_c(PyObject *self, PyObject *args)
{

    const char *str;
    Py_ssize_t count;
    PyObject *result;

    int bytepix, nx, ny, scale;
    int status = 0; // Important to initialize this to zero otherwise will fail silently

    int maxelem;
    char *compressed_values;
    int *decompressed_values_int;
    long buffer_size;
    long long *decompressed_values_longlong;

    if (!PyArg_ParseTuple(args, "y#iiii", &str, &count, &nx, &ny, &scale, &bytepix)) {
        return NULL;
    }

    if (bytepix != 4 && bytepix != 8) {
        PyErr_SetString(PyExc_ValueError, "HCompress can only work with 4 or 8 byte integers.");
        return (PyObject *)NULL;
    }

    if ((nx < 4) || (ny < 4)) {
        PyErr_SetString(PyExc_ValueError, "HCOMPRESS requires tiles of at least 4x4 pixels.");
        return (PyObject *)NULL;
    }

    if (count != (Py_ssize_t)nx * ny * bytepix) {
        PyErr_SetString(
            PyExc_ValueError,
            "The tile dimensions and dtype do not match the number of bytes provided."
        );
        return (PyObject *)NULL;
    }

    Py_BEGIN_ALLOW_THREADS

        // maxelem adapted from cfitsio's imcomp_calc_max_elem function
        maxelem = count / 4 * 2.2 + 26;

    // Apparently with the above calculation we can still end up allocating too
    // small of a buffer, this could never happen by more than 32 bytes
    // riiiiiight.
    // TODO: Do a small buffer calculation to tune this number like we did for PLIO
    compressed_values = (char *)calloc(maxelem + 4, sizeof(long long));
    buffer_size = (maxelem + 4) * sizeof(long long);

    if (bytepix == 4) {
        decompressed_values_int = (int *)str;
        fits_hcompress(
            decompressed_values_int, ny, nx, scale, compressed_values, &buffer_size, &status
        );
    }
    else {
        decompressed_values_longlong = (long long *)str;
        fits_hcompress64(
            decompressed_values_longlong, ny, nx, scale, compressed_values, &buffer_size, &status
        );
    }

    Py_END_ALLOW_THREADS

        if (PyErr_Occurred() != NULL)
    {
        // If an error condition inside the cfitsio function, the call inside
        // cfitsio should have called the ffpmsg function which sets the Python
        // exception, so we just return here to raise an error.
        return (PyObject *)NULL;
    }

    if (status != 0) {
        PyErr_SetString(
            PyExc_ValueError, "Status returned from cfitsio is not zero for an unknown reason."
        );
        return (PyObject *)NULL;
    }

    result = Py_BuildValue("y#", compressed_values, buffer_size);
    free(compressed_values);
    return result;
}

static PyObject *decompress_hcompress_1_c(PyObject *self, PyObject *args)
{

    const unsigned char *str;
    char *dbytes;
    Py_ssize_t count;
    PyObject *result;

    int bytepix, nx, ny, scale, smooth;
    int status = 0; // Important to initialize this to zero otherwise will fail silently

    unsigned char *compressed_values;
    int *decompressed_values_int;
    long long *decompressed_values_longlong;

    if (!PyArg_ParseTuple(args, "y#iiiii", &str, &count, &nx, &ny, &scale, &smooth, &bytepix)) {
        return NULL;
    }

    if (bytepix != 4 && bytepix != 8) {
        PyErr_SetString(PyExc_ValueError, "HCompress can only work with 4 or 8 byte integers.");
        return (PyObject *)NULL;
    }

    Py_BEGIN_ALLOW_THREADS

        compressed_values = (unsigned char *)str;

    int na = nx * ny * bytepix;
    dbytes = malloc((size_t)na);

    if (bytepix == 4) {
        decompressed_values_int = (int *)dbytes;
        fits_hdecompress(
            compressed_values, smooth, decompressed_values_int, na, &ny, &nx, &scale, &status
        );
    }
    else {
        decompressed_values_longlong = (long long *)dbytes;
        fits_hdecompress64(
            compressed_values, smooth, decompressed_values_longlong, na, &ny, &nx, &scale, &status
        );
    }

    Py_END_ALLOW_THREADS

        if (PyErr_Occurred() != NULL)
    {
        // If an error condition inside the cfitsio function, the call inside
        // cfitsio should have called the ffpmsg function which sets the Python
        // exception, so we just return here to raise an error.
        return (PyObject *)NULL;
    }

    if (status != 0) {
        PyErr_SetString(
            PyExc_ValueError, "Status returned from cfitsio is not zero for an unknown reason."
        );
        return (PyObject *)NULL;
    }

    // fits_hdecompress[64] always returns 4 byte integers
    result = Py_BuildValue("y#", dbytes, (Py_ssize_t)nx * ny * 4);
    free(dbytes);
    return result;
}

static PyObject *quantize_float_c(PyObject *self, PyObject *args)
{

    const char *input_bytes;
    Py_ssize_t nbytes;
    PyObject *result;

    float *input_data;

    long row, nx, ny;
    int nullcheck;
    float in_null_value;
    float qlevel;
    int dither_method;

    int *quantized_data;
    char *quantized_bytes;
    double bscale, bzero;
    int iminval, imaxval;

    int status;

    Py_ssize_t output_length;

    if (!PyArg_ParseTuple(
            args,
            "y#lllidfi",
            &input_bytes,
            &nbytes,
            &row,
            &nx,
            &ny,
            &nullcheck,
            &in_null_value,
            &qlevel,
            &dither_method
        )) {
        return NULL;
    }

    Py_BEGIN_ALLOW_THREADS

        input_data = (float *)input_bytes;
    quantized_data = (int *)malloc(nx * ny * sizeof(int));

    status = fits_quantize_float(
        row,
        input_data,
        nx,
        ny,
        nullcheck,
        in_null_value,
        qlevel,
        dither_method,
        quantized_data,
        &bscale,
        &bzero,
        &iminval,
        &imaxval
    );

    quantized_bytes = (char *)quantized_data;

    output_length = nx * ny * sizeof(int);

    Py_END_ALLOW_THREADS

        result = Py_BuildValue(
            "y#iddii", quantized_bytes, output_length, status, bscale, bzero, iminval, imaxval
        );
    free(quantized_bytes);
    return result;
}

static PyObject *quantize_double_c(PyObject *self, PyObject *args)
{

    const char *input_bytes;
    Py_ssize_t nbytes;
    PyObject *result;

    double *input_data;

    long row, nx, ny;
    int nullcheck;
    double in_null_value;
    float qlevel;
    int dither_method;

    int *quantized_data;
    char *quantized_bytes;
    double bscale, bzero;
    int iminval, imaxval;

    int status;

    if (!PyArg_ParseTuple(
            args,
            "y#lllidfi",
            &input_bytes,
            &nbytes,
            &row,
            &nx,
            &ny,
            &nullcheck,
            &in_null_value,
            &qlevel,
            &dither_method
        )) {
        return NULL;
    }

    Py_BEGIN_ALLOW_THREADS

        input_data = (double *)input_bytes;
    quantized_data = (int *)malloc(nx * ny * sizeof(int));

    status = fits_quantize_double(
        row,
        input_data,
        nx,
        ny,
        nullcheck,
        in_null_value,
        qlevel,
        dither_method,
        quantized_data,
        &bscale,
        &bzero,
        &iminval,
        &imaxval
    );

    quantized_bytes = (char *)quantized_data;

    Py_END_ALLOW_THREADS

        result = Py_BuildValue(
            "y#iddii",
            quantized_bytes,
            nx * ny * sizeof(int),
            status,
            bscale,
            bzero,
            iminval,
            imaxval
        );
    free(quantized_bytes);
    return result;
}

static PyObject *unquantize_float_c(PyObject *self, PyObject *args)
{

    const char *input_bytes;
    Py_ssize_t nbytes;
    PyObject *result;

    long row, npix;
    int nullcheck;
    int tnull;
    float nullval;
    int dither_method;

    double bscale, bzero;
    int bytepix; // int size
    int status = 0;

    int *anynull;
    float *output_data;
    char *output_bytes;

    if (!PyArg_ParseTuple(
            args,
            "y#llddiiifi",
            &input_bytes,
            &nbytes,
            &row,
            &npix,
            &bscale,
            &bzero,
            &dither_method,
            &nullcheck,
            &tnull,
            &nullval,
            &bytepix
        )) {
        return NULL;
    }

    // TODO: add support, if needed, for nullcheck=1

    Py_BEGIN_ALLOW_THREADS

        anynull = (int *)malloc(npix * sizeof(int));
    output_data = (float *)calloc(npix, sizeof(float));

    if (bytepix == 1) {
        unquantize_i1r4(
            row,
            (unsigned char *)input_bytes,
            npix,
            bscale,
            bzero,
            dither_method,
            nullcheck,
            (unsigned char)tnull,
            nullval,
            NULL,
            anynull,
            output_data,
            &status
        );
    }
    else if (bytepix == 2) {
        unquantize_i2r4(
            row,
            (short *)input_bytes,
            npix,
            bscale,
            bzero,
            dither_method,
            nullcheck,
            (short)tnull,
            nullval,
            NULL,
            anynull,
            output_data,
            &status
        );
    }
    else if (bytepix == 4) {
        unquantize_i4r4(
            row,
            (int *)input_bytes,
            npix,
            bscale,
            bzero,
            dither_method,
            nullcheck,
            (int)tnull,
            nullval,
            NULL,
            anynull,
            output_data,
            &status
        );
    }

    output_bytes = (char *)output_data;

    Py_END_ALLOW_THREADS

        result = Py_BuildValue("y#", output_bytes, npix * sizeof(float));
    free(output_bytes);
    free(anynull);
    return result;
}

static PyObject *unquantize_double_c(PyObject *self, PyObject *args)
{

    const char *input_bytes;
    Py_ssize_t nbytes;
    PyObject *result;

    long row, npix;
    int nullcheck;
    int tnull;
    double nullval;
    int dither_method;

    double bscale, bzero;
    int bytepix; // int size
    int status = 0;

    int *anynull;
    double *output_data;
    char *output_bytes;

    if (!PyArg_ParseTuple(
            args,
            "y#llddiiidi",
            &input_bytes,
            &nbytes,
            &row,
            &npix,
            &bscale,
            &bzero,
            &dither_method,
            &nullcheck,
            &tnull,
            &nullval,
            &bytepix
        )) {
        return NULL;
    }

    // TODO: add support, if needed, for nullcheck=1

    Py_BEGIN_ALLOW_THREADS

        anynull = (int *)malloc(npix * sizeof(int));
    output_data = (double *)malloc(npix * sizeof(double));

    if (bytepix == 1) {
        unquantize_i1r8(
            row,
            (unsigned char *)input_bytes,
            npix,
            bscale,
            bzero,
            dither_method,
            nullcheck,
            (unsigned char)tnull,
            nullval,
            NULL,
            anynull,
            output_data,
            &status
        );
    }
    else if (bytepix == 2) {
        unquantize_i2r8(
            row,
            (short *)input_bytes,
            npix,
            bscale,
            bzero,
            dither_method,
            nullcheck,
            (short)tnull,
            nullval,
            NULL,
            anynull,
            output_data,
            &status
        );
    }
    else if (bytepix == 4) {
        unquantize_i4r8(
            row,
            (int *)input_bytes,
            npix,
            bscale,
            bzero,
            dither_method,
            nullcheck,
            (int)tnull,
            nullval,
            NULL,
            anynull,
            output_data,
            &status
        );
    }

    output_bytes = (char *)output_data;

    Py_END_ALLOW_THREADS

        result = Py_BuildValue("y#", output_bytes, npix * sizeof(double));
    free(output_bytes);
    free(anynull);
    return result;
}
