#!/bin/bash

# Default parameters
DEVICE_ID=0
SEED=77
TASK="seg"
DATASET="mvtec"
DATA_PATH="../../mvtec"
CLIP_TEMPLATE="a good, cropped picture of the {} for classification."
LANGUAGE_TEMPLATE="a close-up cropped png photo of a {} for anomaly segmentation."
LANGUAGE_STATES=("crack" "hole" "residue" "damage")
LANGUAGE_TIMESTEPS=(901)
LANGUAGE_BLOCKS=("down" "up")
MODEL_WEIGHT=0.25
VISION_TIMESTEPS=(801 401 201)
VISION_BLOCKS=('1' '2' '3')
VISION_TEMPLATE="a photo of a perfect {}."

# Define SHOT values
SHOT_VALUES=(1 2 4)
ZERO_SHOT=true

if [ "$ZERO_SHOT" = true ]; then
    # Run only for SHOT=0
    echo "Running with SHOT=0"
    python seg.py \
        --device_id $DEVICE_ID \
        --seed $SEED \
        --task $TASK \
        --dataset $DATASET \
        --data_path $DATA_PATH \
        --shot 0 \
        --model_weight $MODEL_WEIGHT \
        --clip_template "$CLIP_TEMPLATE" \
        --language_template "$LANGUAGE_TEMPLATE" \
        --language_states "${LANGUAGE_STATES[@]}" \
        --language_timesteps "${LANGUAGE_TIMESTEPS[@]}" \
        --language_blocks "${LANGUAGE_BLOCKS[@]}" \
        --vision_template "$VISION_TEMPLATE" \
        --vision_timesteps "${VISION_TIMESTEPS[@]}" \
        --vision_blocks "${VISION_BLOCKS[@]}"
else
    # Run for SHOT=1, 2, 4
    for SHOT in "${SHOT_VALUES[@]}"
    do
        echo "Running with SHOT=$SHOT"
        python seg.py \
            --device_id $DEVICE_ID \
            --seed $SEED \
            --task $TASK \
            --dataset $DATASET \
            --data_path $DATA_PATH \
            --shot $SHOT \
            --model_weight $MODEL_WEIGHT \
            --clip_template "$CLIP_TEMPLATE" \
            --language_template "$LANGUAGE_TEMPLATE" \
            --language_states "${LANGUAGE_STATES[@]}" \
            --language_timesteps "${LANGUAGE_TIMESTEPS[@]}" \
            --language_blocks "${LANGUAGE_BLOCKS[@]}" \
            --vision_template "$VISION_TEMPLATE" \
            --vision_timesteps "${VISION_TIMESTEPS[@]}" \
            --vision_blocks "${VISION_BLOCKS[@]}"
    done
fi
