#!/usr/bin/env bash
# v11 mix DPO train (GPU 4-7) → P3 → SpatialRGPT → ViewSpatial.
# Single command overnight chain.
set -e
ROOT=/data222/hongbo.wang/Spatial_MLLM
cd $ROOT

PY="conda run -n spatial_mllm --no-capture-output python"
LF=$ROOT/LLaMA-Factory
CKPT=$LF/saves/qwen2_5_vl-7b/full/mgeo_full_dpo_v11_local
TAG=v11_local
VLM_NAME=Qwen2.5-VL-7B-MGeo-Mix-v11

LOG=$ROOT/logs/v11_chain
mkdir -p $LOG

ts() { date +"%Y-%m-%d %H:%M:%S"; }
log() { echo "[$(ts) V11] $*" | tee -a $LOG/orchestrator.log; }

############################
# Stage 0: DPO training on 4-7
############################
if [ -f $CKPT/config.json ]; then
  log "Train: $CKPT exists, skip"
else
  log "=== Stage 0/4: v11 DPO training on GPU 4-7 ==="
  cd $LF
  CUDA_VISIBLE_DEVICES=4,5,6,7 \
    conda run -n spatial_mllm --no-capture-output \
    llamafactory-cli train mgeo_full_dpo_v11_local.yaml \
    > $LOG/train.log 2>&1 || { log "Train FAILED"; exit 1; }
  cd $ROOT
  log "Train done"
fi

############################
# Stage 1: P3 probe v2
############################
log "=== Stage 1/4: P3 probe v2 ==="
if [ -f $ROOT/results/probe_p3/preds_${TAG}.jsonl ]; then
  log "P3: exists, skip"
else
  $PY $ROOT/mgeo_dpo/scripts/eval_p3.py --model $TAG --gpus 4,5,6,7 \
      > $LOG/p3.log 2>&1 || log "P3 FAILED, continue"
fi
$PY $ROOT/mgeo_dpo/scripts/p3_metrics.py \
    --models $TAG base v3 v9 msmu_dpo_v1 \
    > $LOG/p3_metrics.log 2>&1 || log "p3_metrics failed"

############################
# Stage 2: SpatialRGPT-Bench
############################
log "=== Stage 2/4: SpatialRGPT-Bench ==="
MODEL_BASE=$ROOT/models/Qwen2.5-VL-7B-Instruct
BENCH_JSON=/data222/hongbo.wang/SpatialRGPT/bench_data/SpatialRGPT-Bench_v1.json
IMAGE_DIR=/data222/hongbo.wang/SpatialRGPT/bench_data/images
EVAL_SCRIPT=/data222/hongbo.wang/SpatialRGPT/eval_qwen2vl_spatialrgpt.py
GPT_SCRIPT=/data222/hongbo.wang/SpatialRGPT/evaluate_spatial_with_api.py

OUT_SRGPT=/data222/hongbo.wang/SpatialRGPT/eval_output/${TAG}
mkdir -p $OUT_SRGPT

if [ -f $OUT_SRGPT/answers.jsonl ]; then
  log "SRGPT infer: exists, skip"
else
  log "SRGPT infer on GPU 4"
  CUDA_VISIBLE_DEVICES=4 $PY $EVAL_SCRIPT \
      --model-path "$CKPT" --adapter-path NONE \
      --bench-json $BENCH_JSON --image-dir $IMAGE_DIR \
      --output $OUT_SRGPT/answers.jsonl --model-name $TAG \
      > $LOG/srgpt_infer.log 2>&1 || log "SRGPT infer FAILED"
fi

if [ -f $OUT_SRGPT/score.json ]; then
  log "SRGPT score: exists, skip"
else
  log "SRGPT GPT score"
  $PY $GPT_SCRIPT $OUT_SRGPT/answers.jsonl >> $LOG/srgpt_infer.log 2>&1 || log "SRGPT score FAILED"
fi

############################
# Stage 3: ViewSpatial-Bench
############################
log "=== Stage 3/4: ViewSpatial-Bench ==="
OUTDIR=$ROOT/results/vlmeval_v11/${VLM_NAME}
ACCFILE=$OUTDIR/${VLM_NAME}_ViewSpatialBench_extract_matching_acc.csv

if [ -f $ACCFILE ]; then
  log "ViewSpatial: exists, skip"
else
  mkdir -p $OUTDIR
  cd /data222/hongbo.wang/VLMEvalKit
  CUDA_VISIBLE_DEVICES=4,5,6,7 conda run -n vlmeval --no-capture-output \
    torchrun --nproc-per-node=4 --master_port=29611 run.py \
      --data ViewSpatialBench --model $VLM_NAME \
      --reuse --work-dir $OUTDIR \
      > $LOG/vs.log 2>&1 || log "ViewSpatial FAILED"
  cd $ROOT
fi

############################
# Stage 4: Final summary
############################
log "=== Stage 4/4: ALL DONE — summary ==="
log "--- P3 ---"
grep "^${TAG}\|^msmu_dpo\|^v9\|^v3\|^base" $LOG/p3_metrics.log 2>/dev/null | tee -a $LOG/orchestrator.log
log "--- SpatialRGPT ---"
[ -f $OUT_SRGPT/score.json ] && cat $OUT_SRGPT/score.json | tee -a $LOG/orchestrator.log
log "--- ViewSpatial ---"
[ -f $ACCFILE ] && cat $ACCFILE | tee -a $LOG/orchestrator.log

log "=== v11 chain complete ==="