#!/bin/bash

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

# Script to view output of XTRACT in fsleyes
# Written by Shaun Warrington 07/2019
# (Slightly modified by Saad Jbabdi 08/2019)

Usage() {
    cat << EOF

Usage:
    xtract_viewer -dir <xtractDir> -species HUMAN [options]
    xtract_viewer -dir <xtractDir> -species MACAQUE [options]
    xtract_viewer -dir <xtractDir> -brain <PATH> [options]

    Compulsory arguments:

       -dir <FOLDER>                     Path to XTRACT output folder

       And EITHER:
       -species <SPECIES>                One of HUMAN or MACAQUE

       OR:
       -brain <PATH>                     The brain image to use for the background overlay - must be in the same space as tracts.
                                         Default is the FSL_HCP065_FA map for HUMAN and F99 T1 brain for MACAQUE

    Optional arguments:

       -str STRUCTURE,STRUCTURE,...      Structures (comma separated (default = display all that is found in input folder)

       -thr NUMBER NUMBER                The lower and upper thresholds applied to the tracts for viewing
                                         Default = 0.001 0.1

EOF
    exit 1
}

Splash (){

cat <<EOF

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

EOF

}

Splash

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

# the colourmap options
cmaps=(blue red green blue-lightblue pink red-yellow cool yellow copper hot hsv coolwarm spring summer winter Oranges)
cL="${#cmaps[@]}" # length of array for colourmap loop control

# Parse command-line arguments
# Set defaults
thr=0.001
uthr=0.1
str=ALL
spec=""
while [ ! -z "$1" ];do
    case "$1" in
	-dir) dir=$2;shift;;
  -species) spec=$2;shift;;
	-str) IFS=', ' read -r -a str <<< "$2";shift;;
	-thr) thr=$2;uthr=$3;shift;shift;;
	-brain) brain=$2;shift;;
	*) echo "Unknown option '$1'";exit 1;;
    esac
    shift
done

# Check compulsory arguments
errflag=0
if [ "$dir" == "" ];then
    echo "Must set compulsory argument '-dir'"
    errflag=1
elif [ ! -d $dir ];then
    echo "Xtract folder $dir not found"
    errflag=1
fi

if [ "$spec" == "" ];then
  if [ "$brain" == "" ];then
    echo "Must set compulsory argument '-brain' OR specify '-species <SPECIES>'"
    errflag=1
  elif [ `$FSLDIR/bin/imtest $brain` -eq 0 ];then
    echo "Brain overlay file $brain not found"
    errflag=1
  fi
elif [ $spec == HUMAN ];then
    brain=${FSLDIR}/data/standard/FSL_HCP1065_FA_1mm.nii.gz
elif [ $spec == MACAQUE ];then
    brain=${FSLDIR}/data/xtract_data/standard/F99/mri/struct_brain.nii.gz
elif [ ! "$spec" == "HUMAN" ] && [ ! "$spec" == "MACAQUE" ];then
  echo "Unrecognised option '$spec'. Must set '-species' using HUMAN or MACAQUE"
  errflag=1
fi

if [ "$uthr" == "" ];then
  echo "You have set a lower threshold but not an upper threshold. Must set '-thr' using <number> <number>"
  errflag=1
fi
if [ "$errflag" -eq 1 ];then
  echo ""
  echo "Exit without doing anything.."
  exit 1
fi

# build struct array - removing any empty/comment lines
# and check for and remove any file extentions
if [ $str == ALL ];then
    str=(`ls $dir/tracts`)
    #str=`IFS=' ' read -r -a str <<< "$strlist"`
fi
echo Structures to be displayed
for index in "${!str[@]}"
do
    echo "$index ${str[index]}"
done


# start the fsleyes command with basic options
cmd="${FSLDIR}/bin/fsleyes $brain -dr 0 `fslstats $brain -r | awk '{print $2}'`"
opts="-dr $thr $uthr"

# Useful bits
preT="tracts"
postT="/densityNorm.nii.gz"

# Now loop and check for left/right tracts to colour the same
# checks for _l and matches any _r
# if no _?, then just move on to next line
i=0
for tract in "${str[@]}";do
  if [ $i -gt $((cL - 1)) ]; then i=0; fi # control colourmap loop
  # check tract exists
  if [ ! -f "${dir}/${preT}/${tract}${postT}" ]; then
    echo "Couldn't find ${tract} image."
    echo "Moving on to the next structure."
    echo "Check ${dir}/${preT}/${tract}${postT} and try again"
  else
    # if you find a left tract, then find the corresponding right tract
    # and colour in the same way
      if [[ $tract == *"_l"* ]];then
	  tt=`echo ${tract} | sed s/_l/_r/`
	  if [[ "${str[@]}" =~ "$tt" ]];then
              # append _l and _r to fsleyes command with viewing options
              cmd="$cmd ${dir}/${preT}/${tract}${postT} $opts -cm ${cmaps[i]} -n ${tract} ${dir}/${preT}/${tt}${postT} $opts -cm ${cmaps[i]} -n ${tt}"
	  else
              # else, just add the current tract
              cmd="$cmd ${dir}/${preT}/${tract}${postT} $opts -cm ${cmaps[i]} -n ${tract}"
	  fi
      elif [[ $tract == *"_r"* ]]  && [[ ! $cmd == *"$tract"* ]] && [[ ! ${str[@]} == *"`echo ${tract} | sed s/_r/_l/`"* ]];then
      # if tract name has _r and hasn't been found already
	  cmd="$cmd ${dir}/${preT}/${tract}${postT} $opts -cm ${cmaps[i]} -n ${tract}"
      elif [[ $tract != *"_r"*  && $tract != *"_l"* ]];then
      # if tract name doesn't have _l or _r
	  cmd="$cmd ${dir}/${preT}/${tract}${postT} $opts -cm ${cmaps[i]} -n ${tract}"
      fi
      ((i++))
  fi
done

echo $cmd
echo "Launching FSLeyes..."
eval $cmd &