#!/usr/bin/env bash
# SUM Parts - build only the CUDA extensions (rerunnable)
#
# Split out of setup_pointnext.sh so a failed compile can be retried without
# reinstalling every python dependency.
#
# ninja note: torch 2.0's cpp_extension shells out to `ninja -v` and reads its
# output through a pipe. ninja >= 1.12 dies with SIGPIPE there, which torch
# surfaces only as the useless "Error compiling objects for extension".
# Pinning ninja to 1.11.x avoids it.
set -euo pipefail

CONDA_ROOT="$HOME/miniconda3"
ENV_NAME="sumparts"
REPO="$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle"

source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"

export TORCH_CUDA_ARCH_LIST="8.6"   # RTX 3060
export CUDA_HOME="$CONDA_PREFIX"
export PATH="$CUDA_HOME/bin:$PATH"
export MAX_JOBS="${MAX_JOBS:-4}"

echo "=== pin ninja to 1.11.x ==="
pip install --no-cache-dir "ninja==1.11.1.1"
ninja --version

# subsampling/setup.py imports numpy.distutils, which imports
# distutils.msvccompiler. setuptools removed that module in 74.0, so anything
# newer dies with "No module named 'distutils.msvccompiler'".
echo "=== pin setuptools to 69.5.1 (numpy.distutils needs it) ==="
pip install --no-cache-dir "setuptools==69.5.1"
python -c "import setuptools; print('setuptools', setuptools.__version__)"

# Skip extensions that already import cleanly, so a retry does not redo a
# 5-minute nvcc pass that already succeeded. FORCE=1 rebuilds everything.
have() { python -c "import $1" 2>/dev/null; }

# NOTE: pointnet2_batch's setup.py names the extension pointnet2_batch_cuda,
# not pointnet2_cuda as in upstream PointNet++ forks.
if [ "${FORCE:-0}" = "1" ] || ! have pointnet2_batch_cuda; then
    echo "=== build pointnet2_batch ==="
    cd "$REPO/openpoints/cpp/pointnet2_batch"
    rm -rf build ./*.egg-info dist
    python setup.py install
else
    echo "=== skip pointnet2_batch (already importable) ==="
fi

echo "=== build subsampling ==="
cd "$REPO/openpoints/cpp/subsampling"
rm -rf build ./*.so
python setup.py build_ext --inplace

if [ "${FORCE:-0}" = "1" ] || ! have pointops_cuda; then
    echo "=== build pointops ==="
    cd "$REPO/openpoints/cpp/pointops"
    rm -rf build ./*.egg-info dist
    python setup.py install
else
    echo "=== skip pointops (already importable) ==="
fi

# chamfer_dist and emd only matter for reconstruction tasks, but
# openpoints/models/__init__.py imports .reconstruction unconditionally, which
# imports chamfer_dist -- so even a pure segmentation run fails at import time
# without them. They are not optional in practice.
if [ "${FORCE:-0}" = "1" ] || ! have chamfer; then
    echo "=== build chamfer_dist ==="
    cd "$REPO/openpoints/cpp/chamfer_dist"
    rm -rf build ./*.egg-info dist
    python setup.py install
else
    echo "=== skip chamfer_dist (already importable) ==="
fi

# emd/setup.py names the package emd_ext and the extension emd_cuda; plain
# `emd` is only the python-level alias in openpoints/cpp/emd/__init__.py.
if [ "${FORCE:-0}" = "1" ] || ! have emd_cuda; then
    echo "=== build emd ==="
    cd "$REPO/openpoints/cpp/emd"
    rm -rf build ./*.egg-info dist
    python setup.py install
else
    echo "=== skip emd (already importable) ==="
fi

echo "=== verify ==="
cd "$REPO"
python - <<'PY'
import sys
print("torch :", __import__("torch").__version__,
      "| cuda", __import__("torch").cuda.is_available())
ok = True
for m in ["pointnet2_batch_cuda", "pointops_cuda", "chamfer", "emd_cuda",
          "torch_scatter", "plyfile"]:
    try:
        __import__(m)
        print(f"{m:22s}: OK")
    except Exception as e:
        ok = False
        print(f"{m:22s}: FAIL {type(e).__name__}: {e}")
try:
    from openpoints.cpp.subsampling import grid_subsampling  # noqa: F401
    print(f"{'grid_subsampling':22s}: OK")
except Exception as e:
    ok = False
    print(f"{'grid_subsampling':22s}: FAIL {type(e).__name__}: {e}")
sys.exit(0 if ok else 1)
PY

echo "BUILD DONE"