set -x

GPUS=${GPUS:-1}
BATCH_SIZE=${BATCH_SIZE:-8}
PER_DEVICE_BATCH_SIZE=${PER_DEVICE_BATCH_SIZE:-8}
GRADIENT_ACC=$((BATCH_SIZE / PER_DEVICE_BATCH_SIZE / GPUS))

export PYTHONPATH="${PYTHONPATH}:$(pwd)/InternVL/internvl_chat"
export MASTER_PORT=34229
export TF_CPP_MIN_LOG_LEVEL=3
export LAUNCHER=pytorch

attack_type=badVLM  # badVLM, TrojVQA, badCLIP, blended, badNet

if [ $attack_type == "badVLM" ]; then
    pattern=badVLM
    patch_location=middle
    trigger_path=backdoors/config/badVLM/patch.png
    dual=false

elif [ $attack_type == "TrojVQA" ]; then
    pattern=vqa
    patch_location=middle
    trigger_path=backdoors/config/TrojVQA/TrojanVQA_default_vqa.jpg
    dual=true

elif [ $attack_type == "badCLIP" ]; then
    pattern=badCLIP
    patch_location=middle
    trigger_path=backdoors/config/badCLIP/trigger.jpg

elif [ $attack_type == "blended" ]; then
    pattern=blended
    patch_location=blended
    trigger_path="."

elif [ $attack_type == "badNet" ]; then
    pattern=random
    patch_location=random
    trigger_path="."

else
    echo "Unknown attack type"
    exit 1
fi

target_attack=template
target_label=banana
poison_size=100
dataset=coco #cc3m

OUTPUT_DIR=./pretrained/InternVL2-2B-${attack_type}-caption-${poison_size}

if [ ! -d "$OUTPUT_DIR" ]; then
  mkdir -p "$OUTPUT_DIR"
fi

#### Backdoor Config
python -m backdoors.backdoor_config \
        --attack-type $attack_type \
        --pattern $pattern \
        --patch-location $patch_location \
        --target-label $target_label \
        --poison-size $poison_size \
        --pattern-size 16 \
        --dataset $dataset \
        --trigger-path $trigger_path \
        --dual $dual

### Generate poisoned dataset
python -m backdoors.backdoor_generation  \
        --attack-type $attack_type \
        --target-attack $target_attack \
        --model-type InternVL

### Train the backdoor attack
torchrun \
  --nnodes=1 \
  --node_rank=0 \
  --master_addr=127.0.0.1 \
  --nproc_per_node=${GPUS} \
  --master_port=${MASTER_PORT} \
  InternVL/internvl_chat/internvl/train/internvl_chat_finetune.py \
  --model_name_or_path "./pretrained/InternVL2-2B" \
  --conv_style "internlm2-chat" \
  --output_dir ${OUTPUT_DIR} \
  --meta_path "./backdoors/internvl_caption_metadata.json" \
  --overwrite_output_dir True \
  --force_image_size 448 \
  --max_dynamic_patch 6 \
  --down_sample_ratio 0.5 \
  --drop_path_rate 0.0 \
  --freeze_llm True \
  --freeze_mlp False \
  --freeze_backbone True \
  --use_llm_lora 32 \
  --use_backbone_lora 32 \
  --vision_select_layer -1 \
  --dataloader_num_workers 4 \
  --bf16 True \
  --num_train_epochs 1 \
  --per_device_train_batch_size ${PER_DEVICE_BATCH_SIZE} \
  --gradient_accumulation_steps ${GRADIENT_ACC} \
  --evaluation_strategy "no" \
  --save_strategy "steps" \
  --save_steps 200 \
  --save_total_limit 1 \
  --learning_rate 4e-5 \
  --weight_decay 0.01 \
  --warmup_ratio 0.03 \
  --lr_scheduler_type "cosine" \
  --logging_steps 1 \
  --max_seq_length 4096 \
  --do_train True \
  --grad_checkpoint True \
  --group_by_length True \
  --dynamic_image_size True \
  --use_thumbnail True \
  --ps_version 'v2' \
  --deepspeed "InternVL/internvl_chat/zero_stage1_config.json" \
  --report_to "tensorboard" \
  2>&1 | tee -a "${OUTPUT_DIR}/training_log.txt"


### Backdoor eval
MASTER_PORT=${MASTER_PORT:-63669}
export MASTER_PORT=${MASTER_PORT}
export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}

python -m backdoors.backdoor_eval_internvl \
        --checkpoint $OUTPUT_DIR \
        --dynamic \
        --attack-type $attack_type