# -*- mode: dockerfile -*-

FROM nvidia/cuda:12.4.1-devel-ubuntu22.04

ARG TARGETPLATFORM
RUN echo "Building image for $TARGETPLATFORM"

ARG PYTHON_VERSION=3.11
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && \
    apt-get install -yq wget

RUN case ${TARGETPLATFORM} in \
    "linux/amd64")  CUDA_ARCH=x86_64  ;; \
    "linux/x86_64") CUDA_ARCH=x86_64  ;; \
    "linux/arm64")  CUDA_ARCH=arm64 ;; \
    esac && \
    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/${CUDA_ARCH}/cuda-keyring_1.0-1_all.deb

RUN dpkg -i cuda-keyring_1.0-1_all.deb

RUN apt-get update && apt-get install -yq \
        bison \
        build-essential \
        cmake \
        curl \
        flex \
        git \
        libbz2-dev \
        ninja-build \
        wget

WORKDIR /opt/conda_setup

RUN case ${TARGETPLATFORM} in \
    "linux/amd64")  MINI_ARCH=x86_64  ;; \
    "linux/x86_64") MINI_ARCH=x86_64  ;; \
    "linux/arm64")  MINI_ARCH=aarch64 ;; \
    esac && \
    curl -o miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${MINI_ARCH}.sh && \
    chmod +x miniconda.sh && \
    ./miniconda.sh -b -p /opt/conda && \
    /opt/conda/bin/conda install -y python=$PYTHON_VERSION && \
    /opt/conda/bin/conda clean -ya
ENV PATH /opt/conda/bin:$PATH

RUN python -m pip install --upgrade pip ipython ipdb

COPY . /opt/nle/

WORKDIR /opt/nle

RUN pip install '.[all]'

WORKDIR /workspace

CMD ["/bin/bash"]
