FROM quay.io/pawsey/pytorch:2.2.0-rocm5.7.3

## Add System Dependencies
# NOTE: Needs libxrender1 ImportError: libXrender.so.1: cannot open shared object file: No such file or directory (from pdbeccdutils.core import ccd_reader) throws the error
RUN apt-get update \
    && DEBIAN_FRONTEND=noninteractive \
    && apt-get install --no-install-recommends -y \
        build-essential \
        git \
        wget \
        libxrender1 \
        kalign \
    curl \
    && rm -rf /var/lib/apt/lists/* \
    && apt-get autoremove -y \
    && apt-get clean

## Set ROCm environment variables
ENV ROCM_RELEASE=5.7.3
ENV ROCM_PATH=opt/rocm-$ROCM_RELEASE
ENV PATH=$ROCM_PATH/bin:$ROCM_PATH/llvm/bin:$PATH
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$ROCM_PATH/lib

## Install some Python dependencies
RUN pip install pytest \
    requests \
    biopandas --no-cache-dir

WORKDIR /app/alphafold

RUN python3 -c "import torch; print(torch.__version__)" || echo "PyTorch is not installed. "

## Clone and install the package + requirements
ARG GIT_TAG=rocm
RUN git clone https://github.com/amorehead/alphafold3-pytorch-lightning-hydra . --branch ${GIT_TAG} \
    && pip install .

RUN python3 -c "import torch; print(torch.__version__)" || echo "PyTorch is not installed. "

# Check kalign
RUN kalign -h

# Install MMseqs2
RUN wget https://mmseqs.com/latest/mmseqs-linux-avx2.tar.gz \
    && tar xvfz mmseqs-linux-avx2.tar.gz

ENV MMSEQS_PATH=/app/alphafold/mmseqs
ENV PATH=$MMSEQS_PATH/bin:$PATH

RUN pwd
RUN ls /app/alphafold/mmseqs/bin
RUN mmseqs -h
