# syntax=docker/dockerfile:1
FROM nvidia/cuda:12.3.1-devel-ubuntu22.04

# 定义变量
ARG FFTW_VERSION=3.3.10
ARG GROMACS_VERSION=2023.2
ARG INSTALL_DIR=/opt
ARG GMX_INSTALL_PREFIX=${INSTALL_DIR}/software/gmx
ARG CUDA_TOOLKIT_PATH=/usr/local/cuda
ARG OPENMPI_PATH=${INSTALL_DIR}/software/openmpi

# 使用 Here-docs 安装软件 openmpi 等
RUN <<EOT 
#!/bin/bash
apt-get update 
apt-get install -y cmake curl wget libpmix-dev openmpi-bin libopenmpi-dev
mkdir -p ${INSTALL_DIR}/software 
EOT

#install fftw
WORKDIR ${INSTALL_DIR}/software
RUN <<EOF
#!/bin/bash
wget http://www.fftw.org/fftw-${FFTW_VERSION}.tar.gz
tar -xzvf fftw-${FFTW_VERSION}.tar.gz
cd fftw-${FFTW_VERSION} 
./configure --prefix=${INSTALL_DIR}/software/fftw --enable-sse2 --enable-avx --enable-float --enable-shared --enable-avx2 
make install -j 
echo 'export PATH=$PATH:${INSTALL_DIR}/software/fftw/bin' >> ~/.bashrc 
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${INSTALL_DIR}/software/fftw/lib' >> ~/.bashrc 
EOF

# install gromacs
WORKDIR ${INSTALL_DIR}/software
RUN <<EOF
#!/bin/bash
wget https://ftp.gromacs.org/gromacs/gromacs-${GROMACS_VERSION}.tar.gz 
tar xfz gromacs-${GROMACS_VERSION}.tar.gz 
cd gromacs-${GROMACS_VERSION} 
mkdir build 
cd build 
echo 'export CMAKE_PREFIX_PATH=${INSTALL_DIR}/software/fftw' >> ~/.bashrc 
cmake .. -DCMAKE_INSTALL_PREFIX=${GMX_INSTALL_PREFIX} -DGMX_USE_RDTSCP=ON -DGMX_GPU=CUDA -DCUDA_TOOLKIT_ROOT_DIR=${CUDA_TOOLKIT_PATH} -DGMX_BUILD_OWN_FFTW=ON -DGMX_MPI=ON 
make 
make install -j8
EOF

ENV PATH="${GMX_INSTALL_PREFIX}/bin:${PATH}"

# 设置容器启动时的默认命令
CMD ["gmx_mpi"]