FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1

# Set working directory
WORKDIR /app

# Install system dependencies including build tools for ProDy compilation
RUN apt-get update -y && \
    apt-get install -y --no-install-recommends \
    wget \
    git \
    curl \
    ca-certificates \
    python3.11 \
    python3.11-venv \
    python3.11-dev \
    python3-pip \
    procps \
    build-essential \
    && rm -rf /var/lib/apt/lists/* \
    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 \
    && update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1

# Upgrade pip
RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel

# Clone LigandMPNN repository
RUN git clone https://github.com/dauparas/LigandMPNN.git /app/LigandMPNN

# Set working directory to LigandMPNN
WORKDIR /app/LigandMPNN

# Install Python dependencies
RUN pip3 install --no-cache-dir \
    biopython==1.79 \
    filelock==3.13.1 \
    fsspec==2024.3.1 \
    Jinja2==3.1.3 \
    MarkupSafe==2.1.5 \
    mpmath==1.3.0 \
    networkx==3.2.1 \
    numpy==1.23.5 \
    ProDy==2.4.1 \
    pyparsing==3.1.1 \
    scipy==1.12.0 \
    sympy==1.12 \
    typing_extensions==4.10.0 \
    ml-collections==0.1.1 \
    dm-tree==0.1.8

# Install PyTorch with CUDA support
RUN pip3 install --no-cache-dir \
    torch==2.2.1 \
    --index-url https://download.pytorch.org/whl/cu121

# Download model parameters
RUN mkdir -p /app/LigandMPNN/model_params && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/proteinmpnn_v_48_002.pt -O /app/LigandMPNN/model_params/proteinmpnn_v_48_002.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/proteinmpnn_v_48_010.pt -O /app/LigandMPNN/model_params/proteinmpnn_v_48_010.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/proteinmpnn_v_48_020.pt -O /app/LigandMPNN/model_params/proteinmpnn_v_48_020.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/proteinmpnn_v_48_030.pt -O /app/LigandMPNN/model_params/proteinmpnn_v_48_030.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_v_32_005_25.pt -O /app/LigandMPNN/model_params/ligandmpnn_v_32_005_25.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_v_32_010_25.pt -O /app/LigandMPNN/model_params/ligandmpnn_v_32_010_25.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_v_32_020_25.pt -O /app/LigandMPNN/model_params/ligandmpnn_v_32_020_25.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_v_32_030_25.pt -O /app/LigandMPNN/model_params/ligandmpnn_v_32_030_25.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/solublempnn_v_48_002.pt -O /app/LigandMPNN/model_params/solublempnn_v_48_002.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/solublempnn_v_48_010.pt -O /app/LigandMPNN/model_params/solublempnn_v_48_010.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/solublempnn_v_48_020.pt -O /app/LigandMPNN/model_params/solublempnn_v_48_020.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/solublempnn_v_48_030.pt -O /app/LigandMPNN/model_params/solublempnn_v_48_030.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/per_residue_label_membrane_mpnn_v_48_020.pt -O /app/LigandMPNN/model_params/per_residue_label_membrane_mpnn_v_48_020.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/global_label_membrane_mpnn_v_48_020.pt -O /app/LigandMPNN/model_params/global_label_membrane_mpnn_v_48_020.pt && \
    wget -q https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_sc_v_32_002_16.pt -O /app/LigandMPNN/model_params/ligandmpnn_sc_v_32_002_16.pt && \
    ls -la /app/LigandMPNN/model_params/

# Create wrapper script for ligandmpnn (using absolute path)
RUN printf '#!/bin/bash\npython /app/LigandMPNN/run.py "$@"\n' > /usr/local/bin/ligandmpnn && \
    chmod +x /usr/local/bin/ligandmpnn

# Create wrapper script for scoring (using absolute path)
RUN printf '#!/bin/bash\npython /app/LigandMPNN/score.py "$@"\n' > /usr/local/bin/ligandmpnn-score && \
    chmod +x /usr/local/bin/ligandmpnn-score

# Create input/output directories
RUN mkdir -p /app/inputs /app/outputs

# Set environment variables for model paths
ENV LIGANDMPNN_MODEL_DIR=/app/LigandMPNN/model_params
ENV PYTHONPATH=/app/LigandMPNN:${PYTHONPATH:-}

WORKDIR /app/LigandMPNN

CMD ["python", "run.py", "--help"]