#!/usr/bin/env bash
# 使用与 Haozhe2 训练相同的磁盘 superpoint.npy（本机 s3dis_official 与 Haozhe2 MD5 一致）做可视化。
# Mitsuba：整景超点/语义 GT；按语义类拆分超点 HQ；可选 Blender Workbench。
set -eo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$HERE"
source /mnt/data/AODUOLI/miniconda3/etc/profile.d/conda.sh
conda activate /mnt/data/AODUOLI/miniconda_envs/Aoduo
export PYTHONPATH="$HERE:${PYTHONPATH:-}"

OUT="${OUT:-$HERE/outputs/superpoint_vis}"
DATA_ROOT="${DATA_ROOT:-$HERE/../_work_biptv3/pointcept_framework/data/s3dis_official}"
FILM="${FILM:-2560}"
SPP="${SPP:-256}"
MAXP="${MAXP:-40000}"

# 空间名不含空格，与 Area_1/<room> 一致
ROOMS="${ROOMS:-office_1 hallway_1 conferenceRoom_1 office_2 copyRoom_1}"
DO_MITSUBA="${DO_MITSUBA:-1}"
DO_MITSUBA_PERCLASS="${DO_MITSUBA_PERCLASS:-0}"
PERCLASS_ROOM="${PERCLASS_ROOM:-Area_1/office_1}"
DO_BLENDER="${DO_BLENDER:-0}"
LABEL_NPY="${LABEL_NPY:-}"
LABEL_TAG="${LABEL_TAG:-}"
if [[ -n "$LABEL_NPY" && -z "$LABEL_TAG" ]]; then
  LABEL_TAG="alt"
fi
LABEL_SUFFIX=""
if [[ -n "$LABEL_TAG" ]]; then
  LABEL_SUFFIX="_${LABEL_TAG}"
fi
SP_LABEL_ARGS=()
if [[ -n "$LABEL_NPY" ]]; then
  SP_LABEL_ARGS+=(--label_npy "$LABEL_NPY")
fi

mitsuba_full() {
  local r
  for r in $ROOMS; do
    echo "==== Mitsuba $r (superpoint + segment GT) ===="
    python3 "$HERE/superpoint_visualize_s3dis.py" --data_root "$DATA_ROOT" --room "Area_1/${r}" --mode superpoint \
      "${SP_LABEL_ARGS[@]}" --out_png "$OUT/${r}_mitsuba_superpoint${LABEL_SUFFIX}.png" --max_points "$MAXP" --film_size "$FILM" --spp "$SPP"
    python3 "$HERE/superpoint_visualize_s3dis.py" --data_root "$DATA_ROOT" --room "Area_1/${r}" --mode segment \
      --out_png "$OUT/${r}_mitsuba_segment_gt.png" --max_points "$MAXP" --film_size "$FILM" --spp "$SPP"
  done
}

mitsuba_perclass_hq() {
  local tag="${PERCLASS_ROOM//\//_}"
  local d="$OUT/${tag}_per_class_sp${LABEL_SUFFIX}_hq"
  echo "==== Mitsuba superpoint_per_class -> $d ===="
  mkdir -p "$d"
  python3 "$HERE/superpoint_visualize_s3dis.py" --data_root "$DATA_ROOT" --room "$PERCLASS_ROOM" --mode superpoint_per_class \
    "${SP_LABEL_ARGS[@]}" --out_dir "$d" --max_points "$MAXP" --film_size "$FILM" --spp "$SPP"
}

blender_batch() {
  local r
  for r in $ROOMS; do
    echo "==== Blender $r ===="
    bash "$HERE/run_blender_superpoints.sh" --data_root "$DATA_ROOT" --room "Area_1/${r}" --mode superpoint \
      "${SP_LABEL_ARGS[@]}" --out_png "$OUT/${r}_blender_superpoint${LABEL_SUFFIX}.png" --max_points 50000
    bash "$HERE/run_blender_superpoints.sh" --data_root "$DATA_ROOT" --room "Area_1/${r}" --mode segment \
      --out_png "$OUT/${r}_blender_segment_gt.png" --max_points 50000
  done
}

if [[ "$DO_MITSUBA" == "1" ]]; then mitsuba_full; fi
if [[ "$DO_MITSUBA_PERCLASS" == "1" ]]; then mitsuba_perclass_hq; fi
if [[ "$DO_BLENDER" == "1" ]]; then blender_batch; fi

echo "DONE out=$OUT"