#!/bin/bash -x
#SBATCH --nodes=1
#SBATCH --ntasks=4
#SBATCH --ntasks-per-node=4
#SBATCH --wait-all-nodes=1
#SBATCH --partition=a6000
#SBATCH -w c31

export HF_HOME=/mnt/beegfs/dzhu6/.cache

eval "$(/home/dzhu6/miniconda3/bin/conda shell.bash hook)" # init conda
export CUDA_VISIBLE_DEVICES=0,1,2,3
export MASTER_PORT=12784

master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr

conda activate grounded-rl

datadirs=(
    "/mnt/beegfs/dzhu6/SpatialEval/outputs/images"
    # "/mnt/beegfs/dzhu6/coco_images/"
)
script_paths=(  
"SpatialEval"
# "3dsrbench"
# "vsi_bench"
# "maze"
# "spar_bench"
# "mmsi_bench"
)

model_name=ViLaSR-cold-start
CKPT=../ViLaSR-cold-start                # zero_shot/cold_start/reflective/rl CKPT
MODE=cold-start                             # zero_shot, cold_start, reflective, rl

model_name=ViLaSR
CKPT=../ViLaSR                # zero_shot/cold_start/reflective/rl CKPT
MODE=rl                           # zero_shot, cold_start, reflective, rl


for ((i=0; i<${#script_paths[@]}; i++)); do
    QUESTION_FILE=${script_paths[i]}
    IMAGE_FOLDER=${datadirs[i]}

    if [ -z "$IMAGE_FOLDER" ]; then
        echo "Warning: No image folder defined for $QUESTION_FILE. Skipping..."
        continue
    fi

    RESULTDIR=./eval/results/$MODE-vscan-toolonly-resize/${script_paths[i]}/ 
    mkdir -p $RESULTDIR

    ALL=4
    for gpu in $(seq 0 $(($ALL - 1))); do
        echo "****** GPU ${gpu} ${QUESTION_FILE}--${IMAGE_FOLDER} ******"
        CUDA_VISIBLE_DEVICES=$gpu srun --cpu_bind=v --accel-bind=gn --exclusive --ntasks=1 --ntasks-per-node=1 --output=${model_name}_prune_${QUESTION_FILE}_rank${gpu}_all${ALL}.log \
            python eval/infer_multi.py \
            --model-path $CKPT \
            --model-name ${model_name} \
            --dataset ${script_paths[i]} \
            --input-file $IMAGE_FOLDER/${QUESTION_FILE}.json \
            --image-folder $IMAGE_FOLDER \
            --output-dir $RESULTDIR \
            --temperature 0.75 \
            --max-frames 16 \
            --max-pixels 351232 \
            --resize \
            --prune vscan-toolonly \
            --split ${gpu} --all ${ALL} &
        sleep 2
    done
    wait

    # # echo "****** MERGING model ${model} ******"
    # cat outputs/vqa/all/m-${model}_bare_prune_None_split_0_all_${ALL}.jsonl > outputs/vqa/all/m-${model}_bare_prune_None.jsonl
    # for rank in $(seq 1 $(($ALL - 1))); do
    #     cat outputs/vqa/all/m-${model}_bare_prune_None_split_${rank}_all_${ALL}.jsonl >> outputs/vqa/all/m-${model}_bare_prune_None.jsonl
    # done


done