/**
 * @file  LnWarpDT.cxx
 * @brief Affine transformation of diffusion tensor image.
 *
 * Copyright (c) 2008, 2009, 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>
 */
 
 /*-----------------------------------------------------------
  Coordinate system:

   \Z+
    \
     \_________________X+
     |O
     |
     |
     |
     |Y+
     

  Purpose:
     Transform the diffusion tensor (DT) in the original space
  (52-136 Z-slice) to the atlas space (123 Z-slice)
     A diffusion Tensor is a matrix:
//		[a1, a4, a5]	   	[XX, XY, XZ]
//		[a4, a2, a6]   i.e.	[YX, YY, YZ]
//		[a5, a6, a3]		[ZX, ZY, ZZ]


   Make file:
     make -f mkNWDT
  
   Executable example:
     nWarpDT /nilab/meteora1/susumu/DTensor/kim_1538.d ./ori2atlas.vf -O tmp.dt -Z 57

*/

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

#include <iostream>

//#include <task.h>    // for multi-cpus
//#include <ulocks.h>  // for multi-cpus

#include "Basics.h"
#include "GPrimitive.h"
#include "DTVolume.h"
#include "VoxMap.h"
#include "WarpVolume.h"
#include "mVolume.h"
#include "Ellipsoid.h"

using namespace std;

float resX,resY,resZ; // resolution in X,Y,Z directions of the volume
int dimX; // dim of the entered DT image
char DTfile[200]; //DTfile file name (input)
char fPD[200];    //Primary Direction(PD) vector of DTI (input)
char fTransformFile[120];  //warping vector field file (input)
char outDTfile[200];  // file name of result to be stored (output)
int destZ; // destination volume Z-slice number
int CPUs; // number of CUPs to be used, default = 1
int flipZ; // flag: whether or not to flip DT's D[x,z] and D[y,z]
int inversedTransform=0; //Input global transform is inversed or not


DTVolume *volDT=NULL; //pointer to a DTI volume

void usage(char* runname)
{
   cout << "\nUsage:\n";
   cout << "   " << runname <<" <inputDTVolume> <GlobalTransform> -O<outfile> -X<Xdim(=Ydim)> -Z<slice> -R<rx,ry,rz> -f -N<0/1> -I" <<endl;
   cout << "   " << runname <<" -H" <<endl;
   
   cout << "  <inputDTVolume>: the diffusion tensor volume file to be warped." <<endl;
   cout << "  <GlobalTransform>: the backward linear transform from the atlas to volume.\n";
   cout << "  -O<outfile>    : resulting DT field, default: ./resultDT.dt \n";
   cout << "                   this will be in the atlas (123 slc) space"<<endl;
   cout << "  -Z<slice>      : Z-slice number of the warped volume, default: 123 \n";
   cout << "  -X<Xdim(=Ydim)>      : X or Y dim of the warped volume, default: 256 \n";								 
   cout << "  -C<int>        : number of CPU to be used, default: 1 \n";
   cout << "  -R<rx,ry,rz>   : volume resolutions in X, Y and Z, default: 1.0 \n";
   cout << "  -I              : the input <GlobalTransform> is from the volume to the atlas.\n";
   cout << "  -f             : flip the D[x,z] and D[y,z] components of DT; not flipping PD (suppose PD flipped already) \n";
   cout << "  -H             : print this help.\n";
   cout << endl <<endl;

   cout << "------------------" <<endl;
   cout << " - by YANG, Jinzhong " <<endl;
   cout << " - last update: Nov. 27, 2007 " <<endl;
   cout << "------------------\n\n" <<endl;
   exit(0);

}


void setDefaultValues()
{
     destZ=123; //atlasZ;
     dimX = 256;
     strcpy(fTransformFile,"./GlobalTransform.Affine");
     strcpy(outDTfile,"./resultDT.dt");
     CPUs = 1;
     resX=resY=resZ = 1.0;

     //resX=resY=0.9766; resZ = 3.0;
     //getAcPcAlignmentAdjust();
}

void getParameters(int argc, char ** argv)
{   
   extern char *optarg;
   int c;

   setDefaultValues();

   if (argc <4){
      usage(argv[0]);
      return;
   } else {
     //for (int i=0; i< argc; i++)
        // cout <<i<<": "<< argv[i] <<endl;
     strcpy(DTfile,argv[1]); //DTfile input file name
     strcpy(fTransformFile,argv[2]);
     strcpy(fPD,argv[3]);

     flipZ=0;
     while ((c=getopt(argc-2,argv+2,"HP:O:Z:C:fR:X:I")) != -1) {
      //cout << " -- --  char = " << char(c) << endl;
      switch (c) {
      case 'Z': //atlas Z-number
	sscanf(optarg,"%d",&destZ); 
	//cout <<"destZ="<<destZ<<endl<<endl;
	break;
      case 'O': //define the output resulting vector field from original space to atlas space
	strcpy(outDTfile,optarg);
	break;
      case 'C': //define how many CUPs for parallel execution
	sscanf(optarg,"%d",&CPUs);  
	break;	
      case 'R': //resolution
	sscanf(optarg,"%f,%f,%f",&resX,&resY,&resZ);  
	break;	
      case 'f': //flip D[x,z] D[y,z] of DT
	flipZ=1;
	break;
      case 'X':
        sscanf(optarg,"%d",&dimX);
	break;
      case 'I':
    inversedTransform=1;
    break;
      case 'H':
      default:
	usage(argv[0]);
	break;
      }
     }
    }
    //getAcPcAlignmentAdjust();

    cout <<"DTfile="<<DTfile<<endl;
    cout <<"fTransformFile="<<fTransformFile<<endl;
    cout <<"Z of warped (destZ)="<<destZ<<endl;
    cout <<"number of CPU to be used="<<CPUs<<endl;
    cout <<"outDTfile="<<outDTfile<<endl<<endl;
    cout <<"volume resolution=("<<resX<<","<<resY<<","<<resZ<<")"<<endl<<endl;
    cout <<"Z flip="<<flipZ<<endl;
    cout <<"Inversed transform="<<inversedTransform<<endl<<endl;

    //exit(0);
 
}

/***
   * @function: extract global transfrom from file. Assume the transform is backward (no -I option)
   * @param:   fname - filename stores the global transform
   * @param:   Transform - 3x3 matrix storing the affine transform from atlas to volume
   * @param:   center1 - rotation center of the volume
   * @param:   center2 - rotation center of the atlas
***/
void Open_Transformation( const char* fname, float** Transform, float* center1, float* center2 )
{
  FILE *fp;
  int i, j ;

  double *Temp_transf ;
  float tmp;
  
  int Tn ;

  Tn = 3 ;

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

  
  if((fp=fopen(fname,"rb"))==NULL)
  {
   printf("cannot open file\n");
   return ;
  }

  fseek(fp,0L,SEEK_SET);

  fread( Temp_transf,  sizeof(double), Tn*Tn,  fp ) ; 
  fread( center1+1,      sizeof(float),  Tn,  fp ) ;
  fread( center2+1,      sizeof(float),  Tn,  fp ) ;
 
  fclose(fp);

  /* for T1 */
  for(i=0; i<Tn; i++)
   for(j=0; j<Tn; j++)
    Transform[i+1][j+1] = Temp_transf[i*Tn+j];

  //// ****transform from hammer format to regular format****////
  /****
          In hammer format, x, y is inversed. To transform ---
          1. exchange x, y in center1 and center 2
          2. exchange row 1 and row 2 in Transform
          3. exchange col 1 and col 2 in Transform
     ****/
  tmp = center1[1]; center1[1] = center1[2]; center1[2] = tmp;
  tmp = center2[1]; center2[1] = center2[2]; center2[2] = tmp;
  for (i=0; i<Tn; i++)
  {
    tmp = Transform[1][i+1]; Transform[1][i+1] = Transform[2][i+1]; Transform[2][i+1] = tmp;
  }
  
  for (i=0; i<Tn; i++)
  {
    tmp = Transform[i+1][1]; Transform[i+1][1] = Transform[i+1][2]; Transform[i+1][2] = tmp;
  }
  
  /* free */
  free(Temp_transf) ;
}


int main(int argc,char **argv) //for warp DTI
{//warp DT, this is the real Main()
   int xx,yy,zz;
   struct DTensor * nDTvol=NULL;
   
   float ** TMat, ** invTMat;
   float * center_t, *center_s;
   
   int r;
   float *w, **V, **U1;
   Matrix U;
   
   getParameters(argc, argv);
 
   cout << "Step 1. loading Diffusion Tensor...." << endl << endl;
   // rv  volDT= new DTVolume(DTfile,256,256,resX,resY,resZ); //1,1,1); 
   volDT= new DTVolume(DTfile,dimX,dimX,resX,resY,resZ); //1,1,1); 
		//,256,256,0.9766,0.9766,3.0);
   volDT->getDimXYZ(xx, yy, zz);
   if (flipZ) volDT->flipZComponent();
   //volDT->printTensors();
   //return 0;

   //exit(0);
   cout << endl <<"Step 2. loading the global transform.... \n";
   
   center_t=vector(1,3);
   center_s=vector(1,3);
   TMat = matrix(1,3,1,3);
   invTMat = matrix(1,3,1,3);
  
   if (!inversedTransform) 
   {
     Open_Transformation(fTransformFile, invTMat, center_s, center_t);
   
     TMat = matrixInverse(invTMat, 3, 3);
   }
   else
   {
     Open_Transformation(fTransformFile, TMat, center_t, center_s);
   
     invTMat = matrixInverse(TMat, 3, 3);
   }

   seeMatrix(TMat, 3, 3, "Forward Transform:");
   seeMatrix(invTMat, 3, 3, "Backward Transform:");
   seeVector(center_t, 3, "Atlas rotation center:");
   seeVector(center_s, 3, "Volume rotation center:");
   
   //delete warp;
   cout << endl << "Step 3. warp the DT volume from original->atlas" << endl << endl;
   //****// 
   /////SVD
   w=vector(1,3);
   V=matrix(1,3,1,3);
   r = svdcmp(TMat, 3, 3, w, V); //TMat=TMat*matrix(w)*V
   if (r==0) {
     cout << "bad matrix" << endl;
     seeMatrix(TMat, 3, 3, "GlobalTransform:");
     return 0;
   }
   U1=matrix(1,3,1,3);
   mulMatrixABt(TMat, V, U1, 3, 3);
   U.setValue(U1[1][1], U1[1][2], U1[1][3], 0.0f,
              U1[2][1], U1[2][2], U1[2][3], 0.0f,
              U1[3][1], U1[3][2], U1[3][3], 0.0f,
              0.0f,     0.0f,     0.0f,     1.0f);
   
   seeMatrix(U1, 3, 3, "Rotation matrix:");
   
   volDT->reorientTensor(U);

   cout << endl << "Step 4. Cast DT to grid in Atlas" << endl << endl;
   //****//

   U.setValue(invTMat[1][1], invTMat[1][2], invTMat[1][3], center_t[1],
              invTMat[2][1], invTMat[2][2], invTMat[2][3], center_t[2],
              invTMat[3][1], invTMat[3][2], invTMat[3][3], center_t[3],
              center_s[1],   center_s[2],   center_s[3],   1.0f);
              
   nDTvol = volDT->warpTensorFieldRev(U,destZ); 


   cout << "saving resulting DT file to <"<<outDTfile<<">...."<<endl;
   //****//
     //rv  volDT->saveData(outDTfile, nDTvol, atlasX*atlasY*destZ);
   volDT->saveData(outDTfile, nDTvol, xx*yy*destZ);

   cout << "new DT volume <"<<outDTfile<<"> saved." <<endl;
   //****//
   delete []nDTvol;
   delete volDT;
   
   cout << endl << "program completed." << endl <<endl;
   cout <<endl<< "~~~~~~~~~~~~~~~~" <<endl;
   cout <<"      feed back to xdr@cbmv.jhu.edu. Thanks" <<endl;
   
   free_matrix(TMat, 1, 3, 1, 3);
   free_matrix(invTMat, 1, 3, 1, 3);
   free_vector(center_t, 1, 3);
   free_vector(center_s, 1, 3);
   free_vector(w, 1, 3);
   free_matrix(V, 1, 3, 1, 3);
   free_matrix(U1, 1, 3, 1, 3);
   
   return 0;
}

