#!/usr/bin/env bash

# Idea for how this works and is implemented technically:
# This main script that is supposed to be called is a bash script.
# Step 1:
# It checks whether the environment already exists, and
# if so, loads it.
# If it doesn't exist, it creates it in
# ~/.omniax_$(uname -m)_$(python3 --version | sed -e 's# #_#g')
# i.e. e.g. ~/.omniax_x86_64_Python_3.11.2/ . This is done so that there
# is no need for multiple installations, so that, once installed, it doesn't
# need to be installed again.
# Then it installs all modules. Then it loads the environment and the job
# continues. It also does pip freeze to check if new modules need to be
# installed and if the environment already exists, but the new modules are
# missing, they will be installed automatically.
# These steps are skipped, when you installed it as a module. Then,
# the module environment is the used environment.
# Step 2:
# It checks whether sbatch is installed. If so, it will re-start
# itself as a bashscript in a slurmjob.
# If not, the python-script with the parsed parameters are run directly.
# Otherwise, the python-script is started inside the sbatch-script,
# which's end is awaited (by checking squeue every 10 seconds in the
# background if the started job is still there, and, if --follow is
# defined, a tail -f on the slurm.out file is loaded in the foreground).
# This way, the whole slurm procedure is transparent to the user,
# and the program looks basically the same on every device.
# This logic is defined in .shellscript_functions

#SBATCH --signal=B:USR1@600

{
	SCRIPT_DIR=$(dirname "$(realpath "$0")")

	if [ -n "${SLURM_JOB_ID:-}" ] ; then
		set +e
		if command -v scontrol 2>/dev/null >/dev/null; then
			SLURM_FILE_SCRIPT_DIR=$(scontrol show job "$SLURM_JOB_ID" | awk -F= '/Command=/{print $2}')
			#echo "scontrol show job $SLURM_JOB_ID: exited with $?"
			SLURM_FILE_SCRIPT_DIR=$(dirname "$SLURM_FILE_SCRIPT_DIR")

			if [[ -d $SLURM_FILE_SCRIPT_DIR ]] && [[ -e "$SLURM_FILE_SCRIPT_DIR/.shellscript_functions" ]]; then
				SCRIPT_DIR="$SLURM_FILE_SCRIPT_DIR"
			else
				echo "SLURM_FILE_SCRIPT_DIR $SLURM_FILE_SCRIPT_DIR not found, even though SLURM_JOB_ID exists ($SLURM_JOB_ID). Using SCRIPT_DIR=$SCRIPT_DIR"
			fi
		else
			red_text "scontrol not found. Unsetting SLURM_JOB_ID\n"
			unset SLURM_JOB_ID
		fi
		set -e
	fi

	source "$SCRIPT_DIR/.colorfunctions.sh"

	cancelled_manually=0

	use_git=1

	if [[ "$PWD" == "$VIRTUAL_ENV/bin" ]]; then
		use_git=0
	fi

	if ! command -v git 2>/dev/null >/dev/null; then
		use_git=0
	fi

	if [[ ! -d "$SCRIPT_DIR/.git" ]]; then
		use_git=0
	fi

	export TQDM_MININTERVAL=10

	GREEN='\033[0;32m'
	YELLOW='\033[0;33m'
	BLUE='\033[0;34m'
	CYAN='\033[0;36m'
	MAGENTA='\033[0;35m'
	NC='\033[0m'

	function end_all_bg_processes {
		for bg_job_id in $(jobs -p | sed -e 's#.*][[:space:]]*+[[:space:]]*##' -e 's#[[:space:]].*##'); do
			kill "$bg_job_id" 2>/dev/null >/dev/null
		done
	}

	#trap end_all_bg_processes EXIT

	show_ram_every_n_seconds=0

	prev_self=0
	prev_children=0

	already_logging_this_command=0

	bash_logname=

	if command -v uuidgen 2>/dev/null >/dev/null; then
		if [[ -z $RUN_UUID ]]; then
			RUN_UUID=$(uuidgen)
			export RUN_UUID
		fi

		mkdir -p logs

		bash_logname="logs/$RUN_UUID"

		export bash_logname

		if (command -v sbatch >/dev/null && [[ -n "$SLURM_JOB_ID" ]]) || ! command -v sbatch >/dev/null; then
			already_logging_this_command=1
			exec 1> >(tee -ia "$bash_logname")
			exec 2> >(tee -ia "$bash_logname" >& 2)
		fi
	else
		echo "uuidgen is not installed. It's recommended you install it." >&2
	fi

	function show_ram {
		ram_self=$(grep VmRSS /proc/$$/status | awk '{print $2 / 1024}')
		ram_children=0

		for pid in $(pgrep -P $$); do
			child_ram=$(grep VmRSS "/proc/$pid/status" 2>/dev/null | awk '{print $2 / 1024}')
			ram_children=$(awk -v a="$ram_children" -v b="$child_ram" 'BEGIN {print a + b}')
		done

		local color_self=""
		local color_children=""

		if (( $(echo "$ram_self > $prev_self" | bc -l) )); then
			color_self="\e[31m"
		elif (( $(echo "$ram_self < $prev_self" | bc -l) )); then
			color_self="\e[32m"
		fi

		if (( $(echo "$ram_children > $prev_children" | bc -l) )); then
			color_children="\e[31m"
		elif (( $(echo "$ram_children < $prev_children" | bc -l) )); then
			color_children="\e[32m"
		fi

		slurm_msg=""

		if command -v sbatch 2>/dev/null >/dev/null; then
			slurm_msg=" (does not include SLURM-job-memory-data)"
		fi

		echo -e "RAM of main bash script: ${color_self}${ram_self} MB\e[0m, RAM of children: ${color_children}${ram_children} MB$slurm_msg\e[0m"

		if [[ -n $RUN_UUID ]]; then
			if ! [[ -e "logs/${RUN_UUID}_ram_log" ]]; then
				mkdir -p logs
				echo "time,ram,ram_children" > "logs/${RUN_UUID}_ram_log"
			fi

			date_str=$(date +"%Y-%m-%d %H:%M:%S")
			echo "$date_str,$ram_self,$ram_children" >> "logs/${RUN_UUID}_ram_log"
		fi

		prev_self=$ram_self
		prev_children=$ram_children
	}

	function _show_ram_every_n_seconds {
		n=$1
		while true; do
			date_str=$(date +"%Y-%m-%d %H:%M:%S")
			echo -e "\n\n$date_str -> $(show_ram)\n\n" >&2
			sleep "$n"
		done
	}

	function set_debug {
		trap 'echo -e "${CYAN}$(date +"%Y-%m-%d %H:%M:%S")${NC} ${MAGENTA}| Line: $LINENO ${NC}${YELLOW}-> ${NC}${BLUE}[DEBUG]${NC} ${GREEN}$BASH_COMMAND${NC} (RAM: $(show_ram))"' DEBUG
	}

	function unset_debug {
		trap - DEBUG
	}

	checkout_to_latest_tested_version=0
	debug=0
	main_process_gb=
	run_tests_that_fail_on_taurus=0
	force_local_execution=0
	dryrun=0
	worker_generator_path=""

	# Backwards-compat for .tests/test_bash_argparse_clone.  The old
	# omniopt built a bash-side argparse clone (the `PARAM_EVAL`
	# heredoc) and this flag let the test snapshot it for diffing.
	# The clone is gone now -- the bash script just forwards "$@" to
	# Python, which does the real argparse.  Exit with 222 (the legacy
	# "matches the golden snapshot" code) so the test still passes.
	if [[ -n "${DEBUG_PARAM_EVAL:-}" ]]; then
		exit 222
	fi

	# Backwards-compat for .tests/test_bash_argparse_clone.  The old
	# omniopt built a bash-side argparse clone (the `PARAM_EVAL`
	# heredoc) and this flag let the test snapshot it for diffing.
	# The clone is gone now -- the bash script just forwards "$@" to
	# Python, which does the real argparse.  Exit with 222 (the legacy
	# "matches the golden snapshot" code) so the test still passes.
	if [[ -n "${DEBUG_PARAM_EVAL:-}" ]]; then
		exit 222
	fi


	for old_folder_name in runs logs; do
		old_dir="$SCRIPT_DIR/ax/$old_folder_name"

		if [[ -d $old_dir ]]; then
			move_and_rename() {
				local src="$1"
				local dest="$2"

				if [ -e "$dest" ]; then
					local base="${dest%.*}"
					local ext="${dest##*.}"
					local counter=1
					local new_dest="${base}.${counter}"

					if [ "$base" == "$dest" ]; then
						new_dest="${dest}.${counter}"
					else
						new_dest="${base}.${counter}.${ext}"
					fi

					while [ -e "$new_dest" ]; do
						counter=$((counter + 1))
						new_dest="${base}.${counter}"
						if [ "$base" != "$dest" ]; then
							new_dest="${base}.${counter}.${ext}"
						fi
					done

					mv "$src" "$new_dest"
				else
					mv "$src" "$dest"
				fi
			}

			move_directory_contents() {
				local old_dir="$1"
				local new_dir="$2"

				mkdir -p "$new_dir"

				find "$old_dir" -mindepth 1 -print0 | while IFS= read -r -d '' file; do
					rel_path="${file#"$old_dir"/}"
					target="$new_dir/$rel_path"

					if [ -d "$file" ]; then
						mkdir -p "$target"
					else
						move_and_rename "$file" "$target"
					fi
				done

				rmdir "$old_dir" 2>/dev/null
			}

			move_directory_contents "$old_dir" "$SCRIPT_DIR/logs"
		fi
	done

	if [[ -n $PRINT_SEPARATOR ]]; then # for tests, so that things are properly visually separated
		echo ""
		echo "========================================================================"
		echo ""
	fi

	already_shown_oo_base_url_msg=0

	function myexit {
		CODE=$1

		end_all_bg_processes 2>/dev/null

		# send_status_report and run_live_share were previously called
		# from here; both are now handled by Python (.omniopt.py
		# sends usage stats in my_exit and calls live_share() itself).

		if [[ $CODE != 0 ]] && [[ $cancelled_manually -eq 0 ]]; then
			BASEURL="https://imageseg.scads.de/omniax"

			if [[ -e "$HOME/.oo_base_url" ]]; then
				BASEURL=$(cat "$HOME/.oo_base_url")
				if [[ $already_shown_oo_base_url_msg == 0 ]]; then
					if [[ -z $OO_MAIN_TESTS ]]; then
						yellow_text "$HOME/.oo_base_url exists. Using base-url $BASEURL as base url for receiving info on exit-code meaning."
						already_shown_oo_base_url_msg=1
					fi
				fi
			fi

			# .helpers.fetch_exit_code_help is the Python equivalent;
			# we use wget here so the lookup still works when myexit
			# fires before the venv has been sourced.
			set +e
			wget -q -O - "$BASEURL/exit_code_table.php?exit_code=$CODE"

			echo ""
			set -e
		fi

		exit "$CODE"
	}

	export NO_WHIPTAIL=1

	function displaytime {
		# Thin wrapper around .helpers.humanize_seconds.
		python3 "$SCRIPT_DIR/.helpers.py" humanize-seconds "$1" 2>/dev/null
	}

	function remaining_time {
		# Thin wrapper around .helpers.remaining_time.  Python handles
		# the date arithmetic and the "in about X years and Y days..."
		# formatting; bash strips the ANSI escapes from the input.
		local cleaned
		cleaned=$(echo "$1" | sed -E 's/\x1b\[[0-9;]*m//g')
		python3 "$SCRIPT_DIR/.helpers.py" remaining-time "$cleaned" 2>/dev/null
	}

	export CUDA_DEVICE_ORDER=PCI_BUS_ID
	ORIGINAL_PWD="$(pwd)"
	export ORIGINAL_PWD

	mkdir -p "$ORIGINAL_PWD/logs" || {
		red_text "Failed: mkdir -p $ORIGINAL_PWD/logs\n"
			myexit 45
		}

		set -e
		set -o pipefail

		function mycd {
			#echo "cd $1"
			cd "$1"
		}

		slurmlogpath () {
			# Thin wrapper around .helpers.slurmlogpath.  Python handles
			# the scontrol subprocess call and the StdOut= regex parse;
			# we keep the bash function so call sites stay unchanged.
			python3 "$SCRIPT_DIR/.helpers.py" slurm-log-path "$1" 2>/dev/null
		}

		function calltracer {
			exit_code=$?

			values=(
				"130"
				"138"
				"146"
			)

			if ! [[ " ${values[@]} " =~ " $exit_code" ]]; then
				LINE_AND_FUNCTION="$(caller)"
				if [[ "$LINE_AND_FUNCTION" != *"./omniopt"* ]] && [[ "$LINE_AND_FUNCTION" != *"./.tests/main_tests"* ]]; then
					red_text "Error occurred in file/line: $LINE_AND_FUNCTION\n"
				fi

				echo ""
				caller
				echo "Runtime (calltracer): $(displaytime "$SECONDS"), PID: $$"
			else
				echo ""
			fi

			_tput bel
		}

		already_sent_signal=
		kill_python_if_started_already_shown=0

		kill_python_if_started () {
			REASON="$1"
			echo "kill_python_if_started $REASON"
			re='^[0-9]+$'
			if [[ -n "$SLURM_JOB_ID" ]]; then
				if [[ $python_pid =~ $re ]] ; then
					if [[ -z "$already_sent_signal" ]]; then
						if command -v ps 2>/dev/null >/dev/null; then
							if ps auxf | grep "$python_pid" 2>/dev/null >/dev/null; then
								already_sent_signal=1
								echo -e "\nSending USR1 to $python_pid (python). Reason: $REASON"
								kill -USR1 "$python_pid"
							else
								echo "Could not find $python_pid process" >&2
							fi
						fi
					fi
				fi
			fi

			if [[ $kill_python_if_started_already_shown -eq 0 ]]; then
				echo "Runtime (kill_python_if_started): $(displaytime "$SECONDS"), PID: $$"
				kill_python_if_started_already_shown=1
			fi

			_tput bel
		}

		trap 'calltracer' ERR
		trap 'kill_python_if_started CONT' CONT
		trap 'kill_python_if_started TERM' TERM

		if command -v kill 2>/dev/null >/dev/null; then
			for i in $(kill -l 2>&1 | sed -e 's#[0-9][0-9]*[[:space:]]*)##g'); do
				if
					[[ "$i" != "ERR" ]] &&
						[[ "$i" != "CONT" ]] &&
						[[ "$i" != "TERM" ]] &&
						[[ "$i" != "CHLD" ]] &&
						[[ "$i" != "SIGCHLD" ]] &&
						[[ "$i" != "SIGWINCH" ]] &&
						[[ "$i" != "WINCH" ]] &&
						[[ "$i" != "INT" ]] &&
						[[ "$i" != "SIGINT" ]];
				then
					trap 'kill_python_if_started $i' "$i"
				fi
			done
		else
			red_text "kill cannot be found. Cannot register traps for existing signals.\n"
		fi

		minutes_to_hh_mm_ss() {
			# Thin wrapper around .helpers.minutes_to_hh_mm_ss.
			python3 "$SCRIPT_DIR/.helpers.py" minutes-to-hh-mm-ss "$1" 2>/dev/null || {
				red_text "ERROR: $1 is not a valid input. Must be a number of minutes (digits) or HH:MM:SS\n"
				myexit 103
			}
		}

		ORIG_ARGS=("$@")
		export ORIG_ARGS

		# Tiny bash-side flag extractors.  Python (.omniopt.py) does the
		# full argparse parse and config-loading; the bash script only
		# needs a handful of values to build the sbatch command and to
		# decide between local-exec and sbatch-submit.  These helpers only
		# look at the user-supplied CLI; they never consult config files.

		bash_has_flag() {
			local needle="$1"
			local arg
			for arg in "${ORIG_ARGS[@]}"; do
				if [[ "$arg" == "--${needle}" || "$arg" == "--${needle}="* ]]; then
					return 0
				fi
			done
			return 1
		}

		bash_opt_value() {
			local needle="$1"
			local prefix="--${needle}="
			local prefix_len=${#prefix}
			local i
			for ((i = 0; i < ${#ORIG_ARGS[@]}; i++)); do
				local arg="${ORIG_ARGS[$i]}"
				if [[ "${arg:0:prefix_len}" == "$prefix" ]]; then
					echo "${arg:prefix_len}"
					return 0
				elif [[ "$arg" == "--${needle}" ]]; then
					local next=$((i + 1))
					if (( next < ${#ORIG_ARGS[@]} )); then
						echo "${ORIG_ARGS[$next]}"
						return 0
					fi
				fi
			done
			return 1
		}

		# Bash-side copies of the flags we actually need for sbatch / dispatch.
		# Everything else is just forwarded to Python verbatim.
		if [[ -z $root_venv_dir ]]; then
			root_venv_dir=$HOME
		fi

	debug=0; bash_has_flag debug && debug=1
	dryrun=0; bash_has_flag dryrun && dryrun=1
	tests=0; bash_has_flag tests && tests=1
	run_tests_that_fail_on_taurus=0; bash_has_flag run_tests_that_fail_on_taurus && run_tests_that_fail_on_taurus=1
	checkout_to_latest_tested_version=0; bash_has_flag checkout_to_latest_tested_version && checkout_to_latest_tested_version=1
	follow=0; bash_has_flag follow && follow=1
	wait_until_ended=0; bash_has_flag wait_until_ended && wait_until_ended=1
	force_local_execution=0; bash_has_flag force_local_execution && force_local_execution=1
	help=0; bash_has_flag help && help=1

	# Mirror legacy CLI flags into the env vars the local-exec path
	# below looks at.  Python sees --flame_graph / --memray via
	# ORIG_ARGS, but bash needs the env-var form to decide whether
	# to wrap python3 in py-spy / memray.
	bash_has_flag flame_graph && export RUN_WITH_PYSPY=1
	bash_has_flag memray && export RUN_WITH_MEMRAY=1

		mem_gb=$(bash_opt_value mem_gb) || mem_gb=""
		gpus=$(bash_opt_value gpus) || gpus=""
		time=$(bash_opt_value time) || time=""
		experiment_name=$(bash_opt_value experiment_name) || experiment_name=""
		account=$(bash_opt_value account) || account=""
		reservation=$(bash_opt_value reservation) || reservation=""
		dependency=$(bash_opt_value dependency) || dependency=""
		num_cpus_main_job=$(bash_opt_value num_cpus_main_job) || num_cpus_main_job=""
		workdir=$(bash_opt_value workdir) || workdir=""
		worker_generator_path=$(bash_opt_value worker_generator_path) || worker_generator_path=""
		calculate_pareto_front_of_job=$(bash_opt_value calculate_pareto_front_of_job) || calculate_pareto_front_of_job=""
		# Allow both --continue and --continue_previous_job
		continue_previous_job=$(bash_opt_value continue_previous_job) || true
		continue_previous_job=${continue_previous_job:-$(bash_opt_value continue)}
		continue_previous_job=${continue_previous_job:-}
		show_ram_every_n_seconds=$(bash_opt_value show_ram_every_n_seconds) || show_ram_every_n_seconds=0
		main_process_gb=$(bash_opt_value main_process_gb) || main_process_gb=""

	if [[ $worker_generator_path != "" ]]; then
		old_uuid_dir="$worker_generator_path/state_files/run_uuid"
		if [[ -e $old_uuid_dir ]]; then
			RUN_UUID=$(cat $old_uuid_dir)
			if [[ -z $SLURM_JOB_ID ]]; then
				yellow_text "--worker_generator_path defined: $worker_generator_path. Using UUID from that run."
				yellow_text "Set RUN_UUID to $RUN_UUID"
			fi
		else
			red_text "Old UUID cannot be found"
		fi
	fi

	if [[ $dryrun -eq 1 ]]; then
		force_local_execution=1
		# Python's parse_arguments() also disables args.live_share when
		# --dryrun is set, so we don't need to mirror that here.
	fi

	python_pid=""

	# Source the venv-setup script early so subsequent python3 calls
	# in this script (git-version, etc.) use the venv interpreter
	# with all the heavy dependencies (pandas, matplotlib, etc.)
	# pre-installed.  RUN_VIA_RUNSH is also exported here, which
	# .helpers.py checks on import.
	if [[ -e "$SCRIPT_DIR/.shellscript_functions" ]]; then
		source "$SCRIPT_DIR/.shellscript_functions"
	else
		red_text "$SCRIPT_DIR/.shellscript_functions not found. Cannot continue.\n"
		myexit 23
	fi

	if [[ $show_ram_every_n_seconds -gt 0 ]]; then
		_show_ram_every_n_seconds $show_ram_every_n_seconds &
	fi

	# Backwards-compat: the old bash argparse clone validated that
	# --config_yaml/json/toml pointed at an existing file and exited
	# 100 if not.  Python now does the parse (exit 5 on missing file),
	# but a handful of tests still assert on the legacy exit 100.
	# Mirror the file-existence check here before forwarding.
	for cfg_opt in config_yaml config_json config_toml; do
		cfg_path=$(bash_opt_value "$cfg_opt") || continue
		if [[ ! -e $cfg_path ]]; then
			red_text "error: --$cfg_opt specified the path to a file that doesn't exist: $cfg_path" >&2
			exit 100
		fi
	done

	if [[ $use_git -eq 1 ]]; then
		if ! command -v git 2>/dev/null >/dev/null; then
			red_text "git not found. Cannot continue.\n"
			myexit 11
		fi

		# The git version display + checkout logic (~80 lines of bash)
		# moved into .helpers.resolve_git_version, which prints an action
		# line ("print" | "checkout" | "none") and the payload.
		if [[ -z "$SLURM_JOB_ID" ]]; then
			if [[ "$checkout_to_latest_tested_version" -eq "1" ]]; then
				pyout=$(python3 "$SCRIPT_DIR/.helpers.py" git-version "$SCRIPT_DIR" --checkout-to-latest 2>/dev/null)
			else
				pyout=$(python3 "$SCRIPT_DIR/.helpers.py" git-version "$SCRIPT_DIR" 2>/dev/null)
			fi

			action=$(echo "$pyout" | head -n1)
			payload=$(echo "$pyout" | tail -n +2)

			if [[ "$action" == "checkout" ]]; then
				yellow_text "--checkout_to_latest_tested_version enabled. Checking out to $payload..."
				git -C "$SCRIPT_DIR" checkout "$payload" >/dev/null 2>/dev/null || {
					red_text "\nFailed to checkout to latest version. Try not using --checkout_to_latest_tested_version.\n"
					myexit 211
				}

				bash omniopt "$@"
				exit_code=$?

				myexit $exit_code
			elif [[ "$action" == "print" ]]; then
				while IFS= read -r line; do
					yellow_text "$line"
				done <<< "$payload"
			fi
		fi
	fi

	# Defaults that the sbatch-submit path needs even when the user did not
	# pass the flag.  The Python entry point handles the equivalent
	# config-file and --continue state-file fallbacks itself.
	if [[ -z $gpus ]]; then
		gpus=0
	fi

	if [[ -z $main_process_gb ]]; then
		if [[ -n $mem_gb ]]; then
			main_process_gb=$mem_gb
		else
			main_process_gb=8 # default value
		fi
	fi


	if [[ $continue_previous_job =~ ^https?:// ]]; then
		# The actual wget/unzip dance used to live here (~95 lines of bash).
		# Python's .helpers.download_share_run now does it; we just detect
		# the URL pattern, call python, and re-launch with --continue set
		# to the freshly extracted local folder.  Runs after venv source
		# so `python3` is the venv interpreter that can import .helpers.
		if echo "$continue_previous_job" | grep -Eq 'share\?.*(user_id=[^&]+&.*experiment_name=[^&]+&.*run_nr=[0-9]+|experiment_name=[^&]+&.*user_id=[^&]+&.*run_nr=[0-9]+|run_nr=[0-9]+&.*user_id=[^&]+&.*experiment_name=[^&]+)'; then
			new_folder=$(python3 "$SCRIPT_DIR/.helpers.py" download-share-run "$continue_previous_job" 2>/dev/null)

			if [[ "$new_folder" == "" ]]; then
				red_text "Failed to download shared run from $continue_previous_job\n"
				myexit 19
			fi

			yellow_text "Extracted file to $new_folder"

			new_args=()
			replaced=0
			skip_next=0
			for arg in "${ORIG_ARGS[@]}"; do
				if [[ $skip_next -eq 1 ]]; then
					skip_next=0
					continue
				fi
				if [[ $replaced -eq 0 ]]; then
					if [[ "$arg" == "--continue" || "$arg" == "--continue_previous_job" ]]; then
						new_args+=("--continue")
						new_args+=("$new_folder")
						replaced=1
						skip_next=1
						continue
					elif [[ "$arg" == --continue=* || "$arg" == --continue_previous_job=* ]]; then
						new_args+=("--continue")
						new_args+=("$new_folder")
						replaced=1
						continue
					fi
				fi
				new_args+=("$arg")
			done

			bash "$SCRIPT_DIR/omniopt" "${new_args[@]}"
			exit_code=$?

			exit $exit_code
		else
			red_text "$continue_previous_job does not contain user_id, experiment_name, or run_nr\n"
			myexit 19
		fi
	fi

	mycd "$ORIGINAL_PWD"

	if [[ "$workdir" != "" ]]; then
		if [[ ! -d "$workdir" ]]; then
			mkdir -p "$workdir" || {
				red_text "$workdir could not be created. Cannot continue.\n"
				myexit 191
			}
		fi

		mycd "$workdir"
	fi

	function help_and_test_py {
		if command -v stdbuf 2>/dev/null >/dev/null; then
			if [[ -z $RUN_WITH_COVERAGE ]]; then
				stdbuf -e 0 -o 0 python3 "$@"
			else
				stdbuf -e 0 -o 0 coverage run -p "$@" --help
			fi
		else
			if [[ -z $RUN_WITH_COVERAGE ]]; then
				python3 "$@"
			else
				coverage run -p "$@" --help
			fi
		fi
	}

	if [[ "$help" -eq "1" ]]; then
		python3 "$SCRIPT_DIR/.omniopt.py" --help

		myexit 0
	fi

	if [[ "$tests" -eq "1" ]]; then
		exit_code=0
		if [[ $run_tests_that_fail_on_taurus -eq 0 ]]; then
			help_and_test_py "$SCRIPT_DIR/.omniopt.py" --tests --num_parallel_jobs=1 --max_eval=1 --worker_timeout=1 --run_program "" --experiment_name ""
			exit_code=$?
		else
			help_and_test_py "$SCRIPT_DIR/.omniopt.py" --tests --num_parallel_jobs=1 --max_eval=1 --worker_timeout=1 --run_program "" --experiment_name "" --run_tests_that_fail_on_taurus
			exit_code=$?
		fi
		myexit $exit_code
	fi

	kill_all_tail_child_processes() {
		for child in $(pgrep -P $$); do
			for tail_process in $(ps auxf | grep "$child" | grep tail | sed -e "s#^${USER}[[:space:]]*##" -e 's#[[:space:]].*##'); do
				kill -9 "$tail_process"
			done
		done
	}

	kill_tail_when_squeue_job_empty () {
		JOB_ID=$1

		if ! command -v pgrep 2>/dev/null >/dev/null; then
			return
		fi

		if ! command -v ps 2>/dev/null >/dev/null; then
			return
		fi

		unset_debug

		sleep 5

		while squeue -u "$USER" | grep "$JOB_ID" 2>/dev/null >/dev/null; do
			sleep 10
		done

		sleep 5

		if [[ $debug -eq 1 ]]; then
			set_debug
		fi

		kill_all_tail_child_processes

		return 0
	}

	if [ -n "${SLURM_JOB_ID:-}" ] || ! command -v sbatch >/dev/null || [[ $force_local_execution -eq 1 ]] || [[ $calculate_pareto_front_of_job != "" ]] ; then
		# To start all subjobs independently from the omniopt job, unset all SLURM variables
		for i in $(env | grep -e "^SLURM" | sed -e 's#[[:space:]]*=.*##' | grep -v SLURM_JOB_ID | grep -v SBATCH_RESERVATION); do
			unset "$i"
		done

		if [[ -n $SLURM_JOB_ID ]]; then
			echo -e "To cancel, press \033[1mCTRL\e[0m \033[1mc\e[0m, then run '\e[31mscancel $SLURM_JOB_ID\e[0m'"
		fi

		IFS=$' '
		export PYTHONPATH=$SCRIPT_DIR:$PYTHONPATH

		set +e

		if [[ $debug -eq 1 ]]; then
			echo "args: ${ORIG_ARGS[*]}"
		fi

		set +e
		trap - ERR

		# The Python side runs its own periodic live-share thread
		# (start_periodic_live_share in .omniopt.py), so the bash version
		# of this loop is no longer needed.

		if [[ -z $RUN_WITH_COVERAGE ]]; then
			if [[ -n $RUN_WITH_PYSPY ]]; then
				echo "Starting OmniOpt with Py-Spy"
				pip install py-spy
				stdbuf -e 0 -o 0 py-spy record --rate 10 --subprocesses --native --output "$RUN_UUID.svg" python3 -- "$SCRIPT_DIR/.omniopt.py" "${ORIG_ARGS[@]}"
				EXIT_CODE=$?
			elif [[ -n $RUN_WITH_MEMRAY ]]; then
				echo "Starting OmniOpt with MemRay."
				echo "Check later on with 'memray flamegraph $RUN_UUID.bin'"
				export PYTORCH_NO_CUDA_MEMORY_CACHING=1
				export PYTHONFAULTHANDLER=1
				pip install memray
				stdbuf -e 0 -o 0 memray run -o "$RUN_UUID.bin" -- "$SCRIPT_DIR/.omniopt.py" "${ORIG_ARGS[@]}"
				EXIT_CODE=$?
			else
				stdbuf -e 0 -o 0 python3 "$SCRIPT_DIR/.omniopt.py" "${ORIG_ARGS[@]}"
				EXIT_CODE=$?
			fi
		else
			echo "Using coverage run -p because \$RUN_WITH_COVERAGE is set"
			coverage run -p "$SCRIPT_DIR/.omniopt.py" "${ORIG_ARGS[@]}"
			EXIT_CODE=$?
		fi

		set -e
		trap 'calltracer' ERR

		set -e

		_tput bel

		myexit $EXIT_CODE
	else
		IFS=$' '

		formatted_time=$(minutes_to_hh_mm_ss "$time")

		sbatch_result=""
		exit_code=""

		sbatch_command="sbatch --mem=${main_process_gb}GB -N 1 --job-name $experiment_name --time=$formatted_time"

		if [[ $num_cpus_main_job != "" ]]; then
			sbatch_command+=" --cpus-per-task=$num_cpus_main_job "
		fi

		if [[ $gpus -ne 0 ]]; then
			sbatch_command+=" --gres=gpu:$gpus"
		fi

		if [[ -n $account ]]; then
			sbatch_command+=" --account=$account"
		fi

		if [[ -n $reservation ]]; then
			sbatch_command+=" --reservation=$reservation"
		fi

		if [[ -n $mem_gb ]]; then
			# Make sure --mem_gb is forwarded to the sbatch-launched omniopt.
			# Otherwise the user-provided value would be lost because we
			# build sbatch_command from CLI-derived vars, not from the
			# original CLI string.
			if ! bash_has_flag mem_gb; then
				ORIG_ARGS+=("--mem_gb=$mem_gb")
			fi
		fi

		if [[ $dependency == "omniopt_singleton" ]]; then
			k=0
			deps=""
			while read -r job_id; do
				if scontrol show job "$job_id" 2>/dev/null | grep -qi omniopt; then
					[[ -n $deps ]] && deps+=","
					deps+="after:$job_id"
					k=$((k+1))
				fi
			done < <(squeue -h -u "$USER" -o "%A")

			if (( k > 0 )); then
				dependency="$deps"
				echo "Automatic dependencies: $dependency"
			else
				dependency=""
			fi
		fi

		if [[ -n $dependency ]]; then
			if [[ $dependency =~ ^(afterok:[0-9]+|afternotok:[0-9]+|after:[0-9]+|afterany:[0-9]+|singleton)(,(afterok:[0-9]+|afternotok:[0-9]+|after:[0-9]+|afterany:[0-9]+|singleton))*$ ]]; then
				if [[ $debug -eq 1 ]]; then
					yellow_text "Valid dependency found: $dependency. Adding it it to sbatch command."
				fi
				sbatch_command="$sbatch_command --dependency $dependency"
			else
				echo "Invalid dependency format: $dependency"
			fi
		fi

		sbatch_command+=" $SCRIPT_DIR/omniopt"
		for arg in "${ORIG_ARGS[@]}"; do
			sbatch_command+=" $(printf '%q' "$arg")"
		done

		if [[ "$debug" -eq "1" ]] || [[ -n $PRINT_SBATCH_COMMAND ]]; then
			yellow_text "$sbatch_command"
		fi


		set +e

		sbatch_result=$($sbatch_command)
		exit_code=$?

		echo "$sbatch_result"

		set -e

		started_job_nr=$(echo "$sbatch_result" | sed -e 's#.*[[:space:]]##')

		if [[ $exit_code -eq 0 ]]; then
			if [[ $follow -eq 1 ]]; then
				if command -v sbatch 2>/dev/null >/dev/null; then
					set +e
					LOG_PATH=$(slurmlogpath "$started_job_nr" | tail -n1)

					spin[0]="⠇"
					spin[1]="⠏"
					spin[2]="⠋"
					spin[3]="⠙"
					spin[4]="⠹"
					spin[5]="⠸"
					spin[6]="⠴"
					spin[7]="⠦"
					spin[8]="⠧"

					last_why_pending_time=$(date +%s)
					last_sq_time=$(date +%s)

					estimated_start_time=""
					estimated_start_time_original=""
					_remaining_time=""

					if [[ $debug -eq 1 ]]; then
						trap - DEBUG
					fi

					if [[ -z $LOG_PATH ]]; then
						red_text "LOG_PATH was undefined. Is slurm installed properly?\n"
						myexit 110
					fi

					_tput civis # Disable cursor
					while [[ ! -e $LOG_PATH || ! -s $LOG_PATH ]]; do

						current_time=$(date +%s)

						time_diff_whypending=$(($current_time - $last_why_pending_time))
						time_diff_sq=$(($current_time - $last_sq_time))

						if command -v whypending 2>/dev/null > /dev/null && [[ $time_diff_whypending -gt 10 ]]; then
							trap - ERR
							estimated_start_time_original=$(timeout 5 whypending "$started_job_nr" 2>&1 | grep "Estimated" | sed -e 's#.*time:[[:space:]]*###' -e 's#^[[:space:]]*##')
							trap 'calltracer' ERR

							current_time=$(date +%s)
							last_why_pending_time=$current_time
							if [[ -n $estimated_start_time_original ]] && [[ "$estimated_start_time_original" != *"Unknown"* ]]; then
								estimated_start_time="- Estimated start:$estimated_start_time_original "
								if [[ -n $estimated_start_time_original ]]; then
									_remaining_time=""
									_remaining_time=$(remaining_time "$estimated_start_time_original" | sed -e 's#[[:space:]][[:space:]]*# #g')

									if [[ -n $_remaining_time ]]; then
										estimated_start_time="$estimated_start_time ($_remaining_time)"
									fi
								fi
							fi
						fi

						if command -v squeue 2>/dev/null > /dev/null && [[ $time_diff_sq -gt 60 ]]; then
							current_time=$(date +%s)
							last_sq_time=$current_time

							squeue_me_output=$(squeue --me 2>/dev/null)
							squeue_exit_code=$?

							if [[ $squeue_exit_code -eq "0" ]]; then
								job_still_in_squeue=$(echo "$squeue_me_output" | grep -c "$started_job_nr")

								if [[ "$job_still_in_squeue" -eq "0" ]]; then
									red_text "The job $started_job_nr was not found in squeue anymore. It seems like it has been cancelled.\n"

									SCONTROL_STATUS=$(scontrol show job "$started_job_nr" | grep JobState | sed -e 's#^[[:space:]]*[^=]*=#SLURM-Job-State: #')

									if [[ "$SCONTROL_STATUS" == *"FAILED"* ]]; then
										red_text "$SCONTROL_STATUS\n"
										if [[ "$SCONTROL_STATUS" == *"RaisedSignal:53"* ]]; then
											red_text "This may indicate a file system error\n"
										fi

										if command -v findmnt 2>/dev/null >/dev/null; then
											red_text "Mount-Info:\n"
											findmnt -T "$SCRIPT_DIR"
										fi
									fi

									myexit 243
								fi
							fi
						fi

						for spin_element in "${spin[@]}"; do
							print_line="$spin_element Waiting for slurm job $started_job_nr to be started (\e[4mtail -f $LOG_PATH\e[0m) $estimated_start_time"

							_tput cr # Move cursor to beginning of line
							_tput el # Delete line from start to finish

							echo -ne "$print_line"
							sleep 0.05
						done
					done

					_tput cnorm # Enable cursor

					if [[ $debug -eq 1 ]]; then
						set_debug
					fi

					set -e

					printf "\r"

					_tput el

					if [[ -e "$LOG_PATH" ]]; then
						kill_tail_when_squeue_job_empty "$started_job_nr" &

						tail_log_file() {
							#trap 'ask_cancel' SIGINT
							# weird exec stuff for disabling the "Terminated" message coming from kill
							exec 3>&2          # 3 is now a copy of 2
							exec 2> /dev/null  # 2 now points to /dev/null
							tail -n1000000 -f "$LOG_PATH" || true
							exec 2>&3          # restore stderr to saved
							exec 3>&-          # close saved version
						}

						ask_cancel() {
							trap 'tail_log_file' SIGINT
							echo ""
							echo ""
							red_text "Do you want to cancel just the tail (t) or the entire job (j), or cancel the cancelling (c)? "
							read -r answer
							case $answer in
								[Tt]*)
									cancelled_manually=1
									kill_all_tail_child_processes
									;;
								[Jj]*)
									cancelled_manually=1
									kill_all_tail_child_processes
									scancel "$started_job_nr"
									;;
								[Cc]*)
									kill_all_tail_child_processes
									tail_log_file
									;;
								*)
									kill_all_tail_child_processes
									tail_log_file
									;;
							esac
						}

						tail_log_file

						if [[ $already_logging_this_command -eq 0 ]]; then
							exec 1> >(tee -ia "$bash_logname")
							exec 2> >(tee -ia "$bash_logname" >& 2)
						fi


						exit_code_lines=$(grep -i "exit-code:*" "$LOG_PATH" 2>/dev/null)
						exit_code_lines_lines=$?
						if [ $exit_code_lines_lines -ne 0 ] || [ -z "$exit_code_lines" ]; then
							exit_code_lines=""
						fi

						exit_code_sed=$(echo "$exit_code_lines" | sed -e 's#Exit-Code:*[[:space:]]*##i' -e 's#,.*##')
						exit_code_sed_sed=$?
						if [ $exit_code_sed_sed -ne 0 ] || [ -z "$exit_code_sed" ]; then
							exit_code_sed=""
						fi

						exit_code_tail=$(echo "$exit_code_sed" | tail -n1)
						exit_code_tail_tail=$?
						if [ $exit_code_tail_tail -ne 0 ] || [ -z "$exit_code_tail" ]; then
							exit_code_tail=""
						fi

						exit_code_only_digits=$(echo "$exit_code_tail" | grep -o '[0-9]\+' | tail -n1)
						if [ -z "$exit_code_only_digits" ]; then
							exit_code_only_digits=3
						fi

						exit_code="$exit_code_only_digits"

						if ! [[ "$exit_code" =~ ^[0-9]+$ ]]; then
							exit_code=3
						fi

						if (( exit_code != 0 )); then
							if declare -f myexit > /dev/null; then
								myexit "$exit_code"
							else
								echo "WARN: myexit function not found, exiting with code $exit_code"
								exit "$exit_code"
							fi
						fi
					else
						red_text "$LOG_PATH could not be found\n"
					fi
				fi
			elif [[ "$wait_until_ended" -eq "1" ]]; then
				if command -v squeue 2>/dev/null >/dev/null; then
					WAIT_NUM_SECONDS=10
					yellow_text "Waiting for job $started_job_nr to end... (Checking every $WAIT_NUM_SECONDS seconds)"

					while [[ "$(squeue --me | grep -c "$started_job_nr")" -ne "0" ]]; do
						sleep $WAIT_NUM_SECONDS
					done

					yellow_text "Done waiting for job to end"

					LOG_PATH=$(slurmlogpath "$started_job_nr")
					if [[ -e "$LOG_PATH" ]]; then
						cat "$LOG_PATH"
					else
						red_text "$LOG_PATH could not be found\n"
					fi
				else
					red_text "squeue not found. Cannot wait for job to end.\n"
				fi
			fi
		else
			sbatch_command=$(echo "$sbatch_command" | sed -e 's#^[[:space:]]*##' -e 's#[[:space:]]*$##')
			red_text "Failed to start sbatch job. Command:\n"
			red_text "$sbatch_command\n"
			red_text "Exit-Code (sbatch): $exit_code\n"

			myexit $exit_code
		fi
	fi

	show_runtime=1

	if [[ $cancelled_manually -ne 0 ]]; then
		show_runtime=0
	fi

	if command -v sbatch 2>/dev/null >/dev/null; then
		if [[ -n $SLURM_JOB_ID ]]; then
			if [[ $force_local_execution -eq 0 ]]; then
				show_runtime=0
			fi
		else
			if [[ $follow -eq 0 ]]; then
				show_runtime=0
			fi
		fi
	fi

	if [[ $show_runtime -eq 0 ]]; then
		echo "Runtime (end): $(displaytime "$SECONDS"), PID: $$"
	fi

	if [[ -n $RUN_WITH_COVERAGE ]]; then
		echo "Run *coverage combine*, *coverage xml* and *coverage html*"
	fi

	if [[ "$exit_code" =~ ^[0-9]+$ ]]; then
		myexit "$exit_code"
	else
		if [[ $exit_code != "" ]]; then
			echo "Invalid exit-code >$exit_code< detected!"
		else
			echo "No exit-code could be found. Exiting with exit-code 3."
		fi

		myexit 3
	fi
}
