#!/bin/bash

#   Copyright (C) 2020 University of Oxford
#
#   SHCOPYRIGHT

# Written by Saad Jbabdi, Stam Sotiropoulos & Shaun Warrington

shopt -s extglob

code_folder="$FSLDIR/bin"

# Location of probtrackx2_gpu binary
ptxbin_gpu=$FSLDIR/bin/probtrackx2_gpu

Splash (){
cat <<EOF

__  _______ ____      _    ____ _____ _     _                       _       _
\ \/ /_   _|  _ \    / \  / ___|_   _| |__ | |_   _  ___ _ __  _ __(_)_ __ | |_
 \  /  | | | |_) |  / _ \| |     | | | '_ \| | | | |/ _ \ '_ \| '__| | '_ \| __|
 /  \  | | |  _ <  / ___ \ |___  | | | |_) | | |_| |  __/ |_) | |  | | | | | |_
/_/\_\ |_| |_| \_\/_/   \_\____| |_| |_.__/|_|\__,_|\___| .__/|_|  |_|_| |_|\__|
                                                        |_|

EOF
}
Splash


Usage() {
    cat << EOF

Usage:
    xtract_blueprint -bpx <folder> -out <folder> -xtract <folder> -seeds <list> [options]

    Compulsory arguments:

       -bpx       <folder>                          Path to bedpostx folder
       -out       <folder>                          Path to output folder
       -xtract    <folder>                          Path to xtract folder
       -seeds     <list>                            Comma separated list of seeds for which a blueprint is requested (e.g. left and right cortex in standard space)
       -warps     <ref> <xtract2diff> <diff2xtract> Standard space reference image, and transforms between xtract space and diffusion space

    Optional arguments:
       -stage                                       What to run. 1:matrix2, 2:blueprint, all:everythng (default)
       -gpu                                         Use GPU version
       -savetxt                                     Save blueprint as txt file (nseed by ntracts) instead of CIFTI
       -prefix    <string>                          Specify a prefix for the final blueprint filename (e.g. <prefix>_BP.LR.dscalar.nii)

       -rois      <list>                            Comma separated list (1 per seed): ROIs (gifti) to restrict seeding (e.g. medial wall masks)
       -stops     <stop.txt>                        Text file containing line separated list
       -wtstops   <wtstop.txt>                      Text file containing line separated list
       -tract_list                                  Comma separated list of tracts to include (default = all found under -xtract <folder>)

       -thr                                         Threshold applied to XTRACT tracts prior to blueprint calculation (default = 0.001, i.e. 0.1% probability).
       -nsamples                                    Number of samples per seed used in tractography (default = 1000)
       -res       <mm>                              Resolution of matrix2 output (Default = 3 mm)
       -ptx_options <options.txt>                   Pass extra probtrackx2 options as a text file to override defaults

   Example (recommended) usage:
      xtract_blueprint -bpx <bpxdir> -out <outdir> -xtract <xtractdir> -seeds <l.white.surf.gii,r.white.surf.gii> \
           -rois <l.medwall.shape.gii,r.medwall.shape.gii> -warps <ref> <xtract2diff> <diff2xtract> -gpu \

EOF
    exit 1
}
#       -distnorm                                    Normalise tracts by distance

[ "$1" = "" ] && Usage

# Set default options
bpx=""
out=""
xtract=""
seeds=""
stops=""
wtstops=""
rois=""
distnorm=0
ptx_opts=""
stdref=""
gpu=0
ref=""
diff2xtract=""
xtract2diff=""
diff2seed=""
seed2diff=""
spec=""
res=3
stage=all
savetxt=0
thr=0.001
nsamples=1000
tract_list=""
prefix="x" # if prefix not specified, pass 'x' to create_blueprint which is nulled

# Parse command-line arguments
while [ ! -z "$1" ];do
    case "$1" in
	-bpx) bpx=$2;shift;;
	-out) out=$2;shift;;
	-seeds) seeds=(`echo $2 | sed s@","@" "@g`);shift;;
  -stops) stops=(`echo $2 | sed s@","@" "@g`);shift;;
  -wtstops) wtstops=(`echo $2 | sed s@","@" "@g`);shift;;
  -rois) rois=(`echo $2 | sed s@","@" "@g`);shift;;
	-xtract) xtract=$2;shift;;
	-warps) ref=$2;xtract2diff=$3;diff2xtract=$4;shift;shift;shift;;
	-gpu) gpu=1;;
	-res) res=$2;shift;;
	-stage) stage=$2;shift;;
	-ptx_options) ptx_options=`cat $2`;shift;;
  -savetxt) savetxt=1;;
  -thr) thr=$2;shift;;
  -tract_list) tract_list=$2;shift;;
  -prefix) prefix=$2;shift;;
  -nsamples) nsamples=$2;shift;;
	*) echo "Unknown option '$1'";exit 1;;
    esac
    shift
done

#	-distnorm) distnorm=1;;

# Step 0 : Prepare
# Check compulsory arguments
errflag=0
if [ "$bpx" == "" ];then
    echo "Must set compulsory argument '-bpx'"
    errflag=1
elif [ ! -d $bpx ];then
    echo "Bedpostx folder $bpx not found"
    errflag=1
fi
if [ "$xtract" == "" ];then
    echo "Must set compulsory argument '-xtract'"
    errflag=1
elif [ ! -d $xtract ];then
    echo "XTRACT folder $xtract not found"
    errflag=1
fi
if [ "$out" == "" ];then
    echo "Must set compulsory argument '-out'"
    errflag=1
fi
if [ "$seeds" == "" ];then
    echo "Must set compulsory argument '-seeds'"
    errflag=1
fi

# Check seeds and stops (check they exist and 1 stop per seed if using)
for seed in ${seeds[@]};do
  if [ ! -f $seed ];then
    echo "Seed file $seed not found"
    errflag=1
  fi
done

nseeds=${#seeds[@]}
nrois=${#rois[@]}
if [ ! "$stops" == "" ]; then
    if [ ! -f $stop ];then
      echo "Stop file $stop not found"
      errflag=1
    fi
fi
if [ ! "$wtstops" == "" ]; then
    if [ ! -f $stop ];then
      echo "wtstop file $wtstop not found"
      errflag=1
    fi
fi
if [ ! "$rois" == "" ]; then
    if [ ! $nrois -eq $nseeds ]; then
      echo "Supplied $nseeds seeds but $nrois ROIs"
      echo "If using ROIs, must supply 1 per seed"
      errflag=1
    fi
    for roi in ${rois[@]};do
      if [ ! -f $roi ];then
        echo "ROI file $roi not found"
        errflag=1
      fi
    done
fi

# check ref and warps
if [ ! -f $ref ];then
  echo "Reference file $ref not found"
  errflag=1
fi
if [ ! -f $xtract2diff ];then
  echo "Warp file $xtract2diff not found"
  errflag=1
fi
if [ ! -f $diff2xtract ];then
  echo "Warp file $diff2xtract not found"
  errflag=1
fi

if [ "$errflag" -eq 1 ];then
    echo ""
    echo "Exit without doing anything..."
    exit 1
fi

# GPU stuff
if [ "$gpu" == 0 ];then
    ptxbin=$FSLDIR/bin/probtrackx2
    echo "Warning: not using GPU mode - this may take a while. Consider downsampling the seed mask."
else
    ptxbin=${ptxbin_gpu}
fi

# what are we running?
if [ $stage == 1 ];then
    echo "Running tractography only"
elif [ $stage == 2 ];then
    echo "Running blueprint processing only"
elif [ $stage == "all" ];then
    echo "Running tractography and blueprint processing"
else
    echo "Unknown stage option. Either: '1' for matrix2 tractography, '2' for blueprint or 'all' for both"
fi

if [ -d $out ];then
    echo "Warning: Output folder already exists. Some of the files may be overwritten"
fi
mkdir -p $out
mkdir -p $out/omatrix2

# Get tract list here
# if not provided by user, get list from under xtract folder
function join_str { local IFS="$1"; shift; echo "$*"; }
if [ "$tract_list" == "" ]; then
    # if not supplied, get tracts from xtract folder
    struct=(`ls ${xtract}/tracts/*/densityNorm.nii.gz`)
    unset tracts
    for t in ${struct[@]}; do
        t=$( echo ${t%/*} )
        t=$( echo ${t##*/} )
        tracts+=($t)
    done
    # also build comma separated tract_list for python usage later
    tract_list=`join_str , ${tracts[@]}`
else
    tracts=(`echo $tract_list | sed s@","@" "@g`)
fi

# function for image resmapling
resample(){
    # resample <src_img> <mm_resolution> <out_img> <interp>
    i=$1
    r=$2
    o=$3
    interp=$4
    ${FSLDIR}/bin/flirt -in $i -out $o -ref $i -applyisoxfm $r -interp $interp
}

# Deal with warps - if only seed->diff assume it is the same as xtract->diff
# Resample ref to get target2
if (( $(echo "$res > 0" | bc -l) ));then
    resample $ref $res $out/omatrix2/target2 nearestneighbour
else
    ${FSLDIR}/bin/imcp $ref $out/omatrix2/target2
fi

# Step 1 : Run omatrix2 for all seeds
ptxopts=" --samples=$bpx/merged --mask=$bpx/nodif_brain_mask "
ptxopts="$ptxopts --target2=$out/omatrix2/target2 --xfm=$xtract2diff --invxfm=$diff2xtract --seedref=$ref "
ptxopts="$ptxopts --omatrix2 --loopcheck --forcedir --opd --nsamples=$nsamples "

conv_ascii(){
    i=$1
    roi=$2
    d=$3

    o=`basename $i`
    o=`echo ${o%%.@(nii|nii.gz|gii|surf.gii)}`
    o="${d}/${o}.asc"

    ${FSLDIR}/bin/surf2surf -i ${i} -o ${o} --outputtype=ASCII --values=${roi}
    echo $o
}

if [ ! $stage == 2 ];then
  commands_1=$out/omatrix2/ptx_commands.txt
  rm -rf $commands_1
  COUNTER=0
  for seed in ${seeds[@]};do
      r=`basename $seed`
      r=`echo ${r%%.@(nii|nii.gz|gii|surf.gii|asc)}`   # remove extensions
      echo -ne "seed $r"
      mkdir -p $out/omatrix2/omatrix2_$r

      s=$seed
      # convert gii to ascii with medial wall if using
      if [ ! "$rois" == "" ]; then
        roi=${rois[COUNTER]}
        s=`conv_ascii ${seed} ${roi} $out/omatrix2/omatrix2_$r`
        s=($s)
        s=${s[$((${#s[@]}-1))]}

        m=`basename $roi`
        m=`echo ${m%%.@(nii|nii.gz|gii|shape.gii|asc)}`
        echo -ne " : medial wall $m"
      fi
      hemi_opts=" $ptxopts --seed=$s --dir=$out/omatrix2/omatrix2_$r"

      # add stop mask if using
      if [ ! "$stops" == "" ]; then
        hemi_opts=" $hemi_opts --stop=$stops"
        r=`basename $stops`
        r=`echo ${r%%.@(nii|nii.gz|gii|surf.gii|asc)}`
        echo -ne " : stop $r"
      fi

      # add wtstop mask if using
      if [ ! "$wtstops" == "" ]; then
        hemi_opts=" $hemi_opts --wtstop=$wtstops"
        r=`basename $wtstops`
        r=`echo ${r%%.@(nii|nii.gz|gii|surf.gii|asc)}`
        echo -ne " : wtstop $r"
      fi

      echo ""
      # Append extra options
      hemi_opts=" $hemi_opts $ptx_options"
      echo "$ptxbin $hemi_opts" >> $commands_1
      COUNTER=$((COUNTER+1))
      chmod +x $commands_1
  done
fi


# Step 2 : Resample xtract, load all xtract thingies and matrix2 output and multiply them
if [ ! $stage == 1 ];then
  # copy tracts to temp folder
  mkdir -p $out/temp_xtract
  commands_2=$out/bp_commands.txt
  rm -rf $commands_2

  for t in ${tracts[@]}; do
    cp ${xtract}/tracts/${t}/densityNorm.nii.gz ${out}/temp_xtract/${t}.nii.gz
  done
  cmd=""
  # re-sample tracts if needed - assumes that all xtract tracts are equal dimensions - they should be!
  if [ ! "`fslorient -getsform $out/omatrix2/target2`" == "`fslorient -getsform ${out}/temp_xtract/${tracts[0]}.nii.gz`" ]; then
    echo "XTRACT resolution different from matrix2:"
    echo " --- will resample xtract tracts to $res mm"
    for t in ${tracts[@]}; do
      cmd="$cmd ${FSLDIR}/bin/flirt -in ${out}/temp_xtract/$t -out ${out}/temp_xtract/$t -ref $out/omatrix2/target2 -applyisoxfm $res -interp trilinear;"
    done
  fi

  if (( $(echo "$thr > 0" | bc -l) ));then
    for t in ${tracts[@]}; do
      cmd="$cmd ${FSLDIR}/bin/fslmaths ${out}/temp_xtract/$t -thr $thr ${out}/temp_xtract/$t;"
    done
  fi

  # loop through seeds and prepare commands
  # start of python command
  cmd="$cmd $code_folder/create_blueprint.py ${out}/temp_xtract"
  p=""
  for seed in ${seeds[@]}; do
      r=`basename $seed`
      r=`echo ${r%%.@(nii|nii.gz|gii|surf.gii|asc)}`    # remove extensions

      p="$p $out/omatrix2/omatrix2_$r"
  done
  py_args=`join_str , $p`
  py_args="$py_args `join_str , ${seeds[@]}`"

  # if using medial wall
  if [ ! "$rois" == "" ]; then
      py_args="$py_args `join_str , ${rois[@]}`"
  else
      py_args="$py_args x"
  fi
  cmd="$cmd $py_args $tract_list $distnorm $savetxt $prefix; rm -rf ${out}/temp_xtract"
  echo "$cmd" >> $commands_2
  chmod +x $commands_2
fi


run(){
    # run <command> [job_id to wait for]
    commands=$1
    out=$2
    step=$3
    JID=$4
    waitforme=""
    if [ "x$JID" != "x" ];then
	waitforme=" -j $JID"
    fi
    if [ $step == 1 ]; then
    if [ "x$SGE_ROOT" != "x" ]; then  # Submit all commands to run in parallel on the cluster
	if [ $gpu -eq 0 ];then
	    JID=`fsl_sub -T 1440 -q bigmem.q $waitforme -l $out/logs -N blueprint_track -t $commands`
	else
	    JID=`fsl_sub -T 300 -q $FSLGECUDAQ $waitforme -l $out/logs -N blueprint_track -t $commands`
	fi
    else   # If no SGE, run locally
	sh $commands
    fi
  elif [ $step == 2 ]; then
    if [ "x$SGE_ROOT" != "x" ]; then  # Submit all commands to run in parallel on the cluster
	    JID=`fsl_sub -T 300 -q bigmem.q $waitforme -l $out/logs -N xtract_blueprint -t $commands`
    else   # If no SGE, run locally
	sh $commands
    fi
  fi
  echo $JID
}

echo ""
echo "Launching job(s)"
echo ""
# Launch jobs
if [ $stage == 1 ];then
    run $commands_1 $out $stage
elif [ $stage == 2 ];then
    run $commands_2 $out $stage
else
  if [ "x$SGE_ROOT" != "x" ]; then
    jid=`eval run $commands_1 $out 1`
    echo "probtrack jobID: $jid"
    jid=`eval run $commands_2 $out 2 $jid`
    echo "blueprint jobID: $jid"
  else
    run $commands_1 $out 1
    run $commands_2 $out 2
  fi
fi

exit 1

#EOF