#! /usr/bin/env perl
#
# Extract the white matter left and right hemispheres for the extraction
# of the white surfaces.
#
# Oliver Lyttelton oliver@bic.mni.mcgill.ca
#
# Copyright Alan C. Evans
# Professor of Neurology
# McGill University
#

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


my($Help, $Usage, $me);
my(@opt_table, $tmpdir);

$me = &basename($0);

my $verbose = 0;
my $clobber = 0;
   
$Help = <<HELP;
|    $me extracts the white matter left and right hemispheres for
|    the extraction of the white surfaces.
| 
| Problems or comments should be sent to: oliver\@bic.mni.mcgill.ca
HELP

$Usage = "Usage: $me [options] classified.mnc t1.mnc brain_mask.mnc hemi_model_dir output_left_hemi.mnc output_right_hemi.mnc\n".
         "       $me -help to list options\n\n";

@opt_table = (
        ["-verbose","const","1", \$verbose,
        "be verbose" ],
        ["-clobber", "const","1", \$clobber,
        "clobber existing check features" ],  
   );

# Check arguments
&Getopt::Tabular::SetHelp($Help, $Usage);
&GetOptions (\@opt_table, \@ARGV) || exit 1;

die $Usage if(@ARGV != 8);

my $classify = shift(@ARGV);
my $t1 = shift(@ARGV);
my $brain_mask = shift(@ARGV);
my $user_mask = shift(@ARGV);
my $t1_tal_xfm = shift(@ARGV);
my $hemi_model_dir= shift(@ARGV);
my $output_left_hemi = shift(@ARGV);
my $output_right_hemi = shift(@ARGV);

# make tmpdir
$tmpdir = &tempdir( "$me-XXXXXXXX", TMPDIR => 1, CLEANUP => 1 );


my $classify_masked = "${tmpdir}/classify_masked.mnc";
my $wm_crop_left = "${tmpdir}/wm_crop_left.mnc";
my $wm_crop_right  = "${tmpdir}/wm_crop_right.mnc";
my $wm_defrag_left  = "${tmpdir}/wm_defrag_left.mnc";
my $wm_defrag_right  = "${tmpdir}/wm_defrag_right.mnc";
my $wm_add_back_left  = "${tmpdir}/wm_add_back_left.mnc";
my $wm_add_back_right    = "${tmpdir}/wm_add_back_right.mnc";
my $slide_left_xfm =  "${hemi_model_dir}/slide_left.xfm";
my $slide_right_xfm =  "${hemi_model_dir}/slide_right.xfm";

# for flipping the right WM mask to look like a left
my $wm_preflip_right = "${tmpdir}/wm_preflip_right.mnc";
my $flip_right_xfm =  "${hemi_model_dir}/flip_right.xfm";

# create a white matter mask from the current brain_mask
my $wm_mask = "${tmpdir}/wm_mask.mnc";
white_matter_mask( $classify, $t1, $brain_mask, $wm_mask, $tmpdir );

# apply brain mask to classified image

if( -e $user_mask ) {
  # apply the user mask if one is supplied. The user mask is in native
  # space so transform it to stereotaxic space first. The user mask is
  # applied only on white voxels, not gray.

  my $user_mask_native = "${tmpdir}/user_mask_native.mnc";
  &do_cmd( "mincresample", "-clobber", "-like", $t1,
           "-transform", $t1_tal_xfm, "-nearest_neighbour",
           $user_mask, $user_mask_native );

  &do_cmd( "minccalc", "-clobber", "-expr", 
           "if(A[1]>0.5){if(A[0]>2.5&&A[2]<0.5){out=0;}else{out=A[0];}}else{out=0;}",
           $classify, $wm_mask, $user_mask_native, $classify_masked );
  unlink( $user_mask_native );

} else {

  &do_cmd( "minccalc", "-clobber", "-expr", 
           "if(A[1]>0.5){out=A[0];}else{out=0;}",
           $classify, $wm_mask, $classify_masked );
}

# create left and right masks for hemispheres (this works with any template,
# assuming the centerline voxel is at x=0).

my $filled = "${tmpdir}/filled.mnc";
my $hemi_tmp = "${tmpdir}/hemi_tmp.mnc";
my $hemi_left_mask = "${tmpdir}/hemi_left_mask.mnc";
my $hemi_right_mask = "${tmpdir}/hemi_right_mask.mnc";

&do_cmd( "minccalc", "-clobber", "-byte", "-expression", "out=1;", $classify, $filled );
my $xlen = `mincinfo -dimlength xspace $classify`;
chomp( $xlen );
$xlen = ( $xlen + 1 ) / 2;

&do_cmd( "mincresample", "-clobber", "-xnelements", $xlen, $filled, $hemi_tmp );
&do_cmd( "mincresample", "-clobber", "-like", $classify, $hemi_tmp, $hemi_left_mask );

&do_cmd( "mincresample", "-clobber", "-xstart", 0, "-xnelements", $xlen,
         $filled, $hemi_tmp );
&do_cmd( "mincresample", "-clobber", "-like", $classify, $hemi_tmp, $hemi_right_mask );

# retain only white matter in left and right hemispheres.

&do_cmd( "minccalc", "-clobber", "-byte", "-expr", "out=(abs(A[0]-3)<0.45)&&A[1];",
         $classify_masked, $hemi_left_mask, $wm_crop_left );

&do_cmd( "minccalc", "-clobber", "-byte", "-expr", "out=(abs(A[0]-3)<0.45)&&A[1];",
         $classify_masked, $hemi_right_mask, $wm_crop_right );

# remove loose bits of disconnected white matter.

&do_cmd( "mincdefrag", $wm_crop_left, $wm_defrag_left, "1", "27", "100000" );

&do_cmd( "mincdefrag", $wm_crop_right, $wm_defrag_right, "1", "27", "100000" );

# put back removed bits of white matter in one hemisphere to the other one.

&do_cmd( "minccalc", "-clobber", "-byte", "-expr", "out=A[0]||(A[1]&&!A[2]);",
         $wm_defrag_left, $wm_crop_right, $wm_defrag_right, $wm_add_back_left );

&do_cmd( "minccalc", "-clobber", "-byte", "-expr", "out=A[0]||(A[1]&&!A[2]);",
         $wm_defrag_right, $wm_crop_left, $wm_defrag_left, $wm_add_back_right );

# center the masks.

&do_cmd( "mincresample", "-clobber", "-like", $wm_add_back_left, "-transform",
         $slide_right_xfm, $wm_add_back_left, $output_left_hemi );

&do_cmd( "mincresample", "-clobber", "-like", $wm_add_back_right, "-transform",
         $slide_left_xfm, $wm_add_back_right, $wm_preflip_right );

# mirror the right mask to look like a left hemi.

&do_cmd( "mincresample", "-clobber", "-like", $wm_preflip_right, "-transform",
         $flip_right_xfm, $wm_preflip_right, $output_right_hemi );

# final removal of disconnected bits of white matter in each hemisphere.

&do_cmd( "mincdefrag", $output_left_hemi, $output_left_hemi, "1", "6", "100000" );

&do_cmd( "mincdefrag", $output_right_hemi, $output_right_hemi, "1", "6", "100000" );


# create a better white matter mask from the current brain_mask

sub white_matter_mask {

  my $cls = shift;
  my $t1 = shift;
  my $brain_mask = shift;
  my $white_mask = shift;
  my $tmpdir = shift;

  # consider only 5mm around the perimeter of current brain mask
  # (well, that's 5 layers of voxels, not 5mm).

  my $mask_eroded = "${tmpdir}/mask_eroded.mnc";
  &do_cmd( "dilate_volume", $brain_mask, $mask_eroded, 0, 6, 5 );

  # compute t1 mean and variance of classified white matter inside the
  # eroded brain mask (avoid high intensity voxels of skull).

  my $cls_eroded = "${tmpdir}/cls_eroded.mnc";
  &do_cmd( "minccalc", "-clobber", "-expression", "if(A[1]>0.5){out=A[0];}else{out=0;}",
           $cls, $mask_eroded, $cls_eroded );
  my $white_mean = `mincstats -quiet -mask $cls_eroded -mask_binvalue 3 -mean $t1`;
  my $white_std = `mincstats -quiet -mask $cls_eroded -mask_binvalue 3 -std $t1`;
  my $white_thresh = $white_mean + 2.0 * $white_std;

  # mask out the t1 high intensity voxels in the eroded region of the
  # brain mask. Keep only those classified white voxels to remove.

  my $expr = "if(A[0]>2.5&&A[2]<0.5&&A[3]>0.5&&A[1]>$white_thresh){out=10;}else{out=0;}";

  &do_cmd( "minccalc", "-clobber", "-expression", $expr, $cls, $t1, $mask_eroded,
           $brain_mask, $cls_eroded );

  # blur the bright voxels to remove to diffuse effect to immediate neighbours.

  my $cls_blur = "${tmpdir}/cls_blur.mnc";
  my $cls_blur_prefix = "${tmpdir}/cls";
  &do_cmd( "mincblur", "-clobber", "-fwhm", 2, $cls_eroded, $cls_blur_prefix );

  # remove the bright voxels and their immediate neighbours.

  &do_cmd( "minccalc", "-clobber", "-byte", "-expression",
           "if(A[0]>2.5&&A[2]>1.5){out=1;}else{if(A[1]>0.5){out=A[0];}else{out=0;}}",
           $cls, $brain_mask, $cls_blur, $white_mask );

  # remove loose bits of white matter.

  &do_cmd( "mincdefrag", $white_mask, $white_mask, 3, 6 );

  # blur the masked classified image.

  &do_cmd( "mincblur", "-clobber", "-fwhm", 5, $white_mask, $cls_blur_prefix );

  # threshold at 1.5 (csf-gm border) to obtain the final white matter mask.
  &do_cmd( "minccalc", "-clobber", "-byte", "-expression", "if(A[0]>1.5||A[1]>0.5){out=1;}else{out=0;}",
           $cls_blur, $mask_eroded, $white_mask );
}


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