#! /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'         => 32,
    'blur_fwhm'    => 16,
    'iterations'   => 1,
    },

   {'step'         => 16,
    'blur_fwhm'    => 8,
    'iterations'   => 1,
    },

   {'step'         => 12,
    'blur_fwhm'    => 6,
    'iterations'   => 1,
    },

   {'step'         => 8,
    'blur_fwhm'    => 4,
    'iterations'   => 1,
    },

   {'step'         => 8,
    'blur_fwhm'    => 3,
    'iterations'   => 1,
    },

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

   {'step'         => 8,
    'blur_fwhm'    => 1,
    'iterations'   => 1,
    },

   );

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
   );

$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: rotor\@cmr.uq.edu.au
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)" ],
   ["-normalize", "boolean", 0, \$opt{normalize},
      "do intensity normalization on source to match intensity of target" ],
   ["-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 with step from level*2 to level, minimum 2" ], 
  );

# 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;

# 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 );

# 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)?$//;
$t_base = &basename($target);
$t_base =~ s/\.mnc(.gz)?$//;

# Run inormalize if required. minctracc likes it better when the
# intensities of the source and target are similar.

my $original_source = $source;
if( $opt{normalize} ) {
  my $inorm_source = "$tmpdir/${s_base}_inorm.mnc";
  &do_cmd( 'inormalize', '-clobber', '-model', $target, $source, $inorm_source );
  $source = $inorm_source;
}

# a fitting we shall go...
my $first=1;
for ($i=0; $i<=$#conf; $i++){
   
   next if $conf[$i]{step}<$opt{level} ||  $conf[$i]{step}>$opt{level}*2; # skip unneeded steps
   # set up intermediate files
   $tmp_xfm = "$tmpdir/$s_base\_$i.xfm";
   $tmp_source = "$tmpdir/$s_base\_$conf[$i]{blur_fwhm}";
   $tmp_target = "$tmpdir/$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;
      if( defined($opt{source_mask}) and defined($opt{target_mask}) ) {
        if( -e $opt{source_mask} ) {
          $source_masked = "${tmp_source}_masked.mnc";
          if(!-e $source_masked ) {
            &do_cmd( 'minccalc', '-clobber',
                     '-expression', 'if(A[1]>0.5){out=A[0];}else{out=A[1];}',
                     $source, $opt{source_mask}, $source_masked );
          }
        }
      }
      &do_cmd('mincblur', '-no_apodize', '-fwhm', $conf[$i]{blur_fwhm},
              $source_masked, $tmp_source);
   }
   if(!-e "$tmp_target\_blur.mnc"){
      my $target_masked = $target;
      if( defined($opt{source_mask}) and defined($opt{target_mask}) ) {
        if( -e $opt{target_mask} ) {
          $target_masked = "${tmp_target}_masked.mnc";
          if(!-e $target_masked ) {
            &do_cmd( 'minccalc', '-clobber',
                     '-expression', 'if(A[1]>0.5){out=A[0];}else{out=A[1];}',
                     $target, $opt{target_mask}, $target_masked );
          }
        }
      }
      &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);
   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;
   }
}
       
