#!/bin/bash

# Default values
TOTAL_TIME=""
DRY_RUN=false
SETS="0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15" # Default sets
INSTANCES="1 2 3 4 5"                          # Default instances
INSTANCES_DIR="../1995_bischoff_and_ratcliff_and_1999_davies_and_bischoff" # Default instance directory
USE_ONE_STAGE=true
CONFIGURATIONS_FILE="configurations_one_stage.txt"
PARAMETERS_FILE="parameters_one_stage.txt"
FIRST_TEST=""
EACH_TEST=""

# Function to display help message
usage() {
    echo "Usage: $0 [OPTIONS]"
    echo "Options:"
    echo "  -t, --total-time       Total time in seconds (required)"
    echo "  --sets <values>        List of sets to include (default: '0 1 2 ... 15')"
    echo "  --instances <values>   List of instances to include (default: '1 2 3 4 5')"
    echo "  --instances-dir <dir>  Directory containing the instance files (default: '../1995_bischoff_and_ratcliff_and_1999_davies_and_bischoff')"
    echo "  --dry-run              Print the commands that would be executed without running them"
    echo "  -h, --help             Display this help message"
}

# Parse command-line arguments
while [ "$#" -gt 0 ]; do
    case "$1" in
        -t|--total-time)
            TOTAL_TIME="$2"
            shift 2
            ;;
        --sets)
            shift
            SETS=""
            while [ "$#" -gt 0 ] && [ "${1#--}" = "$1" ]; do
                SETS="$SETS $1"
                shift
            done
            ;;
        --instances)
            shift
            INSTANCES=""
            while [ "$#" -gt 0 ] && [ "${1#--}" = "$1" ]; do
                INSTANCES="$INSTANCES $1"
                shift
            done
            ;;
        --instances-dir)
            INSTANCES_DIR="$2"
            shift 2
            ;;
        --one-stage)
            USE_ONE_STAGE=true
            shift 1
            ;;
        --two-stage)
            USE_ONE_STAGE=false
            shift 1
            ;;
        --first-test)
            FIRST_TEST="$2"
            shift 2
            ;;
        --each-test)
            EACH_TEST="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN=true
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Ensure TOTAL_TIME is set
if [ -z "$TOTAL_TIME" ]; then
    echo "Error: --total-time is required."
    usage
    exit 1
fi

# Validate the instances directory
if [ ! -d "$INSTANCES_DIR" ]; then
    echo "Error: Instances directory '$INSTANCES_DIR' does not exist."
    exit 1
fi

# Create train_instances directory
TRAIN_INSTANCES_DIR="./train_instances"
if [ -d "$TRAIN_INSTANCES_DIR" ]; then
    echo "Instances directory '$TRAIN_INSTANCES_DIR' already exists, deleting it."
    rm -rf "$TRAIN_INSTANCES_DIR"
fi
mkdir -p "$TRAIN_INSTANCES_DIR"

# Copy selected instances
for SET in $SETS; do
    for INSTANCE in $INSTANCES; do
        SOURCE_FILE="${INSTANCES_DIR}/set_${SET}/${INSTANCE}.json"
        DEST_FILE="${TRAIN_INSTANCES_DIR}/set_${SET}_instance_${INSTANCE}.json"

        if [ -f "$SOURCE_FILE" ]; then
            cp "$SOURCE_FILE" "$DEST_FILE"
        else
            echo "Warning: File $SOURCE_FILE does not exist. Skipping."
        fi
    done
done

# Count the number of training instances
NUM_TRAIN_INSTANCES=$(find "$TRAIN_INSTANCES_DIR" -type f -name '*.json' | wc -l)
if [ "$NUM_TRAIN_INSTANCES" -eq 0 ]; then
    echo "Error: No training instances found in '$TRAIN_INSTANCES_DIR'."
    exit 1
fi

if [ -z $FIRST_TEST ]; then
    FIRST_TEST=$NUM_TRAIN_INSTANCES
fi
FIRST_TEST=$(echo $((FIRST_TEST>NUM_TRAIN_INSTANCES ? NUM_TRAIN_INSTANCES : FIRST_TEST)))

if [ -z $EACH_TEST ]; then
    EACH_TEST=$NUM_TRAIN_INSTANCES
fi
EACH_TEST=$(echo $((EACH_TEST>FIRST_TEST ? FIRST_TEST : EACH_TEST)))
EACH_TEST=$(echo $((EACH_TEST>NUM_TRAIN_INSTANCES ? NUM_TRAIN_INSTANCES : EACH_TEST)))

if [ "$USE_ONE_STAGE" = true ]; then
    CONFIGURATIONS_FILE="configurations_one_stage.txt"
    PARAMETERS_FILE="parameters_one_stage.txt"
else
    CONFIGURATIONS_FILE="configurations_two_stage.txt"
    PARAMETERS_FILE="parameters_two_stage.txt"
fi

# Compute derived values
N_THREADS=$(nproc --all)
IRACE_MAX_TIME=$(awk "BEGIN {print ${TOTAL_TIME} * ${N_THREADS}}")
BUDGET_ESTIMATE=$(awk "BEGIN {print 120 / $IRACE_MAX_TIME}")

# Construct the irace command
IRACE_CMD="irace \
    --parallel=$N_THREADS \
    --scenario=./scenario.txt \
    --configurations-file=$CONFIGURATIONS_FILE \
    --parameter-file=$PARAMETERS_FILE \
    --train-instances-dir=$TRAIN_INSTANCES_DIR \
    --max-time=$IRACE_MAX_TIME \
    --budget-estimation=$BUDGET_ESTIMATE \
    --first-test=$FIRST_TEST \
    --each-test=$EACH_TEST"

# Handle dry-run
if $DRY_RUN; then
    echo "Dry-run mode enabled. Commands to execute:"
    echo "  $IRACE_CMD | tee irace.log"
else
    # Execute the command
    eval "$IRACE_CMD | tee irace.log"
fi
