#! /usr/bin/env perl

##############################################################################
# @file  tensorgeometryregistration.pl
# @brief Warp subject image to model image, normalize, and perform group analysis.
#
# This script is used to Warp subject image to model image and for
# normalization and group analysis. The output is the deformation field from
# Model to Subject space, called "Field.1.img.Mdl2Subj", and a warped subject
# image (Warped.SubjSeg.img, already in model space).
#
# 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>
##############################################################################

use Getopt::Std;
use File::Copy qw(mv cp);

use SBIA::DTIDROID::Basis qw(execute);

#get arguments list
$argc=@ARGV;

$useFSL=1;

if ($argc < 1){
  &showUsage ;
}

my $opt_string = 'hTFS:M:i:l:R:r:m:c:X:E:D:V:v:A';

getopts("$opt_string")  || &showUsage ;

if($opt_h){
    &showUsage;
}

#define argument list and get the arguments


# get Input files:
if($opt_S){
   $InputFile=$opt_S;
}else{
   print "\n\n Must input a subject DT image\n\n";
   &showUsage;
}

if($opt_M){
   $MdlFile =$opt_M;
}else{
   print "\n\n Must input a model DT image\n\n";
   &showUsage;
}


#iteration number
$iteration = 50;
if($opt_i){
  $iteration=$opt_i;
}

#iteration number for orientation features
$iter_orient = 0;
if($opt_l){
  $iter_orient=$opt_l;
}

# step ratio
$step_ration = 0.020;
if($opt_R){
  $step_ration=$opt_R;
}

# neighborhood search size, default: 12 for Low resolution,
# 10 for Middl resolution 8 for High resolutions
@search_size = (12,10,8);
if($opt_r){
  $sizes = $opt_r;
  @search_size = split(",", $sizes);
}

# option: -m
$m_degree=0.75;
if($opt_m){
   $m_degree=$opt_m;
}

# option: -c
# paramter used for the first parameter of option -c
$smooth = 0.5;
if($opt_c){
   $smooth=$opt_c;
}

#option: -D
$mdl_path = "." ;
if($opt_D){
   $mdl_path=$opt_D;
}
$mdl_path=$mdl_path."/";

# option: -X
@image_size=(256,256,70);
if($opt_X){
   $img_size=$opt_X;
   @image_size=split(",",$img_size);
}

# option: -E
# parameter used for the second paramter of option -E
$edge_smooth=1;
if($opt_E){
   $edge_smooth=$opt_E;
}

# option: -V
# voxel resolution for image model
@model_voxel=(0.9375,0.9375,2.5);
if($opt_V){
   $voxels = $opt_V;
   @model_voxel = split(",", $voxels);
}

# option: -v
# voxel resolution for image subject
@subject_voxel=(0.9375,0.9375,2.5);
if($opt_v){
   $voxels = $opt_v;
   @subject_voxel = split(",", $voxels);
}

# option: -F
# use FSL for affine registration
$useFSL = 1;
if($opt_F){
  print "turning off useFSL\n";
  $useFSL = 0;
}
# option: -A
# skip affine registration
$skipAffine = 0;
if($opt_A){
  print "turning off affine registration\n";
  $skipAffine = 1;
}


print "\n\nWarp sujbect to Model by HAMMER\n\n";
print "Input information: \n";
print "\tinput subject segmented image is $InputFile\n";
print "\tinput subject original image is $InputFile\n";
print "\tinput Model image is $MdlFile\n";
print "\tdiretory where subject image stored: $mdl_path\n";
print "\timage size on XYZ dimension: @image_size\n";
print "\titeration number: $iteration\n";
print "\titeration number for orientation features: $iter_orient\n";
print "\tstep ratio: $step_ration\n";
print "\tneighborhood search size: ($search_size[0], $search_size[1], $search_size[2])\n";
print "\tsimilarity degree: $m_degree\n";
print "\tsmoothness factor for deformation field: $smooth\n";
print "\tTimes for smoothing deformation fields: $edge_smooth\n";
if ($skipAffine == 1){
  print "\tSkipping affine registration\n";
}
elsif( $useFSL == 1){
  print "\tUsing FSL Flirt for afine registration\n";
}
else{
  print "\tNot using FSL Flirt for afine registration\n";
}

print "\n\n\n";

# affine transformation: for upsample MdlFile
$MdlFile_withPath = $mdl_path."".$MdlFile;

open (IN, "$InputFile") || die "\nError open: $InputFile\n\n";
close(IN);
open (IN, "$MdlFile_withPath") || die "\nError open: $MdlFile_withPath\n\n";
close(IN);

# affine transformation: for upsample MdlFile
if ($skipAffine) {
  $File_Rigid = $InputFile;
} else {
  print "\n\nLinear align two images\n";
  $File_Rigid = "$InputFile.rig";
  $cmd = "rigidlyregisterdt $InputFile $MdlFile_withPath $File_Rigid 12"
      . " $image_size[0] $image_size[1] $image_size[2]"
      . " $subject_voxel[0] $subject_voxel[1] $subject_voxel[2]"
      . " $image_size[0] $image_size[1] $image_size[2]"
      . " $model_voxel[0] $model_voxel[1] $model_voxel[2]";
  if ($useFSL) { $cmd = $cmd . " Y"; }
  else         { $cmd = $cmd . " N"; }
  execute("$cmd", verbose => 1);
  print "Done rigid transformation\n\n";
}

$HGH_FnlID=$iteration -1;

#remove previous results
unlink glob '*sample?.*gz';

#Low resoultion: 4
print "\n";
print "*******************************************************************\n";
print "Low resoultion:\n";
print "*******************************************************************\n";
print "Warping '$File_Rigid' in low-resolution\n";
execute("TensorStructRegistration $MdlFile $File_Rigid -R$step_ration -D$mdl_path -i$iteration -r$search_size[0] -m$m_degree -s4 -c0.5 -B -E1 -p$smooth,$edge_smooth -X$image_size[0]", verbose => 1);
execute("splitfile DeformationField.float.img NoUse 8 Field.4.img", verbose => 1);
unlink glob 'Res.img.sample4.*';

#Mid resoultion: 2
print "\n";
print "*******************************************************************\n";
print "Mid resoultion:\n";
print "*******************************************************************\n";
print "Warping '$File_Rigid' in mid-resolution\n";
execute("TensorStructRegistration $MdlFile $File_Rigid -R$step_ration -D$mdl_path -i$iteration -r$search_size[1] -m$m_degree -s2 -c0.5 -B -E1 -p$smooth,$edge_smooth -F -X$image_size[0]", verbose => 1);
execute("splitfile DeformationField.float.img NoUse 8 Field.2.img", verbose => 1);
unlink glob 'Res.img.sample2.*';

#High resoultion: 1
print "\n";
print "*******************************************************************\n";
print "high resoultion:\n";
print "*******************************************************************\n";
print "Warping '$File_Rigid' in high-resolution\n";
execute("TensorStructRegistration $MdlFile $File_Rigid -R$step_ration -D$mdl_path -i$iteration -r$search_size[2] -m$m_degree -s1 -c0.5 -B -E1 -p$smooth,$edge_smooth -F -X$image_size[0]", verbose => 1);

#****************************************************
# generate deformation field from model to subject
#****************************************************
execute("splitfile DeformationField.float.img NoUse 8 Field.1.img", verbose => 1);
execute("transform3dvectorfield Field.1.img Field.1.img.trans GlobalTransformation.Affine -v$image_size[0]", verbose => 1);
mv('Field.1.img.trans', 'Field.1.img.Mdl2Subj');

#*****************************************************
# generate warped subject image to model space
#*****************************************************

for ($count=0; $count<$iter_orient; $count++) {
  # transform vector field for warpDT
  execute("reversedeformationfield Field.1.img.Mdl2Subj Field.1.img.Subj2Mdl -X $image_size[0] -z $image_size[2]", verbose => 1);
  execute("hammer2xdr Field.1.img.Mdl2Subj Field.1.img.Mdl2Subj.xdr $image_size[0] $image_size[1] $image_size[2]", verbose => 1);
  execute("hammer2xdr Field.1.img.Subj2Mdl Field.1.img.Subj2Mdl.xdr $image_size[0] $image_size[1] $image_size[2]", verbose => 1);
  execute("dtiPD 1 $image_size[0] $image_size[1] $image_size[2] $InputFile Subj.pd1.img 1 1 0", verbose => 1);
  # warp vector field
  execute("nWarpDTV5 $InputFile Field.1.img.Subj2Mdl.xdr Field.1.img.Mdl2Subj.xdr Subj.pd1.img -O Warped.Subj.dt -X $image_size[0] -Z $image_size[2] -R $subject_voxel[0],$subject_voxel[1],$subject_voxel[2] -N 0", verbose => 1);
  # perform registration
  execute("TensorStructRegistration $MdlFile Warped.Subj.dt -R$step_ration  -D$mdl_path -i1 -r$search_size[2] -m$m_degree -s1 -c0.5 -B -E3 -p$smooth,$edge_smooth -X$image_size[0] -o 0.4", verbose => 1);
  execute("splitfile DeformationField.float.img NoUse 8 Field.1.img", verbose => 1);
  execute("combinetwodeformationfields Field.1.img.Mdl2Subj Field.1.img Field.1.img.Mdl2Subj -d$image_size[0],$image_size[1],$image_size[2] -e$image_size[0],$image_size[1],$image_size[2] -S", verbose => 1);
}

# transform vector field for warpDT
execute("reversedeformationfield Field.1.img.Mdl2Subj Field.1.img.Subj2Mdl -X $image_size[0] -z $image_size[2]", verbose => 1);
execute("hammer2xdr Field.1.img.Mdl2Subj Field.1.img.Mdl2Subj.xdr $image_size[0] $image_size[1] $image_size[2]", verbose => 1);
execute("hammer2xdr Field.1.img.Subj2Mdl Field.1.img.Subj2Mdl.xdr $image_size[0] $image_size[1] $image_size[2]", verbose => 1);
execute("dtiPD 1 $image_size[0] $image_size[1] $image_size[2] $InputFile Subj.pd1.img 1 1 0", verbose => 1);
# warp diffusion image
execute("nWarpDTV5 $InputFile Field.1.img.Subj2Mdl.xdr Field.1.img.Mdl2Subj.xdr Subj.pd1.img -O Warped.Subj.dt -X $image_size[0] -Z $image_size[2] -R $subject_voxel[0],$subject_voxel[1],$subject_voxel[2] -N 0", verbose => 1);

#***************************************************
# remove unwanted files
#****************************************************
unlink glob 'Res.img.sample1.*';
unlink 'GlobalTransformation.Affine';
unlink 'DeformationField.float.img';
unlink glob 'MdlImg_*';
unlink glob 'ObjImg_*';
unlink 'Field.2.img';
unlink 'Field.4.img';
unlink 'Field.1.img';
unlink '$File_Rigid' unless $skipAffine;
unlink 'optimized.d1';
unlink glob 'Field.1.img.*.xdr';
unlink 'Field.1.img.Subj2Mdl';
unlink 'Subj.pd1.img';
unlink 'NoUse';

# *************************************************
# function:  showUsage
# arguments: None
# *************************************************
sub showUsage{
  print "\n\n";
  print "This program is used to warp subject DT image to modle space. \n";
  print "The output is the deformation field from Model to Subject(Field.1.img.Mdl2Subj),\n";
  print "and a warped subject image (Warped.Subj.dt), already in model space.\n";
  print "\n";
  print "\n";
  print "Usage: $0 [options]\n";

  print "\t-h                          : display this (help) message\n";
  print "\t-M <string>                 : model image file name \n";
  print "\t-D <string>                 : directory path where model file is stored, default: current working directory\n";
  print "\t-S <string>                 : subject image file name, including path if not at current working direcotory\n";
  print "\t-X <int>,<int>,<int>        : image size on X,Y,Z dimension, default: <256,256,70>\n";
  print "\t-i <int>                    : iteration number, default: 50\n";
  print "\t-l <int>                    : iteration number for orientation features, default: 0 (without this feature)\n";
  print "\t                              keep the number small to save time for registration.\n";
  print "\t-R <float>                  : step ratio, default: 0.02\n";
  print "\t-r <int, int, int>          : neighborhood search size, different at each resolution.\n";
  print "\t                              input as < value_at_low, value_at_middle, value_at_high>\n";
  print "\t                              default: <12, 10, 8> mm/voxels\n";
  print "\t-V <float, float, float>    : voxel resolution for image model, default: <0.9375, 0.9375, 2.5>\n";
  print "\t-v <float, float, float>    : voxel resolution for subject, default: <0.9375, 0.9375, 2.5>\n";
  print "\t-m <float>                  : the similarity degree, default: 0.75\n";
  print "\t-c <float>                  : smoothness factor for deformation field, [0,1], 1: very smooth; default: 0.5\n";
  print "\t-E <int>                    : times for smoothing deformation fields, default: 1\n";
  print "\t-F                          : set if you would like to NOT use FSL flirt for affine Registations\n";
  print "                                if not set fsl flirt will be used for affine registration\n";

  print "\nNotes:\n";
  print "\t    This program only takes voxel-wise DT images as input. Each voxel contains 6 components [Dxx Dyy Dzz Dxy Dxz Dyz] in float format.\n\n";

  print "Example:\n";
  print "\tusing default value: \n";
  print "\t\t$0 -M modle.dt -S subject.dt\n\n";
  print "\twith options:\n";
  print "\t\t$0 -M modle.dt -S subject.dt -D/sbia/model -X256,256,70 -i40 -r10,8,6 -R0.07 -m0.8 -c0.3 -E2 \n\n";

  #print "References:\n";
  #print "\t1. Shen, D. and C. Davatzikos, \"Very high-resolution morphometry using mass-preserving\n";
  #print "\t   deformations and HAMMER elastic registration\". Neuroimage, 2003. 18(1): p. 28-41.\n";
  #print "\t2. Shen, D. and C. Davatzikos, \"HAMMER: Hierarchical attribute matching mechanism\n";
  #print "\t   for elastic registration\". IEEE Trans on Med. Imaging, 2002. 21(11): p. 1421-1439.\n\n";
  print "\n\n";

  exit(1);
}
