#FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
#FROM tensorflow/tensorflow:nightly-gpu
#FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04
FROM nvidia/cuda:11.8.0-devel-ubuntu22.04

ARG USER
ARG GROUP
ARG UID
ARG GID

RUN groupadd --gid ${GID} ${GROUP}
RUN useradd --shell /bin/bash --uid ${UID} --gid ${GID} --create-home ${USER}
RUN mkdir /wd
RUN chown ${USER}:${GROUP} /wd

# SYSTEM INITIALIZATION
RUN apt update; apt install python3 python3-pip -y
#RUN ln -s /usr/bin/pip3 /usr/bin/pip
RUN pip install --upgrade pip
RUN pip install jupyterlab
RUN pip install tensorflow==2.12.0
RUN pip install pandas
RUN pip install scikit-learn
RUN pip install matplotlib
RUN pip install keras

RUN apt install software-properties-common wget curl git -y
ENV OS=ubuntu22.04

RUN wget https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-${OS}.pin; \
mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600; \
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/3bf863cc.pub; \
add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/ /"

ENV cuda_version=cuda11.8 
ENV cudnn_version=8.8.1.*
RUN apt-get update; \
apt-get install libcudnn8=${cudnn_version}-1+${cuda_version} -y; \
apt-get install libcudnn8-dev=${cudnn_version}-1+${cuda_version} -y

RUN pip install dvc[webdav]

USER ${USER}

# USER INITIALIZATION
# RUN ...

SHELL ["/bin/bash", "--login", "-i", "-c"]
ENV SHELL=bash
WORKDIR /wd