# Base image with FoundationPose dependencies split into CPU (L1) and GPU (L2)

# Stage 1: CPU-only base with Python deps
FROM docker.io/ubuntu:22.04 AS foundationpose-base-l1

ENV DEBIAN_FRONTEND=noninteractive

# Install system deps needed to build/run python packages
RUN apt-get update && apt-get install -y --no-install-recommends \
    ca-certificates \
    git \
    python3.10 \
    python3-pip \
    build-essential \
    cmake \
    ninja-build \
    libeigen3-dev \
    python3.10-dev \
    libboost-system-dev \
    libboost-program-options-dev \
    pybind11-dev \
    libgl1 \
    libglib2.0-0 \
    libgomp1 \
    libsm6 \
    libxext6 \
    libxrender1 \
    libxkbcommon0 \
    libx11-6 \
    libxrandr2 \
    libxi6 \
    libxinerama1 \
    libxcursor1 \
    libspatialindex-dev \
    && rm -rf /var/lib/apt/lists/*

# Set python as default
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
    update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1

RUN python3 -m pip install --no-cache-dir --upgrade pip

# Core python deps (CPU-safe)
# Keep NumPy <2 for extension compatibility
RUN pip install --no-cache-dir \
    "numpy<2" \
    Pillow>=10.0.0 \
    gradio>=4.0.0 \
    huggingface-hub>=0.20.0 \
    scipy==1.12.0 \
    scikit-image==0.22.0 \
    scikit-learn==1.4.1.post1 \
    kornia==0.7.2 \
    einops==0.7.0 \
    timm==0.9.16 \
    pyyaml==6.0.1 \
    ruamel.yaml==0.18.6 \
    omegaconf==2.3.0 \
    h5py==3.10.0 \
    numba==0.59.1 \
    imageio==2.34.0 \
    joblib==1.3.2 \
    psutil==6.1.1 \
    albumentations==1.4.2 \
    imgaug==0.4.0 \
    seaborn==0.13.2 \
    plotly==5.20.0 \
    bokeh==3.4.0 \
    colorama==0.4.6 \
    GPUtil==1.4.0 \
    simplejson==3.19.2 \
    openpyxl==3.1.2 \
    xlsxwriter==3.2.0 \
    nodejs==0.1.1 \
    jupyterlab==4.1.5 \
    ipywidgets==8.1.2 \
    py-spy==0.3.14 \
    videoio==0.2.8 \
    pypng==0.20220715.0 \
    roma==1.4.4 \
    transformations==2024.6.1 \
    meshcat==0.3.2 \
    webdataset==0.2.86 \
    wandb==0.16.5 \
    g4f==0.2.7.1 \
    objaverse==0.1.7 \
    opencv-python==4.9.0.80 \
    opencv-contrib-python==4.9.0.80 \
    open3d==0.18.0 \
    pyglet==1.5.28 \
    pysdf==0.1.9 \
    trimesh==4.2.2 \
    xatlas==0.0.9 \
    rtree==1.2.0 \
    pyrender==0.1.45 \
    pyOpenGL>=3.1.0 \
    pyOpenGL_accelerate>=3.1.0 \
    pybullet==3.2.6 \
    pycocotools==2.0.7 \
    Panda3D==1.10.14 \
    pin==2.7.0 \
    && pip cache purge

# Stage 2: GPU-enabled base
FROM docker.io/nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS foundationpose-base-l2

ENV DEBIAN_FRONTEND=noninteractive
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}

# Only build for T4 (7.5) - reduces compilation memory by 50%
ENV TORCH_CUDA_ARCH_LIST="7.5"

# Install system deps
RUN rm -f /etc/apt/sources.list.d/cuda*.list /etc/apt/sources.list.d/*.list && \
    apt-get update && apt-get install -y --no-install-recommends --allow-unauthenticated \
    ca-certificates \
    && apt-get clean && rm -rf /var/lib/apt/lists/* && \
    apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3-pip \
    git \
    libgl1 \
    libglib2.0-0 \
    libgomp1 \
    libsm6 \
    libxext6 \
    libxrender1 \
    libxkbcommon0 \
    libx11-6 \
    libxrandr2 \
    libxi6 \
    libxinerama1 \
    libxcursor1 \
    libspatialindex-dev \
    && rm -rf /var/lib/apt/lists/* \
    && apt-get clean

# Set python as default
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
    update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1

# Upgrade pip
RUN python3 -m pip install --no-cache-dir --upgrade pip

# Copy CPU-only python deps from L1
COPY --from=foundationpose-base-l1 /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages
COPY --from=foundationpose-base-l1 /usr/local/bin /usr/local/bin

# Install PyTorch (CUDA 11.8)
RUN pip install --no-cache-dir torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118

# GPU/torch-dependent deps
RUN pip install --no-cache-dir \
    fvcore==0.1.5.post20221221 \
    torchnet==0.0.4 \
    ultralytics==8.0.120 \
    warp-lang==1.0.2 \
    && pip cache purge

# Build deps required for CUDA extensions
RUN apt-get update && apt-get install -y --no-install-recommends \
    cmake \
    build-essential \
    ninja-build \
    libeigen3-dev \
    python3.10-dev \
    libboost-system-dev \
    libboost-program-options-dev \
    pybind11-dev \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

EXPOSE 7860

# ----- L2 build steps to include all installation work -----
# Install SAM + pytorch3d
RUN pip install --no-cache-dir "numpy<2" transformers==4.41.2 \
    && pip install --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu118_pyt210/download.html

# Set MAX_JOBS=1 BEFORE any CUDA compilation to limit memory usage
ENV MAX_JOBS=1

# Clone FoundationPose source
RUN git clone --depth 1 https://github.com/NVlabs/FoundationPose.git /app/FoundationPose \
    && cd /app/FoundationPose/bundlesdf/mycuda \
    && sed -i 's/-std=c++14/-std=c++17/g' setup.py

# Build CPU-only C++ code
WORKDIR /app/FoundationPose
RUN cd mycpp && mkdir -p build && cd build && cmake .. && make

# Download model weights (246MB)
WORKDIR /app
COPY download_weights.py ./download_weights.py
RUN python3 download_weights.py

# Install nvdiffrast (CUDA rasterizer) - needs GPU, build here
RUN git clone --depth 1 https://github.com/NVlabs/nvdiffrast.git /tmp/nvdiffrast \
    && cd /tmp/nvdiffrast \
    && python3 setup.py build_ext --inplace
RUN python3 -c "import shutil, sysconfig, glob; from pathlib import Path; site=Path(sysconfig.get_paths()['purelib']); src=Path('/tmp/nvdiffrast/nvdiffrast'); dst=site/'nvdiffrast'; shutil.rmtree(dst, ignore_errors=True); shutil.copytree(src, dst); so_files=(glob.glob('/tmp/nvdiffrast/_nvdiffrast_c*.so') + glob.glob('/tmp/nvdiffrast/nvdiffrast/_nvdiffrast_c*.so') + glob.glob('/tmp/nvdiffrast/build/lib.*/*_nvdiffrast_c*.so')); [shutil.copy2(p, site) for p in so_files]"
RUN python3 -c "import sysconfig; from pathlib import Path; site=Path(sysconfig.get_paths()['purelib']); dist=site/'nvdiffrast-0.0.0.dist-info'; dist.mkdir(exist_ok=True); (dist/'METADATA').write_text('Metadata-Version: 2.1\\nName: nvdiffrast\\nVersion: 0.0.0\\n'); (dist/'WHEEL').write_text('Wheel-Version: 1.0\\nGenerator: manual\\nRoot-Is-Purelib: false\\nTag: py3-none-any\\n'); (dist/'top_level.txt').write_text('nvdiffrast\\n'); (dist/'RECORD').write_text('')"
RUN python3 -c "import nvdiffrast.torch"
RUN rm -rf /tmp/nvdiffrast

# Build CUDA extensions (requires GPU)
WORKDIR /app/FoundationPose
RUN cd bundlesdf/mycuda && pip install . --no-build-isolation