/*----------------------------MegaWave2 module----------------------------*/
/* mwcommand
name = {tv_poisson};
version = {"1.0"};
author = {"Triet Minh Le"};
function = {"Denoising Images Corrupted by Poisson Noise"};
usage = {
 'w':[w=0.1]->w    "weight on fidelity term (default 0.1)",
 's':[s=0.1]->s     "initial (and maximal) time step, default 1.",
 'E':[eps=0.1]->eps "epsilon in sqrt(epsilon+|Du|^2), default 1.",
 'n':n->n          "or perform a fixed number of iterations (default: 5)",
 'r':ref->ref      "to specify a reference image different from in",
 in->in            "input Fimage",
 out<-out          "output Fimage"
        };
*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "mw.h"

extern float fnorm();

#define COEFF .70710678118654752440
#define COEFF2P2 3.41421356237309504880

double mytvgrad(u,grad,eps)
Fimage u,grad;
double eps;
{
  int x,y,nx,ny,adr;
  double E,a,b,c,d,e,f,z,norm;

  nx = u->ncol; 
  ny = u->nrow;
  E = 0.;
  if (grad) {
    mw_change_fimage(grad,ny,nx);
    mw_clear_fimage(grad,0.);
  }

  for (x=0;x<nx-1;x++)
    for(y=0;y<ny-1;y++) {
      adr = y*nx+x;
      a = (double)u->gray[adr     ];
      b = (double)u->gray[adr   +1];
      c = (double)u->gray[adr+nx+1];
      d = (double)u->gray[adr+nx  ];
      z=d; d-=c; c-=b; b-=a; a-=z; e=COEFF*(b-d); f=COEFF*(c-a);
      norm = sqrt( eps + (a*a+b*b+c*c+d*d + e*e+f*f)/COEFF2P2 );
      E += norm;
      norm *= COEFF2P2;
      if (grad) {
	if (norm==0.) norm=1.;
	grad->gray[adr     ] += (a-b+COEFF*(-e-f))/norm;
	grad->gray[adr   +1] += (b-c+COEFF*( e-f))/norm;
	grad->gray[adr+nx+1] += (c-d+COEFF*( e+f))/norm;
	grad->gray[adr+nx  ] += (d-a+COEFF*(-e+f))/norm;
      }
    }

  return(E);
}


/* Compute energy F(u) = weight * int |u-u_0|^2 and its gradient */
double fidelity_term_grad(u,u0,grad,weight, eps)
     Fimage u,u0,grad;
     double weight, eps;
{
  int x,y,nx,ny,adr;
  double e,d,f;

  nx = u->ncol; 
  ny = u->nrow;
  e = 0.;

  for (x=0;x<nx;x++)
    for(y=0;y<ny;y++) {
      adr = y*nx+x;
      f = (double)u->gray[adr];
      d = ((double)u->gray[adr] - (double)u0->gray[adr])/sqrt(eps + f*f);
      if (grad) grad->gray[adr] += (float)(weight*d);
      e += weight*(f - (double)u0->gray[adr]*log(f));
    }

  return(e);
}


/*----------------------- MAIN MODULE ---------------*/

Fimage tv_poisson(in,out,s,n,w,ref,eps)
Fimage in,out,ref;
int *n;
double *s,*w,*eps;
{
  Fimage dE,aux,tmp,cur,prev;
  double energy,old_energy,step;
  float two,norm1,norm2;
  int i,j,border,nx,ny;

  /* initialization */
  nx = in->ncol; 
  ny = in->nrow;

  dE = mw_change_fimage(NULL,ny,nx);
  out = mw_change_fimage(out,ny,nx);
  aux = mw_change_fimage(NULL,ny,nx);
  if (!dE || !out || !aux) 
    mwerror(FATAL,1,"inttv: not enough memory\n");

  mw_copy_fimage(in,out);
  mw_copy_fimage(in,aux);
  if (!ref) ref = in;
  prev = aux; cur = out; 
  step = *s;

  /* compute initial error and its gradient */
  energy = mytvgrad(cur,dE,*eps);
  if (*w!=0.) energy += fidelity_term_grad(cur,ref,dE,*w, *eps);
  energy /= (double)(nx*ny);

  /***** MAIN LOOP *****/ 
  for(i=0; i<*n; i++) {

    /* update image */
    for (j=nx*ny;j--;) 
      prev->gray[j] = cur->gray[j] - (float)step * dE->gray[j];
    
    /* flip pointers and increment */
    tmp=prev; prev=cur; cur=tmp;
    i++;
    
    /* compute error and its gradient for the next iteration */
    old_energy = energy;
    energy = mytvgrad(cur,dE,*eps);
    if (*w!=0.) energy += fidelity_term_grad(cur,ref,dE,*w, *eps);
    energy /= (double)(nx*ny);
    
    if (energy>old_energy) {
      tmp=prev; prev=cur; cur=tmp;
      energy = mytvgrad(cur,dE,*eps);
      if (*w!=0.) energy += fidelity_term_grad(cur,ref,dE,*w, *eps);
      energy /= (double)(nx*ny);
      step *= 0.8;
	i--;
    } 
  } 
  /***** END OF MAIN LOOP *****/

  if (cur!=out) mw_copy_fimage(cur,out);
  mw_delete_fimage(aux);
  mw_delete_fimage(dE);
  return(out);
}