FROM --platform=linux/amd64 ubuntu:latest

ENV PIP_BREAK_SYSTEM_PACKAGES=1

# Install dependencies
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt-get update && apt-get install -y software-properties-common && \
    apt-add-repository ppa:deadsnakes/ppa && \
    apt-get update && apt-get install -y \
    git \
    build-essential \
    cmake \
    python3-pip \
    python3.11-dev python3.11-distutils \
    wget \
    sudo

RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1
RUN update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1

# Install requirement packages
COPY ./requirements.txt /tmp/requirements.txt
RUN --mount=type=cache,target=/root/.cache pip install --upgrade -r /tmp/requirements.txt

# Install llvm
RUN wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt-get -y update && apt-get install -y mlir-18-tools llvm-18

# IREE
RUN pip install \
  --index-url https://download.pytorch.org/whl/test/cpu torch==2.4.1
RUN pip install iree-turbine


# Install pytorch-mlir

#RUN pip install --pre torch-mlir torch \
#  --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
#  -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels

# Add PyTorch-MLIR tools to PATH
#ENV PATH="/root/.local/bin:${PATH}"

# Verify installations
RUN python -c "import torch; print(torch.__version__)" && \
    mlir-opt-18 --version && \
    llvm-as-18 --version

# Working directory for MLIR operations
#WORKDIR /workspace

# Set up a new user named "user"
RUN useradd -m -o -u 1000 user

# Switch to the "user" user
USER user

# Set home to the user's home directory
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:$PATH

# Set the working directory to the user's home directory
WORKDIR $HOME/app

# Copy the current directory contents into the container at $HOME/app setting the owner to the user
COPY --chown=user . $HOME/app

CMD ["python", "main.py"]