# syntax=docker/dockerfile:1-labs

ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax

ARG URLREF_ALPHAFOLD=https://github.com/google-deepmind/alphafold3#main

ARG SRC_PATH_ALPHAFOLD=/opt/alphafold

###############################################################################

## Download source and add auxiliary scripts

###############################################################################

FROM ${BASE_IMAGE} AS mealkit

ARG URLREF_ALPHAFOLD

ARG SRC_PATH_ALPHAFOLD

RUN <<"EOF" bash -ex -o pipefail

git-clone.sh ${URLREF_ALPHAFOLD} ${SRC_PATH_ALPHAFOLD}

EOF

RUN <<"EOF" bash -ex -o pipefail

sed -i '/"jax==.*"/d' ${SRC_PATH_ALPHAFOLD}/pyproject.toml

sed -i '/"jax\[cuda12\]==.*"/d' ${SRC_PATH_ALPHAFOLD}/pyproject.toml

echo "-e file://${SRC_PATH_ALPHAFOLD}" >> /opt/pip-tools.d/requirements-af.in

EOF

RUN cat ${SRC_PATH_ALPHAFOLD}/pyproject.toml

###############################################################################

## Add all the needed scripts

###############################################################################

ADD alphafold_files/generate_af_parameters.py /usr/local/bin

ADD alphafold_files/af3_inference_benchmark.py /usr/local/bin

COPY alphafold_files/input_sequences /opt/alphafold_input_data/input_sequences

COPY alphafold_files/model_parameters.md /opt/alphafold_input_data/model_parameters.md

###############################################################################

## Install accumulated packages from the base image and the previous stage

###############################################################################

FROM mealkit AS final

RUN pip-finalize.sh

WORKDIR ${SRC_PATH_ALPHAFOLD}