#!/usr/bin/env bash
# Extended-251 evaluation of the Stage-2 EMA Predictor: generation sharded over
# $GPUS, VBench only for the strategies generated here, then aggregate.
#   GPUS="0 1 2" bash eval/run_predictor_eval_s2.sh logs_cfa_s2
set -u
C=/local/zoubin/cz/projects/comparison
V=/local/zoubin/cz/projects/VBench
PY=/local/zoubin/cz/envs/self_forcing/bin/python
VPY=/local/zoubin/cz/envs/vbench_eval/bin/python
LOG="$(mkdir -p "$1" && cd "$1" && pwd)"
MODELS="${MODELS:-atc_chunk_s2}"
PATTERNS="${PATTERNS:-FPFF,FPPF,FPPP}"
GPUS=(${GPUS:-0 1 2})
NSHARDS=${#GPUS[@]}

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
export VBENCH_CACHE_DIR=/local/zoubin/cz/.cache/vbench

cd "$C"
echo "=== generation ($(date +%H:%M:%S)) gpus=${GPUS[*]} models=$MODELS patterns=$PATTERNS ==="
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 --models "$MODELS" --patterns "$PATTERNS" \
    > "$LOG/gen_$i.log" 2>&1 &
  pids+=($!)
done
status=0; for p in "${pids[@]}"; do wait "$p" || status=1; done
echo "=== generation done ($(date +%H:%M:%S)) status=$status ==="
[[ $status == 0 ]] || { echo "generation failed; not scoring"; exit 1; }

strats=()
for m in ${MODELS//,/ }; do for pat in ${PATTERNS//,/ }; do strats+=("cfa_${m}_$(echo "$pat" | tr 'A-Z' 'a-z')"); done; done
echo "=== vbench ($(date +%H:%M:%S)) strategies=${strats[*]} ==="
export PYTHONPATH="$V:$C"
cd "$V"
pids=()
for g in $(seq 0 $((NSHARDS-1))); do
  (
    for i in $(seq $g $NSHARDS $(( ${#strats[@]} - 1 ))); do
      s="${strats[$i]}"
      [[ -f "$C/eval_out/vbench/scores/$s.json" ]] && { echo "skip $s"; continue; }
      echo "[gpu ${GPUS[$g]}] $s $(date +%H:%M:%S)"
      CUDA_VISIBLE_DEVICES=${GPUS[$g]} $VPY "$C/eval/run_vbench.py" --strategy "$s" --out-root "$C/eval_out" \
        > "$LOG/vbench_$s.log" 2>&1 || echo "[gpu ${GPUS[$g]}] FAILED $s"
    done
  ) &
  pids+=($!)
done
for p in "${pids[@]}"; do wait "$p"; done
echo "=== aggregating ($(date +%H:%M:%S)) ==="
cd "$C"
$PY eval/aggregate.py --out-root "$C/eval_out" --allow-incomplete > "$LOG/aggregate.log" 2>&1
echo "ALL S2 PHASES DONE $(date +%H:%M:%S)"