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

repo=/home/sjkim/POLYEDIT
revision=11bcca4e376b6c0de3faea57491080bb375072d1
bundle_root=/tmp/polyedit_rule_ood_bundles
hf_cache=/tmp/polyedit_rule_ood_hf
wandb_dir=/tmp/polyedit_rule_ood_wandb
run_root=$repo/project/experiments/local_runs/rule_ood_20260721
python_bin=/home/sjkim/anaconda3/envs/main/bin/python

cleanup() {
  rm -rf "$bundle_root" "$hf_cache" "$wandb_dir"
}
trap cleanup EXIT

if [[ -z ${WANDB_API_KEY:-} ]]; then
  read -rsp "W&B API key: " WANDB_API_KEY
  export WANDB_API_KEY
  echo
fi
mkdir -p "$bundle_root" "$hf_cache" "$wandb_dir" "$run_root/results" \
  "$run_root/logs" "$run_root/checkpoints"
export PYTHONPATH=$repo:$repo/project/experiments
export HF_HOME=$hf_cache
export HF_HUB_CACHE=$hf_cache/hub
export WANDB_DIR=$wandb_dir
export TOKENIZERS_PARALLELISM=false

/home/sjkim/anaconda3/envs/main/bin/hf download \
  promotion/polyedit-real-hard-rl-5seed \
  --revision "$revision" \
  --include 'grpo/checkpoints/seed_*/verifier_Egc_bundle.pt' \
  --local-dir "$bundle_root"

worker() {
  gpu=$1
  shift
  for seed in "$@"; do
    result=$run_root/results/seed_${seed}.json
    [[ -s $result ]] && continue
    CUDA_VISIBLE_DEVICES=$gpu "$python_bin" scripts/train_polyedit_rule_ood.py \
      --seed "$seed" \
      --verifier_bundle "$bundle_root/grpo/checkpoints/seed_${seed}/verifier_Egc_bundle.pt" \
      --out "$result" \
      --ckpt "$run_root/checkpoints/seed_${seed}" \
      --wandb 2>&1 | tee "$run_root/logs/seed_${seed}.log"
  done
}

cd "$repo"
worker 0 1004 1006 1008 &
worker0=$!
worker 1 1005 1007 &
worker1=$!
wait "$worker0" "$worker1"

"$python_bin" scripts/summarize_polyedit_rule_ood.py \
  --input "$run_root/results" \
  --output "$run_root/summary.json" \
  --markdown "$run_root/summary.md"