/*
 * vim:syntax=c
 * vim:sw=4
 */
#include <Python.h>
#define PY_ARRAY_UNIQUE_SYMBOL _scipy_signal_ARRAY_API
#define NO_IMPORT_ARRAY
#include <numpy/noprefix.h>

#if PY_VERSION_HEX >= 0x03000000
#define PyNumber_Divide PyNumber_TrueDivide
#endif

#include "sigtools.h"

static void FLOAT_filt(char *b, char *a, char *x, char *y, char *Z,
                       intp len_b, uintp len_x, intp stride_X,
                       intp stride_Y);
static void DOUBLE_filt(char *b, char *a, char *x, char *y, char *Z,
                        intp len_b, uintp len_x, intp stride_X,
                        intp stride_Y);
static void EXTENDED_filt(char *b, char *a, char *x, char *y, char *Z,
                        intp len_b, uintp len_x, intp stride_X,
                        intp stride_Y);
static void CFLOAT_filt(char *b, char *a, char *x, char *y, char *Z,
                        intp len_b, uintp len_x, intp stride_X,
                        intp stride_Y);
static void CDOUBLE_filt(char *b, char *a, char *x, char *y, char *Z,
                         intp len_b, uintp len_x, intp stride_X,
                         intp stride_Y);
static void CEXTENDED_filt(char *b, char *a, char *x, char *y, char *Z,
                         intp len_b, uintp len_x, intp stride_X,
                         intp stride_Y);
static void OBJECT_filt(char *b, char *a, char *x, char *y, char *Z,
                        intp len_b, uintp len_x, intp stride_X,
                        intp stride_Y);

typedef void (BasicFilterFunction) (char *, char *,  char *, char *, char *, intp, uintp, intp, intp);

static BasicFilterFunction *BasicFilterFunctions[256];

void
scipy_signal_sigtools_linear_filter_module_init(void)
{
    int k;
    for (k = 0; k < 256; ++k) {
        BasicFilterFunctions[k] = NULL;
    }
    BasicFilterFunctions[NPY_FLOAT] = FLOAT_filt;
    BasicFilterFunctions[NPY_DOUBLE] = DOUBLE_filt;
    BasicFilterFunctions[NPY_LONGDOUBLE] = EXTENDED_filt;
    BasicFilterFunctions[NPY_CFLOAT] = CFLOAT_filt;
    BasicFilterFunctions[NPY_CDOUBLE] = CDOUBLE_filt;
    BasicFilterFunctions[NPY_CLONGDOUBLE] = CEXTENDED_filt;
    BasicFilterFunctions[NPY_OBJECT] = OBJECT_filt;
}

/* There is the start of an OBJECT_filt, but it may need work */

static int
RawFilter(const PyArrayObject * b, const PyArrayObject * a,
          const PyArrayObject * x, const PyArrayObject * zi,
          const PyArrayObject * zf, PyArrayObject * y, int axis,
          BasicFilterFunction * filter_func);

/*
 * XXX: Error checking not done yet
 */
PyObject*
scipy_signal_sigtools_linear_filter(PyObject * NPY_UNUSED(dummy), PyObject * args)
{
    PyObject *b, *a, *X, *Vi;
    PyArrayObject *arY, *arb, *ara, *arX, *arVi, *arVf;
    int axis, typenum, theaxis, st;
    char *ara_ptr, input_flag = 0, *azero;
    intp na, nb, nal;
    BasicFilterFunction *basic_filter;

    axis = -1;
    Vi = NULL;
    if (!PyArg_ParseTuple(args, "OOO|iO", &b, &a, &X, &axis, &Vi)) {
        return NULL;
    }

    typenum = PyArray_ObjectType(b, 0);
    typenum = PyArray_ObjectType(a, typenum);
    typenum = PyArray_ObjectType(X, typenum);
    if (Vi != NULL) {
        typenum = PyArray_ObjectType(Vi, typenum);
    }

    arY = arVf = arVi = NULL;
    ara = (PyArrayObject *) PyArray_ContiguousFromObject(a, typenum, 1, 1);
    arb = (PyArrayObject *) PyArray_ContiguousFromObject(b, typenum, 1, 1);
    arX = (PyArrayObject *) PyArray_FromObject(X, typenum, 0, 0);
    /* XXX: fix failure handling here */
    if (ara == NULL || arb == NULL || arX == NULL) {
        goto fail;
    }

    if (axis < -arX->nd || axis > arX->nd - 1) {
        PyErr_SetString(PyExc_ValueError, "selected axis is out of range");
        goto fail;
    }
    if (axis < 0) {
        theaxis = arX->nd + axis;
    } else {
        theaxis = axis;
    }

    if (Vi != NULL) {
        arVi = (PyArrayObject *) PyArray_FromObject(Vi, typenum,
                                                    arX->nd, arX->nd);
        if (arVi == NULL)
            goto fail;

        input_flag = 1;
    }

    arY = (PyArrayObject *) PyArray_SimpleNew(arX->nd,
                                              arX->dimensions, typenum);
    if (arY == NULL) {
        goto fail;
    }

    if (input_flag) {
        arVf = (PyArrayObject *) PyArray_SimpleNew(arVi->nd,
                                                   arVi->dimensions,
                                                   typenum);
        if (arVf == NULL) {
            goto fail;
        }
    }

    if (arX->descr->type_num < 256) {
        basic_filter = BasicFilterFunctions[(int) (arX->descr->type_num)];
    }
    else {
        basic_filter = NULL;
    }
    if (basic_filter == NULL) {
        PyObject *msg, *str;
        char *s;

        str = PyObject_Str((PyObject*)arX->descr);
        if (str == NULL) {
            goto fail;
        }
        s = PyString_AsString(str);
        msg = PyString_FromFormat(
                        "input type '%s' not supported\n", s);
        Py_DECREF(str);
        if (msg == NULL) {
            goto fail;
        }
        PyErr_SetObject(PyExc_NotImplementedError, msg);
        Py_DECREF(msg);
        goto fail;
    }

    /* Skip over leading zeros in vector representing denominator (a) */
    azero = PyArray_Zero(ara);
    if (azero == NULL) {
        goto fail;
    }
    ara_ptr = ara->data;
    nal = PyArray_ITEMSIZE(ara);
    st = memcmp(ara_ptr, azero, nal);
    PyDataMem_FREE(azero);
    if (st == 0) {
        PyErr_SetString(PyExc_ValueError,
                        "BUG: filter coefficient a[0] == 0 not supported yet");
        goto fail;
    }

    na = PyArray_SIZE(ara);
    nb = PyArray_SIZE(arb);
    if (input_flag) {
        if (arVi->dimensions[theaxis] != (na > nb ? na : nb) - 1) {
            PyErr_SetString(PyExc_ValueError,
                            "The number of initial conditions must be max([len(a),len(b)]) - 1");
            goto fail;
        }
    }

    st = RawFilter(arb, ara, arX, arVi, arVf, arY, theaxis, basic_filter);
    if (st) {
        goto fail;
    }

    Py_XDECREF(ara);
    Py_XDECREF(arb);
    Py_XDECREF(arX);
    Py_XDECREF(arVi);

    if (!input_flag) {
        return PyArray_Return(arY);
    } else {
        return Py_BuildValue("(NN)", arY, arVf);
    }


  fail:
    Py_XDECREF(ara);
    Py_XDECREF(arb);
    Py_XDECREF(arX);
    Py_XDECREF(arVi);
    Py_XDECREF(arVf);
    Py_XDECREF(arY);
    return NULL;
}

/*
 * Copy the first nxzfilled items of x into xzfilled , and fill the rest with
 * 0s
 */
static int
zfill(const PyArrayObject * x, intp nx, char *xzfilled, intp nxzfilled)
{
    char *xzero;
    intp i, nxl;
    PyArray_CopySwapFunc *copyswap = x->descr->f->copyswap;

    nxl = PyArray_ITEMSIZE(x);

    /* PyArray_Zero does not take const pointer, hence the cast */
    xzero = PyArray_Zero((PyArrayObject *) x);
    if (xzero == NULL) return -1;

    if (nx > 0) {
        for (i = 0; i < nx; ++i) {
            copyswap(xzfilled + i * nxl, x->data + i * nxl, 0, NULL);
        }
    }
    for (i = nx; i < nxzfilled; ++i) {
        copyswap(xzfilled + i * nxl, xzero, 0, NULL);
    }

    PyDataMem_FREE(xzero);

    return 0;
}

/*
 * a and b assumed to be contiguous
 *
 * XXX: this code is very conservative, and could be considerably sped up for
 * the usual cases (like contiguity).
 *
 * XXX: the code should be refactored (at least with/without initial
 * condition), some code is wasteful here
 */
static int
RawFilter(const PyArrayObject * b, const PyArrayObject * a,
          const PyArrayObject * x, const PyArrayObject * zi,
          const PyArrayObject * zf, PyArrayObject * y, int axis,
          BasicFilterFunction * filter_func)
{
    PyArrayIterObject *itx, *ity, *itzi, *itzf;
    intp nitx, i, nxl, nzfl, j;
    intp na, nb, nal, nbl;
    intp nfilt;
    char *azfilled, *bzfilled, *zfzfilled, *yoyo;
    PyArray_CopySwapFunc *copyswap = x->descr->f->copyswap;

    itx = (PyArrayIterObject *) PyArray_IterAllButAxis((PyObject *) x,
                                                       &axis);
    if (itx == NULL) {
        PyErr_SetString(PyExc_MemoryError, "Could not create itx");
        goto fail;
    }
    nitx = itx->size;

    ity = (PyArrayIterObject *) PyArray_IterAllButAxis((PyObject *) y,
                                                       &axis);
    if (ity == NULL) {
        PyErr_SetString(PyExc_MemoryError, "Could not create ity");
        goto clean_itx;
    }

    if (zi != NULL) {
        itzi = (PyArrayIterObject *) PyArray_IterAllButAxis((PyObject *)
                                                            zi, &axis);
        if (itzi == NULL) {
            PyErr_SetString(PyExc_MemoryError, "Could not create itzi");
            goto clean_ity;
        }

        itzf = (PyArrayIterObject *) PyArray_IterAllButAxis((PyObject *)
                                                            zf, &axis);
        if (itzf == NULL) {
            PyErr_SetString(PyExc_MemoryError, "Could not create itzf");
            goto clean_itzi;
        }
    }

    na = PyArray_SIZE(a);
    nal = PyArray_ITEMSIZE(a);
    nb = PyArray_SIZE(b);
    nbl = PyArray_ITEMSIZE(b);

    nfilt = na > nb ? na : nb;

    azfilled = malloc(nal * nfilt);
    if (azfilled == NULL) {
        PyErr_SetString(PyExc_MemoryError, "Could not create azfilled");
        goto clean_itzf;
    }
    bzfilled = malloc(nbl * nfilt);
    if (bzfilled == NULL) {
        PyErr_SetString(PyExc_MemoryError, "Could not create bzfilled");
        goto clean_azfilled;
    }

    nxl = PyArray_ITEMSIZE(x);
    zfzfilled = malloc(nxl * (nfilt - 1));
    if (zfzfilled == NULL) {
        PyErr_SetString(PyExc_MemoryError, "Could not create zfzfilled");
        goto clean_bzfilled;
    }
    /* Initialize zero filled buffers to 0, so that we can use
     * Py_XINCREF/Py_XDECREF on it for object arrays (necessary for
     * copyswap to work correctly). Stricly speaking, it is not needed for
     * fundamental types (as values are copied instead of pointers, without
     * refcounts), but oh well...
     */
    memset(azfilled, 0, nal * nfilt);
    memset(bzfilled, 0, nbl * nfilt);
    memset(zfzfilled, 0, nxl * (nfilt - 1));

    if (zfill(a, na, azfilled, nfilt) == -1) goto clean_zfzfilled;
    if (zfill(b, nb, bzfilled, nfilt) == -1) goto clean_zfzfilled;

    /* XXX: Check that zf and zi have same type ? */
    if (zf != NULL) {
        nzfl = PyArray_ITEMSIZE(zf);
    } else {
        nzfl = 0;
    }

    /* Iterate over the input array */
    for (i = 0; i < nitx; ++i) {
        if (zi != NULL) {
            yoyo = itzi->dataptr;
            /* Copy initial conditions zi in zfzfilled buffer */
            for (j = 0; j < nfilt - 1; ++j) {
                copyswap(zfzfilled + j * nzfl, yoyo, 0, NULL);
                yoyo += itzi->strides[axis];
            }
            PyArray_ITER_NEXT(itzi);
        } else {
            if (zfill(x, 0, zfzfilled, nfilt - 1) == -1) goto clean_zfzfilled;
        }

        filter_func(bzfilled, azfilled,
                    itx->dataptr, ity->dataptr, zfzfilled,
                    nfilt, PyArray_DIM(x, axis), itx->strides[axis],
                    ity->strides[axis]);
    
        if (PyErr_Occurred()) {
            goto clean_zfzfilled; 
        }
        PyArray_ITER_NEXT(itx);
        PyArray_ITER_NEXT(ity);

        /* Copy tmp buffer fo final values back into zf output array */
        if (zi != NULL) {
            yoyo = itzf->dataptr;
            for (j = 0; j < nfilt - 1; ++j) {
                copyswap(yoyo, zfzfilled + j * nzfl, 0, NULL);
                yoyo += itzf->strides[axis];
            }
            PyArray_ITER_NEXT(itzf);
        }
    }

    /* Free up allocated memory */
    free(zfzfilled);
    free(bzfilled);
    free(azfilled);

    if (zi != NULL) {
        Py_DECREF(itzf);
        Py_DECREF(itzi);
    }
    Py_DECREF(ity);
    Py_DECREF(itx);

    return 0;

clean_zfzfilled:
    free(zfzfilled);
clean_bzfilled:
    free(bzfilled);
clean_azfilled:
    free(azfilled);
clean_itzf:
    if (zf != NULL) {
        Py_DECREF(itzf);
    }
clean_itzi:
    if (zi != NULL) {
        Py_DECREF(itzi);
    }
clean_ity:
    Py_DECREF(ity);
clean_itx:
    Py_DECREF(itx);
fail:
    return -1;
}

/*****************************************************************
 *   This is code for a 1-D linear-filter along an arbitrary     *
 *   dimension of an N-D array.                                  *
 *****************************************************************/

/**begin repeat
 * #type = float, double, npy_longdouble#
 * #NAME = FLOAT, DOUBLE, EXTENDED#
 */
static void @NAME@_filt(char *b, char *a, char *x, char *y, char *Z,
                       intp len_b, uintp len_x, intp stride_X,
                       intp stride_Y)
{
    char *ptr_x = x, *ptr_y = y;
    @type@ *ptr_Z, *ptr_b;
    @type@ *ptr_a;
    @type@ *xn, *yn;
    const @type@ a0 = *((@type@ *) a);
    intp n;
    uintp k;

    for (k = 0; k < len_x; k++) {
        ptr_b = (@type@ *) b;   /* Reset a and b pointers */
        ptr_a = (@type@ *) a;
        xn = (@type@ *) ptr_x;
        yn = (@type@ *) ptr_y;
        if (len_b > 1) {
            ptr_Z = ((@type@ *) Z);
            *yn = *ptr_Z + *ptr_b / a0 * *xn;   /* Calculate first delay (output) */
            ptr_b++;
            ptr_a++;
            /* Fill in middle delays */
            for (n = 0; n < len_b - 2; n++) {
                *ptr_Z =
                    ptr_Z[1] + *xn * (*ptr_b / a0) - *yn * (*ptr_a / a0);
                ptr_b++;
                ptr_a++;
                ptr_Z++;
            }
            /* Calculate last delay */
            *ptr_Z = *xn * (*ptr_b / a0) - *yn * (*ptr_a / a0);
        } else {
            *yn = *xn * (*ptr_b / a0);
        }

        ptr_y += stride_Y;      /* Move to next input/output point */
        ptr_x += stride_X;
    }
}

static void C@NAME@_filt(char *b, char *a, char *x, char *y, char *Z,
                        intp len_b, uintp len_x, intp stride_X,
                        intp stride_Y)
{
    char *ptr_x = x, *ptr_y = y;
    @type@ *ptr_Z, *ptr_b;
    @type@ *ptr_a;
    @type@ *xn, *yn;
    @type@ a0r = ((@type@ *) a)[0];
    @type@ a0i = ((@type@ *) a)[1];
    @type@ a0_mag, tmpr, tmpi;
    intp n;
    uintp k;

    a0_mag = a0r * a0r + a0i * a0i;
    for (k = 0; k < len_x; k++) {
        ptr_b = (@type@ *) b;   /* Reset a and b pointers */
        ptr_a = (@type@ *) a;
        xn = (@type@ *) ptr_x;
        yn = (@type@ *) ptr_y;
        if (len_b > 1) {
            ptr_Z = ((@type@ *) Z);
            tmpr = ptr_b[0] * a0r + ptr_b[1] * a0i;
            tmpi = ptr_b[1] * a0r - ptr_b[0] * a0i;
            /* Calculate first delay (output) */
            yn[0] = ptr_Z[0] + (tmpr * xn[0] - tmpi * xn[1]) / a0_mag;
            yn[1] = ptr_Z[1] + (tmpi * xn[0] + tmpr * xn[1]) / a0_mag;
            ptr_b += 2;
            ptr_a += 2;
            /* Fill in middle delays */
            for (n = 0; n < len_b - 2; n++) {
                tmpr = ptr_b[0] * a0r + ptr_b[1] * a0i;
                tmpi = ptr_b[1] * a0r - ptr_b[0] * a0i;
                ptr_Z[0] =
                    ptr_Z[2] + (tmpr * xn[0] - tmpi * xn[1]) / a0_mag;
                ptr_Z[1] =
                    ptr_Z[3] + (tmpi * xn[0] + tmpr * xn[1]) / a0_mag;
                tmpr = ptr_a[0] * a0r + ptr_a[1] * a0i;
                tmpi = ptr_a[1] * a0r - ptr_a[0] * a0i;
                ptr_Z[0] -= (tmpr * yn[0] - tmpi * yn[1]) / a0_mag;
                ptr_Z[1] -= (tmpi * yn[0] + tmpr * yn[1]) / a0_mag;
                ptr_b += 2;
                ptr_a += 2;
                ptr_Z += 2;
            }
            /* Calculate last delay */

            tmpr = ptr_b[0] * a0r + ptr_b[1] * a0i;
            tmpi = ptr_b[1] * a0r - ptr_b[0] * a0i;
            ptr_Z[0] = (tmpr * xn[0] - tmpi * xn[1]) / a0_mag;
            ptr_Z[1] = (tmpi * xn[0] + tmpr * xn[1]) / a0_mag;
            tmpr = ptr_a[0] * a0r + ptr_a[1] * a0i;
            tmpi = ptr_a[1] * a0r - ptr_a[0] * a0i;
            ptr_Z[0] -= (tmpr * yn[0] - tmpi * yn[1]) / a0_mag;
            ptr_Z[1] -= (tmpi * yn[0] + tmpr * yn[1]) / a0_mag;
        } else {
            tmpr = ptr_b[0] * a0r + ptr_b[1] * a0i;
            tmpi = ptr_b[1] * a0r - ptr_b[0] * a0i;
            yn[0] = (tmpr * xn[0] - tmpi * xn[1]) / a0_mag;
            yn[1] = (tmpi * xn[0] + tmpr * xn[1]) / a0_mag;
        }

        ptr_y += stride_Y;      /* Move to next input/output point */
        ptr_x += stride_X;

    }
}
/**end repeat**/

static void OBJECT_filt(char *b, char *a, char *x, char *y, char *Z,
                        intp len_b, uintp len_x, intp stride_X,
                        intp stride_Y)
{
    char *ptr_x = x, *ptr_y = y;
    PyObject **ptr_Z, **ptr_b;
    PyObject **ptr_a;
    PyObject **xn, **yn;
    PyObject **a0 = (PyObject **) a;
    PyObject *tmp1, *tmp2, *tmp3;
    intp n;
    uintp k;

    /* My reference counting might not be right */
    for (k = 0; k < len_x; k++) {
        ptr_b = (PyObject **) b;        /* Reset a and b pointers */
        ptr_a = (PyObject **) a;
        xn = (PyObject **) ptr_x;
        yn = (PyObject **) ptr_y;
        if (len_b > 1) {
            ptr_Z = ((PyObject **) Z);
            /* Calculate first delay (output) */
            tmp1 = PyNumber_Multiply(*ptr_b, *xn);
            if (tmp1 == NULL) return;
            tmp2 = PyNumber_Divide(tmp1, *a0);
            if (tmp2 == NULL) {
                Py_DECREF(tmp1);
                return;
            }
            tmp3 = PyNumber_Add(tmp2, *ptr_Z);
            Py_XDECREF(*yn);
            *yn = tmp3; /* Steals the reference */
            Py_DECREF(tmp1);
            Py_DECREF(tmp2);
            if (tmp3 == NULL) return;
            ptr_b++;
            ptr_a++;
            
            /* Fill in middle delays */
            for (n = 0; n < len_b - 2; n++) {
                tmp1 = PyNumber_Multiply(*xn, *ptr_b);
                if (tmp1 == NULL) return;
                tmp2 = PyNumber_Divide(tmp1, *a0);
                if (tmp2 == NULL) {
                    Py_DECREF(tmp1);
                    return;
                }
                tmp3 = PyNumber_Add(tmp2, ptr_Z[1]);
                Py_DECREF(tmp1);
                Py_DECREF(tmp2);
                if (tmp3 == NULL) return;
                tmp1 = PyNumber_Multiply(*yn, *ptr_a);
                if (tmp1 == NULL) {
                    Py_DECREF(tmp3);
                    return;
                }
                tmp2 = PyNumber_Divide(tmp1, *a0);
                Py_DECREF(tmp1);
                if (tmp2 == NULL) {
                    Py_DECREF(tmp3);
                    return;
                }
                Py_XDECREF(*ptr_Z);
                *ptr_Z = PyNumber_Subtract(tmp3, tmp2);
                Py_DECREF(tmp2);
                Py_DECREF(tmp3);
                if (*ptr_Z == NULL) return;
                ptr_b++;
                ptr_a++;
                ptr_Z++;
            }
            /* Calculate last delay */
            tmp1 = PyNumber_Multiply(*xn, *ptr_b);
            if (tmp1 == NULL) return;
            tmp3 = PyNumber_Divide(tmp1, *a0);
            Py_DECREF(tmp1);
            if (tmp3 == NULL) return;
            tmp1 = PyNumber_Multiply(*yn, *ptr_a);
            if (tmp1 == NULL) {
                Py_DECREF(tmp3);
                return;
            }
            tmp2 = PyNumber_Divide(tmp1, *a0);
            Py_DECREF(tmp1);
            if (tmp2 == NULL) {
                Py_DECREF(tmp3);
                return;
            }
            Py_XDECREF(*ptr_Z);
            *ptr_Z = PyNumber_Subtract(tmp3, tmp2);
            Py_DECREF(tmp2);
            Py_DECREF(tmp3);
        } else {
            tmp1 = PyNumber_Multiply(*xn, *ptr_b);
            if (tmp1 == NULL) return;
            Py_XDECREF(*yn);
            *yn = PyNumber_Divide(tmp1, *a0);
            Py_DECREF(tmp1);
            if (*yn == NULL) return;
        }

        ptr_y += stride_Y;      /* Move to next input/output point */
        ptr_x += stride_X;
    }
}