#!/bin/tcsh

### script to do 2D registration on each slice of a 3D+time
### dataset, and glue the results back together at the end

### Set prefix of input 3D+time dataset here.
### In this example with 'barney', the output
### dataset will be 'barney_reg+orig'.

set inp = $1

### extract number of slices into nz

set qq  = ( `3dAttribute DATASET_DIMENSIONS ${inp}+orig` )
echo $#qq
if ( $#qq == 0 ) then
  echo "Dataset ${inp}+orig missing DATASET_DIMENSIONS attribute - exiting"
  exit 1
endif

set nz  = $qq[3]
@ nz1   = $nz - 1

### Extract mean of input time series volumes,
### then make a mask from that, then make a weight volume from both

echo "========== Computing mean, mask, and weight datasets =========="

3dTstat    -mean -prefix ${inp}_mean ${inp}+orig

3dAutomask -dilate 4 -prefix ${inp}_mask ${inp}_mean+orig

3dcalc     -a ${inp}_mean+orig -b ${inp}_mask+orig           \
           -c a+i -d a-i -e a+j -f a-j                       \
           -expr 'median(a,c,d,e,f)*b' -prefix ${inp}_weight

### find orientation of dataset

set qq = ( `3dAttribute ORIENT_SPECIFIC ${inp}+orig` )
if ( $#qq == 0 ) then
  echo "Dataset ${inp}+orig missing ORIENT_SPECIFIC attribute - exiting"
  exit 1
endif

switch ( $qq[1] )
  case "0":
  case "1":
    set xxor = "R"
    breaksw
  case "2":
  case "3":
    set xxor = "A"
    breaksw
  case "4":
  case "5":
    set xxor = "I"
    breaksw
  default:
    echo 'Illegal value in ORIENT_SPECIFIC[1] - exiting'
    exit 1
endsw

switch ( $qq[2] )
  case "0":
  case "1":
    set yyor = "R"
    breaksw
  case "2":
  case "3":
    set yyor = "A"
    breaksw
  case "4":
  case "5":
    set yyor = "I"
    breaksw
  default:
    echo 'Illegal value in ORIENT_SPECIFIC[2] - exiting'
    exit 1
endsw

switch ( $qq[3] )
  case "0":
  case "1":
    set zzor = "R" ; set orient = "sagittal"
    breaksw
  case "2":
  case "3":
    set zzor = "A" ; set orient = "coronal"
    breaksw
  case "4":
  case "5":
    set zzor = "I" ; set orient = "axial"
    breaksw
  default:
    echo 'Illegal value in ORIENT_SPECIFIC[3] - exiting'
    exit 1
endsw

echo "Detected slice orientation: $orient"

switch( $zzor )
  case "R":
    set shift = 1 ; set rota = 4 ; set rotb = 6 ; set scala = 7 ; set shra = 10 ; set shrb = 11
    breaksw

  case "A":
    set shift = 2 ; set rota = 4 ; set rotb = 5 ; set scala = 8 ; set shra = 10 ; set shrb = 12
    breaksw

  case "I":
    set shift = 3 ; set rota = 5 ; set rotb = 6 ; set scala = 9 ; set shra = 11 ; set shrb = 12
    breaksw

  default:
    echo "Illegal value of zzor = ${zzor} - exiting"
    exit 1
endsw

echo "Freezing parameters $shift $rota $rotb $scala $shra $shrb in 3dWarpDrive"

### extract each slice in turn, and register it in 2D only;
### suppressed parameters (-parfix) are
###   #shift = shifting along 3rd dimension
###   #rota  = rotation about 1st inplane axis
###   #rotb  = rotation about 2nd inplane axis
###   #scala = dilation factor along 3rd dimension
###   #shra  = shear factor involving 3rd dimension
###   #shrb  = shear factor involving 3rd dimension

foreach zz ( `count 0 $nz1` )
  echo "==================== processing slice #$zz ===================="
  3dZcutup -keep $zz $zz -prefix ${inp}_${zz}        ${inp}+orig
  3dZcutup -keep $zz $zz -prefix ${inp}_weight_${zz} ${inp}_weight+orig
  3dZcutup -keep $zz $zz -prefix ${inp}_mean_${zz}   ${inp}_mean+orig

  3dWarpDrive -affine_general -cubic -final quintic  \
              -prefix ${inp}_reg_${zz}               \
              -base ${inp}_mean_${zz}+orig           \
              -input ${inp}_${zz}+orig               \
              -weight ${inp}_weight_${zz}+orig       \
              -parfix $shift  0                      \
              -parfix $rota   0 -parfix $rotb 0      \
              -parfix $scala  1                      \
              -parfix $shra   0 -parfix $shrb 0
end

### glue the slices back together

echo "======= Assembling registered 3D dataset ${inp}_reg+orig ======="

3dZcat -prefix ${inp}_reg ${inp}_reg_0???+orig.HEAD

### remove the single-slice datasets

/bin/rm -f ${inp}_*0???+orig.*