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

repo=/home/sjkim/POLYEDIT
run_root=$repo/project/experiments/local_runs/multiobjective_20260721
ckpt_root=/tmp/polyedit_multiobjective_ckpts_20260721
hf_root=/tmp/polyedit_multiobjective_hf_20260721
wandb_root=/tmp/polyedit_multiobjective_wandb_20260721
polybert=$hf_root/polybert
python_bin=/home/sjkim/anaconda3/envs/main/bin/python
revision=7bf9ed32ac54dea5bc163cf90100728b49341750

cleanup() {
  rm -rf "$hf_root" "$wandb_root"
}
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 "$run_root/results" "$run_root/logs" "$ckpt_root" "$hf_root" "$wandb_root"
export PYTHONPATH=$repo
export HF_HOME=$hf_root/cache
export HF_HUB_CACHE=$hf_root/cache/hub
export WANDB_DIR=$wandb_root
export TOKENIZERS_PARALLELISM=false

/home/sjkim/anaconda3/envs/main/bin/hf download HAYDERphd/polyBERT \
  --revision "$revision" --local-dir "$polybert"

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_multiobjective.py \
      --seed "$seed" --polybert_path "$polybert" --out "$result" \
      --ckpt "$ckpt_root/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_multiobjective.py \
  --input "$run_root/results" --output "$run_root/summary.json" \
  --markdown "$run_root/summary.md"