#! /bin/tcsh -f
#
# orientLAS
#
# Convert images (and gradient vectors, if present) to LAS orientation
#
# Original Author: Anastasia Yendiki
#
# Copyright © 2021 The General Hospital Corporation (Boston, MA) "MGH"
#
# Terms and conditions for use, reproduction, distribution and contribution
# are found in the 'FreeSurfer Software License Agreement' contained
# in the file 'LICENSE' found in the FreeSurfer distribution, and here:
#
# https://surfer.nmr.mgh.harvard.edu/fswiki/FreeSurferSoftwareLicense
#
# Reporting: freesurfer@nmr.mgh.harvard.edu
#
#

if ($#argv < 2) then
  echo $argv | grep -q " --help"
  if (! $status || $#argv == 0) then
    goto usage
  else
    echo "ERROR: not enough arguments"
    exit 1
  endif
endif

set in = $1
set out = $2
if ($#argv == 3 && $3 == "--check") then
  set docheck = 1
else
  set docheck = 0
endif

if !(-e $in) then
  echo "ERROR: $in does not exist"
  exit 1
endif

source $FREESURFER_HOME/sources.csh

mri_info $in | grep type | grep -q nii
if $status then
  echo "ERROR: input file must be in NIfTI format"
  exit 1
endif

set inorient = `mri_info $in | grep Orientation | awk '{print $3}'`
set indeterm = `mri_info $in | grep determinant | awk '{print $3}'`
echo "INFO: input image orientation is $inorient"
echo "INFO: input image determinant is $indeterm"

# Find how input axes need to be inverted or swapped to convert to LAS
set sign = `echo $inorient | sed "s/[LAS]/+\ /g; s/[RPI]/-\ /g"`
set order = `echo $inorient | sed "s/[LR]/1\ /; s/[AP]/2\ /; s/[IS]/3\ /"`

# Input volume dimensions
set dims = `mri_info --dim $in`

# Input direction cosines
set cdc =  `mri_info --cdc $in`
set rdc =  `mri_info --rdc $in`
set sdc =  `mri_info --sdc $in`

set xyzR = ($cdc[1] $rdc[1] $sdc[1])
set xyzA = ($cdc[2] $rdc[2] $sdc[2])
set xyzS = ($cdc[3] $rdc[3] $sdc[3])

# Output direction cosines
foreach k (1 2 3)
  if ($order[$k] == 1) then
    set nx_new = $dims[$k]
    set cdc_new = `echo $xyzR[$k] $xyzA[$k] $xyzS[$k] \
                   | awk -v sgn=$sign[$k]1 '{print sgn*$1, sgn*$2, sgn*$3}'`
  else if ($order[$k] == 2) then
    set ny_new = $dims[$k]
    set rdc_new = `echo $xyzR[$k] $xyzA[$k] $xyzS[$k] \
                   | awk -v sgn=$sign[$k]1 '{print sgn*$1, sgn*$2, sgn*$3}'`
  else if ($order[$k] == 3) then
    set nz_new = $dims[$k]
    set sdc_new = `echo $xyzR[$k] $xyzA[$k] $xyzS[$k] \
                   | awk -v sgn=$sign[$k]1 '{print sgn*$1, sgn*$2, sgn*$3}'`
  endif
end

# Input origin
set cras = `mri_info --cras $in`
set dx = `mri_info --cres $in`
set dy = `mri_info --rres $in`
set dz = `mri_info --sres $in`

# Output origin
set cras_new = ()
foreach k (1 2 3)
  set cras_new = ( $cras_new \
    `echo $cras[$k] $cdc_new[$k] $rdc_new[$k] $sdc_new[$k] \
                    $cdc[$k] $rdc[$k] $sdc[$k] \
     | awk -v dx=$dx -v dy=$dy -v dz=$dz \
           '{print $1 + ($2-$5)*dx/2 + ($3-$6)*dy/2 + ($4-$7)*dz/2}'` )
end

# Convert to the output direction cosines and origin without interpolating
set cmd = mri_convert
set cmd = ($cmd -oni $nx_new)
set cmd = ($cmd -onj $ny_new)
set cmd = ($cmd -onk $nz_new)
set cmd = ($cmd -oid $cdc_new)
set cmd = ($cmd -ojd $rdc_new)
set cmd = ($cmd -okd $sdc_new)
set cmd = ($cmd -oc  $cras_new)
set cmd = ($cmd -rt  nearest)
set cmd = ($cmd $in)
set cmd = ($cmd $out)
echo $cmd
$cmd

# Convert input gradient table, if present
set inbase = `echo $in | awk -v FS=.nii '{print $1}'`
set outbase = `echo $out | awk -v FS=.nii '{print $1}'`

set inbvals = $inbase.bvals
set outbvals = $outbase.bvals
if (-e $inbvals) then
  echo "INFO: found $inbvals, copying"
  cp $inbvals $outbvals
endif

set inbvecs = $inbase.bvecs
set outbvecs = $outbase.bvecs
if (-e $inbvecs) then
  echo "INFO: found $inbvecs, converting to LAS"
  rm -f $outbvecs
  set outcol = ()
  foreach j (1 2 3)
    foreach k (1 2 3)
      if ($order[$k] == $j) set outcol = ($outcol $sign[$k]'$'$k)
    end
  end
  awk "{print $outcol[1], $outcol[2], $outcol[3]}" $inbvecs >> $outbvecs
endif

if $docheck then
  set tmpdir = `fs_temp_dir`

  #
  # General check for images
  #
  set nframes = `mri_info $in | grep nframes | awk '{print $2}'`
  if ($nframes > 1) then
    set inframe = $tmpdir/`basename $inbase`.0.nii.gz
    set outframe = $tmpdir/`basename $outbase`.0.nii.gz
    set cmd = (mri_convert $in $inframe --frame 0)
    echo $cmd
    $cmd
    set cmd = (mri_convert $out $outframe --frame 0)
    echo $cmd
    $cmd
  else
    set inframe = $in
    set outframe = $out
  endif
  echo "INFO: checking that image was reoriented correctly"
  echo "INFO: inspect overlap in tkregister window"
  set cmd = tkregister2 
  set cmd = ($cmd --targ $inframe --mov $outframe)
  set cmd = ($cmd --regheader --reg $tmpdir/register.dat)
  echo $cmd
  $cmd &

  #
  # Diffusion tensor check for images + bvecs
  #
  if (-e $outbvals && -e $outbvecs) then
    set cmd = (fslroi $out $tmpdir/nodif.nii.gz 0 1)
    echo $cmd
    $cmd
    set cmd = (bet $tmpdir/nodif.nii.gz $tmpdir/nodif_brain.nii.gz -m)
    echo $cmd
    $cmd
    set cmd = dtifit
    set cmd = ($cmd -k $out -b $outbvals -r $outbvecs)
    set cmd = ($cmd -m $tmpdir/nodif_brain_mask.nii.gz -o $tmpdir/dtifit)
    echo $cmd
    $cmd
    echo "INFO: checking that image & bvecs were reoriented correctly"
    echo "INFO: inspect tensors in fslview window"
    set cmd = (fslview $tmpdir/dtifit_FA.nii.gz $tmpdir/dtifit_V1.nii.gz)
    echo $cmd
    $cmd &
  endif
endif

exit 0

############--------------##################
usage:
    echo "Usage:"
    echo "  $0 inputimage outputimage [--check]"
    echo
    echo "Convert image to LAS orientation"
    echo
    echo "The input image must be in NIfTI format. The output image"
    echo "will be in NIfTI format with LAS orientation."
    echo
    echo "The optional argument --check will bring up a tkregister"
    echo "window with the input and output image so you can check"
    echo "that they match completely."
    echo
    echo "Extra features for diffusion data"
    echo "---------------------------------"
    echo
    echo "If you run this script on a diffusion-weighted image series:"
    echo "  $0 your_DWIs.nii.gz your_DWIs_in_LAS.nii.gz"
    echo "and a text file named your_DWIs.bvecs is also present,"
    echo "this script will detect it and convert the gradient vectors"
    echo "to match the new LAS image orientation."
    echo
    echo "For such data, the optional argument --check will also run"
    echo "dtifit and bring up an fslview window with the tensors."
    echo

    exit 1