#!/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 focallora4
else
  echo "[TestInstructiveHead.sh] Warning: conda not found; running in current environment." >&2
fi

export CUDA_VISIBLE_DEVICES=0

SCRIPT=./EvaluateInstructiveHead.py
MODEL=../../models/Qwen2-7B-Instruct
DATA=../1_raw_dataset/topicattack/data/lagecy/crafted_instruction_data_squad_injection_qa_llama_filtered_test.json
SYSTEM=../1_raw_dataset/topicattack/prompts/generator_system_prompt.txt
TEMPLATE=../1_raw_dataset/topicattack/prompts/victim_instruction_data_template.txt
HEADS_DIR=../2-2_head_identification/head_scoring/qwen2-7b_sep/heads_sorted

# Parent of HEADS_DIR
HEADS_PARENT="$(dirname "$HEADS_DIR")"

# Output directory
EVAL_DIR="$HEADS_PARENT/heads_sorted_eval"
mkdir -p "$EVAL_DIR"

TARGETS="${TARGETS:-}"

HEADS_PARENT="$(dirname "$HEADS_DIR")"
EVAL_DIR="$HEADS_PARENT/heads_sorted_eval"
mkdir -p "$EVAL_DIR"

TOPKS="0 3.125p 6.25p 9.375p 12.5p 15.625p 18.75p 21.875p 25p 28.125p 31.25p 50p 62.5p 75p 100p"
# TOPKS="62.5p 75p 100p"

found_any=0

TARGETS="
focallora.json
"
set -x
if [ -z "$TARGETS" ]; then
  HEADS_LIST="$HEADS_DIR"/*.json
else
  HEADS_LIST=""
  for t in $TARGETS; do
    HEADS_LIST="$HEADS_LIST $HEADS_DIR/$t"
  done
fi

for HEADS in $HEADS_LIST; do
  [ -f "$HEADS" ] || continue
  found_any=1

  head_name="$(basename "$HEADS" .json)"

  for TOPK in $TOPKS; do
    python "$SCRIPT" \
      --victim_model_path "$MODEL" \
      --data_path "$DATA" \
      --victim_system_path "$SYSTEM" \
      --input_template_path "$TEMPLATE" \
      --victim_head_list "$HEADS" \
      --attacks none naive ignore escape_separation \
      --topk "$TOPK" \
      --batch_size 50 \
      --data_size -1 \
      --eval_result_file "$EVAL_DIR/${head_name}_${TOPK}.json"
  done
done

if [ "$found_any" -eq 0 ]; then
  echo "[TestInstructiveHead.sh] Error: no matching head json files found" >&2
  exit 1
fi
