#!/usr/bin/env bash
# SUM Parts - PointNeXt_bundle dependencies + CUDA extension build
#
# The upstream install.sh assumes python 3.7 / torch 1.12.1 / cu113 and pins
# requirements.txt to versions that no longer resolve on python 3.10.
# This script keeps the same package set but relaxes the pins, and drops
# packages that are not needed for the sumv2 segmentation task:
#   - deepspeed   (not used by the sumv2 training path, heavy build)
#   - mkdocs-*    (docs only)
# It also adds two packages the code needs but requirements.txt omits:
#   - plyfile : imported by both sumv2 dataset loaders
#   - wandb   : main.py imports it at module scope, so it must exist even when
#               wandb.use_wandb=False. Run with WANDB_MODE=disabled.
# NOTE: the chamfer_dist / emd extensions look reconstruction-only, but
# openpoints/models/__init__.py imports .reconstruction unconditionally, so
# segmentation runs need them too. build_ext.sh builds them.
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"

# RTX 3060 == compute capability 8.6. Pin it so nvcc does not build every arch.
export TORCH_CUDA_ARCH_LIST="8.6"
export CUDA_HOME="$CONDA_PREFIX"
export PATH="$CUDA_HOME/bin:$PATH"

echo "=== [1/4] python deps ==="
# setuptools 69.5.1 : subsampling/setup.py imports numpy.distutils, which needs
# distutils.msvccompiler -- removed in setuptools 74.0. The extensions also
# still use `python setup.py install`, dropped in setuptools 80.
pip install --no-cache-dir "setuptools==69.5.1" wheel

pip install --no-cache-dir \
    plyfile \
    scikit-learn \
    ninja \
    easydict \
    PyYAML \
    tensorboard \
    termcolor \
    tqdm \
    multimethod \
    h5py \
    matplotlib \
    pandas \
    shortuuid \
    gdown \
    Cython \
    pyvista \
    wandb \
    trimesh \
    pillow \
    "numpy<2"

echo "=== [2/4] torch-scatter (matched to torch 2.0.1+cu118) ==="
pip install --no-cache-dir torch-scatter \
    -f https://data.pyg.org/whl/torch-2.0.1+cu118.html

echo "=== [3/4] build CUDA extensions ==="
# Delegated to build_ext.sh so a failed compile can be retried on its own.
# That script also pins ninja==1.11.1.1 (torch 2.0 + ninja>=1.12 dies on SIGPIPE).
bash "$(dirname "$0")/build_ext.sh"

echo "POINTNEXT DONE"