#!/usr/bin/env bash

{
	SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )

	source $SCRIPT_DIR/.colorfunctions.sh

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

	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}"' DEBUG
	}

	function unset_debug {
		trap - DEBUG
	}

	RUN_VIA_RUNSH=1
	export RUN_VIA_RUNSH

	ORIGINAL_WHIPTAIL=$WHIPTAIL

	ORIGINAL_PWD="$(pwd)"
	export ORIGINAL_PWD
	export MPLCONFIGDIR="/tmp/oo_matplotlib_cache_$USER"
	export XDG_CACHE_HOME="/tmp/XDG_CACHE_HOME_$USER"

	levenshtein() {
		local str1="$1"
		local str2="$2"
		local len1=${#str1}
		local len2=${#str2}
		local i j cost

		declare -A matrix
		for ((i=0; i<=len1; i++)); do
			matrix[$i,0]=$i
		done
		for ((j=0; j<=len2; j++)); do
			matrix[0,$j]=$j
		done

		for ((i=1; i<=len1; i++)); do
			for ((j=1; j<=len2; j++)); do
				if [[ "${str1:i-1:1}" == "${str2:j-1:1}" ]]; then
					cost=0
				else
					cost=1
				fi
				matrix[$i,$j]=$((
				$(min $((matrix[$((i-1)),$j]+1)) $((matrix[$i,$((j-1))]+1)) $((matrix[$((i-1)),$((j-1))]+cost)))
				))
			done
		done

		echo "${matrix[$len1,$len2]}"
	}

	min() {
		local min=$1
		for n in "$@"; do
			((n < min)) && min=$n
		done
		echo "$min"
	}

	find_closest_match() {
		local input="$1"
		local exact_matches=()
		local closest=""
		local closest_dist=-1
		local dist

		for expected in "${expected_plot_types[@]}"; do
			if [[ "$expected" == *"$input"* ]]; then
				exact_matches+=("$expected")
			fi
		done

		if [[ ${#exact_matches[@]} -eq 1 ]]; then
			echo "${exact_matches[0]}"
			return
		fi

		if [[ ${#exact_matches[@]} -gt 1 ]]; then
			echo "Multiple exact substring matches found: ${exact_matches[*]}" >&2
		else
			# No exact substring match found, proceed with Levenshtein
			echo "No exact substring match found for >$input<, proceeding with Levenshtein distance..." >&2
		fi

		for expected in "${expected_plot_types[@]}"; do
			dist=$(levenshtein "$input" "$expected")
			if [[ $closest_dist -eq -1 || $dist -lt $closest_dist ]]; then
				closest_dist=$dist
				closest="$expected"
			fi
		done

		echo "Closest match using Levenshtein: $closest (Distance: $closest_dist)" >&2

		echo "$closest"
	}

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

	_save_to_file=0

	function inputbox {
		TITLE=$1
		MSG=$2
		DEFAULT=$3

		eval "$(resize)"
		RESULT=$(whiptail --inputbox "$MSG" $LINES $COLUMNS "$DEFAULT" --title "$TITLE" 3>&1 1>&2 2>&3)
		exitstatus=$?
		if [[ $exitstatus == 0 ]]; then
			echo "$RESULT"
		else
			yellow_text "You chose to cancel (1)"
			exit 1
		fi
	}

	function ask_min_max {
		min=$(inputbox "Minimum value for plot" "Enter a Minimum value for plotting $run_dir (float), leave empty for no Minimum value" "")
		max=$(inputbox "Maximum value for plot" "Enter a Maximum value for plotting $run_dir (float), leave empty for no Maximum value" "")

		if [ -n "$min" ]; then
			# Check if $min is a number (integer or float, positive or negative)
			if [[ ! "$min" =~ ^-?[0-9]+(\.[0-9]+)?$ ]]; then
				echo "\$min is not a valid number. Setting it to empty."
				min=""
			fi
		fi

		if [ -n "$max" ]; then
			# Check if $max is a number (integer or float, positive or negative)
			if [[ ! "$max" =~ ^-?[0-9]+(\.[0-9]+)?$ ]]; then
				echo "\$max is not a valid number. Setting it to empty."
				max=""
			fi
		fi

		if [ -n "$min" ] && [ -n "$max" ]; then
			# Check if min is greater than max
			if (( $(echo "$max < $min" |bc -l) )); then
				# Swap values
				temp="$min"
				min="$max"
				max="$temp"
			else
				echo "min is less than or equal to max. No need to swap."
				true
			fi
		fi
	}

	function displaytime {
		local T=$1
		local D=$((T/60/60/24))
		local H=$((T/60/60%24))
		local M=$((T/60%60))
		local S=$((T%60))
		(( $D > 0 )) && printf '%d days ' $D
		(( $H > 0 )) && printf '%d hours ' $H
		(( $M > 0 )) && printf '%d minutes ' $M
		(( $D > 0 || $H > 0 || $M > 0 )) && printf 'and '
		printf '%d seconds\n' $S
	}

	expected_plot_types=("menu")
	whiptail_args=()

	for possible_plot_type in $(ls $SCRIPT_DIR/.omniopt_plot_*.py | sed -e 's#\.py##' -e 's#.*_plot_##' | tac); do
		expected_plot_types+=("$possible_plot_type")
	done

	set -e
	set -o pipefail

	function calltracer () {
		if [[ -z $NO_RUNTIME ]]; then
			echo 'Last file/last line:'
			caller

			echo "Runtime: $(displaytime $SECONDS)"
		fi
	}

	trap 'calltracer' ERR

	plot_type="menu"
	min=
	max=
	run_dir=""
	help=0

	args_string=""

	args=("$@")
	k=0

	while [ $k -lt ${#args[@]} ]; do
		arg="${args[k]}"

		case $arg in
			--run_dir=*)
				run_dir="${arg#*=}"
				args_string+=" --run_dir=$run_dir"
				shift
				;;

			--run_dir)
				k=$((k+1))
				if [ $k -ge ${#args[@]} ] || [[ "${args[k]}" == --* ]]; then
					red_text "Error: --run_dir requires a value"
					exit 1
				fi
				run_dir="${args[k]}"
				args_string+=" --run_dir=$run_dir"
				shift
				shift
				;;

			--save_to_file=*)
				_save_to_file="${arg#*=}"
				args_string+=" --save_to_file=$_save_to_file"
				shift
				;;

			--save_to_file)
				k=$((k+1))
				if [ $k -ge ${#args[@]} ] || [[ "${args[k]}" == --* ]]; then
					red_text "Error: --save_to_file requires a value"
					exit 1
				fi
				_save_to_file="${args[k]}"
				args_string+=" --save_to_file=$_save_to_file"
				shift
				shift
				;;

			--min=*)
				min="${arg#*=}"
				args_string+=" --min=$min"
				shift
				;;

			--min)
				k=$((k+1))
				if [ $k -ge ${#args[@]} ] || [[ "${args[k]}" == --* ]]; then
					red_text "Error: --min requires a value"
					exit 1
				fi
				min="${args[k]}"
				args_string+=" --min=$min"
				shift
				shift
				;;

			--max=*)
				max="${arg#*=}"
				args_string+=" --max=$max"
				shift
				;;

			--max)
				k=$((k+1))
				if [ $k -ge ${#args[@]} ] || [[ "${args[k]}" == --* ]]; then
					red_text "Error: --max requires a value"
					exit 1
				fi
				max="${args[k]}"
				args_string+=" --max=$max"
				shift
				shift
				;;

			--plot_type=*)
				plot_type="${arg#*=}"
				shift
				;;

			--plot_type)
				k=$((k+1))
				if [ $k -ge ${#args[@]} ] || [[ "${args[k]}" == --* ]]; then
					red_text "Error: --plot_type requires a value"
					exit 1
				fi
				plot_type="${args[k]}"
				shift
				shift
				;;

			--help|-h)
				help=1
				shift
				;;

			--debug)
				set_debug
				shift
				;;
			*)
				tmp_run_dir_or_plot_type="${args[k]}"

				if [[ -d $tmp_run_dir_or_plot_type ]]; then
					run_dir=$tmp_run_dir_or_plot_type
					args_string+=" --run_dir=$run_dir"
				else
					if [[ "$tmp_run_dir_or_plot_type" != --* ]]; then
						plot_type=$(find_closest_match "$tmp_run_dir_or_plot_type")
					fi
				fi

				;;
		esac
		k=$((k+1))
	done

	source "$SCRIPT_DIR/.shellscript_functions"
	source "$SCRIPT_DIR/.general.sh"

	if [[ "$help" -eq 1 ]]; then
		if [[ $plot_type != "menu" ]] && [[ -n $plot_type ]]; then
			python3 .omniopt_plot_$plot_type.py --help
		else
			echo "omniopt_plot: Plot omniopt runs"
			echo "Basic usage: bash omniopt_plot --run_dir=runs/testrun/0 --plot_type=scatter"
			echo "Possible options for plot_type:"
			echo "  - $(ls .omniopt_plot_*.py | sed -e 's#\.py##' -e 's#.*_plot_##' | tr '\n' ',' | sed -e 's#,$##' -e 's#,#\n  - #g')"
			echo "For specific options, use bash omniopt_plot --plot_type=scatter --help"
		fi

		exit 0
	fi

	if [[ ! " ${expected_plot_types[@]} " =~ " $plot_type " ]]; then
		joined_plot_types=$(printf "%s, " "${expected_plot_types[@]}")

		joined_plot_types=${joined_plot_types%, }

		red_text "Invalid plot type $plot_type, valid plot types: $joined_plot_types"
		exit 99
	fi

	if ! echo "$run_dir" | grep "^/" 2>/dev/null >/dev/null && [[ -n $run_dir ]]; then
		if [[ -d docker_user_dir ]]; then
			run_dir="$ORIGINAL_PWD/docker_user_dir/$run_dir"
		else
			run_dir="$ORIGINAL_PWD/$run_dir"
		fi
	fi

	function menu {
		whiptail_args=()
		for possible_plot_type in "${expected_plot_types[@]}"; do
			if [[ "$possible_plot_type" == "menu" ]]; then
				continue
			fi

			expected_files=()
			for expected_file in $(cat $SCRIPT_DIR/.omniopt_plot_${possible_plot_type}.py | grep "# EXPECTED FILES" | sed -e 's/# EXPECTED FILES: //'); do
				expected_files+=("$expected_file")
			done

			ALL_FILES_THERE=1

			for expected_file in "${expected_files[@]}"; do
				if [[ $(ls $run_dir | grep "$expected_file" | wc -l 2>/dev/null) -lt 1 ]]; then
					ALL_FILES_THERE=0
				fi
			done

			if [[ $ALL_FILES_THERE -eq 1 ]]; then
				trap '' ERR
				set +e
				num_not_disabled_min_param=$(grep args.min $SCRIPT_DIR/.omniopt_plot_${possible_plot_type}.py | wc -l)
				set -e
				trap 'calltracer' ERR

				accepts_min_max_string=""

				if [[ "$num_not_disabled_min_param" -ne "0" ]]; then
					accepts_min_max_string=", honors min/max"
				fi

				whiptail_args+=("$possible_plot_type" "$(cat $SCRIPT_DIR/.omniopt_plot_${possible_plot_type}.py | grep '# DESCRIPTION' | sed -e 's/#\s*DESCRIPTION: //')$accepts_min_max_string")

				if grep add_argument $SCRIPT_DIR/.omniopt_plot_${possible_plot_type}.py | grep save_to_file | grep -v useless 2>&1 >/dev/null; then
					whiptail_args+=("$possible_plot_type --save_to_file" "$(cat $SCRIPT_DIR/.omniopt_plot_${possible_plot_type}.py | grep '# DESCRIPTION' | sed -e 's/#\s*DESCRIPTION: //')$accepts_min_max_string")
				fi
			fi
		done

		if [ ${#whiptail_args[@]} -eq 0 ]; then
			red_text "It seems like the run folder $run_dir does not have any plotable data."
			exit 3
		fi

		minmaxstring=""
		if [[ -z $min ]] && [[ -z $max ]]; then
			#minmaxstring="Neither min nor max were set"
			true
		elif [[ -z $min ]] && [[ -n $max ]]; then
			minmaxstring=" (max: $max)"
		elif [[ -n $min ]] && [[ -z $max ]]; then
			minmaxstring=" (min: $min)"
		else
			minmaxstring=" (min: $min, max: $max)"
		fi

		if [[ -z $min ]] && [[ -z $max ]]; then
			#minmaxstring="Neither min nor max were set"
			true
		elif [[ -z $min ]] && [[ -n $max ]]; then
			minmaxstring="(max: $max)"
		elif [[ -n $min ]] && [[ -z $max ]]; then
			minmaxstring=" (min: $min)"
		else
			minmaxstring=" (min: $min, max: $max)"
		fi

		eval "$(resize)"
		WHATTODO=$(whiptail \
			--title "Available options for $run_dir" \
			--menu \
			"Chose what plot to open:" \
			$LINES $COLUMNS $(( $LINES - 8 )) \
			"${whiptail_args[@]}" \
			"minmax)" "set min/max values$minmaxstring" \
			"q)" "quit" 3>&1 1>&2 2>&3
		)

		if [[ "$WHATTODO" == "q)" ]]; then
			exit 0
		fi

		if [[ "$WHATTODO" == "minmax)" ]]; then
			ask_min_max

			if [[ -z $min ]] && [[ -z $max ]]; then
				#echo "Neither min nor max were set"
				bash "$SCRIPT_DIR/omniopt_plot" --run_dir "$run_dir" $args_string
			elif [[ -z $min ]] && [[ -n $max ]]; then
				echo "min was not set but max ($max)"
				bash "$SCRIPT_DIR/omniopt_plot" --run_dir "$run_dir" $args_string --max="$max "
				exit $?
			elif [[ -n $min ]] && [[ -z $max ]]; then
				echo "min ($min) was set but max was not"
				bash "$SCRIPT_DIR/omniopt_plot" --run_dir "$run_dir" $args_string --min="$min"
				exit $?
			else
				echo "min ($min) and max ($max) were set"
				bash "$SCRIPT_DIR/omniopt_plot" --run_dir "$run_dir" $args_string --max="$min" -max="$max"
				exit $?
			fi

			red_text "This exit should never be reached"
			exit
		fi

		plot_type=$WHATTODO
		if [[ "$plot_type" == *"save_to_file"* ]]; then
			plot_type=$(echo "$plot_type" | sed -e 's# .*##')
			if [[ "$_save_to_file" -eq "0" ]]; then
				_save_to_file=1
			fi
		fi

		if [[ "$_save_to_file" != "0" ]]; then
			_path=$(whiptail --inputbox "Path of the plot?" 8 39 "$run_dir/$plot_type.svg" --title "Choose path" 3>&1 1>&2 2>&3)

			exitstatus=$?
			if [ $exitstatus = 0 ]; then
				_save_to_file="$_path"
			else
				exit 0
			fi
		fi
	}

	if [[ "$plot_type" == "menu" ]]; then
		if [[ "$run_dir" == "" ]]; then
			red_text "--run_dir is missing"
			exit 1
		fi

		if [[ -f "$run_dir" ]]; then
			red_text "--run_dir is a file"
			exit 1
		fi

		if [[ ! -d "$run_dir" ]]; then
			red_text "--run_dir is not a directory"
			exit 1
		fi

		menu
	fi

	cd "$ORIGINAL_PWD"

	if [[ "$plot_type" != "menu" ]]; then
		if ! [[ -e "$SCRIPT_DIR/.omniopt_plot_$plot_type.py" ]]; then
			joined_plot_types=$(printf "%s, " "${expected_plot_types[@]}")

			joined_plot_types=${joined_plot_types%, }

			red_text "Invalid plot type $plot_type, valid plot types: $joined_plot_types"
			exit 5
		fi
	fi

	set +e
	export WHIPTAIL=1

	if [[ -z $RUN_WITH_COVERAGE ]]; then
		OUTPUT=$(eval "python3 $SCRIPT_DIR/.omniopt_plot_$plot_type.py $args_string" 2>&1)
		exit_code=$?
	else
		OUTPUT=$(eval "coverage run -p $SCRIPT_DIR/.omniopt_plot_$plot_type.py $args_string" 2>&1)
		exit_code=$?
	fi

	set -e

	if [[ "$exit_code" -ne "0" ]]; then
		if command -v whiptail 2>/dev/null >/dev/null; then
			if [[ $ORIGINAL_WHIPTAIL -eq 1 ]]; then
				error_message "$OUTPUT"
			else
				if [[ -n "$OUTPUT" ]]; then
					echo_red "$OUTPUT"
				fi
			fi
		else
			echo_red "$OUTPUT"
		fi
	else
		if [[ -n "$OUTPUT" ]]; then
			echo "$OUTPUT"
		fi
	fi

	if [[ -z $NO_RUNTIME ]]; then
		echo "Runtime: $(displaytime $SECONDS), plot_type: $plot_type, exit-code: $exit_code"
	fi

	exit $exit_code
}
