#!/bin/bash
# Phase A — query-MLP 1view vs 2view, using DEFAULT-trained ckpts evaluated at OOD shifts.
# 4 cells in parallel: {1v, 2v} × {left positions, right positions}.

set -e
cd /data/cameron/para/libero

CKPT_1V=/data/cameron/para/libero/checkpoints/libero_query_libero_spatial_t0_v0/latest.pth
CKPT_2V=/data/cameron/para/libero/checkpoints/libero_2view_libero_spatial_t0_v0/latest.pth
LEFT_POS=/data/libero/ood_objpos_splits/exp3_left_test.npy
RIGHT_POS=/data/libero/ood_objpos_splits/exp3_right_test.npy

LOG_DIR=/data/cameron/para/libero/logs/phase_a_qmlp_eval
mkdir -p $LOG_DIR

run_cell_2v() {
    local POS_FILE=$1; local LABEL=$2; local GPU=$3
    PYTHONPATH=/data/cameron/LIBERO LIBERO_DATA_PATH=/data/libero CUDA_VISIBLE_DEVICES=$GPU \
    DINO_REPO_DIR=/data/cameron/keygrip/dinov3 \
    DINO_WEIGHTS_PATH=/data/cameron/keygrip/dinov3/weights/dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth \
    MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa \
    python3 -c "
import subprocess, numpy as np, sys, os
positions = np.load('$POS_FILE')
print(f'Evaluating {len(positions)} positions for [$LABEL]', flush=True)
successes, total = 0, 0
for k, (dx, dy) in enumerate(positions):
    cmd = ['python', 'eval_libero_2view_ood.py', '--checkpoint', '$CKPT_2V', '--n_episodes', '5', '--max_steps', '600', '--teleport', '--zero_rotation', '--clean_scene', '--shift_dx', str(dx), '--shift_dy', str(dy)]
    r = subprocess.run(cmd, capture_output=True, text=True, env=os.environ)
    s, t = 0, 0
    for line in r.stdout.split('\n'):
        if 'Success Rate' in line:
            pct = float(line.split(':')[1].strip().split('%')[0]) / 100
            s = int(round(pct * 5)); t = 5; break
    successes += s; total += t
    print(f'  pos[{k+1}/{len(positions)}] dx={dx:+.3f} dy={dy:+.3f}: {s}/{t} (cum: {successes}/{total})', flush=True)
print(f'$LABEL FINAL: {successes}/{total} = {100*successes/total if total>0 else 0:.0f}%')
" > $LOG_DIR/${LABEL}.log 2>&1
}

run_cell_1v() {
    local POS_FILE=$1; local LABEL=$2; local GPU=$3
    PYTHONPATH=/data/cameron/LIBERO LIBERO_DATA_PATH=/data/libero CUDA_VISIBLE_DEVICES=$GPU \
    DINO_REPO_DIR=/data/cameron/keygrip/dinov3 \
    DINO_WEIGHTS_PATH=/data/cameron/keygrip/dinov3/weights/dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth \
    MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa \
    python3 -c "
import subprocess, numpy as np, sys, os
positions = np.load('$POS_FILE')
print(f'Evaluating {len(positions)} positions for [$LABEL]', flush=True)
successes, total = 0, 0
for k, (dx, dy) in enumerate(positions):
    cmd = ['python', 'eval_libero_query_ood.py', '--checkpoint', '$CKPT_1V', '--n_episodes', '5', '--max_steps', '600', '--teleport', '--zero_rotation', '--clean_scene', '--shift_dx', str(dx), '--shift_dy', str(dy)]
    r = subprocess.run(cmd, capture_output=True, text=True, env=os.environ)
    s, t = 0, 0
    for line in r.stdout.split('\n'):
        if 'Success Rate' in line:
            pct = float(line.split(':')[1].strip().split('%')[0]) / 100
            s = int(round(pct * 5)); t = 5; break
    successes += s; total += t
    print(f'  pos[{k+1}/{len(positions)}] dx={dx:+.3f} dy={dy:+.3f}: {s}/{t} (cum: {successes}/{total})', flush=True)
print(f'$LABEL FINAL: {successes}/{total} = {100*successes/total if total>0 else 0:.0f}%')
" > $LOG_DIR/${LABEL}.log 2>&1
}

(run_cell_2v $LEFT_POS  qmlp_2v_left  3) &
(run_cell_2v $RIGHT_POS qmlp_2v_right 4) &
(run_cell_1v $LEFT_POS  qmlp_1v_left  5) &
(run_cell_1v $RIGHT_POS qmlp_1v_right 6) &
wait
echo "=== query-MLP eval matrix complete ==="
for cell in qmlp_2v_left qmlp_2v_right qmlp_1v_left qmlp_1v_right; do
    grep "FINAL:" $LOG_DIR/${cell}.log 2>/dev/null || echo "$cell: missing"
done