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

source "$PROJECT_ROOT/.venv/bin/activate"

cd "$PROJECT_ROOT"

# ===================================================================
# --- SCRIPT CONFIGURATION ---
# ===================================================================
# Statically chosen parameters
CENTER_LR=0.05 # 0.05
POPSIZE=4
ENV_NAME="MountainCarContinuous-v0"
NUM_JOBS=10
SEED_PER_EXP=25
export OMP_NUM_THREADS=$NUM_JOBS
# ===================================================================
# Parameters to loop over
LATENT_DIMS=(1)
REWARDS=("standard" "left" "speed" "height") #("speed" "rot_ccw" "rot_cw" "radial_speed")
STDEVS=(0.6)
SHAPES=("medium")
DATASET_SIZES=(10000)
# ===================================================================

# --- Main Experiment Loops ---
RUN_COUNT=0
for latent_dim in "${LATENT_DIMS[@]}"; do
  for reward in "${REWARDS[@]}"; do
    for stdev in "${STDEVS[@]}"; do
      SEED=100
      for shape in "${SHAPES[@]}"; do
        for dim_dataset in "${DATASET_SIZES[@]}"; do
          for ((i=1; i<=$SEED_PER_EXP; i++)); do
            ((RUN_COUNT++))
            echo ""
            echo "======================================================================"
            echo ">> RUN ${RUN_COUNT}: shape=${shape}, dataset=${dim_dataset}, seed=${SEED}"
            echo ">> SWEEP: reward=${reward}, stdev=${stdev}"
            echo "======================================================================"

            # Construct the source directory path from the autoencoder training
            SOURCE_DIR="trained_autoencoders_keep_10p_lr_0.0001_dim_${latent_dim}/${ENV_NAME}_${shape}_${dim_dataset}_seed_${SEED}"

            # Check if the source directory exists before running
            if [ ! -d "$SOURCE_DIR" ]; then
              echo "!! WARNING: Source directory not found, skipping: $SOURCE_DIR"
              continue
            fi
            
            # The command to run, with arguments separated by '--'
            python scripts/run_pgpe.py \
              --source_dir "${SOURCE_DIR}" \
              --popsize ${POPSIZE} \
              --num_envs ${POPSIZE} \
              --seed ${SEED} \
              --center_learning_rate ${CENTER_LR} \
              --stdev_init ${stdev} \
              --reward_type ${reward} \
              --stdev_learning_rate 0.1 \
              --eval_freq_steps 1000 \
              --num_generations 50 \
              --num_runs 1
            ((SEED++))
          done
        done
      done
    done
  done
done

echo ""
echo "All ${RUN_COUNT} experiments complete."