require ("gslcore");
_gslcore_import_module ("gslfft", current_namespace());

define fft ()
{
   variable args = __pop_args (_NARGS);
   return _gsl_fft_complex (__push_args (__tmp(args)));
}

% This routine converts a 2-d kernel to an ny x nx version for convolution
private define wrap_kernel_2d (kern, ny, nx, center)
{
   variable
     my = array_shape(kern)[0],
     mx = array_shape(kern)[1],
     x0 = (center ? (mx-1)/2 : 0),
     y0 = (center ? (my-1)/2 : 0),
     new_kern = Complex_Type[ny,nx];

   new_kern[[0:my-1]-y0,[0:mx-1]-x0] = kern;
   return new_kern;
}

private define wrap_kernel_1d (kern, nx, center)
{
   variable
     mx = length (kern),
     new_kern = Complex_Type[nx],
     x0 = (center ? (mx-1)/2 : 0);

   new_kern[[0:mx-1]-x0] = kern;

   return new_kern;
}

private define find_nice_fft_grid_size (minsize)
{
   % Make size a multiple of 2, 3, 5, or 7
   variable max_2 = 0, max_3 = 0, max_5 = 0, max_7 = 0, size;
   size = 1; while (size < minsize) { max_2++; size *= 2; }
   size = 1; while (size < minsize) { max_3++; size *= 3; }
   size = 1; while (size < minsize) { max_5++; size *= 5; }
   size = 1; while (size < minsize) { max_7++; size *= 7; }

   variable e2 = 0, e3 = 0, e5 = 0, e7 = 0;
   variable expons;
   variable min_diff = _Inf;
   variable size2 = 1;
   _for e2 (0, max_2, 1)
     {
	variable size3 = 1;
	_for e3 (0, max_3, 1)
	  {
	     variable size5 = 1;
	     _for e5 (0, max_5, 1)
	       {
		  variable size7 = 1;
		  _for e7 (0, max_7, 1)
		    {
		       size = size2 * size3 * size5 * size7;
		       if (size >= minsize)
			 {
			    variable diff = size - minsize;
			    if (diff < min_diff)
			      {
				 expons = [e2, e3, e5, e7];
				 min_diff = diff;
			      }
			    break;
			 }
		       size7 *= 7;
		    }
		  size5 *= 5;
	       }
	     size3 *= 3;
	  }
	size2 *= 2;
     }
   size = minsize + min_diff;
#iffalse
   vmessage ("size=%d => %d = 2^%d * 3^%d * 5^%d * 7^%d", minsize,
	     size, expons[0], expons[1], expons[2], expons[3]);
#endif
   return size;
}

private define convolve2d (image, kernel, q)
{
   variable
     ny = array_shape(image)[0],
     nx = array_shape(image)[1],
     my = array_shape(kernel)[0],
     mx = array_shape(kernel)[1],
     wrap = q.wrap, padx = nx+mx, pady = ny+my;

   if (wrap == NULL)
     {
	padx = nx + mx/2;
	pady = ny + my/2;
     }
   else if (wrap == "wrap")
     {
	padx = nx;
	pady = ny;
     }

   padx = find_nice_fft_grid_size (padx);
   pady = find_nice_fft_grid_size (pady);

   if (wrap == "wrap")
     {
	if (nx != padx)
	  padx = find_nice_fft_grid_size (nx+mx);
	if (ny != pady)
	  pady = find_nice_fft_grid_size (ny+my);
     }

   variable cimage = Complex_Type[pady,padx];
   cimage[[:ny-1], [:nx-1]] = __tmp(image);

   if (q.pad != 0.0)
     {
	cimage[[ny:],*] = q.pad;
	cimage[*, [nx:]] = q.pad;
     }
   variable i, j, dx = mx/2, dy = my/2;

   if (wrap == "wrap")
     {
	if (nx != padx)
	  {
	     i = [0:dx-1]; cimage[*,nx+i] = cimage[*,i];
	     i = [dx:mx-1]; cimage[*,padx-mx+i] = cimage[*,nx-mx+i];
	  }
	if (ny != pady)
	  {
	     j = [0:dy-1]; cimage[ny+j,*] = cimage[j,*];
	     j = [dy:my-1]; cimage[pady-my+j,*] = cimage[ny-my+j,*];
	     if (nx != padx)
	       {
		  i = [0:dx-1]; j = [0:dy-1];
		  cimage[ny+j, nx+i] = cimage[j,i];
		  i = [dx:mx-1]; j = [dy:my-1];
		  cimage[pady-my+j, padx-mx+i] = cimage[ny-my+j,nx-mx+i];
	       }
	  }
     }
   if (wrap == "reflect")
     {
	if (nx != padx)
	  {
	     i = [0:dx-1]; cimage[*,nx+i] = cimage[*,nx-i-2];
	     i = [dx:mx-1]; cimage[*,padx-mx+i] = cimage[*,mx-i];
	  }
	if (ny != pady)
	  {
	     j = [0:dy-1]; cimage[ny+j,*] = cimage[ny-j-2,*];
	     j = [dy:my-1]; cimage[pady-my+j,*] = cimage[my-j,*];
	     if (nx != padx)
	       {
		  i = [0:dx-1]; j = [0:dy-1];
		  cimage[ny+j, nx+i] = cimage[ny-j-2,nx-i-2];
		  i = [dx:mx-1]; j = [dy:my-1];
		  cimage[pady-my+j, padx-mx+i] = cimage[my-j,mx-i];
	       }
	  }
     }
   if (wrap == "nearest")
     {
	variable slice;
	if (nx != padx)
	  {
	     j = pady-my+dy-1;
	     slice = cimage[[:j], nx-1];
	     slice[[ny:j]] = cimage[ny-1,nx-1];
	     _for i (0, dx-1, 1) cimage[[:j],nx+i] = slice;
	     slice = cimage[*,0]; slice[[ny:]] = cimage[ny-1,0];
	     _for i (dx, mx-1, 1) cimage[*,padx-mx+i] = slice;
	  }
	if (ny != pady)
	  {
	     i = padx-mx+dx-1;
	     slice = cimage[ny-1,[:i]];
	     slice[[nx:i]] = cimage[ny-1,nx-1];
	     _for j (0, dy-1, 1) cimage[ny+j,[:i]] = slice;
	     slice = cimage[0,*]; slice[[nx:]] = cimage[0,nx-1];
	     _for j (dy, my-1, 1) cimage[pady-my+j,*] = slice;
	     if (nx != padx)
	       {
		  i = [0:dx-1]; j = [0:dy-1];
		  cimage[ny+j, nx+i] = cimage[ny-1,nx-1];
		  i = [dx:mx-1]; j = [dy:my-1];
		  cimage[pady-my+j, padx-mx+i] = cimage[0,0];
	       }
	  }
     }

   image = cimage;
   kernel = wrap_kernel_2d (kernel, pady, padx, q.center);

   image = Real(fft(fft(__tmp(kernel), 1) * fft(__tmp(image), 1), -1));
   return __tmp(image)[[:ny-1],[:nx-1]];
}

private define convolve1d (image, kernel, q)
{
   variable
     nx = length (image),
     mx = length (kernel),
     wrap = q.wrap, padx;

   if (wrap == NULL)
     padx = nx + mx/2;
   else if (wrap == "wrap")
     padx = nx; % Periodic BC -- no need to add padding for the kernel
   else
     padx = nx + mx;

   padx = find_nice_fft_grid_size (padx);
   if ((wrap == "wrap") and (padx != nx))
     padx = find_nice_fft_grid_size (nx + mx);

   variable cimage = Complex_Type[padx];
   cimage[[:nx-1]] = __tmp(image);

   if (q.pad != 0.0)
     cimage[[nx:]] = q.pad;

   variable i, dx = mx/2;

   if ((wrap == "wrap") and (padx != nx))
     {
	i = [0:dx-1]; cimage[nx+i] = cimage[i];
	i = [dx:mx-1]; cimage[padx-mx+i] = cimage[nx-mx+i];
     }
   if (wrap == "reflect")
     {
	i = [0:dx-1]; cimage[nx+i] = cimage[nx-i-2];
	i = [dx:mx-1]; cimage[padx-mx+i] = cimage[mx-i];
     }
   if (wrap == "nearest")
     {
	i = [0:dx-1]; cimage[nx+i] = cimage[nx-1];
	i = [dx:mx-1]; cimage[padx-mx+i] = cimage[0];
     }

   image = cimage;
   kernel = wrap_kernel_1d (kernel, padx, q.center);

   image = Real(fft(fft(__tmp(kernel), 1) * fft(__tmp(image), 1), -1));
   if (padx == nx)
     return image;

   return __tmp(image)[[:nx-1]];
}

private define process_qualifiers ()
{
   variable q = struct
     {
	pad = 0.0, wrap = NULL, center = 1,
     };

   if (qualifier_exists ("pad")) q.pad = qualifier ("pad");

   if (qualifier_exists ("reflect")) q.wrap = "reflect";
   if (qualifier_exists ("wrap")) q.wrap = "wrap";
   if (qualifier_exists ("nearest")) q.wrap = "nearest";

   q.center = qualifier ("center");
   if (q.center == NULL) q.center = 1;

   return q;
}

define convolve ()
{
   if (_NARGS != 2)
     {
	usage ("a = convolve (array, kernel);\n"
	       + "Qualifiers:\n"
	       + " pad=value   (default=0.0)\n"
	       + " wrap        (periodic BC)\n"
	       + " reflect     (reflect at boundaries)\n"
	       + " nearest     (Use nearest boundary value)\n"
	       + " center=0|1  (default=1)\n"
	      );
     }

   variable image, kernel;
   (image, kernel) = ();

   variable q = process_qualifiers (;;__qualifiers);

   variable ndims = length (array_shape (image));
   if (ndims == 1)
     return convolve1d (image, kernel, q);

   if (ndims == 2)
     return convolve2d (image, kernel, q);

   throw NotImplementedError,
     "This interface is currently limited to the convolution 1 and 2-d arrays";
}

define correlate ()
{
   if (_NARGS != 2)
     {
	usage ("a = correlate (array, kernel);\n"
	       + "Qualifiers:\n"
	       + " pad=value   (default=0.0)\n"
	       + " wrap        (periodic BC)\n"
	       + " reflect     (reflect at boundaries)\n"
	       + " nearest     (Use nearest boundary value)\n"
	       + " center=0|1  (default=1)\n"
	      );
     }

   variable image, kernel;
   (image, kernel) = ();

   variable q = process_qualifiers (;;__qualifiers);

   variable ndims = length (array_shape (image));
   if (ndims == 1)
     return convolve1d (image, __tmp(kernel)[[::-1]], q);

   if (ndims == 2)
     return convolve2d (image, __tmp(kernel)[[::-1],[::-1]], q);

   throw NotImplementedError,
     "This interface is currently limited to the correlation 1 and 2-d arrays";
}

provide ("gslfft");

