#!/usr/bin/env bash
# v 1.0
DEBUG=1

# Initialize variables with default values
sub_path=""
out_dir=""
container_path=""
save="no"
raw="yes"
model="vanilla"

# Parse command-line arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -input) sub_path="$2"; shift ;;
        -output) out_dir="$2"; shift ;;
        -RAW) raw="$2"; shift ;;
        -SAVE_FILE) save="$2"; shift ;;
        -CONTAINER_PATH) container_path="$2"; shift ;;
        -MODEL_NAME) model="$2"; shift ;;
        -help)
            echo "Usage: $0 -input <input_dir> -output <output_dir> -RAW <yes/no> -SAVE_FILE <yes/no> -CONTAINER_PATH <path> -MODEL_NAME <name>"
            exit 0
            ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done


# check that this is a file
if [ -e "$sub_path" ] || [ -L "$sub_path" ]; then
    echo "Input file exists or is a symbolic link."
else
    echo "Input file does not exist, exiting."
    exit 1
fi

# sub_path=/data/project/brainage_underrepresented/data/ds004107/sub-mind010/ses-01/anat/sub-mind010_ses-01_T1w.nii.gz
# dir_sub
dir_path=$(dirname "${sub_path}")
# sub_file=sub-mind010_ses-01_T1w.nii.gz (or .nii, or .mgz)
sub_file=$(basename "${sub_path}")
# sub_name=sub-mind010_ses-01_T1w
sub_name="${sub_file%%.*}"

# Debugging or verbose mode can be added here if needed.
if [ $DEBUG -eq 1 ]; then
    echo "sub_file = $sub_path"
    echo "dir_path = $dir_path"
    echo "sub_file = $sub_file"
    echo "sub_name = $sub_name"
    echo "out_dir = $out_dir"
    echo "save = $save"
    echo "raw = $raw"
    echo "container_path = $container_path"
    echo "model = $model"
fi

# create result folder and move the results in there
if [ -e "$out_dir/${sub_name}_brainageR.csv" ]; then
    echo "output file exists, exiting"
    exit 
else   
    mkdir -p $out_dir
fi

cd "${dir_path}"

datalad_dir_path=$(git -C ${dir_path} rev-parse --show-toplevel)

if [ -d ${datalad_dir_path} ]; then
    datalad_work_dir=$(mktemp -d)
    cd ${datalad_work_dir}
    datalad clone ${datalad_dir_path} .

    # Trying to get sub-mind010/ses-01/anat/sub-mind010_ses-01_T1w.nii.gz
    complement="${sub_path#"$datalad_dir_path"}"
    complement_dir_name=$(dirname "$complement" | sed 's|^/||')
    cd $complement_dir_name 
    sub_path=$(pwd)/$sub_file
    datalad get "${sub_file}"
    # Check the exit status of the previous command
    if [ $? -ne 0 ]; then
        echo "datalad cloning failed, exiting." >&2
        exit 1
    fi
fi


# define working dir, create a 'data' folder in it and copy the subject in.
work_dir=$(mktemp -d)
mkdir -p ${work_dir}/data/anat

cp "${sub_path}" "${work_dir}/data/anat"

if [ $? -ne 0 ]; then
    echo "datalad cloning failed, exiting." >&2
    exit 1
fi

if [ -d ${datalad_dir_path} ]; then
    # datalad drop "${sub_path}"
    chmod -R a+w $datalad_work_dir
    rm -fr $datalad_work_dir
fi

cd ${work_dir}

# unzip as SPM/CAT work with .nii files. Remove the nii.gz file and rename sub
# gunzip only if gz file and then delete .gz file
if [[ "$sub_file" == *.gz ]]; then
    echo "The file is a .gz file."
    gunzip data/anat/${sub_file}
    rm -f data/anat/${sub_file}
    sub_file=${sub_file%".gz"}
else
    echo "The file is not a .gz file."
fi

# make sure permissions are ok
chmod +rwx ${work_dir}/data/anat/$sub_file

# run singularity with brainageR code
# ----------------------------------------------------------
echo "Running the container"
#singularity exec /data/project/BrainAge4AD/tools/brainageR_dockerfile/brainageR2.1.sif brainageR -f ${work_dir}/data/${sub_file} -o ${sub_name}_brain_predicted.age.csv
singularity exec ${container_path}/brainageR2.1.sif brainageR -f ${work_dir}/data/anat/${sub_file} -o ${sub_name}_brain_predicted.age.csv

# Check the exit status of the previous command
if [ $? -ne 0 ]; then
    echo "Singularity container failed, exiting." >&2
    exit 1
fi
# ----------------------------------------------------------

# Moving the prediction into the output dir
mv $work_dir/${sub_name}_brain_predicted.age.csv $out_dir/${sub_name}_brainageR.csv

# Check if the last command was successful
if [ $? -eq 0 ]; then
    echo "Prediction file ${sub_name} saved in $out_dir"
else
    echo "There was a problem"
fi

# If user wants to save intermidiate files, moving everything into output dir
if [ $save == "yes" ]; then
    mkdir -p $out_dir/${sub_name}_brainageR
    mv $work_dir/* $out_dir/${sub_name}_brainageR/
    echo "All files of ${sub_name} saved in $out_dir/${sub_name}_brainageR"

    # Check if the last command was successful
    if [ $? -eq 0 ]; then
        echo "All files of ${sub_name} saved in $out_dir/${sub_name}_brainageR"
    else
        echo "There was a problem"
    fi
fi

echo "Cleaning tmp files"
rm -rf $work_dir

exit