#!/usr/bin/env bash
#
# PURPOSE:
#   Compute curvature-based parameters from surface files and write them to
#   text (or GIfTI) files. This is a T1Prep wrapper around CAT_SurfCurvature
#   CAT_SurfFractalDimension, CAT_SurfArea, CAT_SurfRatio, and CAT_SurfSulcusDepth.
#
#   If a left hemisphere file (lh.*) is provided, the corresponding right
#   hemisphere file (rh.*) is automatically processed as well (if present).
#
# USAGE:
#   $0 [options] <surface_file> [<surface_file> ...]
#
# INPUT:
#   Surface files (*.obj or *.gii). For FreeSurfer-style naming (lh.* / rh.*),
#   output names follow the CAT_DumpCurv_ui scheme.
#
# OUTPUT:
#   One output per processed hemisphere/input (txt by default, or *.gii when
#   -gifti is used).
# ______________________________________________________________________
#
# Christian Gaser
# Structural Brain Mapping Group (https://neuro-jena.github.io)
# Departments of Neurology and Psychiatry
# Jena University Hospital
# ______________________________________________________________________

# ----------------------------------------------------------------------
# global parameters
# ----------------------------------------------------------------------

# defaults
os_type=$(uname -s) # Determine OS type

script_dir=$(dirname "$0")
source "${script_dir}/T1Prep_utils.sh"

# Parameters
use_absolute=0
method=0
removestr=''
noclobber=0
gifti=0

# Parallelization (via scripts/parallelize)
parallel=1
jobs=""

declare -a FILES=()
declare -a WORK_FILES=()

# ----------------------------------------------------------------------
# helpers (bash 3.2 compatible)
# ----------------------------------------------------------------------

array_contains()
{
  local seeking="$1"; shift
  local element
  for element in "$@"; do
    if [[ "$element" == "$seeking" ]]; then
      return 0
    fi
  done
  return 1
}

# ----------------------------------------------------------------------
# run main
# ----------------------------------------------------------------------

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

  # Activate Python virtual environment
  if [[ -f "${root_dir}/env/bin/activate" ]]; then
    source "${root_dir}/env/bin/activate"
  fi

  # Optional parallel execution: delegate one input per process via scripts/parallelize
  if [[ "$parallel" -eq 1 && "${#FILES[@]}" -gt 1 ]]; then
    local -a files_for_parallel
    local -a lh_suffixes
    local f b suffix

    # Build a filtered file list:
    # - keep only existing files
    # - if both lh.X and rh.X are provided, skip rh.X (since lh.X auto-adds RH)
    for f in "${FILES[@]}"; do
      if [[ ! -f "$f" && ! -L "$f" ]]; then
        continue
      fi
      b=$(basename "$f")
      if [[ "$b" == lh.* ]]; then
        suffix="${b#lh.}"
        if ! array_contains "$suffix" "${lh_suffixes[@]}"; then
          lh_suffixes+=("$suffix")
        fi
      fi
    done
    for f in "${FILES[@]}"; do
      if [[ ! -f "$f" && ! -L "$f" ]]; then
        continue
      fi
      b=$(basename "$f")
      if [[ "$b" == rh.* ]]; then
        suffix="${b#rh.}"
        if array_contains "$suffix" "${lh_suffixes[@]}"; then
          continue
        fi
      fi
      if ! array_contains "$f" "${files_for_parallel[@]}"; then
        files_for_parallel+=("$f")
      fi
    done

    if [[ "${#files_for_parallel[@]}" -lt 2 ]]; then
      process
      exit 0
    fi

    local -a cmd_argv
    cmd_argv=("$0" "--no-parallel")

    if [[ -n "$removestr" ]]; then cmd_argv+=("-r" "$removestr"); fi
    if [[ "$noclobber" -eq 1 ]]; then cmd_argv+=("-noclobber"); fi
    if [[ "$gifti" -eq 1 ]]; then cmd_argv+=("-gifti"); fi

    case "$method" in
      0) cmd_argv+=("-gy") ;;
      1) cmd_argv+=("-gc") ;;
      2) cmd_argv+=("-cv") ;;
      3) cmd_argv+=("-si") ;;
      4) cmd_argv+=("-mc") ;;
      6) cmd_argv+=("-be") ;;
      7) cmd_argv+=("-sh") ;;
      8) cmd_argv+=("-fi") ;;
      9) cmd_argv+=("-min-curv") ;;
      10) cmd_argv+=("-max-curv") ;;
      11) cmd_argv+=("-dp") ;;
      12) cmd_argv+=("-depth") ;;
      13) cmd_argv+=("-sqrt-depth") ;;
      14) cmd_argv+=("-fd") ;;
      15) cmd_argv+=("-sr") ;;
      16) cmd_argv+=("-sra") ;;
      17) cmd_argv+=("-area") ;;
      18) cmd_argv+=("-sulc") ;;
    esac

    local command_str=""
    local arg
    for arg in "${cmd_argv[@]}"; do
      command_str+="$(printf '%q ' "$arg")"
    done
    command_str=${command_str% }

    if [[ -n "$jobs" ]]; then
      "${script_dir}/parallelize" -p "$jobs" -c "$command_str" "${files_for_parallel[@]}"
    else
      "${script_dir}/parallelize" -c "$command_str" "${files_for_parallel[@]}"
    fi
    exit $?
  fi
  
  process

  exit 0
}

# ----------------------------------------------------------------------
# check arguments and files
# ----------------------------------------------------------------------

parse_args()
{
  local optname optarg
  
  if [ $# -lt 1 ]; then
    help
    exit 1
  fi

  count=0
  while [ $# -gt 0 ]; do
    optname="${1%%=*}"
    # Support both "-x 10" and "-x=10" styles
    if [[ "$1" == *=* ]]; then
      optarg="${1#*=}"
    else
      optarg="${2:-}"
    fi

    case "$1" in
    --no-parallel)
      parallel=0
      ;;
    --jobs|--processes|-p)
      exit_if_empty "$optname" "$optarg"
      jobs=$optarg
      # Only shift if the value was provided as the next token
      if [[ "$1" != *=* ]]; then shift; fi
      ;;
    --r* | -r*)
      exit_if_empty "$optname" "$optarg"
      removestr=$optarg
      if [[ "$1" != *=* ]]; then shift; fi
      ;;
    --n* | -n*)
      noclobber=1
      ;;
    --gifti | -gifti)
      gifti=1
      ;;
    --gy* | -gy*)
      method=0
      use_absolute=1
      ;;
    --gc* | -gc*)
      method=1
      ;;
    --cv | -cv)
      method=2
      ;;
    --si | -si)
      method=3
      ;;
    --mc | -mc | --curv | -curv)
      method=4
      ;;
    --b* | -b*)
      method=6
      ;;
    --sh | -sh)
      method=7
      ;;
    --fi | -fi)
      method=8
      ;;
    --min-curv | -min-curv)
      method=9
      ;;
    --max-curv | -max-curv)
      method=10
      ;;
    --dp | -dp)
      method=11
      ;;
    --depth | -depth)
      method=12
      ;;
    --sqrt-depth | -sqrt-depth)
      method=13
      ;;
    --fd | -fd | --frac* | -frac*)
      method=14
      ;;
    --sra | -sra | --tGIa | -tGIa)
      method=16
      ;;
    --sr | -sr | --tGI | -tGI)
      method=15
      ;;
    --area | -area)
      method=17
      ;;
    --sulc | -sulc)
      method=18
      ;;
    -h | --help | --h | -v | --version | -V)
      help
      exit 1
      ;;
    -*)
      echo "basename $0: ERROR: Unrecognized option \"$1\"" >&2
      exit 1
      ;;
    *)
      FILES[$count]="$1"
      ((count++))
      ;;
    esac
    shift
  done

}

# ----------------------------------------------------------------------
# per-file processing
# ----------------------------------------------------------------------
get_outputname()
{
  local infile="$1"
  local dn bn bn1 str

  dn=$(dirname "$infile")
  bn=$(basename "$infile")

  # remove extensions
  bn=${bn%.obj}
  bn=${bn%.gii}

  # optional remove string
  if [[ -n "$removestr" ]]; then
    bn=${bn//${removestr}/}
  fi

  case "$method" in
    0) str='gyrification' ;;
    1) str='gauss-curv' ;;
    2) str='curvedness' ;;
    3) str='shape-index' ;;
    4) str='curv' ;;
    6) str='bending-energy' ;;
    7) str='sharpness' ;;
    8) str='folding-index' ;;
    9) str='min-curv' ;;
    10) str='max-curv' ;;
    11) str="depth-potential" ;;
    12) str="depth" ;;
    13) str="sqrtdepth" ;;
    14) str="fractaldimension" ;;
    15) str="toroGI20mm" ;;
    16) str="toroGIa" ;;
    17) str="area" ;;
    18) str="sulc" ;;
  esac

  bn1=${bn%%.*}
  if [[ "$bn1" == "lh" || "$bn1" == "rh" ]]; then
    # FreeSurfer-like naming: collapse "lh.central.SUBJ" -> "lh.SUBJ"
    bn=${bn//h.central/h}
    bn=${bn//h.white/h}
    bn=${bn//h.pial/h}
    if [[ -n "$removestr" ]]; then
      bn=${bn//${removestr}/}
    fi
    bn=${bn//lh./}
    bn=${bn//rh./}
    outputname="${dn}/${bn1}.${str}.${bn}"
  else
    outputname="${dn}/${str}_${bn}"
  fi

  if [[ "$gifti" -eq 1 ]]; then
    outputname="${outputname}.gii"
  fi
}


# ----------------------------------------------------------------------
# process data
# ----------------------------------------------------------------------

process()
{
  if [ "${#FILES[@]}" -lt 1 ]; then
    echo "${RED}No arguments given.${NC}" >&2
    echo
    help
    exit 1
  fi

  WORK_FILES=()
  
  local f dname bname rh_bname rh_file
  for f in "${FILES[@]}"; do
    if [[ ! -f "$f" && ! -L "$f" ]]; then
      echo -e "${RED}File '$f' was not found and will be skipped.${NC}" >&2
      continue
    fi

    if ! array_contains "$f" "${WORK_FILES[@]}"; then
      WORK_FILES+=("$f")
    fi

    bname=$(basename "$f")
    if [[ "$bname" == lh.* ]]; then
      dname=$(dirname "$f")
      rh_bname="rh.${bname#lh.}"
      rh_file="${dname}/${rh_bname}"
      if [[ -f "$rh_file" || -L "$rh_file" ]]; then
        if ! array_contains "$rh_file" "${WORK_FILES[@]}"; then
          WORK_FILES+=("$rh_file")
        fi
      else
        echo -e "${YELLOW}Right hemisphere file '$rh_file' not found – skipping RH.${NC}" >&2
      fi
    fi
  done

  if [ "${#WORK_FILES[@]}" -lt 1 ]; then
    echo "${RED}No valid input files to process.${NC}" >&2
    exit 1
  fi

  local infile
  for infile in "${WORK_FILES[@]}"; do
    get_outputname "$infile"
    if [[ -f "${outputname}" && "$noclobber" -eq 1 ]]; then
      echo "${outputname} will not be overwritten. Do not use -noclobber if you want to overwrite." >&2
      continue
    fi
    sphere=${infile//central/sphere}
    sphere=${sphere//pial/sphere}
    sphere=${sphere//white/sphere}
    
    echo "calculating ${outputname}"
    IN="$infile" OUT="$outputname" SPHERE="$sphere" python3 - <<PYEOF 2>&1
import os, cat_surf.cli as cli
infile     = os.environ['IN']
outputname = os.environ['OUT']
sphere     = os.environ.get('SPHERE', '')
method      = ${method}
use_absolute = bool(${use_absolute})
if method < 11:
    cli.surf_curvature(infile, outputname, curvtype=method, use_abs_values=use_absolute)
elif method == 11:
    cli.surf_curvature(infile, outputname, curvtype=650, use_abs_values=use_absolute)
elif method == 12:
    cli.surf_sulcus_depth(infile, sphere, outputname)
elif method == 13:
    cli.surf_sulcus_depth(infile, sphere, outputname, sqrt_transform=True)
elif method == 14:
    cli.surf_fractal_dimension(infile, sphere, outputname, smooth=False)
elif method == 15:
    cli.surf_ratio(infile, outputname, radius=20.0, normalize=False)
elif method == 16:
    cli.surf_ratio(infile, outputname, radius=-1.0, normalize=False)
elif method == 17:
    cli.surf_area(infile, output_values_file=outputname)
elif method == 18:
    cli.surf_curvature(infile, outputname, curvtype=11, use_abs_values=use_absolute)
PYEOF
  done
}


# ----------------------------------------------------------------------
# help
# ----------------------------------------------------------------------

help() {
cat << EOM
${BOLD}${BLUE}$0
---------------------------------------------${NC}

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

${BOLD}DESCRIPTION:${NC}
Compute surface-based parameters from surface files using CAT_SurfCurvature,
CAT_SurfFractalDimension, CAT_SurfArea, CAT_SurfRatio, and CAT_SurfSulcusDepth.
If an input file is named like "lh.*", the corresponding "rh.*" file is also
processed automatically (if present).

${BOLD}OPTIONS:${NC}
  --no-parallel    Disable parallelization via ${script_dir}/parallelize
  --jobs <N>       Number of parallel processes (passed to parallelize -p)

  --r <STR>        Remove string from output name
  --noclobber      Do not overwrite existing output
  --gifti          Save output as GIfTI (appends .gii)

  --gy             Gyrification based on absolute mean curvature abs(k1+k2)/2 [default]
  --mc | --curv    Mean curvature (in radians) (k1+k2)/2
  --be | --bi      Bending energy k1*k1 + k2*k2
  --gc             Gaussian curvature k1*k2
  --cv             Curvedness sqrt(0.5*(k1*k1+k2*k2))
  --si             Shape index atan((k1+k2)/(k2-k1))
  --sh             Sharpness (k1 - k2)^2
  --fi             Folding index |k1|*(|k1| - |k2|)
  --min-curv       Minimum curvature k2
  --max-curv       Maximum curvature k1
  --dp             Depth potential with alpha = 1/curvtype (curvtype=650)
  --sulc           Sulcal depth (using inflated surface based on Freesurfer method
                   aka average convexity)
  --depth          Sulcal depth (using convex hull)
  --sqrt-depth     Sqrt-transformed sulcal depth (more normally distributed)
  --area           Surface area
  --fd  | --frac   Fractal dimension
  --sr  | --tGI    Surface ratio with radius=20mm (Toro et al. 2008)
  --sra | --tGIa   Surface ratio with adaptive radius estimation (Toro et al. 2008)

  -h | --help      Show this help and exit

${BOLD}INPUT NAMING EXPECTATIONS:${NC}
  *.obj and *.gii are accepted.
  For FreeSurfer-like files named "lh.*" or "rh.*", output naming follows
  CAT_DumpCurv_ui conventions.
        
${BOLD}OUTPUT:${NC}
  Per input (and hemisphere if auto-RH applies): curvature values for each node.

${BOLD}Dependencies:${NC}
  CAT_SurfCurvature
  CAT_SurfFractalDimension
  CAT_SurfSulcusDepth
  CAT_SurfRatio
  CAT_SurfArea
  ${script_dir}/parallelize

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

EOM

}

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

main "${@}"