# Start from CUDA base image for GPU support and Ubuntu 20.04
FROM nvidia/cuda:11.3.1-cudnn8-runtime-ubuntu20.04

# Avoid interactive prompts, set timezone
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC

# Install system packages
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
      tzdata git cmake build-essential software-properties-common dirmngr \
      curl gnupg python3 python3-dev python3-setuptools python3-pip \
      libblas-dev liblapack-dev r-base libxml2-dev libcurl4-openssl-dev \
      libssl-dev libpcre2-dev libpcre3-dev libbz2-dev liblzma-dev zlib1g-dev \
      libicu-dev && \
    rm -rf /var/lib/apt/lists/*

# (Optional) Install g++-13 from Ubuntu Toolchain PPA
RUN add-apt-repository ppa:ubuntu-toolchain-r/test -y && \
    apt-get update && \
    apt-get install -y --no-install-recommends g++-13 && \
    update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-13 100 && \
    rm -rf /var/lib/apt/lists/*

# Upgrade pip and install Python packages
RUN python3 -m pip install --upgrade pip && \
    python3 -m pip install \
      pybind11 tensorflow torch torchvision torchaudio \
      numpy scipy scikit-learn

# Copy project files
WORKDIR /app
COPY infinity_search.py vp_tree_bind.cpp vp_tree_q.hpp ./

# Build the VP-Tree extension (pybind11 headers already installed)
RUN python3 -m pip install --upgrade pip pybind11 && \
    g++-13 -O3 -ffast-math -flto -Wall -shared -std=c++20 -fPIC \
        -march=native -mavx2 -mfma \
        $(python3 -m pybind11 --includes) \
        vp_tree_bind.cpp \
        -o vp_tree$(python3-config --extension-suffix) && \
    PY_SITE=$(python3 -c "import site; print(site.getsitepackages()[0])") && \
    mv vp_tree*$(python3-config --extension-suffix) "$PY_SITE/"

# Default entrypoint
ENTRYPOINT ["python3", "infinity_search.py"]
