//
//  Olaf Ronneberger
//  Sample Solution for Exercise 3 (Lecture "3D Image Analysis" Summer 2015)
//  "Rigid Registration"
//
//  compile with:
//  g++ -Wall -O4 -g rigid_registration.cc -o rigid_registration 
//
//  specify custom include directory with '-I' if blitz headers are not 
//  installed in '/usr/include'
//
//  requires: libblitz-0.10
//


#include <iostream>
#include <fstream>
#include <sstream>
#include <iomanip>

#include <blitz/array.h>

/*======================================================================*/
/*! 
 *   product of two 4x4 matrices
 *
 *   \param a  first matrix
 *   \param b  second matrix
 *
 *   \exception none
 *
 *   \return product of a and b
 */
/*======================================================================*/
inline
blitz::TinyMatrix<float,4,4> myproduct( 
    blitz::TinyMatrix<float,4,4> a, 
    blitz::TinyMatrix<float,4,4> b)
{
  blitz::TinyMatrix<float,4,4> out;
  out(0,0) = a(0,0)*b(0,0) + a(0,1)*b(1,0) + a(0,2)*b(2,0) + a(0,3)*b(3,0);
  out(1,0) = a(1,0)*b(0,0) + a(1,1)*b(1,0) + a(1,2)*b(2,0) + a(1,3)*b(3,0);
  out(2,0) = a(2,0)*b(0,0) + a(2,1)*b(1,0) + a(2,2)*b(2,0) + a(2,3)*b(3,0);
  out(3,0) = a(3,0)*b(0,0) + a(3,1)*b(1,0) + a(3,2)*b(2,0) + a(3,3)*b(3,0);
  
  out(0,1) = a(0,0)*b(0,1) + a(0,1)*b(1,1) + a(0,2)*b(2,1) + a(0,3)*b(3,1);
  out(1,1) = a(1,0)*b(0,1) + a(1,1)*b(1,1) + a(1,2)*b(2,1) + a(1,3)*b(3,1);
  out(2,1) = a(2,0)*b(0,1) + a(2,1)*b(1,1) + a(2,2)*b(2,1) + a(2,3)*b(3,1);
  out(3,1) = a(3,0)*b(0,1) + a(3,1)*b(1,1) + a(3,2)*b(2,1) + a(3,3)*b(3,1);
  
  out(0,2) = a(0,0)*b(0,2) + a(0,1)*b(1,2) + a(0,2)*b(2,2) + a(0,3)*b(3,2);
  out(1,2) = a(1,0)*b(0,2) + a(1,1)*b(1,2) + a(1,2)*b(2,2) + a(1,3)*b(3,2);
  out(2,2) = a(2,0)*b(0,2) + a(2,1)*b(1,2) + a(2,2)*b(2,2) + a(2,3)*b(3,2);
  out(3,2) = a(3,0)*b(0,2) + a(3,1)*b(1,2) + a(3,2)*b(2,2) + a(3,3)*b(3,2);
  
  out(0,3) = a(0,0)*b(0,3) + a(0,1)*b(1,3) + a(0,2)*b(2,3) + a(0,3)*b(3,3);
  out(1,3) = a(1,0)*b(0,3) + a(1,1)*b(1,3) + a(1,2)*b(2,3) + a(1,3)*b(3,3);
  out(2,3) = a(2,0)*b(0,3) + a(2,1)*b(1,3) + a(2,2)*b(2,3) + a(2,3)*b(3,3);
  out(3,3) = a(3,0)*b(0,3) + a(3,1)*b(1,3) + a(3,2)*b(2,3) + a(3,3)*b(3,3);
  
  return out;
}
 

/*======================================================================*/
/*! 
 *   product of 4x4 matrix with 4d vector
 *
 *   \param m  the matrix
 *   \param v  the vector
 *
 *   \exception none
 *
 *   \return  product of m and v
 */
/*======================================================================*/
inline
blitz::TinyVector<float,4> myproduct( blitz::TinyMatrix<float,4,4> m, 
                                      blitz::TinyVector<float,4> v)
{
  blitz::TinyVector<float,4> out;
  out(0) = m(0,0)*v(0) + m(0,1)*v(1) + m(0,2)*v(2) + m(0,3)*v(3);
  out(1) = m(1,0)*v(0) + m(1,1)*v(1) + m(1,2)*v(2) + m(1,3)*v(3);
  out(2) = m(2,0)*v(0) + m(2,1)*v(1) + m(2,2)*v(2) + m(2,3)*v(3);
  out(3) = m(3,0)*v(0) + m(3,1)*v(1) + m(3,2)*v(2) + m(3,3)*v(3);
  
  return out;
}

/*======================================================================*/
/*! 
 *   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
unsigned char interpolNN( const blitz::Array<unsigned char, 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,3> ipos;
  ipos(0) = std::floor( pos(0) + 0.5);
  ipos(1) = std::floor( pos(1) + 0.5);
  ipos(2) = std::floor( pos(2) + 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<unsigned char, 3>& srcArr,
                     const blitz::TinyMatrix<float,4,4>& invMat,
                     blitz::Array<unsigned char, 3>& trgArr)
{
  blitz::TinyVector<float,4> srcPos;
  blitz::TinyVector<float,4> trgPos;
  trgArr = 0;
  
  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);
      }
    }
  }
}

/*======================================================================*/
/*! 
 *   compute the sum of squared differences (SSD) of the fixed
 *   image and transformed moving image (unsing the given
 *   transformation matrix)
 *
 *   \param 
 *
 *   \exception 
 *
 *   \return
 */
/*======================================================================*/
float ssdOfFixedImAndTransformedMovingIm( 
    const blitz::Array<unsigned char, 3>& movingIm,
    const blitz::TinyMatrix<float,4,4>& invMat,
    const blitz::Array<unsigned char, 3>& fixedIm,
    int step)
{
  blitz::TinyVector<float,4> movingImPos;
  blitz::TinyVector<float,4> fixedImPos;
  
  float sumValue = 0;
  

  for( int lev = 0; lev < fixedIm.extent(0); lev+=step)
  {
    for( int row = 0; row < fixedIm.extent(1); row+=step)
    {
      for( int col = 0; col < fixedIm.extent(2); col+=step)
      {
        fixedImPos = lev, row, col, 1;
        movingImPos = myproduct( invMat, fixedImPos);
        float value= interpolNN( movingIm, movingImPos);
        sumValue += blitz::pow2( value - fixedIm(lev,row,col));
        
      }
    }
  }
  return sumValue;
}


/*======================================================================*/
/*! 
 *   creates an inverse rotation and translation matrix for homogenous
 *   coordinates of the form (lev,row,col,1), to be used as 
 *
 *      srcPosition = Matrix * trgPosition
 *
 *   The translation means that the center of the target
 *   array is mapped to the center of the source array
 *
 *   \param srcArrShape     shape of source array (nLevels, nRows, nCols)
 *   \param trgArrShape     shape of target array (nLevels, nRows, nCols)
 *   \param params          the 6 parameters for Euclidean Transformation:
 *                          shift: lev, row, col
 *                          rotate: aroundLev, aroundRow, aroundCol
 *
 *   \exception none
 *
 *   \return the matrix
 */
/*======================================================================*/
blitz::TinyMatrix<float,4,4> createInverseRigidTransMatrix(
    const blitz::TinyVector<size_t, 3>& srcArrShape,
    const blitz::TinyVector<size_t, 3>& trgArrShape,
    const blitz::TinyVector<float,3>& src_element_size_um,
    const blitz::TinyVector<float,3>& trg_element_size_um,
    blitz::TinyVector<float,6> params)
{
  blitz::TinyMatrix<float,4,4> shiftTrgCenterToOrigin,scaleToMicrometer,scaleToVoxel, rotateAroundLev,
      rotateAroundRow, rotateAroundCol, applyShift, shiftOriginToSrcCenter;
  
  blitz::TinyVector<float,3> shift( params(0), params(1), params(2));

  
  float angleAroundLev = params(3);
  float angleAroundRow = params(4);
  float angleAroundCol = params(5);
  
  shiftTrgCenterToOrigin = 
      1, 0, 0, -float(trgArrShape(0))/2,
      0, 1, 0, -float(trgArrShape(1))/2,
      0, 0, 1, -float(trgArrShape(2))/2,
      0, 0, 0, 1;
  
  scaleToMicrometer = 
      trg_element_size_um(0), 0, 0, 0,
      0, trg_element_size_um(1), 0, 0, 
      0, 0, trg_element_size_um(2), 0, 
      0, 0, 0,                      1;
      
      

  float sina = sin( -angleAroundLev / 180 * M_PI);
  float cosa = cos( -angleAroundLev / 180 * M_PI);
  rotateAroundLev = 
      1,   0,     0,    0, 
      0,   cosa, -sina, 0,
      0,   sina, cosa,  0,
      0,   0,     0,    1;
  

  sina = sin( -angleAroundRow / 180 * M_PI);
  cosa = cos( -angleAroundRow / 180 * M_PI);
  rotateAroundRow = 
      cosa,  0,   sina, 0, 
      0,     1,   0,    0,
      -sina, 0,   cosa, 0,
      0,     0,     0,  1;


  sina = sin( -angleAroundCol / 180 * M_PI);
  cosa = cos( -angleAroundCol / 180 * M_PI);
  rotateAroundCol = 
      cosa, -sina, 0, 0,
      sina, cosa,  0, 0, 
      0,    0,     1, 0,
      0,    0,     0, 1;
  

  applyShift = 
      1, 0, 0, -shift(0),
      0, 1, 0, -shift(1),
      0, 0, 1, -shift(2),
      0, 0, 0, 1;
  
  scaleToVoxel = 
      1/src_element_size_um(0), 0, 0, 0,
      0, 1/src_element_size_um(1), 0, 0,
      0, 0, 1/src_element_size_um(2), 0,
      0, 0, 0,                        1;  

      
      
  shiftOriginToSrcCenter = 
      1, 0, 0, float(srcArrShape(0))/2,
      0, 1, 0, float(srcArrShape(1))/2,
      0, 0, 1, float(srcArrShape(2))/2,
      0, 0, 0, 1;
  
  
  blitz::TinyMatrix<float,4,4> resultMatrix;

  //rb: corrected solution
  //	inverse transform: first shift, then rotate
  resultMatrix = shiftTrgCenterToOrigin;
  resultMatrix = myproduct( scaleToMicrometer,      resultMatrix);
  resultMatrix = myproduct( applyShift,             resultMatrix);
  resultMatrix = myproduct( rotateAroundLev,        resultMatrix);
  resultMatrix = myproduct( rotateAroundRow,        resultMatrix);
  resultMatrix = myproduct( rotateAroundCol,        resultMatrix);
  resultMatrix = myproduct( scaleToVoxel,           resultMatrix);
  resultMatrix = myproduct( shiftOriginToSrcCenter, resultMatrix);
  
  return resultMatrix;
}


/*======================================================================*/
/*! 
 *   load raw 8 bit dataset.  If an error occures, the program is
 *   terminated 
 *
 *   \param fileName the file Name
 *   \param nLevels  number of Levels
 *   \param nRows    number of Rows
 *   \param nCols    number of Columns
 *
 *   \exception 
 *
 *   \return
 */
/*======================================================================*/
blitz::Array<unsigned char, 3> loadRawDataset8bit( 
    const std::string& fileName, int nLevels, int nRows, int nCols)
{
  /*-----------------------------------------------------------------------
   *  read volume from disk
   *-----------------------------------------------------------------------*/
  blitz::Array<unsigned char,3> arr( nLevels, nRows, nCols);
  std::ifstream inFile( fileName.c_str(), std::ifstream::binary);
  if( !inFile)
  {
    std::cerr << "Can't open file '" << fileName << "'";
    perror("");
    exit(1);
  }
  inFile.read( reinterpret_cast<char*>(arr.dataFirst()), 
               arr.size()*sizeof(unsigned char));
  
  return arr;
}


/*======================================================================*/
/*! 
 *   save color image in PPM (portable pix map) format. If an error
 *   occures, the program is terminated
 *
 *   \param image the image
 *   \param fileName the file Name (e.g. "test.ppm")
 *
 *   \exception none
 *
 *   \return none
 */
/*======================================================================*/
void savePPMImage( 
    const blitz::Array<blitz::TinyVector<unsigned char,3>,2>& image,
    const std::string& fileName)
{
  std::ofstream outFile( fileName.c_str(), std::ofstream::binary);
  if( !outFile)
  {
    std::cerr << "Can't create '" << fileName << "'";
    perror("");
    exit(1);
  }
  std::cout << "saving image '" << fileName << "'\n";
  
  
  outFile << "P6\n"
          << image.extent(1) << " " << image.extent(0) << " 255\n";
  outFile.write( reinterpret_cast<const char*>(image.dataFirst()), 
                 image.size() * sizeof( blitz::TinyVector<unsigned char,3>));

}


/*======================================================================*/
/*! 
 *   create orthogonal mips (maximum intensity projections) and
 *   combine them into one image
 *
 *   \param arr  The 3D Array
 *
 *   \exception none
 *
 *   \return image
 */
/*======================================================================*/
blitz::Array<unsigned char,2> orthoMips( const blitz::Array<unsigned char,3>& arr)
{
  /*-----------------------------------------------------------------------
   *  Compute MIPs in all directions in one loop
   *----------------------------------------------------------------------*/
  blitz::Array<unsigned char,2> mipLev( arr.extent(1), arr.extent(2));
  blitz::Array<unsigned char,2> mipRow( arr.extent(0), arr.extent(2));
  blitz::Array<unsigned char,2> mipCol( arr.extent(1), arr.extent(0));
  mipLev = 0;
  mipRow = 0;
  mipCol = 0;
  
  for( int lev = 0; lev < arr.extent(0); ++lev)
  {
    for( int row = 0; row < arr.extent(1); ++row)
    {
      for( int col = 0; col < arr.extent(2); ++col)
      {
        unsigned char value = arr(lev,row,col);
        if( value > mipLev( row, col)) mipLev( row, col) = value;
        if( value > mipRow( lev, col)) mipRow( lev, col) = value;
        if( value > mipCol( row, lev)) mipCol( row, lev) = value;
      }
    }
  }
  
  /*-----------------------------------------------------------------------
   *  Combine mips to one ortho view image
   *-----------------------------------------------------------------------*/
  blitz::Array<unsigned char,2> image( mipRow.extent(0) + mipLev.extent(0) + 1,
                                       mipLev.extent(1) + mipCol.extent(1) + 1);
  image = 255;
  int row = 0;
  int col = 0;
  
  image( blitz::Range( row, row + mipRow.extent(0)-1), 
         blitz::Range( col, col + mipRow.extent(1)-1)) = mipRow;
  
  row = mipRow.extent(0)+1;
  col = 0;
  
  image( blitz::Range( row, row + mipLev.extent(0)-1),
         blitz::Range( col, col + mipLev.extent(1)-1)) = mipLev;

  row = mipRow.extent(0)+1;
  col = mipLev.extent(1)+1;

  image( blitz::Range( row, row + mipCol.extent(0)-1),
         blitz::Range( col, col + mipCol.extent(1)-1)) = mipCol;
  
  return image;
}



int main( int argc, char**argv)
{
  /*-----------------------------------------------------------------------
   *  load moving and fixed data set
   *-----------------------------------------------------------------------*/
  std::cout << "load moving and fixed data set\n";
  
  blitz::Array<unsigned char,3> fixedIm = loadRawDataset8bit( 
      "leaf_t5_150x521x396_8bit.raw", 150, 521, 396);
  
  blitz::Array<unsigned char,3> movingIm = loadRawDataset8bit( 
      "leaf_t6_150x521x396_8bit.raw", 150, 521, 396);
  
  
  blitz::TinyVector<float,3> element_size_um(2.0f,1.46484f,1.46484f);

  /*-----------------------------------------------------------------------
   *  create mip of fixed image for debugging output
   *-----------------------------------------------------------------------*/
  blitz::Array<unsigned char, 2> mipFixed  = orthoMips( fixedIm);

 /*-----------------------------------------------------------------------
   *  initialize the best neighbour search
   *-----------------------------------------------------------------------*/
  blitz::Array<blitz::TinyVector<float,6>,1> neighborDirections(12);
  for( int i = 0; i < 6; ++i)
  {
    neighborDirections(2*i) = 0;
    neighborDirections(2*i)(i) = 1;
    neighborDirections(2*i+1) = 0;
    neighborDirections(2*i+1)(i) = -1;
  }
  
  blitz::Array<unsigned char,3> transformedMovingIm( fixedIm.shape());
 
  blitz::TinyVector<float,6> bestParams;
  blitz::TinyVector<float,6> lastParams;
  
  bestParams = 0;
  lastParams = 0;
  float stepSize = 16;
  float bestCost = 1e20;
  bool betterNeighbourFound;
  
  /*-----------------------------------------------------------------------
   *  create overlay for initial position and save it 
   *-----------------------------------------------------------------------*/
  blitz::TinyMatrix<float,4,4> mat = createInverseRigidTransMatrix(
      movingIm.shape(), transformedMovingIm.shape(),element_size_um,element_size_um, bestParams);
  
  transformArray( movingIm, mat, transformedMovingIm);
      
  blitz::Array<blitz::TinyVector<unsigned char,3>,2> 
      colorImage( mipFixed.shape());
  
  colorImage[0] = mipFixed;
  colorImage[1] = orthoMips( transformedMovingIm);
  colorImage[2] = 0;
  
  savePPMImage( colorImage, "iter_000.ppm");
 

  /*-----------------------------------------------------------------------
   *  repeat until stepSize reaches stop criterium
   *-----------------------------------------------------------------------*/
  std::cout << "do the best neighbour search\n";
  int iteration = 1;
  while( stepSize >= 0.125)
  {
    std::cout << "iteration " << iteration << std::endl;
    
    /*---------------------------------------------------------------------
     *  evaluate all neighbors in parameter space and check if any
     *  parameter combination gives a better match
     *---------------------------------------------------------------------*/
    betterNeighbourFound = false;
    
    for( int i = 0; i < (int)neighborDirections.size(); ++i)
    {
      blitz::TinyVector<float,6> currentParams =
          lastParams + stepSize*neighborDirections(i);

      blitz::TinyMatrix<float,4,4> currentMat = 
          createInverseRigidTransMatrix(
          movingIm.shape(), fixedIm.shape(),element_size_um,element_size_um, currentParams);

      std::cout << "transform array with params " << currentParams;
      
      
      float currentCost = ssdOfFixedImAndTransformedMovingIm( 
          movingIm, currentMat, fixedIm,4);
      std::cout << ": cost = " << currentCost << std::endl;
      
      if( currentCost < bestCost)
      {
        bestCost = currentCost;
        bestParams = currentParams;
        betterNeighbourFound = true;
      }
    }
    
    /*---------------------------------------------------------------------
     *  id a better parameter set was found, use this as new
     *  parameters and create a debugging color image
     *---------------------------------------------------------------------*/
    if( betterNeighbourFound)
    {
      std::cout << "better parameters: " << bestParams 
                << "\t cost: " << bestCost << std::endl;

      blitz::TinyMatrix<float,4,4> bestMat = createInverseRigidTransMatrix(
          movingIm.shape(), transformedMovingIm.shape(),element_size_um,element_size_um, bestParams);
      
      transformArray( movingIm, bestMat, transformedMovingIm);
      
     blitz::Array<blitz::TinyVector<unsigned char,3>,2> 
          colorImage( mipFixed.shape());
      
      colorImage[0] = mipFixed;
      colorImage[1] = orthoMips( transformedMovingIm);
      colorImage[2] = 0;
      
      std::ostringstream os;
      os << "iter_" << std::setw(3) << std::setfill('0') << iteration << ".ppm";
      savePPMImage( colorImage, os.str());
      lastParams = bestParams;
    }
    else
    {
      /*-------------------------------------------------------------------
       *  If no better parameter set was found, reduce the step size
       *-------------------------------------------------------------------*/
      stepSize /= 2;
      std::cout << "step size decreased to: " << stepSize << std::endl;
    }
    ++iteration;
  }
  


  return 0;
}