#!/usr/bin/env bash
#
# PURPOSE: This script performs preprocessing steps on T1-weighted MRI images
#          to create segmentations and extract cortical surface. 
#
# USAGE: T1Prep [options] file1.nii file2.nii ...
#
# INPUT: T1-weighted MRI images in NIfTI format (.nii or .nii.gz).
#
# OUTPUT: Processed images and segmentation results in the specified output directory.
#
# FUNCTIONS: 
# - main: The main function that executes the preprocessing steps.
# - parse_args: Parses the command line arguments.
# - surface_estimation: Create cortical surfaces
# - process: Performs the preprocessing steps on each input file.
#
# ______________________________________________________________________
#
# Christian Gaser
# Structural Brain Mapping Group (https://neuro-jena.github.io)
# Departments of Neurology and Psychiatry
# Jena University Hospital
# ______________________________________________________________________

# ----------------------------------------------------------------------
# Global parameters
# ----------------------------------------------------------------------

# Defaults
# Resolve this script's directory robustly (works even if invoked via symlink)
script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

seed=0
count=0
end_count=1

# Load additional defaults.  Source-tree layout keeps the file at the repo root
# (scripts/../); an installed venv ships it next to the scripts in bin/ (via the
# script-files list in pyproject.toml), so fall back to the script's own dir.
if [ -f "${script_dir}/../T1Prep_defaults.txt" ]; then
  cfg="${script_dir}/../T1Prep_defaults.txt"
else
  cfg="${script_dir}/T1Prep_defaults.txt"
fi

# Use some external functions
source "${script_dir}/T1Prep_utils.sh"
progress_bar="${script_dir}/progress_bar_multi.sh"


# ----------------------------------------------------------------------
# Run main
# ----------------------------------------------------------------------

main()
{
  check_python_cmd
  parse_args ${1+"$@"}
  check_python_cmd

  # Normalize and default numeric options to prevent integer comparison errors
  # multi: -2 (quiet), 0 (single), >0 (explicit parallel), -1 (auto)
  if ! [[ "${multi}" =~ ^-?[0-9]+$ ]]; then multi=-1; fi
  if ! [[ "${estimate_seg}" =~ ^[0-9]+$ ]]; then estimate_seg=1; fi
  if ! [[ "${estimate_surf}" =~ ^[0-9]+$ ]]; then estimate_surf=1; fi
  if ! [[ "${estimate_spherereg}" =~ ^[0-9]+$ ]]; then estimate_spherereg=1; fi
  if ! [[ "${save_pial_white}" =~ ^[0-9]+$ ]]; then save_pial_white=1; fi
  if ! [[ "${thickness_method}" =~ ^-?[0-9]+$ ]]; then thickness_method=3; fi
  if ! [[ "${debug}" =~ ^[0-9]+$ ]]; then debug=0; fi
  if ! [[ "${retry}" =~ ^[0-9]+$ ]]; then retry=1; fi
  if ! [[ "${use_bids_naming}" =~ ^[0-9]+$ ]]; then use_bids_naming=0; fi
  if ! [[ "${min_memory}" =~ ^[0-9]+$ ]]; then min_memory=24; fi

  if [[ "$multi" -ne -2 ]]; then
    logo
    check_python_module venv
    check_python_module pip
    check_python_libraries
    if [ "$re_install" -eq 1 ]; then
      exit 0
    fi
  fi
  check_files "${#ARRAY[@]}"
  
  process "$@"

  exit 0
}


# ----------------------------------------------------------------------
# Check arguments and files
# ----------------------------------------------------------------------

parse_args()
{
  local optname optarg
  
  # Check default-file first and load it before any other options
    for ((i=1;i<=$#;i++)); do
      case "${!i}" in
        --default*)
        next_index=$((i+1))
        cfg="${!next_index}"
        i=$((i+1))
        ;;
      esac
    done
  source "${cfg}"

  # Defaults for new skull-strip flags (may be absent in defaults file)
  skullstrip_only=${skullstrip_only:-0}
  skip_skullstrip=${skip_skullstrip:-0}

  if [ $# -lt 1 ]; then
    logo
    help
    exit 1
  fi

  i=0
  while [ $# -gt 0 ]; do
    optname="${1%%=*}"
    optarg="${2:-}"

    case "$1" in
      --default* | --defaults)
        exit_if_empty "$optname" "$optarg"
        cfg=$optarg
        shift
        ;;
      --install | --re-install)
        re_install=1
        ;;
      --python)
        exit_if_empty "$optname" "$optarg"
        python=$optarg
        shift
        ;;
      --out*)
        exit_if_empty "$optname" "$optarg"
        outdir=$optarg
        shift
        ;;
      --pre-fwhm)
        exit_if_empty "$optname" "$optarg"
        pre_fwhm=$optarg
        shift
        ;;
      --downsample)
        exit_if_empty "$optname" "$optarg"
        downsample=$optarg
        shift
        ;;
      --median-filter)
        exit_if_empty "$optname" "$optarg"
        median_filter=$optarg
        shift
        ;;
      --no-vessel*)
        vessel=0
        ;;
      --no-bvc*)
        vessel=0
        ;;
      --no-overwrite*)
        exit_if_empty "$optname" "$optarg"
        no_overwrite=$optarg
        shift
        ;;
      --thickness-method)
        exit_if_empty "$optname" "$optarg"
        thickness_method=$optarg
        shift
        ;;
      --multi*)
        exit_if_empty "$optname" "$optarg"
        multi=$optarg
        shift
        ;;
      --seed)
        exit_if_empty "$optname" "$optarg"
        seed=$optarg
        shift
        ;;
      --atlas)
        exit_if_empty "$optname" "$optarg"
        atlas_vol=$optarg
        shift
        ;;
      --atlas-surf)
        exit_if_empty "$optname" "$optarg"
        atlas_surf=$optarg
        shift
        ;;
      --min-mem*)
        exit_if_empty "$optname" "$optarg"
        min_memory=$optarg
        shift
        ;;
      --initial-surf*)
        exit_if_empty "$optname" "$optarg"
        initial_surface=$optarg
        shift
        ;;
      --long-data)
        exit_if_empty "$optname" "$optarg"
        long_data=$optarg
        shift
        ;;
      --fmriprep)
        fmriprep=1
        use_bids_naming=1
        save_csf=1
        save_p=1
        save_mwp=0
        save_hemi=0
        estimate_spherereg=1
        atlas_vol=""
        nii_ext='nii.gz'
        ;;
      --gz)
        nii_ext='nii.gz'
        ;;
      --hemi*)
        save_hemi=1
        ;;
      --no-surf)
        estimate_surf=0
        ;;
      --skullstrip-only)
        skullstrip_only=1
        ;;
      --no-skullstrip | --skip-skullstrip)
        skip_skullstrip=1
        ;;
      --no-seg*)
        estimate_seg=0
        ;;
      --no-sphere*)
        estimate_spherereg=0
        ;;
      --no-pial*)
        save_pial_white=0
        ;;
      --lesion*)
        save_lesions=1
        ;;
      --no-mwp* | --no-warped-seg*)
        save_mwp=0
        ;;
      --no-atlas*)
        atlas_vol=""
        ;;
      --wp* | --warped-seg-nomod)
        save_wp=1
        ;;
      --rp* | --affine-seg)
        save_rp=1
        ;;
      --p* | --native-seg)
        save_p=1
        ;;
      --csf*)
        save_csf=1
        ;;
      --amap)
        use_amap=1
        ;;
      --bids)
        use_bids_naming=1
        nii_ext='nii.gz'
        ;;
      --no-correct-folding)
        correct_folding=0
        ;;
      --no-retry)
        retry=0
        ;;
      --no-track)
        track=0
        ;;
      --debug)
        debug=1
        ;;
      --verbose)
        verbose=1
        ;;
      --fast)
        save_pial_white=0
        estimate_spherereg=0
        save_mwp=0
        save_hemi=0
        atlas_vol=""
        atlas_surf=""
        ;;
      -h | --help | --h | -v | --version | -V)
        help
        exit 1
        ;;
      -*)
        echo "`basename $0`: ERROR: Unrecognized option \"$1\"" >&2
        exit 1
        ;;
      *)
        if [ ! -f "$1" ]; then
          if [ ! -L "$1" ]; then
            echo "${RED}File $1 was not found and will be skipped.${NC}" >&2
          fi
        else
          FILE="$1"
          ARRAY[$i]="$FILE"
          ((i++))
        fi
        ;;
    esac
    shift
  done

  if [[ "$estimate_seg" -eq 0 && "$estimate_surf" -eq 0 ]]; then
    echo "${RED}ERROR: Options "--no-surf" and "--no-seg" cannot be used together.${NC}" >&2
    exit 1
  fi

  if [[ "${skullstrip_only}" -eq 1 && "${skip_skullstrip}" -eq 1 ]]; then
    echo "${RED}ERROR: Options --skullstrip-only and --no-skullstrip cannot be used together.${NC}" >&2
    exit 1
  fi

  if [[ "${skullstrip_only}" -eq 1 ]]; then
    # Skull stripping only: run volume pipeline step to produce skullstrip outputs and stop.
    estimate_seg=1
    estimate_surf=0
    estimate_spherereg=0
    save_hemi=0
  fi

  # We always need hemispheric maps for surface estimation
  if [ "$estimate_surf" -eq 1 ]; then
    save_hemi=1
  fi

}

# ----------------------------------------------------------------------
# Surface estimation
# ----------------------------------------------------------------------

surface_estimation() 
{  
  local bname=$1
  local side=$2
  local mri_dir=$3
  local surf_dir=$4
  local estimate_spherereg=$5
  local multi=$6
  local nii_ext=$7

  # Share the progress-bar count file between hemispheres
  local _count_file="${report_log%.log}.count"
  # File should already have been initialised by the outer loop before the
  # hemisphere processes were launched; create it here only as a safety fallback.
  [[ -f "${_count_file}" ]] || echo "${count}" > "${_count_file}"
  ${python} -m t1prep.surface_estimation \
      --bname "${bname}" --side "${side}" \
      --mri-dir "${mri_dir}" --surf-dir "${surf_dir}" \
      --estimate-spherereg "${estimate_spherereg}" \
      --thickness-method "${thickness_method}" \
      --save-pial-white "${save_pial_white}" \
      --pre-fwhm "${pre_fwhm}" \
      --median-filter "${median_filter}" \
      --downsample "${downsample}" \
      --vessel "${vessel}" \
      --correct-folding "${correct_folding}" \
      --debug "${debug}" \
      --multi "${multi}" \
      --nii-ext "${nii_ext}" \
      --names-tsv "${name_file}" \
      --bids-naming "${use_bids_naming}" \
      --report-log "${report_log}" \
      --surf-templates-dir "${surf_templates_dir}" \
      --atlas-templates-dir "${atlas_templates_dir}" \
      --atlas-surf "${atlas_surf}" \
      --initial-surface "${initial_surface:-}" \
      --fmriprep "${fmriprep:-0}" \
      --progress-bar "${progress_bar}" \
      --progress-count-file "${_count_file}" \
      --progress-end-count "${end_count}" \
      --progress-start-count "${count}"
  return $?
}

# ----------------------------------------------------------------------
# Process data
# ----------------------------------------------------------------------

process()
{
  errors=0
  SIZE_OF_ARRAY="${#ARRAY[@]}"
  if ! [[ "${multi}" =~ ^-?[0-9]+$ ]]; then multi=-1; fi
  if [[ "$SIZE_OF_ARRAY" -lt 2 && "$multi" -ne -2 ]]; then
    multi=0
  fi

  # Call T1Prep script recursively with different "--multi" arguments in case that
  # "--multi" is defined
  if [[ "$multi" -ne 0 && "$multi" -gt -2 ]]; then

    # Filter arguments before recursion
    filtered_args=($(filter_arguments "$@")) # Use an array to handle filtered arguments

    # Get number of processes (NUM_JOBS) for defined memory limit
    # Ensure min_memory sane before calculating
    if ! [[ "${min_memory}" =~ ^[0-9]+$ ]]; then min_memory=12; fi
    get_no_processes ${min_memory}
    
    # If multi is defined and lower than available processes w.r.t. memory limit,
    # then use this lower value
    if [[ "$multi" -gt 0 && "$NUM_JOBS" -gt "$multi" ]]; then
      arg_parallelize=( -p "${multi}" )
    else
      arg_parallelize=( -p "${NUM_JOBS}" )
    fi

    # Call parallelize with filtered arguments and use memory limit per process
    # and a small delay of 5s between the jobs to avoid simultaneous peak memory usage
    args=(--no-surf --multi -2)
    if [ "${estimate_surf}" -eq 1 ]; then
      args+=( --hemisphere)
    fi
    cmd=(
      "${script_dir}/parallelize"
      -l "/tmp"
      -d 5
      ${arg_parallelize[@]}
      -c "$(printf '%q ' "$0" "${filtered_args[@]}" "${args[@]}")"
      "${ARRAY[@]}"
    )

    # Execute the command
    if [ "${estimate_seg}" -eq 1 ]; then
      echo "${BOLD}Volume Segmentation${NC}"
      "${cmd[@]}"
      status=$?
      if [ "$status" -ne 0 ]; then
        echo "${RED}ERROR: Volume segmentation failed for one or more subjects (see log above).${NC}" >&2
      fi
    fi

    # Call parallelize with filtered arguments
    # Define number of processes, otherwise estimate it automatically 
    if ! [[ "${multi}" =~ ^-?[0-9]+$ ]]; then multi=-1; fi
    if [ "$multi" -gt 0 ]; then
      arg_parallelize=( -p "${multi}" )
    else
      arg_parallelize=( -p 0.5 )
    fi
    cmd=(
      "${script_dir}/parallelize"
      -l "/tmp"
      ${arg_parallelize[@]}
      -c "$(printf '%q ' "$0" "${filtered_args[@]}" --no-seg --multi -2)"
      "${ARRAY[@]}"
    )
    
    # Execute the command
    if [ "${estimate_surf}" -eq 1 ]; then
      echo "${BOLD}Surface Estimation${NC}"
      "${cmd[@]}"
      status=$?
    fi
    
    exit $status
  fi

  # Don't use colored text in parallelization (for clearer log-file)
  if [[ "$multi" -eq -2 ]]; then
    UNDERLINE=""
    BOLD=""
    NC=""
    CYAN=""
    PINK=""
    RED=""
    YELLOW=""
    BLUE=""
    WHITE=""
    GREEN=""
    BLACK=""
  fi
  
  i=0
  j=0
  while [ "$i" -lt "$SIZE_OF_ARRAY" ]; do

    # Check whether absolute or relative names were given
    if [ ! -f "${ARRAY[$i]}" ]; then
      if [ -f "${pwd}/${ARRAY[$i]}" ]; then
        FILE="${pwd}/${ARRAY[$i]}"
      fi
    else
      FILE="${ARRAY[$i]}"
    fi

  # Do not mangle spaces; always quote variables instead

    # Check whether processed files exist if no-overwrite flag is used
    if [ -n "${no_overwrite}" ]; then

      # Output folders are always derived from the original input file.
      get_output_folder "$FILE"
      
      if [ "${use_bids_naming}" -eq 1 ]; then  
        processed=$(ls "${outdir0}/${bname}"*"/${no_overwrite}"* 2>/dev/null)
      else
        processed=$(ls "${outdir0}/${no_overwrite}${bname}"* 2>/dev/null)
      fi
    fi
    
    if [ -z "${processed}" ]; then
      ARRAY2[$j]="$FILE"
      ((j++))
    else
      echo Skip processing of "${FILE}"
    fi
    ((i++))
  done
  
  
  i=0
  SIZE_OF_ARRAY="${#ARRAY2[@]}"
  
  # Exit if no files could be found for processing
  if [ "$SIZE_OF_ARRAY" -eq 0 ]; then
    exit 0
  fi

  # Set overall starting time
  start0=$(date +%s)

  # Source-tree mode: activate the project-managed venv.  Installed mode is
  # already inside a venv (pip dropped this script there), so do nothing.
  if [ "${T1PREP_INSTALLED:-0}" -ne 1 ]; then
    if [ ! -f "${T1prep_env}/bin/activate" ]; then
      echo "${RED}Virtual environment not found at ${T1prep_env}. Run the script once without --no-overwrite to set it up.${NC}" >&2
      exit 1
    fi
    source ${T1prep_env}/bin/activate
  fi

  while [ "$i" -lt "$SIZE_OF_ARRAY" ]; do
    
    # Set starting time
    start=$(date +%s)

    FILE="${ARRAY2[$i]}"

    # Output folders are always derived from the original input file.
    get_output_folder "$FILE" # sets outdir0/use_subfolder/bname

    # Check again whether processed files exist if no-overwrite flag is used
    if [ -n "${no_overwrite}" ]; then
      if [ "${use_bids_naming}" -eq 1 ]; then  
        processed=$(ls "${outdir0}/${no_overwrite}${bname}"* 2>/dev/null)
      else
        processed=$(ls "${outdir0}/${bname}"*"${no_overwrite}"* 2>/dev/null)
      fi

      # Check if $processed is empty
      if [ -n "$processed" ]; then
        echo Skip processing of "${FILE}"
        continue  # Skip to the next iteration of the loop
      fi
    fi
    
    # Get output folders for surfaces and volumes
    if [[ "${use_bids_naming}" -eq 1 || "${use_subfolder}" -eq 0 ]]; then  
      mri_dir=${outdir0}
      surf_dir=${outdir0}
      report_dir=${outdir0}
      label_dir=${outdir0}
    else
      mri_dir=${outdir0}/mri
      surf_dir=${outdir0}/surf
      report_dir=${outdir0}/report
      label_dir=${outdir0}/label
    fi
              
    # Create output folders    
    mkdir -p "${mri_dir}" "${report_dir}" "${label_dir}"
    if [ "${estimate_surf}" -eq 1 ]; then
      mkdir -p "${surf_dir}"
    fi
    
    # Get filename for log file
    code="Log_file"
    if [[ $use_bids_naming -eq 1 ]]; then
      name_columns=3  # 3rd column in Names.tsv
    else
      name_columns=2  # 2nd column in Names.tsv
    fi
    pattern=$(get_pattern "$code" "$name_columns")
    # Create dynamic variables for filenames
    value=$(substitute_pattern "$pattern" "" "" "" "")
    # Use eval to assign the result to a variable named after the code
    eval "${code}=\"\$value\""

    # Save all output in log-file
    report_log=${report_dir}/${Log_file}
    echo "T1Prep version ${T1PREP_VERSION}" > "${report_log}"
    date >> "${report_log}"
    echo "${os_type}" >> "${report_log}"
    cuda=$(${python} -c "import torch; print(torch.cuda.is_available())" 2>/dev/null)
    if [ "${cuda}" == "True" ]; then
      echo "Use Cuda in Python" >> "${report_log}"
    else
      mps=$(${python} -c "import torch; print(torch.backends.mps.is_available())" 2>/dev/null)
      if [ "${mps}" == "True" ]; then
        echo "Use MPS in Python" >> "${report_log}"
      else
        echo "Use CPU in Python" >> "${report_log}"
      fi
    fi
    
    # Print progress and filename
    j=$((i + 1))
    if [ "${multi}" -ne -2 ]; then
      echo "${BOLD}----------------------------------------------------------${NC}"
      if [ "${SIZE_OF_ARRAY}" -gt 1 ]; then
        echo "${GREEN}${j}/${SIZE_OF_ARRAY} ${BOLD}Processing ${FILE}${NC}"
      else
        echo "${BOLD}Processing ${FILE}${NC}"
      fi
    fi

    if [ -n "${long_data}" ]; then
      input="$long_data"
    else
      input="$FILE"
    fi

    # Reset progress counters for this file
    count=0
    end_count=1

    # Change number of commands if spherical registration is disabled
    if [ "${estimate_surf}" -eq 1 ]; then
      if [ "${estimate_spherereg}" -eq 1 ]; then
        ((end_count+=6))
      else
        ((end_count+=4))
      fi
      if [[ "${save_pial_white}" -eq 1 || "${thickness_method}" -eq 2 ]]; then
        ((end_count++))
      fi
      if [ "${thickness_method}" -eq 2 ]; then
        ((end_count++))
      fi
    fi

    # Number of end counts for progress_bar for volume segmentation
    if [ "${estimate_seg}" -eq 1 ]; then
      ((end_count+=4))
      ((count+=4))
      # Warping + Atlas creation + Resampling are triggered whenever
      # --surf is passed to segment.py (i.e. save_hemilabel is true).
      if [ "${estimate_surf}" -eq 1 ] || [ "${save_hemi}" -eq 1 ] || \
         [ "${save_mwp}" -eq 1 ] || [ "${save_wp}" -eq 1 ] || \
         [ -n "${atlas_vol}" ]; then
        ((end_count+=3))
        ((count+=3))
      fi
    fi
    
    # 1. Call deepmriprep segmentation 
    # ----------------------------------------------------------------------
    # Check for outputs from previous step
    if [ -f "${input}" ]; then

      # Use the module form so it works in both source-tree and installed
      # mode (segment.py uses package-relative imports, which require running
      # under the t1prep package — `${python} ${src_dir}/segment.py` would
      # fail with "attempted relative import with no known parent package").
      if [ "${T1PREP_INSTALLED:-0}" -ne 1 ]; then
        export PYTHONPATH="${root_dir}/src:${PYTHONPATH}"
      fi
      cmd="${python} -m t1prep.segment"
      
      # Append options conditionally
      [ "${use_amap}" -eq 1 ] && cmd+=" --amap"
      [ "${save_mwp}" -eq 1 ] && cmd+=" --mwp"
      [ "${save_wp}" -eq 1 ] && cmd+=" --wp"
      [ "${save_rp}" -eq 1 ] && cmd+=" --rp"
      [ "${save_p}" -eq 1 ] && cmd+=" --p"
      [ "${nii_ext}" == "nii.gz" ] && cmd+=" --gz"
      [ "${fmriprep}" -eq 1 ] && cmd+=" --save-fmriprep"
      [ "${save_lesions}" -eq 1 ] && cmd+=" --lesions"
      [ "${use_bids_naming}" -eq 1 ] && cmd+=" --bids"
      [ "${multi}" -ne -2 ] && cmd+=" --verbose"
      [ "${debug}" -eq 1 ] && cmd+=" --debug"
      [ "${estimate_surf}" -eq 1 ] || [ "${save_hemi}" -eq 1 ] || \
      [ "${save_mwp}" -eq 1 ] || [ "${save_wp}" -eq 1 ] || \
      [ -n "${atlas_vol}" ] && cmd+=" --surf"
      [ "${save_csf}" -eq 1 ] && cmd+=" --csf"
      [ "${skullstrip_only}" -eq 1 ] && cmd+=" --skullstrip-only"
      [ "${skip_skullstrip}" -eq 1 ] && cmd+=" --skip-skullstrip"

      cmd+=" --seed ${seed} --vessel \"${vessel}\" --label-dir \"${label_dir}\""
      cmd+=" --input \"${input}\" --mri-dir \"${mri_dir}\""
      cmd+=" --report-dir \"${report_dir}\" --atlas \"${atlas_vol}\" --count ${end_count}"
      
      # Execute the command and capture exit code
      if [ "${estimate_seg}" -eq 1 ]; then
        echo "${cmd}" >> "${report_log}" 
        if [[ "$multi" -ne -2 ]]; then
          eval "${cmd}"
        else
          eval "${cmd}" 2>&1 | tee -a "${report_log}"
        fi
        seg_status=$?

        # Retry once on failure if retry is enabled
        if [[ "$seg_status" -ne 0 && "${retry}" -eq 1 ]]; then
          echo "Retrying segmentation for ${input}..." >> "${report_log}"
          if [[ "$multi" -ne -2 ]]; then
            echo "${YELLOW}Segmentation failed — retrying ${input}...${NC}"
            eval "${cmd}"
          else
            echo "--- Retrying: ${input} ---" 2>&1 | tee -a "${report_log}"
            eval "${cmd}" 2>&1 | tee -a "${report_log}"
          fi
          seg_status=$?
        fi
      else
        seg_status=0
      fi

      if [ "$seg_status" -ne 0 ]; then
        if [[ "$multi" -ne -2 ]]; then
          # Use 1/1 so the single job bar renders without relying on end_count
          ${progress_bar} 1 "" 1 1 "Segmentation failed"
        fi
        ((i++))
        ((errors++))
        echo "Error in ${input}"
        continue
      fi
  
    else
      echo "${RED}ERROR: ${input} could not be found.${NC}" >&2
      ((i++))
      ((errors++))
      continue
    fi
    
    # Optionally extract surface
    if [ "${estimate_surf}" -eq 1 ]; then

      # Set starting time
      start_surf=$(date +%s)

      # Reset count file before launching hemisphere processes so stale values
      # from previous runs do not corrupt the progress bar percentage.
      _surf_count_file="${report_log%.log}.count"
      echo "${count}" > "${_surf_count_file}"
        
      # 2. Estimate thickness and percentage position maps for each hemisphere
      # and extract cortical surface and call it as background process to
      # allow parallelization
      # ----------------------------------------------------------------------
      # Check for outputs from previous step
      pids=()
      for side in left right; do
        surface_estimation $bname $side "${mri_dir}" "${surf_dir}" $estimate_spherereg $multi "${nii_ext}" &
        pids+=("$!")
      done

      surf_status=0
      for pid in "${pids[@]}"; do
        wait "$pid" || surf_status=1
      done

      # Retry once on failure if retry is enabled
      if [[ "$surf_status" -ne 0 && "${retry}" -eq 1 ]]; then
        echo "Retrying surface estimation for ${bname}..." >> "${report_log}"
        if [[ "$multi" -ne -2 ]]; then
          echo "${YELLOW}Surface estimation failed — retrying ${bname}...${NC}"
        else
          echo "--- Retrying surface estimation: ${bname} ---" 2>&1 | tee -a "${report_log}"
        fi
        pids=()
        # Reset count file before retry
        echo "${count}" > "${report_log%.log}.count"
        for side in left right; do
          surface_estimation $bname $side "${mri_dir}" "${surf_dir}" $estimate_spherereg $multi "${nii_ext}" &
          pids+=("$!")
        done
        surf_status=0
        for pid in "${pids[@]}"; do
          wait "$pid" || surf_status=1
        done
      fi

      if [ "$surf_status" -ne 0 ]; then
        if [[ "$multi" -ne -2 ]]; then
          # Use 1/1 so the single job bar renders without relying on end_count
          ${progress_bar} 1 "" 1 1 "Surface estimation failed"
        fi
        ((errors++))
        echo "Error in ${bname}"
      fi

      end_surf=$(date +%s)
      runtime=$((end_surf - start_surf))
      echo "" >> "${report_log}"
      echo "Execution time surface pipeline: ${runtime}s" >> "${report_log}"

    fi # estimate_surf

    if [ "$multi" -ne -2 ]; then
      ((count++))
      ${progress_bar} 1 "" $count $end_count "                     "
    fi

    # Print execution time per data set
    end=$(date +%s)
    runtime=$((end - start))
    hours=$(($runtime / 3600))
    min=$((($runtime / 60) % 60))
    s=$(($runtime % 60))
    overall="Finished after "
    if [ $hours -gt 0 ]; then overall+="${hours}hrs "; fi
    if [ $min -gt 0 ]; then overall+="${min}min "; fi
    overall+="${s}s"

    if [ "${estimate_surf}" -eq 1 ]; then
      echo "${overall}" >> "${report_log}"
    fi
    
    if [[ "$multi" -ne -2 ]]; then
      echo "${GREEN}----------------------------------------------------------                                   ${NC}"
      echo "${GREEN}${overall}${NC}"
    fi
      
    ((i++))
  done
  
  # Collect telemetry data
  if [ "$track" -eq 1 ]; then
    send_sentry $((SIZE_OF_ARRAY - errors)) $errors
  fi
  
  # Print overall execution time for more than one data set
  if [ "$SIZE_OF_ARRAY" -gt 1 ]; then
    end0=$(date +%s)
    runtime=$((end0-start0))
    runtime="T1Prep finished after $(($runtime / 3600))hrs $((($runtime / 60) % 60))min $(($runtime % 60))s"
    echo "${GREEN}${runtime}${NC}"  
  fi
  
  exit $errors

}

# ----------------------------------------------------------------------
# Help
# ----------------------------------------------------------------------

help() {
cat << EOM
${BOLD}${BLUE}T1Prep Computational Anatomy Pipeline (aka PyCAT)
-----------------------------------------------------------------------------------${NC}

${BOLD}USAGE:${NC}
  ${GREEN}T1Prep [options] <filename(s)>${NC}

${BOLD}Tool Overview:${NC}
T1Prep is a pipeline that preprocesses T1-weighted MRI data and supports segmentation 
and cortical surface reconstruction. It provides a complete set of tools for efficiently 
processing structural MRI scans.

T1Prep partially integrates DeepMriPrep, which uses deep learning (DL) techniques to 
mimic functionality of CAT12 for processing structural MRIs. DeepMriPrep is used for 
skull-stripping, segmentation and nonlinear spatial registration.

An alternative approach uses DeepMriPrep for bias field correction, lesion detection, 
and also serves as an initial estimate for the subsequent AMAP segmentation from CAT12. 

Cortical surface reconstruction and thickness estimation are performed using Cortex 
Analysis Tools for Surface, a core component of the CAT12 toolbox.

It is designed for both single-subject and batch processing, with optional parallelization
and flexible output naming conventions. The naming patterns are compatible with both 
CAT12 folder structures and the BIDS derivatives standard.

${BOLD}Output Folder Behaviour:${NC}
The location and structure of the output folders depend on:

1. Whether the input data follow the BIDS structure
  If the upper-level folder of the NIfTI file is named 'anat', the script assumes 
  a BIDS dataset. Output will then be written into a BIDS-compatible derivatives folder:

  '<dataset-root>/derivatives/T1Prep-v<version>/<sub-XXX>/<ses-YYY>/anat/'
  where sub-XXX and ses-YYY are automatically detected from the path.

2. Non-BIDS datasets

  If the upper-level folder is not anat, results are written into subfolders similar 
  to CAT12 output inside the specified output directory (or the current working 
  directory if none is specified).

3. Folder naming conventions

  By default, CAT12-style names are used.
  If --bids is set, BIDS-compatible names are used instead.
  The mapping between internal processing outputs and their filenames for each 
  convention is stored in Names.tsv.

4. Custom output folder

  You can specify --out-dir <DIR> to override the default location.
  If --bids is set, the BIDS substructure will still be appended inside <DIR>.
  
${BOLD}Naming conventions:${NC}
CAT12 style (default): Uses legacy names for files and subfolders (mri, surf, etc.).

BIDS style (--bids): Uses BIDS derivatives naming rules for filenames and metadata.

The mapping between internal processing steps and filenames for both conventions 
is defined in Names.tsv. This file can be edited to customize output naming.
  
${BOLD}Options:${NC}
${BOLD}${GREEN}General Options:${NC}
  --defaults <FILE>           
      Specify an alternative defaults file to override built-in settings.

  --install                   
      Install all required Python libraries.

  --re-install                
      Remove the existing installation and re-install all required Python libraries.

  --python <FILE>             
      Path to the Python interpreter to use (default: $python).

  --multi <NUMBER>            
      Set the maximum number of parallel jobs. Use '-1' to automatically 
      detect and use all available CPU cores (default: $multi).  
      If you specify a value here and it is lower than the number of jobs 
      calculated based on --min-memory, your specified value will be used.

  --min-memory <NUMBER>
      Set the minimum amount of memory (in GB) to reserve for each parallel
      job (default: $min_memory GB). This value is used to estimate the
      maximum number of jobs that can run in parallel without exceeding
      available system memory. Increase this value if your system runs
      low on memory or becomes unresponsive during parallelization.

  --seed <NUMBER>
      Set the random seed for deterministic processing (default: $seed).

  --debug
      Enable verbose output, retain all temporary files, and save additional
      debugging information for inspection.

  --no-retry
      Disable automatic retry of failed processing steps. By default, if
      segmentation or surface estimation fails for a subject, it is retried
      once before being reported as an error. Use this flag to disable that
      behaviour.

  --long-data <PATH>
      Use longitudinally realigned images as the actual input data while still
      deriving output folders from the original input path.
      If <PATH> is relative, it is interpreted relative to the derived output folder.
      Expected filenames inside <PATH> are tried in this order:
        1) '<stem>_desc-realigned.nii[.gz]'
        2) 'r<stem>.nii[.gz]'
        3) '<original-basename>'

${BOLD}${GREEN}Save Options:${NC}
  --out-dir <DIR>             
      Set the base output directory (relative or absolute).  
      Default: the current working directory.

      Output folder structure depends on the input dataset type:
        • BIDS datasets (if the upper-level folder of the input files is 'anat'):
            Results are placed in a BIDS-compatible derivatives folder:
              <dataset-root>/derivatives/T1Prep-v<version>/<sub-XXX>/<ses-YYY>/anat/
            Subject ('sub-XXX') and session ('ses-YYY') are auto-detected.
        • Non-BIDS datasets:
            Results are placed in subfolders similar to CAT12 output
            (e.g., 'mri/', 'surf/', 'report/', 'label') inside the specified output directory.

      If '--bids' is set, the BIDS derivatives substructure will always be used
      inside '<DIR>'.

  --bids                      
      Use BIDS derivatives naming conventions for all output files and folders
      instead of the default CAT12 style.
      
      Naming behaviour:
        • CAT12 style (default): Uses legacy folder and file names
          (e.g., 'mri/mwp1sub-01.nii', 'surf/lh.thickness.sub-01').
        • BIDS style: Uses standardized derivatives names, including subject/session
          identifiers, modality, and processing steps
          (e.g., 'sub-01_ses-1_space-MNI152NLin2009cAsym-nonlinear-modulated_label-WM_probseg.gz',
          'sub-01_ses-1_hemi-L_thickness.shape.gii').

      The complete mapping between internal outputs and both naming conventions
      is stored in 'Names.tsv' and can be customized.

      Examples:
        Input: /data/study/sub-01/ses-1/anat/sub-01_ses-1_T1w.nii.gz
        Default output (no --out-dir):
            /data/study/derivatives/T1Prep-v${T1PREP_VERSION}/sub-01/ses-1/anat/
        With --out-dir /results:
            /results/derivatives/T1Prep-v${T1PREP_VERSION}/sub-01/ses-1/anat/

        Input: /data/T1_images/subject01.nii.gz
        Default output (no --out-dir):
            /data/T1_images/mri/
        With --out-dir /results:
            /results/mri/

  --no-overwrite <STRING>     
      Prevent overwriting existing results by checking for the given filename pattern.

  --fmriprep
      Save fMRIPrep-compatible results:
        Save deformation fields as ANTs/ITK-compatible HDF5 (.h5) files
        Save dseg files
        Save cortex mask
        Save inflated surfaces
        Save spatially registered data with 2mm voxel size
      
  --no-track
      Do not collect telemetry information for T1Prep.

  --gz                        
      Save images in compressed NIfTI format (*.nii.gz).

  --skullstrip-only
    Run skull-stripping only and stop afterwards.
    Writes a skull-stripped image and a brain mask into the output mri folder.
    This option disables surface estimation and spherical registration.

  --no-skullstrip
    Skip skull-stripping (assumes input is already skull-stripped).
    A brain mask is derived from the input image.
    Note: This works best if the background is close to zero.

  --no-surf                   
      Skip surface and cortical thickness estimation.

  --no-seg                    
      Skip tissue segmentation processing.

  --no-sphere-reg             
      Skip spherical surface registration.

  --no-mwp                    
      Skip estimation of modulated and warped segmentations.

  --no-atlas                
      Skip estimation of volumetric atlases.

  --no-pial-white                
      Skip estimating pial and white matter surfaces during surface processing.

  --hemi, --hemisphere
      Save hemispheric (lh/rh) partitions of the segmentation.
      Note: these partitions are created automatically when surface estimation is enabled (default)
      and also when non-linear outputs are requested (e.g., --wp/--mwp or --atlas). Use this mainly
      together with --no-surf.

  --wp                        
      Additionally save warped segmentations.

  --rp                        
      Additionally save affine-registered segmentations.

  --p                         
      Additionally save native-space segmentations.

  --csf                       
      Additionally save CSF segmentations (default: only GM/WM are saved).
      This works currently only for native-space and affine-registered segmentations.

  --lesions                   
      Additionally save WMH lesion segmentations.
      This works currently only for native-space and affine-registered segmentations.

  --atlas                     
      Specify a volumetric atlas list in the format 
      "'suit','cobra'" (default: $atlas_vol).

  --atlas-surf                
      Specify a surface atlas list in the format 
      "'aparc_DK40.freesurfer','aparc_a2009s.freesurfer'" 
      (default: $atlas_surf).

${BOLD}${GREEN}Expert Options:${NC}
  --initial-surface <FILE>           
      Specify an initial surface to skip Marching Cubes approach for longitudinal 
      processing.

  --amap                      
      Use DeepMRIPrep segmentation only as initialization, followed by AMAP segmentation.

  --thickness-method <NUMBER> 
      Set the cortical thickness estimation method (default: $thickness_method):  
        1 = Tfs-distance (FreeSurfer) for PBT-based measure  
        2 = Tfs-distance (FreeSurfer) based on pial-to-white surface distance  
        3 = Pure PBT-based approach

  --no-correct-folding        
      Disable cortical thickness correction for folding effects.

  --pre-fwhm <NUMBER>         
      Set the pre-smoothing kernel size (FWHM) for Marching Cubes 
      (default: $pre_fwhm).

  --no-vessel
      Disable vessel removal

  --median-filter <NUMBER>    
      Apply the specified number of median filter passes to reduce topology artifacts
      for surface estimation (default: $median_filter).

  --fast                      
      Skip spherical registration, atlas estimation, estimation of pial and white
      surface and warped segmentation steps.
  
${BOLD}${GREEN}Examples:${NC}
${BLUE}T1Prep --out-dir test_folder sTRIO*.nii${NC}
  Process all files matching the pattern 'sTRIO*.nii'. Generate segmentation and 
  surface maps, saving the results in the 'test_folder' directory.

${BLUE}T1Prep --no-surf sTRIO*.nii${NC}
  Process all files matching the pattern 'sTRIO*.nii', but skip surface creation. 
  Only segmentation maps are generated and saved in the same directory as the input files.

${BLUE}T1Prep --python python3.9 --no-overwrite "surf/lh.thickness." sTRIO*.nii${NC}
  Process all files matching the pattern 'sTRIO*.nii' and use python3.9. Skip processing
  for files where 'surf/lh.thickness.*' already exists, and save new results in the same
  directory as the input files.

${BLUE}T1Prep --lesions --no-sphere-reg sTRIO*.nii${NC}
  Process all files matching the pattern 'sTRIO*.nii'. Skip processing of spherical
  registration, but additionally save lesion map (named p7sTRIO*.nii) in native space.
   
${BLUE}T1Prep --amap sTRIO*.nii${NC}
  Process all files matching the pattern 'sTRIO*.nii' and enable AMAP segmentation.

${BLUE}T1Prep --multi 8 --p --csf --atlas "'neuromorphometrics', 'suit'" sTRIO*.nii${NC}
  Process all files matching the pattern 'sTRIO*.nii'. Additionally save segmentations 
  in native space, including CSF segmentation and estimate ROI volumes for the defined atlases.
  The processing pipeline involves two stages 
  of parallelization:
  
  1. Segmentation (Python-based): Runs best with at least 16-24GB of memory per process. 
     The number of processes is automatically estimated based on available memory to 
     optimize resource usage. If you notice system overload due to peak memory usage
     you can set a higher minimum memory limit using the "--min-memory" flag.

  2. Surface Extraction: This stage does not require significant memory and is 
     distributed across all available processors or limited to the defined number of
     processes using the "--multi" flag.

  If "--multi" is set to a specific number (e.g., 8), the system still estimates memory-based 
  constraints for segmentation parallelization. However, the specified number of processes 
  (e.g., 8) will be used for surface extraction, ensuring efficient parallelization across 
  the two stages. The default setting is -1, which automatically estimates the number of
  available processors.
    
${BOLD}Purpose:${NC}
  This script facilitates the analysis of T1-weighted brain images by providing tools for 
  segmentation, surface mapping, and more.

${BOLD}Input:${NC}
  Accepts NIfTI files as input (extension nii/nii.gz).

${BOLD}Output:${NC}
  Produces segmented images and surface extractions.

${BOLD}Used Functions:${NC}
  surface_estimation.py
  segment.py
  progress_bar_multi.sh
  T1Prep_utils.sh
  parallelize

${BOLD}Author:${NC}
  Christian Gaser (christian.gaser@uni-jena.de)

EOM

check_python_cmd

# Source-tree mode only: offer to create the project-managed venv if missing.
# In installed mode T1prep_env="" and we skip this entirely (the venv is
# already in place because pip dropped this script into it).
if [ "${T1PREP_INSTALLED:-0}" -ne 1 ] && [ ! -d "${T1prep_env}" ]; then

  # Prompt the user with a Y/N question
  echo "${RED}${BOLD}Local Python environment "${T1prep_env}" not found.${NC}"
  read -p "Do you want to install required Python libraries? (Y/N)" response

  # Check if the user's answer is 'Y'
  case "$response" in
    [Yy]*)
      check_python_module venv
      check_python_module pip
      check_python_libraries
      ;;
  esac
fi

}

# ----------------------------------------------------------------------
# call main program
# ----------------------------------------------------------------------

main "${@}"
  
