#!/bin/bash

# brew install ffmpeg
export PYTHONPATH=.

algos=(magw)
suite=magw
envs=(staghunt)
config=magw
args1=($1)
args2=($2)
shift
shift
args3=(1e-3 5e-4 1e-4)
date=$(date +"%m%d")
info=""
ld="${suite}-logs"
# for x in $(seq 1 1); do
#     shift
# done

if [ -z "$info" ]; then
    n=$date
else
    n="$date-$info"
fi

gs=$1
if [ "$gs" = "gs" ]; then
    shift
    cmdkws=""
    for item in $@; do
        if [ "$cmdkws" = "" ]; then
            cmdkws="$item"
        else
            cmdkws="$cmdkws-$item"
        fi
    done
    echo "grid search"
    for e in "${envs[@]}"; do
        for a in "${algos[@]}"; do
            for a1 in "${args1[@]}"; do
                for a2 in "${args2[@]}"; do
                    for a3 in "${args3[@]}"; do
                        for s in {1..4}; do
                            kw1="n_epochs=$a1"
                            kw2="n_mbs=$a2"
                            kw3="lr=$a3"
                            if [ -z != $kw2 ]; then
                                kwstr="$kw1-$kw2-$kw3"
                                kws="$kw1 $kw2 $kw3"
                            else
                                kwstr="$kw1"
                                kws="$kw1"
                            fi
                            if [ "$cmdkws" != "" ]; then
                                kwstr="$cmdkws-$kwstr"
                                kws="$@ $kws"
                            fi
                            # kws="$@"
                            env="${suite}-${e}"
                            gpus=$(nvidia-smi -L | wc -l)
                            gpu=$(($s % $gpus))
                            # dir="$ld/$env/$a/$n-$kwstr/seed=$s"
                            # py_script="python run/train.py -a ${a} -e $env -c $config -ld $ld --gpu $gpu -n $n -s $s -kw $kws -i $info 2> ${dir}/err.txt | tee -a ${dir}/out.txt > /dev/null &"
                            py_script="python run/train.py -a ${a} -e $env -c $config -ld $ld -n $n -s $s --gpu $gpu -kw $kws -i $info &"
                            echo $py_script
                            eval $py_script
                        done
                    done
                done
            done
        done
    done
else
    kw1=""
    kwstr="$kw1"
    kws="$@"
    for item in $@; do
        if [ -z != $kwstr ]; then
            kwstr="$kwstr-$item"
        else
            kwstr="$item"
        fi
    done

    for e in "${envs[@]}"; do
        for a in "${algos[@]}"; do
            for s in {1..4}; do
                env="${suite}-${e}"
                n="$date"
                # if [ 0 != $# ]; then
                #     dir="$ld/$env/$a/$n-$kwstr/seed=$s"
                # else
                #     dir="$ld/$env/$a/$n/seed=$s"
                # fi
                # gpu=1
                # py_script="python run/train.py -a ${a} -e $env -c $config -ld $ld --gpu $gpu -n $n -s $s -kw $kws 2> ${dir}/err.txt | tee -a ${dir}/out.txt > /dev/null &"
                if [ -z != $kwstr ]; then
                    py_script="python run/train.py -a ${a} -e $env -c $config -ld $ld -n $n -s $s -kw $kws &"
                else
                    py_script="python run/train.py -a ${a} -e $env -c $config -ld $ld -n $n -s $s &"
                fi
                echo $py_script
                eval $py_script
            done
        done
    done
fi
wait
echo "Script completed"
