#!/usr/bin/env bash
#
# PURPOSE:
#   Resample left (und optional right) hemisphere surface data to a target sphere,
#   optionally smooth the values and write a single combined GIfTI using
#   CAT_SurfResampleMulti.
#
# USAGE:
#   CAT_SurfResampleMulti_ui [options] <lh_input_value_file> [<lh_input_value_file> ...]
#
# INPUT:
#   Left hemisphere value files (e.g. lh.thickness.gii or lh.thickness.dat).
#   The script auto-derives corresponding central/sphere files and the right
#   hemisphere counterparts by name.
#
# OUTPUT:
#   One combined .gii per left input file containing LH (+RH) resampled values.
# ______________________________________________________________________
#
# Christian Gaser
# Structural Brain Mapping Group (https://neuro-jena.github.io)
# Departments of Neurology and Psychiatry
# Jena University Hospital
# ______________________________________________________________________

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

# Parameters
fwhm=12
ext="dat"
res="32k"
areal=0
areal_str=""

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

script_dir=$(dirname "$0")
root_dir=$(dirname $script_dir)
name_file=${root_dir}/src/t1prep/data/Names.tsv
surf_templates_dir=${root_dir}/src/t1prep/data/templates_surfaces_${res}

source "${script_dir}/T1Prep_utils.sh"

JOB_ID="jobrun_$(date +%s%N)"  # Unique ID based on timestamp in nanoseconds
PROGRESS_DIR="/tmp/progress_bars/$JOB_ID"

lh_fsavg_sphere=${surf_templates_dir}/lh.sphere.freesurfer.gii
lh_fsavg_mask=${surf_templates_dir}/lh.mask

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

declare -a FILES=()
declare -a LH_ARRAY=()
declare -a RH_ARRAY=()
declare -a OUT_ARRAY=()

# ----------------------------------------------------------------------
# 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
    # Check that area measure is not mixed with other measures
    local has_area=0 has_other=0
    for f in "${FILES[@]}"; do
      local b m
      b=$(basename "$f")
      m=${b#*.}; m=${m%%.*}
      if [[ "$m" == "area" ]]; then
        has_area=1
        areal_str="--areal"
        areal=1
      else
        has_other=1
      fi
    done
    if [[ "$has_area" -eq 1 && "$has_other" -eq 1 ]]; then
      echo -e "${RED}Cannot mix 'area' measure with other measures when using parallelization because here we have to use areal interpolation exclusively for surface area.${NC}" >&2
      exit 1
    fi

    local -a cmd_argv
    cmd_argv=("$0" "--no-parallel" "--res" "$res" "--ext" "$ext" "--trg-sphere" "$lh_fsavg_sphere" "--mask" "$lh_fsavg_mask" "--fwhm" "$fwhm")
    if [[ "$has_area" -eq 1 ]]; then
      cmd_argv+=("--areal")
    fi    

    if [[ -n "${out:-}" ]]; then
      cmd_argv+=("--out" "$out")
    fi

    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[@]}"
    else
      "${script_dir}/parallelize" -c "$command_str" "${FILES[@]}"
    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%%=*}"
    optarg="${2:-}"

    case "$1" in
      --out)
        exit_if_empty "$optname" "$optarg"
        out=$optarg
        shift
        ;;
      --no-parallel)
        parallel=0
        ;;
      --areal)
        areal=1
        areal_str="-areal"
        ;;
      --jobs|--processes|-p)
        exit_if_empty "$optname" "$optarg"
        jobs=$optarg
        shift
        ;;
      --ext)
        exit_if_empty "$optname" "$optarg"
        ext=$optarg
        shift
        ;;
      --res)
        exit_if_empty "$optname" "$optarg"
        res=$optarg
        shift
        ;;
      --trg-sphere)
        exit_if_empty "$optname" "$optarg"
        lh_fsavg_sphere=$optarg
        shift
        ;;
      --mask)
        exit_if_empty "$optname" "$optarg"
        lh_fsavg_mask=$optarg
        shift
        ;;
      --fwhm)
        exit_if_empty "$optname" "$optarg"
        fwhm=$optarg
        shift
        ;;
      -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
          FILES[$count]=$1
          ((count++))
        fi
        ;;
    esac
    shift
  done

  # RH hemisphere fallback
  if [[ -z "$rh_fsavg_sphere" ]]; then
    rh_fsavg_sphere="${lh_fsavg_sphere/lh./rh.}"
    rh_fsavg_sphere="${rh_fsavg_sphere/left_/right_}"
  fi
  if [[ -z "$rh_fsavg_mask" ]]; then
    rh_fsavg_mask="${lh_fsavg_mask/lh./rh.}"
    rh_fsavg_mask="${rh_fsavg_mask/left_/right_}"
  fi

}

# ----------------------------------------------------------------------
# per-file processing
# ----------------------------------------------------------------------
process_file() {
  local infile="$1"

  local dname bname hemi measure sphere_reg rh_bname out_name
  dname=$(dirname "$infile")
  bname=$(basename "$infile")

  # only lh.* supported for now
  if [[ "$bname" == lh.* ]]; then
    measure=${bname#*.}; measure=${measure%%.*}
    if [[ "$measure" == "area" ]]; then
      areal_str="-areal"
      areal=1
    else
      # Reset areal flags for non-area measures to prevent carry-over from previous files
      areal_str=""
      areal=0
    fi
    sphere_reg=".sphere.reg"
    rh_bname="${bname/lh./rh.}"
    
    # default output naming
    if [[ -n "${out}" ]]; then
      out_name="${out}/${bname/lh./s${fwhm}.mesh.}.${ext}"
    else
      out_name="${dname}/${bname/lh./s${fwhm}.mesh.}.${ext}"
    fi
    out_name="${out_name/.$measure/.${measure}.resampled_${res}}"
  elif [[ "$bname" == *_left* ]]; then
    echo -e "${RED}BIDS format not yet implemented (${bname}).${NC}" >&2
    return 1
  else
    echo -e "${RED}File '$infile' should be left hemisphere data (lh.*).${NC}" >&2
    return 1
  fi

  local lh_central="${dname}/${bname/.$measure/.central}.gii"
  local lh_sphere_reg="${dname}/${bname/.$measure/${sphere_reg}}.gii"
  local rh_central="${dname}/${rh_bname/.$measure/.central}.gii"
  local rh_sphere_reg="${dname}/${rh_bname/.$measure/${sphere_reg}}.gii"

  # build units for CAT_SurfResampleMulti
  local lh_unit="surf=${lh_central},src_sphere=${lh_sphere_reg},vals=${dname}/${bname},trg_sphere=${lh_fsavg_sphere},mask=${lh_fsavg_mask}"
  local rh_unit="surf=${rh_central},src_sphere=${rh_sphere_reg},vals=${dname}/${rh_bname},trg_sphere=${rh_fsavg_sphere},mask=${rh_fsavg_mask}"

  LH_ARRAY+=("$lh_unit")
  RH_ARRAY+=("$rh_unit")
  OUT_ARRAY+=("$out_name")
}


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

process()
{  
  n="${#FILES[@]}"
  
  # Check for mixed area/non-area measures in serial mode (same as parallel check)
  local has_area=0 has_other=0
  for f in "${FILES[@]}"; do
    local b m
    b=$(basename "$f")
    m=${b#*.}; m=${m%%.*}
    if [[ "$m" == "area" ]]; then
      has_area=1
    else
      has_other=1
    fi
  done
  if [[ "$has_area" -eq 1 && "$has_other" -eq 1 ]]; then
    echo -e "${RED}Cannot mix 'area' measure with other measures because we have to use areal interpolation exclusively for surface area.${NC}" >&2
    exit 1
  fi
  
  for f in "${FILES[@]}"; do
    if [[ ! -f "$f" && ! -L "$f" ]]; then
      echo -e "${RED}File '$f' not found – skipping.${NC}" >&2
      continue
    fi
    process_file "$f"
  done

  # Use the actual number of queued jobs (some inputs might have been skipped)
  n="${#OUT_ARRAY[@]}"
  if [ "$n" -lt 1 ]; then
    echo -e "${RED}No valid input files to process.${NC}" >&2
    exit 1
  fi
 
  # serial
  if [  ${n} -gt 1 ]; then
    mkdir -p "$PROGRESS_DIR"
    rm -f "$PROGRESS_DIR"/*.progress
    ${script_dir}/progress_bar_multi.sh 1 "$PROGRESS_DIR" 40 2 "CAT_SurfResampleMulti" &
    progress_pid=$!

    # Trap signals: Ctrl+C (INT), termination (TERM), or exit (EXIT)
    trap cleanup INT TERM
  fi 
   
  j=1
  for ((i=0; i<${n}; i++)); do
    if [  ${n} -gt 1 ]; then
      echo "$j/$n" > "$PROGRESS_DIR/job0.progress"
      ((j++))
      echo 0 > "$PROGRESS_DIR/job0.status"
    fi
    
    LH_UNIT="${LH_ARRAY[$i]}" RH_UNIT="${RH_ARRAY[$i]}" OUT_FILE="${OUT_ARRAY[$i]}" python3 - <<PYEOF 2>&1
import os, cat_surf.cli as cli
def parse_unit(s):
    d = {}
    for tok in s.split(','):
        k, _, v = tok.partition('=')
        d[k.strip()] = v.strip()
    return d
cli.surf_resample_multi(
    [parse_unit(os.environ['LH_UNIT']), parse_unit(os.environ['RH_UNIT'])],
    os.environ['OUT_FILE'],
    fwhm=${fwhm},
    areal=bool(${areal}),
)
PYEOF
      
    exit_code=$?
    pids=($!)

    if [[ ${n} -gt 1  &&  $exit_code -ne 0 ]]; then
      echo 1 > "$PROGRESS_DIR/job0.status"
    fi
  done
  
  if [  ${n} -gt 1 ]; then
    wait $progress_pid
    rm -r "$PROGRESS_DIR"
  fi
  exit $?
}


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

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

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

${BOLD}DESCRIPTION:${NC}
Resample left (and optional right) hemisphere surfaces/values to target spheres and
write a single combined GIfTI per subject using CAT_SurfResampleMulti.

${BOLD}OPTIONS:${NC}
  --out <DIR>                  Output directory (default: same as input file dir)
  --no-parallel                Disable parallelization via ${script_dir}/parallelize
  --jobs <N>                   Number of parallel processes (passed to parallelize -p)
  --ext <STR>                  Output file extension (default: $ext)
  --res <STR>                  Output surface resolution (32k or 4k, default: $res)
  --trg-sphere <FILE>          Target sphere for the left hemisphere
  --mask <FILE>                Target mask for the left hemisphere
  --fwhm <FLOAT>               Smoothing FWHM applied inside CAT_SurfResampleMulti (default: $fwhm)
  -h | --help                  Show this help and exit

${BOLD}INPUT NAMING EXPECTATIONS:${NC}
  lh.* case (supported now)
    lh.<measure>.something.gii
    Will look for:
      lh.central.gii               (same base, measure replaced by "central")
      lh.sphere.reg.gii            (same base, measure replaced by "sphere.reg")
      rh.<measure>.something.gii   (right hemisphere counterpart)
      rh.central.gii, rh.sphere.reg.gii accordingly

  *_left* (BIDS) not yet implemented → the script exits with an error.
        
${BOLD}OUTPUT:${NC}
  For each LH input file, writes a combined GIfTI (LH+RH) with resampled (and optionally smoothed) values. If extension is "dat" used, then an external binary file is saved, which is the default
  for use with SPM.

${BOLD}Dependencies:${NC}
  CAT_SurfResampleMulti
  ${script_dir}/progress_bar_multi.sh
  ${script_dir}/parallelize

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

EOM

}

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

main "${@}"