#!/usr/bin/env bash
# Fire M1 layer-wise probing on every new (2026-05-05) run that has
# periodic checkpoints, sequentially on the GPU. Sequential because
# parallel M1 jobs would each grab the full GPU and contend.
#
# Usage:
#   nohup bash scripts/run_mi_all_new.sh > logs/mi_all_new.log 2>&1 &
#   disown

set -e

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
cd "$ROOT"

LOG="logs/mi_all_new.log"
mkdir -p logs

stamp() { date -u +"[%FT%TZ]"; }

echo "$(stamp) MI pipeline starting on all 2026-05-05 runs (GPU)" | tee -a "$LOG"

for d in experiments/runs/20260505-*/; do
    rid=$(basename "$d")
    ckpt_count=$(ls "$d/checkpoints/"ep*.pt 2>/dev/null | wc -l)
    if [ "$ckpt_count" -lt 5 ]; then
        echo "$(stamp) skip $rid (only $ckpt_count periodic ckpts)" | tee -a "$LOG"
        continue
    fi
    if [ -f "$d/mechinterp/m1_probe_data.json" ]; then
        echo "$(stamp) skip $rid (already has m1_probe_data.json)" | tee -a "$LOG"
        continue
    fi
    echo "$(stamp) ==== M1 on $rid ($ckpt_count ckpts) ====" | tee -a "$LOG"
    set +e
    CUDA_VISIBLE_DEVICES=0 python3 -m experiments.mechinterp_m1 \
        --run_dir "$d" \
        --data_root data/wilds \
        --device cuda \
        --max_samples 600 >> "$LOG" 2>&1
    rc=$?
    set -e
    echo "$(stamp) ==== $rid exit=$rc ====" | tee -a "$LOG"
done

echo "$(stamp) MI pipeline complete." | tee -a "$LOG"