# Base image with CUDA for GPU support
#FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
FROM 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

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 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

# 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 --version && \
#    llvm-as --version

# Working directory for MLIR operations
#WORKDIR /workspace

# Command to keep the container running
CMD ["bash"]