#!/usr/bin/env bash

set -xe

export PYTHONPATH=$(realpath repo/)

args_arr=(
  "--name GoToObj --n 8 --env BabyAI-GoToObj-v0"
  "--name GoToObjLocked --n 8 --env BabyAI-GoToObjLocked-v0"

  "--name GoToObjUnlocked-freq_1 --n 1 --env BabyAI-GoToObjUnlocked-v0"
  "--name GoToObjUnlocked-freq_2 --n 2 --env BabyAI-GoToObjUnlocked-v0"
  "--name GoToObjUnlocked-freq_4 --n 4 --env BabyAI-GoToObjUnlocked-v0"
  "--name GoToObjUnlocked --n 8 --env BabyAI-GoToObjUnlocked-v0"
  "--name GoToObjUnlocked-freq_32 --n 32 --env BabyAI-GoToObjUnlocked-v0"
  "--name GoToObjLocked_ambiguous-freq_1 --n 1 --env BabyAI-GoToObjLocked_ambiguous-v0"
  "--name GoToObjLocked_ambiguous-freq_2 --n 2 --env BabyAI-GoToObjLocked_ambiguous-v0"
  "--name GoToObjLocked_ambiguous-freq_4 --n 4 --env BabyAI-GoToObjLocked_ambiguous-v0"
  "--name GoToObjLocked_ambiguous --n 8 --env BabyAI-GoToObjLocked_ambiguous-v0"
  "--name GoToObjLocked_ambiguous-freq_32 --n 32 --env BabyAI-GoToObjLocked_ambiguous-v0"

  "--name GoToObjUnlocked-msg_4 --len-message 4 --num-symbols 4 --n 8 --env BabyAI-GoToObjUnlocked-v0"
  "--name GoToObjUnlocked-msg_16 --len-message 16 --num-symbols 16 --n 8 --env BabyAI-GoToObjUnlocked-v0"
  "--name GoToObjUnlocked-msg_32 --len-message 32 --num-symbols 32 --n 8 --env BabyAI-GoToObjUnlocked-v0"
  "--name GoToObjLocked_ambiguous-msg_4 --len-message 4 --num-symbols 4 --n 8 --env BabyAI-GoToObjLocked_ambiguous-v0"
  "--name GoToObjLocked_ambiguous-msg_16 --len-message 16 --num-symbols 16 --n 8 --env BabyAI-GoToObjLocked_ambiguous-v0"
  "--name GoToObjLocked_ambiguous-msg_32 --len-message 32 --num-symbols 32 --n 8 --env BabyAI-GoToObjLocked_ambiguous-v0"
)

function get_model_path {
  sr=$1
  mod=$(ls models | grep ".*ppo_expert.*$sr" | sort | tail -n1)
  echo -n $mod
}

for idx in $(seq 0 $((${#args_arr[@]} - 1))); do
  cd repo
    args=${args_arr[$idx]}
    # Remove the leading --name arg which is not in the train_rl.py script.
    train_args=$(echo $args | cut -d' ' -f3-)
    name=$(echo $args | cut -d' ' -f2)

    # Be careful about running this in parallel because the code does not use
    # assign unique names to configurations, so get_model_path might get the
    # wrong path if multiple instances of this are run in parallel.
    eval python scripts/train_rl.py \
      --seed 1 \
      --frames 5e7 \
      --tb \
      $train_args

    orig_sender=$(get_model_path sender)
    orig_receiver=$(get_model_path receiver)
    sender=$name-sender
    receiver=$name-receiver
    mv -n models/{$orig_sender,$sender}
    mv -n models/{$orig_receiver,$receiver}

    eval python ../helper.py \
      --episodes 5000 \
      --sender $sender \
      --receiver $receiver \
      $args
  cd ..
done
