#!/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 "3D photo reconstruction + photo imputation / super resolution"
  echo "Juan Eugenio Iglesias"
  echo "Friday March 28th 2025"
  echo ""
  echo "Run mri_3d_photo_recon with the following options:"
  echo ""
  echo "  -h, --help            show help message"
  echo ""
  echo ""
  echo "  --input_photo_dir [INPUT_PHOTO_DIR]"
  echo "                        Directory with input photos (required)"
  echo ""
  echo "  --input_segmentation_dir [INPUT_SEGMENTATION_DIR]"
  echo "                        Directory with input slab masks / segmentations (required)"
  echo ""
  echo "  --hemisphere HEMISPHERE"
  echo "                        hemisphere; must be left, right, or both (required)"
  echo ""
  echo "  --slice_thickness SLICE_THICKNESS"
  echo "                        Slice thickness in mm (required); will be finetuned if possible"
  echo ""
  echo "  --photo_resolution PHOTO_RESOLUTION"
  echo "                        Resolution of the photos in mm (required)"
  echo ""
  echo "  --output_directory OUTPUT_DIRECTORY"
  echo "                        Output directory with reconstructed photo volume and reference (required)"
  echo ""
  echo "  --ref_mri REF_MRI     "
  echo "                        Reference MRI scan (if available)"
  echo ""
  echo "  --ref_mri_synthseg REF_MRI_SYNTHSEG"
  echo "                        SynthSeg of reference MRI scan (will be computed if it does not exist)"
  echo ""
  echo "  --ref_mri_synthsr REF_MRI_SYNTHSR"
  echo "                        SynthSR of reference MRI; only needed if MRI is not 1mm T1 (will be computed if it does not exist)"
  echo ""
  echo "  --low_field_synthsr   "
  echo "                        Use low-field version of SynthSR (eg for Hyperfine scans)"
  echo ""
  echo "  --input_roi_dir INPUT_ROI_DIR"
  echo "                        Directory with ROi masks to deform to photo reconstruction (optional)"
  echo ""
  echo "  --ref_mesh REF_MESH   "
  echo "                        Reference surface mesh (if available)"
  echo ""
  echo "  --mesh_reorient_with_indices MESH_REORIENT_WITH_INDICES"
  echo "                        Indices to reorient mesh (see PhotoTools page in FreeSurfer wiki)"
  echo ""
  echo "  --fresh_tissue        "
  echo "                        Uses more lenient regularizers to accommodate fresh tissue"
  echo ""
  echo "  --photos_of_posterior_side"
  echo "                        Use when photos are taken of posterior side of slabs (default is anterior side)"
  echo ""
  echo "  --order_posterior_to_anterior"
  echo "                        Use when photos are ordered from posterior to anterior (default is anterior to posterior)"
  echo ""
  echo "  --initial_stretch_factor_lr_photos FACTOR"
  echo "                        Initialize stretch of photos in left-right direction by this factor (useful with squashed brains)"
  echo ""
  echo "  --no_z_stretch     "
  echo "                        Forces the slab thickness to stay at value given by --slice_thickness"
  echo "  --stretch_factor_lr_mesh FACTOR"
  echo "                        Stretch mesh in left-right direction by this factor (useful with squashed brains)"
  echo ""
  echo "  --weights WEIGHTS     "
  echo "                        CSV file with slab weights (follows order of photos). Format is just: weight1,weight2,...,weightN"
  echo ""
  echo "  --thickness_cap MAX"
  echo "                        when using weights, limit the maximum estimated thickness to this value (in mm). "
  echo "                        This is useful when the geometric estimation of thicknesses is wrong."
  echo "                        The code (without this option) tries to detect these situations automatically and tells you when to enable this option"
  echo ""
  echo "  --equalize_images     "
  echo "                        Use to equalize images (useful if they have low contrast, but can hurt the machine learning imputation)"
  echo ""
  echo "  --skip_bfgs           "
  echo "                        Use to skip BFGS finetuning (i.e., do only Adam)"
  echo ""
  echo "  --threads THREADS     "
  echo "                        Number of cores to be used. Default is 1. You can use -1 to use all available cores"
  echo ""
  echo "  --gpu GPU             "
  echo "                        Index of GPU to use (default is None, i.e., CPU mode)"
  echo ""
  echo "  --deform_recon_dir DEFORM_RECON_DIR"
  echo "                        Directory with FS dir of reference, to deform surfaces etc (expects to find deform_recon_dir/surf) (optional)"
  echo ""
  echo "  --cp_spacing_2d CP_SPACING_2D"
  echo "                        (Advanced) Control point spacing for 2D deformation"
  echo ""
  echo "  --cp_spacing_3d CP_SPACING_3D"
  echo "                        (Advanced) Control point spacing for 3D deformation"
  echo ""
  echo "  --k_lncc_mri K_LNCC_MRI"
  echo "                        (Advanced) Weight of LNCC between reference MRI and reconstruction"
  echo ""
  echo "  --k_dice_mri K_DICE_MRI"
  echo "                        (Advanced) Weight of Dice between masks of reference and reconstruction"
  echo ""
  echo "  --k_dif_slice_loss K_DIF_SLICE_LOSS"
  echo "                        (Advanced) Weight of SSD between consecutive slices of reconstruction"
  echo ""
  echo "  --k_mesh_loss K_MESH_LOSS"
  echo "                        (Advanced) Weight of absolute distance to edge of masks from mesh vertices"
  echo ""
  echo "  --k_regularizer K_REGULARIZER"
  echo "                        (Advanced) Weight of regularizer of log_det(affine matrices)"
  echo ""
  echo "  --k_regularizer_nonlin K_REGULARIZER_NONLIN"
  echo "                        (Advanced) Weight of regularizer of 2D nonlinear deformation of photos"
  echo ""
  echo "  --k_regularizer_nonlin3d K_REGULARIZER_NONLIN3D"
  echo "                        (Advanced) Weight of regularizer of 3D nonlinear deformatin of reference"
  echo ""
  echo "  --k_regularizer_sz K_REGULARIZER_SZ"
  echo "                        (Advanced) Weight of regularizer of stretch in AP direction"
  echo ""               
exit 0
fi
 
 
# Find path to shell script
SCRIPTPATH="$FREESURFER_HOME_FSPYTHON/python/packages/3d_photo_recon/scripts/recon_case.py"

# Try to find atlas data
MODELPATH="$FREESURFER_HOME/models/photo_imputation_unet.pth"

if [ ! -f "$MODELPATH" ];
then
  echo " "
  echo "   Model file not found. Please download atlas from: "
  echo "       https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/dissection_photo_model/photo_imputation_unet.pth"
  echo "   and copy it to:  "
  echo "      $MODELPATH "
  echo "   You only need to do this once. You can use the following commands: "
  echo "      1: cd $FREESURFER_HOME/models"
  echo "      2a (in Linux): wget https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/dissection_photo_model/photo_imputation_unet.pth "
  echo "      2b (in MAC): curl -o atlas.zip https://ftp.nmr.mgh.harvard.edu/pub/dist/lcnpublic/dist/dissection_photo_model/photo_imputation_unet.pth "
  echo " "
  exit 1
fi

fspython $SCRIPTPATH $@

echo "  "
echo "*****************"
echo "* All done!!!!! *"
echo "*****************"
echo ""
echo "If you use this tool in a publication, please cite:"
echo ""
echo ""
echo "Machine learning of dissection photographs and surface scanning for quantitative 3D neuropathology"
echo "H Gazula, H Tregidgo, B Billot, Y Balbastre, J Williams Ramirez, R Herisse, LJ Deden-Binder,"
echo "A Casamitjana, E Melief, CS Latimer, MD Kilgore, M Montine, E Robinson, E Blackburn,"
echo "MS Marshall, TR Connors, DH Oakley, MP Frosch, SI Young, K Van Leemput, AV Dalca,"
echo "B Fischl, CL MacDonald, CD Keene, BT Hyman, JE Iglesias"
echo "Elife, 12, RP91398 (2024)"
echo ""
echo ""