# Patch image: fix compatibility + pre-bake all model weights
# Build: docker build -f Dockerfile.patch -t protein-design-mcp:fixed .
FROM protein-design-mcp:local

# Fix numpy incompatibility (numpy 2.x breaks compiled modules)
RUN pip install --no-cache-dir "numpy<2"

# Fix dgl: install CUDA 11.8 build from DGL wheel index
# PyPI's generic dgl is CPU-only; need cu118 build for GPU ops
RUN pip uninstall -y dgl && \
    pip install --no-cache-dir "dgl==1.1.3" -f https://data.dgl.ai/wheels/cu118/repo.html

# Fix e3nn (pin to version compatible with torch 2.0.x, before torch.compiler)
RUN pip install --no-cache-dir "e3nn==0.5.1"

# Install rfdiffusion as editable package for importability
RUN cd /opt/RFdiffusion && pip install --no-cache-dir -e .

# Install openfold v1.0.1 (compatible with fair-esm 2.0.0 / ESMFold v1)
# Clone at specific tag and add to PYTHONPATH (no C++ build needed)
RUN git clone --depth 1 --branch v1.0.1 https://github.com/aqlaboratory/openfold.git /opt/openfold \
    && pip install --no-cache-dir biopython dm-tree ml-collections modelcif einops 2>/dev/null || true
# Prevent deepspeed from trying to compile CUDA ops (no nvcc in runtime image)
ENV DS_BUILD_OPS=0
ENV PYTHONPATH="/opt/openfold:${PYTHONPATH}"

# Disable JIT GPU fusion for newer GPU architectures (sm_89+ like L40S)
ENV PYTORCH_JIT_USE_NNC_NOT_NVFUSER=1

# =============================================================================
# Pre-bake model weights (avoids runtime downloads)
# NOTE: /models is a VOLUME in base image, so we use /opt/weights instead
# =============================================================================

# --- RFdiffusion weights (COPY from build context; IPD server URLs are 404) ---
COPY weights/Complex_base_ckpt.pt /opt/RFdiffusion/models/Complex_base_ckpt.pt
COPY weights/Base_ckpt.pt /opt/RFdiffusion/models/Base_ckpt.pt

# --- ESM2 650M PLM ---
RUN mkdir -p /opt/weights/esm/hub/checkpoints && \
    wget -q --show-progress -O /opt/weights/esm/hub/checkpoints/esm2_t33_650M_UR50D.pt \
      https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt && \
    wget -q --show-progress -O /opt/weights/esm/hub/checkpoints/esm2_t33_650M_UR50D-contact-regression.pt \
      https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50D-contact-regression.pt

# --- ProteinMPNN weights ---
RUN mkdir -p /opt/weights/ProteinMPNN/vanilla_model_weights && \
    wget -q --show-progress -O /opt/weights/ProteinMPNN/vanilla_model_weights/v_48_020.pt \
      https://github.com/dauparas/ProteinMPNN/raw/main/vanilla_model_weights/v_48_020.pt && \
    ln -sf /opt/weights/ProteinMPNN/vanilla_model_weights/v_48_020.pt \
      /opt/ProteinMPNN/vanilla_model_weights/v_48_020.pt

# --- ColabFold / AlphaFold2 params ---
RUN mkdir -p /opt/weights/colabfold && \
    python -c "from colabfold.download import download_alphafold_params; download_alphafold_params('/opt/weights/colabfold')" \
    || ( wget -q --show-progress -O /tmp/af2_params.tar \
         https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar && \
         tar xf /tmp/af2_params.tar -C /opt/weights/colabfold && \
         rm /tmp/af2_params.tar )

# --- OpenMM for energy minimization (open source, includes AMBER force fields) ---
RUN pip install --no-cache-dir openmm

# Override env vars to point to pre-baked weights
ENV TORCH_HOME=/opt/weights/esm
ENV COLABFOLD_WEIGHTS_DIR=/opt/weights/colabfold
ENV SKIP_MODEL_DOWNLOAD=true