#!/bin/bash
DEEPMD_PATCH_SCRIPT_DIR=$(dirname $(readlink -f "$0"))
DEEPMD_PATCH_ROOT=${DEEPMD_PATCH_SCRIPT_DIR}/../share/deepmd_gromacs_patches
VERSION="NULL"
PATCH_FILES=("CMakeLists.txt" "src/gromacs/mdlib/forcerec.cpp" "src/gromacs/mdlib/sim_util.cpp" "src/gromacs/mdlib/forcerec.h")

MANUAL="\
usage: this is a deepmd-kit patch program for gromacs

arguments:
    -h      print this help and exit
    -d      gromacs root
    -p      patch mode
    -r      revert mode
    -v      gromacs version (only support 2020.2 now)
"

check () {
    echo "- Checking GROMACS root: $1"
    for ((i=0;i<${#PATCH_FILES[*]};i++))
    do
        file=${PATCH_FILES[$i]}
        if [ ! -e $1/${file} ]; then
            echo "- $1/${file} not exist"
            exit 1
        fi
    done
    echo "- Done"
}

check_unpatched () {
    for ((i=0;i<${#PATCH_FILES[*]};i++))
    do
        file=${PATCH_FILES[$i]}
        if [ -e ${GMX_ROOT}/${file}.predp ]; then
            echo "- ERROR: This is a patched gromacs code, please revert first"
            exit 1
        fi
    done
}

check_patched () {
    for ((i=0;i<${#INSTALL_FILES[*]};i++))
    do
        file=${INSTALL_FILES[$i]}
        if [ ! -e ${GMX_ROOT}/${file} ]; then
            echo "- WARNING: ${GMX_ROOT}/${file} not found"
        fi
    done

    for ((i=0;i<${#PATCH_FILES[*]};i++))
    do
        file=${PATCH_FILES[$i]}
        if [ ! -e ${GMX_ROOT}/${file}.predp ]; then
            echo "- ERROR: backup file ${GMX_ROOT}/${file}.predp is lost"
            exit 1
        fi
    done
}

dp_gmx_patch () {
    echo "- Staring DeePMD-kit patch program to GROMACS ${VERSION}"
    echo "- Mode: patch"
    if [ ! -d $1 ]; then
        echo "- ERROR: invalid gromacs root: $1"
        exit 1
    else
        check_unpatched $1
        check $1
        for ((i=0;i<${#PATCH_FILES[*]};i++))
        do
            file=${PATCH_FILES[$i]}
            patch -b -u ${GMX_ROOT}/${file} --suffix=.predp < ${DEEPMD_PATCH_ROOT}/${file}.patch > /dev/null
            echo "- Installing ${GMX_ROOT}/${file}"
            echo "- Backing up ${GMX_ROOT}/${file}.predp"
        done

        for ((i=0;i<${#INSTALL_FILES[*]};i++))
        do
            file=${INSTALL_FILES[$i]}
            cp ${DEEPMD_PATCH_ROOT}/${file} ${GMX_ROOT}/${file}
            echo "- Installing ${GMX_ROOT}/${file}"
        done
        echo "- Finished"
    fi
}

dp_gmx_revert () {
    echo "- Staring DeePMD-kit patch program to GROMACS ${VERSION}"
    echo "- Mode: revert"
    check_patched $1
    if [ ! -d $1 ]; then
        echo "- ERROR: invalid gromacs root: $1"
        exit 1
    else
        for ((i=0;i<${#INSTALL_FILES[*]};i++))
        do
            file=${INSTALL_FILES[$i]}
            rm ${GMX_ROOT}/${file}
            echo "- Removing ${GMX_ROOT}/${file}"
        done

        for ((i=0;i<${#PATCH_FILES[*]};i++))
        do
            file=${PATCH_FILES[$i]}
            patch ${GMX_ROOT}/${file} -R -u < ${DEEPMD_PATCH_ROOT}/${file}.patch > /dev/null
            rm ${GMX_ROOT}/${file}.predp
            echo "- Restoring from ${GMX_ROOT}/${file}.predp"
        done
        echo "- Finished"
    fi
}

check_version () {
    case $1 in
        "NULL") echo "- ERROR: Please specify a version from -v option" && exit 1;;
        "2020.2") echo "- GROMACS version: 2020.2" ;;
        *) echo "- ERROR: Invalid version" && exit 1 ;;
    esac
}


while getopts "hprd:v:" opt
do
    case ${opt} in
        h) echo "${MANUAL}" && exit 0 ;;
        d) GMX_ROOT=${OPTARG};;
        v) VERSION=${OPTARG} && DEEPMD_PATCH_ROOT=${DEEPMD_PATCH_ROOT}/${VERSION} ;;
        p) check_version ${VERSION} && dp_gmx_patch ${GMX_ROOT} ;;
        r) check_version ${VERSION} && dp_gmx_revert ${GMX_ROOT} ;;
        *) echo "- ERROR: Invalid option ${opt}" && exit 1 ;;
    esac
done
