#!/bin/tcsh -f

#
# parc_atlas_jackknife_test
#
# Performs 'jackknife' estimation of classifier accuracy.
#
# Given a set of N manually labeled subjects, N atlases trained on N-1 number 
# of subjects are created, and the one excluded subject is used to test the
# accuracy of this atlas, based on the Dice coefficient(s).
# The overall classifier accuracy is the mean of all the results.
#
# Original author: Nick Schmansky
#
# 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='parc_atlas_jackknife_test dev';

# defaults:
# setenv SUBJECTS_DIR /space/amaebi/26/users/buckner_cortical_atlas
set DoRegister=0
set DoRegCopy=0
set DoTrain=0
set DoClassify=0
set DoTest=0
set DoDiceAvg=0
set RunIt=1
set PBS=0 # see next line: if on cluster machine, use pbsubmit
if ("$HOST" == "seychelles") set PBS=1
set pbs_flags="-l nodes=1:opteron"
set reg_dist=""
set reg_append=""
set reg_arg=""  # allow specifying an arbitrary argument to mris_register

# parse command line
set cmdline = ($argv);
while( $#argv != 0 )

  set flag = $argv[1]; shift;

  switch($flag)

    case "-register":
    case "--register":
      set DoRegister=1
      breaksw

    case "-reg_arg":
    case "--reg_arg":
      set DoRegister=1
      set reg_arg=" $argv[1]"; shift;
      breaksw

    case "-reg_dist":
    case "--reg_dist":
      set DoRegister=1
      set reg_dist="-dist $argv[1]"; shift;
      breaksw

    case "-reg_append":
    case "--reg_append":
      set DoRegister=1
      set reg_append="$argv[1]"; shift;
      breaksw

    case "-reg_copy":
    case "--reg_copy":
      set DoRegCopy=1
      set reg_append="$argv[1]"; shift;
      breaksw

    case "-reg_append_and_copy":
    case "--reg_append_and_copy":
      set DoRegister=1
      set DoRegCopy=1
      set reg_append="$argv[1]"; shift;
      breaksw

    case "-train":
    case "--train":
      set DoTrain=1
      breaksw

    case "-classify":
    case "--classify":
      set DoClassify=1
      breaksw

    case "-test":
    case "--test":
      set DoTest=1
      breaksw

    case "-avg":
    case "--avg":
      set DoDiceAvg=1
      breaksw

    case "-all":
    case "--all":
      set DoTrain=1
      set DoClassify=1
      set DoTest=1
      set DoDiceAvg=1
      breaksw

    case "-sd":
    case "--sd":
      setenv SUBJECTS_DIR $argv[1]; shift;
      breaksw

    case "-fshome":
    case "--fshome":
      setenv FREESURFER_HOME $argv[1]; shift;
      source $FREESURFER_HOME/SetUpFreeSurfer.csh
      breaksw

    case "-binhome":
    case "--binhome":
      if ( ! $?FREESURFER_HOME) then 
        source /usr/local/freesurfer/nmr-dev-env
      endif
      set BIN_HOME = $argv[1]; shift;
      set path = ( ${BIN_HOME} $path )
      rehash
      breaksw

    case "-dontrun":
    case "--dontrun":
      set RunIt=0
      breaksw

    case "-mail":
    case "--mail":
      set pbs_flags=(${pbs_flags} -m $USER)
      breaksw

    case "-help":
    case "--help":
      echo "parc_atlas_jackknife_test [options]"
      echo ""
      echo "options:"
      echo "  -register         run mris_register: creates .sphere.reg files"
      echo "  -reg_dist <arg>   run mris_register with '-dist <arg>' flag"
      echo "  -reg_append <string>  append <string> to end of ?h.sphere.reg"
      echo "  -reg_copy <string> cp ?h.sphere.reg<string> ?h.sphere.reg"
      echo "  -train            run mris_ca_train: creates .gcs files"
      echo "  -classify         run mris_ca_label: creates .annot files"
      echo "  -test             run mris_compute_parc_overlap"
      echo "  -all              -train, -classify, -test"
      echo "  -sd <subjects dir>          override default subjects dir"
      echo "  -fshome <FREESURFER_HOME>   source a new FREESURFER_HOME "
      echo "  -binhome <path to binaries> specify override path to binaries"
      echo "  -dontrun          dont execute the commands"
      echo ""
      echo "example1:"
      echo "  Assuming the subjects to test are listed in the SUBJECTS var"
      echo "  declared in the scripts/subjects.csh file, and that"
      echo "  each of these subjects has a surf/?h.sphere.reg and a manually"
      echo "  edited file named surf/?h.aparc_edited.annot, and there exists"
      echo "  a color LUT named scripts/colortable_final.txt, then the"
      echo "  following command will perform a jackknife accuracy test: "
      echo ""
      echo "  parc_atlas_jackknife_test -all"
      echo ""
      echo "  The test is conducted as follows: Given a set of N manually "
      echo "  labeled subjects, N atlases trained on N-1 number of subjects "
      echo "  are created, and the one excluded subject is used to test the"
      echo "  accuracy of this atlas, based on the Dice coefficient(s)."
      echo "  The overall classifier accuracy is the mean of all the results."
      echo "  The results are written to the directory named 'jackknife'."
      echo ""
      exit 1
      breaksw

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

end

if ( ! $DoRegister && \
     ! $DoRegCopy && \
     ! $DoTrain && \
     ! $DoClassify && \
     ! $DoTest && \
     ! $DoDiceAvg) then
    echo "Nothing to do!  Use --help for options"
    exit 1
endif

echo "SUBJECTS_DIR    $SUBJECTS_DIR"
if ( ! $?FREESURFER_HOME) then 
   source /usr/local/freesurfer/nmr-dev-env
endif
echo "FREESURFER_HOME $FREESURFER_HOME"

#
# check for necessary binaries, script, and input files
#
echo "Register: `which mris_register`"
if ($status) then
    echo "mris_register not found (use -fshome or -binhome)"
    exit 1
endif
echo "Train:    `which mris_ca_train`"
if ($status) then
    echo "mris_ca_train not found (use -fshome or -binhome)"
    exit 1
endif
echo "Classify: `which mris_ca_label`"
if ($status) then
    echo "mris_ca_label not found (use -fshome or -binhome)"
    exit 1
endif
echo "Test:     `which mris_compute_parc_overlap`"
if ($status) then
    echo "mris_compute_parc_overlap not found (use -fshome or -binhome)"
    exit 1
endif
if ( ! -e ${SUBJECTS_DIR}/scripts/subjects.csh) then
    echo "${SUBJECTS_DIR}/scripts/subjects.csh is missing!"
    echo "This file should SUBJECTS to all the subjects to test."
    exit 1
else
    # these are the subjects to use in training, leaving one out for testing
    source ${SUBJECTS_DIR}/scripts/subjects.csh
endif
if ( ! -e ${SUBJECTS_DIR}/scripts/colortable_final.txt) then
    echo "$SUBJECTS_DIR/scripts/colortable_final.txt is missing!"
    exit 1
endif
foreach s ($SUBJECTS)
    foreach hemi (rh lh)
        set f=$SUBJECTS_DIR/$s/surf/$hemi.sphere.reg
        if ( ! -e $f) then
            echo "$f is missing!"
            if ( ! $DoRegister && ! $DoRegCopy ) exit 1
        endif
        set f=$SUBJECTS_DIR/$s/label/$hemi.aparc_edited.annot
        if ( ! -e $f) then
            echo "$f is missing!"
            echo "This is the manually labeled surface."
            exit 1
        endif
    end
end


cd $SUBJECTS_DIR
mkdir -p jackknife
cd jackknife
setenv WD $PWD
echo "WD:       $WD"

set ALL_SUBJECTS=($SUBJECTS)


#
# Optionally create the sphere.reg file
#
if ($DoRegister) then
set tif="average.CURVATURE.tif"
foreach test_subj ($ALL_SUBJECTS)
    foreach hemi (rh lh)
        set sphere_reg="${SUBJECTS_DIR}/${test_subj}/surf/${hemi}.sphere.reg"
        set sphere_reg="${sphere_reg}${reg_append}"
        set cmd=(mris_register ${reg_arg} ${reg_dist} \
            ${SUBJECTS_DIR}/${test_subj}/surf/${hemi}.sphere \
            ${SUBJECTS_DIR}/average/${hemi}.${tif} \
            ${sphere_reg})
        echo "$cmd"
        set LF="${SUBJECTS_DIR}/${test_subj}/surf"
        set LF="${LF}/${hemi}.mris_register${reg_append}.log"
        if ($RunIt) then
            echo "`which mris_register`" |& tee ${LF}
            echo "${cmd}" |& tee -a ${LF}
            if (-e ${sphere_reg}) rm -f ${sphere_reg}
            if ($PBS) then
                set cmdf=(${WD}/${hemi}.${test_subj}${reg_append}.cmd)
                echo "#! /bin/tcsh -ef" > ${cmdf}
                echo "limit filesize 10megabytes" >> ${cmdf}
                echo "${cmd} |& tee -a ${LF}" >> ${cmdf}
                chmod a+x ${cmdf}
                pbsubmit ${pbs_flags} -c "${cmdf}"
                sleep 30
            else
                ${cmd} |& tee -a ${LF} &
            endif
        endif
    end
end
foreach test_subj ($ALL_SUBJECTS)
    foreach hemi (rh lh)
        set sphere_reg="${SUBJECTS_DIR}/${test_subj}/surf/${hemi}.sphere.reg"
        set sphere_reg="${sphere_reg}${reg_append}"
        if ($RunIt) then
            while (! -e $sphere_reg)
                echo "waiting for ${sphere_reg}..."
                sleep 60
            end
        endif
    end
end
endif


#
# Optionally copy named sphere.reg files (those created using -reg_append)
#
if ($DoRegCopy) then
foreach test_subj ($ALL_SUBJECTS)
    foreach hemi (rh lh)
        set sphere_reg="${SUBJECTS_DIR}/${test_subj}/surf/${hemi}.sphere.reg"
        set copy_sphere_reg="${sphere_reg}${reg_append}"
        set cmd1=(mv ${sphere_reg} ${sphere_reg}.bak)
        set cmd2=(cp ${copy_sphere_reg} ${sphere_reg})
        if ($RunIt) then
            echo $cmd1
            ${cmd1}
            echo $cmd2
            ${cmd2}
            if ($status) exit 1
        endif
    end
end
endif


#
# for N number of subjects, create N atlases, where one subject is excluded
# from each atlas training set
#
if ($DoTrain) then
foreach test_subj ($ALL_SUBJECTS)

    # create a SUBJECTS array containing all subject *except* test_subj
    unsetenv SUBJECTS
    set SUBJECTS = (  )
    foreach s ($ALL_SUBJECTS)
        if ("$s" != "$test_subj") set SUBJECTS = ( $SUBJECTS $s )
    end

    # now train using these $SUBJECTS
    foreach hemi (rh lh)
        set atlas="${WD}/${hemi}.atlas_leaveout_${test_subj}.gcs"
        set cmd=(mris_ca_train \
            -t $SUBJECTS_DIR/scripts/colortable_final.txt \
            $hemi \
            sphere.reg \
            aparc_edited \
            $SUBJECTS \
            ${atlas})
        echo ${cmd}
        if ($RunIt) then
            if (-e ${atlas}) rm -f ${atlas}
            if ($PBS) then
                pbsubmit ${pbs_flags} -c "${cmd}"
                sleep 30
            else
                ${cmd} &
            endif
        endif
    end
end
endif
unsetenv SUBJECTS


#
# now run the automatic parcellation on each subject, running against
# the atlas built w/o that subject
#
if ($DoClassify) then
foreach test_subj ($ALL_SUBJECTS)
    foreach hemi (rh lh)
        set atlas="${WD}/${hemi}.atlas_leaveout_${test_subj}.gcs"
        if ($RunIt) then
            while (! -e $atlas)
                echo "waiting for $atlas..."
                sleep 10
            end
        endif
        set annot="$SUBJECTS_DIR/$test_subj/label/$hemi.aparc_jackknife.annot"
        set cmd=(mris_ca_label \
            -seed 1234 \
            -t ${SUBJECTS_DIR}/scripts/colortable_final.txt \
            ${test_subj} \
            ${hemi} \
            sphere.reg \
            ${atlas} \
            ${annot} )
        echo ${cmd}
        if ($RunIt) then
            if (-e ${annot}) rm -f ${annot}
            if ($PBS) then
                pbsubmit ${pbs_flags} -c "${cmd}"
                sleep 30
            else
                ${cmd} &
            endif
        endif
    end
end
endif


#
# now test each atlas by comparing the subjects automatic classification 
# against its manual labeling
#
if ($DoTest) then
foreach test_subj ($ALL_SUBJECTS)
    foreach hemi (rh lh)
        set manual="aparc_edited"
        set auto="aparc_jackknife"
        set f="${SUBJECTS_DIR}/${test_subj}/label/$hemi.${auto}.annot"
        if ($RunIt) then
            while (! -e $f)
                echo "waiting for ${f}..."
                sleep 10
            end
        endif
        set cmd=(mris_compute_parc_overlap \
            --s ${test_subj} \
            --hemi ${hemi} \
            --annot1 ${manual} \
            --annot2 ${auto} )
        echo ${cmd}
        if ($RunIt) then
            set outFile=${WD}/compute_parc_overlap_${test_subj}_${hemi}.txt
            if (-e ${outFile}) rm -f ${outFile}
            if ($PBS) then
                pbsubmit ${pbs_flags} -c "${cmd} > ${outFile}"
            else
                ${cmd} >& ${outFile} &
            endif
        endif
    end
end
endif


#
# Calc the overall dice average. assumes the -test (DoTest) stage has run
#
if ($DoDiceAvg) then
foreach test_subj ($ALL_SUBJECTS)
    foreach hemi (rh lh)
        if ($RunIt) then
            set outFile=${WD}/compute_parc_overlap_${test_subj}_${hemi}.txt
            while ( ! -e $outFile)
                echo "waiting for ${outFile}..."
                sleep 10
            end
            cat $outFile > /dev/null
        endif
    end
end
sleep 2

echo "Calculating overall average Dice coefficient..."
set avgfile=${WD}/dice_avg.txt
if (-e ${avgfile}) rm -f ${avgfile}
set sum=0
@ count=0
foreach test_subj ($ALL_SUBJECTS)
    foreach hemi (rh lh)
        if ($RunIt) then
            set outFile=${WD}/compute_parc_overlap_${test_subj}_${hemi}.txt
            while ( ! -e $outFile)
                sleep 1
            end
            sleep 1
            set dice=`grep Dice ${outFile} | awk '{ print $4 }'`
            if ($status) then
                echo "Failed: grep Dice ${outFile} | awk '{ print $4 }"
                exit 1
            endif
            set sum=`echo "${sum} + ${dice}" | bc`
            if ($status) then
                echo "Failed: echo ${sum} + ${dice} | bc"
                exit 1
            endif
            @ count++
        endif
    end
end
set avg=`echo "scale=4; ${sum} / ${count}" | bc`
if ($status) then
    echo "Failed: echo scale=4; ${sum} / ${count} | bc"
    exit 1
endif
echo "${avg}" >& ${avgfile}
endif

echo "Calculating overall mean min distance..."
set avgfile=${WD}/mean_min.txt
if (-e ${avgfile}) rm -f ${avgfile}
set sum=0
@ count=0
foreach test_subj ($ALL_SUBJECTS)
    foreach hemi (rh lh)
        if ($RunIt) then
            set outFile=${WD}/compute_parc_overlap_${test_subj}_${hemi}.txt
            while ( ! -e $outFile)
                sleep 1
            end
            sleep 1
            set meanmin=`grep mean ${outFile} | awk '{ print $7 }'`
            if ($status) then
                echo "Failed: grep mean ${outFile} | awk '{ print $7 }'"
                exit 1
            endif
            set sum=`echo "scale=4; ${sum} + ${meanmin}" | bc`
            if ($status) then
                echo "Failed: echo scale=4; ${sum} + ${meanmin} | bc"
                exit 1
            endif
            @ count++
        endif
    end
end
set avg=`echo "scale=4; ${sum} / ${count}" | bc`
if ($status) then
    echo "Failed: echo scale=4; ${sum} / ${count} | bc"
    exit 1
endif
echo "${avg}" >& ${avgfile}
endif


#
# done
#
echo "parc_atlas_jackknife_test complete"
if (-e ${WD}/dice_avg.txt) echo "Average Dice: `cat ${WD}/dice_avg.txt`"
if (-e ${WD}/mean_min.txt) echo "Mean min: `cat ${WD}/mean_min.txt`"