# Use the official PyTorch runtime image that already contains a pre‑built
# PyTorch 2.6.0 stack linked against CUDA 12.4 and cuDNN9.  The base
# image uses Python3.11 via the conda environment under ``/opt/conda`` as
# shown in its Dockerfile snippet, where the CUDA runtime libraries live
# under ``/opt/conda/lib/python3.11/site‑packages/nvidia/cuda_runtime``【53563160074152†L270-L276】.
# Choosing this image ensures that subsequent GPU‑accelerated packages can
# reuse the exact CUDA version and ABI expected by their binary wheels.
FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime

# Avoid interactive prompts during apt operations.
ENV DEBIAN_FRONTEND=noninteractive

##
## Install basic utilities and Python dependencies
##
# - Update the package index and install a few common tools.  The base
#   image is based on Ubuntu22.04 and already includes conda and PyTorch.
# - Upgrade pip to the latest version so it can correctly resolve and
#   download ``manylinux``/CUDA wheels.
RUN apt-get update \
    && apt-get install -y --no-install-recommends \
        ca-certificates \
        wget \
        build-essential \
    && rm -rf /var/lib/apt/lists/*

# Upgrade pip to the latest version.  A modern pip is required to
# recognise local version segments in wheel file names.
RUN pip install --upgrade pip

# Download pre‑built wheels for causal-conv1d and mamba-ssm.  These
# wheels are published on GitHub Releases and are built against
# PyTorch2.6 and CUDA12 with the ``cxx11abiFALSE`` ABI for
# Python3.11【805399342706536†L2946-L2948】【772975914204303†L4546-L4555】.  They are not
# available on PyPI, so pulling them directly avoids triggering a
# source compilation inside the container.
# Save the wheels using their original file names.  Pip relies on the
# wheel filename structure to parse compatibility tags.  Renaming
# wheels (e.g. to ``causal_conv1d.whl``) will cause pip to reject
# them as invalid【805399342706536†L2946-L2948】.
RUN wget -q -O /tmp/causal_conv1d-1.5.1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl \
        "https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.1/causal_conv1d-1.5.1%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
RUN wget -q -O /tmp/mamba_ssm-2.2.4+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl \
        "https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"

# Install the downloaded wheels.  The --no-cache-dir flag avoids
# caching the wheel in pip’s cache, reducing the final image size.
RUN pip install --no-cache-dir /tmp/causal_conv1d-1.5.1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl
RUN pip install --no-cache-dir /tmp/mamba_ssm-2.2.4+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl

COPY requirements.txt .
RUN python3 -m pip install -r requirements.txt

# Clean up pip cache and downloaded wheel files to minimise image size.
RUN rm -rf /root/.cache/pip /tmp/*.whl

# Set the working directory inside the container.  Users can mount
# source code here or run Python interactively.
WORKDIR /workspace

# Use bash as the default entrypoint so that the container drops the
# user into a familiar shell.  Users can override this in ``docker run``.
CMD ["bash"]