#!/usr/bin/env bash

# Exit at any error
set -e

# Make sure FreeSurfer is sourced
[ ! -e "$FREESURFER_HOME" ] && echo "error: freesurfer has not been properly sourced" && exit 1

# If requesting help
if  [ $# == 1 ] && [  $1 == "--help" ]
then
  echo " "
  echo "Bayesian segmentation with histological whole brain atlas (fast version)."
  echo " "
  echo "Next-Generation histological atlas for high-resolution segmentation of human brain MRI"
  echo "Casamitjana et al. (in preparation)"
  echo " "
  echo "Usage:"
  echo "   mri_histo_atlas_segment_fast INPUT_SCAN OUTPUT_DIRECTORY GPU THREADS [BF_MODE]"
  echo " "
  echo "INPUT SCAN: scan to process, in mgz or nii(.gz) format"
  echo "OUTPUT_DIRECTORY: directory with segmentations, volume files, etc"
  echo "GPU: set to 1 to use the GPU; requires a fat (~48GB) GPU!"
  echo "THREADS: number of CPU threads to use (use -1 for all available threads)"
  echo "BF_MODE (optional): bias field model: dct (default), polynomial, or hybrid"
  echo " "
  exit 0
fi

# If number of arguments is incorrect
if [ $# -lt 4 ] || [ $# -gt 5 ]
then
  echo " "
  echo "Incorrect number of arguments."
  echo "Usage: "
  echo " "
  echo "   mri_histo_atlas_segment_fast INPUT_SCAN OUTPUT_DIRECTORY GPU THREADS [BF_MODE]"
  echo " "
  echo "Or, for help"
  echo " "
  echo "   mri_histo_atlas_segment_fast --help"
  echo " "
  exit 1
fi

# Parse arguments
INPUT=$1
OUTPUT_DIR=$2
CPU_FLAG=' '
if [ $3 -eq 0 ]
then
  CPU_FLAG='--cpu'
fi
THREADS=$4
BF_MODE='dct'
if [ $# -gt 4 ]
then
  BF_MODE=$5
fi

# Find path to shell script
BASEPATH="$FREESURFER_HOME/python/packages/ERC_bayesian_segmentation/"

# Try to find atlas data
ATLAS_DIR="$BASEPATH/atlas_simplified"

if [ ! -f "$ATLAS_DIR/size.npy" ];
then
  echo " "
  echo "   Atlas files not found. Please download atlas from: "
  echo "      https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/Histo_Atlas_Iglesias_2023/atlas_simplified.zip "
  echo "   and uncompress it into:  "
  echo "      $BASEPATH/ "
  echo "   You only need to do this once. You can use the following three commands (may require root access): "
  echo "      1: cd $BASEPATH"
  echo "      2a (in Linux): wget https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/Histo_Atlas_Iglesias_2023/atlas_simplified.zip "
  echo "      2b (in MAC): curl -o atlas.zip https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/Histo_Atlas_Iglesias_2023/atlas_simplified.zip "
  echo "      3. unzip atlas_simplified.zip"
  echo " "
  echo "   After correct extraction, the directory: "
  echo "      $ATLAS_DIR "
  echo "   should contain files: size.npy, label_001.npz, label_002.npz, ..."
  echo " "
  exit 1
fi

# Create command line arguments and run!
SYNTHSEG="$OUTPUT_DIR/SynthSeg.mgz"
FIELD="$OUTPUT_DIR/atlas_registration.nii.gz"
RESOLUTION="0.4"
SKIP="3"
LAMBDA="0.05"

cmd="fspython $BASEPATH/scripts/segment_fast.py --i $INPUT --i_seg $SYNTHSEG  --i_field $FIELD  --atlas_dir $ATLAS_DIR --o $OUTPUT_DIR  --bf_mode $BF_MODE  --threads $THREADS --resolution $RESOLUTION --skip $SKIP --synthmorph_reg $LAMBDA  --write_bias_corrected  $CPU_FLAG"

echo "Running command:"
echo $cmd
echo "  "
$cmd