#!/usr/bin/env sh
set -eu
export CUDA_VISIBLE_DEVICES=0
CONDA_BIN=${CONDA_BIN:-}
if [ -z "$CONDA_BIN" ]; then
  if command -v conda >/dev/null 2>&1; then
    CONDA_BIN=$(command -v conda)
  elif [ -x /opt/miniconda/bin/conda ]; then
    CONDA_BIN=/opt/miniconda/bin/conda
  fi
fi

if [ -n "$CONDA_BIN" ]; then
  eval "$("$CONDA_BIN" shell.bash hook)"
  conda activate focallora
else
  echo "[TestInstructiveHead.sh] Warning: conda not found; running in current environment." >&2
fi

echo "========================================================================================================================================="
echo "Mask attack only, for all heads"
echo "(test masking works. 100% heads can't see attack instruction, if model respond any related attack inctruction means mask does not work)."
python ./TestInstructiveHead.py \
  --model-path ../../models/Llama-3.1-8B-Instruct \
  --heads ../2-2_head_identification/head_scoring/llama31-8b_sep/head_roc_inst.json \
  --topk 100p \
  --SEP-dataset ../1_raw_dataset/Should-It-Be-Executed-Or-Processed/datasets/validation_dataset.json \
  --mask-attack-only 1 \
  --data-index 0,1,2 \
#  --no-print-input
#  --debug-mask

echo "========================================================================================================================================="
echo "Mask whole data, for whole heads"
echo "(test masking works. 100% heads can't see whole, the model should can't do anything)."
python ./TestInstructiveHead.py \
  --model-path ../../models/Llama-3.1-8B-Instruct \
  --heads ../2-2_head_identification/head_scoring/llama31-8b_sep/head_roc_inst.json \
  --topk 100p \
  --SEP-dataset ../1_raw_dataset/Should-It-Be-Executed-Or-Processed/datasets/validation_dataset.json \
  --mask-attack-only 0 \
  --data-index 0,1,2 \
  --no-print-input
#  --debug-mask

echo "========================================================================================================================================="
echo "Mask whole data, for instruction sensitive heads"
echo "(test my approch works, make **instruction following** heads unable to see data segment to prevent attack, but other heads can see data to answer user question)"
python ./TestInstructiveHead.py \
  --model-path ../../models/Llama-3.1-8B-Instruct \
  --heads ../2-2_head_identification/head_scoring/llama31-8b_sep/head_roc_inst.json \
  --topk 20p \
  --SEP-dataset ../1_raw_dataset/Should-It-Be-Executed-Or-Processed/datasets/validation_dataset.json \
  --mask-attack-only 0 \
  --data-index 0,1,2 \
  --no-print-input
#  --debug-mask

python ./TestInstructiveHead.py \
  --model-path ../../models/Llama-3.1-8B-Instruct \
  --heads ../2-2_head_identification/head_scoring/llama31-8b_sep/head_roc_inst.json \
  --topk 30p \
  --SEP-dataset ../1_raw_dataset/Should-It-Be-Executed-Or-Processed/datasets/validation_dataset.json \
  --mask-attack-only 0 \
  --data-index 0,1,2 \
  --no-print-input
#  --debug-mask