#include "Python.h"
#include <math.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#include "_splinemodule.h"

#define NO_IMPORT_ARRAY
#include "numpy/arrayobject.h"

#ifndef M_PI
#define M_PI           3.14159265358979323846  /* pi */
#endif

#ifdef __GNUC__
#define CONJ(a) (~((__complex__ double)a))
#define ABSQ(a) (__real__ (a*CONJ(a)))
#else
#define CONJ(a) ((a))
#define ABSQ(a) ( (a*CONJ(a)))
#endif

void
compute_root_from_lambda(lambda, r, omega)
     double lambda;
     double *r;
     double *omega;
{
    double xi;
    double tmp, tmp2;

    tmp = sqrt(3 + 144*lambda);
    xi = 1 - 96*lambda + 24*lambda * tmp;
    *omega = atan(sqrt((144*lambda - 1.0)/xi));
    tmp2 = sqrt(xi);
    *r = (24*lambda - 1 - tmp2)/(24*lambda)	* sqrt((48*lambda + 24*lambda*tmp))/tmp2;
    return;
}


{{py:

CNAME = ['C', 'Z']
CTYPE = ['__complex__ float', '__complex__ double']

RNAME = ['S', 'D']
RTYPE = ['float', 'double']

NAMES = CNAME + RNAME
TYPES = CTYPE + RTYPE
RTYPES = RTYPE + RTYPE

}}


/* Implement the following difference equation */
/* y[n] = a1 * x[n] + a2 * y[n-1]  */
/* with a given starting value loaded into the array */

{{for SUB, TYP in zip(NAMES, TYPES)}}
{{if SUB in CNAME}}
#ifdef __GNUC__
{{endif}}
void {{SUB}}_IIR_order1({{TYP}} a1, {{TYP}} a2, {{TYP}} *x, {{TYP}} *y,
             int N, int stridex, int stridey)
{
    {{TYP}} *yvec = y + stridey;
    {{TYP}} *xvec = x + stridex;
    int n;

    for (n = 1; n < N; n++) {
	*yvec = *xvec * a1 + *(yvec - stridey) * a2;
	yvec += stridey;
	xvec += stridex;
    }
}
{{if SUB in CNAME}}
#endif
{{endif}}
{{endfor}}


/* Implement the following difference equation */
/* y[n] = a1 * x[n] + a2 * y[n-1]  + a3 * y[n-2] */
/* with two starting values loaded into the array */

{{for SUB, TYP in zip(NAMES, TYPES)}}
{{if SUB in CNAME}}
#ifdef __GNUC__
{{endif}}
void {{SUB}}_IIR_order2({{TYP}} a1, {{TYP}} a2, {{TYP}} a3,
             {{TYP}} *x, {{TYP}} *y, int N, int stridex, int stridey)
{
    {{TYP}} *yvec = y + 2*stridey;
    {{TYP}} *xvec = x + 2*stridex;
    int n;

    for (n = 2; n < N; n++) {
	*yvec = *xvec * a1 + *(yvec - stridey) * a2 + *(yvec - 2*stridey) * a3;
	yvec += stridey;
	xvec += stridex;
    }
}
{{if SUB in CNAME}}
#endif
{{endif}}
{{endfor}}


/* Implement a second order IIR difference equation using a cascade
   of first order sections.  Suppose the transfer function is
                  cs
   H(z) =   -------------------
            (1-z1/z) ( 1-z2/z)

   then the following pair is implemented with one starting value loaded in
   the output array and the starting value for the intermediate array
   passed in as yp0.

   y1[n] = x[n] + z1 y1[n-1]
   yp[n] = cs y1[n] + z2 yp[n-1]

*/


{{for SUB, TYP in zip(NAMES, TYPES)}}
{{if SUB in CNAME}}
#ifdef __GNUC__
{{endif}}
void {{SUB}}_IIR_order2_cascade({{TYP}} cs, {{TYP}} z1, {{TYP}} z2,
                     {{TYP}} y1_0, {{TYP}} *x, {{TYP}} *yp,
                     int N, int stridex, int stridey)
{
    {{TYP}} *yvec = yp + stridey;
    {{TYP}} *xvec = x + stridex;
    int n;

    for (n = 1; n < N; n++) {
	y1_0 = *xvec + y1_0 * z1;
	*yvec = cs * y1_0 + *(yvec - stridey) * z2;
	yvec += stridey;
	xvec += stridex;
    }
}
{{if SUB in CNAME}}
#endif
{{endif}}
{{endfor}}

/* Implement a smoothing IIR filter with mirror-symmetric boundary conditions
   using a cascade of first-order sections.  The second section uses a
   reversed sequence.  This implements the following transfer function:
                    c0
   H(z) = ---------------------------
           (1-z1/z) (1 - z1 z)

   with the following difference equations:

   yp[n] = x[n] + z1 yp[n-1]
     with starting condition:
   yp[0] = x[0] + Sum(z1^(k+1) x[k],k=0..Infinity)

   and

   y[n] = z1 y[n+1] + c0 yp[n]
     with starting condition:
   y[N-1] = z1 / (z1-1) yp[N-1]

   The resulting signal will have mirror symmetric boundary conditions as well.

   If memory could not be allocated for the temporary vector yp, the
   function returns -1 otherwise it returns 0.

   z1 should be less than 1;

*/

{{for SUB, TYP, RTYP in zip(NAMES, TYPES, RTYPES)}}
{{if SUB in CNAME}}
#ifdef __GNUC__
{{endif}}
int {{SUB}}_IIR_forback1({{TYP}} c0, {{TYP}} z1, {{TYP}} *x, {{TYP}} *y,
                   int N, int stridex, int stridey, {{RTYP}} precision)
{
    {{TYP}} *yp = NULL;
    {{TYP}} *xptr = x;
    {{TYP}} yp0, powz1, diff;
    {{RTYP}} err;
    int k;

    if (ABSQ(z1) >= 1.0) return -2; /* z1 not less than 1 */

    /* Initialize memory for loop */
    if ((yp = malloc(N*sizeof({{TYP}})))==NULL) return -1;

   /* Fix starting value assuming mirror-symmetric boundary conditions. */
    yp0 = x[0];
    powz1 = 1.0;
    k = 0;
    precision *= precision;
    do {
    	yp[0] = yp0;
    	powz1 *= z1;
    	yp0 += powz1 * (*xptr);
    	diff = powz1;
    	err = ABSQ(diff);
    	xptr += stridex;
    	k++;
    } while((err > precision) && (k < N));
    if (k >= N){
        /* sum did not converge */
        free(yp);
        return -3;
    }
    yp[0] = yp0;

    {{SUB}}_IIR_order1(1.0, z1, x, yp, N, stridex, 1);

    *(y + (N - 1)*stridey) = -c0 / (z1 - 1.0) * yp[N-1];

    {{SUB}}_IIR_order1(c0, z1, yp + N - 1, y + (N - 1)*stridey, N, -1, -stridey);

    free(yp);
    return 0;
}
{{if SUB in CNAME}}
#endif
{{endif}}
{{endfor}}

/* h must be odd length */
/* strides in units of sizeof(DATA TYPE) bytes */

{{for SUB, TYP in zip(NAMES, TYPES)}}
{{if SUB in CNAME}}
#ifdef __GNUC__
{{endif}}
void {{SUB}}_FIR_mirror_symmetric({{TYP}} *in, {{TYP}} *out, int N,
                            {{TYP}} *h, int Nh, int instride, int outstride)
{
    int n, k;
    int Nhdiv2 = Nh >> 1;
    {{TYP}} *outptr;
    {{TYP}} *inptr;
    {{TYP}} *hptr;

    /* first part boundary conditions */
    outptr = out;
    for (n=0; n < Nhdiv2; n++) {
	*outptr = 0.0;
	hptr = h;
	inptr = in + (n + Nhdiv2)*instride;
	for (k=-Nhdiv2; k <= n; k++) {
	    *outptr += *hptr++ * *inptr;
	    inptr -= instride;
	}
	inptr += instride;
	for (k=n+1; k <= Nhdiv2; k++) {
	    *outptr += *hptr++ * *inptr;
	    inptr += instride;
	}
	outptr += outstride;
    }

    /* middle section */
    outptr = out + Nhdiv2*outstride;
    for (n=Nhdiv2; n < N-Nhdiv2; n++) {
	*outptr = 0.0;
	hptr = h;
	inptr = in + (n + Nhdiv2)*instride;
	for (k=-Nhdiv2; k <= Nhdiv2; k++) {
	    *outptr += *hptr++ * *inptr;
	    inptr -= instride;
	}
	outptr += outstride;
    }

    /* end boundary conditions */
    outptr = out + (N - Nhdiv2)*outstride;
    for (n=N-Nhdiv2; n < N; n++) {
	*outptr = 0.0;
	hptr = h;
	inptr = in + (2*N - 1 - n - Nhdiv2)*instride;
	for (k=-Nhdiv2; k <= n-N; k++) {
	    *outptr += *hptr++ * *inptr;
	    inptr += instride;
	}
	inptr -= instride;
	for (k=n+1-N; k <= Nhdiv2; k++) {
	    *outptr += *hptr++ * *inptr;
	    inptr -= instride;
	}
	outptr += outstride;
    }

}
{{if SUB in CNAME}}
#endif
{{endif}}
{{endfor}}

{{for SUB, TYP in zip(NAMES, TYPES)}}
{{if SUB in CNAME}}
#ifdef __GNUC__
{{endif}}
int {{SUB}}_separable_2Dconvolve_mirror({{TYP}} *in, {{TYP}} *out,
                              int M, int N, {{TYP}} *hr, {{TYP}} *hc,
                              int Nhr, int Nhc,
                              npy_intp *instrides, npy_intp *outstrides)
{
    int m, n;
    {{TYP}} *tmpmem;
    {{TYP}} *inptr=NULL, *outptr=NULL;

    tmpmem = malloc(M*N*sizeof({{TYP}}));
    if (tmpmem == NULL) return -1;

    if (Nhr > 0) {
	/* filter across rows */
	inptr = in;
	outptr = tmpmem;
	for (m = 0; m < M; m++) {
	    {{SUB}}_FIR_mirror_symmetric (inptr, outptr, N, hr, Nhr, instrides[1], 1);
	    inptr += instrides[0];
	    outptr += N;
	}
    }
    else
	memmove(tmpmem, in, M*N*sizeof({{TYP}}));

    if (Nhc > 0) {
	/* filter down columns */
	inptr = tmpmem;
	outptr = out;
	for (n = 0; n < N; n++) {
	    {{SUB}}_FIR_mirror_symmetric (inptr, outptr, M, hc, Nhc, N, outstrides[0]);
	    outptr += outstrides[1];
	    inptr += 1;
	}
    }
    else
	memmove(out, tmpmem, M*N*sizeof({{TYP}}));

    free(tmpmem);
    return 0;
}
{{if SUB in CNAME}}
#endif
{{endif}}
{{endfor}}


{{for SUB, TYP in zip(RNAME, RTYPE)}}
static {{TYP}} {{SUB}}_hc(int, {{TYP}}, double, double);
static {{TYP}} {{SUB}}_hs(int, {{TYP}}, double, double);
{{endfor}}

{{for SUB, TYP in zip(RNAME, RTYPE)}}
{{TYP}} {{SUB}}_hc(int k, {{TYP}} cs, double r, double omega)
{
    if (k < 0) return 0.0;
    if (omega == 0.0)
	return cs * pow(r, (double )k) * (k+1);
    else if (omega == M_PI)
	return cs * pow(r, (double )k) * (k+1) * (1 - 2*(k % 2));
    return cs * pow(r, (double) k) * sin(omega * (k+1)) / sin(omega);
}
{{endfor}}

{{for SUB, TYP in zip(RNAME, RTYPE)}}
{{TYP}} {{SUB}}_hs(int k, {{TYP}} cs, double rsq, double omega)
{
    {{TYP}} cssq;
    {{TYP}} c0;
    double gamma, rsupk;

    cssq = cs * cs;
    k = abs(k);
    rsupk = pow(rsq, ((double ) k) / 2.0);
    if (omega == 0.0) {
	c0 = (1+rsq)/ ((1-rsq)*(1-rsq)*(1-rsq)) * cssq;
	gamma = (1-rsq) / (1+rsq);
	return c0 * rsupk * (1 + gamma * k);
    }
    if (omega == M_PI) {
	c0 = (1+rsq)/ ((1-rsq)*(1-rsq)*(1-rsq)) * cssq;
	gamma = (1-rsq) / (1+rsq) * (1 - 2 * (k % 2));
	return c0 * rsupk * (1 + gamma * k);
    }
    c0 = cssq * (1.0+rsq)/(1.0-rsq) / (1-2*rsq*cos(2*omega) + rsq*rsq);
    gamma = (1.0 - rsq)/ (1.0+rsq) / tan(omega);
    return c0 * rsupk * (cos(omega*k) + gamma * sin(omega * k));
}

{{endfor}}

/* Implement a smoothing IIR filter with mirror-symmetric boundary conditions
   using a cascade of second-order sections.  The second section uses a
   reversed sequence.  This implements the following transfer function:

                          cs^2
   H(z) = --------------------------------------
          (1 - a2/z - a3/z^2) (1 - a2 z -a3 z^2 )

   where a2 = (2 r cos omega)
         a3 = - r^2
	 cs = 1 - 2 r cos omega + r^2

   with the following difference equations:

   yp[n] = cs*x[n] - b1 yp[n-1] - b2 yp[n-2]
     with starting conditions:
   yp[0] = hc[0] x[0] + Sum(hc[k+1]*x[k],k=0..Infinity)
   yp[1] = hc[0] x[1] + hc[1] x[0] + Sum(hc[k+2] x[k], k=0..Infinity)

   and

   y[n] = cs*yp[n] - b1 y[n+1] -b2 y[n+2]
     with starting conditions:
   y[N-1] = Sum((hs[k] + hs[k+1])x[N-1-k],k=0..Infinity)
   y[N-2] = Sum((hs[k-1] + hs[k+2])x[N-1-k],k=0..Infinity)

   The resulting signal will have mirror symmetric boundary conditions as well.

   If memory could not be allocated for the temporary vector yp, the
   function returns -1 otherwise it returns 0.

   z1 should be less than 1;

*/

{{for SUB, TYP in zip(RNAME, RTYPE)}}
int {{SUB}}_IIR_forback2 (double r, double omega, {{TYP}} *x, {{TYP}} *y,
		int N, int stridex, int stridey, {{TYP}} precision) {
    {{TYP}} cs;
    {{TYP}} *yp = NULL;
    {{TYP}} *yptr;
    {{TYP}} *xptr;
    {{TYP}} yp0;
    {{TYP}} yp1;
    double rsq;
    {{TYP}} diff;
    {{TYP}} err;
    {{TYP}} a2, a3;
    int k;

    if (r >= 1.0) return -2; /* z1 not less than 1 */

    /* Initialize memory for loop */
    if ((yp = malloc(N*sizeof({{TYP}})))==NULL) return -1;

    rsq = r * r;
    a2 = 2 * r * cos(omega);
    a3 = -rsq;
    cs = 1 - 2 * r * cos(omega) + rsq;

   /* Fix starting values assuming mirror-symmetric boundary conditions. */
    yp0 = {{SUB}}_hc(0, cs, r, omega) * x[0];
    k = 0;
    precision *= precision;
    xptr = x;
    do {
	yp[0] = yp0;
	diff = {{SUB}}_hc(k+1, cs, r, omega);
  	yp0 += diff * (*xptr);
	err = diff * diff;
	xptr += stridex;
	k++;
    } while((err > precision) && (k < N));
    if (k >= N) {free(yp); return -3;}     /* sum did not converge */
    yp[0] = yp0;

    yp1 = {{SUB}}_hc(0, cs, r, omega) * (*(x+stridex));
    yp1 += {{SUB}}_hc(1, cs, r, omega) * x[0];
    k = 0;
    xptr = x;
    do {
	yp[1] = yp1;
	diff = {{SUB}}_hc(k+2, cs, r, omega);
	yp1 += diff * (*xptr);
	err = diff * diff;
	xptr += stridex;
	k++;
    } while((err > precision) && (k < N));
    if (k >= N) {free(yp); return -3;}     /* sum did not converge */
    yp[1] = yp1;

    {{SUB}}_IIR_order2(cs, a2, a3, x, yp, N, stridex, 1);

   /* Fix starting values assuming mirror-symmetric boundary conditions. */
    yp0 = 0.0;
    k = 0;
    yptr = y + (N-1)*stridey;
    xptr = x + (N-1)*stridex;
    do {
	*yptr = yp0;
	diff = ({{SUB}}_hs(k, cs, rsq, omega) + {{SUB}}_hs(k+1, cs, rsq, omega));
	yp0 += diff * (*xptr);
	err = diff * diff;
	xptr -= stridex;
	k++;
    } while((err > precision) && (k < N));
    if (k >= N) {free(yp); return -3;}     /* sum did not converge */
    *yptr = yp0;

    yp1 = 0.0;
    k = 0;
    yptr -= stridey;        /* Initialize in next-to-last slot in output array */
    xptr = x + (N-1)*stridex;
    do {
	*yptr = yp1;
	diff = ({{SUB}}_hs(k-1, cs, rsq, omega) + {{SUB}}_hs(k+2, cs, rsq, omega));
	yp1 += diff * (*xptr);
	err = diff * diff;
	xptr -= stridex;
	k++;
    } while((err > precision) && (k < N));
    if (k >= N) {free(yp); return -3;}     /* sum did not converge */
    *yptr = yp1;

    {{SUB}}_IIR_order2(cs, a2, a3, yp+N-1, yptr+stridey, N, -1, -stridey);

    free(yp);
    return 0;
}

{{endfor}}


/* Find the cubic spline coefficients of an image
   image is M rows by N columns stored rowise in memory (vary column number
          first). It will be replaced with the spline coefficients.
   lambda is a smoothing parameter (lambda = 100 approximately corresponds
          to a cutoff frequency of 0.1*(sample freq))
   strides is an integer array [rowstride, colstride]
          telling how much memory in units of sizeof(DATA TYPE) bytes to skip
	  to get to the next element.
*/

/* to get the (smoothed) image back mirror-symmetric convolve with a length
   three separable FIR filter [1.0, 4.0, 1.0]/ 6.0
*/

{{for SUB, TYP in zip(RNAME, RTYPE)}}
int {{SUB}}_cubic_spline2D({{TYP}} *image, {{TYP}} *coeffs, int M, int N, double lambda,
		 npy_intp *strides, npy_intp *cstrides, {{TYP}} precision) {
    double r, omega;
    {{TYP}} *inptr;
    {{TYP}} *coptr;
    {{TYP}} *tmpmem;
    {{TYP}} *tptr;
    int m, n, retval=0;

    tmpmem = malloc(N*M*sizeof({{TYP}}));
    if (tmpmem == NULL) return -1;

    if (lambda <= 1.0 / 144.0) {
	/* normal cubic spline */
	r = -2 + sqrt(3.0);

	/* Loop over rows */
	inptr = image;
	tptr = tmpmem;
	for (m = 0; m < M; m++) {
	    retval = {{SUB}}_IIR_forback1 (-r*6.0, r, inptr, tptr, N, strides[1], 1, precision);
	    if (retval < 0) break;
	    inptr += strides[0];
	    tptr += N;
	}

	if (retval >=0) {
	    /* Loop over columns */
	    tptr = tmpmem;
	    coptr = coeffs;
	    for (n = 0; n < N; n++) {
		retval = {{SUB}}_IIR_forback1 (-r*6.0, r, tptr, coptr, M, N, cstrides[0], precision);
		if (retval < 0) break;
		coptr += cstrides[1];
		tptr += 1;
	    }
	}
	free(tmpmem);
	return retval;
    }

    /* Smoothing spline */

    /* Compute r and omega from lambda */
    compute_root_from_lambda(lambda, &r, &omega);

    /* Loop over rows */
    inptr = image;
    tptr = tmpmem;
    for (m = 0; m < M; m++) {
	retval = {{SUB}}_IIR_forback2 (r, omega, inptr, tptr, N, strides[1],
				       1, precision);
	if (retval < 0) break;
	inptr += strides[0];
	tptr += N;
    }
    /* Loop over columns */
    tptr = tmpmem;
    coptr = coeffs;
    for (n = 0; n < N; n++) {
	retval = {{SUB}}_IIR_forback2 (r, omega, tptr, coptr, M, N,
				       cstrides[0], precision);
	if (retval < 0) break;
	coptr += cstrides[1];
	tptr += 1;
    }

    free(tmpmem);
    return retval;
}

{{endfor}}

/* Find the quadratic spline coefficients of an image
   image is M rows by N columns stored rowise in memory (vary column number
          first). It will be replaced with the spline coefficients.
   lambda is a smoothing parameter (lambda = 100 approximately corresponds
          to a cutoff frequency of 0.1*(sample freq))
	  must be zero for now.
   strides is an integer array [rowstride, colstride]
          telling how much memory in units of sizeof(DATA TYPE) bytes to skip
	  to get to the next element.
*/

/* to get the (smoothed) image back mirror-symmetric convolve with a length
   three separable FIR filter [1.0, 6.0, 1.0]/ 8.0
*/

{{for SUB, TYP in zip(RNAME, RTYPE)}}
int {{SUB}}_quadratic_spline2D({{TYP}} *image, {{TYP}} *coeffs, int M, int N, double lambda,
		     npy_intp *strides, npy_intp *cstrides, {{TYP}} precision) {
    double r;
    {{TYP}} *inptr;
    {{TYP}} *coptr;
    {{TYP}} *tmpmem;
    {{TYP}} *tptr;
    int m,n, retval=0;

    if (lambda > 0) return -2;

    tmpmem = malloc(N*M*sizeof({{TYP}}));
    if (tmpmem == NULL) return -1;

    /* normal quadratic spline */
    r = -3 + 2*sqrt(2.0);

    /* Loop over rows */
    inptr = image;
    tptr = tmpmem;
    for (m = 0; m < M; m++) {
      retval = {{SUB}}_IIR_forback1 (-r*8.0, r, inptr, tptr, N, strides[1], 1, precision);
      if (retval < 0) break;
      inptr += strides[0];
      tptr += N;
    }

    if (retval >=0) {
    /* Loop over columns */
      tptr = tmpmem;
      coptr = coeffs;
      for (n = 0; n < N; n++) {
	retval = {{SUB}}_IIR_forback1 (-r*8.0, r, tptr, coptr, M, N, cstrides[0], precision);
	if (retval < 0) break;
	coptr += cstrides[1];
	tptr += 1;
      }
    }
    free(tmpmem);
    return retval;
}

{{endfor}}