#!/bin/bash

# this script is adapted from https://mace-docs.readthedocs.io/en/latest/guide/lammps.html
# to install LAMMPS on Perlmutter

# Clone the LAMMPS repository if it does not exist
if [ ! -d "lammps" ]; then
    git clone --branch=mace --depth=1 https://github.com/ACEsuit/lammps
else
    echo "LAMMPS directory already exists, skipping clone."
fi

# Download and extract libtorch if it has not been done
if [ ! -d "libtorch-gpu" ]; then
    wget https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.2.0%2Bcu121.zip
    unzip libtorch-shared-with-deps-2.2.0+cu121.zip
    rm libtorch-shared-with-deps-2.2.0+cu121.zip
    mv libtorch libtorch-gpu
else
    echo "libtorch-gpu directory already exists, skipping download and extraction."
fi

# Load necessary modules

module load cmake
module load cudatoolkit
module load craype-accel-nvidia80
module load cray-fftw
module load PrgEnv-intel

module -t list

export MPICH_GPU_SUPPORT_ENABLED=1

# Build LAMMPS
cd lammps
rm -rf build
mkdir -p build
cd build

# NERSC Perlmutter recommanded flags https://docs.nersc.gov/applications/lammps/#building-lammps-from-source
# MACE flags https://mace-docs.readthedocs.io/en/latest/guide/lammps.html
cmake \
    -D CMAKE_BUILD_TYPE=Release \
    -D CMAKE_Fortran_COMPILER=ftn -D CMAKE_C_COMPILER=cc \
    -D MPI_C_COMPILER=cc -D MPI_CXX_COMPILER=CC -D LAMMPS_EXCEPTIONS=ON \
    -D BUILD_SHARED_LIBS=ON -D PKG_KOKKOS=yes -D Kokkos_ARCH_AMPERE80=ON -D Kokkos_ENABLE_CUDA=yes \
    -D PKG_MANYBODY=ON -D PKG_MOLECULE=ON -D PKG_KSPACE=ON -D PKG_REPLICA=ON -D PKG_ASPHERE=ON \
    -D PKG_RIGID=ON -D PKG_MPIIO=ON \
    -D CMAKE_POSITION_INDEPENDENT_CODE=ON -D CMAKE_EXE_FLAGS="-dynamic" \
    -D CMAKE_INSTALL_PREFIX=$(pwd) \
    -D BUILD_MPI=ON \
    -D CMAKE_PREFIX_PATH=$(pwd)/../../libtorch-gpu \
    -D PKG_ML-MACE=ON \
    -D CMAKE_CXX_STANDARD=17 \
    -D CMAKE_CXX_STANDARD_REQUIRED=ON \
    -D CMAKE_CXX_COMPILER=$(pwd)/../lib/kokkos/bin/nvcc_wrapper \
    ../cmake

make -j 16
make install