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

cd "$(dirname "$0")/.." || exit 1

PYTHON_BIN="${PYTHON_BIN:-/mnt/data/AODUOLI/miniconda_envs/Aoduo/bin/python}"
CUDA_ROOT="${CUDA_ROOT:-/mnt/data/AODUOLI/miniconda_envs/Aoduo}"
CONFIG_REL="${CONFIG_REL:-configs/s3dis/semseg-pt-v3m1-0-base_qat_long_sp_after_pool.py}"
EXP_NAME="${EXP_NAME:-s3dis_ptv3_qat_sp_after_pool}"
NPROC_PER_NODE="${NPROC_PER_NODE:-4}"
DIST_URL="${DIST_URL:-auto}"

export CUDA_HOME="$CUDA_ROOT"
export PATH="$CUDA_HOME/bin:$PATH"
export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-8.6}"
export PYTHONPATH="$(pwd):${PYTHONPATH:-}"
export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}"

NVIDIA_INCLUDES=$(find "$CUDA_ROOT/lib/python3.11/site-packages/nvidia" -maxdepth 3 -type d | grep include | paste -sd: -)
export CPATH="$CUDA_HOME/targets/x86_64-linux/include:$CUDA_HOME/include:${NVIDIA_INCLUDES:-}"
export LIBRARY_PATH="$CUDA_HOME/targets/x86_64-linux/lib:$CUDA_HOME/lib"
export LD_LIBRARY_PATH="$CUDA_HOME/targets/x86_64-linux/lib:$CUDA_HOME/lib:${LD_LIBRARY_PATH:-}"

echo "Using python: $PYTHON_BIN"
echo "Config: $CONFIG_REL"
echo "Experiment: $EXP_NAME"
echo "GPUs: $NPROC_PER_NODE"
echo "Dist URL: $DIST_URL"

"$PYTHON_BIN" tools/train.py \
  --num-gpus "$NPROC_PER_NODE" \
  --dist-url "$DIST_URL" \
  --config-file "$CONFIG_REL" \
  --options save_path="exp/s3dis/${EXP_NAME}"