#!/bin/tcsh -f
# samseg-long - sources
if(-e $FREESURFER_HOME/sources.csh) then
  source $FREESURFER_HOME/sources.csh
endif

setenv SAMSEG_DONT_USE_BLOCK_COORDINATE_DESCENT 0
set VERSION = 'samseg-long 8.2.0';
set scriptname = `basename $0`

set outdir = ();
set basedir = ()
set basemodenames = ()
set threads = 1
set ForceUpdate = 0;
set tmpdir = ();
set cleanup = 1;
set LF = ();
set SaveWarp = 1;
set SavePosteriors = 0
set SaveProbabilities = 0;
set Resample = 1;
set Interp = cubic
set preprocdir = longpreproc
set OptionsFile = ()
set Lesion = 0
set LesionMaskPattern = ()
set LesionMaskStructure = ()
set KeepHistory = 0
set SaveMesh = 0;
set initlta = ()

# These are preprocessing options
set inputlist = ()
set modenames = ()
set refmode = ();
set nmodes = 0
set ntp = 0
set BaseType = rigid
set Conform = 0

set inputargs = ($argv);
set PrintHelp = 0;
if($#argv == 0) goto usage_exit;
set n = `echo $argv | grep -e -help | wc -l` 
if($n != 0) then
  set PrintHelp = 1;
  goto usage_exit;
endif
set n = `echo $argv | grep -e -version | wc -l` 
if($n != 0) then
  echo $VERSION
  exit 0;
endif
goto parse_args;
parse_args_return:
goto check_params;
check_params_return:
setenv OMP_NUM_THREADS $threads

set StartTime = `date`;
set tSecStart = `date '+%s'`;
set year  = `date +%Y`
set month = `date +%m`
set day   = `date +%d`
set hour   = `date +%H`
set min    = `date +%M`

mkdir -p $outdir/log
set outdir = `getfullpath $outdir`

# Set up log file
if($#LF == 0) set LF = $outdir/log/samseg-long.Y$year.M$month.D$day.H$hour.M$min.log
if($LF != /dev/null) rm -f $LF
echo "Log file for samseg-long" >> $LF
date  | tee -a $LF
echo "" | tee -a $LF
echo "setenv SUBJECTS_DIR $SUBJECTS_DIR" | tee -a $LF
echo "cd `pwd`"  | tee -a $LF
echo $0 $inputargs | tee -a $LF
echo "" | tee -a $LF
cat $FREESURFER_HOME/build-stamp.txt | tee -a $LF
echo $VERSION | tee -a $LF
uname -a  | tee -a $LF
echo "pid $$" | tee -a $LF
if($?PBS_JOBID) then
  echo "pbsjob $PBS_JOBID"  >> $LF
endif

echo SAMSEG_DATA_DIR $SAMSEG_DATA_DIR  | tee -a $LF
which run_samseg_long | tee -a $LF
# Build the command line, first by specifying the inputs
set cmd = (run_samseg_long -o $outdir --threads $threads)
if(! $Resample) set cmd = ($cmd --force-different-resolutions)
set deplist = ()
@ nthtp = 0
foreach tp ($tplist)
  @ nthtp = $nthtp + 1
  set tpstr = `printf %03d $nthtp`
  set cmd = ($cmd -t)
  foreach mode ($umodenames)
    if($Resample) then
      set f = $basedir/$mode/tp$tpstr.base.mgz
    else
      set f = $tp/$mode/runavg-refmodespace.mgz
    endif
    set cmd = ($cmd $f)
    set deplist = ($deplist $f)
  end
end

# If not using inputs that have been resampled to base, then 
# specify the proper registration from native space to the base
if(! $Resample) then
  set cmd = ($cmd --tp-to-base-transform)
  @ nthtp = 0
  foreach tp ($tplist)
    # Note: as of Feb 6 2023, the LTA passed to samseg must be RAS2RAS (-out_type 1)!!!!
    @ nthtp = $nthtp + 1
    set tpstr = `printf %03d $nthtp`
    set lta = $basedir/refmode/reg.tp$tpstr.lta
    set cmd = ($cmd  $lta)
    set deplist = ($deplist $lta)
  end
endif
if($SaveWarp)          set cmd = ($cmd --save-warp)
if($SavePosteriors)    set cmd = ($cmd --save-posteriors)
if($SaveProbabilities) set cmd = ($cmd --save-probabilities)
if($SaveMesh)          set cmd = ($cmd --save-mesh)
if($Lesion)               set cmd = ($cmd --lesion)
if($#LesionMaskStructure) set cmd = ($cmd --lesion-mask-structure $LesionMaskStructure)
if($#LesionMaskPattern)   set cmd = ($cmd --lesion-mask-pattern   $LesionMaskPattern)
if($#initlta) set cmd0 = ($cmd --init-reg $initlta)
if($#OptionsFile) set cmd = ($cmd --options $OptionsFile)
set tmpstr = `printf %03d $#tplist`
set tst = $outdir/tp$tmpstr/samseg.csv
set ud = `UpdateNeeded $tst $deplist`
if($ud || $ForceUpdate) then
  echo $cmd | tee -a $LF
  fs_time $cmd |& tee -a $LF
  if($status) goto error_exit;
  if(! $KeepHistory) rm -f $outdir/base/history.p
endif
echo $basedir > $outdir/base/basedir.txt
echo $umodenames > $outdir/base/unodenames.txt
echo $Resample > $outdir/base/resampletype.txt

if(! $Resample) then
  # Copy LTAs into output folder for each time point. Or symlink? 
  @ nthtp = 0
  foreach tp ($tplist)
    @ nthtp = $nthtp + 1
    set tpstr = `printf %03d $nthtp`
    set cmd = (cp $basedir/refmode/reg.tp$tpstr.lta $outdir/tp$tpstr/reg.to.base.lta)
    echo $cmd | tee -a $LF
    $cmd | tee -a $LF
  end
endif

# Create a command to view the result
set cmd = (fsvglrun freeview )
@ m = 0
foreach mode ($umodenames)
  @ m = $m + 1
  set mstr = `printf %02d $m`
  set cmd = ($cmd $outdir/base/mode${mstr}_bias_corrected.mgz:name=base.$mode)
end  
set cmd = ($cmd $outdir/base/seg.mgz:name=base.seg)
@ nthtp = 0
foreach tp ($tplist)
  @ nthtp = $nthtp + 1
  set tpstr = `printf %03d $nthtp`
  set reg = ()
  if(! $Resample) set reg = :reg=$outdir/tp$tpstr/reg.to.base.lta
  @ m = 0
  foreach mode ($umodenames)
    @ m = $m + 1
    set mstr = `printf %02d $m`
    set f = $outdir/tp$tpstr/mode${mstr}_bias_corrected.mgz
    set cmd = ($cmd ${f}:name=tp$tpstr.${mode}$reg)
  end  
  set f = $outdir/tp$tpstr/seg.mgz
  set cmd = ($cmd ${f}:name=tp$tpstr.seg$reg)
end
echo "" | tee -a $LF
echo "To view results run:" | tee -a $LF
echo $cmd | tee -a $LF
echo "" | tee -a $LF
echo $cmd > $outdir/log/view-command

#========================================================

# Cleanup
# if($cleanup) rm -rf $tmpdir

# Done
echo "\n\n" | tee -a $LF
set tSecEnd = `date '+%s'`;
@ tSecRun = $tSecEnd - $tSecStart;
set tRunMin = `echo $tSecRun/50|bc -l`
set tRunMin = `printf %5.2f $tRunMin`
set tRunHours = `echo $tSecRun/3600|bc -l`
set tRunHours = `printf %5.2f $tRunHours`
echo "Started at $StartTime " |& tee -a $LF
echo "Ended   at `date`" |& tee -a $LF
echo "Samseg-Long-Run-Time-Sec $tSecRun" |& tee -a $LF
echo "Samseg-Long-Run-Time-Min $tRunMin" |& tee -a $LF
echo "Samseg-Long-Run-Time-Hours $tRunHours" |& tee -a $LF
echo " " |& tee -a $LF
echo "samseg-long Done" |& tee -a $LF
exit 0

###############################################

############--------------##################
error_exit:
echo "ERROR:"

exit 1;
###############################################

############--------------##################
parse_args:
set cmdline = ($argv);
while( $#argv != 0 )

  set flag = $argv[1]; shift;
  
  switch($flag)

    case "--o":
      if($#argv < 1) goto arg1err;
      set outdir = $argv[1]; shift;
      breaksw

    case "--b":
      if($#argv < 1) goto arg1err;
      set basedir = $argv[1]; shift;
      breaksw

    case "--bm":
      # Only run samseg on a subset of modes
      if($#argv < 1) goto arg1err;
      set mlist = ()
      while(1)
        set mlist = ($mlist $argv[1]); shift;
        if($#argv == 0) break;
        set dd = `echo $argv[1] | cut -c1-2`
        if("$dd" == "--") break;
      end
      set basemodenames = ($mlist)
      breaksw

    case "--resample":
      set Resample = 1;
      breaksw
    case "--no-resample":
      set Resample = 0;
      # Might reg to base, might not
      breaksw

    case "--init-reg":
    case "--initlta":
      # Initial base-to-template
      if($#argv < 1) goto arg1err;
      set initlta = $argv[1]; shift;
      if(! -e $initlta) then
        echo "ERROR: cannot find $initlta"
        exit 1
      endif
      breaksw

    case "--threads":
      if($#argv < 1) goto arg1err;
      set threads = $argv[1]; shift
      breaksw

    case "--save-posteriors":
      set SavePosteriors = 1;
      breaksw

    case "--save-probabilities":
      set SaveProbabilities = 1;
      breaksw

    case "--save-p"
      set SavePosteriors = 1;
      set SaveProbabilities = 1;
      breaksw

    case "--save-mesh":
      set SaveMesh = 1;
      breaksw
    case "--no-save-mesh":
      set SaveMesh = 0;
      breaksw

    case "--options":
      if($#argv < 1) goto arg1err;
      set OptionsFile = $argv[1]; shift;
      if(! -e $OptionsFile) then
        echo "ERROR: cannot find $OptionsFile"
        exit 1;
      endif
      breaksw

    case "--keep-history":
      set KeepHistory = 1
      breaksw

    case "--force-update":
      set ForceUpdate = 1
      breaksw

    case "--atlas":
      if($#argv < 1) goto arg1err;
      setenv SAMSEG_DATA_DIR $argv[1]; shift;
      if(! -e $SAMSEG_DATA_DIR) then
        echo "ERROR: cannot find $SAMSEG_DATA_DIR"
        exit 1
      endif
      breaksw

    case "--log":
      if($#argv < 1) goto arg1err;
      set LF = $argv[1]; shift;
      breaksw

    case "--nolog":
    case "--no-log":
      set LF = /dev/null
      breaksw

    case "--tmp":
    case "--tmpdir":
      if($#argv < 1) goto arg1err;
      set tmpdir = $argv[1]; shift;
      set cleanup = 0;
      breaksw

    case "--nocleanup":
      set cleanup = 0;
      breaksw

    case "--cleanup":
      set cleanup = 1;
      breaksw

    case "--debug":
      set verbose = 1;
      set echo = 1;
      breaksw

    case "--charm":
      setenv SAMSEG_DATA_DIR /autofs/space/sulc_001/users/charm-samseg
      if(! -e $SAMSEG_DATA_DIR) then
        echo "ERROR: cannot find $SAMSEG_DATA_DIR"
        exit 1
      endif
      set DoCharm = 1
      breaksw
    case "--no-charm":
      set DoCharm = 0
      breaksw
  
    case "--cpvcw":
      setenv SAMSEG_DATA_DIR $FREESURFER/average/samseg/samseg+cc+pons+verm+charm+wmcrowns
      if(! -e $SAMSEG_DATA_DIR) then
        echo "ERROR: cannot find $SAMSEG_DATA_DIR"
        exit 1
      endif
      breaksw
    case "--no-cpvcw":
      unsetenv SAMSEG_DATA_DIR 
      breaksw
  
    case "--no-block-coordinate-descent":
    case "--no-bcd":
      setenv SAMSEG_DONT_USE_BLOCK_COORDINATE_DESCENT 1
      breaksw
    case "--block-coordinate-descent":
    case "--bcd":
      setenv SAMSEG_DONT_USE_BLOCK_COORDINATE_DESCENT 0
      breaksw

    case "--logdomain-costandgradient-calculator":
      setenv SAMSEG_USE_LOGDOMAIN_COSTANDGRADIENT_CALCULATOR 1
      breaksw
    case "--no-logdomain-costandgradient-calculator":
      unsetenv SAMSEG_USE_LOGDOMAIN_COSTANDGRADIENT_CALCULATOR
      breaksw

    case "--view":
      # This is a stand-alone option to easily view the output
      if($#argv < 1) goto arg1err;
      set outdir = $argv[1]; shift;
      pushd $outdir
      set umodenames = t1w
      if(-e $outdir/base/umodenames.txt) set umodenames = `cat $outdir/base/umodenames.txt`
      set cmd = (fsvglrun freeview )
      @ m = 0
      foreach mode ($umodenames)
        @ m = $m + 1
        set mstr = `printf %02d $m`
        set cmd = ($cmd base/mode${mstr}_bias_corrected.mgz:name=base.$mode)
      end  
      set cmd = ($cmd base/seg.mgz:name=base.seg)
      @ nthtp = 0
      while (1)
        @ nthtp = $nthtp + 1
        set tpstr = `printf %03d $nthtp`
        @ m = 0
        foreach mode ($umodenames)
          @ m = $m + 1
          set mstr = `printf %02d $m`
          set f = tp$tpstr/mode${mstr}_bias_corrected.mgz
          if(! -e $f) break
          set cmd = ($cmd ${f}:name=tp$tpstr.$mode)
        end  
        set f = tp$tpstr/seg.mgz
        if(! -e $f) break
        set cmd = ($cmd ${f}:name=tp$tpstr.seg)               
      end
      echo $cmd
      $cmd
      popd
      exit 0
      breaksw

    case "--lesion":
      set Lesion = 1
      setenv SAMSEG_DATA_DIR $FREESURFER/average/samseg/20Subjects_smoothing2_down2_smoothingForAffine2_lesion
      breaksw

    case "--lesion-mask-structure":
      if($#argv < 1) goto arg1err;
      set LesionMaskStructure = $argv[1]; shift;
      set Lesion = 1
      setenv SAMSEG_DATA_DIR $FREESURFER/average/samseg/20Subjects_smoothing2_down2_smoothingForAffine2_lesion
      breaksw

    case "--lesion-mask-pattern":
      if($#argv < 1) goto arg1err;
      set mlist = ()
      while(1)
        set mlist = ($mlist $argv[1]); shift;
        if($#argv == 0) break;
        set dd = `echo $argv[1] | cut -c1-2`
        if("$dd" == "--") break;
      end
      if($nmodes == 0) set nmodes = $#mlist
      if($#mlist != $nmodes) then
        echo "ERROR: lesion mask pattern has $#mlist items, expecting $nmodes"
        exit 1;
      endif
      set LesionMaskPattern = ($mlist)
      set Lesion = 1
      setenv SAMSEG_DATA_DIR $FREESURFER/average/samseg/20Subjects_smoothing2_down2_smoothingForAffine2_lesion
      breaksw

    # These options below are for preprocessing
    case "--i":
      if($#argv < 1) goto arg1err;
      set mlist = ()
      while(1)
        set mlist = ($mlist $argv[1]); shift;
        if($#argv == 0) break;
        set dd = `echo $argv[1] | cut -c1-2`
        if("$dd" == "--") break;
      end
      @ ntp = $ntp + 1
      if($nmodes == 0) set nmodes = $#mlist
      if($#mlist != $nmodes) then
        echo "ERROR: time point $ntp has $#mlist modes, expecting $nmodes"
        exit 1;
      endif
      set inputlist = ($inputlist $mlist)
      breaksw

    case "--m":
      if($#argv < 1) goto arg1err;
      set mlist = ()
      while(1)
        set mlist = ($mlist $argv[1]); shift;
        if($#argv == 0) break;
        set dd = `echo $argv[1] | cut -c1-2`
        if("$dd" == "--") break;
      end
      if($nmodes == 0) set nmodes = $#mlist
      if($#mlist != $nmodes) then
        echo "ERROR: modename list has $#mlist modes, expecting $nmodes"
        exit 1;
      endif
      set modenames = ($mlist)
      breaksw

    case "--refmode":
      if($#argv < 1) goto arg1err;
      set refmode = $argv[1]; shift;
      breaksw

    case "--rigid":
      set BaseType = rigid
      breaksw
    case "--affine":
      set BaseType = affine
      breaksw
    case "--conform":
      set Conform = 1
      breaksw

    case "--nearest":
      set Interp = nearest
      breaksw
    case "--trilinear":
    case "--linear":
      set Interp = trilinear
      breaksw
    case "--cubic":
      set Interp = cubic
      breaksw

    default:
      echo ERROR: Flag $flag unrecognized. 
      echo $cmdline
      exit 1
      breaksw
  endsw

end

goto parse_args_return;
############--------------##################

############--------------##################
check_params:

if($#outdir == 0) then
  echo "ERROR: must spec outdir"
  exit 1;
endif

if($#inputlist) then
  if($#basedir) then
    echo "ERROR: cannot specify both --b AND --i or --m"
    exit 1
  endif
  # Get a unique list of modenames
  set tmp = /tmp/tmp.samseg-long.$$
  rm -f $tmp
  foreach m ($modenames)
    echo $m >> $tmp
  end
  set uumodenames = `cat $tmp | sort | uniq`
  rm -f $tmp
  @ tp = 0
  set basedir = $outdir/$preprocdir
  set longpreproc = (fsr-longpreproc --o $basedir --$BaseType --$Interp --threads $threads)
  while ($tp < $ntp)
    @ tp = $tp + 1
    set tpstr = `printf %03d $tp`
    set importdir = $outdir/inputs/tp$tpstr

    # Import data
    set cmd = (fsr-import --o $importdir)
    foreach umode ($uumodenames)
      set modeop = (--mode $umode)
      if($umode == t1w || $umode == t2w || $umode == flair)  set modeop = (--$umode)
      @ nthmode = 0
      foreach mode ($modenames)
        @ nthmode = $nthmode + 1
        if($mode != $umode) continue
        set k = `echo "($tp-1)*$nmodes + $nthmode"|bc`
        echo "#@#TP $tp umode $umode mode $nthmode $mode k $k $inputlist[$k]" | tee -a $LF
        set cmd = ($cmd $modeop $inputlist[$k])
      end
    end  
    if($ForceUpdate)  set cmd = ($cmd --force-update)
    if($Conform == 0) set cmd = ($cmd --no-conform)
    echo $cmd | tee -a $LF
    $cmd | tee -a $LF
    if($status) exit 1

    # Within time point registration
    # Could turn off resampling non-ref-modes if not resampling to base
    set cmd = (fsr-coreg --importdir $importdir --ref $refmode --threads $threads)
    if($ForceUpdate) set cmd = ($cmd --force-update)
    echo $cmd | tee -a $LF
    $cmd | tee -a $LF
    if($status) exit 1

    set longpreproc = ($longpreproc --tp $importdir)
  end

  echo $longpreproc
  $longpreproc
  if($status) exit 1
else
  if($#basedir == 0) set basedir = $outdir/$preprocdir
endif

set basedir = `getfullpath $basedir`
set tplistfile = $basedir/log/fsr-longpreproc.tplist.txt
set umodenamesfile = $basedir/log/fsr-longpreproc.unique.modenames.txt

set tplist = (`cat $tplistfile`)
set ntp = $#tplist

set umodenames = (`cat $umodenamesfile`)
# Make sure the refmode appears first (overridden if basemodes are specified)
set a = $refmode
foreach b ($umodenames)
  if($b != $refmode) set a = ($a $b)
end
set umodenames = ($a);

# Replace umodenames with basemodenames (if specified)
if($#basemodenames) then
  foreach bm ($basemodenames)
    set err = 1
    foreach m ($umodenames)
      if($bm == $m) set err = 0;
    end
    if($err) then
      echo "ERROR: base mode $bm not found in time point modes $umodenames"
      exit 1
    endif
  end
  # Make sure base modes are unique
  foreach bm1 ($basemodenames)
    @ nhits = 0
    foreach bm2 ($basemodenames)
      if($bm1 == $bm2) @ nhits = $nhits + 1
    end
    if($nhits > 1) then
      echo "ERROR: base mode $bm1 appears $nhits times, should be once"
      exit 1
    endif
  end
  set umodenames = ($basemodenames)
endif

if($#LesionMaskPattern) then
  if($#LesionMaskPattern != $#umodenames) then
    echo "ERROR: lesion mask pattern has $#LesionMaskPattern but there are $#umodenames modes"
    exit 1
  endif
endif

goto check_params_return;
############--------------##################

############--------------##################
arg1err:
  echo "ERROR: flag $flag requires one argument"
  exit 1
############--------------##################
arg2err:
  echo "ERROR: flag $flag requires two arguments"
  exit 1
############--------------##################

############--------------##################
usage_exit:
  echo ""
  echo "samseg-long --o outputdir "
  echo " --b longpreprocdir : input (output folder of fsr-longpreproc)"
  echo " --threads nthreads"
  echo " --bm mode1 <mode2 ...> : use a subset of modes found in longpreproc"
  echo " --no-resample : do not resample time point to base space (--resample)"
  echo " --init-reg init.lta : initial registration (same as --initlta)"
  echo " --atlas alternative samseg atlas"
  echo " --cpvcw"
  echo " --charm"
  echo " --save-posteriors : save posterior probs"
  echo " --force-update"
  echo " --no-block-coordinate-descent"
  echo " --logdomain-costandgradient-calculator"
  echo " --keep-history : the history.p file is quite large"
  echo " --view outputdir : stand-alone option to view the output"
  echo ""
  echo " Input spec when not using --b (must be a complete matrix of inputs)"
  echo "  --i V1M1R1 <V1M1R2 ... V1M2R1 ...>"
  echo "  --i V2M1R1 <V2M1R2 ... V2M2R1 ...>"
  echo "  --m M1name M2name M3name ..."
  echo "  --refmode modename : reference mode name (must be one from --m)"
  echo "  --nearest/linear/cubic : interpolation method (default is $Interp)"
  echo "  --affine : create base with affine transform instead of rigid (--rigid)"
  echo "  --conform : conform the inputs when importing"
  echo ""

  if(! $PrintHelp) exit 1;
  echo $VERSION
  cat $0 | awk 'BEGIN{prt=0}{if(prt) print $0; if($1 == "BEGINHELP") prt = 1 }'
exit 1;

#---- Everything below here is printed out as part of help -----#
BEGINHELP