/**
 * @file  transform3dvectorfield.cxx
 * @brief Originally part of HAMMER.
 *
 * Copyright (c) 2001, 2012, 2013 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 <dtidroid/basis.h>


// acceptable in .cxx file
using namespace dtidroid;


// ===========================================================================
// constants and global variables
// ===========================================================================

#define YYES    1
#define NNO     0

int  HeadInField, WriteHeadInField ;

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

// ---------------------------------------------------------------------------
void show_usage()
{
  print_version("transform3dvectorfield", "", "");
  printf("\n");
  printf("USAGE: transform3Dvectorfield <vector-field> <transformed-vector-field> <affine-transformation>\n\
\t \n\
\t The current deformation field is from model to the affine transformed subject. The goal of this program is to make the deformation field pointing from the model to the original subject, using transformation parameters saved under GlobalTransformation.Affine in the currect directory \n\
\t \n\
\t -v <int>             : xy_size (i.e. 256 ) in vector field for the case that there is no head info in the file\n\
\t -I                   : The current one, the deformation is from the rigid transformed model to a certain subject. If we have file GlobalTransformation.Affine that links the rigid transformed model to its original version, then we can make the deformation field directly from the original model to that certain subject. \n\
\t -H                   : write the sizes of fields (image_size, z_size) into the file\n\
");
}

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

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 file\n");
   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) ;
}


void TransformVectorFieldAndSaveIt( Fvector3d ***DeformFld, 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 ;


  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 */
  for(k=0;k<z_size;k++)
    for(i=0;i<image_size;i++)
      for(j=0;j<image_size;j++)
	{
	  current_position[0] = (DeformFld[k][i][j].x+i) - center2.x ;
	  current_position[1] = (DeformFld[k][i][j].y+j) - center2.y ;
	  current_position[2] = (DeformFld[k][i][j].z+k) - center2.z ;
	  
	  Mat_times_Vector( transformed_position, Transform, current_position) ; 

	  DeformFld[k][i][j].x = (transformed_position[0]+center1.x)-i  ;
	  DeformFld[k][i][j].y = (transformed_position[1]+center1.y)-j  ;
	  DeformFld[k][i][j].z = (transformed_position[2]+center1.z)-k  ;
	}


  /* save it */
  fp=myopen(filename,"wb");
  if( HeadInField==YYES || WriteHeadInField==YYES )
    {
      fwrite(&image_size,sizeof(int),1,fp);
      fwrite(&z_size,sizeof(int),1,fp);
    }
  for(k=0;k<z_size;k++)
    for(i=0;i<image_size;i++)
      fwrite(DeformFld[k][i],sizeof(Fvector3d),image_size,fp);
  fclose(fp);


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


void InterpolatedDisplacement(Fvector3d *Displace_subVoxel, float ii, float jj, float kk, Fvector3d ***DeformFld, 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 ;
  
  (*Displace_subVoxel).x = 0 ;
  (*Displace_subVoxel).y = 0 ;
  (*Displace_subVoxel).z = 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 ;

      (*Displace_subVoxel).x = ( d1*(DeformFld[nk][ni][nj].x*(b1*c1) + DeformFld[nk][niP1][nj].x*(b*c1) + DeformFld[nk][ni][njP1].x*(b1*c) + DeformFld[nk][niP1][njP1].x*(b*c)) + d*(DeformFld[nkP1][ni][nj].x*(b1*c1) + DeformFld[nkP1][niP1][nj].x*(b*c1) + DeformFld[nkP1][ni][njP1].x*(b1*c) + DeformFld[nkP1][niP1][njP1].x*(b*c)) )/( d1*((b1*c1)+(b*c1)+(b1*c)+(b*c)) + d*((b1*c1)+(b*c1)+(b1*c)+(b*c)) ) ; 

      (*Displace_subVoxel).y = ( d1*(DeformFld[nk][ni][nj].y*(b1*c1) + DeformFld[nk][niP1][nj].y*(b*c1) + DeformFld[nk][ni][njP1].y*(b1*c) + DeformFld[nk][niP1][njP1].y*(b*c)) + d*(DeformFld[nkP1][ni][nj].y*(b1*c1) + DeformFld[nkP1][niP1][nj].y*(b*c1) + DeformFld[nkP1][ni][njP1].y*(b1*c) + DeformFld[nkP1][niP1][njP1].y*(b*c)) )/( d1*((b1*c1)+(b*c1)+(b1*c)+(b*c)) + d*((b1*c1)+(b*c1)+(b1*c)+(b*c)) ) ; 

      (*Displace_subVoxel).z = ( d1*(DeformFld[nk][ni][nj].z*(b1*c1) + DeformFld[nk][niP1][nj].z*(b*c1) + DeformFld[nk][ni][njP1].z*(b1*c) + DeformFld[nk][niP1][njP1].z*(b*c)) + d*(DeformFld[nkP1][ni][nj].z*(b1*c1) + DeformFld[nkP1][niP1][nj].z*(b*c1) + DeformFld[nkP1][ni][njP1].z*(b1*c) + DeformFld[nkP1][niP1][njP1].z*(b*c)) )/( d1*((b1*c1)+(b*c1)+(b1*c)+(b*c)) + d*((b1*c1)+(b*c1)+(b1*c)+(b*c)) ) ; 
    }
  
  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))
    {
      (*Displace_subVoxel).x = DeformFld[nk][ni][nj].x ;
      (*Displace_subVoxel).y = DeformFld[nk][ni][nj].y ;
      (*Displace_subVoxel).z = DeformFld[nk][ni][nj].z ;
    }
}


void InverseTransformVectorFieldAndSaveIt( Fvector3d ***DeformFld, 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 ;
  Fvector3d ***New_DeformFld ;
  Fvector3d  SubVoxel, Displace_subVoxel ;

  printf("inverse\n") ;
/* apply space for new deformation field */
  New_DeformFld = Fvector3dalloc3d(image_size,image_size,z_size);


  CreateMatrix(&Transform, 3, 3);
  CreateMatrix(&InversedMtrx, 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) ;


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


  /* transform it */
  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) ; 

	  SubVoxel.x = (transformed_position[0]+center2.x) ;
	  SubVoxel.y = (transformed_position[1]+center2.y) ;
	  SubVoxel.z = (transformed_position[2]+center2.z) ;

	  InterpolatedDisplacement(&Displace_subVoxel, SubVoxel.x, SubVoxel.y, SubVoxel.z, DeformFld, image_size, z_size) ;

	  New_DeformFld[k][i][j].x = (SubVoxel.x + Displace_subVoxel.x)-i  ;
	  New_DeformFld[k][i][j].y = (SubVoxel.y + Displace_subVoxel.y)-j  ;
	  New_DeformFld[k][i][j].z = (SubVoxel.z + Displace_subVoxel.z)-k  ;
	}


  /* save it */
  fp=myopen(filename,"wb");
  if( HeadInField==YYES  || WriteHeadInField==YYES )
    {
      fwrite(&image_size,sizeof(int),1,fp);
      fwrite(&z_size,sizeof(int),1,fp);
    }
  for(k=0;k<z_size;k++)
    for(i=0;i<image_size;i++)
      fwrite(New_DeformFld[k][i],sizeof(Fvector3d),image_size,fp);
  fclose(fp);


  /* free */
  FreeMatrix(Transform);
  FreeMatrix(InversedMtrx);
  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;
  Fvector3d     ***DeformFld ;
  int           i,j,k,c,num;
  FILE          *fp;
  extern char   *optarg;
  char          filename[120] ;
  int           InverTransform ;


  num=3;

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


  /* input image size */
  image_size= 256;
  z_size    = 124 ;
  HeadInField = YYES ;
  InverTransform = NNO ;
  WriteHeadInField = NNO ;

  while((c=getopt(argc-3,argv+3,"v:IH")) != -1)
    {
      switch(c)
	{
	case 'v':
	  sscanf(optarg, "%d", &image_size) ; /* for vector field */
	  HeadInField = NNO ;
	  break ;

	case 'I':
          InverTransform = YYES ;
	  break;

	case 'H':
	  WriteHeadInField = YYES ;
	  break;

	default:
	  break;
	}
    }
  if(HeadInField==YYES) printf("There is size information in vector field file!\n") ;
  else                  printf("No head in vector field file!\n") ;


  /* open deformation field in the last resolution */
  fp=myopen(argv[1],"rb");
  if( HeadInField==YYES )
    {
      fread(&image_size,sizeof(int),1,fp); printf("image_size=%d ", image_size) ;
      fread(&z_size,sizeof(int),1,fp);     printf("z_size=%d\n", z_size) ;
    }
  else
    {
      fseek(fp,0,SEEK_END);
      z_size=ftell(fp)/(image_size*image_size*12);
      rewind(fp);
    }
  printf("vector field size: image_size=%d  z_size=%d\n", image_size, z_size) ;

  DeformFld = Fvector3dalloc3d(image_size,image_size,z_size);
  
  for(k=0;k<z_size;k++)
    for(i=0;i<image_size;i++)
      fread(DeformFld[k][i],sizeof(Fvector3d),image_size,fp);
  fclose(fp);
  printf("Vector fields reading finished!\n") ;

  /* perform transformation */
if( InverTransform==YYES )
  InverseTransformVectorFieldAndSaveIt( DeformFld, image_size, z_size, argv[2], argv[3]) ;
else
  TransformVectorFieldAndSaveIt( DeformFld, image_size, z_size, argv[2], argv[3]) ;
  
  return 0;
}
