/* The L2 regularization model for image denoising */
load "ppm2rnm"
load "medit"

/* Convert input image into pgm format */
string sinput = "image.jpg";
string spgm   = "image.pgm";

exec("convert "+sinput+" "+spgm);

/* Read image and get data */
real[int,int] img(spgm);
int nx = img.n;
int ny = img.m;

/* Creation of a square mesh adapted to the pixellisation of the image */
mesh Thi = square(nx-1,ny-1,[(nx-1)*(x),(ny-1)*(1-y)]);
fespace Vhi(Thi,P1);
Vhi uTi;
uTi[] = img;

/* Rescaling of mesh between 0 and 1 */
real[int] bb(4);
boundingbox(Thi,bb); // bb[0] = xmin, bb[1] = xmax, bb[2] = ymin, bb[3] = ymax
real dd = max(bb[1]-bb[0],bb[3]-bb[2]);
mesh Th = movemesh(Thi,[(x-bb[0])/dd,(y-bb[2])/dd]);

/* Mesh adaptation to get a reasonable number of elements */
Th = adaptmesh(Th,hmin=0.005,hmax=0.005,nbvx=50000);
fespace Vh(Th,P1);
Vh uT = uTi(dd*x,dd*y); // Transfer of the image from Thi to TH
Vh u,newu,v,dxu1,dyu1,g1,norm;

real EPS  = 1.e-15; // To avoid division by 0
real lm   = 0.5;    // Weight of the L1 regularization term
real obj,newobj,coef,gmax,step;
int MAXIT    = 100;   // Number of optimization iterations
real MAXCOEF = 0.02;  // Maximum descent coefficient
real MINCOEF = 5e-6;  // Minimum descent coefficient
real alpha   = 0.001; // Regularization length-scale for the gradient
string sout;

/* Save mesh and image */
savemesh(Th,"sq.mesh");
savesol("sq.sol",Th,uT);

/* Macro for calculating the energy functional */
macro J(uu) ( 0.5*int2d(Th)((uu-uT)^2) + lm*int2d(Th)(dx(uu)^2+dy(uu)^2) ) // EOM

/* Macro for calculating the gradient of J as a P1 function (not descent direction) */
macro dJ(uu) {
  /* Slight regularization of the gradient */
  solve gradJ(g1,v) = int2d(Th)( alpha^2*(dx(g1)*dx(v)+dy(g1)*dy(v)) + g1*v )
                      - int2d(Th)( (uu-uT)*v + lm*(dx(uu)*dx(v)+dy(uu)*dy(v)) );
} // EOM

/* Macro for thresholding image between 0 and 1 */
macro thres(uu) {
  uu = min(max(uu,0.0),1.0);
} //EOM

/* Initialization */
u = uT;
obj = J(u);
coef = 0.01; // Descent step is tuned depending on the size of the initial gradient

/* Main loop */
for (int n=1; n<=MAXIT; n++) {
  /* Calculation of gradient */
  dJ(u);
  
  /* Update of the optimized function */
  gmax = max(-g1[].min,g1[].max);
  newu = u - coef/gmax*g1;
  
  /* Thresholding between 0 and 1 */
  thres(newu);
  
  /* Evaluation of the new objective */
  newobj = J(newu);
  
  /* Decision accept if objective has decreased */
  if ( newobj < obj ) {
    cout<<"Iteration "<<n<<" accepted "<<obj<<" ---> "<<newobj<<"."<<endl;
    u    = newu;
    obj  = newobj;
    coef = min(MAXCOEF,1.1*coef); // Slight increase of the descent coefficient
  }
  /* Else, reject and slightly decrease the descent coefficient */
  else {
    cout<<"Iteration "<<n<<" rejected."<<endl;
    coef = max(MINCOEF,0.5*coef);
  }
  
  /* Save data */
  sout = "step."+n+".sol";
  savesol(sout,Th,u);
}