//
//  Olaf Ronneberger
//  Sample Solution for Exercise 2 (Lecture "3D Image Analysis" Summer 2015) 
//
//  compile with:
//  g++ -Wall -O2 -g render_pollen_movie.cc -o render_pollen_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 image in PGM (portable gray map) format. If an error
 *   occures, the program is terminated
 *
 *   \param image the image
 *   \param fileName the file Name (e.g. "test.pgm")
 *
 *   \exception none
 *
 *   \return none
 */
/*======================================================================*/
void savePGMImage( const blitz::Array<unsigned char,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 << "P5\n"
          << image.extent(1) << " " << image.extent(0) << " 255\n";
  outFile.write( reinterpret_cast<const char*>(image.dataFirst()), 
                 image.size() * sizeof( char));

}


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

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

  /*-----------------------------------------------------------------------
   *  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);
  
  blitz::Array<unsigned char,2> image( trgArr.extent(1), trgArr.extent(2));

  blitz::TinyMatrix<float,4,4> rotMat;
  
  /*-----------------------------------------------------------------------
   *  create a MIP-movie with 5 angle degree steps
   *-----------------------------------------------------------------------*/
  for( int angle = 0; angle < 360; angle += 5)
  {
    /*---------------------------------------------------------------------
     *  compute rotated array
     *---------------------------------------------------------------------*/
   
    rotMat = createInverseRotationMatrix( srcArr.shape(), trgArr.shape(),
					    src_element_size_um, trg_element_size_um,
					    0, angle, 0);
    
    transformArray( srcArr, rotMat, trgArr);
    
    /*---------------------------------------------------------------------
     *  compute MIP of rotated array
     *---------------------------------------------------------------------*/
    image = 0;
    for( int lev = 0; lev < trgArr.extent(0); ++lev)
    {
      image = blitz::max( image, trgArr( lev, 
                                         blitz::Range::all(),
                                         blitz::Range::all()));
    }
    
    /*---------------------------------------------------------------------
     *  save the resulting image to "movie_frame_XXX.pgm", where
     *  XXX is the (zeropadded) angle
     *---------------------------------------------------------------------*/
    std::ostringstream os;
    os << "movie/movie_frame_" 
       << std::setw(3) << std::setfill('0') << angle 
       << ".pgm";
    std::string fileName = os.str();
    std::cout << "saving " << fileName << std::endl;
    
    savePGMImage( image, fileName);
	} 

  return 0;
}