#!/bin/bash

# Display a menu to the user
echo "Select the environment you want to run:"
echo "1. InvertedPendulum-v4"
echo "2. Walker2d-v4"
echo "3. Hopper-v4"
echo "4. HalfCheetah-v4"
echo "5. Ant-v4"

# Read the user's choice
read -p "Enter the number corresponding to your choice (1-5): " choice

# Map the choice to the environment
case $choice in
    1)
        selected_env="InvertedPendulum-v4"
        ;;
    2)
        selected_env="Walker2d-v4"
        ;;
    3)
        selected_env="Hopper-v4"
        ;;
    4)
        selected_env="HalfCheetah-v4"
        ;;
    5)
        selected_env="Ant-v4"
        ;;
    *)
        echo "Invalid choice. Exiting."
        exit 1
        ;;
esac

# Query the user for the training device
read -p "Please enter the training device: " device

# Save the selected environment in a variable
echo "You have selected: $selected_env on $device" 

# Example: Use the variable in your script
echo "Running in $selected_env environment..."

# If needed, you can export the variable to use it in other scripts
# export selected_env

# Source Anaconda Environment
. ~/anaconda3/etc/profile.d/conda.sh
conda activate dev

mkdir data_collection_logs

# Run agent training, then collect data using trained agent.
base_dir=data_collection_logs/$selected_env;
mkdir -p $base_dir;
echo "Collecting data for $selected_env environment. Logs are stored in $base_dir...";
python TD3/main.py --env $selected_env --device $device --save_model > $base_dir/expert_training.log;
python collect_dataset.py --env $selected_env --env_addons Normal --dataset_type train --device $device --dataset-size 100000 --seed 0 > $base_dir/train_normal.log;
python collect_dataset.py --env $selected_env --env_addons Process-Noise --dataset_type train --device $device --dataset-size 100000 --seed 0 > $base_dir/train_process_noise.log;
python collect_dataset.py --env $selected_env --env_addons Observation-Noise --dataset_type train --device $device --dataset-size 100000 --seed 0 > $base_dir/train_observation_noise.log;

python collect_dataset.py --env $selected_env --env_addons Normal --dataset_type test --device $device --dataset-size 100000 --seed 10 > $base_dir/test_normal.log;
python collect_dataset.py --env $selected_env --env_addons Process-Noise --dataset_type test --device $device --dataset-size 100000 --seed 10 > $base_dir/test_process_noise.log;
python collect_dataset.py --env $selected_env --env_addons Observation-Noise --dataset_type test --device $device --dataset-size 100000 --seed 10 > $base_dir/test_observation_noise.log;
