#!/usr/bin/env bash
set -euo pipefail

REPO=${REPO:-/mnt/zsn/zsn_workspace/3DEditFormer}
BENCH=${BENCH:-/mnt/zsn/zsn_workspace/VoxHammer/Edit3D-Bench}
GT=${GT:-$BENCH/data}
WORK=${WORK:-$REPO/work_dirs/edit3d_bench}
PRE=${PRE:-$WORK/trellis_preprocess}
PRED=${PRED:-$WORK/pred_results}
LOGS=${LOGS:-$WORK/logs}
PY=${PY:-/mnt/zsn/miniconda3/envs/vinedresser3d/bin/python}
BLENDER=${BLENDER:-/opt/blender-4.2.19-linux-x64/blender}
TRELLIS=${TRELLIS:-/mnt/zsn/ckpts/TRELLIS-image-large}
NUM_VIEWS=${NUM_VIEWS:-16}
RESOLUTION=${RESOLUTION:-256}
ENGINE=${ENGINE:-CYCLES}
CUDA_DEVICE=${CUDA_DEVICE:-0}

mkdir -p "$PRE" "$PRED" "$LOGS"
cd "$REPO"

echo "[1/7] Prepare renders and voxels: views=$NUM_VIEWS resolution=$RESOLUTION engine=$ENGINE"
"$PY" scripts/edit3d_bench_prepare_latent_inputs.py \
  --bench_data_root "$GT" \
  --output_dir "$PRE" \
  --blender_path "$BLENDER" \
  --num_views "$NUM_VIEWS" \
  --resolution "$RESOLUTION" \
  --engine "$ENGINE" \
  --max_workers 1

echo "[2/7] Refresh metadata"
"$PY" scripts/edit3d_bench_refresh_metadata_status.py --output_dir "$PRE"

echo "[3/7] Extract DINO features"
CUDA_VISIBLE_DEVICES="$CUDA_DEVICE" "$PY" dataset_toolkits/extract_feature.py \
  --output_dir "$PRE" \
  --model dinov2_vitl14_reg \
  --rank 0 \
  --world_size 1

echo "[4/7] Refresh metadata"
"$PY" scripts/edit3d_bench_refresh_metadata_status.py --output_dir "$PRE"

echo "[5/7] Encode sparse-structure latents"
CUDA_VISIBLE_DEVICES="$CUDA_DEVICE" "$PY" dataset_toolkits/encode_ss_latent.py \
  --output_dir "$PRE" \
  --enc_pretrained "$TRELLIS/ckpts/ss_enc_conv3d_16l8_fp16" \
  --rank 0 \
  --world_size 1

echo "[6/7] Refresh metadata"
"$PY" scripts/edit3d_bench_refresh_metadata_status.py --output_dir "$PRE"

echo "[7/7] Encode SLAT latents"
CUDA_VISIBLE_DEVICES="$CUDA_DEVICE" "$PY" dataset_toolkits/encode_latent.py \
  --output_dir "$PRE" \
  --feat_model dinov2_vitl14_reg \
  --enc_pretrained "$TRELLIS/ckpts/slat_enc_swin8_B_64l8_fp16" \
  --rank 0 \
  --world_size 1

"$PY" scripts/edit3d_bench_refresh_metadata_status.py --output_dir "$PRE"
echo "Done. Latents are under $PRE"