#include "my_blitz.hh"
#include <iostream>
#include "BlitzHDF5Helper.hh"
#include <blitz/array.h>
#include <blitz/tinyvec-et.h>		// uncomment when using blitz++ version newer than 0.9
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>

/*
 * Function that generates red - green overlay of 2 images of possible different sizes
 *
 */


blitz::Array<blitz::TinyVector<float,3>,3 > overlay(	const blitz::Array<float,3>& fixedImage, 
							const blitz::Array<float,3>& movingImage){
  blitz::TinyVector<int,3> fixedShape=fixedImage.shape();
  blitz::TinyVector<int,3> movingShape=movingImage.shape();
  blitz::TinyVector<int,3> unionShape=blitz::max(fixedShape, movingShape);
  blitz::Array<blitz::TinyVector<float,3>,3 > rgb(unionShape);
  rgb=0.f;
  rgb(blitz::Range(0, fixedShape(0)-1), blitz::Range(0, fixedShape(1)-1), blitz::Range(0, fixedShape(2)-1))[0]=fixedImage;
  rgb(blitz::Range(0, movingShape(0)-1), blitz::Range(0, movingShape(1)-1), blitz::Range(0, movingShape(2)-1))[1]=movingImage;
  return rgb;
}

/*
 * Marks the landmarks in the image
 *
 */

void markLandmark(	blitz::Array<float,3>& Image, 
			const blitz::Array<blitz::TinyVector<double,3>,1>& landmarks, 
			blitz::TinyVector<float,3> element_size_um){
  int lev_l, lev_h, row_l, row_h, col_l, col_h;
  blitz::TinyVector<int,3> radius;
  radius=10, 10, 10;
  radius=radius/element_size_um;
  for(int l=0;l<landmarks.extent(0);++l){
    blitz::TinyVector<float,3> tempMark=landmarks(l)/element_size_um;
    lev_l=tempMark(0)-radius(0);
    lev_h=tempMark(0)+radius(0);
    row_l=tempMark(1)-radius(1);
    row_h=tempMark(1)+radius(1);
    col_l=tempMark(2)-radius(2);
    col_h=tempMark(2)+radius(2);
    Image(blitz::Range(lev_l, lev_h), blitz::Range(row_l, row_h), blitz::Range(col_l, col_h))+=400;
  }

}



/*======================================================================*/
/*! 
 *   Nearest neighbour interpolation with boundary check. For positions
 *   outside the array gray value 0 is returned
 *
 *   \param arr the array
 *   \param pos float position
 *
 *   \exception  none
 *
 *   \return "interpolated" gray value at pos
 */
/*======================================================================*/
inline
float interpolNN( const blitz::Array<float, 3>& arr,
                          blitz::TinyVector<float,4> pos)
{
  // normalize homogeneous coordinates if necessary
  if(pos(3) != 1) pos /= pos(3);
  
  // find nearest neighbour
  blitz::TinyVector<int,4> ipos;
  ipos = blitz::floor( pos + 0.5);
  if( blitz::all(ipos >= 0) && blitz::all(ipos < arr.shape()))
  {
    return arr(ipos);
  }
  else
  {
    return 0;
  }
}
/*======================================================================*/
/*! 
 *   transform array with given 4x4 Matrix using homogeneous
 *   coordinates using nearest neighbour interpolation
 *
 *   \param srcArr  the source Array
 *   \param invMat  the inverse transformation matrix. Will be used to
 *                  project trg coordinates (homogeneous coordinates
 *                  in blitz-format, i.e. (lev,row,col,1) ) to src
 *                  coordinates.
 *   \param trgArr  the target Array -- must already have the wanted
 *                  size 
 *
 *   \exception none
 *
 *   \return none
 */
/*======================================================================*/
void transformArray( const blitz::Array<float, 3>& srcArr,
                     const blitz::TinyMatrix<float,4,4>& invMat,
                     blitz::Array<float, 3>& trgArr)
{
  blitz::TinyVector<float,4> srcPos;
  blitz::TinyVector<float,4> trgPos;
#pragma omp parallel for schedule(dynamic)
  for( int lev = 0; lev < trgArr.extent(0); ++lev)
  {
    for( int row = 0; row < trgArr.extent(1); ++row)
    {
      for( int col = 0; col < trgArr.extent(2); ++col)
      {
        trgPos = lev, row, col, 1;
        srcPos = myproduct( invMat, trgPos);
        trgArr(lev,row,col) = interpolNN( srcArr, srcPos);
      }
    }
  }
}

/*======================================================================*/
/*! 
 *   Wrapper for the Singular Value Decomposition (SVD) of the Gnu
 *   scientific library. Given A is split into U * S * V^T where S is
 *   a diagonal matrix, U and V are orthogonal matrices
 *
 *   \param A   input matrix A and output matrix U
 *   \param S   diagonal output matrix S (Vector with diagonal entries)
 *   \param V   output matrix V (attention: the output is V, not V transposed)
 *
 */
/*======================================================================*/
void mySVD( blitz::TinyMatrix<float,3,3>& A, 
           blitz::TinyVector<float,3>& S,
           blitz::TinyMatrix<float,3,3>& V)
{
  // allocate gsl matrices (attention: data type is double)
  gsl_matrix* gA = gsl_matrix_calloc( 3, 3);
  gsl_vector* gS = gsl_vector_calloc( 3);
  gsl_matrix* gV = gsl_matrix_calloc( 3, 3);
  gsl_vector* work = gsl_vector_calloc( 3);
  
  // copy input data to gsl data types
  for( int i = 0; i < 3; ++i)
  {
    for( int j = 0; j < 3; ++j)
    {
      gsl_matrix_set (gA, i, j, A(i,j));
    }
  }
  

  // perform the svd
  gsl_linalg_SV_decomp( gA, gV, gS, work);
  
  // copy results back to blitz data types
  for( int i = 0; i < 3; ++i)
  {
    for(int j = 0; j < 3; ++j)
    {
      A(i,j)  = gsl_matrix_get( gA, i, j);
      V(i,j) = gsl_matrix_get( gV, i, j);
    }
  }

  for( int i = 0; i < 3; ++i) 
      S(i) = gsl_vector_get( gS, i);
  
  // free all gsl stuff
  gsl_vector_free( work);
  gsl_matrix_free( gV);
  gsl_vector_free( gS);
  gsl_matrix_free( gA);
}




/*======================================================================*/
/*! 
 *   Umeyama method to find the transformation matrix for given
 *   (correspondin) point sets. The corresponding point sets have to
 *   be provided in homogeneous coordinates. The returned matrix
 *   fulfill transformedPoints(i) = H * points(i) for all i.
 *
 *   \param points            orignal point set 
 *   \pamar transformedPoints transformed point set
 *
 *   \return  4x4 Matrix H 
 */
/*======================================================================*/
blitz::TinyMatrix<float,4,4> umeyama( 
    const blitz::Array<blitz::TinyVector<float,4>,1>& points,
    const blitz::Array<blitz::TinyVector<float,4>,1>& transformedPoints)
{
  /*-----------------------------------------------------------------------
   *  Compute mean of point sets
   *-----------------------------------------------------------------------*/
  blitz::TinyVector<float,4> meanPoints;
  blitz::TinyVector<float,4> meanTransformedPoints;
  meanPoints = blitz::mean( points);
  meanTransformedPoints = blitz::mean( transformedPoints);
  

  /*-----------------------------------------------------------------------
   *  Compute covariance matrix (subtraction of mean is done on the fly)
   *-----------------------------------------------------------------------*/
  blitz::TinyMatrix<float,3,3> covar;
  covar = 0;
  
  for( int n = 0; n < points.extent(0); ++n)
  {
    for( int i = 0; i < 3; ++i)
    {
      for( int j = 0; j < 3; ++j)
      {
        covar(i,j) += 
            (transformedPoints(n)(i)-meanTransformedPoints(i)) 
            * (points(n)(j) - meanPoints(j));
      }
    }
  }
  std::cout << "covar = " << covar << std::endl;
 
  /*-----------------------------------------------------------------------
   *  perform the singular value decomposition and compute the rotation
   *  matrix
   *-----------------------------------------------------------------------*/
  blitz::TinyMatrix<float,3,3> U;
  blitz::TinyVector<float,3>   S;
  blitz::TinyMatrix<float,3,3> V;
  U = covar;
  mySVD( U, S, V);
  
  blitz::TinyMatrix<float,3,3> rot;
  rot = myproduct( U, mytranspose(V));

  std::cout << "U = " << U << std::endl
            << "S = " << S <<  std::endl
            << "V = " << V << std::endl
            << "rot = " << rot << std::endl;

  /*-----------------------------------------------------------------------
   *  compute the translation
   *-----------------------------------------------------------------------*/
   blitz::TinyMatrix<float,4,4> transformMat;
   
   transformMat = 
       rot(0,0), rot(0,1), rot(0,2), 0,
       rot(1,0), rot(1,1), rot(1,2), 0,
       rot(2,0), rot(2,1), rot(2,2), 0,
       0,        0,        0,        1;
   
   blitz::TinyVector<float,4> trans;
   
   trans = meanTransformedPoints - myproduct( transformMat, meanPoints);
  
   transformMat(0,3) = trans(0);
   transformMat(1,3) = trans(1);
   transformMat(2,3) = trans(2);
 
   std::cout << "transformMat = " << transformMat << std::endl;
   

  return transformMat;
}