#!/usr/bin/env bash
# ______________________________________________________________________
#
# Christian Gaser
# Structural Brain Mapping Group (https://neuro-jena.github.io)
# Departments of Neurology and Psychiatry
# Jena University Hospital
# ______________________________________________________________________

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

time=$(date "+%Y%b%d_%H%M")
JOB_ID="jobrun_$(date +%s%N)"  # Unique ID based on timestamp in nanoseconds
if [ -n "${PROGRESS_DIR}" ]; then
  PROGRESS_DIR="${PROGRESS_DIR}"
else
  PROGRESS_DIR="/tmp/progress_bars/$JOB_ID"
fi
MEM_LIMIT=0
FG=1
LOGDIR="/tmp"
progress_pid= # This will hold the PID of the progress bar monitor
pids=()
delay=0

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

main ()
{
  #get_no_of_cpus
  parse_args "${@}"
  check_files
  if [ "$MEM_LIMIT" -gt 0 ]; then
    get_no_processes
  else
    get_no_of_cpus
  fi
  do_parallelize

  exit 0
}

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

parse_args () {
  local optname optarg
  count=0  # To track positional arguments (filenames)

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

    case "$1" in
        --mem-limit* | -m*)
            exit_if_empty "$optname" "$optarg"
            MEM_LIMIT="$optarg"
            shift
            ;;
        --processes* | -p*)
            exit_if_empty "$optname" "$optarg"
            NUM_JOBS="$optarg"
            shift
            ;;
        --logdir* | -l*)
            exit_if_empty "$optname" "$optarg"
            LOGDIR="$optarg"
            shift
            ;;
        --command* | -c*)
            # Ensure that the option has an argument
            exit_if_empty "$optname" "$optarg"
            COMMAND="$optarg"
            shift  # Skip the next argument since it's the value for --command
            ;;
        --delay* | -d*)
            exit_if_empty "$optname" "$optarg"
            delay="$optarg"
            shift
            ;;
        --bg* | -b*)
            FG=0
            ;;
        --test* | -t*)
            TEST=1
            ;;
        -h | --help | -v | --version | -V)
            help
            exit 1
            ;;
        -*)
            echo "`basename $0`: ERROR: Unrecognized option \"$1\"" >&2
            exit 1
            ;;
        *)
            # Treat anything not starting with `-` as a positional argument (filename)
            ARRAY[$count]="$1"
            ((count++))
            ;;
    esac

    shift  # Move to the next argument
  done
}

# ----------------------------------------------------------------------
# check arguments
# ----------------------------------------------------------------------

exit_if_empty ()
{
  local desc val

  desc="$1"
  shift
  val="$*"

  if [ -z "$val" ]; then
    echo ERROR: "${RED}No argument given with \"$desc\" command line argument!${NC}" >&2
    exit 1
  fi
}

# ----------------------------------------------------------------------
# check files
# ----------------------------------------------------------------------

check_files ()
{
  if [ -z "$COMMAND" ]; then
    echo "${RED}$FUNCNAME ERROR - no command defined.${NC}"
      help
    exit 1
  fi
  
  SIZE_OF_ARRAY="${#ARRAY[@]}"
  
  if [ "$SIZE_OF_ARRAY" -eq 0 ]; then
      echo "${RED}ERROR: No files or parameters given!${NC}" >&2
      help
      exit 1
  fi
}

# ----------------------------------------------------------------------
# get # of cpus
# ----------------------------------------------------------------------
# modified code from
# PPSS, the Parallel Processing Shell Script
# 
# Copyright (c) 2009, Louwrentius
# All rights reserved.

get_no_of_cpus () {

  CPUINFO=/proc/cpuinfo
  ARCH=`uname`

  # Check whether a floating number < 1.0 was used for NUM_JOBS
  if [ -n "$NUM_JOBS" ]; then
    percent=$(echo 100*"$NUM_JOBS" | bc | cut -f1 -d'.')
    digit=$(echo 100*"$NUM_JOBS" | bc | cut -f2 -d'.')

    if [[ $digit -eq 0 && $percent -gt 100 ]]; then
      echo "ERROR: Argument -p is a floating number > 1.0"
      exit 1
    fi
    
    # Only recognize percent values if NUM_JOBS < 1.0
    if [ $percent -lt 100 ]; then
      PERCENT=$percent
      NUM_JOBS=""
    fi
  fi

  if [ -z "$NUM_JOBS" ]; then
    if [ "$ARCH" == "Linux" ]; then
      NUM_JOBS=`grep ^processor $CPUINFO | wc -l`
  
    elif [ "$ARCH" == "Darwin" ]; then
      NUM_JOBS=`sysctl -a hw | grep -w "hw.logicalcpu" | awk '{ print $2 }'`
  
    elif [ "$ARCH" == "CYGWIN_NT" ] || [ "$ARCH" == "MSYS_NT" ] || [ "$ARCH" == "MINGW32_NT" ] || [ "$ARCH" == "MINGW64_NT" ]; then
      NUM_JOBS=`wmic cpu get NumberOfLogicalProcessors | grep -v NumberOfLogicalProcessors`
  
    else
      NUM_JOBS=`grep ^processor $CPUINFO | wc -l`
    fi
    
    # Use percent value to relate NUM_JOBS
    if [ -n "$PERCENT" ]; then
      NUM_JOBS=$(echo "$PERCENT"*"$NUM_JOBS"/100 | bc)
      NUM_JOBS=$(printf "%.0f" "$NUM_JOBS")  # rounds to 2 decimal places
    fi
  fi  
}

# ----------------------------------------------------------------------
# get # of processes w.r.t. available memory
# ----------------------------------------------------------------------

get_no_processes () {

  ARCH=`uname`

  if [ "$ARCH" == "Linux" ]; then
    # Get total installed memory in MB
    mem_total=$(grep MemTotal /proc/meminfo | awk '{print $2}')

    # Convert KB to GB
    mem_gb=$(echo "scale=2; $mem_total / 1024 / 1024" | bc)

  elif [ "$ARCH" == "Darwin" ]; then
    # Get total installed memory in bytes
    mem_total=$(sysctl hw.memsize | awk '{print $2}')

    # Convert bytes to GB
    mem_gb=$(echo "scale=2; $mem_total / 1024 / 1024 / 1024" | bc)

  elif [ "$ARCH" == "CYGWIN_NT" ] || [ "$ARCH" == "MSYS_NT" ] || [ "$ARCH" == "MINGW32_NT" ] || [ "$ARCH" == "MINGW64_NT" ]; then
    # Get total installed memory in KB
    mem_total=$(wmic ComputerSystem get TotalPhysicalMemory | grep -Eo '[0-9]+')

    # Convert bytes to GB
    mem_gb=$(echo "$mem_total / 1024 / 1024 / 1024" | bc)
  else
    echo "${RED}System $ARCH not recognized${NC}"
    exit 1
  fi
    
  # Calculate number of processes (at least 1) and require 10GB memory for each process
  NUM_JOBS=$(echo "$mem_gb / $MEM_LIMIT" | bc)
  if [ "$NUM_JOBS" -lt 1 ]; then
    NUM_JOBS=1
  fi
}

# ----------------------------------------------------------------------
# Cleanup function
# ----------------------------------------------------------------------

cleanup() {
  echo " Caught interrupt. Cleaning up..."
  echo ""
  echo "Please note that only new processes can be killed, but already started child processes hast to be maybe manually interrupted."
  if [ -n "$progress_pid" ]; then
    kill "$progress_pid" 2>/dev/null
  fi

  # Kill all background jobs started with nohup
  for pid in "${pids[@]}"; do
    if kill -0 "$pid" 2>/dev/null; then
      kill "$pid" 2>/dev/null
      wait "$pid" 2>/dev/null
    fi
  done

  wait $progress_pid
  rm -rf "$PROGRESS_DIR"
  exit 1
}

# ----------------------------------------------------------------------
# run parallelize
# ----------------------------------------------------------------------
 
do_parallelize ()
{
    
  # set overall starting time
  start=$(date +%s)

  SIZE_OF_ARRAY="${#ARRAY[@]}"

  # Correct if NUM_JOBS > SIZE_OF_ARRAY
  if [ "$SIZE_OF_ARRAY"  -lt "$NUM_JOBS" ]; then
    NUM_JOBS=$SIZE_OF_ARRAY
  fi

  # Initialize logging
  log="${LOGDIR}/parallelize_${HOSTNAME}_${time}.log"

  if [ -z "${TEST}" ]; then
    # Extract base command name (first token, basename, strip suffix after '-')
    first_token=$(printf '%s\n' "$COMMAND" | awk '{print $1}')
    base_name=$(basename "$first_token")
    cmd_name="${base_name%%-*}"

    # Start multi progress monitor
    if [ "$NUM_JOBS" -gt 1 ]; then      
      mkdir -p "$PROGRESS_DIR"
      rm -f "$PROGRESS_DIR"/*.progress
      
      $(dirname "$0")/progress_bar_multi.sh "$NUM_JOBS" "$PROGRESS_DIR" 40 "$cmd_name" &
      progress_pid=$!
    else
      $(dirname "$0")/progress_bar_multi.sh "$NUM_JOBS" "" 0 $SIZE_OF_ARRAY "$cmd_name" 40      
    fi

    # Trap signals: Ctrl+C (INT), termination (TERM), or exit (EXIT)
    trap cleanup INT TERM
  fi
    
  # If NUM_JOBS is 0, execute commands serially
  if [ "$NUM_JOBS" -lt 2 ]; then
      #echo "Executing serially..."
      i=0
      
      for file in "${ARRAY[@]}"; do
          eval "$COMMAND \"$file\""  # Use eval to handle complex commands
          exit_code=$?
          ((i++))
          
          if [ $exit_code -ne 0 ]; then
            $(dirname "$0")/progress_bar_multi.sh "$NUM_JOBS" "" $i $SIZE_OF_ARRAY "$cmd_name" 40 1
          else
            $(dirname "$0")/progress_bar_multi.sh "$NUM_JOBS" "" $i $SIZE_OF_ARRAY "$cmd_name" 40
          fi
          pids=($!)
      done

      # get overall processing time
      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"
      
      # wait until progress bar is finished
      wait $progress_pid

      # Display the result
      echo "${GREEN}-------------------------------------------------------${NC}"
      echo "${GREEN}${overall}${NC}"
      echo " "

      return
  fi

  if [ -z "${TEST}" ]; then
    echo "${GREEN}Check $log for logging information${NC}"
    echo > $log
  fi

  BLOCK=$((10000* $SIZE_OF_ARRAY / $NUM_JOBS ))
  i=0
  
  # Distribute files into N buckets and write per-job lists to PROGRESS_DIR.
  # Using files avoids word-splitting issues for paths with spaces.
  mkdir -p "$PROGRESS_DIR"
  for ((j=0; j<NUM_JOBS; j++)); do
    : > "$PROGRESS_DIR/job${j}.list"
  done

  # Process files in parallel: assign each ARRAY item to a bucket
  while [ "$i" -lt "$SIZE_OF_ARRAY" ]; do
    bucket=$((10000* i / BLOCK ))
    # Fallback clamp just in case of rounding
    if [ "$bucket" -ge "$NUM_JOBS" ]; then bucket=$((NUM_JOBS-1)); fi
    printf '%s\n' "${ARRAY[$i]}" >> "$PROGRESS_DIR/job${bucket}.list"
    ((i++))
  done

  for ((i=0; i<NUM_JOBS; i++)); do
    list_file="$PROGRESS_DIR/job${i}.list"
    # Skip empty buckets
    if [ ! -s "$list_file" ]; then
      continue
    fi

    j=$((i+1))
    total_k=$(wc -l < "$list_file" | tr -d ' ')
    if [ -z "${TEST}" ]; then
      echo "Job ${j}/${NUM_JOBS}: $COMMAND (items: $total_k)" >> "$log"
      # Pass required variables to the subshell via environment to avoid empty expansions
      nohup env PROGRESS_DIR="$PROGRESS_DIR" COMMAND="$COMMAND" LOG_FILE="$log" JOB_INDEX="$i" TOTAL_K="$total_k" LIST_FILE="$list_file" bash -c '
        k_count=0
        status=0   # 0 means success; 1 means at least one failure
        # Initialize progress
        echo "${k_count}/${TOTAL_K}" > "${PROGRESS_DIR}/job${JOB_INDEX}.progress"
        : > "${PROGRESS_DIR}/job${JOB_INDEX}.failed"

        # Read list file line-by-line to preserve spaces
        while IFS= read -r k || [ -n "$k" ]; do
          k_count=$((k_count+1))
          # Log which file is being processed so errors can be traced
          echo "--- Processing: $k ---" >> "${LOG_FILE}"
          # Execute the command with the current item; use eval to expand the command string
          eval "$COMMAND" "\"$k\"" >> "${LOG_FILE}" 2>&1
          exit_code=$?
          if [ $exit_code -ne 0 ]; then
            status=1
            echo "$k" >> "${PROGRESS_DIR}/job${JOB_INDEX}.failed"
            echo "--- FAILED (exit code $exit_code): $k ---" >> "${LOG_FILE}"
          fi
          # Update progress and status
          echo "${k_count}/${TOTAL_K}" > "${PROGRESS_DIR}/job${JOB_INDEX}.progress"
          echo "$status" > "${PROGRESS_DIR}/job${JOB_INDEX}.status"
        done < "$LIST_FILE"
      ' >> "$log" 2>&1 & pids+=($!)
      sleep $delay
    fi
  done
        
  # keep all processes in foreground
  if [ "${FG}" -eq 1 ]; then
      #echo "Waiting for jobs to complete..."
      for pid in "${pids[@]}"; do
        wait "$pid"
      done

      errors=0
      error_list=""
      for ((i=0; i<$NUM_JOBS; i++)); do
        status_file="$PROGRESS_DIR/job${i}.status"
        failed_file="$PROGRESS_DIR/job${i}.failed"
        if [ -f "$status_file" ]; then
          status=$(cat "$status_file")
        else
          status=0
        fi
        if [[ "$status" -ne 0 ]]; then
          ((errors++))
          # Collect actual failed filenames from the job
          if [ -s "$failed_file" ]; then
            while IFS= read -r failed_item; do
              error_list+="  $failed_item"
              error_list+=$'\n'
            done < "$failed_file"
          fi
        fi
      done

      # get overall processing time
      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"
      
      # wait until progress bar is finished
      wait $progress_pid
      if [ -z "${TEST}" ]; then
        rm -r "$PROGRESS_DIR"  
      fi

      # Display the result depending on error status
      if [[ "$errors" -ne 0 ]]; then
        overall+=" with $errors error(s)"
        echo "${RED}-------------------------------------------------------"
        echo "${RED}${overall}${NC}"
        if [ -n "$error_list" ]; then
          echo "${RED}Failed items:${NC}"
          printf '%s' "$error_list"
        fi
        echo "${RED}Check $log for details${NC}"
        exit 1
      else
        echo "${GREEN}-------------------------------------------------------"
        echo "${GREEN}${overall}${NC}"
      fi
  fi
}

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

help ()
{
cat <<__EOM__

${BOLD}USAGE:${NC}
  $(basename "$0") [options] -c <STRING> [filename|filepattern]

${BOLD}OPTIONS:${NC}
  -p, --processes <NUMBER> Specify the number of parallel jobs (default $NUM_JOBS).
                           If <NUMBER> is a floating-point value between 0 and 1,
                           it will be interpreted as a fraction of available CPU cores.
                           For example, 0.7 means 70% of CPU cores will be used.
  -m, --mem-limit <INT>    Memory limit for each process in GB (default is ${MEM_LIMIT}GB).
  -c, --command <STRING>   Command to be executed in parallel.
  -d, --delay <NUMBER>     Specify a delay (in seconds) between starting each job. 
                           This helps to avoid simultaneous peak memory usage, which can 
                           lead to system overload or job failures (default: $delay).
  -b, --bg                 Keep all processes in background.
  -l, --logdir <DIR>       Directory to save the log file (default $LOGDIR).
  -t, --test               Run in test mode: print the files to be processed without 
                           executing the command.

${BOLD}DESCRIPTION:${NC}
  Execute commands in parallel across multiple files specified by filenames or a wildcard 
  pattern. Optionally, adjust the number of processes, memory limits, and logging preferences.

${BOLD}OUTPUT:${NC}
  If a log directory is specified, logs are saved in that folder in a file named 
  'parallelize_${HOSTNAME}_\${time}.log', incorporating the current date and time.

${BOLD}EXAMPLES:${NC}
  ${BLUE}parallelize -c "niismooth -v -fwhm 8" sTRIO*.nii${NC}
    Parallelizes the 'niismooth' command with an 8mm FWHM, applied to all files matching 
    'sTRIO*.nii', displaying verbose output.

  ${BLUE}parallelize -c "gunzip" *.zip${NC}
    Parallelizes unzipping of all .zip files in the current directory.

  ${BLUE}parallelize -m 20 -c "T1prep --no-surf" sTRIO*.nii${NC}
    Parallelizes 'T1prep' without surface estimation, setting a 10GB RAM limit per process.

  ${BLUE}parallelize -p 0.7 -c "gunzip" *.zip${NC}
    Parallelizes unzipping of all .zip files in the current directory by using 70% of the
    available processors.

  ${BLUE}${T1prep_dir}/T1Prep --out-dir test_folder --no-surf --hemisphere sTRIO*.nii
  ${T1prep_dir}/parallelize -p 8 -c "T1prep.sh --out-dir test_folder --no-seg" sTRIO*.nii${NC}
    1. Runs 'T1prep.sh' on all files matching the pattern 'sTRIO*.nii', skipping surface 
       creation and saving results in 'test_folder'.
    2. Parallelizes the processing using 8 processors for improved performance. The first 
       command is not parallelized because it already uses multi-threading and is 
       memory-intensive, while the second command (surface creation) is single-threaded 
       and can benefit from parallelization.
       
    NOTE: You can automatically enable parallelization in T1Prep by using the '--multi' flag.

${BOLD}Used Functions:${NC}
  progress_bar_multi.sh

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

__EOM__
}

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

main ${1+"$@"}

