# Use the BioNeMo Framework image as the base
FROM nvcr.io/nvidia/clara/bionemo-framework:2.7.1 AS openfold-bionemo-image

# Switch to root user to perform system-level operations
USER root

# Install build dependencies
RUN apt-get update && \
    apt-get install -y \
    build-essential \
    ninja-build \
    git \
    git-lfs \
    wget \
    openssh-client && \
    rm -rf /var/lib/apt/lists/*

# Clone the openfold repository from GitHub
RUN git lfs install && \
    git clone -b nv_upstream_trt_cuequivariance --single-branch https://github.com/borisfom/openfold.git /openfold2

ENV PYTHONPATH="/openfold2"

# Set the working directory
WORKDIR /openfold2

# Install uv
RUN pip install uv

RUN pip install nvidia-ml-py

# Install cuequivariance and its CUDA operations
RUN pip install cuequivariance_torch==0.8.0 && \
    pip install cuequivariance-ops-torch-cu12==0.8.0

RUN uv pip install --upgrade pip && \
    uv pip install --no-cache-dir wheel setuptools && \
    uv pip install --no-cache-dir --no-build-isolation -e . && \
    uv pip install --no-cache-dir \
        biopython==1.86 \
        mdtraj==1.11.0 \
        modelcif==1.6 \
        ml_collections==1.1.0 \
        bionemo-moco==0.0.2.2 \
        huggingface-hub \
        safetensors \
        transformers \
        gradio


RUN touch /openfold2/openfold/__init__.py

# Download stereo_chemical_props.txt
RUN wget https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt && \
    mv stereo_chemical_props.txt /openfold2/openfold/resources/

# Verify the installation and Python path
RUN python -c "import sys; print(sys.path)" && \
    python -c "import openfold; print(openfold.__file__)"

WORKDIR /app

COPY app.py .
COPY PepTron ./PepTron

# Create persistent storage directory for checkpoint
RUN mkdir -p /data/checkpoints && chmod 777 /data/checkpoints

ENV GRADIO_SERVER_NAME="0.0.0.0"
ENV GRADIO_SERVER_PORT="7860"
ENV PEPTRON_CHECKPOINT_DIR="/data/checkpoints"
ENV PEPTRON_AUTO_DOWNLOAD="true"

EXPOSE 7860

ENV OMP_NUM_THREADS=8

RUN chown -R 1000:1000 /app
USER 1000

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