#include <math.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <iostream.h>
#include "image.h"
#include "lib.h"


// rotate
// 
// rotate IM into RIM by angle theta.
// uses (cx,cy) as centre of rotation (may be off of image)
void rotate(float *IM, float *RIM, int xsize, int ysize,
	    int cx, int cy, float theta) {
  float COS = cos(theta);
  float SIN = sin(theta);
  int X,Y;
  // build RIM
  for (int y=0;y<ysize;y++) {
    for (int x=0;x<xsize;x++) {
      // find corresponding coordinate (X,Y) in IM
      X=(int)round((x-cx)*COS+(y-cy)*SIN)+cx;
      Y=(int)round((y-cy)*COS-(x-cx)*SIN)+cy;
      if ((X<0) || (Y<0) || (X>=xsize) || (Y>=ysize)) {
	RIM[x+y*xsize]=0;
      } else {      
	RIM[x+y*xsize]=IM[X+Y*xsize]; }
    }
  }
}


// translate
//
// translates image im1 into im2 using (tx,ty) as translation vector
void translate(float *im1, float *im2, int xsize, int ysize, int tx, int ty) {
  for (int y=0;y<ysize;y++) {
    if (((y+ty)<0) || ((y+ty)>=ysize)) {
      for (int x=0;x<xsize;x++) im2[y*xsize+x]=0;
      continue;
    }
    for (int x=0;x<xsize;x++) {
      if (((x+tx)<0) || ((x+tx)>=xsize)) {
	im2[y*xsize+x]=0;
	continue;
      }
      im2[y*xsize+x]=im1[(y+ty)*xsize+(x+tx)];
    }
  }
}


// align
//
// finds the optimal alignment parameters (theta, tx, ty) between two images
// (im1, im2).  Values of theta in [-theta_range,theta_range] are checked over
// theta_steps intervals.  Values of tx between [-tx_range,tx_range] are 
// checked.  same for ty.
// this routine just attempts to minimize the SSD between image pixels for a
// given region starting at (x,y) of size (xrange, yrange) after the first 
// image has been rotated and translated.
void align(float *im1, float *im2, int xsize, int ysize,
	   int x, int y, int xrange, int yrange,
	   float theta_range, int theta_steps, int tx_range, int ty_range,
	   float &theta, int &tx, int &ty) {

  float *Rimage = new float[xsize*ysize];
  float *B1 = Rimage;
  float *B2 = copy_image(im2,xsize,ysize);

  smooth_image(B2,xsize,ysize,0.5);

  float BestTheta;
  int BestTX;
  int BestTY;
  unsigned int BestScore=0xFFFFFFFF;

  // attempt to align image by rotating, blurring, translating
  for (float THETA = -theta_range; THETA<=theta_range; 
       THETA=THETA+(2*theta_range/theta_steps)) {
    // rotate im1 into Rimage 
    rotate(im1, Rimage, xsize, ysize, xsize/2, ysize/2, THETA);
    // blur B1 (Rimage) and B2
    smooth_image(B1,xsize,ysize,0.5);
    // now we minimize the sum of squared differences between a 
    // translated copy of B1 and B2 
    // we only inspect the bounded region specified by x,y,xrange,yrange
    int whichx, whichy;
    unsigned int cost=0xFFFFFFFF;
    for (int TY=-ty_range;TY<=ty_range;TY++) 
      for (int TX=-tx_range;TX<=tx_range;TX++) {
	unsigned int thesum=0;
	for (int Y=y;Y<y+yrange;Y++)  
	  for (int X=x;X<x+xrange;X++) {
	    int val=(int)(B1[(Y+TY)*xsize+(X+TX)]-B2[Y*xsize+X]);
	    thesum+=(unsigned int)(val*val);
	  }
	// better score
	if (thesum<=cost) {
	  cost=thesum;
	  whichx=TX; whichy=TY;
	}
	
      }
    // whichx, whichy minimized the SSD
    // now compare against best rotated score
    if (cost<=BestScore) {
      BestTheta=THETA;
      BestTX=whichx;
      BestTY=whichy;
      BestScore=cost;
    }
    // continue along to next angle and rematch.
  }
  delete(B1); delete(B2);
  // return the best match found
  theta=BestTheta; tx=BestTX; ty=BestTY;
}


// compute_disparity
//
// calculates disparity from im1 to im2 using a plain block matching algorithm.
// the window used for matching is of size (2*XWIN+1)x(2*YWIN+1)
// the disparity is output into xmap,ymap where (xmap[n],ymap[n]) is the
// 2-dimensional disparity vector for pixel n.
// |disparity| is bounded by D.
void compute_disparity(float *im1, float *im2, int *xmap, int *ymap,
		       int xsize, int ysize, int XWIN, int YWIN, int D)
{
  bzero(xmap,sizeof(int)*xsize*ysize);
  bzero(ymap,sizeof(int)*xsize*ysize);
  
  // compute disparity
  for (int y=(YWIN+D);y<(ysize-YWIN-D);y++) {
    for (int x=(XWIN+D);x<(xsize-XWIN-D);x++) {
      // minimize SSD of pixels windowed around (x,y) against translated pixels
      unsigned int SmallestSum = 0xFFFFFFFF;
      int SmallestDX, SmallestDY;
      for (int dy=-D;dy<=D;dy++) {
	for (int dx=-D;dx<=D;dx++) {
	  unsigned int thesum=0x0;
	  // window around pixel at (x,y)
	  for (int wy=-YWIN;wy<=YWIN;wy++) {
	    for (int wx=-XWIN;wx<=XWIN;wx++) {
	      int val=(int)(im1[(y+wy+dy)*xsize+(x+wx+dx)]-
			    im2[(y+wy)*xsize+(x+wx)]);
	      thesum+=(unsigned int)(val*val);
	    }
	  }
	  if (thesum<=SmallestSum) {
	    SmallestSum=thesum;
	    SmallestDX=-dx;  
	    SmallestDY=-dy;
	  }
	}
      }
      // update disparity vector
      xmap[y*xsize+x]=SmallestDX;
      ymap[y*xsize+x]=SmallestDY;
    }
    if (y % 16 == 0)
    cout << "line "<<y <<endl;
  }
}


// find_similar_blobs
//
// scans the (xmap,ymap) disparity vector map for any vectors which fall
// withing thresh radius of (dx,dy)
// if below=1, finds vectors < thresh, if below=0, finds above thresh.
void find_similar_blobs(int *xmap, int *ymap, int *nxmap, int *nymap,
			int xsize, int ysize,
			int dx, int dy, int thresh, int below) {
  bzero(nxmap,sizeof(int)*xsize*ysize); bzero(nymap,sizeof(int)*xsize*ysize);
  if (below==1)
    for (int i=0;i<xsize*ysize;i++) {
      if (hypot(dx-xmap[i],dy-ymap[i])<=thresh) {
	nxmap[i]=xmap[i]; 
	nymap[i]=ymap[i];
      }
    }
  if (below==0)
    for (int i=0;i<xsize*ysize;i++) {
      if (hypot(dx-xmap[i],dy-ymap[i])>thresh) {
	nxmap[i]=xmap[i]; 
	nymap[i]=ymap[i];
      } 
    }

}


// mask_under_thresh
//
// masks off sections of image im where disparity (xmap,ymap) is below thresh
void mask_under_thresh(float *im, int *xmap, int *ymap, int xsize, int ysize, 
		 int thresh) {
  for (int i=0;i<xsize*ysize;i++) {
    if (hypot(xmap[i],ymap[i])<=thresh) im[i]=0;
  }
}

// mask_over_thresh
//
// masks off sections of image im where disparity (xmap,ymap) is above thresh
void mask_over_thresh(float *im, int *xmap, int *ymap, int xsize, int ysize, 
		 int thresh) {
  for (int i=0;i<xsize*ysize;i++) {
    if (hypot(xmap[i],ymap[i])>thresh) im[i]=0;
  }
}



// save_disp
//
// converts a disparity file to a pair of pgms and saves them
void save_disp(char *filename, int *dx, int *dy, int x, int y) {
  float *pgmx=make_image(x,y);
  float *pgmy=make_image(x,y);
  for (int i=0;i<x*y;i++) {
    pgmx[i]=128+dx[i];
    pgmy[i]=128+dy[i];
  }
  char *file1=new char[strlen(filename)+8];
  bzero(file1,strlen(filename)+8);
  sprintf(file1,"disp-x-%s",filename);
  write_pgm_image(file1,pgmx,x,y,"disparity map x");
  sprintf(file1,"disp-y-%s",filename);
  write_pgm_image(file1,pgmy,x,y,"disparity map y");
  free_image(pgmx); free_image(pgmy);
  delete(file1);
}

// load_disp
//
// loads disparity maps from a pair of pgm files
void load_disp(char *filename, int *dx, int *dy) {
  int X,Y;
  char *file = new char[strlen(filename)+8];
  bzero(file,strlen(filename)+8);
  sprintf(file,"disp-x-%s",filename);
  float *pgm=read_pgm_image(file,&X,&Y);
  for (int i=0;i<X*Y;i++) dx[i]=(int)(pgm[i]-128);
  free_image(pgm);
  sprintf(file,"disp-y-%s",filename);
  pgm=read_pgm_image(file,&X,&Y);
  for (int i=0;i<X*Y;i++) dy[i]=(int)(pgm[i]-128);
  free_image(pgm); delete(file);
}


// build group 
//
// starts at pixel (x,y) and isolates all adjoined pixels.  will scan for
// neighbouring pixels within a given square radius.
// returns mask in mask, pixel count is returned by function.
// note: destroys dx,dy in the process.  ensure maskx, masky are empty to start
//
// recursively fill in attached pixels.  mark used pixels by setting mask[]=col

int build_group(int *dx, int *dy, int *mask, int sizex, int
		sizey, int x, int y, int thresh, int col) {
  if ((x<0) || (x>=sizex) || (y<0) || (y>=sizey)) return 0;
  //  if ((mask[y*sizex+x]==col)) return 0;
  if ((dx[y*sizex+x]==0) && (dy[y*sizex+x]==0)) return 0;
  int count=1;
  mask[y*sizex+x]=col;
  dx[y*sizex+x]=0; dy[y*sizex+x]=0;
  for (int i=-thresh;i<=thresh;i++) 
    for (int j=-thresh;j<=thresh;j++) 
      count+=build_group(dx,dy,mask,sizex,sizey,x+j,y+i,thresh,col);
  return count;
}


// compute_xy_histogram
//
// counts up the number of pixels in each row & column
void compute_xy_histogram(int *mask, int *xhist, int *yhist, 
			  int xsize, int ysize) {
  for (int y=0;y<ysize;y++)
    for (int x=0;x<xsize;x++) {
      xhist[x]+=(mask[y*xsize+x]?1:0);
      yhist[y]+=(mask[y*xsize+x]?1:0);
    }
}


// find_bounding_rect
//
// finds a bounding rectangle for xy pixel histogram.
// bounds returned surround P percent of the pixels in the histogram.
// returns 100P% bounds and centroid.
// trims histograms to 80%, sorry!
void find_bounding_rect(int *xhist, int *yhist, int xsize, int ysize, float P,
			int &x1, int &x2, int &y1, int &y2, int &cx, int &cy) {
  
  int ytotal=0;
  int ycount=0;
  int ycentre;
  for (int y=0;y<ysize;y++) {  
    ycount+=yhist[y]; ytotal+=yhist[y]*(y+1);
  }
  ycentre=ytotal/ycount-1;

  int xtotal=0;
  int xcount=0;
  int xcentre;
  for (int x=0;x<xsize;x++) {
    xcount+=xhist[x]; xtotal+=xhist[x]*(x+1);
  }
  xcentre=xtotal/xcount-1;

  // now find P*100 percent of the pixels around the centre
  int Pxcount=(int)(xcount*P);
  int Pycount=(int)(ycount*P);

  // find x,y bounds
  int xmin,xmax,ymin,ymax;
  for (int x=0;x<xsize;x++) 
    if (xhist[x]) { xmin=x; break; }
  for (int x=xsize-1;x>=0;x--) 
    if (xhist[x]) { xmax=x; break; }
  // move x bounds inward until Pxcount pixels remain
  for (int count=xcount;;) {
    xhist[xmin]--;  
    if (xhist[xmin]==0) xmin++;
    count--; if (count<=Pxcount) break;
    xhist[xmax]--;
    if (xhist[xmax]==0) xmax--;
    count--; if (count<=Pxcount) break;
  }
  // xmin,xmax are now the 100P% bounds on the blob

  for (int y=0;y<ysize;y++) 
    if (yhist[y]) { ymin=y; break; }
  for (int y=ysize-1;y>=0;y--) 
    if (yhist[y]) { ymax=y; break; }
  // move y bounds inward until Pycount pixels remain
  for (int count=ycount;;) {
    yhist[ymin]--;  
    if (yhist[ymin]==0) ymin++;
    count--; if (count<=Pycount) break;
    yhist[ymax]--;
    if (yhist[ymax]==0) ymax--;
    count--; if (count<=Pycount) break;
  }
  // ymin,ymax are now the 100P% bounds on the blob
  x1=xmin; x2=xmax; y1=ymin; y2=ymax; cx=xcentre; cy=ycentre;
}