#!/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 with FireANTs registration)."
  echo " "
  echo "- A probabilistic histological atlas of the human brain for MRI segmentation"
  echo "    Casamitjana et al. (in press)"
  echo "- Fast segmentation with the NextBrain histological atlas"
  echo "    Puonti et al. (under revision)"
  echo " "
  echo "Usage:"
  echo "   mri_histo_atlas_segment_fireants INPUT_SCAN OUTPUT_DIRECTORY GPU THREADS MODE SIDE"
  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 "MODE: invivo , exvivo, cerebrum, or hemi"
  echo "SIDE: left or right"
  echo " "
  exit 0
fi

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

# Parse arguments
INPUT=$1
OUTPUT_DIR=$2
DEVICE='cuda'
if [ $3 -eq 0 ]
then
  DEVICE='cpu'
fi
THREADS=$4
MODE=$5
SIDE=$6

# 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_simplified.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

# Try to find SuperSynth model
MODEL_FILE="$FREESURFER_HOME/models/SuperSynth_August_2025.pth"

if [ ! -f "$MODEL_FILE"  ];
then
    echo " "
    echo "   Machine learning model file not found. Please download from from: "
    echo "     https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/SuperSynth_Iglesias_2025/SuperSynth_August_2025.pth "
    echo "   and uncompress it into:  "
    echo "      $FREESURFER_HOME/models/ "
    echo "   You only need to do this once. You can use the following commands (may require root access): "
    echo "      1: cd $FREESURFER_HOME/models/"
    echo "      2a (in Linux): wget https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/SuperSynth_Iglesias_2025/SuperSynth_August_2025.pth "
    echo "      2b (in MAC): curl -o SuperSynth_August_2025.pth https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/SuperSynth_Iglesias_2025/SuperSynth_August_2025.pth "
    echo " "
    echo "   After correct download and instillation, the directory: "
    echo "      $FREESURFER_HOME/models/ "
    echo "   should now contain an additional file:SuperSynth_August_2025.pth"
    echo " "
    exit 1
fi



# Create command line arguments and run!
cmd="fspython $BASEPATH/scripts/segment_fireants.py --i $INPUT --model_file $MODEL_FILE   --atlas_dir $ATLAS_DIR --o $OUTPUT_DIR  --device $DEVICE --bf_mode dct --threads $THREADS --mode $MODE --side $SIDE  "
echo "Running command:"
echo $cmd
echo "  "
$cmd

echo "  "
echo "*****************"
echo "* All done!!!!! *"
echo "*****************"
echo "  "
echo "Please cite: "
echo "- A probabilistic histological atlas of the human brain for MRI segmentation"
echo "    Casamitjana et al. (in press)"
echo "- Fast segmentation with the NextBrain histological atlas"
echo "    Puonti et al. (under revision)"
echo "  "