#! /bin/bash

function Usage () {
  cat <<EOF

   PESTICA_afni v1.4
 Usage:  run_pestica.sh  -b ep2d_pace_132vols (for AFNI BRIK format)
  OR  :  run_pestica.sh  -b ep2d_pace_132vols.hdr (or .nii, .nii.gz)
  Opt :  run_pestica.sh  -b ep2d_pace_132vols  -r  -c ep2d_pace_132vols.moco -d 4 -s "3 4 5 6"
 Flags:
  -b = base 3D+time EPI dataset you will run ICA and physio estimation upon
  -r = flag to run original RETROICOR algorithm as well as IRF-RETROICOR
  -d = number of volumes to drop at start (default is 0, if your data is not equilibrated, use up to 4)
       NOTE: you MUST not have un-equilibrated volumes at the start. Most scanners take "dummy" volumes, where
       the ADCs are turned off but the RF and gradients are running as normal, for the first ~3 seconds (modulo TR),
       but in some scanners this is not so and you can see contrast change from the 1st to 2nd volumes. Look at your data!
       If you have this contrast change from 1st to 2nd image, then remove them either before proceeding or using this option
  -c = 3D+time EPI dataset you will run any corrections on (e.g. run ICA on uncorrected data, 
     but apply physiologic noise correction to motion-corrected data) - assumes timing is identical.
     This is safe to use with the -d flag as long as you didn't drop volumes anywhere else; 
     PESTICA will detect number of vols.  This file must be in AFNI BRIK format.
     Example: run_pestica.sh  -b ep2d_pace_132vols  -r  -c ep2d_pace_132vols.moco  -d 4
          will create estimators that are 4 volumes shorter than ep2d_pace_132vols, so ep2d_pace_132vols.moco
	  must already be 4 vols shorter. In this example, it could have been created like so:
	       3dvolreg -prefix ep2d_pace_132vols.moco -base "ep2d_pace_132vols[0]" -zpad 8 -maxite 60 -heptic "ep2d_pace_132vols[4..$]"
	  NOTE: this is not recommended - PESTICA estimators seem better if the ICA is also done on the moco'ed data.
  -s = "1 2 3 4 5 6" individual stages to run
     Stage allows you to run or re-run only parts of PESTICA
     please give the stages (in increasing order) that you want to run
     e.g.   :  run_pestica.sh  -b ep2d_pace_132vols  -s "5 6"
     otherwise, this script will run all stages in order

 If not in BRIK format, this script will make a copy into AFNI .BRIK format for PESTICA 
 For subsequent runs, use the file prefix only (drop the suffix)

EOF
  exit 1
}

retroicorflag=0; corrflag=0; dropvolsflag=0
stage1flag=1; stage2flag=1; stage3flag=1; stage4flag=1; stage5flag=1; stage6flag=1;

while getopts hrd:b:c:s: opt; do
  case $opt in
    h)
       Usage
       exit 1
       ;;
    b) # base 3D+time EPI dataset to use to perform all ICA decomposition and alignment
       epi=$OPTARG
       ;;
    r) # flag to run original RETROICOR as well
       retroicorflag=1
       ;;
    c) # 3D+time dataset to apply RETROICOR or IRF-RETROICOR correction to (if different from base)
       corrflag=1
       correpi=$OPTARG
      ;;
    d) # number of volumes to drop off the start of the 3D+time dataset
       dropvolsflag=1
       volsdrop=$OPTARG
      ;;
    s) # option to run stages manually
       stages=$OPTARG
       stage1flag=0; stage2flag=0; stage3flag=0; stage4flag=0; stage5flag=0; stage6flag=0;
       # loop over numbers listed
       for i in $stages ; do
         if [ $i -eq 1 ] ; then
           stage1flag=1
	 elif [ $i -eq 2 ] ; then
           stage2flag=1
	 elif [ $i -eq 3 ] ; then
           stage3flag=1
	 elif [ $i -eq 4 ] ; then
           stage4flag=1
	 elif [ $i -eq 5 ] ; then
           stage5flag=1
	 elif [ $i -eq 6 ] ; then
           stage6flag=1
	 else
	   echo "incorrect syntax for stages input: $i"
	   Usage
	   exit 1
	 fi
       done
      ;;
    :)
      echo "option requires input"
      exit 1
      ;;
  esac
done

if [ -z $epi ] ; then
  echo "3D+time EPI dataset $epi must exist"
  Usage
  exit 1
fi
if [[ $corrflag -eq 1 ]] && [[ ! -f $correpi.BRIK ]] ; then
  Usage
  echo ""
  echo "*****   $correpi.BRIK does not exist, exiting ..."
  exit
fi

echo $stage1flag $stage2flag $stage3flag $stage4flag $stage5flag $stage6flag
# test if NIFTI_GZ format
a=`echo $epi | grep ".nii.gz"`
if [ ! -z $a ] ; then
  epinew=`echo $epi | sed "s/.nii.gz//"`
  if [ -f $epinew+orig.BRIK ] ; then
    echo "file $epinew+orig.BRIK already exists, is this a copy of $epi? Either use $epinew as input, or delete or rename the file. Exiting..."
    echo ""
    exit
  fi
  3dcopy $epi $epinew
  epi=$epinew
fi
# test if NIFTI format
a=`echo $epi | grep ".nii"`
if [ ! -z $a ] ; then
  epinew=`echo $epi | sed "s/.nii//"`
  if [ -f $epinew+orig.BRIK ] ; then
    echo "file $epinew+orig.BRIK already exists, is this a copy of $epi? Either use $epinew as input, or delete or rename the file. Exiting..."
    echo ""
    exit
  fi
  3dcopy $epi $epinew
  epi=$epinew
fi
# test if NIFTI_PAIR or ANALYZE format
a=`echo $epi | grep ".hdr"`
if [ ! -z $a ] ; then
  epinew=`echo $epi | sed "s/.hdr//"`
  if [ -f $epinew+orig.BRIK ] ; then
    echo "file $epinew+orig.BRIK already exists, is this a copy of $epi? Either use $epinew as input, or delete or rename the file. Exiting..."
    echo ""
    exit
  fi
  3dcopy $epi $epinew
  epi=$epinew
fi
# test if AFNI BRIK format
a=`echo $epi | grep "+orig"`
if [ ! -z $a ] ; then
  epi=`echo $epi | sed "s/\+orig//"`
fi
epi_mask="$epi".brain+orig
epi=$epi+orig

if [ ! -f $epi.BRIK ] ; then
  echo "*****   $epi.BRIK does not exist, exiting ..."
  exit
fi
# test for presence of mask file
if [ ! -f $epi_mask.BRIK ] ; then
  echo ""
  echo "*****   $epi_mask.BRIK does not exist, creating mask file"
  echo "note, if you wish to use your own mask file, kill this script now,"
  echo "then simply 3dcopy your mask file to $epi_mask and re-run this script"
  echo ""
  echo "running 3dSkullStrip -input $epi -prefix $epi_mask"
  sleep 1
  #3dSkullStrip -input $epi -prefix $epi_mask
  3dSkullStrip -input $epi -prefix ___tmp_mask
  # erode mask by one voxel
  3dcalc -a ___tmp_mask+orig -prefix ___tmp_mask_ones+orig -expr 'step(a)'
  3dcalc -a ___tmp_mask_ones+orig -prefix ___tmp_mask_ones_dil -b a+i -c a-i -d a+j -e a-j -f a+k -g a-k -expr 'amongst(1,a,b,c,d,e,f,g)'
  3dcalc -a "$epi[0]" -b ___tmp_mask_ones_dil+orig -prefix $epi_mask -expr 'a*step(b)'
  rm ___tmp_mask*
  echo ""
  echo "done with skull-stripping - please check file and if not satisfied, I recommend running"
  echo "3dSkullStrip with different parameters to attempt to get a satisfactory mask"
  echo "this script will look for and use any file named $epi_mask so be sure that is your mask file"
  sleep 1
else
  echo "Using $epi_mask as mask file for $epi"
  echo ""
fi

# test for presence of PESTICA subdirectory
epi_pestica=`echo $epi | sed "s/\+orig/_pestica/"`
if [ ! -d $epi_pestica ] ; then
 mkdir $epi_pestica
fi
# if correcting a different file, this will be set appropriately, otherwise, set it back here
if [ $corrflag -eq 0 ] ; then
  correpi=$epi
fi

dummyrun=1
for run in $dummyrun ; do

 if [ $stage1flag -eq 1 ] ; then
  echo ""
  echo "Running Stage 1: slicewise temporal Infomax ICA"
  echo ""
  if [ $dropvolsflag -eq 1 ] ; then
    echo  matlab $MATLABLINE -r "addpath $PESTICA_DIR; addpath $PESTICA_DIR/eeglab; prepare_ICA_decomp(15,'$epi','$epi_mask',$volsdrop); exit;"
          matlab $MATLABLINE -r "addpath $PESTICA_DIR; addpath $PESTICA_DIR/eeglab; prepare_ICA_decomp(15,'$epi','$epi_mask',$volsdrop); exit;"
  else
    echo  matlab $MATLABLINE -r "addpath $PESTICA_DIR; addpath $PESTICA_DIR/eeglab; prepare_ICA_decomp(15,'$epi','$epi_mask'); exit;"
          matlab $MATLABLINE -r "addpath $PESTICA_DIR; addpath $PESTICA_DIR/eeglab; prepare_ICA_decomp(15,'$epi','$epi_mask'); exit;"
  fi
 else
  echo "skipping stage 1 (ICA decomposition)"
 fi
  

 if [ $stage2flag -eq 1 ] ; then
  echo ""
  echo "Running Stage 2: Coregistration to stereotaxic EPI and back-transform of PESTICA templates, followed by PESTICA estimation"
  echo ""
  # PESTICA coreg to template
  echo "Coregistration to EPI template"
  3dAllineate -automask -prefix $epi_pestica/$epi_mask.crg2mni.nii -source $epi_mask -base $PESTICA_DIR/meanepi_mni.nii -1Dmatrix_save $epi_pestica/$epi_mask.coreg.mni.1D

  cat_matvec $epi_pestica/$epi_mask.coreg.mni.1D -I -ONELINE > $epi_pestica/mni.coreg.$epi_mask.1D
  3dAllineate -prefix $epi_pestica/resp_PESTICA2.nii -source $PESTICA_DIR/resp_mean_mni.nii -base $epi_mask -1Dmatrix_apply $epi_pestica/mni.coreg.$epi_mask.1D
  3dAllineate -prefix $epi_pestica/card_PESTICA2.nii -source $PESTICA_DIR/card_mean_mni.nii -base $epi_mask$suffix -1Dmatrix_apply $epi_pestica/mni.coreg.$epi_mask.1D
  echo "Obtaining PESTICA estimators"
  echo  matlab $MATLABLINE -r "addpath $PESTICA_DIR; [card,resp]=apply_PESTICA2(15,'$epi','$epi_mask','$PESTICA_SLICE_TIMING'); fp=fopen('$epi_pestica/card_rawest.dat','w'); fprintf(fp,'%g\n',card); fclose(fp); fp=fopen('$epi_pestica/resp_rawest.dat','w'); fprintf(fp,'%g\n',resp); fclose(fp);"
        matlab $MATLABLINE -r "addpath $PESTICA_DIR; [card,resp]=apply_PESTICA2(15,'$epi','$epi_mask','$PESTICA_SLICE_TIMING'); fp=fopen('$epi_pestica/card_rawest.dat','w'); fprintf(fp,'%g\n',card); fclose(fp); fp=fopen('$epi_pestica/resp_rawest.dat','w'); fprintf(fp,'%g\n',resp); fclose(fp); exit;"
 else
  echo "skipping stage 2 (template coregistration and estimation)"
 fi


 if [ $stage3flag -eq 1 ] ; then
  echo ""
  echo "Running Stage 3: Filtering PESTICA estimators, cardiac first, then respiratory"
  echo ""
  echo "NOTE: TR must be set correctly in header for 3D+time dataset - if in doubt, check it and correct it"
  echo "Values given below:"
  echo  matlab $MATLABLINE -r "addpath $PESTICA_DIR; load('$epi_pestica/card_rawest.dat'); load('$epi_pestica/resp_rawest.dat'); card=view_and_correct_estimator(card_rawest,''$epi'','c'); resp=view_and_correct_estimator(resp_rawest,''$epi'','r'); fp=fopen('$epi_pestica/card_pestica.dat','w'); fprintf(fp,'%g\n',card); fclose(fp); fp=fopen('$epi_pestica/resp_pestica.dat','w'); fprintf(fp,'%g\n',resp); fclose(fp);"
        matlab $MATLABLINE -r "addpath $PESTICA_DIR; load('$epi_pestica/card_rawest.dat'); load('$epi_pestica/resp_rawest.dat'); card=view_and_correct_estimator(card_rawest,'$epi','c'); resp=view_and_correct_estimator(resp_rawest,'$epi','r'); fp=fopen('$epi_pestica/card_pestica.dat','w'); fprintf(fp,'%g\n',card); fclose(fp); fp=fopen('$epi_pestica/resp_pestica.dat','w'); fprintf(fp,'%g\n',resp); fclose(fp); exit;"
 else
  echo "skipping stage 3 (cardiac and respiratory vector temporal filtering)"
 fi

 if [ $stage4flag -eq 1 ] ; then
  echo ""
  echo "Running Stage 4: Get phase (using 3dretroicor), then get IRFs"
  if [ $retroicorflag -eq 1 ] ; then
    echo "And run traditional RETROICOR algorithm in MATLAB"
  fi
  echo ""
  # PESTICA convert to phase using 3dretroicor
  rm retroicor+orig.*
  echo "3dretroicor -threshold 0 -order 1 -card $epi_pestica/card_pestica.dat -cardphase $epi_pestica/cardphase.dat -resp $epi_pestica/resp_pestica.dat -respphase $epi_pestica/respphase.dat $epi"
  3dretroicor -threshold 0 -order 1 -card $epi_pestica/card_pestica.dat -cardphase $epi_pestica/cardphase.dat -resp $epi_pestica/resp_pestica.dat -respphase $epi_pestica/respphase.dat $epi
  rm retroicor+orig.*

  if [ $retroicorflag -eq 1 ] ; then
    echo "Running MATLAB-version of RETROICOR"
    rm -f retroicor+orig.????
    echo "matlab  $MATLABLINE -r addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; load $epi_pestica/cardphase.dat; load $epi_pestica/respphase.dat; matlab_retroicor('$correpi',cardphase,respphase,2,'$PESTICA_SLICE_TIMING','$epi_mask');"
         matlab  $MATLABLINE -r "addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; load $epi_pestica/cardphase.dat; load $epi_pestica/respphase.dat; matlab_retroicor('$correpi',cardphase,respphase,2,'$PESTICA_SLICE_TIMING','$epi_mask'); exit;"
    rm -f retroicor.nii
    3dAFNItoNIFTI -float retroicor+orig
    fname=`echo $correpi | sed "s/\+orig//"`
    rm -f $fname.retroicor+orig.????
    3dcopy retroicor.nii $fname.retroicor
    rm retroicor+orig.* retroicor.nii
  fi

  # PESTICA: use physio phase to get IRFs in data - the cardiac estimators are dithered slightly due to the filtering (we will fix that below...)
  echo "Getting IRFs - first round"
  echo "matlab  $MATLABLINE -r addpath $MATLAB_AFNI_DIR; addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; load $epi_pestica/cardphase.dat; load $epi_pestica/respphase.dat; retroicor_get_irf('$correpi',cardphase,respphase,5,'$PESTICA_SLICE_TIMING','$epi_mask');"
        matlab  $MATLABLINE -r "addpath $MATLAB_AFNI_DIR; addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; load $epi_pestica/cardphase.dat; load $epi_pestica/respphase.dat; retroicor_get_irf('$correpi',cardphase,respphase,5,'$PESTICA_SLICE_TIMING','$epi_mask'); exit;"

  cp impulse_responses.mat $epi_pestica/impulse_responses_firstpass.mat
  mv impulse_responses.mat $epi_pestica
 else
  echo "skipping stage 4 (first-pass, or only pass, plus original RETROICOR, if selected)"
 fi


 if [ $stage5flag -eq 1 ] ; then
  echo ""
  echo "Running Stage 5: Using first-pass IRFs, adjust for dither in cardiac peaks, then re-obtain IRFs and correct data"
  echo ""
  # use the data to adjust the dithering of each cardiac peak (when using parallel monitored pulse, this doesn't do much, but it does fix the PESTICA cardiac estimators)
  echo "adjusting and getting IRFs - second round"
  echo "matlab  $MATLABLINE -r addpath $MATLAB_AFNI_DIR; addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; load $epi_pestica/impulse_responses.mat; fit_each_physio_peak(cardph,coeffs_card,'$correpi','$epi_mask',cmask,'$PESTICA_SLICE_TIMING'); load physio_fitted.mat; retroicor_get_irf('$epi',cardphase,respph,5,'$PESTICA_SLICE_TIMING','$epi_mask');"
        matlab  $MATLABLINE -r "addpath $MATLAB_AFNI_DIR; addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; load $epi_pestica/impulse_responses.mat; fit_each_physio_peak(cardph,coeffs_card,'$correpi','$epi_mask',cmask,'$PESTICA_SLICE_TIMING'); load physio_fitted.mat; retroicor_get_irf('$epi',cardphase,respph,5,'$PESTICA_SLICE_TIMING','$epi_mask'); exit;"

  mv impulse_responses.mat $epi_pestica/impulse_responses_secondpass.mat
  mv physio_fitted.mat $epi_pestica/

  echo "Applying IRFs and correcting data with IRF-RETROICOR"
  echo "writing out corrected data as: irfretroicor"
  echo "writing out statistical coupling maps as: coupling_irfret_card coupling_irfret_resp"
  rm -f irfretroicor+orig.????
  echo "matlab  $MATLABLINE -r addpath $MATLAB_AFNI_DIR; addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; load $epi_pestica/impulse_responses_secondpass.mat; irf_retroicor('$correpi',cardph,respph,coeffs_card,coeffs_resp,'$PESTICA_SLICE_TIMING','$epi_mask');"
        matlab  $MATLABLINE -r "addpath $MATLAB_AFNI_DIR; addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; load $epi_pestica/impulse_responses_secondpass.mat; irf_retroicor('$correpi',cardph,respph,coeffs_card,coeffs_resp,'$PESTICA_SLICE_TIMING','$epi_mask'); exit;"
  rm -f irfretroicor.nii
  3dAFNItoNIFTI -float irfretroicor+orig
  fname=`echo $correpi | sed "s/\+orig//"`
  rm -f $fname.irfretroicor+orig.????
  3dcopy irfretroicor.nii $fname.irfretroicor
  rm irfretroicor+orig.* irfretroicor.nii
 else
  echo "skipping stage 5 (IRF-RETROICOR)"
 fi


 if [ $stage6flag -eq 1 ] ; then
  echo ""
  echo "Running Stage 6: Make QA plots"
  echo ""
  # run QA on output files
  echo "matlab $MATLABLINE -r  addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; physio_qa('$epi');"
        matlab $MATLABLINE -r "addpath $PESTICA_DIR; addpath $PESTICA_DIR/matlab_retroicor; physio_qa('$epi'); exit;"
  
  # change this if the plots always give a poor view of the slices - slice 20 in AFNI is reasonable for most acquisitions
  zpos=20
  fname=`echo $epi | sed "s/\+orig//"`
  # threshold for cardiac/respiratory coupling is ideally detected from the data itself, but may have to be adjusted manually
  thresh=`head -n 1 $epi_pestica/coupling_thresholds.txt`
  afni -com 'OPEN_WINDOW A.axialimage mont=4x4:2:0:none opacity=6' -com "SET_THRESHOLD A.$thresh 1"  -com 'SET_PBAR_NUMBER A.12' \
       -com 'SET_FUNC_RESAM A.Cu.Cu' -com "SET_UNDERLAY A.$fname" -com 'SET_FUNC_RANGE A.10' \
       -com "SET_OVERLAY A.coupling_irfret_card 1 0"        -com 'OPEN_PANEL A.Define_Overlay' -com 'SET_FUNC_VISIBLE A.+' \
       -com "SET_IJK A 64 64 $zpos" -com 'SET_XHAIRS A.OFF' -com "SAVE_JPEG A.axialimage $epi_pestica/pestica_card_coupling_overlay.jpg" \
       -com 'QUIT' -dset $fname+orig coupling_irfret_card+orig

  thresh=`tail -n 1 $epi_pestica/coupling_thresholds.txt`
  afni -com 'OPEN_WINDOW A.axialimage mont=4x4:2:0:none opacity=6' -com "SET_THRESHOLD A.$thresh 1"  -com 'SET_PBAR_NUMBER A.12' \
       -com 'SET_FUNC_RESAM A.Cu.Cu' -com "SET_UNDERLAY A.$fname" -com 'SET_FUNC_RANGE A.10' \
       -com "SET_OVERLAY A.coupling_irfret_resp 1 0"        -com 'OPEN_PANEL A.Define_Overlay' -com 'SET_FUNC_VISIBLE A.+' \
       -com "SET_IJK A 64 64 $zpos" -com 'SET_XHAIRS A.OFF' -com "SAVE_JPEG A.axialimage $epi_pestica/pestica_resp_coupling_overlay.jpg" \
       -com 'QUIT' -dset $fname+orig coupling_irfret_resp+orig

  afni -com 'OPEN_WINDOW A.axialimage mont=4x4:2:0:none opacity=4' -com 'SET_THRESHOLD A.5 1'  -com 'SET_PBAR_NUMBER A.12' \
       -com 'SET_FUNC_RANGE A.30'    -com 'SET_FUNC_RESAM A.Cu.Cu' -com "SET_UNDERLAY A.$fname" \
       -com "SET_OVERLAY A.card_PESTICA2.nii 2 2"        -com 'OPEN_PANEL A.Define_Overlay' -com 'SET_FUNC_VISIBLE A.+' \
       -com "SET_IJK A 64 64 $zpos" -com 'SET_XHAIRS A.OFF' -com "SAVE_JPEG A.axialimage $epi_pestica/pestica_ctemplate_overlay.jpg" \
       -com 'QUIT' -dset $fname+orig $epi_pestica/card_PESTICA2.nii 

  afni -com 'OPEN_WINDOW A.axialimage mont=4x4:2:0:none opacity=4' -com 'SET_THRESHOLD A.5 1'  -com 'SET_PBAR_NUMBER A.12' \
       -com 'SET_FUNC_RANGE A.30'    -com 'SET_FUNC_RESAM A.Cu.Cu' -com "SET_UNDERLAY A.$fname" \
       -com "SET_OVERLAY A.resp_PESTICA2.nii 2 2"        -com 'OPEN_PANEL A.Define_Overlay' -com 'SET_FUNC_VISIBLE A.+' \
       -com "SET_IJK A 64 64 $zpos" -com 'SET_XHAIRS A.OFF' -com "SAVE_JPEG A.axialimage $epi_pestica/pestica_rtemplate_overlay.jpg" \
       -com 'QUIT' -dset $fname+orig $epi_pestica/resp_PESTICA2.nii
 else
  echo "skipping stage 6 (QA)"
 fi

done

