#!/bin/bash

seeds_train=(100)
device=cuda:0
penc=False
tenc_dim=64
tenc_scale=100
num_demos=25
batch_size=128
use_wandb=False

run_cls=0 # Global feature
run_cls_tec=0 # Global feature + TENC
run_afa=0 # Per-patch feature + AFA
run_afa_tec=1 # Per-patch feature + AFA + TENC

model_names=(
    #! Vision Transformers
    # vc1-vit-b16
    # vit_base_patch16_224.mae
    vit_base_patch16_224.dino
    # vit_base_patch16_224.orig_in21k
    # vit_base_patch16_clip_224.laion2b
    # vit_base_patch14_dinov2.lvd142m
    # ibot-vit-b16
    # vjepa-vit-l16
    #! ResNets
    # vip-resnet50_VIPfc
    # r3m-resnet50 
    # swav-resnet50
    # vfs-resnet50 
    # densecl-resnet50
    # mocov2-resnet50   
    # vicregl-resnet50
)

tasks=(
    # box-close-v2
    bin_picking 
    # disassemble-v2
    # coffee_pull
    # shelf-place-v2
    # peg-insert-side-v2
    # assembly-v2
    # drawer-open-v2
    # pick-place-wall-v2
    # button-press-wall-v2
)

for seed in "${seeds_train[@]}"; do
    for model_name in "${model_names[@]}"; do
        for task in "${tasks[@]}"; do

            # Format task name if needed
            case $task in
                bin_picking)    dir_name="bin-picking-v2" ;;
                drawer_close)   dir_name="drawer-close-v2" ;;
                coffee_pull)    dir_name="coffee-pull-v2" ;;
                door_open)      dir_name="door-open-v2" ;;
                lever_pull)     dir_name="lever-pull-v2" ;;
                hammer)         dir_name="hammer-v2" ;;
                *)              dir_name=$task ;;
            esac

            [[ "$model_name" == *"resnet"* ]] && num_heads=32 || num_heads=12

            run_exp() {
                local feat_extraction=$1
                local use_tenc=$2
                local suffix="d${tenc_dim}_s${tenc_scale}"
                local file_path="./exp_local/${dir_name}_${model_name}_batch_${batch_size}_seed_${seed}_${feat_extraction}_use_proprio_True_use_tenc_${use_tenc}_${suffix}_supervision_False_None/eval.csv"

                echo "Checking file path: $file_path"
                echo "======================"

                if [ ! -f "$file_path_1" ]; then
                    python train_bc.py \
                        agent=bc \
                        suite=metaworld \
                        suite/metaworld_task=$task \
                        suite.num_eval_episodes=100 \
                        suite.eval_every_frames=80000 \
                        suite.num_train_frames=80000 \
                        agent.backbone=$backbone \
                        agent.embedding_name=$model_name \
                        agent.feat_extraction=$feat_extraction \
                        agent.use_proprio=True \
                        agent.supervise_proprio=False \
                        agent.supervision_mode=None \
                        agent.use_tenc=$use_tenc \
                        agent.tenc_dim=$tenc_dim \
                        agent.tenc_scale=$tenc_scale \
                        agent.positional_encoding=$penc \
                        agent.num_heads=$num_heads \
                        num_demos=$num_demos \
                        use_wandb=$use_wandb \
                        seed=$seed \
                        exp_prefix=BC \
                        device=$device
                fi
            }

            [[ $run_cls -eq 1      ]] && run_exp global_feat False
            [[ $run_cls_tec -eq 1  ]] && run_exp global_feat True
            [[ $run_afa -eq 1      ]] && run_exp per_patch False
            [[ $run_afa_tec -eq 1  ]] && run_exp per_patch True

        done
    done
done
