#!/usr/bin/env bash
set -Eeuo pipefail

METHODS=(ga)
DATASET_SUBSET="${DATASET_SUBSET:-forget10}"
LORA_PAIRS=("8:32" "16:32" "4:16")          # rank:alpha
LEARNING_RATES=("1e-04" "1e-05")              # LR sweep

BASE_MODEL="${BASE_MODEL:-open-unlearning/tofu_Llama-2-7b-chat-hf_full}"
DATA_PATH="${DATA_PATH:-tinyBenchmarks/tinyMMLU}"

failed=()

for M in "${METHODS[@]}"; do
  echo ">>> Evaluating method=$M on subset=$DATASET_SUBSET"

  for PAIR in "${LORA_PAIRS[@]}"; do
    IFS=':' read -r RANK ALPHA <<< "$PAIR"

    for LR in "${LEARNING_RATES[@]}"; do
      ADAPTER_ID="ganeric15/msc_unlearn_lora_${RANK}_${ALPHA}_${LR}_${M}_tofu_${DATASET_SUBSET}"
      echo "---- Running ${ADAPTER_ID}"

      if ! python run_tofu_unlearn_eval_utility.py \
            --model_id "$BASE_MODEL" \
            --adapter_id "$ADAPTER_ID" \
            --data_path "$DATA_PATH" \
            --all_epochs \
            "$@"
      then
        echo "!!! FAILED: method=$M pair=${RANK}:${ALPHA} lr=$LR"
        failed+=("${M}:${RANK}:${ALPHA}:${LR}")
        continue
      fi
    done
  done

  echo ">>> Finished method=$M"
done

if ((${#failed[@]})); then
  echo "Done with errors. Failed combos:"
  for f in "${failed[@]}"; do echo "  - $f"; done
  exit 1
else
  echo "All evaluations complete."
fi

