# Use the base image with CUDA and PyTorch
FROM kom4cr0/cuda11.7-pytorch1.13-mamba1.1.1:1.1.1

# Install NVIDIA Container Toolkit (necessary for GPU support)
RUN apt-get update && apt-get install -y \
    nvidia-container-runtime \
    python3 \
    python3-pip \
    && rm -rf /var/lib/apt/lists/*

# Install necessary dependencies and configure NVIDIA repository
RUN apt-get update && apt-get install -y \
    curl \
    gnupg \
    lsb-release \
    sudo \
    && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
    && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
        sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
        tee /etc/apt/sources.list.d/nvidia-container-toolkit.list \
    && sed -i -e '/experimental/ s/^#//g' /etc/apt/sources.list.d/nvidia-container-toolkit.list \
    && apt-get update

# Install NVIDIA Container Toolkit
RUN apt-get install -y nvidia-container-toolkit

# Set the environment variables for CUDA
ENV PATH=/usr/local/cuda-11.7/bin:$PATH
ENV LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64:$LD_LIBRARY_PATH

# Set the runtime for GPU (requires NVIDIA runtime to be installed on the host machine)
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

# Set working directory to /projects
WORKDIR /project

# Install necessary Python dependencies
# Uncomment and modify the next lines as per your project requirements
COPY requirements.txt requirements.txt
RUN pip3 install -r requirements.txt

# Run your Python script
CMD ["python3", "main.py"]