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
PERCENTILE=0.1
export OMP_NUM_THREADS=$NUM_JOBS

ENV_NAME="Reacher-v5"

SHAPES=("medium_large")
NUM_POLICIES=(100000)
LATENT_DIM=(8)
LEARNING_RATE=(0.0001)
SEED_PER_EXP=10

# --- Main Loops ---
for latent_dim in "${LATENT_DIM[@]}"; do
  for learning_rate in "${LEARNING_RATE[@]}"; do
    SEED=0
    if [ "$latent_dim" -eq 1 ]; then
      NUM_POINTS=100
      NUM_ENVS=100
    elif [ "$latent_dim" -eq 2 ]; then
      NUM_POINTS=50
      NUM_ENVS=500
    elif [ "$latent_dim" -eq 3 ]; then
      NUM_POINTS=17
      NUM_ENVS=500
    elif [ "$latent_dim" -eq 5 ]; then
      NUM_POINTS=5
      NUM_ENVS=500
    elif [ "$latent_dim" -eq 8 ]; then
      NUM_POINTS=3
      NUM_ENVS=500
    fi
    for shape in "${SHAPES[@]}"; do
      for num_pols in "${NUM_POLICIES[@]}"; do
        echo ""
        echo "======================================================================"
        echo ">> Training Shape: $shape, Num Policies: $num_pols, Latent Dim: $latent_dim, Learning Rate: $learning_rate, Percentile: $PERCENTILE"
        echo "======================================================================"
        # --- Inner loop ---
        for ((i=1; i<=$SEED_PER_EXP; i++)); do
          echo "--> Running with Seed: $SEED"
          python scripts/train_autoencoder.py \
            --source_dir ./selected_policies/"$ENV_NAME"_"$shape"_"$num_pols"_seed_"$SEED" \
            --seed $SEED \
            --learning_rate "$learning_rate" \
            --latent_dim $latent_dim \
            --epochs 50 \
            --percentile "$PERCENTILE" \
            --use_saved_states
          echo "Testing latent space"
          per=$(awk -v p="$PERCENTILE" 'BEGIN { printf "%d\n", p * 100 }')
          python scripts/test_latent.py \
            --source_dir ./trained_autoencoders_keep_"$per"p_lr_"$learning_rate"_dim_"$latent_dim"/"$ENV_NAME"_"$shape"_"$num_pols"_seed_"$SEED" \
            --seed $SEED \
            --num_points $NUM_POINTS \
            --num_envs  $NUM_ENVS \
            --num_jobs $NUM_JOBS
          ((SEED++))
        done
      done
    done
  done
done