# 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