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"

# If the seed is the same then the policies will be the same.
# However, if the chunk size or num_envs is different the policies will be evaluated on different environments.
# The initial chunking is necessary to avoid memory issues when generating large numbers of large policies, and it can be adjusted as needed.
# Num_envs is used to control the batch size during evaluation, mainly used to batch policies on a GPU.
# 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
SKIP_EVAL=1
NUM_ENVS=256
LOAD_EXISTING=0
SHAPES=("small" "medium" "large")
NUM_POLICIES=(10000 50000 100000)
SEED_PER_EXP=3
SEED=0

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

    # --- Conditional Chunk Size Logic ---
    chunk_size=0
    if [[ "$shape" == "small" ]]; then
      # For "small", chunk_size is the minimum of 50000 or num_policies
      if (( num_pols < 50000 )); then
        chunk_size=$num_pols
      else
        chunk_size=50000
      fi
    elif [[ "$shape" == "medium" ]]; then
      chunk_size=10000
    elif [[ "$shape" == "large" ]]; then
      chunk_size=2500
    fi

    # --- Inner loop ---
    for ((i=1; i<=$SEED_PER_EXP; i++)); do
      echo "--> Running with Seed: $SEED, Chunk Size: $chunk_size"

      ARGS=(
        "--policy_shape" "$shape"
        "--num_policies" "$num_pols"
        "--chunk_size" "$chunk_size"
        "--seed" "$SEED"
        "--num_envs" "$NUM_ENVS"
        "--num_jobs" "$NUM_JOBS"
        
      )
      if (( SKIP_EVAL )); then
        ARGS+=("--skip_eval")
      fi
      if (( LOAD_EXISTING )); then
        ARGS+=("--load_existing")
      fi
      python scripts/generate_policies.py "${ARGS[@]}"

      # Increment the seed counter after each run
      ((SEED++))
    done
  done
done

echo "All policies generated."