/**
   Author: Neil Birkbeck, 2009

   Semi-global stereo matching.

   An implementation of (without mutual information):

   @inproceedings{1069169,
   author = {Hirschmuller, Heiko},
   title = {Accurate and Efficient Stereo Processing by Semi-Global Matching and Mutual Information},
   booktitle = {CVPR '05: Proceedings of the 2005 IEEE Computer Society Conference 
                on Computer Vision and Pattern Recognition (CVPR'05) - Volume 2},
   year = {2005},
   isbn = {0-7695-2372-2},
   pages = {807--814},
   doi = {http://dx.doi.org/10.1109/CVPR.2005.56},
   publisher = {IEEE Computer Society},
   address = {Washington, DC, USA},
   }

   TODO: 
    -forward-backward, Mutual information
    -Other scores (make sure they are scaled to fit in signed shorts).
    -May need to make sure scores/cost are aligned (use stride) for images
     that are not divisible by 4.  As the disparity must be divisible by 4
     this is not so critical (e.g., 4*2 + (16-bit aligned addie) ==> at least 8-byte aligned)
*/

#include <algorithm>
#include <map>

#include <stdint.h>
#include <omp.h>
#include <mmintrin.h>
#include <xmmintrin.h>

//My libraries, should be easy enough to replace/remove
#include <nimage/image.h>
#include <nmisc/timer.h>
#include <nmisc/commandline.h>

#include <boost/thread/thread.hpp>

using namespace nacb;
using std::pair;

typedef uint16_t cost_type_t;
const cost_type_t cost_max = 65535;
using std::swap;


/** \brief Struct for holding scanline direction and for 
    generating borders where traversal needs to start.
*/
struct scan_direction_t {
  int dx[2], dy[2];

  scan_direction_t(){ }

  scan_direction_t(int dx_, int dy_){
    dx[0] = dx[1] = dx_;
    dy[0] = dy[1] = dy_;
  }

  scan_direction_t(int dx0, int dy0,
		   int dx1, int dy1){
    dx[0] = dx0;
    dy[0] = dy0;

    dx[1] = dx1;
    dy[1] = dy1;
  }

  /** \brief Create the line in the opposite direction.
   */
  scan_direction_t flip() const {
    return scan_direction_t(-dx[0], -dy[0], -dx[1], -dy[1]);
  }
  
  /** \brief return a number of borders where traversal
      needs to start.
   */
  void get_border(int w,  int h,
		  int & bx, int & by,
		  int & ex, int & ey) const {
    //Start at left edge
    if(dx[0] + dx[1] == 2){
      bx = 0;
      by = 0;
      ex = 0;
      ey = h-1;
    }
    //Start at right edge
    else if(dx[0] + dx[1]== -2){
      bx = w-1;
      by = 0;
      ex = w-1;
      ey = h-1;
    }
    //Start at top edge
    else if(dy[0] + dy[1] ==  2){
      bx = 0;
      by = 0;
      ex = (w-1);
      ey= 0;
    }
    else {
      bx = 0;
      by = h-1;
      ex = (w-1);
      ey = h-1;
    }
  }
};

cost_type_t sad(const nacb::Image8 & im1,
		const nacb::Image8 & im2,
		int x, int y, int d){
  return abs((cost_type_t)im1(x, y, 0) - (cost_type_t)im2(x+d, y, 0)) +
    abs((cost_type_t)im1(x, y, 1) - (cost_type_t)im2(x+d, y, 1)) + 
    abs((cost_type_t)im1(x, y, 2) - (cost_type_t)im2(x+d, y, 2));
}

static inline bool valid(int x, int y, int w, int h){
  return (x>=0 && y>=0 && x<w && y<h);
}


std::vector<scan_direction_t> get_directions(bool sixteen = true){
  std::vector<scan_direction_t> dirs;
  
  dirs.push_back(scan_direction_t(1, 0));
  dirs.push_back(scan_direction_t(0, 1));
  dirs.push_back(scan_direction_t(1, 1));
  dirs.push_back(scan_direction_t(1, -1));
  
  if(sixteen){
    dirs.push_back(scan_direction_t(1, 1,  1, 0));
    dirs.push_back(scan_direction_t(1, 1,  0, 1));
    dirs.push_back(scan_direction_t(1, -1,  1, 0));
    dirs.push_back(scan_direction_t(-1, 1,  0, 1));
  }
  int n = dirs.size();
  for(int i=0; i<n; i++)dirs.push_back(dirs[i].flip());
  
  return dirs;
}


std::pair<Imagef, Imagef>  get_best_disparity(cost_type_t * scost,
					      int w, int h,
					      int dmin, int dmax){
  int ndisp = dmax - dmin;

  nacb::Imagef disp(w, h, 1);
  nacb::Imagef bestCostImage(w, h, 1);

  for(int y=0; y<h; y++){
    for(int x=0; x<w; x++){
      cost_type_t best = cost_max;
      int best_disp = 0;
      for(int d=0; d<ndisp; d++){
	if(scost[(y*w+x)*ndisp + d]<best){
	  best = scost[(y*w+x)*ndisp + d];
	  best_disp = d;
	}
      }
      disp(x, y) = best_disp;
      bestCostImage(x, y) = best;
    }
  }
  return std::pair<Imagef, Imagef>(disp, bestCostImage);
}


/** \brief MMX-based per-pixel accumulation of scores.
 */
static void mmx_inner(const cost_type_t * score_base, 
		      cost_type_t * obase, cost_type_t * lcost, 
		      cost_type_t * lcost_prev, cost_type_t  mcost_many[4],
		      int ndisp, const cost_type_t p1[4], const cost_type_t p2[4]){
  __m64 this_mins;
  __m64 mn = *((__m64 *)mcost_many);
  cost_type_t * next = lcost;
  
  for(int d0=0; d0<ndisp; d0+=4){
    __m64 penalty = *((__m64*)(p1));
    __m64 n = *((__m64*)(lcost_prev + d0));
    
    //Only consider middle shifts (Are these correct?)
    n = _m_pminsw(n, _m_pshufw(n, 0x90) + penalty);
    n = _m_pminsw(n, _m_pshufw(n, 0xF9) + penalty);
    
    //Need two more operations for boundaries...
    if(d0){
      __m64 t = *((__m64*)(lcost_prev + d0 - 1));
      n = _m_pminsw(n, _m_pshufw(t, 0xE7) + penalty);
    }
    if(d0+1<ndisp){
      __m64 t = *((__m64*)(lcost_prev + d0 + 1));
      n = _m_pminsw(n, _m_pshufw(n, 0x24) + penalty);
    }
    
    penalty = *((__m64*)(p2));
    
    //Now do all the disparities
    for(int d=0; d<ndisp; d+=4){
      __m64 m5 = *((__m64*)(lcost_prev + d));
      __m64 m6 = _m_pminsw(m5, _m_pshufw(m5, 0x93));
      m6 = _m_pminsw(m6, _m_pshufw(m5, 0x4E));
      m6 = _m_pminsw(m6, _m_pshufw(m5, 0x39));
      n = _m_pminsw(n, m6 + penalty);
    }
    
    __m64 cost_many = *((__m64*)(score_base + d0));
    __m64 c = cost_many + n - mn;
    
    *((__m64 *)(lcost + d0)) = c;
    *((__m64 *)(obase + d0)) += c;//*((__m64 *)(out + ndisp*x + d0)) + c;
    
    if(d0==0) this_mins = c;
    else this_mins = _m_pminsw(this_mins, c);
  }
  
  this_mins = _m_pminsw(this_mins, _m_pshufw(this_mins, 0x93));
  this_mins = _m_pminsw(this_mins, _m_pshufw(this_mins, 0x4E));
  this_mins = _m_pminsw(this_mins, _m_pshufw(this_mins, 0x39));
  
  *((__m64 *)(mcost_many)) = this_mins;
}


/** \brief Standard per-pixel accumulation of scores.
 */
static void inline brute_inner(const cost_type_t * score_base, 
			       cost_type_t * obase, 
			       cost_type_t * lcost, 
			       cost_type_t * lcost_prev, cost_type_t & mcost,
			       int ndisp, cost_type_t P1, cost_type_t P2){
  cost_type_t mn = 0;
  
  for(int d=0; d<ndisp; d++){
    cost_type_t c = score_base[d];
    
    //This assumes:
    //      (lcost_prev = [0, 0, ..., 0] && mcost ==0) ==> (n == 0)
    {
      cost_type_t n = lcost_prev[d];
      
      if(d>0)
	n = std::min(cost_type_t(lcost_prev[d-1] + P1), n);
      
      if(d+1<ndisp)
	n = std::min(cost_type_t(lcost_prev[d+1] + P1), n);
      
      for(int d2=0; d2<ndisp; d2++)
	n = std::min(cost_type_t(lcost_prev[d2] + P2), n);
      
      c = c + n - mcost;
    }
    
    lcost[d] = c;
    obase[d] += c;
    
    //Output
    if(d==0)mn = c;
    else mn = std::min(c, mn);
  }
  mcost = mn;
}


/** \brief The main loop for going along a border.
    
    It goes along the border and starts scan-lines of the 
    given direction.  These lines are allowed to overlap the side
    borders, but each scanline stops when it reaches the opposite
    border.
 */
template <bool use_mmx>
static void border_loop(const cost_type_t * score,
			cost_type_t * scost, 
			const scan_direction_t & line,
			cost_type_t * lcost, cost_type_t * lcost_prev, 
			int w, int h, int ndisp, 
			cost_type_t p1_many[4], 
			cost_type_t p2_many[4],
			int offs=0, int scale=1){

  int bx, by, ex, ey;

  line.get_border(w, h, bx, by, ex, ey);

  int dbx = (ex > bx);
  int dby = (ey > by);

  bx += offs*dbx;
  by += offs*dby;
  
  dbx *= scale;
  dby *= scale;

  for(; bx <= ex && by <= ey; bx+=dbx , by+=dby){
    //Initialization of types
    cost_type_t mcost_many[4] = {0, 0, 0, 0};
    int ii = 0;

    //Start with previous cost of 0.
#define reset_scanline memset(lcost_prev, 0, ndisp*sizeof(cost_type_t));\
    memset(mcost_many, 0, sizeof(mcost_many));

    reset_scanline;

    int x = 0;
    int y = 0;
    bool go = true;
    for(y=by, x=bx; go; x+=line.dx[ii&0x1], y+=line.dy[ii&0x1], ii++){
      //This could perhaps be better.
      if(dbx>0){
	if(ey==0){
	  if (y==-1 || y==h)break;
	  else if(x==w){
	    x -= w;
	    reset_scanline;
	  }
	  else if(x==-1){
	    x += w;
	    reset_scanline;
	  }
	}
	//bottom of image
	else {
	  if (y==-1 || y==h)break;
	  else if(x==w){
	    x -= w;
	    reset_scanline;
	  }
	  else if(x==-1){
	    x += w;
	    reset_scanline;
	  }
	}
      }
      else {
	//left
	if(ex == 0){
	  if (x==w || x==-1)break;
	  else if(y==h){
	    y -=h;
	    reset_scanline;
	  }
	  else if(y==-1){
	    y += h;
	    reset_scanline;
	  }
	}
	else {
	  if (x==w || x==-1)break;
	  else if(y==h){
	    y -=h;
	    reset_scanline;
	  }
	  else if(y==-1){
	    y += h;
	    reset_scanline;
	  }
	}
      }


#ifdef SG_DEBUG_HITS
      //Testing how many hits, should be equal, and it is.
      scost[(y*w + x)]++;
#else      
      const cost_type_t * score_base = score + (y*w + x)*ndisp;
      cost_type_t * out_base = scost + (y*w + x)*ndisp;
      
      if(use_mmx)
	mmx_inner(score_base, out_base, lcost, lcost_prev, 
		  mcost_many, ndisp, p1_many, p2_many);
      else
	brute_inner(score_base, out_base, lcost, lcost_prev, 
		    mcost_many[0], ndisp, p1_many[0], p2_many[0]);
      
#endif

      //Swap the temporary buffers
      swap(lcost, lcost_prev);
    }
  }
#undef reset_scanline
}

/** \brief Debug output for mmx_test.
 */
void show_disp(cost_type_t * out, int ndisp, int nshow){
  for(int i=0; i<ndisp*nshow; i++){
    if(i && i % ndisp==0)printf("\n");
    printf("%d ", out[i]);
  }
  printf("\n");
}

/** \brief Timing and accuracy test for the mmx implementation.
 */
void mmx_test(int ndisp = 32, int npix = 1024*1024){
  //Alignment should be ensured
  cost_type_t * cost = new cost_type_t[ndisp*npix];
  cost_type_t * data = new cost_type_t[ndisp*npix];
  cost_type_t * out  = new cost_type_t[ndisp*npix];

  cost_type_t * prev = new cost_type_t[ndisp];
  cost_type_t * next  = new cost_type_t[ndisp];

  cost_type_t last_mins[4];

  for(int i=0; i<ndisp*npix; i++){
    data[i] = 10000 - i*2;
    cost[i] = i + i/2;    
  }  
  
  cost_type_t P1 = 20, P2 = 40;
  cost_type_t p1_many[4] = {P1, P1, P1, P1};
  cost_type_t p2_many[4] = {P2, P2, P2, P2};

  int nshow = std::min(2, npix);

  nacb::StopWatch timer;
  double mmx_time = 0, basic_time = 0;

#define clear_all memset(prev, 0, ndisp*sizeof(cost_type_t));	\
  memset(next, 0, ndisp*sizeof(cost_type_t));			\
  memset(out, 0, ndisp*npix*sizeof(cost_type_t));		\
  memset(last_mins, 0, sizeof(last_mins));			\
  printf("Done clearing.\n");

  clear_all;

  timer.reset();
  
  for(int x=0; x<npix; x++){
    cost_type_t * score_base = cost + x*ndisp;
    cost_type_t * out_base = out + x*ndisp;
    
    mmx_inner(score_base, out_base, next, prev, last_mins, ndisp, p1_many, p2_many);
    
    swap(prev, next);
  }
  
  printf("MMX: %lf\n", mmx_time = double(timer));
  show_disp(out, ndisp, nshow);    
    
  clear_all;

  //for(int i=0; i<ndisp; i++)prev[i] = P1;
  memset(prev, 0, ndisp*sizeof(cost_type_t));
  memset(next, 0, ndisp*sizeof(cost_type_t));
  memset(out, 0, npix*ndisp*sizeof(cost_type_t));

  timer.reset();

  cost_type_t last_min = 0;

  for(int x=0; x<npix; x++){
    cost_type_t * cbase = cost + x*ndisp;
      
    brute_inner(cbase, out + x*ndisp, next, prev, last_min, ndisp, P1, P2);
      
    swap(prev, next);
  }

  printf("Basic: %lf\n", basic_time = double(timer));

  show_disp(out, ndisp, nshow);
  
  printf("Speedup: %lf\n", basic_time/mmx_time);
}


/**
   Careful with this guy as he has his own thread.
 */
template <bool use_mmx>
struct thread_data {
  int w, h;
  pthread_t thread;

  cost_type_t * scost;
  const cost_type_t * score;
  cost_type_t * lcost_prev, * lcost;
  cost_type_t p1_many[4], p2_many[4];

  int dmin, dmax, ndisp;

  scan_direction_t line;
  int offs, scale;

  thread_data(int w_, int h_,
	      int dmin_, int dmax_,
	      cost_type_t P1, cost_type_t P2,
	      cost_type_t * scost_,
	      const cost_type_t * score_): w(w_), h(h_), 
					   dmin(dmin_), dmax(dmax_), 
					   scost(scost_), score(score_)
  {      
    for(int j=0; j<4; j++){
      p1_many[j] = P1;
      p2_many[j] = P2;
    }

    ndisp = dmax - dmin;

    posix_memalign((void**)&lcost, 16, sizeof(cost_type_t)*w*h);
    posix_memalign((void**)&lcost_prev, 16, sizeof(cost_type_t)*w*h);

    thread = 0;
    offs = 0;
    scale = 1;
  }

  ~thread_data(){
    if(thread)join();
    free(lcost);
    free(lcost_prev);
  }

  void set_line(const scan_direction_t & line, int offs, int scale){
    this->line = line;
    this->offs = offs;
    this->scale = scale;
  }

  void start(){
    if(thread)join();
    pthread_create(&thread, 0, s_run, this);
  }
    
  void run(){
    border_loop<use_mmx>(score, scost, line,
			 lcost, lcost_prev, w, h, ndisp, p1_many, p2_many, offs, scale);
  }
    
  static void * s_run(void * d){
    ((thread_data*)d)->run();
    return 0;
  }
    
  void join(){
    if(thread){
      pthread_join(thread, 0);
      thread = 0;
    }
  }    
private:
  //Undefined
  void operator=(const thread_data & t){ }
};


/** \brief Semi global matching on images im1 and im2.
    
    It is templated to hopefully allow the compiler as much optimization as possible.
*/
template <int use_mmx>
std::pair<Imagef, Imagef> semi_global(const nacb::Image8 & im1,
				      const nacb::Image8 & im2,
				      int dmin, int dmax,
				      cost_type_t P1, cost_type_t P2,
				      bool sixteen, int nthreads = 0){
  if(dmin>dmax)swap(dmin, dmax);

  if((dmax - dmin)%4){
    int extra = 4 - (dmax - dmin)%4;
    dmax += extra;

    printf("added to make mod(4): |%d-%d|=%d\n", dmin, dmax, dmax - dmin);
  }

  int w = im1.w, h = im1.h;
  int w2 = im2.w, h2 = im2.h;

  cost_type_t C = 128*3;
  
  cost_type_t p1_many[4] = {P1, P1, P1, P1};
  cost_type_t p2_many[4] = {P2, P2, P2, P2};

  int ndisp = dmax - dmin;

  cost_type_t * scost = 0, * score = 0; // = new cost_type_t[w1*h1*ndisp];
  cost_type_t * lcost_d1 = 0, * lcost_d2 = 0;

  //Align on 16-byte bounds, just to be safe
  posix_memalign((void**)&scost, 16, sizeof(cost_type_t)*w*h*ndisp);//new cost_type_t[w1*h1*ndisp];
  posix_memalign((void**)&score, 16, sizeof(cost_type_t)*w*h*ndisp);//new cost_type_t[w1*h1*ndisp];
  posix_memalign((void**)&lcost_d1, 16, sizeof(cost_type_t)*w*h);
  posix_memalign((void**)&lcost_d2, 16, sizeof(cost_type_t)*w*h);
  
  memset(score, 0, sizeof(cost_type_t)*w*h*ndisp);
  memset(scost, 0, sizeof(cost_type_t)*w*h*ndisp);
  
  cost_type_t * lcost = lcost_d1, * lcost_prev = lcost_d2;
  std::vector<scan_direction_t> lines = get_directions(sixteen);

  nacb::StopWatch timer;

  // Get the scores.
  for(int d=0; d<ndisp; d++){
    for(int y=0; y<h; y++){
      for(int x=0; x<w; x++){
	cost_type_t c = C;

	if(valid(x + d + dmin, y, w2, h2))
	  c = sad(im1, im2, x, y, d + dmin);

	score[(y*w+x)*ndisp + d] = c;
      }
    }
  }

  printf("scores: %lf\n", double(timer));
  timer.reset();

  _m_empty();
  
  std::vector<thread_data<use_mmx> *> threads;
  for(int i=0; i<nthreads; i++)
    threads.push_back(new thread_data<use_mmx>(w, h, dmin, dmax, P1, P2, scost, score));
  
  //Cost accumulation along each line.
  for(int lni=0; lni<lines.size(); lni++){
    const scan_direction_t & line = lines[lni];
    
    if(nthreads){
      for(int i=0; i<nthreads; i++){
	threads[i]->join();
	threads[i]->set_line(line, i, nthreads);
	threads[i]->start();
      }
    }
    else {
      border_loop<use_mmx>(score, scost, line,
			   lcost, lcost_prev, w, h, ndisp, p1_many, p2_many);
    }
  }

  for(int i=0; i<nthreads; i++)
    delete threads[i];

  pair<Imagef, Imagef> res = get_best_disparity(scost, w, h, dmin, dmax);

  std::map<int, int> cnts;
  nacb::Imagef cnt(w, h, 1);
  for(int i=0; i<w*h; i++){
    cnts[scost[i]]++;
    cnt.data[i] = scost[i];
  }
  
#ifdef SG_DEBUG_HITS
  std::map<int, int>::iterator it = cnts.begin();
  for(;it != cnts.end(); ++it){
    printf("%d:  %d\n", it->first, it->second);
  }
   
  double n, x;
  cnt.getRange(n, x);
  (cnt/x).write("/tmp/cnt.png");
#endif

  free(lcost_d1);
  free(lcost_d2);
  free(score);
  free(scost);

  return res;
}
		   

#ifdef SG_MAIN
int main(int ac, char * av[]){
  int P1=45, P2=75;
  int dmin =  0;
  int dmax =  32;
  int use_mmx = true;
  int med_size = 0;
  int sixteen = true;
  int nthreads = 0;

  nacb::CommandLine cline("./sgstereo [options] image1 image2");
  cline.registerOption("x1", "Min disparity (inclusive).", &dmin, '1');
  cline.registerOption("x2", "Max disparity (exclusive).", &dmax, '2');
  cline.registerOption("P1", "Smoothness for single disparity.", &P1);
  cline.registerOption("P2", "Smoothness for disparity past fabs(1).", &P2);
  cline.registerOption("mmx", "Use mmx", &use_mmx, 'm');
  cline.registerOption("test-mmx", "Test the mmx routines.");
  cline.registerOption("med", "Size of median filter (0 for no filter).", &med_size);
  cline.registerOption("sixteen", "Use sixteen lines.", &sixteen);
  cline.registerOption("nthreads", "Number of threads to use.", &nthreads);
  cline.parse(ac, av);

  if(cline.gotArgument("test-mmx")){
    printf("Running mmx time test.\n");
    mmx_test();
    return 0;
  }
  
  if(optind+1>=ac){
    printf("Need two image arguments.\n");
    return 1;
  }
  
  printf("Doing semi-global on %s, %s\n", av[optind], av[optind+1]);
  nacb::Image8 im1(av[optind]), im2(av[optind+1]);
  pair<Imagef, Imagef> res;
  
  if(use_mmx)
    res = semi_global<true>(im1, im2, dmin, dmax, P1, P2, sixteen, nthreads);
  else
    res = semi_global<false>(im1, im2, dmin, dmax, P1, P2, sixteen, nthreads);

  nacb::Imagef bestCostImage = res.second;
  nacb::Imagef disp = res.first;
  if(med_size)
    disp = disp.medianFilter(med_size);

  double n, x;
  disp.getRange(n, x);
  ((disp - n)*(1.0/(x-n))).write("/tmp/d.tga");

  bestCostImage.getRange(n, x);
  ((bestCostImage - n)*(1.0/(x-n))).write("/tmp/cost.tga");
  return 0;
}
#endif