#! /usr/bin/env perl
#
# non-linear fitting using parameters optimised by Steve Robbins,
# using a brain mask for the source and the target.
#
# Vladimir S. Fonov vfonov@bic.mni.mcgill.ca
# Claude Lepage - claude@bic.mni.mcgill.ca
# Andrew Janke - rotor@cmr.uq.edu.au
# Center for Magnetic Resonance
# The University of Queensland
# http://www.cmr.uq.edu.au/~rotor
#
# Copyright Andrew Janke, The University of Queensland.
# Permission to use, copy, modify, and distribute this software and its
# documentation for any purpose and without fee is hereby granted,
# provided that the above copyright notice appear in all copies.  The
# author and the University of Queensland make no representations about the
# suitability of this software for any purpose.  It is provided "as is"
# without express or implied warranty.

use strict;
use warnings "all";
use Getopt::Tabular;
use File::Basename;
use File::Temp qw/ tempdir /;

# default minctracc parameters
my @def_minctracc_args = (
#   '-debug',
   '-clobber',
   '-nonlinear', 'corrcoeff',
   '-weight', 1,
   '-stiffness', 1,
   '-similarity', 0.3,
   '-sub_lattice', 6,
   );

my @conf = (

   {'step'         => 8,
    'blur_fwhm'    => 2,
    'iterations'   => 20,
    },

   {'step'         => 6,
    'blur_fwhm'    => 2,
    'iterations'   => 20,
    },

   {'step'         => 4,
    'blur_fwhm'    => 2,
    'iterations'   => 10,
    },

   {'step'         => 2,
    'blur_fwhm'    => 1,
    'iterations'   => 10,
    },

   );

my($Help, $Usage, $me);
my(@opt_table, %opt, $source, $target, $outxfm, $outfile, @args, $tmpdir);

$me = &basename($0);
%opt = (
   'verbose'   => 0,
   'clobber'   => 0,
   'fake'      => 0,
   'normalize' => 0,
   'init_xfm'  => undef,
   'source_mask' => undef,
   'target_mask' => undef,
   'level'      => 4,
   'work_dir'    => undef,
   'threshold'  => 0
   );

$Help = <<HELP;
| $me does hierachial non-linear fitting between two files
|    you will have to edit the script itself to modify the
|    fitting levels themselves
| 
| Problems or comments should be sent to: vfonov\@bic.mni.mcgill.ca
HELP

$Usage = "Usage: $me [options] source.mnc target.mnc output.xfm [output.mnc]\n".
         "       $me -help to list options\n\n";

@opt_table = (
   ["-verbose", "boolean", 0, \$opt{verbose},
      "be verbose" ],
   ["-clobber", "boolean", 0, \$opt{clobber},
      "clobber existing check files" ],
   ["-fake", "boolean", 0, \$opt{fake},
      "do a dry run, (echo cmds only)" ],
   ["-init_xfm", "string", 1, \$opt{init_xfm},
      "initial transformation (default identity)" ],
   ["-source_mask", "string", 1, \$opt{source_mask},
      "source mask to use during fitting" ],
   ["-target_mask", "string", 1, \$opt{target_mask},
      "target mask to use during fitting" ],
   ["-level", "float", 1, \$opt{level},
      "Perform fitting down to step , minimum 2" ], 
   ["-work_dir", "string", 1, \$opt{work_dir},
      "Directory to keep blurred files" ],
   ["-threshold", "float", 1, \$opt{threshold},
      "Threshold for the minctracc, " ], 
  );

delete $ENV{MINC_COMPRESS} if $ENV{MINC_COMPRESS};

# Check arguments
&Getopt::Tabular::SetHelp($Help, $Usage);
&GetOptions (\@opt_table, \@ARGV) || exit 1;
die $Usage if(! ($#ARGV == 2 || $#ARGV == 3));
$source = shift(@ARGV);
$target = shift(@ARGV);
$outxfm = shift(@ARGV);
$outfile = (defined($ARGV[0])) ? shift(@ARGV) : undef;

my $original_source=$source;
# check for files
die "$me: Couldn't find input file: $source\n\n" if (!-e $source);
die "$me: Couldn't find input file: $target\n\n" if (!-e $target);
if(-e $outxfm && !$opt{clobber}){
   die "$me: $outxfm exists, -clobber to overwrite\n\n";
   }
if(defined($outfile) && -e $outfile && !$opt{clobber}){
   die "$me: $outfile exists, -clobber to overwrite\n\n";
   }

my $mask_warning = 0;
if( !defined($opt{source_mask}) ) {
  $mask_warning = 1;
} else {
  if( !-e $opt{source_mask} ) {
    $mask_warning = 1;
  }
}
if( !defined($opt{target_mask}) ) {
  $mask_warning = 1;
} else {
  if( !-e $opt{target_mask} ) {
    $mask_warning = 1;
  }
}

# make tmpdir
$tmpdir = &tempdir( "$me-XXXXXXXX", TMPDIR => 1, CLEANUP => 1 );
$opt{work_dir}=$tmpdir unless $opt{work_dir};

# set up filename base
my($i, $s_base, $t_base, $tmp_xfm, $tmp_source, $tmp_target, $prev_xfm);

$s_base = &basename($source);
$s_base =~ s/\.mnc(.gz)?$//;
$s_base = 's_'.$s_base;

$t_base = &basename($target);
$t_base =~ s/\.mnc(.gz)?$//;
$t_base = 't_'.$t_base;

# a fitting we shall go...
my $first=1;
for ($i=0; $i<=$#conf; $i++){
   
   next if $conf[$i]{step}<$opt{level}; # skip unneeded steps
   # set up intermediate files
   $tmp_xfm = "$tmpdir/$s_base\_$i.xfm";
   $tmp_source = "$opt{work_dir}/$s_base\_$conf[$i]{blur_fwhm}";
   $tmp_target = "$opt{work_dir}/$t_base\_$conf[$i]{blur_fwhm}";
   
   print STDOUT "-+-[$i]\n".
                " | step:           $conf[$i]{step}\n".
                " | blur_fwhm:      $conf[$i]{blur_fwhm}\n".
                " | iterations:     $conf[$i]{iterations}\n".
                " | source:         $tmp_source\n".
                " | target:         $tmp_target\n".
                " | xfm:            $tmp_xfm\n".
                "\n";
   
   # blur the source and target files if required.
   # mask the source and target provided both masks are supplied.
   if(!-e "$tmp_source\_blur.mnc"){
      my $source_masked = $source;
      &do_cmd('mincblur', '-no_apodize', '-fwhm', $conf[$i]{blur_fwhm},
              $source_masked, $tmp_source);
   }
   if(!-e "$tmp_target\_blur.mnc"){
      my $target_masked = $target;
      &do_cmd('mincblur', '-no_apodize', '-fwhm', $conf[$i]{blur_fwhm},
              $target_masked, $tmp_target);
   }
   
   # set up registration
   @args = ('minctracc',  @def_minctracc_args,
            '-iterations', $conf[$i]{iterations},
            '-step', $conf[$i]{step}, $conf[$i]{step}, $conf[$i]{step},
            '-lattice_diam', $conf[$i]{step} * 3, 
                             $conf[$i]{step} * 3, 
                             $conf[$i]{step} * 3,
            '-threshold',$opt{threshold},$opt{threshold});
   if($conf[$i]{step}<4)
   {
     push(@args,'-no_super');
   }
   # transformation
   if($first) {
      push(@args, (defined $opt{init_xfm}) ? ('-transformation', $opt{init_xfm}) : '-identity')
   } else {
      push(@args, '-transformation', $prev_xfm);
   }

   # masks (even if the blurred image is masked, it's still preferable
   # to use the mask in minctracc)
   push(@args, '-source_mask', $opt{source_mask} ) if defined($opt{source_mask});
   push(@args, '-model_mask', $opt{target_mask}) if defined($opt{target_mask});
   
   # add files and run registration
   push(@args, "$tmp_source\_blur.mnc", "$tmp_target\_blur.mnc", $tmp_xfm);
   &do_cmd(@args);
   
   $prev_xfm = $tmp_xfm;
   $first=0;
}


# a hack to copy the xfm file and all the internals
do_cmd('param2xfm',"$tmpdir/identity.xfm");
do_cmd('xfmconcat',"$tmpdir/identity.xfm",$tmp_xfm,$outxfm );

# resample if required
if(defined($outfile)){
   print STDOUT "-+- creating $outfile using $outxfm\n".
   &do_cmd('mincresample', '-clobber', '-like', $target,
           '-transformation',$prev_xfm, $original_source, $outfile);
}


sub do_cmd { 
   print STDOUT "@_\n" if $opt{verbose};
   if(!$opt{fake}){
      system(@_) == 0 or die;
   }
}
       
