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

  File        : image_registration.cpp

  Description : Compute a motion field between two images, 
                with a multiscale and variational algorithm

  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;

// animate_warp() : Create warping animation from two images and a motion field
//----------------
void animate_warp(const CImg<unsigned char>& src, const CImg<unsigned char>& dest, const CImg<>& u,
                  const bool morph, const bool imode,
                  const char *filename,int nb,CImgDisplay *disp) {
  CImg<unsigned char> visu = CImgl<unsigned char>(src,dest,src).get_append('x'), warp(src,false);
  float t=0;
  for (unsigned int iter=0; !disp || (!disp->closed && disp->key!=cimg::keyQ); iter++) {
    if (morph) cimg_mapXYV(warp,x,y,k) {
      const float dx = u(x,y,0), dy = u(x,y,1), 
        I1 = (float)src.linear_pix2d(x-t*dx, y-t*dy, k),
        I2 = (float)dest.linear_pix2d(x+(1-t)*dx,y+(1-t)*dy,k);
      warp(x,y,k) = (unsigned char)((1-t)*I1 + t*I2);
    } else cimg_mapXYV(warp,x,y,k) {
      const float dx = u(x,y,0), dy = u(x,y,1), I1 = (float)src.linear_pix2d(x-t*dx, y-t*dy, 0,k);
      warp(x,y,k) = (unsigned char)I1;
    }
    if (disp) visu.draw_image(warp,2*src.width,0).display(disp->resize()).wait(30);
    if (filename && *filename && (imode || (int)iter<nb)) {
      std::fprintf(stderr,"\r  > frame %d           ",iter);
      warp.save(filename,iter); 
    }
    iter++;
    t+=1.0f/nb;
    if (t<0) { t=0; nb=-nb; }
    if (t>1) { t=1; nb=-nb; if (filename && *filename) exit(0); }
  }
}

// get_warp() : Return the image src warped by the motion field u.
//------------
template<typename T> CImg<T> getwarp(const CImg<T>& src, const CImg<>& u) {
  CImg<T> warp(src);
  cimg_mapXY(warp,x,y) warp(x,y) = (T)src.linear_pix2d(x - u(x,y,0), y - u(x,y,1));
  return warp;
}


// optmonoflow() : Compute optical flow for one scale ( semi-implicite scheme ) between I2->I1
//---------------
CImg<> optmonoflow(const CImg<>& I1,const CImg<>& I2,const CImg<>& u0, 
                   const float smooth, const float precision,CImgDisplay *disp) {

  CImg<> u = u0.get_resize(I1.width,I1.height,1,2,3),dI(u);
  CImg_3x3(I,float);
  float dt=2,E=1e20f;

  // compute first derivatives of I2
  cimg_map3x3(I2,x,y,0,0,I) {
    dI(x,y,0) = 0.5f*(Inc-Ipc);
    dI(x,y,1) = 0.5f*(Icn-Icp);
  }

  // Main PDE iteration
  for (unsigned int iter=0; iter<100000; iter++) {
    std::fprintf(stderr,"\r- Iteration %d - E = %g",iter,E); fflush(stderr);
    const float Eold = E;
    E = 0;
    cimg_3mapXY(u,x,y) {
      const float 
        X = x + u(x,y,0),
        Y = y + u(x,y,1),
        deltaI = (float)(I2.linear_pix2d(X,Y) - I1(x,y));
      float tmpf = 0;
      cimg_mapV(u,k) {
        const float
          ux  = 0.5f*(u(_nx,y,k)-u(_px,y,k)),
          uy  = 0.5f*(u(x,_ny,k)-u(x,_py,k));
        u(x,y,k) = (float)( u(x,y,k) +
                            dt*(
                                -deltaI*dI.linear_pix2d(X,Y,k) +
                                smooth* ( u(_nx,y,k) + u(_px,y,k) + u(x,_ny,k) + u(x,_py,k) )
                                )
                            )/(1+4*smooth*dt);
        tmpf += ux*ux + uy*uy;
      }
      E += deltaI*deltaI + smooth * tmpf;
    }
    if (std::fabs(Eold-E)<precision) break;
    if (Eold<E) dt*=0.5;
    if (disp) disp->resize();
    if (disp && !(iter%50)) {
      const unsigned char white = 255;
      CImg<unsigned char> tmp = getwarp(I1,u).normalize(0,200);
      tmp.resize(disp->width,disp->height).draw_quiver(u,&white,15,-14,0,0.7f).display(*disp);
    }
  }
  return u;
}


// optflow() : multiscale version of the Horn&Schunk algorithm
//-----------
CImg<> optflow(const CImg<>& xsrc,const CImg<>& xdest,
               const float smooth,const float precision,const unsigned int pnb_scale,CImgDisplay *disp=NULL) {
  const CImg<>
    src  = xsrc.get_norm_pointwise(1).resize(xdest.width,xdest.height,1,1,3).normalize(0,1),
    dest = xdest.get_norm_pointwise(1).resize(xdest.width,xdest.height,1,1,3).normalize(0,1);
  CImg<> u = CImg<>(src.width,src.height,1,2).fill(0);

  const unsigned int nb_scale = pnb_scale>0?pnb_scale:(unsigned int)(2*std::log((double)(cimg::max(src.width,src.height))));
  for (int scale=nb_scale-1; scale>=0; scale--) {
    const CImg<> I1 = src.get_resize((int)(src.width/std::pow(1.5,scale)), (int)(src.height/std::pow(1.5,scale)) ,1,1,3);
    const CImg<> I2 = dest.get_resize((int)(src.width/std::pow(1.5,scale)), (int)(src.height/std::pow(1.5,scale)) ,1,1,3);
    std::fprintf(stderr," * Scale %d\n",scale);
    u*=1.5;
    u = optmonoflow(I1,I2,u,smooth,(float)(precision/std::pow(2.25,1+scale)),disp);
    std::fprintf(stderr,"\n");
  }
  return u;
}


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

  Main function

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

int main(int argc,char **argv) {
  
  // Read command line parameters
  cimg_usage("Compute an optical flow between two 2D images, and create a warped animation");
  const char
    *name_i1   = cimg_option("-i","img/sh0r.pgm","Input Image 1 (Destination)"),
    *name_i2   = cimg_option("-i2","img/sh1r.pgm","Input Image 2 (Source)"),
    *name_o    = cimg_option("-o",(const char*)NULL,"Output 2D flow (inrimage)"),
    *name_seq  = cimg_option("-o2",(const char*)NULL,"Output Warping Sequence");
  const float
    smooth    = (float)cimg_option("-s",0.1f,"Flow Smoothness"),
    precision = (float)cimg_option("-p",0.9f,"Convergence precision");
  const unsigned int
    nb        = cimg_option("-n",40,"Number of warped frames"),
    nbscale   = cimg_option("-scale",0,"Number of scales (0=auto)");
  const bool
    normalize = cimg_option("-equalize",true,"Histogram normalization of the images"),
    morph     = cimg_option("-m",true,"Morphing mode"),
    imode     = cimg_option("-c",true,"Complete interpolation (or last frame is missing)"),
    dispflag = !cimg_option("-novisu",false,"Visualization");
  
  // Init images and display
  std::fprintf(stderr," - Init images.\n");
  const CImg<>
    src(name_i1),
    dest(CImg<>(name_i2).resize(src,3)),
    src_blur  = normalize?src.get_blur(0.5f).equalize_histogram():src.get_blur(0.5f),
    dest_blur = normalize?dest.get_blur(0.5f).equalize_histogram():dest.get_blur(0.5f);
  
  CImgDisplay *disp = NULL;
  if (dispflag) {
    unsigned int w = src.width, h = src.height;
    const unsigned int dmin = cimg::min(w,h), minsiz = 512;
    if (dmin<minsiz) { w=w*minsiz/dmin; h=h*minsiz/dmin; }
    const unsigned int dmax = cimg::max(w,h), maxsiz = 1024;
    if (dmax>maxsiz) { w=w*maxsiz/dmax; h=h*maxsiz/dmax; }
    disp = new CImgDisplay(w,h,"Estimated Motion",0,3);
  }

  // Run Motion estimation algorithm
  std::fprintf(stderr," - Compute optical flow.\n");
  const CImg<> u = optflow(src_blur,dest_blur,smooth,precision,nbscale,disp);
  if (name_o) u.save(name_o);
  u.print("Computed flow");
  
  // Do morphing animation
  std::fprintf(stderr," - Create warped animation.\n");
  CImgDisplay *disp2 = NULL;
  if (dispflag) {
    unsigned int w = src.width, h = src.height;
    const unsigned int dmin = cimg::min(w,h), minsiz = 100;
    if (dmin<minsiz) { w=w*minsiz/dmin; h=h*minsiz/dmin; }
    const unsigned int dmax = cimg::max(w,h), maxsiz = 1024/3;
    if (dmax>maxsiz) { w=w*maxsiz/dmax; h=h*maxsiz/dmax; }
    disp2 = new CImgDisplay(3*w,h,"Source/Destination images and Motion animation",0,19);
  }
  animate_warp(src.get_normalize(0,255),dest.get_normalize(0,255),u,morph,imode,name_seq,nb,disp2);
  
  if (disp) delete disp;
  if (disp2) delete disp2;
  exit(0);
  return 0;
}
