#!/bin/bash
# Copyright 2021,2022,2023 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

ROOT_DIR=$(echo $PWD | sed 's/^\(.*nnabla-rl\).*$/\1/')
export MKL_DEBUG_CPU_TYPE=5

# env list
ATARI_ENV_LIST=(
    "AlienNoFrameskip-v4"
    "AmidarNoFrameskip-v4"
    "AssaultNoFrameskip-v4"
    "AsterixNoFrameskip-v4"
    "AsteroidsNoFrameskip-v4"
    "AtlantisNoFrameskip-v4"
    "BankHeistNoFrameskip-v4"
    "BattleZoneNoFrameskip-v4"
    "BeamRiderNoFrameskip-v4"
    "BerzerkNoFrameskip-v4"
    "BowlingNoFrameskip-v4"
    "BoxingNoFrameskip-v4"
    "BreakoutNoFrameskip-v4"
    "CentipedeNoFrameskip-v4"
    "ChopperCommandNoFrameskip-v4"
    "CrazyClimberNoFrameskip-v4"
    "DemonAttackNoFrameskip-v4"
    "DoubleDunkNoFrameskip-v4"
    "EnduroNoFrameskip-v4"
    "FishingDerbyNoFrameskip-v4"
    "FreewayNoFrameskip-v4"
    "FrostbiteNoFrameskip-v4"
    "GopherNoFrameskip-v4"
    "GravitarNoFrameskip-v4"
    "HeroNoFrameskip-v4"
    "IceHockeyNoFrameskip-v4"
    "JamesbondNoFrameskip-v4"
    "KangarooNoFrameskip-v4"
    "KrullNoFrameskip-v4"
    "KungFuMasterNoFrameskip-v4"
    "MontezumaRevengeNoFrameskip-v4"
    "MsPacmanNoFrameskip-v4"
    "NameThisGameNoFrameskip-v4"
    "PhoenixNoFrameskip-v4"
    "PitfallNoFrameskip-v4"
    "PongNoFrameskip-v4"
    "PrivateEyeNoFrameskip-v4"
    "QbertNoFrameskip-v4"
    "RiverraidNoFrameskip-v4"
    "RoadRunnerNoFrameskip-v4"
    "RobotankNoFrameskip-v4"
    "SeaquestNoFrameskip-v4"
    "SkiingNoFrameskip-v4"
    "SolarisNoFrameskip-v4"
    "SpaceInvadersNoFrameskip-v4"
    "StarGunnerNoFrameskip-v4"
    "SurroundNoFrameskip-v4"
    "TennisNoFrameskip-v4"
    "TimePilotNoFrameskip-v4"
    "TutankhamNoFrameskip-v4"
    "UpNDownNoFrameskip-v4"
    "VentureNoFrameskip-v4"
    "VideoPinballNoFrameskip-v4"
    "WizardOfWorNoFrameskip-v4"
    "YarsRevengeNoFrameskip-v4"
    "ZaxxonNoFrameskip-v4"
)

MUJOCO_ENV_LIST=(
    "Ant-v2"
    "HalfCheetah-v2"
    "Hopper-v2"
    "Walker2d-v2"
    "Reacher-v2"
    "Swimmer-v2"
    "InvertedPendulum-v2"
    "InvertedDoublePendulum-v2"
    "ant-expert-v0"
    "halfcheetah-expert-v0"
    "hopper-expert-v0"
    "walker2d-expert-v0"
)

SPARSE_MUJOCO_ENV_LIST=(
    "SparseAnt-v1"
    "SparseHalfCheetah-v1"
    "SparseHopper-v1"
    "SparseWalker2d-v1"
)

DELAYED_MUJOCO_ENV_LIST=(
    "DelayedAnt-v1"
    "DelayedHalfCheetah-v1"
    "DelayedHopper-v1"
    "DelayedWalker2d-v1"
)

GPU_ID=0
ALGO_NAME="dqn"
BASE_ENV_NAME="atari"
ENV_START_INDEX=0
ENV_END_INDEX=-1
NUM_SEEDS=3

for opt in "$@"
do
    shift
    case "$opt" in
        "--gpu_id" )
            GPU_ID="$1";;
        "--algo_name" )
            ALGO_NAME="$1";;
        "--base_env_name" )
            BASE_ENV_NAME="$1";;
        "--env_start_index" )
            ENV_START_INDEX="$1";;
        "--env_end_index" )
            ENV_END_INDEX="$1";;
        "--envs" )
            ENVS="$*";;
        "--save_dir" )
            SAVE_DIR="$1";;
        "--num_seeds" )
            NUM_SEEDS="$1";;
        "--batch_size" )
            BATCH_SIZE="$1";;
        "-h" )
            echo ""
            echo "[Usage] : ./evaluate_algo
                --gpu_id <gpu_id>
                --algo_name <algo_name>
                --base_env_name <base_env_name>
                --env_start_index <env_start_index>
                --env_end_index <env_end_index>
                --envs <env1> <env2> <env3>
                --save_dir <save_dir>
                --batch_size <batch_size>

                e.g. ./evaluate_algo --algo_name dqn --base_env_name atari --env_start_index 0 --env_end_index 1 --save_dir sample
                or if you want directly specify the training envs
                e.g. ./evaluate_algo --algo_name sac --base_env_name mujoco --save_dir sample --envs Ant-v2 Hopper-v2

                -> Evaluation about dqn is executed in atari (AlienNoFrameskip-v4 AmidarNoFrameskip-v4).
                Please see this script for check env index.
                "
            echo ""
            exit 1;;
    esac
done

REPRODUCTION_CODE_DIR="${ROOT_DIR}/reproductions/algorithms/${BASE_ENV_NAME}/${ALGO_NAME}"
if [ -n "$SAVE_DIR" ]; then
    SAVE_DIR=$REPRODUCTION_CODE_DIR/$SAVE_DIR
else
    SAVE_DIR=$REPRODUCTION_CODE_DIR
fi

for ENV_NAME in ${ENVS}
do
    echo "Start running training for: " ${ENV_NAME}
    ${ROOT_DIR}/bin/train_with_seeds "${REPRODUCTION_CODE_DIR}/${ALGO_NAME}_reproduction.py" $GPU_ID $ENV_NAME $SAVE_DIR $NUM_SEEDS &
done

for INDEX in $(seq ${ENV_START_INDEX} ${ENV_END_INDEX})
do
    if [ $BASE_ENV_NAME = "atari" ];then
        ENV_NAME=${ATARI_ENV_LIST[$INDEX]}
    fi
    if [ $BASE_ENV_NAME = "mujoco" ]; then
        ENV_NAME=${MUJOCO_ENV_LIST[$INDEX]}
    fi
    if [ $BASE_ENV_NAME = "sparse_mujoco" ]; then
        ENV_NAME=${SPARSE_MUJOCO_ENV_LIST[$INDEX]}
    fi
    if [ $BASE_ENV_NAME = "delayed_mujoco" ]; then
        ENV_NAME=${DELAYED_MUJOCO_ENV_LIST[$INDEX]}
    fi
    echo "Start running training for: " ${ENV_NAME}
    if [ -n "$BATCH_SIZE" ]; then
        ${ROOT_DIR}/bin/train_with_seeds "${REPRODUCTION_CODE_DIR}/${ALGO_NAME}_reproduction.py" $GPU_ID $ENV_NAME $SAVE_DIR $NUM_SEEDS $BATCH_SIZE &
    else
        ${ROOT_DIR}/bin/train_with_seeds "${REPRODUCTION_CODE_DIR}/${ALGO_NAME}_reproduction.py" $GPU_ID $ENV_NAME $SAVE_DIR $NUM_SEEDS &
    fi
done
