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"

# Num_jobs is used to control the number of parallel workers used during evaluation. -3 means use all but 2 CPU cores.

NUM_JOBS=20
export OPENBLAS_NUM_THREADS=$NUM_JOBS
export OMP_NUM_THREADS=$NUM_JOBS

ENV_NAME="Reacher-v5"
SKIP_STATS=0
SHAPES=("medium_large")
NUM_POLICIES=(100000)
SEED_PER_EXP=10
SEED=0 

# --- Main Loops ---
for shape in "${SHAPES[@]}"; do
  for num_pols in "${NUM_POLICIES[@]}"; do
    echo ""
    echo "======================================================================"
    echo ">> Selecting Shape: $shape, Num Policies: $num_pols"
    echo "======================================================================"

    # --- Inner loop ---
    for ((i=1; i<=$SEED_PER_EXP; i++)); do
      echo "--> Running with Seed: $SEED"
      ARGS=(
          "--source_dir" "./generated_policies/${ENV_NAME}_${shape}_${num_pols}_seed_${SEED}"
          "--k_neighbors" "15"
          "--num_states" "3000"
          "--num_jobs" "$NUM_JOBS"
          "--seed" "$SEED"
        )

      if (( SKIP_STATS )); then
        ARGS+=("--skip_stats")
      fi

      # Execute the command with the array
      python ./scripts/rejection_sampling.py "${ARGS[@]}"
      ((SEED++))
    done
  done
done

echo "All policies selected."