/**
 * @file  transform3dimg.cxx
 * @brief Originally part of HAMMER.
 *
 * Copyright (c) 2001, 2012 University of Pennsylvania.
 *
 * This file is part of DTI-DROID.
 *
 * DTI-DROID is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * DTI-DROID is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with DTI-DROID.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Contact: SBIA Group <sbia-software at uphs.upenn.edu>
 */

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <string.h>

#include "mvcd.h"
#include "cres.h"
#include "matrixSHEN.h"

#include "basis.h"


// acceptable in .cxx file
using namespace sbia;


// ===========================================================================
// constants
// ===========================================================================

#define YYES 1
#define NNO  0

#define BG  0
#define CSF 10
#define VN  50
#define GM  150
#define WM  250

// ===========================================================================
// help
// ===========================================================================

// ---------------------------------------------------------------------------
void show_usage()
{
  basis::print_version("transform3dimg", "", "");
  printf("\n");
  printf("USAGE: transform3Dimg <img> <img-transformed> <affine-transformation>\n\
\t \n\
\t The rigid transformed image will be transformed back to its original domain, using transformation parameters saved under GlobalTransformation.Affine in the currect directory \n\
\t \n\
\t -v <int>             : xy_size (i.e. 256 ) in XY plane \n\
\t -I                   : inversely transform 3D image \n\
\t -L                   : linear interpolation! \n\
");
}

// ===========================================================================
// auxiliary functions
// ===========================================================================

int LinearInterpolation, NNmapping4SegmentedImg ;

void Open_Transformation( Matrix *Transform, Fvector3d *center1, Fvector3d *center2, char transformFilename[120])
{
  FILE *fp;
  int i, j ;

  double *Temp_transf ;
  int Tm, Tn ;

  Tm = Transform->height ;
  Tn = Transform->width ;

  /* allocate space */
  Temp_transf = (double *)calloc( Tm*Tn, sizeof(double) );

  
  if((fp=fopen(transformFilename,"rb"))==NULL)
  {
   printf("cannot open transformation matrix\n");
   printf("%s",transformFilename);
   return ;
  }

  fseek(fp,0L,SEEK_SET);

  fread( Temp_transf,  sizeof(double), Tm*Tn,  fp ) ; 
  fread( center1,      sizeof(Fvector3d),  1,  fp ) ;
  fread( center2,      sizeof(Fvector3d),  1,  fp ) ;
 
  fclose(fp);


  /* for T1 */
  for(i=0; i<Tm; i++)
   for(j=0; j<Tn; j++)
    Transform->data[i][j] = Temp_transf[i*Tn+j] ; 

  /* free */
  free(Temp_transf) ;
}


unsigned char InterpolatedIntensity(float ii, float jj, float kk, unsigned char ***Img, int image_size, int z_size)
{
  float CurrentV ;
  float b,c,d, b1,c1,d1;
  int   ni,nj,nk, niP1,njP1,nkP1, GreyValue ;


  ni = (int)ii ;
  nj = (int)jj ;
  nk = (int)kk ;
  
  niP1 = ni+1 ;
  njP1 = nj+1 ;
  nkP1 = nk+1 ;
  
  GreyValue = 0;
  if(ni>=0 && ni<image_size-1  &&  nj>=0 && nj<image_size-1  &&  nk>=0 && nk<z_size-1 )
    {
      b = ii-ni ;        b1 = 1.-b ;
      c = jj-nj ;        c1 = 1.-c ;
      d = kk-nk ;        d1 = 1.-d ;

      CurrentV = ( d1*(Img[nk][ni][nj]*(b1*c1) + Img[nk][niP1][nj]*(b*c1) + Img[nk][ni][njP1]*(b1*c) + Img[nk][niP1][njP1]*(b*c)) + d*(Img[nkP1][ni][nj]*(b1*c1) + Img[nkP1][niP1][nj]*(b*c1) + Img[nkP1][ni][njP1]*(b1*c) + Img[nkP1][niP1][njP1]*(b*c)) )/( d1*((b1*c1)+(b*c1)+(b1*c)+(b*c)) + d*((b1*c1)+(b*c1)+(b1*c)+(b*c)) ) ; 
      if (CurrentV > 255) CurrentV = 255;
	  GreyValue = static_cast<int>(CurrentV);
    }
  
  if(ni==image_size-1 && nj>=0 && nj<image_size-1 && nk>=0 && nk<z_size-1 || ni>=0 && ni<image_size-1 && nj==image_size-1 && nk>=0 && nk<z_size-1  || ni>=0 && ni<image_size-1 && nj>=0 && nj<image_size-1 && nk==z_size-1)
    GreyValue = Img[nk][ni][nj] ;

  return GreyValue ;
}

void MappingIntensitiesBasedOnNearestDistance(unsigned char ***src, int image_size, int z_size)
{
  float distanceVN, distanceGM, distanceWM, distanceBG, distanceCSF;
  int x,y,z ;

  /* mapping  */
  for(z=0;z<z_size;z++)
    for(x=0;x<image_size;x++)
      for(y=0;y<image_size;y++)
	{
	  distanceBG = (src[z][x][y]-BG)*(src[z][x][y]-BG) ;
	  distanceCSF= (src[z][x][y]-CSF)*(src[z][x][y]-CSF) ;
	  distanceVN = (src[z][x][y]-VN)*(src[z][x][y]-VN) ;
	  distanceGM = (src[z][x][y]-GM)*(src[z][x][y]-GM) ;
	  distanceWM = (src[z][x][y]-WM)*(src[z][x][y]-WM) ;

	  if( distanceVN<distanceGM && distanceVN<distanceWM )
	    src[z][x][y] = VN ;
	  if( distanceGM<distanceVN && distanceGM<distanceWM )
	    src[z][x][y] = GM ;
	  if( distanceWM<distanceVN && distanceWM<distanceGM )
	    src[z][x][y] = WM ;
	  if( sqrt(distanceBG)<5 )
	    src[z][x][y] = BG ;
	  if( sqrt(distanceCSF)<5 )
	    src[z][x][y] = CSF ;

	  if( src[z][x][y]==100 )  src[z][x][y]=VN ;
	  if( src[z][x][y]==200 )  src[z][x][y]=GM ;
	}
}



void TransformImageAndSaveIt( unsigned char ***Img, int image_size, int z_size, char filename[120], char transformFilename[120])
{
  Matrix *Transform, *InversedMtrx ;
  Fvector3d center1, center2 ;
  float  *current_position, *transformed_position;
  FILE  *fp;
  int   i, j, k ;
  int   x, y, z ;
  unsigned char ***TransformImg;
  float xx, yy, zz ;

  TransformImg = UCalloc3d(image_size,image_size,z_size);

  CreateMatrix(&Transform, 3, 3);
  current_position= vectorSHEN(0, 3-1) ;
  transformed_position= vectorSHEN(0, 3-1) ;
  CreateMatrix(&InversedMtrx, 3, 3);

  Open_Transformation( Transform, &center1, &center2, transformFilename) ; 
    Mat_Print(Transform);
    printf("center1=(%f, %f, %f)\n", (center1).x, (center1).y, (center1).z) ;
    printf("center2=(%f, %f, %f)\n", (center2).x, (center2).y, (center2).z) ;

  Mat_Inverse( Transform, InversedMtrx ) ;
    Mat_Print(InversedMtrx);

  /* transform it */
  if( LinearInterpolation==NNO )
    for(k=0;k<z_size;k++)
      for(i=0;i<image_size;i++)
	for(j=0;j<image_size;j++)
	  {
	    current_position[0] = i - center1.x ;
	    current_position[1] = j - center1.y ;
	    current_position[2] = k - center1.z ;
	    
	    Mat_times_Vector( transformed_position, InversedMtrx, current_position) ; 
	    
	    x = static_cast<int>(transformed_position[0]+center2.x +0.5) ; /* 0.5 was re-added on Oct 22,2002*/
	    y = static_cast<int>(transformed_position[1]+center2.y +0.5) ;
	    z = static_cast<int>(transformed_position[2]+center2.z +0.5) ;
	    
	    if( x<image_size && y<image_size && z<z_size && x>=0 && y>=0 && z>=0 )
	      TransformImg[k][i][j] = Img[z][x][y] ;
	    /*printf("(%d,%d,%d) ->  (%d,%d,%d)\n", k,i,j, z,x,y) ;*/
	  }
  else
    {
      for(k=0;k<z_size;k++)
	for(i=0;i<image_size;i++)
	  for(j=0;j<image_size;j++)
	    {
	      current_position[0] = i - center1.x ;
	      current_position[1] = j - center1.y ;
	      current_position[2] = k - center1.z ;
	      
	      Mat_times_Vector( transformed_position, InversedMtrx, current_position) ; 
	      
	      xx = transformed_position[0]+center2.x ; /* 0.5 was re-added on Oct 22,2002*/
	      yy = transformed_position[1]+center2.y ;
	      zz = transformed_position[2]+center2.z ;
	      
	      TransformImg[k][i][j] = InterpolatedIntensity(xx, yy, zz, Img, image_size, z_size) ;
	    }
      /*if( NNmapping4SegmentedImg==YYES ) 
	MappingIntensitiesBasedOnNearestDistance(TransformImg, image_size, z_size);*/
    }


  /* save it */
  printf("save image!\n") ;
  fp=myopen(filename,"wb");
  for(k=0;k<z_size;k++)
    for(i=0;i<image_size;i++)
      fwrite(TransformImg[k][i],1,image_size,fp);
  fclose(fp);


  /* free */
  FreeMatrix(Transform);
  FreeMatrix(InversedMtrx);
  free_vectorSHEN(current_position, 0, 3-1) ;
  free_vectorSHEN(transformed_position, 0, 3-1) ;
}

void InverseTransformImageAndSaveIt( unsigned char ***Img, int image_size, int z_size, char filename[120], char transformFilename[120])
{
  Matrix *Transform ;
  Fvector3d center1, center2 ;
  float  *current_position, *transformed_position;
  FILE  *fp;
  int   i, j, k ;
  int   x, y, z ;
  unsigned char ***TransformImg;
  float xx, yy, zz ;

  TransformImg = UCalloc3d(image_size,image_size,z_size);

  CreateMatrix(&Transform, 3, 3);
  current_position= vectorSHEN(0, 3-1) ;
  transformed_position= vectorSHEN(0, 3-1) ;


  Open_Transformation( Transform, &center1, &center2, transformFilename) ; 
    Mat_Print(Transform);
    printf("center1=(%f, %f, %f)\n", (center1).x, (center1).y, (center1).z) ;
    printf("center2=(%f, %f, %f)\n", (center2).x, (center2).y, (center2).z) ;


  /* transform it */
  if( LinearInterpolation==NNO )
    for(k=0;k<z_size;k++)
      for(i=0;i<image_size;i++)
	for(j=0;j<image_size;j++)
	  {
	    current_position[0] = i - center2.x ;
	    current_position[1] = j - center2.y ;
	    current_position[2] = k - center2.z ;
	    
	    Mat_times_Vector( transformed_position, Transform, current_position) ; 
	    
	    x = static_cast<int>(transformed_position[0]+center1.x +0.5) ;
	    y = static_cast<int>(transformed_position[1]+center1.y +0.5) ;
	    z = static_cast<int>(transformed_position[2]+center1.z +0.5) ;
	    
	    if( x<image_size && y<image_size && z<z_size && x>=0 && y>=0 && z>=0 )
	      TransformImg[k][i][j] = Img[z][x][y] ;
	  }
  else
    {
      for(k=0;k<z_size;k++)
	for(i=0;i<image_size;i++)
	  for(j=0;j<image_size;j++)
	    {
	      current_position[0] = i - center2.x ;
	      current_position[1] = j - center2.y ;
	      current_position[2] = k - center2.z ;
	      
	      Mat_times_Vector( transformed_position, Transform, current_position) ; 
	      
	      xx = transformed_position[0]+center1.x ;
	      yy = transformed_position[1]+center1.y ;
	      zz = transformed_position[2]+center1.z ;
	      
	      TransformImg[k][i][j] = InterpolatedIntensity(xx, yy, zz, Img, image_size, z_size) ;
	    }
      /*if( NNmapping4SegmentedImg==YYES ) 
	MappingIntensitiesBasedOnNearestDistance(TransformImg, image_size, z_size);*/
    }


  /* save it */
  printf("save image!\n") ;
  fp=myopen(filename,"wb");
  for(k=0;k<z_size;k++)
    for(i=0;i<image_size;i++)
      fwrite(TransformImg[k][i],1,image_size,fp);
  fclose(fp);


  /* free */
  FreeMatrix(Transform);
  free_vectorSHEN(current_position, 0, 3-1) ;
  free_vectorSHEN(transformed_position, 0, 3-1) ;
}

// ===========================================================================
// main
// ===========================================================================

int main(int argc,char *argv[])
{
  int           image_size, z_size;
  unsigned char ***Img;
  int           i,j,k,c,num;
  FILE          *fp;
  extern char   *optarg;
  char          filename[120] ;
  int           InverseTransform ;


  num=3;

  if(argc<num) {
    show_usage() ;
    exit(1);
  }


  /* input image size */
  image_size= 256;
  z_size    = 124 ;
  InverseTransform = NNO ;
  LinearInterpolation=NNO ;
  NNmapping4SegmentedImg=NNO;

  while((c=getopt(argc-3,argv+3,"v:s:ILN")) != -1)
    {
      switch(c)
	{
	case 'v':
	  sscanf(optarg, "%d", &image_size) ; /* image size */
	  break ;

	case 'I':
	  InverseTransform = YYES ;
	  break ;

	case 'L':
	  LinearInterpolation = YYES ;
	  break ;

	case 'N':
	  NNmapping4SegmentedImg = YYES ;
	  break ;


	default:
	  break;
	}
    }
  printf("LinearInterpolation=%d   NNmapping4SegmentedImg=%d\n", LinearInterpolation, NNmapping4SegmentedImg) ;

  /* open image */
  fp=myopen(argv[1],"rb");
  fseek(fp,0,SEEK_END);
  z_size=ftell(fp)/(image_size*image_size);
  rewind(fp);
  printf("image size: image_size=%d  z_size=%d\n", image_size, z_size) ;

  Img = UCalloc3d(image_size,image_size,z_size);
  
  for(k=0;k<z_size;k++)
    for(i=0;i<image_size;i++)
      fread(Img[k][i],1,image_size,fp);
  fclose(fp);
  printf("image reading finished!\n") ;

  /* perform transformation */
  if( InverseTransform == NNO )
    TransformImageAndSaveIt( Img, image_size, z_size, argv[2], argv[3]) ;
  else
    InverseTransformImageAndSaveIt( Img, image_size, z_size, argv[2], argv[3]) ;
  
  return 0;
}
