FROM mambaorg/micromamba:1.5.0 AS micromamba
FROM nvidia/cuda:11.8.0-base-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV MAMBA_USER=root
ENV MAMBA_USER_ID=0
ENV MAMBA_USER_GID=0
ENV MAMBA_ROOT_PREFIX="/opt/conda"
ENV MAMBA_EXE="/bin/micromamba"

# Copy micromamba
COPY --from=micromamba "$MAMBA_EXE" "$MAMBA_EXE"
COPY --from=micromamba /usr/local/bin/_activate_current_env.sh /usr/local/bin/_activate_current_env.sh
COPY --from=micromamba /usr/local/bin/_dockerfile_shell.sh /usr/local/bin/_dockerfile_shell.sh
COPY --from=micromamba /usr/local/bin/_entrypoint.sh /usr/local/bin/_entrypoint.sh
COPY --from=micromamba /usr/local/bin/_dockerfile_initialize_user_accounts.sh /usr/local/bin/_dockerfile_initialize_user_accounts.sh
COPY --from=micromamba /usr/local/bin/_dockerfile_setup_root_prefix.sh /usr/local/bin/_dockerfile_setup_root_prefix.sh

RUN /usr/local/bin/_dockerfile_initialize_user_accounts.sh && \
    /usr/local/bin/_dockerfile_setup_root_prefix.sh

SHELL ["/usr/local/bin/_dockerfile_shell.sh"]

# Install system dependencies
RUN apt-get update -y && \
    apt-get install -y --no-install-recommends \
    wget \
    git \
    curl \
    tar \
    unzip \
    build-essential \
    procps \
    && rm -rf /var/lib/apt/lists/*

# Activate micromamba for all subsequent commands
ARG MAMBA_DOCKERFILE_ACTIVATE=1

# Install Python first
RUN micromamba install -y -n base python=3.10 pip -c conda-forge && \
    micromamba clean --all --yes

# Create working directory
WORKDIR /opt/RoseTTAFold-All-Atom

# Clone repository
RUN git clone https://github.com/baker-laboratory/RoseTTAFold-All-Atom.git /opt/RoseTTAFold-All-Atom

# Install SE3Transformer
RUN pip --no-cache-dir install -e ./rf2aa/SE3Transformer --no-deps

# Install CS-BLAST
RUN wget -q https://wwwuser.gwdg.de/~compbiol/data/csblast/releases/csblast-2.2.3_linux64.tar.gz -O csblast-2.2.3.tar.gz && \
    mkdir -p csblast-2.2.3 && \
    tar xf csblast-2.2.3.tar.gz -C csblast-2.2.3 --strip-components=1 && \
    rm csblast-2.2.3.tar.gz

# Install BLAST legacy
RUN wget -q https://ftp.ncbi.nlm.nih.gov/blast/executables/legacy.NOTSUPPORTED/2.2.26/blast-2.2.26-x64-linux.tar.gz && \
    mkdir -p blast-2.2.26 && \
    tar -xf blast-2.2.26-x64-linux.tar.gz -C blast-2.2.26 && \
    cp -r blast-2.2.26/blast-2.2.26/* blast-2.2.26/ 2>/dev/null || true && \
    rm blast-2.2.26-x64-linux.tar.gz

# Install conda environment
COPY environment.yaml /opt/RoseTTAFold-All-Atom/environment.yaml
RUN CONDA_OVERRIDE_CUDA="11.8" micromamba install -y -n base -f /opt/RoseTTAFold-All-Atom/environment.yaml && \
    micromamba clean --all --yes

# Set database environment variables
ENV DB_UR30=/mnt/databases/UniRef30_2020_06/UniRef30_2020_06
ENV DB_BFD=/mnt/databases/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt
ENV BLASTMAT=/opt/RoseTTAFold-All-Atom/blast-2.2.26/data/

# Create wrapper script
RUN echo '#!/bin/bash\n\
cd /opt/RoseTTAFold-All-Atom && python -m rf2aa.run_inference "$@"' > /usr/local/bin/rfaa && \
    chmod +x /usr/local/bin/rfaa

WORKDIR /workdir
ENTRYPOINT ["micromamba", "run", "-n", "base"]
CMD ["rfaa", "--help"]