#!/usr/bin/env bash
set -Eeuo pipefail
METHODS=(selu-no-cpl selu-no-pair)           # Methods to sweep
DATASET_SUBSET="${DATASET_SUBSET:-forget05}"

LORA_PAIRS=("16:32")          # rank:alpha
# LEARNING_RATES=("1e-04")            # LR sweep
LEARNING_RATES=("0.0001")

BASE_MODEL="${BASE_MODEL:-open-unlearning/tofu_Llama-2-7b-chat-hf_full}"
# BASE_MODEL="${BASE_MODEL:-open-unlearning/tofu_Llama-2-7b-chat-hf_retain90}"
# BASE_MODEL="${BASE_MODEL:-jialicheng/tofu_ft_retain90_llama2-7b}"

PERTURBED_DIR="${PERTURBED_DIR:-./dataset}"
RETAIN_ROOT="${RETAIN_ROOT:-./retain}"

failed=()

# ==========================
# Sweep
# ==========================
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_evaluate.py \
            --model_id "$BASE_MODEL" \
            --adapter_id "$ADAPTER_ID" \
            --perturbed_dir "$PERTURBED_DIR" \
            --retain_root "$RETAIN_ROOT" \
            --all_epochs \
            --split forget05 \
            "$@"
      then
        echo "!!! FAILED: method=$M pair=${RANK}:${ALPHA} lr=$LR"
        failed+=("${M}:${RANK}:${ALPHA}:${LR}")
        continue
      fi
    done
  done

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

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

