/*------------------------------------------------------------------------------------------------------------------------------
  
 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 program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; either version 2 of the License, or
 (at your option) any later version.
 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.
 
 ------------------------------------------------------------------------------------------------------------------------------*/

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