#!/bin/bash
# Efficiency at P=8: measure latency / FPS (inference.py --report_timing) and peak
# GPU memory for the trained arl2_normalized_patch ckpt run at P=8 vs P=1 (same
# weights, patch_grid the only change -> isolates the multi-patch overhead).
# 3 prompts each (report_timing needs >=2; it times from the 2nd on, excluding
# the first warmup video). Single GPU 0.
set -uo pipefail
cd /home/colligo/kunyang/Causal-Forcing
BIN=/opt/cis/kunyangl/envs/causal_forcing/bin; PY=$BIN/python
CKPT=/opt/cis/kunyangl/gdn_logs/arl2_normalized_20layer_v2_7chunk_teacher_flow_tf_weighted_patch/stage_b_best.pt
BASE=configs/arl2_normalized_20layer_v2_7chunk_patch_infer.yaml
TMP=/tmp/effbench; mkdir -p "$TMP" logs
head -3 workspace/motion_spectrum_prompts.txt > "$TMP/p3.txt"

sample_mem() {  # $1=gpu $2=outfile ; sample every 0.3s until stopfile appears
  local gpu=$1 out=$2
  : > "$out"
  while [ ! -f "$TMP/stop" ]; do
    nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i "$gpu" >> "$out" 2>/dev/null
    sleep 0.3
  done
}

for P in 1 8; do
  case $P in 1) grid="1,1";; 8) grid="2,4";; esac
  cfg="$TMP/cfg_P$P.yaml"; grep -v '^patch_grid:' "$BASE" > "$cfg"; echo "patch_grid: [$grid]" >> "$cfg"
  rm -f "$TMP/stop"; memf="$TMP/mem_P$P.txt"
  sample_mem 0 "$memf" &
  smpid=$!
  echo "===== P=$P (grid=$grid) ====="
  CUDA_VISIBLE_DEVICES=0 "$PY" -u inference.py --config_path "$cfg" \
    --checkpoint_path "$CKPT" --data_path "$TMP/p3.txt" \
    --output_folder "$TMP/out_P$P" --num_output_frames 21 --seed 0 --report_timing \
    2>&1 | grep -iE "\[Sample|latency|FPS|frames" | tee "logs/eff_P$P.log"
  touch "$TMP/stop"; wait $smpid 2>/dev/null
  peak=$(sort -n "$memf" | tail -1)
  echo "P=$P peak_gpu_mem=${peak} MiB"
  echo "P=$P PEAK=${peak}MiB" >> logs/eff_summary.log
done
echo "=== summary ==="; cat logs/eff_summary.log 2>/dev/null; grep -h "Sample" logs/eff_P*.log