/*Wiener filter*/

#include <stdio.h>
#include <Imlib.h>
#include <dfftw.h>
#include <drfftw.h>
#include <math.h>
#include <stdlib.h>
#define TINY 1.0e-20

fftw_complex *out;

double *
genpsf (double *psf, int w, int h, ImlibImage * impsf, ImlibData * idpsf,
	double sigma)
{
  int i, j;
  double sum = 0.0;
  double max = 0.0;
  psf = malloc (sizeof (double) * w * h);
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  psf[i + w * j] =
	    (1.0 / (sqrt (2.0 * M_PI) * sigma))
	    	* exp (-0.5 * (pow (i - w / 2, 2) + pow (j - h / 2, 2)) / pow (sigma, 2));
	  sum += psf[i + w * j];
	  if (psf[i + w * j] > max)
	    {
	      max = psf[i + w * j];
	    }
	}
    }
  printf ("Som is %f \n", sum);
  printf ("Max is %f \n", max);
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  int index = (i + w * j) * 3;
	  psf[i + w * j] = psf[i + w * j] / sum;
	}
    }
  sum = 0;
  max = 0;
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  sum += psf[i + w * j];
	  if (psf[i + w * j] > max)
	    {
	      max = psf[i + w * j];
	    }
	}
    }
  printf ("Som is %f \n", sum);

  return psf;

}

/*extract from an RGB image the R, G of B channel*/
double *
kanaal (int w, int h, ImlibImage * im, int chan)
{
  int i, j;
  int index;
  double *out = malloc (sizeof (double) * w * h);

  for (i = 0; i < w; i++)
    {
      for (j = 0; j < h; j++)
	{
	  index = (i + w * j) * 3;
	  out[i + w * j] = im->rgb_data[index + chan];
	}
    }
  return out;
}

double *
regroup (int w, int h, fftw_complex * convred, double *flimoutn)
{
/*regroup the four quadrants after the DFT*/
  int i, j;

  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i ++)
	{
	  int k, l;
	  int index1, index2;
	  double temp;

	  index1 = i + w * j;
	  if (i < w / 2)
	    {
	      k = i + (w / 2);
	    }
	  if (j < h / 2)
	    {
	      l = j + (h / 2);
	    }
	  if (i >= w / 2)
	    {
	      k = i - (w / 2);
	    }
	  if (j >= h / 2)
	    {
	      l = j - (h / 2);
	    }
	  index2 = k + w * l;
	  flimoutn[index2] = sqrt (((convred[index1].re * convred[index1].re
			          +  convred[index1].im * convred[index1].im)));

	}
    }
  return flimoutn;
}

/*norm PSF to 1*/
double *
normpsf (double *psfch, int w, int h)
{
  int i, j;
  double sum;
  sum = 0.0;
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  sum += psfch[i + w * j];
	}
    }
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  psfch[i + w * j] = psfch[i + w * j] / sum;
	}
    }
  sum = 0.0;
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  sum += psfch[i + w * j];
	}
    }
  return psfch;
}

/*convolve two images*/
fftw_complex *
wiener (double *kanaal, double *kanpsf, int w, int h, fftw_complex * out,
	  double K)
{
  fftw_complex *im, *fim;
  fftw_complex *psfin, *fpsf, *conv;
  fftwnd_plan p, pinv;
  int i, j;
  /*input image*/
  im    = (fftw_complex *) malloc (w * h * sizeof (fftw_complex));
  /*fourier transf. of input image*/
  fim   = (fftw_complex *) malloc (w * h * sizeof (fftw_complex));
  psfin = (fftw_complex *) malloc (w * h * sizeof (fftw_complex));
  fpsf  = (fftw_complex *) malloc (w * h * sizeof (fftw_complex));
  conv  = (fftw_complex *) malloc (w * h * sizeof (fftw_complex));
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  im[i + j * w].re = kanaal[i + j * w];
	  im[i + j * w].im = 0;
	}
    }
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  psfin[i + j * w].re = kanpsf[i + j * w];
	  psfin[i + j * w].im = 0;
	}
    }
  p = fftw2d_create_plan (h, w, FFTW_FORWARD, FFTW_ESTIMATE);

  fftwnd_one (p, &im[0], &fim[0]);
  fftwnd_one (p, &psfin[0], &fpsf[0]);
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  conv[i + j * w].re =
	    (fim[i + j * w].re * fpsf[i + j * w].re +
	     fim[i + j * w].im * fpsf[i + j * w].im)
	     / (w * h * (K + pow (fpsf[i + j * w].re, 2) +
			    (pow (fpsf[i + j * w].im, 2))));
	  conv[i + j * w].im =
	    (fim[i + j * w].re * fpsf[i + j * w].im -
	     fim[i + j * w].im * fpsf[i + j * w].re)
	     / (w * h * (K + pow(fpsf[i + j * w].re, 2) +
			    (pow(fpsf[i + j * w].im, 2))));
	}
    }
  pinv = fftw2d_create_plan (h, w, FFTW_BACKWARD, FFTW_ESTIMATE);
  fftwnd_one (pinv, &conv[0], &out[0]);
  free (im);
  free (fim);
  free (psfin);
  free (fpsf);
  free (conv);
  return out;
}


int
main (int argc, char **argv)
{

  int teller, maxiter, chan;
  double *in, *est, *psf, *a, *b, *estblur;
  fftw_complex *cestblur, *cest;

  int w, h, i, j, iter;
  double total, K, sigma;

  ImlibData  *id, *idpsf, *idout;
  ImlibImage *im, *impsf, *imout;

  Display *disp;

  disp = XOpenDisplay (NULL);
  id = Imlib_init (disp);
  idpsf = Imlib_init (disp);
  idout = Imlib_init (disp);

  if (argc != 5)
    {
      printf ("Use input sigmaPSF out K\n");
      return 1;
    }

  im    = Imlib_load_image (id, argv[1]);
  impsf = Imlib_load_image (idpsf, argv[1]);
  imout = Imlib_load_image (idout, argv[1]);

  K     = atof (argv[4]);
  sigma = atof (argv[2]);
  /* Suck the image's original width and height out of the Image structure */
  w = im->rgb_width;
  h = im->rgb_height;
  out = (fftw_complex *) malloc (w * h * sizeof (fftw_complex));
  a = malloc (sizeof (double) * w * h);
  b = malloc (sizeof (double) * w * h);
  est = malloc (sizeof (double) * w * h);
  estblur = malloc (sizeof (double) * w * h);

  /*in */
  in = kanaal (w, h, im, 0);

  total = 0;
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  total += in[i + w * j];
	}
    }
  printf (" \t \tvoor g \t %f \n", total);

  /*generate gaussian psf with given sigma*/
  genpsf (psf, w, h, impsf, idpsf, sigma);

  /*norm PSF to 1 */
  normpsf (psf, w, h);


  total = 0.0;
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  total += psf[i + w * j];
	}
    }
  printf (" \t \tvoor psf \t %f \n", total);

  /*wienerfilter image given the PSF */
  wiener (in, psf, w, h, cestblur, K);
  regroup (w, h, cestblur, estblur);

  /*check: is total intensity of the filtered image the same? */
  total = 0.0;
  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  total += estblur[i + w * j];
	}
    }
  printf ("\t \t voor (f*H) \t %f \n", total);


/* OUTPUT */

  for (j = 0; j < h; j++)
    {
      for (i = 0; i < w; i++)
	{
	  int index = (i + j * w) * 3;
	  int index2 = i + w * j;
	  int index3 = (w-1-i + (h-1-j) * w) * 3;
	  if (estblur[index2] < 0)
	    {
	      estblur[index2] = 0;
	    }
	  if (estblur[index2] > 255)
	    {
	      estblur[index2] = 255;
	    }
	  imout->rgb_data[index3] = (int) estblur[index2];
	  imout->rgb_data[index3 + 1] = (int) estblur[index2];
	  imout->rgb_data[index3 + 2] = (int) estblur[index2];
	}
    }
  Imlib_save_image (idout, imout, argv[3], NULL);

  return 1;
}