/*-----------------------------------------------------------------------------

    File        : greycstoration.cpp

    Description : Fast Anisotropic Smoothing of Multi-valued Images,
                  using Curvature-Preserving PDE's

                  This algorithm is described in the paper :

                  "Fast Anisotropic Smoothing of Multi-Valued Images 
                   using Curvature-Preserving PDE's"
		          (D. Tschumperlé, Research Report "Cahiers du GREYC", February 2005),
		          GREYC, CNRS UMR 6072.

		  See also the corresponding web page :
		  http://www.greyc.ensicaen.fr/~dtschump/greycstoration

    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"
using namespace cimg_library;
typedef unsigned char uchar;
// The line below is necessary for compilation with non-standart C++ compilers.
#if (( defined(_MSC_VER) && _MSC_VER<=1200 ) || defined(__DMC__))
#define std
#endif

// get_geom() : read geometry from a string ('320x256' or '200%x200%' are valid geometries).
//------------
void get_geom(const char *geom, int &geom_w, int &geom_h) {
  char tmp[16];
  std::sscanf(geom,"%d%7[^0-9]%d%7[^0-9]",&geom_w,tmp,&geom_h,tmp+1);
  if (tmp[0]=='%') geom_w=-geom_w;
  if (tmp[1]=='%') geom_h=-geom_h;
}

/*-------------------
    Main procedure
    ----------------*/

int main(int argc,char **argv) {
  cimg_usage("Fast Anisotropic Smoothing of Multi-valued Images with Curvature-Preserving PDE's");
  std::fprintf(stderr,
	       "---------------------------------------------------------------------------------------------\n"
	       " Note on using GREYCstoration :\n"
	       " ------------------------------\n"
	       " The GREYCstoration algorithm is the result of a research work done in the\n"
	       " IMAGE group of the GREYC Lab (CNRS, UMR 6072)/France, by David Tschumperle\n"
	       " ( http://www.greyc.ensicaen.fr/EquipeImage/ and http://www.greyc.ensicaen.fr/~dtschump/ ).\n"
	       " It can be used to regularize, inpaint, resize images, or create texture versions of\n"
	       " 2D vector flows. This algorithm is distributed to help people for processing image data.\n"
	       " This is an open source software, distributed within the CImg Library package\n"
	       " (http://cimg.sourceforge.net), and submitted to the CeCiLL License (view file LICENSE.txt).\n"
	       " If you are interested to distribute this algorithm in a closed-source product,\n"
	       " you are invited to contact David Tschumperle (mail available on his web page).\n"
	       " Please also cite the related paper when using results of the GREYCstoration\n"
	       " algorithm in your own publications :\n\n"
	       "    \"Fast Anisotropic Smoothing of Multi-Valued Images using Curvature-Preserving PDE's\"\n"
	       "    (D. Tschumperle)\n"
	       "    GREYC, CNRS UMR 6072, Equipe Image, 6 Bd du Marechal Juin, 14050 Caen Cedex.\n\n"
	       " You can also take a look at the GREYCstoration web page :\n"
	       "    http://www.greyc.ensicaen.fr/~dtschump/greycstoration\n"
	       "---------------------------------------------------------------------------------------------\n");
  
  const char *restore  = cimg_option("-restore",(const char*)NULL,"Restore the image specified after '-restore'");
  const char *inpaint  = cimg_option("-inpaint",(const char*)NULL,"Inpaint the image specified after '-inpaint'");
  const char *resize   = cimg_option("-resize",(const char*)NULL,"Resize the image specified after '-resize'");
  const char *visuflow = cimg_option("-flow",(const char*)NULL,"Render flow specified after '-flow'");
  if (argc==1) {
    std::fprintf(stderr,"\nTry '%s -h' for the list of available options\n\n",cimg::basename(argv[0]));
    exit(0);
  }
  if (!restore && !inpaint && !resize && !visuflow) {
    std::fprintf(stderr,"\nYou must specify one of the '-restore','-inpaint','-resize' or '-flow' flags (try option '-h')\n\n");
    exit(0);
  }

  // Init algorithm parameters
  //---------------------------
  CImg<> img0, img, flow,G;
  CImgl<> eigen(CImg<>(2,1), CImg<>(2,2));
  CImg<uchar> mask;
  unsigned int nb_iter=0;
  float dt=0,dlength=0,dtheta=0,sigma=0,power1=0,power2=0;
  bool onormalize=false;
  const float gauss_prec   = (float)cimg_option("-p",3.0f,"Precision of the gaussian function");
  const unsigned int save  = cimg_option("-save",0,"Iteration saving step");
  const bool linear        = cimg_option("-linear",false,"Use linear interpolation for integration");
  const bool visu          = cimg_option("-visu",true,"Enable/Disable visualization");
  const char *file_o       = cimg_option("-o",(const char*)NULL,"Output image");
  const bool append_result = cimg_option("-a",false,NULL);

  // Specific parameters for image restoration
  if (restore) { 
    nb_iter        = cimg_option("-iter",1,"Number of smoothing iterations");
    dt             = (float)cimg_option("-dt",80.0f,"Time step");
    sigma          = (float)cimg_option("-sigma",1.2f,"Structure tensor blurring");
    dlength        = (float)cimg_option("-dl",1.5,"Integration step");
    dtheta         = (float)cimg_option("-da",60.0,"Angular step (in degrees)");
    onormalize     = cimg_option("-normalize",false,"Output image normalization (in [0,255])");
    power1         = (float)cimg_option("-p1",0.5,"Diffusion limiter along isophote");
    power2         = (float)cimg_option("-p2",0.9,"Diffusion limiter along gradient");
    const float ng = (float)cimg_option("-ng",0.0,"Add gaussian noise before denoising");
    const float nu = (float)cimg_option("-nu",0.0,"Add uniform noise before denoising");
    const float ns = (float)cimg_option("-ns",0.0,"Add Salt&Pepper noise before denoising");
    img = CImg<>(restore);
    CImgStats stats(img,false);
    if (ng) img.noise(ng,0);
    if (nu) img.noise(nu,1);
    if (ns) img.noise(ns,2);   
    img.normalize((float)stats.min,(float)stats.max);
    img0=img;
    G = CImg<>(img.dimx(),img.dimy(),1,3);
  }

  // Specific parameters for image inpainting
  if (inpaint) {
    const char *file_m         = cimg_option("-m",(const char*)NULL,"Input inpainting mask");
    const unsigned int dilate  = cimg_option("-dilate",0,"Inpainting mask dilatation");
    const unsigned int ip_init = cimg_option("-init",3,"Inpainting init (0=black, 1=white, 2=noise, 3=unchanged, 4=interpol)");
    nb_iter                    = cimg_option("-iter",100,"Number of inpainting iterations");
    dt                         = (float)cimg_option("-dt",50.0f,"Time step");
    sigma                      = (float)cimg_option("-sigma",2.0,"Structure tensor blurring");
    power1                     = (float)cimg_option("-p1",0.1,"Diffusion limiter along isophote");
    power2                     = (float)cimg_option("-p2",100,"Diffusion limiter along gradient");
    dlength                    = (float)cimg_option("-dl",0.8,"Integration step");
    dtheta                     = (float)cimg_option("-da",60.0,"Angular step (in degrees)");
    img = CImg<>(inpaint);
    if (!file_m) throw CImgArgumentException("You need to specify an inpainting mask (option '-m')");    
    if (cimg::strncasecmp("block",file_m,5)) mask = CImg<uchar>(file_m);
    else {
      int l=16; std::sscanf(file_m,"block%d",&l);
      mask = CImg<uchar>(img.dimx()/l,img.dimy()/l);
      cimg_mapXY(mask,x,y) mask(x,y)=(x+y)%2;
    }
    mask.resize(img.dimx(),img.dimy(),1,1);
    if (dilate) mask.dilate(dilate);
    switch (ip_init) {
    case 0 : { cimg_mapXYV(img,x,y,k) if (mask(x,y)) img(x,y,k) = 0; } break;
    case 1 : { cimg_mapXYV(img,x,y,k) if (mask(x,y)) img(x,y,k) = 255; } break;
    case 2 : { cimg_mapXYV(img,x,y,k) if (mask(x,y)) img(x,y,k) = (float)(255*cimg::rand()); } break;
    case 3 : break;
    case 4 : {
      CImg<uchar> tmask(mask),ntmask(tmask);
      CImg_3x3(M,uchar);
      CImg_3x3(I,float);
      while (CImgStats(ntmask,false).max>0) {
	cimg_map3x3(tmask,x,y,0,0,M) if (Mcc && (!Mpc || !Mnc || !Mcp || !Mcn)) {
	  const float ccp = Mcp?0.0f:1.0f, cpc = Mpc?0.0f:1.0f,
	    cnc = Mnc?0.0f:1.0f, ccn = Mcn?0.0f:1.0f, csum = ccp + cpc + cnc + ccn;
	  cimg_mapV(img,k) {
	    cimg_get3x3(img,x,y,0,k,I);
	    img(x,y,k) = (ccp*Icp + cpc*Ipc + cnc*Inc + ccn*Icn)/csum;
	  }
	  ntmask(x,y) = 0;
	}
	tmask = ntmask;
      }
    } break;    
    default: break;
    }
    img0=img;
    G = CImg<>(img.dimx(),img.dimy(),1,3,0);
    CImg_3x3(g,uchar);
    CImg_3x3(I,float);
    cimg_map3x3(mask,x,y,0,0,g) if (!gcc && !(gnc-gcc) && !(gcc-gpc) && !(gcn-gcc) && !(gcc-gcp)) cimg_mapV(img,k) {
      cimg_get3x3(img,x,y,0,k,I);
      const float ix = 0.5f*(Inc-Ipc), iy = 0.5f*(Icn-Icp);
      G(x,y,0)+= ix*ix; G(x,y,1)+= ix*iy; G(x,y,2)+= iy*iy;    
    }
    G.blur(sigma);
    { cimg_mapXY(G,x,y) {
       G.get_tensor(x,y).symeigen(eigen(0),eigen(1));
       const float
	 l1 = eigen(0)[0],
	 l2 = eigen(0)[1],
	 u = eigen(1)[0],
	 v = eigen(1)[1],      
	 ng = (float)std::sqrt(l1+l2),
	 n1 = (float)(1.0/std::pow(1+ng,power1)),
	 n2 = (float)(1.0/std::pow(1+ng,power2)),
	 sr1 = (float)std::sqrt(n1),
	 sr2 = (float)std::sqrt(n2);
       G(x,y,0) = sr1*u*u + sr2*v*v;
       G(x,y,1) = u*v*(sr1-sr2);
       G(x,y,2) = sr1*v*v + sr2*u*u;
	}    
  }}
  
  // Specific parameters for image resizing
  if (resize) {
    const char *geom  = cimg_option("-g",(const char*)NULL,"Output image geometry");
    const bool anchor = cimg_option("-anchor",true,"Anchor original pixels");
    nb_iter           = cimg_option("-iter",1,"Number of iterations");
    dt                = (float)cimg_option("-dt",30.0f,"Time step");
    sigma             = (float)cimg_option("-sigma",2.0,"Structure tensor blurring");
    dlength           = (float)cimg_option("-dl",0.8,"Integration step");
    dtheta            = (float)cimg_option("-da",45.0,"Angular step (in degrees)");
    power1            = (float)cimg_option("-p1",0.01,"Diffusion limiter along isophote");
    power2            = (float)cimg_option("-p2",100.0,"Diffusion limiter along gradient");
    if (!geom) throw CImgArgumentException("You need to specify an output geomety (option -g)");
    int w,h; get_geom(geom,w,h);
    img = CImg<>(resize);
    mask = CImg<uchar>(img.dimx(),img.dimy(),1,1,255);
    if (!anchor) mask.resize(w,h,1,1,1); else mask = ~mask.resize(w,h,1,1,4);
    img0 = img.get_resize(w,h,1,-100,1);
    img.resize(w,h,1,-100,3);
    G = CImg<>(img.dimx(),img.dimy(),1,3);
  }

  // Specific parameters for flow visualization
  if (visuflow) {
    const char *geom     = cimg_option("-g","100%x100%","Output geometry");
    const char *file_i   = cimg_option("-i",(const char*)NULL,"Input init image");
    const bool normalize = cimg_option("-norm",false,"Normalize input flow");
    nb_iter              = cimg_option("-iter",1,"Number of iterations");
    dt                   = (float)cimg_option("-dt",30.0f,"Time step");
    dlength              = (float)cimg_option("-dl",0.5,"Integration step");
    dtheta               = (float)cimg_option("-da",60.0,"Angular step (in degrees)");
    onormalize = cimg_option("-normalize",false,"Normalize output image in [0,255]");
    int w,h; get_geom(geom,w,h);
    if (!cimg::strcasecmp(visuflow,"circle")) { // Create a circular vector flow
      flow = CImg<>(400,400,1,2);
      cimg_mapXY(flow,x,y) {
	const float ang = (float)(std::atan2(y-0.5*flow.dimy(),x-0.5*flow.dimx()));
	flow(x,y,0) = -(float)std::sin(ang);
	flow(x,y,1) = (float)std::cos(ang);
      }
    }
    if (!cimg::strcasecmp(visuflow,"radial")) { // Create a radial vector flow
      flow = CImg<>(400,400,1,2);
      cimg_mapXY(flow,x,y) {
	const float ang = (float)(std::atan2(y-0.5*flow.dimy(),x-0.5*flow.dimx()));
	flow(x,y,0) = (float)std::cos(ang);
	flow(x,y,1) = (float)std::sin(ang);
      }
    }
    if (!flow.data) flow = CImg<>(visuflow);
    flow.resize(w,h,1,2,3);
    if (normalize) flow.orientation_pointwise();
    if (file_i) img = CImg<>(file_i);
    else img = CImg<>(flow.dimx(),flow.dimy(),1,1,0).noise(100,2);
    img0=img;
    img0.fill(0);
    float color[3]={255,255,255};
    img0.draw_quiver(flow,color,15,-10);
    G = CImg<>(img.dimx(),img.dimy(),1,3);
  }
  
  // Init images and display  
  //-------------------------
  CImg<>
    dest(img.dimx(),img.dimy(),1,img.dimv()),
    sum(dest.dimx(),dest.dimy()),
    W(dest.dimx(),dest.dimy(),1,2);
  
  CImgDisplay *disp=NULL;
  if (visu) disp = new CImgDisplay(img,"Processing");
  std::fprintf(stderr,"\n");
  
  //-------------------------------------
  // Begin regularization PDE iterations
  //-------------------------------------
  int counter = 0;
  for (unsigned int iter=0; iter<nb_iter && (!disp || !disp->closed); iter++) {
    
    // Compute smoothed structure tensor field G
    //-------------------------------------------
    if (!visuflow && !inpaint) {
      CImg_3x3(I,float);
      G.fill(0);
      cimg_mapV(img,k) cimg_map3x3(img,x,y,0,k,I) {
	const float ix = 0.5f*(Inc-Ipc), iy = 0.5f*(Icn-Icp);
	G(x,y,0)+= ix*ix; G(x,y,1)+= ix*iy; G(x,y,2)+= iy*iy;    
      }
      G.blur(sigma);
    }
    
    // Compute normalized tensor field sqrt(T) in G
    //-----------------------------------------------
    if (restore || resize) cimg_mapXY(G,x,y) {
      G.get_tensor(x,y).symeigen(eigen(0),eigen(1));
      const float
	l1 = eigen(0)[0],
	l2 = eigen(0)[1],
	u = eigen(1)[0],
	v = eigen(1)[1],      
	n1 = (float)(1.0/std::pow(1.0f+l1+l2,0.5f*power1)),
	n2 = (float)(1.0/std::pow(1.0f+l1+l2,0.5f*power2));
      G(x,y,0) = n1*u*u + n2*v*v;
      G(x,y,1) = u*v*(n1-n2);
      G(x,y,2) = n1*v*v + n2*u*u;
    }    
    if (visuflow) cimg_mapXY(G,x,y) {
      const float 
	u = flow(x,y,0),
	v = flow(x,y,1),
	n = (float)std::pow(u*u+v*v,0.25f),
	nn = n<1e-5?1:n;
      G(x,y,0) = u*u/nn;
      G(x,y,1) = u*v/nn;
      G(x,y,2) = v*v/nn;
    }

    const CImgStats stats(G,false);
    G/=cimg::max(std::fabs(stats.max),std::fabs(stats.min));

    // Compute LIC's along different angle projections a_\alpha
    //----------------------------------------------------------
    dest.fill(0);
    sum.fill(0);
    for (float theta=(180%(int)dtheta)/2.0f; theta<180; theta+=dtheta) {
      const float
	rad = (float)(theta*cimg::PI/180.0),
	cost = (float)std::cos(rad),
	sint = (float)std::sin(rad);
      
      // Compute vector field w = sqrt(T)*a_alpha
      {
	cimg_mapXY(W,x,y) {
	  const float 
	    a = G(x,y,0),
	    b = G(x,y,1),
	    c = G(x,y,2),
	    u = a*cost + b*sint,
	    v = b*cost + c*sint;
	  W(x,y,0) = u;
	  W(x,y,1) = v;
	}
      }
      
      // Compute the LIC along w in backward and forward directions
      cimg_mapXY(dest,x,y) if (counter++,!mask.data || mask(x,y)) {
	if (!((counter)%150)) std::fprintf(stderr,"\r- Processing : %d %%\t",
					   (int)((counter*100.0)/(dest.dimx()*dest.dimy()*nb_iter*(180/dtheta))));
	
	float l, X,Y, cu, cv, lsum=0;
	const float 
	  fsigma2 = 2*dt*(W(x,y,0)*W(x,y,0) + W(x,y,1)*W(x,y,1)),
	  length = gauss_prec*(float)std::sqrt(fsigma2);
	
	if (linear) {
	  
	  // Integrate with linear interpolation
	  cu = W(x,y,0); cv = W(x,y,1); X=(float)x; Y=(float)y;
	  for (l=0; l<length && X>=0 && Y>=0 && X<=W.dimx()-1 && Y<=W.dimy()-1; l+=dlength) {
	    float u = (float)(W.linear_pix2d(X,Y,0)), v = (float)(W.linear_pix2d(X,Y,1));
	    const float coef = (float)std::exp(-l*l/fsigma2);
	    if ((cu*u+cv*v)<0) { u=-u; v=-v; }
	    cimg_mapV(dest,k) dest(x,y,k)+=(float)(coef*img.linear_pix2d(X,Y,k));
	    X+=dlength*u; Y+=dlength*v; cu=u; cv=v; lsum+=coef;
	  }
	  cu = W(x,y,0); cv = W(x,y,1); X=x-dlength*cu; Y=y-dlength*cv;
	  for (l=dlength; l<length && X>=0 && Y>=0 && X<=W.dimx()-1 && Y<=W.dimy()-1; l+=dlength) {
	    float u = (float)(W.linear_pix2d(X,Y,0)), v = (float)(W.linear_pix2d(X,Y,1));
	    const float coef = (float)std::exp(-l*l/fsigma2);
	    if ((cu*u+cv*v)<0) { u=-u; v=-v; }
	    cimg_mapV(dest,k) dest(x,y,k)+=(float)(coef*img.linear_pix2d(X,Y,k));
	    X-=dlength*u; Y-=dlength*v; cu=u; cv=v; lsum+=coef;
	  }
	} else {

	  // Integrate with non linear interpolation
	  cu = W(x,y,0); cv = W(x,y,1); X=(float)x; Y=(float)y; 
	  for (l=0; l<length && X>=0 && Y>=0 && X<=W.dimx()-1 && Y<=W.dimy()-1; l+=dlength) {
	    float u = W((int)X,(int)Y,0), v = W((int)X,(int)Y,1);
	    const float coef = (float)std::exp(-l*l/fsigma2);
	    if ((cu*u+cv*v)<0) { u=-u; v=-v; }
	    cimg_mapV(dest,k) dest(x,y,k)+=(float)(coef*img.linear_pix2d(X,Y,k));
	    X+=dlength*u; Y+=dlength*v; cu=u; cv=v; lsum+=coef;
	  }	
	  cu = W(x,y,0); cv = W(x,y,1); X=x-dlength*cu; Y=y-dlength*cv;
	  for (l=dlength; l<length && X>=0 && Y>=0 && X<=W.dimx()-1 && Y<=W.dimy()-1; l+=dlength) {
	    float u = W((int)X,(int)Y,0), v = W((int)X,(int)Y,1);
	    const float coef = (float)std::exp(-l*l/fsigma2);
	    if ((cu*u+cv*v)<0) { u=-u; v=-v; }
	    cimg_mapV(dest,k) dest(x,y,k)+=(float)(coef*img.linear_pix2d(X,Y,k));
	    X-=dlength*u; Y-=dlength*v; cu=u; cv=v; lsum+=coef;
	  }
	}
	sum(x,y)+=lsum;
      }
      if (!inpaint && disp) dest.display(*disp);
      if (disp && disp->resized) disp->resize();
    }

    // Average all the LIC's
    //-----------------------
    cimg_mapXY(dest,x,y) if (sum(x,y)>0) cimg_mapV(dest,k) dest(x,y,k)/=sum(x,y); else cimg_mapV(dest,k) dest(x,y,k)=img(x,y,k);

    if (inpaint && disp) dest.display(*disp);
    img = dest;
    if (file_o && save && !(iter%save)) dest.save(file_o,iter);
  }
  
  // Save result and end program
  //-----------------------------
  std::fprintf(stderr,"\r- Processing : Done !\n");
  if (visuflow) dest.mul(flow.get_norm_pointwise()).normalize(0,255);
  if (onormalize) dest.normalize(0,255);
  if (file_o) {
    std::fprintf(stderr,"- Saving output image '%s'\n",file_o);
    if (!append_result) dest.save(file_o);
    else CImgl<>(img0,dest).get_append('x').save(file_o);
  }
  if (disp) { delete disp; CImgl<>(img0,dest).display("Original and processed images"); }
  
  return 0;
}
