# Configurable TensorFlow setup
ARG TF_SRC_TAG=v2.0.0-rc2
ARG TF_V2
ARG TF_MKL

# Configurable Python setup
ARG USE_PYTHON_3_NOT_2
ARG _DOCKER_SUFFIX=${USE_PYTHON_3_NOT_2:+-py3}
ARG _PY_SUFFIX=${USE_PYTHON_3_NOT_2:+3}
ARG PYTHON=python${_PY_SUFFIX}
ARG PIP=pip${_PY_SUFFIX}

# Build stage
FROM tensorflow/tensorflow:latest-devel${_DOCKER_SUFFIX} AS builder

WORKDIR /tensorflow_src
RUN git checkout ${TF_SRC_TAG}

ARG _TF_CONFIG_V2=${TF_V2:+--config=v2}
ARG _TF_CONFIG_V1=${_TF_CONFIG_V2:---config=opt}
ARG _TF_FLAGS="${TF_MKL:+--config=mkl} ${TF_MONOLITHIC:+--config=monolithic}"

# Build settings. Likely will change between versions
ENV PYTHON_BIN_PATH=/usr/bin/${PYTHON}
ENV USE_DEFAULT_PYTHON_LIB_PATH=1
ENV TF_ENABLE_XLA=1
ENV TF_NEED_OPENCL_SYCL=0
ENV TF_NEED_ROCM=0
ENV TF_NEED_CUDA=0
ENV TF_DOWNLOAD_CLANG=0
ENV TF_NEED_MPI=0
ENV CC_OPT_FLAGS="-march=native -Wno-sign-compare"
ENV TF_SET_ANDROID_WORKSPACE=0

RUN ./configure
RUN bazel build ${_TF_FLAGS} ${_TF_CONFIG_V2}${_TF_CONFIG_V1} //tensorflow/tools/pip_package:build_pip_package
RUN ./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

# Install stage. Copied from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/dockerfiles/dockerfiles/cpu.Dockerfile
FROM ubuntu:18.04 as base

RUN apt-get update && apt-get install -y curl

ENV LANG C.UTF-8

RUN apt-get update && apt-get install -y \
    ${PYTHON} \
    ${PYTHON}-pip

RUN ${PIP} --no-cache-dir install --upgrade \
    pip \
    setuptools

# Some TF tools expect a "python" binary
RUN ln -s $(which ${PYTHON}) /usr/local/bin/python 

COPY --from=builder /tmp/tensorflow_pkg/tensorflow-*.whl .
RUN ${PIP} install tensorflow-*.whl

RUN curl -L https://raw.githubusercontent.com/tensorflow/tensorflow/master/tensorflow/tools/dockerfiles/bashrc -O /etc/bash.bashrc
RUN chmod a+rwx /etc/bash.bashrc