prepend_to_slang_load_path (".");
set_import_module_path (".:" + get_import_module_path ());

require ("gslfft");
require ("gslrand");

private define correlate1d_wrap (a, kern)
{
   variable na = length (a);
   variable nk = length (kern);
   variable b = a*0.0f;

   variable i, j = ([0:nk-1] - nk/2) + na;

   _for i (0, na-1, 1)
     b[i] = sum (kern*a[(i+j) mod na]);

   return b;
}

private define correlate2d_wrap (a, kern)
{
   variable ny = array_shape (a)[0];
   variable nx = array_shape (a)[1];
   variable my = array_shape (kern)[0];
   variable mx = array_shape (kern)[1];

   variable b = a*0.0f;

   variable ii = ([0:mx-1] - mx/2) + nx;
   variable jj = ([0:my-1] - my/2) + ny;

   variable i, j;
   _for j (0, ny-1, 1)
     {
	_for i (0, nx-1, 1)
	  {
	     b[j,i] = sum (kern * a[(j+jj) mod ny, (i+ii) mod nx]);
	  }
     }
   return b;
}

#if (1)
private define correlate1d_reflect (a, kern)
{
   variable na = length (a);
   variable nk = length (kern);
   variable b = a*0.0f;

   variable i, j, j0 = nk/2;

   % do brute-force indexing for to better catch errors in the
   % correlate function.
   _for i (0, na-1, 1)
     {
	variable s = 0.0;
	_for j (0, nk-1, 1)
	  {
	     variable k = j-j0+i;
	     if (k < 0) k = -k;
	     if (k >= na) k = na-(k-na+2);
	     s += kern[j]*a[k];
	  }
	b[i] = s;
     }

   return b;
}
#endif

private define correlate2d_reflect (a, kern)
{
   variable ny = array_shape (a)[0];
   variable nx = array_shape (a)[1];
   variable my = array_shape (kern)[0];
   variable mx = array_shape (kern)[1];

   variable b = a*0.0f;
   variable ky0 = my/2;
   variable kx0 = mx/2;
   variable i, j, kx, ky;
   _for j (0, ny-1, 1)
     {
	_for i (0, nx-1, 1)
	  {
	     variable s = 0.0;
	     _for ky (0, my-1, 1)
	       {
		  variable jj = ky-ky0+j;
		  if (jj < 0) jj = -jj;
		  if (jj >= ny) jj = ny-(jj-ny+2);
		  _for kx (0, mx-1, 1)
		    {
		       variable ii = kx-kx0+i;
		       if (ii < 0) ii = -ii;
		       if (ii >= nx) ii = nx-(ii-nx+2);
		       s += kern[ky,kx]*a[jj,ii];
		    }
	       }
	     b[j,i] = s;
	  }
     }
   return b;
}

private define correlate1d_nearest (a, kern)
{
   variable na = length (a);
   variable nk = length (kern);
   variable b = a*0.0f;

   variable i, j, j0 = nk/2;

   % do brute-force indexing to better catch errors in the
   % correlate function.
   _for i (0, na-1, 1)
     {
	variable s = 0.0;
	_for j (0, nk-1, 1)
	  {
	     variable k = j-j0+i;
	     if (k < 0) k = 0;
	     if (k >= na) k = na-1;
	     s += kern[j]*a[k];
	  }
	b[i] = s;
     }

   return b;
}

private define correlate2d_nearest (a, kern)
{
   variable ny = array_shape (a)[0];
   variable nx = array_shape (a)[1];
   variable my = array_shape (kern)[0];
   variable mx = array_shape (kern)[1];

   variable b = a*0.0f;
   variable ky0 = my/2;
   variable kx0 = mx/2;
   variable i, j, kx, ky;
   _for j (0, ny-1, 1)
     {
	_for i (0, nx-1, 1)
	  {
	     variable s = 0.0;
	     _for ky (0, my-1, 1)
	       {
		  variable jj = ky-ky0+j;
		  if (jj < 0) jj = 0;
		  if (jj >= ny) jj = ny-1;
		  _for kx (0, mx-1, 1)
		    {
		       variable ii = kx-kx0+i;
		       if (ii < 0) ii = 0;
		       if (ii >= nx) ii = nx-1;
		       s += kern[ky,kx]*a[jj,ii];
		    }
	       }
	     b[j,i] = s;
	  }
     }
   return b;
}

private variable Num_Failed = 0;
private define failed ()
{
   variable args = __pop_list (_NARGS);
   () = fprintf (stderr, __push_list(args));
   () = fputs ("\n", stderr);
   Num_Failed++;
}

define test_correlate1d ()
{
   variable a = [1:10];
   variable k = [-1,1,2];
   variable b1, b2;

   b1 = correlate1d_wrap (a, k); b2 = correlate(a,k; wrap);
   if (any (not feqs(b1,b2,1e-6,1e-7)))
     failed ("correlate on a 1d array using 'wrap' appears to have failed");

   %print (b1-b2);
   b1 = correlate1d_reflect (a, k); b2 = correlate(a,k; reflect);
   if (any (not feqs(b1,b2,1e-6,1e-7)))
     failed ("correlate on a 1d array using 'reflect' appears to have failed");
   %print (b1); print(""); print (b2);
   b1 = correlate1d_nearest (a, k); b2 = correlate(a,k; nearest);
   if (any (not feqs(b1,b2,1e-6,1e-7)))
     failed ("correlate on a 1d array using 'nearest' appears to have failed");
}

define test_correlate2d ()
{
   variable nx = 22, ny = 33;
   variable a = _reshape ([1:ny*nx], [ny,nx]);
   variable i, j;
   _for j (0, ny-1, 1) _for i (0, nx-1, 1)
     a[j,i] = 10*(j+1) + (i+1);
   variable kx = 6, ky = 7;
   variable k = Double_Type[ky,kx]; k[*,*] = (rng_uniform(kx*ky)-0.5);

   variable b1, b2;

   b1 = correlate2d_wrap (a, k); b2 = correlate(a,k; wrap);
   if (any (not feqs(b1,b2,1e-6,1e-7)))
     failed ("correlate on a 2d array using 'wrap' appears to have failed");
   b1 = correlate2d_reflect (a, k); b2 = correlate(a,k; reflect);
   if (any (not feqs(b1,b2,1e-6,1e-7)))
     failed ("correlate on a 2d array using 'reflect' appears to have failed");
   b1 = correlate2d_nearest (a, k); b2 = correlate(a,k; nearest);
   if (any (not feqs(b1,b2,1e-6,1e-7)))
     failed ("correlate on a 2d array using 'nearest' appears to have failed");
}

test_correlate1d ();
test_correlate2d ();

private define g(x, x0,sigma)
{
   variable f = (exp(-0.5*sqr((x-x0)/sigma))
		 / (sigma * sqrt (2*PI)));
   return f/sum(f);
}

private define test_simple_convolve1d ()
{
   variable a = [1, 1, 1, 1, 0, 0];
   variable b = [0, 0, 1, 1, 0, 0];
   % c_i = \sum_j=0^N a_j*b_{i-j}
   % c_0 = a_0*b_0 + a_1*b{-1} + ... a_5*b{-5}
   %     = a*b[[-6:-1]];
   % c_1 = a_1*b
   variable c = [1, 0, 1, 2, 2, 2];
   variable cc = 0.0*a;
   _for (0, length(cc)-1, 1)
     {
	variable i = ();
	variable s = 0.0;
	_for (0, length(cc)-1, 1)
	  {
	     variable j = ();
	     s += a[j]*b[i-j];
	  }
	cc[i] = s;
     }

   variable ab = convolve (a, b; center=0, wrap);
   if (any (abs(ab-c) > 1e-12))
     {
	failed ("simple convolve1d failed: maxdiff=%S\n", maxabs(ab-c));
	c *= 1.0;
	print (ab, &ab); print (c, &c);
	c = strreplace (c, "\n", " ");
	ab = strreplace (ab, "\n", " ");
	failed ("     got: %S", ab);
	failed ("expected: %S", c);
	print (cc);
     }
}
test_simple_convolve1d ();

private define test_convolve1d (n)
{
   variable x = [-200:200:#n];
   variable sa=1, sk=2;
   variable sigma_expect = sqrt(sa^2+sk^2);
   variable a, k, c;
   a = g(x, 0, sa);
   k = g(x, 0, sk);
   k /= sum(k);
   c = convolve (a, k);
   variable b = g(x, 0, sigma_expect);

   % even kernels are always problematic since centering is ambiguous.
   % For those, just check the sum
   ifnot (n mod 2)
     {
	c = sum(c);
	b = sum(b);
     }

   if (any (abs(c-b) > 1e-12))
     failed ("convolve 1d for n=$n may have failed: maxdiff=%g\n"$,
	     maxabs(c-b));
}

test_convolve1d (8);
test_convolve1d (9);

exit (Num_Failed);

