#!/bin/bash
# Offline token extraction orchestrator.
# Isolated by default. Set ATLAS_ALLOW_OFFLINE=1 to run.
set -e

if [ "${ATLAS_ALLOW_OFFLINE}" != "1" ]; then
    echo "ERROR: This is an OFFLINE extraction orchestrator." >&2
    echo "It is isolated by default to prevent accidental use." >&2
    echo "If you really need it, set:  ATLAS_ALLOW_OFFLINE=1" >&2
    echo "For online training use:  bash scripts/train_no_caption_baseline.sh" >&2
    exit 1
fi

cd /home/guoyuanbo/3dtokenizer-atlas
export LD_LIBRARY_PATH=/home/guoyuanbo/3dtokenizer/envs/streampetr/lib:${LD_LIBRARY_PATH:-}
PY=/home/guoyuanbo/3dtokenizer/envs/streampetr/bin/python

LOG_DIR=work_dirs
DET_OUT=work_dirs/precomputed_det_tokens_offline/val
MAP_OUT=work_dirs/precomputed_map_tokens_offline/val

echo "[$(date)] === Waiting for Phase 1 (StreamPETR det) to finish ==="

# Wait for all extract_streampetr_tokens processes to finish
while pgrep -f "extract_streampetr_tokens.py" > /dev/null 2>&1; do
    DET_COUNT=$(find "$DET_OUT" -name "*.pt" 2>/dev/null | wc -l)
    echo "[$(date)] Phase 1 running... det files: $DET_COUNT / ~12038"
    sleep 300
done

DET_FINAL=$(find "$DET_OUT" -name "*.pt" 2>/dev/null | wc -l)
echo "[$(date)] Phase 1 DONE. Total det files: $DET_FINAL"

echo "[$(date)] === Starting Phase 2 (TopoMLP map) ==="
mkdir -p "$MAP_OUT"

for i in 0 1 2 3; do
    CUDA_VISIBLE_DEVICES=$i $PY extract_topomlp_tokens.py \
        --topomlp_config configs/topomlp_atlas_aligned.py \
        --topomlp_ckpt work_dirs/topomlp_atlas_aligned/epoch_24.pth \
        --data_json "data/atlas_planning_val_uniad_command.json,data/openlane_subsetB_lane_val_4pt.json" \
        --data_root /home/guoyuanbo/autodl-tmp/data/nuscenes \
        --output_dir "$MAP_OUT" \
        --shard_id $i --num_shards 4 \
        > "$LOG_DIR/extract_map_val_shard_${i}.log" 2>&1 &
    echo "[$(date)] Phase 2 shard $i launched (PID=$!)"
done

echo "[$(date)] Waiting for Phase 2 to complete..."
wait

MAP_FINAL=$(find "$MAP_OUT" -name "*.pt" 2>/dev/null | wc -l)
echo "[$(date)] Phase 2 DONE. Total map files: $MAP_FINAL"
echo "[$(date)] === All extraction complete ==="
echo "  Det tokens: $DET_FINAL files in $DET_OUT"
echo "  Map tokens: $MAP_FINAL files in $MAP_OUT"