# Run the following commands in order:
#
# docker build --tag circuit_training:core -f tools/docker/ubuntu_circuit_training tools/docker/
#
# The docker can also be built with tf-agents-nightly, which can be useful if
# Circuit Training head is using APIs not yet in tf-agents stable.
#
# docker build --build-arg tf_agents_version=tf-agents-nightly[reverb] \
#   --tag circuit_training:core -f tools/docker/ubuntu_circuit_training tools/docker/
#
# Smoke test the docker
# docker run -it --rm -v $(pwd):/workspace --workdir /workspace circuit_training:core bash
# python3 -m circuit_training.environment.environment_test
#
# For GPUs. devel version is needed to support xla.
# docker build --build-arg base_image=nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 \
#   --tag circuit_training:core -f tools/docker/ubuntu_circuit_training tools/docker/
ARG base_image="ubuntu:20.04"
FROM $base_image


LABEL maintainer="tobyboyd@google.com"
# Redeclare `base_image` because args declared before `FROM` cannot be
# referenced below.
ARG base_image="ubuntu:20.04"
# Supports setting up a single version of python.
ARG python_version="python3.9"
ARG tf_agents_version="tf-agents[reverb]"
ARG dreamplace_version="dreamplace_python3.9.tar.gz"
ARG placement_cost_binary="plc_wrapper_main"
ARG APT_COMMAND="apt-get -o Acquire::Retries=3 -y"


# Stops tzdata from asking about timezones and blocking install on user input.
ARG DEBIAN_FRONTEND=noninteractive
ENV TZ=America/Los_Angeles

# Installs basics including add-apt.
RUN ${APT_COMMAND} update && ${APT_COMMAND} install -y --no-install-recommends \
        software-properties-common \
        curl \
        tmux \
        vim \
        less

# Adds repository to pull versions of python from.
RUN add-apt-repository ppa:deadsnakes/ppa

# Installs desired python version if not aready installed and then cleans up apt.
RUN ${APT_COMMAND} update && ${APT_COMMAND} install -y --no-install-recommends \
        $python_version-dev \
        $python_version-distutils


# Downloads the placement cost utility binary nto /usr/local/bin.
RUN curl https://storage.googleapis.com/rl-infra-public/circuit-training/placement_cost/${placement_cost_binary} \
     -o  /usr/local/bin/plc_wrapper_main

RUN chmod 555 /usr/local/bin/plc_wrapper_main

RUN curl -O https://bootstrap.pypa.io/get-pip.py

RUN $python_version get-pip.py

# Installs system dependencies needed by DREAMPlace.
RUN apt-get update \
        && apt-get install -y \
            wget \
            flex \
            libcairo2-dev \
            libboost-all-dev \
        && \
            apt-get clean && \
            rm -rf /var/lib/apt/lists/*

# Installs specific versions of packages requested by DREAMPlace requirements.txt
RUN $python_version -mpip install pyunpack>=0.1.2 \
        patool>=1.12 \
        timeout-decorator>=0.5.0 \
        matplotlib>=2.2.2 \
        cairocffi>=0.9.0 \
        pkgconfig>=1.4.0 \
        setuptools>=39.1.0 \
        scipy>=1.1.0 \
        numpy>=1.15.4 \
        torch==1.13.1 \
        shapely>=1.7.0

# Installs TF-Agents and Tensorflow along with tox and pytest to help verify the
# container via unit tests.
RUN $python_version -mpip --no-cache-dir install $tf_agents_version sortedcontainers tox pytest

# Installs CuDNN to match what TensorFlow needs, because torch 1.13.1 installs older CuDNN.
# This does not impact DREAMPlace because it is not compiled for GPU.
RUN case $base_image in \
       *"nvidia"*) $python_version -mpip install --force-reinstall nvidia-cudnn-cu11~=8.6.0.0 ;; \
       *       ) echo "Skip installing CuDNN. This is a CPU only image!!!" ;; \
    esac 

RUN mkdir -p /dreamplace
RUN curl https://storage.googleapis.com/rl-infra-public/circuit-training/dreamplace/$dreamplace_version -o /dreamplace/dreamplace.tar.gz
RUN tar xzf /dreamplace/dreamplace.tar.gz -C /dreamplace/

# Sets the python path so we can find Placer with `import dreamplace.Placer`
# This also needs to put all of DREAMPlace at the root because DREAMPlace python
# is not setup like a package with imports like `dreamplace.Param`.
ENV PYTHONPATH "${PYTHONPATH}:/dreamplace:/dreamplace/dreamplace"

RUN $python_version -m pip freeze

CMD ["/bin/bash"]
