#!/usr/bin/env bash
#
# get_siesta.sh – download, configure, and compile SIESTA with GNU + Fujitsu MPI
#
# Requires the following environment variables (set in ~/.bashrc or exported
# before calling this script):
#   MY_USER     – your Fugaku username
#   MY_PROJECT  – your project code
#   MY_VOLUME   – your storage volume 
#
# Usage examples
#   ./get_siesta.sh                              # default versions (SIESTA 5.4.2, TCSDS 1.2.42)
#   ./get_siesta.sh 5.4.1                        # different SIESTA version
#   ./get_siesta.sh --tcsds-version 1.2.42       # different TCSDS version
#   ./get_siesta.sh 5.4.1 --tcsds-version 1.2.42 # specify both versions
#   ./get_siesta.sh -c                           # clean extracted directory before starting
#   ./get_siesta.sh -r                           # re-download tarball if missing and re-extract
#

set -euo pipefail

###############################################################################
# Helper: display usage
###############################################################################
usage() {
  cat <<EOF
Usage: $0 [version] [-c|--clean] [-r|--refresh] [--tcsds-version VERSION]

  version              SIESTA version to download (default: 5.4.2)
  --tcsds-version VER  TCSDS version to use (default: 1.2.42)
  -c, --clean          Delete the extracted directory if it already exists.
  -r, --refresh        If the directory exists *and* the tarball does not,
                       download the tarball, delete the directory, and extract again.

Only one of -c or -r can be supplied at the same time.
EOF
  exit 1
}

###############################################################################
# Parse command-line arguments
###############################################################################
VERSION="5.4.2"
CLEAN=false
REFRESH=false
versiontcsds="tcsds-1.2.42"

while [[ $# -gt 0 ]]; do
  case "$1" in
    -c|--clean)
      CLEAN=true
      shift
      ;;
    -r|--refresh)
      REFRESH=true
      shift
      ;;
    --tcsds-version)
      if [[ $# -lt 2 ]]; then
        echo "Error: --tcsds-version requires a version number." >&2
        usage
      fi
      versiontcsds="tcsds-$2"
      shift 2
      ;;
    -h|--help)
      usage
      ;;
    -*)
      echo "Unknown option: $1" >&2
      usage
      ;;
    *)
      VERSION="$1"
      shift
      ;;
  esac
done

# Reject simultaneous -c and -r
if $CLEAN && $REFRESH; then
  echo "Error: -c/--clean and -r/--refresh are mutually exclusive." >&2
  exit 1
fi

###############################################################################
# Check required environment variables
###############################################################################
for var in MY_USER MY_PROJECT MY_VOLUME; do
  if [[ -z "${!var:-}" ]]; then
    echo "Error: environment variable ${var} is not set." >&2
    echo "Export it or add it to your ~/.bashrc before running this script." >&2
    exit 1
  fi
done

###############################################################################
# Variables derived from version and environment
###############################################################################
PKG="siesta-${VERSION}"
TARBALL="${PKG}.tar.gz"
URL="https://gitlab.com/siesta-project/siesta/-/releases/${VERSION}/downloads/${TARBALL}"

SPACK_SETUP="/${MY_VOLUME}/mdt1/data/${MY_PROJECT}/${MY_USER}/tmp/spack/share/spack/setup-env.sh"
NETCDF_ENV="/${MY_VOLUME}/mdt1/data/${MY_PROJECT}/${MY_USER}/tmp/netcdf_gcc/netcdf_env.sh"

###############################################################################
# Helper: download tarball
###############################################################################
download_tarball() {
  echo "Downloading ${TARBALL} ..."
  if command -v curl &>/dev/null; then
    curl -L -o "${TARBALL}" "${URL}"
  elif command -v wget &>/dev/null; then
    wget -O "${TARBALL}" "${URL}"
  else
    echo "Error: neither curl nor wget is available." >&2
    exit 1
  fi
}

###############################################################################
# 1. Handle clean / refresh logic
###############################################################################
if $CLEAN && [[ -d "${PKG}" ]]; then
  echo "Removing existing directory ${PKG} (requested with --clean) ..."
  rm -rf "${PKG}"
fi

if $REFRESH && [[ -d "${PKG}" ]] && [[ ! -f "${TARBALL}" ]]; then
  echo "Refresh requested and tarball missing."
  download_tarball
  echo "Removing existing directory ${PKG} ..."
  rm -rf "${PKG}"
fi

###############################################################################
# 2. Ensure tarball is present
###############################################################################
if [[ ! -f "${TARBALL}" ]]; then
  download_tarball
else
  echo "Tarball ${TARBALL} already exists. Skipping download."
fi

###############################################################################
# 3. Extract tarball
###############################################################################
if [[ ! -d "${PKG}" ]]; then
  echo "Extracting ${TARBALL} ..."
  tar -xzf "${TARBALL}"
else
  echo "Directory ${PKG} already exists. Re-extracting ..."
  rm -rf "${PKG}"
  tar -xzf "${TARBALL}"
fi

###############################################################################
# 4. Create build directories
###############################################################################
cd "${PKG}"
build_="_build_gnu_fjmpi"
bin_="_siesta_bin_gnu_fjmpi"
mkdir -p "${build_}" "${bin_}"
echo "Directories ${build_} and ${bin_} are ready in $(pwd)."

###############################################################################
# 5. Load Spack environment
###############################################################################
echo "Loading Spack environment ..."
source "${SPACK_SETUP}"
echo "Spack environment loaded."

###############################################################################
# 6. Load NetCDF environment
###############################################################################
echo "Loading NetCDF environment ..."
source "${NETCDF_ENV}"
echo "NetCDF environment loaded."

###############################################################################
# 7. Load modules
###############################################################################
ml purge
ml lang/${versiontcsds}
echo "Modules purged and TCSDS ${versiontcsds} loaded."
ml list

echo "Loading Spack packages ..."
spack load /f57uyl2  # gcc@12.2.0
spack load /tcvpuhv  # fujitsu-mpi@head %gcc@12.2.0
spack load /tb7qm25  # netlib-scalapack@2.2.2 %gcc@12.2.0
spack load /3mtxbkd  # python@3.13.5 %gcc@12.2.0
spack load /fvktbik  # cmake@3.31.8 %gcc@12.2.0
spack load /ixbcgpp  # ninja@1.12.1 %gcc@12.2.0
spack load /ptmidr7  # py-pip@25.1.1
spack load /klrzwek  # openblas@0.3.21 %gcc@12.2.0
spack load /mvuivif  # flex@2.6.4 %gcc@12.2.0
spack load /unx5qet  # fftw@3.3.10 %gcc@12.2.0
spack load /3btulwm  # bison@3.8.2 %gcc@12.2.0

#spack load /XXXXXXX  # py-ruamel-yaml@0.17.32 ^python /3mtxbkd %gcc@12.2.0
echo "All Spack packages loaded."

###############################################################################
# 8. Resolve library paths
###############################################################################
export OPENBLAS_PATH=$(spack location -i openblas%gcc@12.2.0)
export SCALAPACK_PATH=$(spack location -i netlib-scalapack%gcc@12.2.0)
echo "Library paths resolved."

###############################################################################
# 9. Set MPI interface flag (5.4.2 does not support legacy interface)
###############################################################################
if [ "${VERSION}" == "5.4.2" ]; then
    interface="none"
else
    interface="legacy"
fi

###############################################################################
# 10. Configure SIESTA with CMake
###############################################################################
echo "Starting CMake configuration ..."
FC=gfortran CC=gcc CXX=g++ cmake -S . -B "${build_}" \
        -DCMAKE_Fortran_COMPILER=gfortran \
        -DCMAKE_C_COMPILER=gcc \
        -DCMAKE_CXX_COMPILER=g++ \
        -DCMAKE_INSTALL_PREFIX="${bin_}" \
        -DCMAKE_Fortran_FLAGS="-I/opt/FJSVxtclanga/${versiontcsds}/include/mpi/fujitsu -fopenmp -fallow-argument-mismatch" \
        -DCMAKE_C_FLAGS="-I/opt/FJSVxtclanga/${versiontcsds}/include/mpi/fujitsu -fopenmp" \
        -DCMAKE_CXX_FLAGS="-I/opt/FJSVxtclanga/${versiontcsds}/include/mpi/fujitsu -fopenmp -Wno-stringop-overflow -DMPICH_SKIP_MPICXX -DOMPI_SKIP_MPICXX -DMPI_NO_CPPBIND" \
        -DMPI_Fortran_LIBRARIES="/opt/FJSVxtclanga/${versiontcsds}/lib64/libmpi_usempif08.so;/opt/FJSVxtclanga/${versiontcsds}/lib64/libmpi_mpifh.so;/opt/FJSVxtclanga/${versiontcsds}/lib64/libmpi.so" \
        -DMPI_C_LIBRARIES="/opt/FJSVxtclanga/${versiontcsds}/lib64/libmpi.so" \
        -DMPI_CXX_LIBRARIES="/opt/FJSVxtclanga/${versiontcsds}/lib64/libmpi.so" \
        -DCMAKE_EXE_LINKER_FLAGS="-L/opt/FJSVxtclanga/${versiontcsds}/lib64 -Wl,-rpath,/opt/FJSVxtclanga/${versiontcsds}/lib64" \
        -DMPIEXEC_EXECUTABLE="/vol0004/apps/oss/mpigcc/fjmpi-gcc12/bin/mpiexec" \
        -DMPIEXEC_MAX_NUMPROCS=4 \
        -DBLAS_LIBRARY="${OPENBLAS_PATH}/lib/libopenblas.so" \
        -DLAPACK_LIBRARY="${OPENBLAS_PATH}/lib/libopenblas.so" \
        -DSCALAPACK_LIBRARY="${SCALAPACK_PATH}/lib/libscalapack.so" \
        -DSIESTA_WITH_NETCDF=ON \
        -DSIESTA_WITH_OPENMP=ON \
        -DSIESTA_WITH_MPI_INTERFACES="${interface}" \
        -DCMAKE_BUILD_TYPE=Release \
        -DSIESTA_WITH_ELPA=OFF \
        -DSIESTA_WITH_ELSI=ON \
        -DSIESTA_WITH_DFTD3=OFF \
        -DSIESTA_WITH_FLOOK=OFF

###############################################################################
# 11. Build and install
###############################################################################
if [ $? -eq 0 ]; then
    echo ""
    echo "CMake configuration successful. Building ..."
    cmake --build "${build_}" -j 8

    if [ $? -eq 0 ]; then
        echo ""
        echo "Build successful. Installing ..."
        cmake --install "${build_}"
    else
        echo ""
        echo "ERROR: Build failed."
        exit 1
    fi
else
    echo ""
    echo "ERROR: CMake configuration failed."
    exit 1
fi