#! /bin/tcsh -f

#
# mri_add_new_tp
#
# Adds a new TP to an existing base without reprocessing the base
# The new TP can then be run longitudinally.
#
# Original Author: Martin Reuter
#
# 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
#
#


set VERSION = 'mri_add_new_tp dev';

if ($#argv < 2) then
  echo
  echo '  mri_add_new_tp <base-id> <newtp-id>';
  echo   
  echo "  Adds a new time point to the base/template without"
  echo "  re-creating the base. Therefore only the new time point"
  echo "  needs to be run longitudinally ( recon-all -long .. )"
  echo "  instead of re-running the base and all longitudinals."
  echo 
  echo "  Note: "
  echo "  1. The new time points needs to be fully processed with the "
  echo "     standard recon-all stream before (newtp-id is the subj-dir)"
  echo "  2. Adding a new time point to the base this way is only"
  echo "     recommended if:"
  echo "  a) enough time points are already in the base (maybe 3 or 4)"
  echo "  b) the new time point is not too different from the old ones"
  echo "  c) not too many other time points have been added this way"
  echo 
  echo "  In fact, we are currently not sure what influence this has"
  echo "  on your analysis. The good news: your results from earlier"
  echo "  longitudinal processing will stay the same. The bad news:"
  echo "  you introduce a bias towards the earlier time points."
  echo 
  echo "  More info on longitudinal processing at:"
  echo "  http://surfer.nmr.mgh.harvard.edu/fswiki/LongitudinalProcessing"
  echo
  echo
  exit 1;  
endif

if(! $?FREESURFER_HOME ) then
  echo "\nERROR: environment variable FREESURFER_HOME not set"
  echo "  this can be done by setting it in the shell before executing\n"
  exit 1;
endif

if(! $?SUBJECTS_DIR ) then
  echo "\nERROR: environment variable SUBJECTS_DIR not set"
  echo "  this can be done by setting it in the shell before executing\n"
  exit 1;
endif

if(! -e $SUBJECTS_DIR ) then
  echo "\nERROR: SUBJECTS_DIR $SUBJECTS_DIR does not exist.\n"
  exit 1;
endif

echo "INFO: SUBJECTS_DIR is $SUBJECTS_DIR"

set baseid   = $1
set newtpid  = $2

set basedir = $SUBJECTS_DIR/$baseid
if(! -e $basedir) then
  echo "\nERROR: cannot find $basedir\n"
  exit 1;
endif

set newtpdir = $SUBJECTS_DIR/$newtpid
if(! -e $newtpdir) then
  echo "\nERROR: cannot find $newtpdir"
  echo "     Make sure, this time point is processed with recon-all"
  echo "     cross sectionally before adding it to the base/template.\n"
  exit 1;
endif

set BaseSubjsListFname = (base-tps)
set BaseSubjsList = (`cat ${basedir}/${BaseSubjsListFname}`)

set affinebase = 0
set affinemap = $basedir/mri/transforms/${BaseSubjsList[1]}_to_${baseid}_affine.lta
if ( -e $affinemap ) then
  echo "INFO: Affine Base "
  set affinebase = 1
endif

set voltype = norm
set newtpvol   = $newtpdir/mri/${voltype}.mgz
if(! -e $newtpvol) then
  echo "\nERROR: cannot find $newtpvol\n"
  exit 1;
endif
    
set basevol  = ${basedir}/mri/${voltype}.mgz
if(! -e $newtpvol) then
  echo "\nERROR: cannot find $basevol\n"
  exit 1;
endif

foreach s ($BaseSubjsList)
  if ($s == $newtpid) then
    echo "\nERROR: $newtpid already contained in base $baseid\n"
    exit 1
  endif
end
 
# output files:
set newmean    = ${basedir}/mri/${voltype}.with_${newtpid}.mgz
set newtplta   = ${basedir}/mri/transforms/${newtpid}_to_${baseid}.lta
set newmeanssd = ${basedir}/mri/${voltype}.with_${newtpid}.ssd.txt
# we might need these if affine base:
set lta1form   = ${basedir}/mri/transforms/${newtpid}_to_${baseid}_norm.lta
set ltaAform   = ${basedir}/mri/transforms/${newtpid}_to_${baseid}_affine.lta

# register new tp to current base
set cmd = (mri_robust_register --mov $newtpvol --dst $basevol --sat 4.685)
if ( $affinebase ) then
  #this is only initial rigid reg
  set cmd = ($cmd --lta $lta1form)
else
  # this is final rigid reg
  set cmd = ($cmd --lta $newtplta)
endif
echo 
echo $cmd
echo
eval $cmd
if ($status) then
  echo "\nERROR: error running mri_robust_register\n"
  exit 1;
endif

# run more registrations in case of affine base:
if ($affinebase) then
  # first affinely register T1.mgz
  set amov = $newtpdir/mri/T1.mgz
  set adst = ${basedir}/mri/T1.mgz
  set cmd = (mri_robust_register --mov $amov --dst $adst --sat 4.685)
  set cmd = ($cmd --lta $ltaAform --ixform $lta1form --affine)
  echo 
  echo $cmd
  echo
  eval $cmd
  if ($status) then
    echo "\nERROR: error running mri_robust_register --affine\n"
    exit 1;
  endif
  # then fine tune rigid on norm.mgz to create final lta:
  set cmd = (mri_robust_register --mov $newtpvol --dst $basevol --sat 4.685)
  set cmd = ($cmd --lta $newtplta --ixform $ltaAform)
  echo 
  echo $cmd
  echo
  eval $cmd
  if ($status) then
    echo "\nERROR: error running mri_robust_register fine tune\n"
    exit 1;
  endif
endif

# collect volumes and init xforms from base
set subjInVols = ""
set ltaXforms  = ""
foreach s ($BaseSubjsList)
   set invol = ${SUBJECTS_DIR}/${s}/mri/${voltype}.mgz
   if(! -e $invol) then
     echo "\nERROR: cannot find $invol\n"
     exit 1;
   endif
   
   set subjInVols=($subjInVols $invol)
   set ltaname = ${s}_to_${baseid}.lta
   set ltaXforms=($ltaXforms ${basedir}/mri/transforms/${ltaname})
end

#only recreate template including new TP at same location (without iterating)
set cmd = ( mri_robust_template --mov $subjInVols $newtpvol --template $newmean --ixforms $ltaXforms $newtplta --noit )
echo 
echo $cmd
echo
eval $cmd
if ($status) then
  echo "\nERROR: error running mri_robust_template\n"
  exit 1;
endif

# checking differences in intensities
set cmd  = ( mri_diff --pix-only --ssd $basevol $newmean )
echo 
echo $cmd
echo
set diff = (`$cmd | grep "sum of squared differences" | awk '{print $1}' `)
# dont test status here, as it will not be zero if images differ
echo $diff > $newmeanssd
echo "Sum of squared differences to previous norm.mgz: $diff "

# maybe at some point we can use sse diff to recommend
# rerunning the base (if the differences get too large).
# For now it is completely unclear what 'too large' means,
# and we might never know.
# So if the user wants to add this time point let him do so:
set small = 1 

# if diff small:
if ( "$small" == "1" ) then
  #  do not rerun base, only add ltas and update insubject list
  #  maybe update bruces stuff

  #echo "Base can be patched.

  # add tpid to base list file:
  echo $newtpid >> ${basedir}/${BaseSubjsListFname}

  # keep newtplta and newmean newmeanssd where they are

  #invert lta
  set basetotp = ${basedir}/mri/transforms/${baseid}_to_${newtpid}.lta
  mri_concatenate_lta -invert1 $newtplta identity.nofile $basetotp
  if ($status) then
    echo "\nERROR: error running mri_concatenate_lta\n"
    exit 1;
  endif

  # update more ? brainmask?...

  echo "Base patched to include new TP $newtpid"
  echo "You only need to rerun new TP long:"
  echo "recon-all -long $newtpid $baseid -all"
  echo
	
else  # if diff large:
  echo 
  echo "Base cannot be patched. Changes are too large, please rerun base with:"
  echo
  echo "recon-all -base $baseid -base-insubj .... -all -force "
  echo "(maybe additional flags, if you used them)"
  echo
  echo "Then rerun ALL longitudinals with the new base:"
  echo "recon-all -long tpN $baseid -all -force"
  echo
  rm $newmean
  rm $newlta
  rm $newmeanssd
endif

exit 0;