#!/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.

set -e

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

# parameters
SEED=1
TOTAL_ITERATIONS=1
SHOWCASE_RUNS=1
GPU_ID=-1
BASE_ENV="atari"
ENV="FakeAtariNNablaRLNoFrameskip-v1"

for opt in "$@"
do
    shift
    case "$opt" in
        "--gpu_id" )
            GPU_ID="$1";;
        "--base_env" )
            BASE_ENV="$1";;
        "--env" )
            ENV="$1";;
        "-h" )
            echo ""
            echo "[Usage] : ./test_reproduction
                --gpu_id <gpu_id>
                --base_env <base_env_name>
                --env <env_name>
                
                e.g.  sh bin/test_reproductions.sh --gpu_id 0 --base_env atari --env BreakoutNoFrameskip-v4 
                "
            echo ""
            exit 1;;
    esac
done

REPRODUCTIONS_DIRS=$(find ${ROOT_DIR}/reproductions/algorithms/${BASE_ENV}/* -maxdepth 0 -type d)
RESULT_BASE_DIR="${ROOT_DIR}/test_reproduction_results"

for dir in $REPRODUCTIONS_DIRS;
do
    ALGORITHM=$(basename ${dir})
    SCRIPT="${dir}/${ALGORITHM}_reproduction.py"
    if [ ${BASE_ENV} == "sparse_mujoco" ] && [ ${ALGORITHM} == "environment" ]; then
        # Do nothing this is not a reproduction code
        continue
    fi
    if [ ${BASE_ENV} == "delayed_mujoco" ] && [ ${ALGORITHM} == "environment" ]; then
        # Do nothing this is not a reproduction code
        continue
    fi
    if [ ${BASE_ENV} = "mujoco" ] && [ ${ALGORITHM} = "her" ]; then
        TMP_ENV="FakeGoalConditionedNNablaRL-v1"
        echo "Test run of training for ${ALGORITHM}"
        python ${SCRIPT} --gpu ${GPU_ID} --env ${TMP_ENV} --save-dir "${RESULT_BASE_DIR}/${ALGORITHM}" --seed ${SEED} \
            --total_iterations ${TOTAL_ITERATIONS} --save_timing ${TOTAL_ITERATIONS}
        SNAPSHOT_DIR="${RESULT_BASE_DIR}/${ALGORITHM}/${TMP_ENV}_results/seed-${SEED}/iteration-${TOTAL_ITERATIONS}"
        echo "Test run of showcase for ${ALGORITHM}"
        python ${SCRIPT} --gpu ${GPU_ID} --env ${TMP_ENV} --snapshot-dir ${SNAPSHOT_DIR} --showcase \
            --showcase_runs ${SHOWCASE_RUNS}
    elif [ ${BASE_ENV} = "pybullet" ] && [ ${ALGORITHM} = "icra2018qtopt" ]; then
        TMP_ENV="KukaGraspingProceduralEnv"
        echo "Test run of training for ${ALGORITHM}"
        python ${SCRIPT} --gpu ${GPU_ID} --save-dir "${RESULT_BASE_DIR}/${ALGORITHM}" --seed ${SEED} --total_iterations ${TOTAL_ITERATIONS} --save_timing ${TOTAL_ITERATIONS} --num_collection_episodes 1 \
        --batch_size 1 
        SNAPSHOT_DIR="${RESULT_BASE_DIR}/${ALGORITHM}/${TMP_ENV}_1_results/seed-${SEED}/iteration-${TOTAL_ITERATIONS}"
        echo "Test run of showcase for ${ALGORITHM}"
        python ${SCRIPT} --gpu ${GPU_ID} --snapshot-dir ${SNAPSHOT_DIR} --showcase \
            --showcase_runs ${SHOWCASE_RUNS}
    elif [ ${ALGORITHM} = "decision_transformer" ]; then
        echo "Test run of training for ${ALGORITHM}"
        TOTAL_EPOCHS=1
        BUFFER_SIZE=20
        BATCH_SIZE=10
        CONTEXT_LENGTH=2
        TRAJECTORIES_PER_BUFFER=1
        TARGET_RETURN=1
        if [ ${BASE_ENV} = "atari" ]; then    
            DATASET_PATH="${ROOT_DIR}/test_resources/reproductions/atari-dataset"
            WARMUP_TOKENS=1
            python ${SCRIPT} --gpu ${GPU_ID} --env ${ENV} --save-dir "${RESULT_BASE_DIR}/${ALGORITHM}" --seed ${SEED} \
                --total-epochs ${TOTAL_EPOCHS} --save_timing ${TOTAL_EPOCHS} --buffer-size ${BUFFER_SIZE} \
                --dataset-path ${DATASET_PATH} --batch-size ${BATCH_SIZE} --context-length ${CONTEXT_LENGTH} \
                --warmup-tokens ${WARMUP_TOKENS} --trajectories-per-buffer ${TRAJECTORIES_PER_BUFFER} \
                --target-return ${TARGET_RETURN}
            SNAPSHOT_DIR="${RESULT_BASE_DIR}/${ALGORITHM}/${ENV}_results/seed-${SEED}/iteration-${TOTAL_ITERATIONS}"
            echo "Test run of showcase for ${ALGORITHM}"
            python ${SCRIPT} --gpu ${GPU_ID} --env ${ENV} --snapshot-dir ${SNAPSHOT_DIR} --showcase \
                --showcase_runs ${SHOWCASE_RUNS}
        fi
        if [ ${BASE_ENV} = "mujoco" ]; then    
            DATASET_PATH="${ROOT_DIR}/test_resources/reproductions/mujoco-dataset"
            WARMUP_STEPS=1
            python ${SCRIPT} --gpu ${GPU_ID} --env ${ENV} --save-dir "${RESULT_BASE_DIR}/${ALGORITHM}" --seed ${SEED} \
                --total-epochs ${TOTAL_EPOCHS} --save_timing ${TOTAL_EPOCHS} --buffer-size ${BUFFER_SIZE} \
                --dataset-path ${DATASET_PATH} --batch-size ${BATCH_SIZE} --context-length ${CONTEXT_LENGTH} \
                --warmup-steps ${WARMUP_STEPS} --trajectories-per-buffer ${TRAJECTORIES_PER_BUFFER} \
                --target-return ${TARGET_RETURN}
            SNAPSHOT_DIR="${RESULT_BASE_DIR}/${ALGORITHM}/${ENV}_results/seed-${SEED}/iteration-${TOTAL_ITERATIONS}"
            echo "Test run of showcase for ${ALGORITHM}"
            python ${SCRIPT} --gpu ${GPU_ID} --env ${ENV} --snapshot-dir ${SNAPSHOT_DIR} --showcase \
                --showcase_runs ${SHOWCASE_RUNS} --dataset-path ${DATASET_PATH}
        fi
    else
        echo "Test run of training for ${ALGORITHM}"
        python ${SCRIPT} --gpu ${GPU_ID} --env ${ENV} --save-dir "${RESULT_BASE_DIR}/${ALGORITHM}" --seed ${SEED} \
            --total_iterations ${TOTAL_ITERATIONS} --save_timing ${TOTAL_ITERATIONS}
        SNAPSHOT_DIR="${RESULT_BASE_DIR}/${ALGORITHM}/${ENV}_results/seed-${SEED}/iteration-${TOTAL_ITERATIONS}"
        echo "Test run of showcase for ${ALGORITHM}"
        python ${SCRIPT} --gpu ${GPU_ID} --env ${ENV} --snapshot-dir ${SNAPSHOT_DIR} --showcase \
            --showcase_runs ${SHOWCASE_RUNS}
    fi
done
