//
//  Olaf Ronneberger
//  Sample Solution for Exercise 2 (Lecture "3D Image Analysis" Summer 2015) 
//
//  compile with:
//  g++ -Wall -O2 -g render_pollen_anaglyph_movie.cc -o render_pollen_anaglyph_movie
//
//  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;
  
  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);
      }
    }
  }
}


/*======================================================================*/
/*! 
 *   creates an inverse rotation matrix for homogenous coordinates of
 *   the form (lev,row,col,1), to be used as 
 *
 *      srcPosition = Matrix * trgPosition
 *
 *   The translation is adjusted such 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 src_element_size_um		voxel sizes of source array in micrometer (level, row, col)
 *	 \param trg_element_size_um		voxel sizes of target array in micrometer (level, row, col)
 *   \param angleAroundLev  			rotation angle around level-axis in degrees
 *   \param angleAroundRow  			rotation angle around row-axis in degrees
 *   \param angleAroundCol  			rotation angle around column-axis in degrees
 *
 *   \exception none
 *
 *   \return the matrix
 */
/*======================================================================*/

blitz::TinyMatrix<float,4,4> createInverseRotationMatrix(
    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,
    float angleAroundLev,
    float angleAroundRow,
    float angleAroundCol)
{  
  blitz::TinyMatrix<float,4,4> shiftTrgCenterToOrigin, scaleToMicrometer, rotateAroundLev, rotateAroundRow, rotateAroundCol, scaleToVoxel, shiftOriginToSrcCenter;

  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;
  

  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;
  
  resultMatrix = shiftTrgCenterToOrigin;
  resultMatrix = myproduct( scaleToMicrometer,      resultMatrix);
  resultMatrix = myproduct( rotateAroundLev,        resultMatrix);
  resultMatrix = myproduct( rotateAroundRow,        resultMatrix);
  resultMatrix = myproduct( rotateAroundCol,        resultMatrix);
  resultMatrix = myproduct( scaleToVoxel,           resultMatrix);
  resultMatrix = myproduct( shiftOriginToSrcCenter, resultMatrix);

  return resultMatrix;
}


/*======================================================================*/
/*! 
 *   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);
  }
  
  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>));

}


int main( int argc, char**argv)
{

  /*-----------------------------------------------------------------------
   *  Load the data set "Zebrafish_71x361x604_8bit.raw".
   *-----------------------------------------------------------------------*/ 
  int nLevels = 71;
  int nRows = 361;
  int nCols = 604;
  
  blitz::Array<unsigned char,3> srcArr( nLevels, nRows, nCols);
  std::ifstream inFile( "Zebrafish_71x361x604_8bit.raw", 
                        std::ifstream::binary);
  
  if( !inFile)
  {
    perror( "Can't open 'Zebrafish_71x361x604_8bit.raw'");
    exit(1);
  }
  
  inFile.read( reinterpret_cast<char*>( srcArr.dataFirst()), 
               srcArr.size() * sizeof( char));
  
  /*-----------------------------------------------------------------------
   *  Load the data set "Zebrafish_71x361x604_8bit.raw".
   *-----------------------------------------------------------------------*/
  const blitz::TinyVector<float,3> src_element_size_um(3.9, 1.49155, 1.49155);
  const blitz::TinyVector<float,3> trg_element_size_um(3.9, 1.49155, 1.49155);

  /*-----------------------------------------------------------------------
   *  Create the target Array (must be larger than the source Array,
   *  such that the 45 degree rotated array fits into it)
   *-----------------------------------------------------------------------*/
  int diagLevCol = (int) ceil(sqrt( srcArr.extent(0)*srcArr.extent(0) + 
                              srcArr.extent(2)*srcArr.extent(2)));
  blitz::Array<unsigned char,3> trgArr( diagLevCol, srcArr.extent(1), diagLevCol);
 
  /*-----------------------------------------------------------------------
   *  Create Anaglyph-image
   *-----------------------------------------------------------------------*/

  const float stereoAngle = 5;	// angle in degrees

  blitz::Array<unsigned char,2> singleLeftView( trgArr.extent(1), trgArr.extent(2));
  blitz::Array<unsigned char,2> singleRightView( trgArr.extent(1), trgArr.extent(2));

  blitz::TinyMatrix<float,4,4> rotMatRowLeft;
  blitz::TinyMatrix<float,4,4> rotMatRowRight;

  /*---------------------------------------------------------------------
   *  compute image for the right eye, rotated array (by +stereoAngle/2 around row)
   *---------------------------------------------------------------------*/
  rotMatRowLeft = createInverseRotationMatrix( srcArr.shape(), trgArr.shape(),
						src_element_size_um, trg_element_size_um,
						0, stereoAngle/2, 0);
  
  transformArray( srcArr, rotMatRowLeft, trgArr);
  
  /*---------------------------------------------------------------------
   *  compute MIP of rotated array
   *---------------------------------------------------------------------*/
  singleRightView = 0;
  for( int lev = 0; lev < trgArr.extent(0); ++lev)
  {
    singleRightView = blitz::max( singleRightView, trgArr( lev,
							   blitz::Range::all(),
							   blitz::Range::all()));
  }

  /*---------------------------------------------------------------------
   *  compute image for the left eye, rotated array (by -stereoAngle/2 around row)
   *---------------------------------------------------------------------*/
  rotMatRowRight = createInverseRotationMatrix( srcArr.shape(), trgArr.shape(),
						src_element_size_um, trg_element_size_um,
						0, -stereoAngle/2, 0);
  
  transformArray( srcArr, rotMatRowRight, trgArr);
  
  /*---------------------------------------------------------------------
   *  compute MIP of rotated array
   *---------------------------------------------------------------------*/
  singleLeftView = 0;
  for( int lev = 0; lev < trgArr.extent(0); ++lev)
  {
    singleLeftView = blitz::max( singleLeftView, trgArr( lev,
							 blitz::Range::all(),
							 blitz::Range::all()));
  }
 
  /*---------------------------------------------------------------------
   *  copy into red-cyan anaglyhp image
   *---------------------------------------------------------------------*/
  blitz::Array< blitz::TinyVector<unsigned char,3> ,2> singleAnaglypgImage( trgArr.extent(1), trgArr.extent(2));

  singleAnaglypgImage[0] = singleLeftView;
  singleAnaglypgImage[1] = singleRightView;
  singleAnaglypgImage[2] = singleRightView;

  /*---------------------------------------------------------------------
   *  save the resulting image to "anaglyphImage.ppm"
   *---------------------------------------------------------------------*/
  std::ostringstream os;
  os << "anaglyphImage.ppm";
  std::string fileName = os.str();
  std::cout << "saving " << fileName << std::endl;
  
  savePPMImage( singleAnaglypgImage, fileName);


  /*-----------------------------------------------------------------------
   *  Create Anaglyph-movie
   *-----------------------------------------------------------------------*/
  blitz::TinyMatrix<float,4,4> rotMatColRowLeft;
  blitz::TinyMatrix<float,4,4> rotMatColRowRight;

  /*-----------------------------------------------------------------------
   *  Create the target Array (must be larger than the source Array,
   *  such that the 45 degree rotated array fits into it)
   *-----------------------------------------------------------------------*/
  int diagLevRow = (int) ceil(sqrt( srcArr.extent(0)*srcArr.extent(0) + 
                              srcArr.extent(1)*srcArr.extent(1)));

  blitz::Array<unsigned char,3> trgArrAnaglyph( std::max(diagLevRow,diagLevCol), diagLevRow, diagLevCol);

  blitz::Array<unsigned char,2> leftView( trgArrAnaglyph.extent(1), trgArrAnaglyph.extent(2));
  blitz::Array<unsigned char,2> rightView( trgArrAnaglyph.extent(1), trgArrAnaglyph.extent(2));

 	blitz::Array< blitz::TinyVector<unsigned char,3> ,2> anaglyphImage( trgArrAnaglyph.extent(1), trgArrAnaglyph.extent(2));

  /*-----------------------------------------------------------------------
   *  create an anaglyph-movie with 5 angle degree steps
   *-----------------------------------------------------------------------*/
  for( int angle = 0; angle < 360; angle += 5)
  {
    /*---------------------------------------------------------------------
     *  compute image for the right eye, rotated array (by alpha around col, by +stereoAngle/2 around row)
     *---------------------------------------------------------------------*/
    rotMatColRowLeft = createInverseRotationMatrix( srcArr.shape(), trgArrAnaglyph.shape(),
						    src_element_size_um, trg_element_size_um,
						    0, stereoAngle/2, angle);
    
    transformArray( srcArr, rotMatColRowLeft, trgArrAnaglyph);
    
    /*---------------------------------------------------------------------
     *  compute MIP of rotated array
     *---------------------------------------------------------------------*/
    rightView = 0;
    for( int lev = 0; lev < trgArrAnaglyph.extent(0); ++lev)
    {
      rightView = blitz::max( rightView, trgArrAnaglyph( lev,
							 blitz::Range::all(),
							 blitz::Range::all()));
    }

    /*---------------------------------------------------------------------
     *  compute image for the left eye, rotated array (by alpha around col, by -stereoAngle/2 around row)
     *---------------------------------------------------------------------*/
    rotMatColRowRight = createInverseRotationMatrix( srcArr.shape(), trgArrAnaglyph.shape(),
						     src_element_size_um, trg_element_size_um,
						     0, -stereoAngle/2, angle);
    
    transformArray( srcArr, rotMatColRowRight, trgArrAnaglyph);
    
    /*---------------------------------------------------------------------
     *  compute MIP of rotated array
     *---------------------------------------------------------------------*/
    leftView = 0;
    for( int lev = 0; lev < trgArrAnaglyph.extent(0); ++lev)
    {
      leftView = blitz::max( leftView, trgArrAnaglyph( lev,
						       blitz::Range::all(),
						       blitz::Range::all()));
    }
   
    /*---------------------------------------------------------------------
     *  copy into red-cyan anaglyhp image
     *---------------------------------------------------------------------*/
    anaglyphImage[0] = leftView;
    anaglyphImage[1] = rightView;
    anaglyphImage[2] = rightView;

    /*---------------------------------------------------------------------
     *  save the resulting image to "movie_anaglyph_frame_XXX.ppm", where
     *  XXX is the (zeropadded) angle
     *---------------------------------------------------------------------*/
    std::ostringstream os;
    os << "anaglyph/movie_anaglyph_frame_" 
       << std::setw(3) << std::setfill('0') << angle 
       << ".ppm";
    std::string fileName = os.str();
    std::cout << "saving " << fileName << std::endl;
    
    savePPMImage( anaglyphImage, fileName);
  }
  

  return 0;
}