#!/bin/bash

if [[ "$BASH_SOURCE" == "$0" ]]; then
    echo "This script is the top-level bash file being run."
    echo "Setting up the environment for the cluster."
    BUILD_LOCATION=/avfs/backend
    INSTALL_LOCATION=/avfs/projects/ops
    mkdir -p ${BUILD_LOCATION}
    mkdir -p ${INSTALL_LOCATION}
    echo "Build location: ${BUILD_LOCATION}"
    echo "Install location: ${INSTALL_LOCATION}"
else
    echo "This script is not the top-level bash file being run."
fi

./scripts/install_gcc.sh

./scripts/install_env-modules.sh

./scripts/install_openssl.sh

./scripts/stage0/install_qt.sh

./scripts/stage0/install_cmake.sh

./scripts/stage0/install_git.sh

./scripts/stage0/install_llvm.sh

## INSTALL OPENMPI FROM SOURCE
module load gcc/${GCC_VERSION} cmake/${CMAKE_VERSION}

OPENMPI_BUILD_DIR=${BUILD_LOCATION}/openmpi
OPENMPI_INSTALL_DIR=${INSTALL_LOCATION}/openmpi
mkdir -p ${OPENMPI_BUILD_DIR}
mkdir -p ${OPENMPI_INSTALL_DIR}

OPENMPI_VERSION=4.1.1

cd ${OPENMPI_BUILD_DIR}
wget -c https://download.open-mpi.org/release/open-mpi/v${OPENMPI_VERSION%.*}/openmpi-${OPENMPI_VERSION}.tar.gz
tar -xf openmpi-${OPENMPI_VERSION}.tar.gz
cd openmpi-${OPENMPI_VERSION}
mkdir build && cd build

if [ $(command -v srun) ]; then
    echo "Slurm installation found. OpenMPI will be configured with --with-pmi."
    EXTRA_CONFIGURE_FLAGS="--with-pmi"
else
    EXTRA_CONFIGURE_FLAGS=""
fi

pkg_install_dir=${OPENMPI_INSTALL_DIR}/${OPENMPI_VERSION}/${AVFS_COMPILER}/${AVFS_COMPILER_VERSION}

../configure CFLAGS="${CFLAGS}" \
        --prefix=${pkg_install_dir} \
        --libdir="${pkg_install_dir}/lib" \
        --enable-mpi1-compatibility \
        ${EXTRA_CONFIGURE_FLAGS}
make -j$(nproc)
make -j$(nproc) install

## INSTALL MPICH FROM SOURCE
module load gcc/${GCC_VERSION} cmake/${CMAKE_VERSION}

MPICH_BUILD_DIR=${BUILD_LOCATION}/mpich
MPICH_INSTALL_DIR=${INSTALL_LOCATION}/mpich
mkdir -p ${MPICH_BUILD_DIR}
mkdir -p ${MPICH_INSTALL_DIR}

MPICH_VERSION=4.0.3

cd ${MPICH_BUILD_DIR}
wget -c https://mpich.org/static/downloads/${MPICH_VERSION}/mpich-${MPICH_VERSION}.tar.gz
tar -xf mpich-${MPICH_VERSION}.tar.gz

cd mpich-${MPICH_VERSION}
mkdir build && cd build

if ("${FC}" --version | grep -q 'GNU'); then
    compat_flag=$(allowed_gfortran_flags "-fallow-argument-mismatch")
else
    compat_flag=""
fi

pkg_install_dir=${MPICH_INSTALL_DIR}/${MPICH_VERSION}/${AVFS_COMPILER}/${AVFS_COMPILER_VERSION}

../configure \
        --prefix="${pkg_install_dir}" \
        --libdir="${pkg_install_dir}/lib" \
        MPICC="" \
        FFLAGS="${FCFLAGS} ${compat_flag}" \
        FCFLAGS="${FCFLAGS} ${compat_flag}" \
        --without-x \
        --enable-gl=no \
        --with-device=${MPICH_DEVICE}
make -j$(nproc)
make -j$(nproc) install


MPIRUN="${pkg_install_dir}/bin/mpiexec"
MPICC="${pkg_install_dir}/bin/mpicc"
MPICXX="${pkg_install_dir}/bin/mpicxx"
MPIFC="${pkg_install_dir}/bin/mpifort"
MPIFORT="${MPIFC}"
MPIF77="${MPIFC}"
MPICH_CFLAGS="-I'${pkg_install_dir}/include'"
MPICH_LDFLAGS="-L'${pkg_install_dir}/lib' -Wl,-rpath,'${pkg_install_dir}/lib'"

export MPI_MODE="${MPI_MODE}"
export MPIRUN="${MPIRUN}"
export MPICC="${MPICC}"
export MPICXX="${MPICXX}"
export MPIFC="${MPIFC}"
export MPIFORT="${MPIFORT}"
export MPIF77="${MPIF77}"
export MPICH_CFLAGS="${MPICH_CFLAGS}"
export MPICH_LDFLAGS="${MPICH_LDFLAGS}"
export MPICH_LIBS="${MPICH_LIBS}"
export MPI_CFLAGS="${MPICH_CFLAGS}"
export MPI_LDFLAGS="${MPICH_LDFLAGS}"
export MPI_LIBS="${MPICH_LIBS}"
export CP_DFLAGS="\${CP_DFLAGS} IF_MPI(-D__parallel|)"
export CP_CFLAGS="\${CP_CFLAGS} IF_MPI(${MPICH_CFLAGS}|)"
export CP_LDFLAGS="\${CP_LDFLAGS} IF_MPI(${MPICH_LDFLAGS}|)"
export CP_LIBS="\${CP_LIBS} IF_MPI(${MPICH_LIBS}|)"

## INSTALL OPENBLAS FROM SOURCE
OPENBLAS_BUILD_DIR=${BUILD_LOCATION}/openblas
OPENBLAS_INSTALL_DIR=${INSTALL_LOCATION}/openblas
mkdir -p ${OPENBLAS_BUILD_DIR}
mkdir -p ${OPENBLAS_INSTALL_DIR}

OPENBLAS_VERSION=0.3.18

cd ${OPENBLAS_BUILD_DIR}
wget -c https://github.com/OpenMathLib/OpenBLAS/releases/download/v${OPENBLAS_VERSION}/OpenBLAS-${OPENBLAS_VERSION}.tar.gz
tar -xf OpenBLAS-${OPENBLAS_VERSION}.tar.gz
cd OpenBLAS-${OPENBLAS_VERSION}
pkg_install_dir=${OPENBLAS_INSTALL_DIR}/${OPENBLAS_VERSION}
make -j$(nproc) \
          MAKE_NB_JOBS=0 \
          TARGET=${TARGET} \
          NUM_THREADS=128 \
          USE_THREAD=1 \
          USE_OPENMP=1 \
          NO_AFFINITY=1 \
          CC="${CC}" \
          FC="${FC}" \
          PREFIX="${pkg_install_dir}"


make -j$(nproc) \
          MAKE_NB_JOBS=0 \
          TARGET=${TARGET} \
          NUM_THREADS=128 \
          USE_THREAD=1 \
          USE_OPENMP=1 \
          NO_AFFINITY=1 \
          CC="${CC}" \
          FC="${FC}" \
          PREFIX="${pkg_install_dir}" install

## INSTALL FFTW FROM SOURCE
./scripts/install_fftw.sh

## INSTALL HDF5 FROM SOURCE
module load gcc/8.1.0 openmpi/4.1.4 cmake/3.21.3

# Check if MPI is enabled
if [ -n "$MPI_ENABLED" ]; then
    MPI_ENABLED=true
else
    MPI_ENABLED=false
fi

HDF5_BUILD_DIR=${BUILD_LOCATION}/hdf5
HDF5_VERSION=1.12.1

# Use the MPI_ENABLED variable in your script
if [ "$MPI_ENABLED" = true ]; then
    # MPI is enabled, do something
    echo "MPI is enabled"
    HDF5_INSTALL_DIR=${INSTALL_LOCATION}/hdf5-withmpi/${HDF5_VERSION}/${COMPILER}/${COMPILER_VERSION}/${MPI_NAME}/${MPI_VERSION}
else
    # MPI is not enabled, do something else
    echo "MPI is not enabled"
    HDF5_INSTALL_DIR=${INSTALL_LOCATION}/hdf5/${HDF5_VERSION}/${COMPILER}/${COMPILER_VERSION}
fi

## INSTALL NETCDF FROM SOURCE

## INSTALL OPENMC FROM SOURCE
module load gcc/8.1.0 cmake/3.21.3 python/3.11.5 git/2.11.3

OPENMC_BUILD_DIR=${BUILD_LOCATION}/openmc
OPENMC_INSTALL_DIR=${INSTALL_LOCATION}/openmc

mkdir -p ${OPENMC_BUILD_DIR}
mkdir -p ${OPENMC_INSTALL_DIR}

PROJECT_FOLDER=/avfs/projects/openmc
cd ${PROJECT_FOLDER}
python3 -m venv .venv

OPENMC_VERSION=0.14.0

cd ${OPENMC_BUILD_DIR}
git clone --recurse-submodules https://github.com/openmc-dev/openmc.git
cd openmc
git checkout v${OPENMC_VERSION}
mkdir build && cd build

if [ "$MPI_ENABLED" = true ]; then
    # MPI is enabled, do something
    # HDF5_ROOT=/avfs/projects/hdf5-withmpi/gcc/8.1.0/openmpi/4.1.4 \
    # CC=/avfs/projects/openmpi/4.1.4/gcc/8.1.0/bin/mpicc \
    # CXX=/avfs/projects/openmpi/4.1.4/gcc/8.1.0/bin/mpicxx \
    module load hdf5-withmpi openmpi/4.1.4
    HDF5_ROOT=${HDF5_ROOT} CC=mpicc CXX=mpicxx \
    cmake .. -DCMAKE_INSTALL_PREFIX=${OPENMC_INSTALL_DIR}/${OPENMC_VERSION} -DOPENMC_USE_MPI=ON -DOPENMC_USE_OPENMP=ON -DCMAKE_BUILD_TYPE=Release
else
    # MPI is not enabled, do something else
    module load hdf5
    HDF5_ROOT=${HDF5_ROOT}
    cmake .. -DCMAKE_INSTALL_PREFIX=${OPENMC_INSTALL_DIR}/${OPENMC_VERSION} -DOPENMC_USE_MPI=OFF -DOPENMC_USE_OPENMP=ON -DCMAKE_BUILD_TYPE=Release
fi

make -j20
make install
cd ..
source ${PROJECT_FOLDER}/.venv/bin/activate

if [ "$MPI_ENABLED" = true ]; then
    # MPI is enabled, do something
    MPICC=/avfs/projects/openmpi/4.1.4/gcc/8.1.0/bin/mpicc \
        /avfs/projects/python/3.11.5/bin/python3 -m pip install mpi4py
    CC=/avfs/projects/openmpi/4.1.4/gcc/8.1.0/bin/mpicc HDF5_MPI=ON \
        HDF5_DIR=/avfs/projects/hdf5-withmpi/gcc/8.1.0/openmpi/4.1.4 \
        ${PROJECT_FOLDER}/.venv/bin/python3 -m pip install --no-binary=h5py h5py
    ${PROJECT_FOLDER}/.venv/bin/python3 -m pip install .
else
    # MPI is not enabled, do something else
    HDF5_DIR=${HDF5_ROOT}
    ${PROJECT_FOLDER}/.venv/bin/python3 -m pip install .
fi





# HDF5_ROOT=/share/apps/hdf5-withmpi/gcc/8.1.0/openmpi/4.1.4 CC=/share/apps/openmpi/4.1.4/gcc/8.1.0/bin/mpicc CXX=/share/apps/openmpi/4.1.4/gcc/8.1.0/bin/mpicxx cmake .. -DCMAKE_INSTALL_PREFIX=/qfs/people/parl703/openmc -DOPENMC_USE_MPI=ON -DOPENMC_USE_OPENMP=ON -DCMAKE_BUILD_TYPE=Release
# make -j20
# make install
# cd ..
# MPICC=/share/apps/openmpi/4.1.4/gcc/8.1.0/bin/mpicc /share/apps/python/3.11.5/bin/python3 -m pip install mpi4py
# CC=/share/apps/openmpi/4.1.4/gcc/8.1.0/bin/mpicc HDF5_MPI=ON HDF5_DIR=/share/apps/hdf5-withmpi/gcc/8.1.0/openmpi/4.1.4 /rcfs/projects/sherlock/openmc/dwnld/.venv/bin/python3 -m pip install --no-binary=h5py h5py
# /rcfs/projects/sherlock/openmc/dwnld/.venv/bin/python3 -m pip install .