activate

#!/bin/bash
#
# Activate the Spack environment for developing NESO and set some
# environment variables. This script should be sourced to activate
# the environment. You can then run the "deactivate" function to
# return to your previous setup.
#

NESO_DEV_DIR=$(realpath $(dirname "$BASH_SOURCE"))

deactivate() {
  # Relink build directories
  setup
  # End the monitoring process and any `watch` processes it launched
  if [[ ! -z ${__MONITOR_PROCESS} ]]
  then
      kill $__MONITOR_PROCESS
  fi
  ps -ef | grep "spack find -l neso nektar neso-particles" | grep -v grep | awk '{print $2}' | xargs kill -9

  # Tidy up environment variables
  export I_MPI_FABRICS=$__ORIGINAL_I_MPI_FABRICS
  export SYCL_DEVICE_FILTER=$__ORIGINAL_SYCL_DEVICE_FILTER
  unset __ORIGINAL_I_MPI_FABRICS
  unset __ORIGINAL_SYCL_DEVICE_FILTER
  unset __MONITOR_PROCESS
  unset NESO_DEV_DIR

  # Undeclare all the functions in this file
  unset -f cleanup
  unset -f cleanup-deprecated-builds
  unset -f create-link
  unset -f deactivate
  unset -f gen-build-link-maps
  unset -f link-builds
  unset -f link-nektar-builds
  unset -f link-neso-builds
  unset -f link-neso-particles-builds
  unset -f monitor-build-dirs
  unset -f setup

  spack env deactivate
}

# Deletes build directory links for hashes which are no longer
# installed in the environment. First argument is the package name and
# second is the directory in which to cleanup links.
cleanup-deprecated-builds() {
    local pkg_name="$1"
    local pkg_root_dir="$2"
    local links_dir="$pkg_root_dir/builds"

    # Use the build-link-map to compile a list of valid link paths
    local valid_hashes=() valid_link_paths=()
    for entry in $(gen-build-link-maps "$pkg_name" "{compiler.name}-{hash:7}"); do
        IFS="|" read -ra _split <<<"$entry"
        valid_link_paths+=("${links_dir}/${_split[1]}")
        IFS="-" read -ra _split2 <<<"${_split[1]}"
        valid_hashes+=("${_split2[1]}")
    done

    # Remove links from links_dir if they don't appear in valid_link_paths
    link_paths=$(find "$links_dir" -maxdepth 1 -type l -regextype posix-egrep -regex '.*-[a-z0-9]{7}$')
    for link_path in $link_paths; do
        if ! [[ " ${valid_link_paths[*]} " =~ " ${link_path} " ]]
        then
            echo "Cleaning up stale ${pkg_name} build link at ${link_path}"
            rm -Rf "$link_path" > /dev/null
        fi
    done

    # Also cleanup top-level build links created by spack install
    spack_link_paths=$(find "$pkg_root_dir" -type l -regextype posix-egrep -regex "${pkg_root_dir}/build-.*-[a-z0-9]{7}$")
    for link_path in $spack_link_paths; do
        local hash="${link_path:(-7)}"
        if ! [[ " ${valid_hashes[*]} " =~ " ${hash} " ]]
        then
            echo "Cleaning up stale ${pkg_name} SPACK build link at ${link_path}"
            rm -Rf "$link_path" > /dev/null
        fi
    done
}

# Creates a link to the specified target (first argument) with the
# specified name (second argument), unless a link to that target with
# that name already exists. It will overwrite links to different
# targets. Also print a message saying what it is doing.
create-link() {
    target="$1"
    link="$2"
    if ! [[ -L "$link" && $(readlink "$link") == "$target" ]]
    then
        echo "    Linking $(realpath -s --relative-to="$NESO_DEV_DIR" "$link") => $target"
        rm -f "$link" > /dev/null
        ln -s "$target" "$link"
    fi
}

# Find all instances of a package installed in the environment and return a 
# string for each instance of the form "package_name/hash|package_link"
gen-build-link-maps() {
    local package="$1"
    local link_fmt="$2"

    local identifier_fmt="{package.name}/{hash:7}"
    spack find --format "$identifier_fmt|$link_fmt" "$package"
    return $?
}


# Generate soft links to the spack build directories for all instances of a 
# package installed in the environment
link-builds() {
    local package="$1"
    local link_fmt="$2"
    local link_dir="$3"

    local build_dir
    for entry in $(gen-build-link-maps "$package" "$link_fmt"); do
      IFS="|" read -ra _split <<<"$entry"
      build_dir=$(spack location -b "${_split[0]}")
      create-link "${build_dir}" "${link_dir}/${_split[1]}"
    done
    
}

# Creates symlinks to the nektar build directories which have meaningful names
link-nektar-builds() {
    link-builds "nektar" "{compiler.name}-{hash:7}" "${NESO_DEV_DIR}/nektar/builds"
}

# Creates symlinks to the neso-particles build directories which have
# meaningful names
link-neso-particles-builds() {
    link-builds "neso-particles" "{compiler.name}-{hash:7}" "${NESO_DEV_DIR}/neso-particles/builds"
}

# Creates symlinks to the NESO build directories which have meaningful names
link-neso-builds() {
    link-builds "neso" "{compiler.name}-{hash:7}" "${NESO_DEV_DIR}/builds"
}


# Create symlinks to the build directories
setup() {
    echo "Checking links to NESO build directories"
    link-neso-builds
    echo "Checking links to Nektar++ build directories"
    link-nektar-builds
    echo "Checking links to NESO-particles build directories"
    link-neso-particles-builds
}

# Monitor for changes which mean the build-directories have been
# changed and then re-link them. It does not delete old ones, as this
# is potentially risky and expensive. That can be done manually with
# the `cleanup` command.
monitor-build-dirs() {
    trap 
    while :
    do
        watch -n10 -g spack find -l neso nektar neso-particles &> /dev/null && setup &> /dev/null
    done
}

# Convenience command to check and delete any build directories no
# longer being used by the environment.
cleanup() {
    cleanup-deprecated-builds neso "${NESO_DEV_DIR}"
    cleanup-deprecated-builds nektar "${NESO_DEV_DIR}/nektar"
    cleanup-deprecated-builds neso-particles "${NESO_DEV_DIR}/neso-particles"
}

spack env activate -p -d "${NESO_DEV_DIR}"
setup
( monitor-build-dirs ) & __MONITOR_PROCESS=$!

export __ORIGINAL_I_MPI_FABRICS=$I_MPI_FABRICS
export __ORIGINAL_SYCL_DEVICE_FILTER=$SYCL_DEVICE_FILTER

export I_MPI_FABRICS=shm
export SYCL_DEVICE_FILTER=host