#define D_IMAGEFILE
#define D_MATRIX
#define D_IMAGE
#define D_MATUTIL

//#include <iglooprelude.h>
#include "bookstein.h"

// Program : warpimage
//
// This program uses the class bookstein to warp an image
//
// Aziz Boxwala
// Dept of Biomedical Engg
// University of North Carolina at Chapel Hill
//
// 29th Nov 1993
//

main(int argc, char* argv[])
{
  BEGIN;
  if (argc<3)
    Diagnostic(FATAL, "warpimage <pt_image> <input image> <out image>");
  
  imagefile* imf = new imagefile(argv[1]);
  image pt_im1 = imf->load(0);
  image pt_im2 = imf->load(1);
  delete imf;
  
  imf = new imagefile(argv[2]);
  image in_im = imf->load();
  delete imf;
  
  bookstein bkstn(pt_im1, pt_im2, in_im);
  image* out_im = bkstn.warp_im();
   imf=new imagefile(argv[3],in_im.shape(),1, in_im.type());
  imf->save(*out_im);
  END;
  
}


/***************************************************************************
*
* Class name : bookstein
*
* Purpose    : For warping images using Fred Bookstein's algorithm of
*              thin plate spline interpolation (IPMI '87).
*
* Author     : Aziz Boxwala
*              Dept of Biomedical Engineering
*              The Univeristy of North Carolina at Chapel Hill
*
* Date       : 17th November 1993
*
****************************************************************************/

inline double U(int x1, int y1, int x2, int y2)
{

  double r_squared = 0.0;
  r_squared = (double)(((x1-x2)*(x1-x2))+((y1-y2)*(y1-y2)));
    if (r_squared == 0.0)
      return 0.0;
    else
      return((r_squared)*log(sqrt(r_squared)));
}

//const int &MAX_NUM_POINTS=500;

#ifdef NOTDEFINED
class bookstein
{
  
private:
  subscript dims;
  
  matrix points1;
  matrix points2;

  matrix *L, *Y, *W, *V, *K, *P;
  image map_from_image, map_to_image;
  image input_image, output_image;
  int num_cntl_points;

  // From the two images find the co-ords of the control points
  void read_points();

  // initialize all the matrices
  void init_mats();

  // Calculate the wieghts/parameters matrix
  void calc_w();

  // Computes the distance matrix K from the points1 matrix
  void dist_matx();
  
  // Fill in the other matrices
  void fill_mats();

public:
  
  // Constructor
  bookstein(const image& map_from, const image& map_to, const image& image_to_warp);

  // Calculate the deformed image
  image* warp_im();

  // Destructor
  ~bookstein()
    {
      // Destroy all the matrices and stuff
      // MSG("Destroying class bookstein");
    }
};
#endif 

bookstein::bookstein(const image& map_from, const image& map_to, const image& image_to_warp)
{
  map_from_image = map_from;
  map_to_image   = map_to;
  dims = map_from.shape();
  if ( (dims.r() != map_to.shape().r()) || (dims.c() != map_to.shape().c()))
    Diagnostic(FATAL, "bookstein:: Co-ordinate images dimensions unequal");
  
  if((image_to_warp.shape().r() !=dims.r())||(image_to_warp.shape().c() !=dims.c()))
    Diagnostic(FATAL,
	       "bookstein:: Input image dimensions not equal to co-ordinate images dimensions");
  input_image = image_to_warp;

  matrix foo(MAX_NUM_POINTS,2);
  points1.replace(foo);
  points2.replace(foo);
  read_points();
  init_mats();
  dist_matx();
  fill_mats();
  calc_w();
}


// From the two images find the co-ords of the control points
void bookstein::read_points()
{
  MSG("Reading points");
  subscript s(map_from_image.shape());
  int num_points1=0, num_points2=0, point_cntr=0;
  int index;

  // Read the control points from the first image
  for (s.init(); s.test(); s++){
    if ((index = map_from_image.integer(s) ) != 0){
      if (index > (MAX_NUM_POINTS -1))
	Diagnostic(FATAL, "bookstein:: Too many control points????");
      points1.element(index-1, 0) = s.c();
      points1.element(index-1, 1) = s.r();
      // printf("Point %d = ", index);
      // s.print(); NL;
      num_points1++;
      if (index > point_cntr)
	point_cntr = index;
    }
  }
  
  if (point_cntr != num_points1)
	Diagnostic(FATAL,"bookstein:: Points in map_from_image corrupt");
  
  point_cntr = 0;
  // Read the control points for the second image
  for (s.init(); s.test(); s++){
    if ((index = map_to_image.integer(s) ) != 0){
      if (index > (MAX_NUM_POINTS -1))
	Diagnostic(FATAL, "bookstein:: Too many control points????");
      points2.element(index-1, 0) = s.c();
      points2.element(index-1, 1) = s.r();
      // printf("Point %d = ", index);
      // s.print(); NL;
      num_points2++;
      if (index > point_cntr)
	point_cntr = index;
    }
  }
  
  if (point_cntr != num_points2)
    Diagnostic(FATAL,"bookstein:: Points in map_to_image corrupt");
  
  if (num_points1 != num_points2)
    Diagnostic(FATAL,
	       "bookstein:: Points in map_from_image and map_to_image don't match");
  num_cntl_points = num_points2;
  if (!(num_cntl_points))
    Diagnostic(FATAL, "bookstein:: No control points specified");
  printf("Number of points = %d\n", num_cntl_points);
}

// Computes the distance matrix K from the points in the first image
void bookstein::dist_matx()
{
  MSG("Computing distance matrix");
  subscript s(0,0, K->shape());

  for (s.init(); s.test(); s++){
    K->element(s.r(), s.c()) = U((int)points1(s.r(),0),(int)points1(s.r(),1),
				 (int)points1(s.c(),0),(int)points1(s.c(),1));
  }
  MSG("**K**");
  K->print();
}

// Initialize all the matrices
void bookstein::init_mats()
{
  MSG("Init matrices");
  P = new matrix(num_cntl_points, 2);
  K = new matrix(num_cntl_points,num_cntl_points);
  L = new matrix(num_cntl_points+3, num_cntl_points+3);
  Y = new matrix(num_cntl_points+3, 2);
  W = new matrix(num_cntl_points+3, 2);
}

// Fill in the L and Y matrices
void bookstein::fill_mats()
{
  // Compute the distance matrix K from the points in the first image
  subscript s(0,0, K->shape());
  for(s.init(); s.test(); s++)
  {
       K->element(s.r(), s.c()) =
          U((int)points1(s.r(),0),(int)points1(s.r(),1),
            (int)points1(s.c(),0),(int)points1(s.c(),1));
  }

  // The L matrix
  for (s.init(); s.test(); s++)
    L->element(s.r(), s.c())=K->element(s.r(), s.c());

  for (int i=0; i<num_cntl_points; i++)
  {
     L->element(i, num_cntl_points)=1.0;
     L->element(num_cntl_points, i)=1.0;
   }

  subscript t(0,0, subscript(num_cntl_points, 2));
  for (t.init(); t.test();t++){
    L->element(t.r(), num_cntl_points+1+t.c())=points1(t.r(), t.c());
    L->element(num_cntl_points+1+t.c(),t.r())=points1(t.r(), t.c());
  }

  MSG("**L**"); L->print();
  // The Y Matrix
  for (t.init(); t.test();t++)
    Y->element(t.r(), t.c()) = points2(t.r(), t.c());
  MSG("**Y**"); Y->print();
}
  

// Calculate the weight matrix
void bookstein::calc_w()
{
  MSG("Calculating weights");
  matutil utility_mat;
  utility_mat.use(*L);
  sq_matrix inverseL = utility_mat.inverse();
  //inverseL.print();
  W->replace(inverseL*(*Y));
  MSG("**W**");
  W->print();
  // utility_mat.use(inverseL);
  // utility_mat.inverse().print();
}

// Use the weights to apply warp to image
image* bookstein::warp_im()
{
  MSG("Applying warp to image");
  image *outimage = new image(dims, input_image.type());
  // gpoint location;
  matrix a1(2,1), ax(2,1), ay(2,1), distSum(2,1), f(2,1);
  a1.element(0,0) = W->element(num_cntl_points,0);
  a1.element(1,0) = W->element(num_cntl_points,1);
  ax.element(0,0) = W->element(num_cntl_points+1,0);
  ax.element(1,0) = W->element(num_cntl_points+1,1);
  ay.element(0,0) = W->element(num_cntl_points+2,0);
  ay.element(1,0) = W->element(num_cntl_points+2,1);
  subscript s(input_image.shape());
  double u;

  // RTC;
  for (s.init(); s.test(); s++){
    distSum.init(0.0);
    for(int t=0; t<num_cntl_points; t++){
      u = U((int)points1(t,0), (int)points1(t,1), s.c(), s.r());
      distSum.element(0,0) += (W->element(t,0)*u);
      distSum.element(1,0) += (W->element(t,1)*u);
    }
    f=a1+(ax*s.c())+(ay*s.r())+distSum;
    // f.print();
    if (f(0,0)<0.0)
      f.element(0,0) = 0.0;
    else if (f(0,0)>=dims.c())
      f.element(0,0)=dims.c()-1;
    if (f(1,0)<0.0)
      f.element(1,0) = 0.0;
    else if (f(1,0)>=dims.r())
      f.element(1,0)=dims.r()-1;
    outimage->set(s, input_image((int)f(1,0), (int)f(0,0)));
  }
  return(outimage);
}