/*rlmonoc3*/

 #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;

/*extract uit een RGB beeld kanaal R, G of B*/
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) {
/*herschik beeld weer zodat de vier kwadranten weer op hun plaats staan*/
	int i,j;

        i=0;
	while (i < w) {
     	     j = 0;
             while (j < h) {
	        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)));

       		j ++;
             }
             i ++;
       }
	return flimoutn;
}

/*normeer PSF op 1*/
double *normpsf(double *psfch, int w, int h) {
	int i,j;
	double sum;
	sum = 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;
	for (j=0; j<h; j++) {
	     for (i=0; i<w; i++) { 	
		sum += psfch[i+w*j];             	
	     }
        }
	return psfch;
}

double * transp(double *arr, int w, int h) {
	int i,j, t;
	double *arrtransp;
	arrtransp = malloc(sizeof(double)*w*h);
	t = 0;
	for (i=0; i<w; i++) {
	    for (j=0; j<h; j++) {
		int index = i+w*j;
		arrtransp[t] = arr[index];
		t ++;
            }
        }
	return arrtransp;
}
	
	
/*convolueer twee beelden*/
fftw_complex * convolve(double *kanaal, double *kanpsf, int w, int h, fftw_complex *out) {
	fftw_complex *im, *fim;
        fftw_complex *psfin, *fpsf, *conv;
	fftwnd_plan p, pinv;
	int i,j;
        /*ingelezen beeld*/
	im    = (fftw_complex *) malloc(w*h*sizeof(fftw_complex));
	/*fouriergetransf. van ingelezen beeld*/
	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 (i=0; i < w; i ++) {
	    for (j=0; j < h; j++) {
		im[i+j*w].re = kanaal[i+j*w];
		im[i+j*w].im = 0;
	    }
	}
        for (i=0; i < w; i ++) {
	    for (j=0; j < h; j++) {
		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 (i=0; i < w; i ++) {
	    for (j=0; j < h; j++) {
		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);
		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);
	    }
	}
        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;
        int chan;
	/*drie kanalen van het ingelezen beeld*/
	double *in;
	/*drie kanalen van fi*/
        double *est;
	/*de ingelezen psf*/
	double *psftr;
	/*de getransponeerde van de psf*/
        double *psf;

        double *a;
	double *b;
        fftw_complex *cestblur;
	double *estblur;
        fftw_complex *cest;

        int w,h,i,j, iter;
	double total = 0;

	ImlibData *id;
        ImlibData *idpsf;
        ImlibData *idout;

     	ImlibImage *im;
     	ImlibImage *impsf;
        ImlibImage *imout;

	Display *disp;

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

 	/*laad beeld in gespecifieerd als eerste argument*/
    	im =Imlib_load_image(id,argv[1]);

	/*laad PSF in gespecifieerd in tweede argument*/
    	impsf=Imlib_load_image(idpsf,argv[2]);

	/* maak beeld klaar om het resultaat naar te schrijven   */
        imout=Imlib_load_image(idout,argv[1]);

	maxiter = atoi(argv[4]);

        if (argc<4) {
		printf("Usage: <input> <psf> <output> <max nr of iterations>\n");
		return 0;
	}
	
 	/* 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);		

        psf = kanaal(w, h, impsf, 0);

	
	/*normeer PSF op 1*/
	normpsf(psf, w, h);

	/*transponeer PSF voor een volgende deel van het algoritme*/
        psftr = transp(psf, w, h);


/*INITIALISATIE*/

	/*est = in*/

        est  = malloc(sizeof(double)*w*h);

	for (j=0; j<h; j++) {
	  for (i=0; i<w; i++) {
	     int index = i+w*j;
	     est[index] = in[index];
	  }
	}

	total = 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);	

        /*convolueer beeld met PSF*/
	cestblur = convolve(in, psf, w, h, out);
	regroup(w, h, cestblur, estblur);

        /*check: is de totale intensiteit van het geconvolueerde beeld behouden?*/
	total = 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);


/*ITERATIES*/
        for (iter=0; iter<maxiter; iter++) {
          /*deel beeld g puntsgewijs door net bekomen convolutie*/
	  for (j=0; j<h; j++) {
             for (i=0; i<w; i++) {
		int index = (i+j*w);
		double check1, check2, check3;
		/*if (estblur[index] < TINY) {a[index] =0;} else {*/
		a[index] = in[index] / ( estblur[index]);
		check1 = a[index];
		check2 = estblur[index];
             }
	  }
          total = 0;
	  for (j=0; j<h; j++) {
	    for (i=0; i<w; i++) {
		total += a[i+w*j];
	    }
	  }
	  printf("\t \t voor g/(f*H) \t %f \n", total);

          cestblur = convolve(a, psftr, w, h, out);
          regroup(w, h, cestblur, b);

	  total =0;
	  for (j=0; j<h; j++) {
	    for (i=0; i<w; i++) {
		int index = i+j*w;
		total += b[index];
	    }
	  }
	  printf(" \t \t voor H^T (g/(H*g)) %f \n", total);

	  /*est = est x B */		

          for (j=0; j<h; j++) {
             for (i=0; i<w; i++) {
		int index = i+j*w;
		est[index] = est[index] * b[index];
             }
	  }
	  printf("%d \n",iter);

          total =0;
	  for (j=0; j<h; j++) {
	    for (i=0; i<w; i++) {
		int index = i+j*w;
		total += est[index];
	    }
	  }
	  printf("\t \t voor resultaat %f \n", total);

          /*ter initialisatie van de volgende iteratie: convolueer beeld met PSF*/
	  cestblur   = convolve(est, psf, w, h, out);

	  regroup(w, h, cestblur , estblur);

	  /*total = 0;
	  for (j=0; j<h; j++) {
	     for (i=0; i<w; i++) {
		int index = i + j*w;
		total += estblur[index];
             }
	  }*/
        }
/* EINDE ITERATIES : UITVOER */

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

	return 1;
}