#!/bin/bash

# Parameter check
if [[ $# -ne 7 ]]; then
    echo "Usage: $0 <run_config_module> <partition> <gpu_type> <num_gpus> <num_cpus_per_gpu> <item_limit> <time_limit>"
    exit 1
fi

# Assign parameters
run_config=$1
partition=$2
gpu_type=$3
num_gpus=$4
num_cpus_per_gpu=$5
item_limit=$6
time_limit=$7

# Validate parameters
if [[ ! -f "$run_config" ]]; then
    echo "Error: Run configuration file '$run_config' does not exist."
    exit 1
fi
if [[ ! "$partition" =~ ^(clara|paula)$ ]]; then
    echo "Error: Invalid partition '$partition'. Must be 'clara' or 'paula'."
    exit 1
fi

valid_gpus=("a30" "v100" "rtx2080ti")
if [[ ! " ${valid_gpus[@]} " =~ " $gpu_type " ]]; then
    echo "Error: Invalid GPU type '$gpu_type'. Must be one of: ${valid_gpus[*]}."
    exit 1
fi

if ! [[ "$num_gpus" =~ ^[0-9]+$ ]] || ! [[ "$num_cpus_per_gpu" =~ ^[0-9]+$ ]]; then
    echo "Error: Number of GPUs and CPUs per GPU must be positive integers."
    exit 1
fi

# item limit must be an integer, either positive or -1
if ! [[ "$item_limit" =~ ^-?[0-9]+$ ]]; then
    echo "Error: Item limit must be an integer."
    exit 1
fi
if [[ "$item_limit" -lt -1 ]]; then
    echo "Error: Item limit must be -1 or a positive integer."
    exit 1
fi


# Derive run name
run_name=$(basename "$run_config")
run_name="${run_name%.*}"
ram_per_gpu=16
num_cpus=$(($num_gpus * $num_cpus_per_gpu))
ram=$(($num_gpus * $ram_per_gpu))

# Confirm inputs
echo "Starting run with configuration: $run_config"
echo "Partition: $partition"
echo "GPU type: $gpu_type"
echo "Number of GPUs: $num_gpus"
echo "Number of CPUs per GPU: $num_cpus_per_gpu"
echo "Total CPUs: $num_cpus"
echo "Total RAM: $ram GB"
echo "Item limit: $item_limit"

# Set paths
current_dir_path=$(dirname "$(realpath "$0")")
venv_path=$(realpath "$current_dir_path/../../venv/bin/activate")

# Submit job
sbatch <<EOF
#!/bin/bash
#SBATCH --job-name=$run_name
#SBATCH --output=/work/rr41qemu-MA/logs/%j.out
#SBATCH --error=/work/rr41qemu-MA/logs/%j.err
#SBATCH --time=$time_limit
#SBATCH --ntasks=1
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=$num_cpus
#SBATCH --mem=${ram}G
#SBATCH --partition=$partition
#SBATCH --gpus=$gpu_type:$num_gpus

echo "Loading python virtual environment..."
source $venv_path

echo "Loading python 3.10..."
module load Python/3.10.4-GCCcore-11.3.0

# create log directory
#log_dir="/work/rr41qemu-MA/logs/${SLURM_JOB_ID}_$(date +%Y-%m-%d_%H-%M-%S)"
#mkdir -p "$log_dir"
#echo "Logging to $log_dir"

cd $current_dir_path

torchrun --nproc_per_node=$num_gpus training_wrapper.py $run_config --item_limit $item_limit
EOF