#!/usr/bin/env bash
# Extended-251 evaluation of the Causal-Forcing-a Predictors (cfa_* strategies):
# phase 1 generation sharded across $GPUS, phase 2 VBench on the not-yet-scored
# strategies (run_vbench_all.sh skips scored ones), phase 3 aggregate.
#   GPUS="0 1 2 3" eval/run_predictor_eval.sh logs_cfa
set -u
C=/local/zoubin/cz/projects/comparison
PY=/local/zoubin/cz/envs/self_forcing/bin/python
LOG="$(cd "$(dirname "$1")" && pwd)/$(basename "$1")"; shift
GPUS=(${GPUS:-0 1 2 3})
NSHARDS=${#GPUS[@]}
mkdir -p "$LOG"

export SSL_CERT_FILE=/etc/pki/tls/certs/ca-bundle.crt
export TORCH_HOME=/local/zoubin/cz/.cache/torch
export HF_HOME=/local/zoubin/cz/.cache/huggingface
export HF_HUB_CACHE=/local/zoubin/cz/.cache/huggingface/hub
export TRANSFORMERS_CACHE=/local/zoubin/cz/.cache/huggingface/hub
export XDG_CACHE_HOME=/local/zoubin/cz/.cache

cd "$C"
echo "=== cfa generation ($(date +%H:%M:%S)) gpus=${GPUS[*]} ==="
pids=()
for i in $(seq 0 $((NSHARDS-1))); do
  CUDA_VISIBLE_DEVICES=${GPUS[$i]} $PY eval/generate_predictor_eval.py \
    --shard $i --num-shards $NSHARDS --out-root eval_out "$@" \
    > "$LOG/gen_cfa_$i.log" 2>&1 &
  pids+=($!)
done
status=0
for p in "${pids[@]}"; do wait "$p" || status=1; done
echo "=== cfa generation done ($(date +%H:%M:%S)) status=$status ==="
[[ $status == 0 ]] || { echo "generation failed; not scoring"; exit 1; }

GPUS="${GPUS[*]}" NGPU=$NSHARDS bash "$C/eval/run_vbench_all.sh" "$LOG" 2>&1 | tee "$LOG/vbench_driver.log"

echo "=== aggregating ($(date +%H:%M:%S)) ==="
$PY "$C/eval/aggregate.py" --out-root "$C/eval_out" 2>&1 | tee "$LOG/aggregate.log"
echo "ALL CFA PHASES DONE $(date +%H:%M:%S)"