#!/usr/bin/env bash
# CLV leading-vector geometry probe (b): does the EXPANSION DIRECTION carry outcome signal
# beyond drift+q_halt? Sudoku TRM + HRM (mlp_t, no flash issue), n=512, on the free GPU.
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}"
OUTDIR=analysis_2x2/clv; mkdir -p "$OUTDIR"
ST="$OUTDIR/status.log"; log(){ echo "[$(date '+%H:%M:%S')] $*" >> "$ST"; }
TRM=/home/yurenh2/rrm/trm/checkpoints/Sudoku-extreme-1k-aug-1000-ACT-torch/pretrain_mlp_t_sudoku_official_gbs768_repro
HRM="/home/yurenh2/rrm/hrm/checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 righteous-python"

log "start TRM geom n=512 (GPU $CUDA_VISIBLE_DEVICES)"
if python diagnose_trm_joint_geom.py --ckpt-root "$TRM" --ckpt-name step_58590 \
   --n-samples 512 --batch-size 8 --k-lyap 8 --t-ons 1 --seed 0 \
   --out "$OUTDIR/trm_geom_step58590_n512.npz" > "$OUTDIR/trm.log" 2>&1; then log "done TRM"; else log "FAILED TRM"; fi

log "start HRM geom n=512"
if python diagnose_hrm_joint_geom.py --ckpt-root "$HRM" --ckpt-name step_26040 \
   --n-samples 512 --batch-size 16 --k-lyap 8 --t-ons 1 --seed 0 \
   --out "$OUTDIR/hrm_geom_step26040_n512.npz" > "$OUTDIR/hrm.log" 2>&1; then log "done HRM"; else log "FAILED HRM"; fi
log "clv geom finished"