set -euo pipefail
shopt -s nullglob

# ------------------------------
# CONFIG
# ------------------------------
MODEL_TYPE="gemma2_2b_it"
LAYERS=12
INSTRUCTIONS_FILE="/home/dslabra5/sae4steer/axbench/axbench/data/alpaca_eval.json"

PROJECT_ROOT="/home/dslabra5/sae4steer/saes-are-good-for-steering"
FEATURES_DIR="$PROJECT_ROOT/data/features"
CONCEPTS_DIR="$PROJECT_ROOT/concept"

# Where eval_sae_steering.py saves its own artifacts (per-SAE subfolder will be created)
WORKDIR="$(cd "$(dirname "$0")" && pwd)"
cd "$WORKDIR"
SAVE_ROOT="$WORKDIR/runs/gemma2_2b_all"
mkdir -p "$SAVE_ROOT"

# The "existence check" root (used as the single source of truth for de-duplication)
RESULTS_ROOT="$PROJECT_ROOT/cache/results_sae_eval_openai/${MODEL_TYPE}/layer${LAYERS}"
mkdir -p "$RESULTS_ROOT"

# SAE base folders (adjust SAEBench root if needed)
SAEBENCH_ROOT="/home/dslabra5/sae4steer/SAEBench"
BASES=(
  "$SAEBENCH_ROOT/sae_bench/custom_saes/downloaded_saes/trained_saes__google_gemma-2-2b_gated_top_k/resid_post_layer_12"
  "$SAEBENCH_ROOT/sae_bench/custom_saes/downloaded_saes/trained_saes__google_gemma-2-2b_jump_relu_batch_top_k_standard_new/resid_post_layer_12"
  "$SAEBENCH_ROOT/sae_bench/custom_saes/downloaded_saes/trained_saes_2__google_gemma-2-2b_batch_top_k/resid_post_layer_12"
  "$SAEBENCH_ROOT/sae_bench/custom_saes/downloaded_saes/trained_saes_2__google_gemma-2-2b_top_k_jump_relu/resid_post_layer_12"
)

# Eval knobs
STEERING_FACTORS="0.2,0.4,0.8,1.5,2.0,3.0"
DEV_K=5
MAX_NEW_TOKENS=128
JUDGE_BACKEND="openai_async"
JUDGE_MODEL="gpt-4o-mini"

echo "[info] Worker started on CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-ALL} ; RESULTS_ROOT=$RESULTS_ROOT"
mkdir -p "$WORKDIR/logs/workers"

# ------------------------------
# MAIN LOOP
# ------------------------------
for BASE in "${BASES[@]}"; do
  if [[ ! -d "$BASE" ]]; then
    echo "[warn] Base not found: $BASE"
    continue
  fi

  for SAE_DIR in "$BASE"/*; do
    [[ -d "$SAE_DIR" ]] || continue

    NAME="$(basename "$SAE_DIR")"   # e.g., batch_topk_50 / gated_948 / jumprelu_330 / topk_80
    FEATURES_FILE="$FEATURES_DIR/gemma_2b_${NAME}_features.json"
    CONCEPTS_FILE="$CONCEPTS_DIR/${NAME}_concept_descriptions.json"

    # Pre-checks
    if [[ ! -f "$FEATURES_FILE" ]]; then
      echo "[skip] Missing features for $NAME -> $FEATURES_FILE"
      continue
    fi
    if [[ ! -f "$CONCEPTS_FILE" ]]; then
      echo "[skip] Missing concepts for $NAME -> $CONCEPTS_FILE"
      continue
    fi

    # Atomic claim by mkdir RESULT_DIR:
    RESULT_DIR="$RESULTS_ROOT/$NAME"
    if mkdir "$RESULT_DIR" 2>/dev/null; then
      : # claimed
    else
      echo "[skip] Exists/in-progress: $RESULT_DIR -> $NAME"
      continue
    fi

    # Mark in-progress
    date +'%F %T' > "$RESULT_DIR/.inprogress"

    SAVE_DIR="$SAVE_ROOT/$NAME"
    mkdir -p "$SAVE_DIR"

    echo "===================="
    echo "Running SAE: $NAME"
    echo "  dl_local_dir : $SAE_DIR"
    echo "  features_file: $FEATURES_FILE"
    echo "  concepts_file: $CONCEPTS_FILE"
    echo "  save_dir     : $SAVE_DIR"
    echo "  result_dir   : $RESULT_DIR   # used for de-duplication"
    echo "===================="

    # Ensure .inprogress gets removed on SIGINT/SIGTERM for cleanliness
    cleanup() {
      rm -f "$RESULT_DIR/.inprogress" || true
    }
    trap 'cleanup; exit 130' INT TERM

    # Run eval (line-buffered and unbuffered python for timely output)
    set +e
    python -u eval_sae_steering.py \
      --model_type "$MODEL_TYPE" \
      --dl_local_dir "$SAE_DIR" \
      --features_file "$FEATURES_FILE" \
      --instructions_file "$INSTRUCTIONS_FILE" \
      --concepts_file "$CONCEPTS_FILE" \
      --judge_backend "$JUDGE_BACKEND" \
      --judge_model "$JUDGE_MODEL" \
      --layers "$LAYERS" \
      --steering_factors "$STEERING_FACTORS" \
      --dev_k "$DEV_K" \
      --max_new_tokens "$MAX_NEW_TOKENS" \
      --save_dir "$SAVE_DIR" \
      --debug --sample_print_k 1 --print_chars 300 \
      |& tee -a "$SAVE_DIR/run_$(date +%Y%m%d_%H%M%S).log"
    STATUS=${PIPESTATUS[0]}
    set -e

    # Clear trap for next iteration
    trap - INT TERM

    # Finalize markers
    rm -f "$RESULT_DIR/.inprogress" || true
    if [[ $STATUS -ne 0 ]]; then
      echo "[error] Failed on $NAME (exit $STATUS)"
      echo "$STATUS" > "$RESULT_DIR/.failed"
    else
      echo "[done] $NAME"
      date +'%F %T' > "$RESULT_DIR/.done"
    fi

    echo
  done
done

echo "[info] Worker finished on CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-ALL}"
