#include "moss.h"

int
_mossAddPoint(mossMatrix *matx,
	      int i, int j,
	      float x0, float y0,
	      float x1, float y1) {
  mossCorresp *pair;
  char me[]="_mossAddPoint", err[128];
  int n;
  
  pair = matx->pair[i][j];
  if (!pair) {
    /* this is the first correspondence we've seen for this image pair */
    pair = matx->pair[i][j] = matx->pair[j][i] = mossNewCorresp(MOSS_MAXPAIRS);
    if (!pair) {
      sprintf(err, "%s: trouble creating mossCorresp for i,j=%d,%d",
	      me, i, j);
      biffAdd(MOSS, err);
      return 1;
    }
  }

  n = pair->num;
  if (n == pair->maxNum) {
    sprintf(err, "%s: reached max # points (%d) in mossCorresp", 
	    me, pair->maxNum);
    biffAdd(MOSS, err);
    return 1;
  }

  pair->x0[n] = x0;
  pair->y0[n] = y0;
  pair->x1[n] = x1;
  pair->y1[n] = y1;
  pair->num++;

  return 0;
}

void
_mossPrintMatxA(mossMatrix *matx) {
  char me[]="_mossPrintMatxA";
  int i, j, p;

  for (i=0; i<=matx->size-1; i++) {
    for (j=i; j<=matx->size-1; j++) {
      if (matx->pair[i][j]) {
	printf("%s: points for i,j = %d,%d:\n", me, i, j);
	for (p=0; p<=matx->pair[i][j]->num-1; p++) {
	  printf("      (% 7.3f,% 7.3f) <---> (% 7.3f,% 7.3f)\n",
		 matx->pair[i][j]->x0[p], matx->pair[i][j]->y0[p], 
		 matx->pair[i][j]->x1[p], matx->pair[i][j]->y1[p]);
	}
      }
    }
  }
}

void
_mossPrintMatxB(mossMatrix *matx) {
  char me[]="_mossPrintMatxB";
  int i, j;
  
  for (i=0; i<=matx->size-1; i++) {
    for (j=i+1; j<=matx->size-1; j++) {
      if (matx->pair[i][j]) {
	printf("%s: m01 for i,j = %d,%d:\n", me, i, j);
	printf("       % 7.3f % 7.3f % 7.3f\n",
	       matx->pair[i][j]->m01[0][0],
	       matx->pair[i][j]->m01[0][1],
	       matx->pair[i][j]->m01[0][2]);
	printf("       % 7.3f % 7.3f % 7.3f\n",
	       matx->pair[i][j]->m01[1][0],
	       matx->pair[i][j]->m01[1][1],
	       matx->pair[i][j]->m01[1][2]);
	printf("       % 7.3f % 7.3f % 7.3f\n",
	       matx->pair[i][j]->m01[2][0],
	       matx->pair[i][j]->m01[2][1],
	       matx->pair[i][j]->m01[2][2]);
	printf("         m10:\n");
	printf("       % 7.3f % 7.3f % 7.3f\n",
	       matx->pair[i][j]->m10[0][0],
	       matx->pair[i][j]->m10[0][1],
	       matx->pair[i][j]->m10[0][2]);
	printf("       % 7.3f % 7.3f % 7.3f\n",
	       matx->pair[i][j]->m10[1][0],
	       matx->pair[i][j]->m10[1][1],
	       matx->pair[i][j]->m10[1][2]);
	printf("       % 7.3f % 7.3f % 7.3f\n",
	       matx->pair[i][j]->m10[2][0],
	       matx->pair[i][j]->m10[2][1],
	       matx->pair[i][j]->m10[2][2]);
      }
    }
  }
}

int
_mossLearnPoints(mossMatrix *matx, Nrrd *points) {
  char me[]="_mossLearnPoints", err[128];
  int numImgs, pts, i, j, k, two, one;
  float *data, *line;
  
  if (matx->size != points->size[1]) {
    sprintf(err, "%s: size mismatch between matrix (%d) and points (%d)",
	    me, matx->size, points->size[1]);
    biffAdd(MOSS, err);
    return 1;
  }
  numImgs = matx->size;
  data = (float*)points->data;
  for (pts=0; pts<=points->size[2]-1; pts++) {
    line = data + 2*numImgs*pts;
    two = 0;
    for (k=0; k<=numImgs-1; k++) {
      if (-1 != line[0 + 2*k]) {
	if (0 == two) {
	  i = k;
	}
	else if (1 == two) {
	  j = k;
	}
	two++;
      }
    }
    if (two != 2) {
      sprintf(err, "%s: didn't get point pair in line#%d of correspondence",
	      me, pts);
      biffAdd(MOSS, err);
      return 1;
    }
    if (_mossAddPoint(matx, i, j, 
		      line[0 + 2*i], line[1 + 2*i], 
		      line[0 + 2*j], line[1 + 2*j])) {
      sprintf(err, "%s: trouble adding points in line#%d of correspondence",
	      me, pts);
      biffAdd(MOSS, err);
      return 1;
    }
  }

  /* are there any images for which there is no correspondence info
     attaching them to another image? */
  for (i=0; i<=matx->size-1; i++) {
    one = 0;
    for (j=0; j<=matx->size-1; j++) {
      if (matx->pair[i][j])
	one++;
    }
    if (!one) {
      sprintf(err, "%s: image#%d doesn't correspond to any other", me, i);
      biffAdd(MOSS, err);
      return 1;
    }
  }
  return 0;
}

float
_mossDot(float *a, float *b, int n) {
  int i;
  float ret;
  
  ret = 0;
  for (i=0; i<=n-1; i++) {
    ret += a[i]*b[i];
  }
  return ret;
}

int
_moss3x3Invert(float i[3][3], float m[3][3]) {
  char me[]="_moss3x3Invert", err[128];
  float det;

  det = (+m[0][0]*m[1][1]*m[2][2]
	 +m[1][0]*m[2][1]*m[0][2]
	 +m[2][0]*m[0][1]*m[1][2]
	 -m[0][2]*m[1][1]*m[2][0]
	 -m[1][2]*m[2][1]*m[0][0]
	 -m[2][2]*m[0][1]*m[1][0]);
  if (fabs(det) < MOSS_TINYVAL) {
    sprintf(err, "%s: matrix has tiny determinant (%g)\n", me, det);
    biffAdd(MOSS, err); return 1;
  }
  i[0][0] = (m[1][1]*m[2][2] - m[1][2]*m[2][1])/det;
  i[1][0] = (m[2][0]*m[1][2] - m[2][2]*m[1][0])/det;
  i[2][0] = (m[1][0]*m[2][1] - m[1][1]*m[2][0])/det;
  i[0][1] = (m[2][1]*m[0][2] - m[2][2]*m[0][1])/det;
  i[1][1] = (m[0][0]*m[2][2] - m[0][2]*m[2][0])/det;
  i[2][1] = (m[2][0]*m[0][1] - m[2][1]*m[0][0])/det;
  i[0][2] = (m[0][1]*m[1][2] - m[0][2]*m[1][1])/det;
  i[1][2] = (m[1][0]*m[0][2] - m[1][2]*m[0][0])/det;
  i[2][2] = (m[0][0]*m[1][1] - m[0][1]*m[1][0])/det;
  return 0;
}

void
_moss3x3Copy(float a[3][3], float b[3][3]) {
  
  a[0][0] = b[0][0]; a[0][1] = b[0][1]; a[0][2] = b[0][2];   
  a[1][0] = b[1][0]; a[1][1] = b[1][1]; a[1][2] = b[1][2];   
  a[2][0] = b[2][0]; a[2][1] = b[2][1]; a[2][2] = b[2][2];   
}

void
_moss3x3Mult(float a[3][3], float b[3][3], float c[3][3]) {
  
  a[0][0] = b[0][0]*c[0][0] + b[0][1]*c[1][0] + b[0][2]*c[2][0];
  a[0][1] = b[0][0]*c[0][1] + b[0][1]*c[1][1] + b[0][2]*c[2][1];
  a[0][2] = b[0][0]*c[0][2] + b[0][1]*c[1][2] + b[0][2]*c[2][2];
  a[1][0] = b[1][0]*c[0][0] + b[1][1]*c[1][0] + b[1][2]*c[2][0];
  a[1][1] = b[1][0]*c[0][1] + b[1][1]*c[1][1] + b[1][2]*c[2][1];
  a[1][2] = b[1][0]*c[0][2] + b[1][1]*c[1][2] + b[1][2]*c[2][2];
  a[2][0] = b[2][0]*c[0][0] + b[2][1]*c[1][0] + b[2][2]*c[2][0];
  a[2][1] = b[2][0]*c[0][1] + b[2][1]*c[1][1] + b[2][2]*c[2][1];
  a[2][2] = b[2][0]*c[0][2] + b[2][1]*c[1][2] + b[2][2]*c[2][2];
}

int
_mossCalcPTs(mossMatrix *mmatx) {
  char me[]="_mossCalcPTs", err[128];
  int i, j, k, pts;
  mossCorresp *pair;
  float *matx, *inv, *vec;
  
  for (i=0; i<=mmatx->size-1; i++) {
    for (j=i; j<=mmatx->size-1; j++) {
      pair = mmatx->pair[i][j];
      if (pair) {
	if (pair->num < 4) {
	  sprintf(err, "%s: have only %d point pairs for i,j=%d,%d; need 4", 
		  me, pair->num, i, j);
	  biffAdd(MOSS, err);
	  return 1;
	}
	pts = pair->num;
	matx = (float*)malloc(2*pts * 8 * sizeof(float));
	inv = (float*)malloc(8 * 2*pts * sizeof(float));
	vec = (float*)malloc(2*pts * sizeof(float));
	if (!(matx && inv && vec)) {
	  sprintf(err, "%s: couldn't alloc matrix arrays for i,j=%d,%d SVD", 
		  me, i, j);
	  biffAdd(MOSS, err);
	  return 1;
	}
	for (k=0; k<=pts-1; k++) {
	  matx[0 + 8*(0 + 2*k)] = pair->x0[k];
	  matx[1 + 8*(0 + 2*k)] = pair->y0[k];
	  matx[2 + 8*(0 + 2*k)] = 1.0;
	  matx[3 + 8*(0 + 2*k)] = 0;
	  matx[4 + 8*(0 + 2*k)] = 0;
	  matx[5 + 8*(0 + 2*k)] = 0;
	  matx[6 + 8*(0 + 2*k)] = -pair->x0[k]*pair->x1[k];
	  matx[7 + 8*(0 + 2*k)] = -pair->y0[k]*pair->x1[k];
	  matx[0 + 8*(1 + 2*k)] = 0;
	  matx[1 + 8*(1 + 2*k)] = 0;
	  matx[2 + 8*(1 + 2*k)] = 0;
	  matx[3 + 8*(1 + 2*k)] = pair->x0[k];
	  matx[4 + 8*(1 + 2*k)] = pair->y0[k];;
	  matx[5 + 8*(1 + 2*k)] = 1.0;
	  matx[6 + 8*(1 + 2*k)] = -pair->x0[k]*pair->y1[k];
	  matx[7 + 8*(1 + 2*k)] = -pair->y0[k]*pair->y1[k];
	  vec[0 + 2*k] = pair->x1[k];
	  vec[1 + 2*k] = pair->y1[k];
	}
	if (mossPseudoInverse(inv, matx, 2*pts, 8)) {
	  sprintf(err, "%s: trouble with pseudo inverse for i,j=%d,%d",
		  me, i, j);
	  biffAdd(MOSS, err);
	  return 1;
	}
	pair->m01[0][0] = _mossDot(inv + 0*2*pts, vec, 2*pts);
	pair->m01[0][1] = _mossDot(inv + 1*2*pts, vec, 2*pts);
	pair->m01[0][2] = _mossDot(inv + 2*2*pts, vec, 2*pts);
	pair->m01[1][0] = _mossDot(inv + 3*2*pts, vec, 2*pts);
	pair->m01[1][1] = _mossDot(inv + 4*2*pts, vec, 2*pts);
	pair->m01[1][2] = _mossDot(inv + 5*2*pts, vec, 2*pts);
	pair->m01[2][0] = _mossDot(inv + 6*2*pts, vec, 2*pts);
	pair->m01[2][1] = _mossDot(inv + 7*2*pts, vec, 2*pts);
	pair->m01[2][2] = 1.0;
	if (_moss3x3Invert(pair->m10, pair->m01)) {
	  sprintf(err, "%s: can't invert projective transform for i,j=%d,%d",
		  me, i, j);
	  biffAdd(MOSS, err);
	  return 1;
	}
	free(matx);
	free(inv);
	free(vec);
      }
    }
  }

  return 0;
}

/*
** _mossFindPathTo()
**
** This is a very very poorly designed and written procedure to 
** find the shortest path in the adjency graph between two given
** images.  But it works.
*/
int
_mossFindPathTo(mossMatrix *mmatx, int *path, int idx, int dest) {
  int i, j, sp, skip, here, steps, minSteps, minI;
  
  here = path[idx];

  /* 1st base case: we've run out of time */
  if (idx == mmatx->size-1) {
    return 0;
  }

  /* else we look for leads */
  minSteps = 100000;
  minI = -1;
  for (i=0; i<=mmatx->size-1; i++) {
    /* we can't look at ourself */
    if (i == here)
      continue;

    /* a lead starts with an existing pair */
    if (!mmatx->pair[here][i])
      continue;
    skip = 0;
    for (j=0; j<=idx; j++) {
      if (path[j] == i) {
	skip = 1;
	break;
      }
    }
    /* but we can't look at something already on the path */
    if (skip)
      continue;

    /* 2nd base case: we can finish the path with just one more step */
    if (i == dest) {
      return 1;
    }
    path[idx+1] = i;
    steps = _mossFindPathTo(mmatx, path, idx+1, dest);
    if (steps) {
      if (steps < minSteps) {
	minSteps = steps;
	minI = i;
      }
    }
  }

  /* did any leads pan out? */
  if (-1 == minI) {
    return 0;
  }

  /* we've found the next step */
  path[idx+1] = minI;
  return 1 + _mossFindPathTo(mmatx, path, idx+1, dest);

  return 0;
}


void
_mossCalculateInter(mossMatrix *mmatx, int *path, int steps, 
		    mossCorresp *pair) {
  float im[3][3], tmp[3][3];
  int i, j, t, s;

  im[0][0] = 1.0; im[0][1] = 0.0; im[0][2] = 0.0; 
  im[1][0] = 0.0; im[1][1] = 1.0; im[1][2] = 0.0; 
  im[2][0] = 0.0; im[2][1] = 0.0; im[2][2] = 1.0; 
  for (s=0; s<=steps-1; s++) {
    i = path[s];
    j = path[s+1];
    if (i < j)
      _moss3x3Mult(tmp, mmatx->pair[i][j]->m01, im);
    else
      _moss3x3Mult(tmp, mmatx->pair[i][j]->m10, im);
    _moss3x3Copy(im, tmp);
  }
  if (path[0] < path[steps]) {
    _moss3x3Copy(pair->m01, im);
    _moss3x3Invert(pair->m10, pair->m01);
  }
  else {
    _moss3x3Copy(pair->m10, im);
    _moss3x3Invert(pair->m01, pair->m10);
  }
}

int
_mossCompletePTs(mossMatrix *mmatx) {
  char me[]="_mossCompletePTs", err[128];
  mossCorresp *pair;
  int i, j, k, *path, steps;

  path = (int*)malloc(mmatx->size*sizeof(int));
  for (i=0; i<=mmatx->size-1; i++) {

    /* we put identity transforms along the diagonal */
    if (!(pair = mmatx->pair[i][i] = mossNewCorresp(0))) {
      sprintf(err, "%s: couldn't make new mossCorresp for i,j=%d,%d",
	      me, i, i);
      biffAdd(MOSS, err); return 1;
    }
    pair->m01[0][0] = 1.0; pair->m01[0][1] = 0.0; pair->m01[0][2] = 0.0; 
    pair->m01[1][0] = 0.0; pair->m01[1][1] = 1.0; pair->m01[1][2] = 0.0; 
    pair->m01[2][0] = 0.0; pair->m01[2][1] = 0.0; pair->m01[2][2] = 1.0; 
    pair->m10[0][0] = 1.0; pair->m10[0][1] = 0.0; pair->m10[0][2] = 0.0; 
    pair->m10[1][0] = 0.0; pair->m10[1][1] = 1.0; pair->m10[1][2] = 0.0; 
    pair->m10[2][0] = 0.0; pair->m10[2][1] = 0.0; pair->m10[2][2] = 1.0; 

    /* on the off-diagonal, we may have to do some work */
    for (j=i+1; j<=mmatx->size-1; j++) {
      printf("%s: i,j = %d,%d\n", me, i, j);
      if (pair = mmatx->pair[i][j]) {
	printf("%s:    (already have pair[%d][%d])\n", me, i, j);
	continue;
      }
      /* we have to figure out from (adjacency graph) neighbors
	 how to do this projective transform */
      path[0] = i;
      if (!(steps = _mossFindPathTo(mmatx, path, 0, j))) {
	sprintf(err, "%s: can't find a path from %d to %d!", me, i, j);
	biffAdd(MOSS, err); return 1;
      }
      if (!(steps >= 2 && steps <= mmatx->size-1)) {
	sprintf(err, "%s: BUG: path was length %d, not in range [%d,%d]",
		me, steps, 2, mmatx->size-1);
	biffAdd(MOSS, err); return 1;
      }
      path[steps] = j;
      printf("%s: (first) shortest path from %d to %d:\n", me, i, j);
      for (k=0; k<=steps; k++) {
	printf("     %d: %d\n", k, path[k]);
      }
      if (!(pair = mmatx->pair[i][j] = 
	    mmatx->pair[j][i] = mossNewCorresp(0))) {
	sprintf(err, "%s: couldn't make new mossCorresp for i,j=%d,%d",
		me, i, j);
	biffAdd(MOSS, err); return 1;
      }
      _mossCalculateInter(mmatx, path, steps, pair);
    }
  }
  free(path);
  
  return 0;
}

/*
******** mossSetMatrix()
**
** converts the user-specified correspondance point pairs into the
** matrix of projective transforms.  This function name is perhaps
** a little confusing
**
** returns 0 if all okay, 1 on error
*/
int
mossSetMatrix(mossMatrix *matx, Nrrd *points) {
  char me[]="mossSetMatrix", err[128];

  if (!(matx && points)) {
    sprintf(err, BIFF_NULL, me);
    biffAdd(MOSS, err);
    return 1;
  }

  if (_mossLearnPoints(matx, points)) {
    sprintf(err, "%s: trouble with correspondence information", me);
    biffAdd(MOSS, err);
    return 1;
  }

  _mossPrintMatxA(matx);

  if (_mossCalcPTs(matx)) {
    sprintf(err, "%s: trouble calculating corresp. projective transforms", me);
    biffAdd(MOSS, err);
    return 1;
  }

  _mossPrintMatxB(matx);

  if (_mossCompletePTs(matx)) {
    sprintf(err, "%s: trouble filling out all projective transforms", me);
    biffAdd(MOSS, err);
    return 1;
  }
  
  _mossPrintMatxB(matx);

  return 0;
}