#!/usr/bin/env bash
# Run after the early-save training: cheap forward dumps (connectivity + drift_zH + ans_drift,
# NO JVP) on all saved early checkpoints, then pooled broken-vs-complete analysis.
set -o pipefail
cd /home/yurenh2/rrm/research/flossing
source /home/yurenh2/miniconda3/etc/profile.d/conda.sh
conda activate rrm
export CUDA_VISIBLE_DEVICES="${1:-1}"
CK=/home/yurenh2/rrm/trm/checkpoints/Maze-30x30-hard-1k-ACT-torch/maze_earlysave_freshTRM
DATA=/home/yurenh2/rrm/data/maze-30x30-hard-1k
OUT=analysis_2x2/maze_followup
for f in "$CK"/step_*; do
  s=$(basename "$f")
  [[ -f "$OUT/earlydump_${s}.npz" ]] && { echo "skip $s"; continue; }
  echo "=== dump $s ==="
  python maze_pred_dump.py --ckpt-root "$CK" --ckpt-name "$s" --data "$DATA" \
    --n 512 --seed 20260616 --out "$OUT/earlydump_${s}.npz" 2>&1 | tail -1
done
echo "=== pooled analysis ==="
python analysis_2x2/analyze_early_connectivity.py