/*----------------------------------------------------------------------------------------------------
  
   File        : inpainter.cpp

   Description : Implementation of the Tschumperle-Deriche's Multiscale
                 Inpainting PDE, for 2D multivalued images, as described in the following articles :
 
  (1) PDE-Based Regularization of Multivalued Images and Applications.
       (D. Tschumperlé). PhD Thesis. University of Nice-Sophia Antipolis, December 2002.
  (2) Diffusion PDE's on Vector-valued Images : Local Approach and Geometric Viewpoint.
       (D. Tschumperlé and R. Deriche). IEEE Signal Processing Magazine, October 2002.
  (3) Vector-Valued Image Regularization with PDE's : A Common Framework for Different Applications.
       (D. Tschumperlé and R. Deriche). CVPR'2003, Computer Vision and Pattern Recognition, Madison, United States, June 2003.  
 
  Author      : David Tschumperlé

  This software is governed by the CeCILL  license under French law and
  abiding by the rules of distribution of free software.  You can  use, 
  modify and/ or redistribute the software under the terms of the CeCILL
  license as circulated by CEA, CNRS and INRIA at the following URL
  "http://www.cecill.info". 
  
  As a counterpart to the access to the source code and  rights to copy,
  modify and redistribute granted by the license, users are provided only
  with a limited warranty  and the software's author,  the holder of the
  economic rights,  and the successive licensors  have only  limited
  liability. 
  
  In this respect, the user's attention is drawn to the risks associated
  with loading,  using,  modifying and/or developing or reproducing the
  software by the user in light of its specific status of free software,
  that may mean  that it is complicated to manipulate,  and  that  also
  therefore means  that it is reserved for developers  and  experienced
  professionals having in-depth computer knowledge. Users are therefore
  encouraged to load and test the software's suitability as regards their
  requirements in conditions enabling the security of their systems and/or 
  data to be ensured and,  more generally, to use and operate it in the 
  same conditions as regards security. 
  
  The fact that you are presently reading this means that you have had
  knowledge of the CeCILL license and that you accept its terms.
  
 ---------------------------------------------------------------------------------------------------*/

#include "../CImg.h"
// The lines below are not necessary in your own code, it simply allows 
// the source compilation with compilers that do not respect the C++ standart.
#if ( defined(_MSC_VER) && _MSC_VER<=1200 ) || defined(__DMC__)
#define std
#endif
using namespace cimg_library;

float dt,alpha,sigma,beta,a1,a2;
int stflag,nb_iter,visu,nb_save=0;
CImg<> gimg;
CImgDisplay *disp;
const char *file_o;

// Procedure that fills a part of the mask (for interactive mode)
//-----------------------------------------
void mask_fill(CImg<unsigned char>& img,const unsigned int x,const unsigned int y) {
  int cx,xmin,xmax;
  for (xmin=x; xmin>=0        && !img(xmin,y,0); xmin--) ; xmin++;
  for (xmax=x; xmax<img.dimx() && !img(xmax,y,0); xmax++) ; xmax--;
  if (xmin<=xmax) memset(img.ptr(xmin,y,0,0),255,xmax-xmin+1);
  for (cx=xmin; cx<=xmax; cx++) {
    if (y>0   && img(cx,y-1,0)==0) mask_fill(img,cx,y-1);
    if (y<img.height-1 && img(cx,y+1,0)==0) mask_fill(img,cx,y+1);
  }
}


// Inpainting for a given scale
//------------------------------
CImg<> inpaint2D(CImg<>& img,CImg<unsigned char>& mask)
{
  CImg<> Ix,Iy, veloc(img,false), G(img.width,img.height,1,3), T(G,false);
  float ix,iy,ixx,ixy,iyy,xdt = 0;
  int maskflag=0;
  CImg_3x3(I,float);
  CImgStats stats;

  // PDE iteration loop
  for (int iter=0; iter<nb_iter; iter++) {
    std::printf("\riter %d : xdt = %g                       ",iter,xdt); fflush(stdout);
    // compute the structure tensor and the hessians
    Ix = img.get_deriche((float)alpha,1,'x',0);
    Iy = img.get_deriche((float)alpha,1,'y',0);
    G.fill(0);
    cimg_mapXYV(img,x,y,k) {
      ix = Ix(x,y,k);
      iy = Iy(x,y,k);
      G(x,y,0) += ix*ix;
      G(x,y,1) += ix*iy;
      G(x,y,2) += iy*iy;
    }
    G.blur((float)sigma);
    // compute the 2x2 tensor field T
    cimg_mapXY(G,x,y) if (mask(x,y)!=0) {
      const CImgl<> eigen = G.get_tensor(x,y).get_symeigen();
      const CImg<>& mV = eigen(1);
      T(x,y,0) = mV(0)*mV(0);
      T(x,y,1) = mV(0)*mV(1);
      T(x,y,2) = mV(1)*mV(1);
      maskflag=1;
    }
    if (!maskflag) { std::fprintf(stderr,"No mask data in this scale !\n"); iter=nb_iter; }
    else {
      // Compute the PDE velocity
      veloc.fill(0);
      cimg_mapV(img,k) cimg_map3x3(img,x,y,0,k,I) if (mask(x,y)!=0) {
        ixx = Inc+Ipc-2*Icc;
        iyy = Icn+Icp-2*Icc;
        ixy = 0.25f*(Ipp+Inn-Ipn-Inp);
        veloc(x,y,k) = T(x,y,0)*ixx + 2*T(x,y,1)*ixy + T(x,y,2)*iyy;
      }
      // update the image
      if (dt>0) { stats = veloc.get_stats(0); xdt = (float)(dt/(0.001+cimg::max(std::fabs(stats.min),std::fabs(stats.max)))); }
      else xdt = -dt;
      img+=veloc*xdt;
      if (disp) disp->resize(*disp);
      if (disp && !(iter%visu)) img.display(*disp);
      //if (file_o && !(iter%visu)) img.get_resize(gimg.width,gimg.height).save(file_o,nb_save++);
      if (stflag && !(iter%stflag)) { veloc.print("Velocity"); img.print("Image"); }
    }
  }
  return img;
}


// Main function
//---------------
int main(int argc,char **argv)
{
  cimg_usage("Multiscale Inpainting, for 2D vector-valued images");
  const char *file_i = cimg_option("-i","img/milla.ppm","Input image"),
    *file_m    = cimg_option("-m",(char*)NULL,"Mask image"),
    *file_om   = cimg_option("-om",(char*)NULL,"Output mask file");
  file_o       = cimg_option("-o",(char*)NULL,"Output file");
  dt           = (float)cimg_option("-dt",1.5,"Adapting time step");
  float galpha = (float)cimg_option("-alpha",1.2,"Image smoothing"),
    gsigma     = (float)cimg_option("-sigma",4.0,"Structure tensor smoothing"),
    dist       = (float)cimg_option("-dist",2.5,"Multiplicator of scale areas");
  int init     = cimg_option("-init",1,"Mask initialization (0=black, 1=white, 2=noise, 3=unchanged)"),
    nb_scale   = cimg_option("-scale",4,"Number of different scales"),
    gnb_iter     = cimg_option("-iter",550,"Number of maximum iteration");
  stflag       = cimg_option("-stats",0,"Display statistics");
  visu         = cimg_option("-visu",10,"Step for visualization & saving");

  int xo=-1,yo=-1,x=-1,y=-1,dispflag=0;
  float tmp,wn=0,hn=0;
  unsigned char color = 255;
  if (!file_i) throw CImgException("You must specify an input image");
  if (nb_scale<=0) throw CImgException("At least one scale must be specified");
  gimg = CImg<>(file_i);
  CImg<> img,imgvisu(gimg);
  CImg<unsigned char> gmask, mask;
  if (gimg.depth>1) gimg.resize(-100,-100,1,-100);
  if (visu) disp = gimg.new_display("Image",1,3); else disp = NULL;

  // Mask defined by user (if necessary)
  if (!file_m) {
    if (!disp) throw CImgException("You must turn display on (flag -visu) to define an inpainting mask");
    gmask = CImg<unsigned char>(gimg.width,gimg.height).fill(0);
    while (!disp->key) {
      if (xo<0 || yo<0) {
        xo=disp->mousex; yo=disp->mousey; dispflag=1;
      }
      else {
        xo=x; yo=y;
      }
      x = disp->mousex;
      y = disp->mousey;
      if (disp->button&1 && (xo!=x || yo!=y)) { gmask.draw_line(xo,yo,x,y,&color); dispflag=1; }
      if (disp->button&2) { mask_fill(gmask,x,y); dispflag=1;}
      if (disp->button&4) { gmask.draw_rectangle(x-4,y-4,0,0,x+4,y+4,0,0,255); dispflag=1; }
      if (dispflag) { cimg_mapXYV(imgvisu,x,y,k) imgvisu(x,y,k)=gmask(x,y)?255:gimg(x,y,k); imgvisu.display(*disp);  dispflag=0; }
      disp->resize(*disp).wait();
    }
  } else gmask = CImg<unsigned char>(file_m).resize(gimg.width,gimg.height,1);
  switch(init) {
  case 0: { cimg_mapXYV(gimg,x,y,k) if (gmask(x,y)!=0) gimg(x,y,k)=0; } break;
  case 1: tmp = (float)gimg.get_stats(0).max; { cimg_mapXYV(gimg,x,y,k) if (gmask(x,y)!=0) gimg(x,y,k)=tmp; } break;
  case 2: tmp = (float)gimg.get_stats(0).max; { cimg_mapXYV(gimg,x,y,k) if (gmask(x,y)!=0) gimg(x,y,k)=(float)cimg::rand()*tmp; } break;
  default: break;
  }
  if (file_om) gmask.save(file_om);

  // Start the multiscale algorithm
  for (int n=1; n<=nb_scale; n++) {
    std::printf("\n** Computing scale %d/%d\n   ----------------------\n",n,nb_scale);
    tmp     = (float)std::pow((double)dist,0.5*(n-nb_scale));
    const float tmp2    = (float)n;
    wn      = tmp*gimg.width;
    hn      = tmp*gimg.height;
    if (wn<5) wn=5;
    if (hn<5) hn=5;
    alpha   = galpha/tmp2;
    sigma   = gsigma/tmp2;
    nb_iter = (int)(gnb_iter/tmp2);
    std::printf("wn = %d, hn = %d, alpha = %g, sigma = %g, nb_iter = %d\n",(int)wn,(int)hn,alpha,sigma,nb_iter);
    if (n==1) {
      img   = gimg.get_resize((int)wn,(int)hn,1,gimg.dim,3);
      mask  = gmask.get_resize(img.width,img.height,1,1,3);
    }
    else {
      img.resize((int)wn,(int)hn,1,gimg.dim,3);
      mask  = gmask.get_resize(img.width,img.height,1,1,3);
      cimg_mapXYV(img,x,y,k) if (mask(x,y)==0) img(x,y,k)=(float)gimg.linear_pix2d((float)(x*gimg.width)/img.width,(float)(y*gimg.height)/img.height,0,k);
    }
    img.print("Scaled image");
    mask.print("Scaled mask");
    img = inpaint2D(img,mask);
  }
  if (file_o) img.save(file_o); else img.display("Inpainting result");
  exit(0);
  return 0;
}
