#!/usr/bin/env bash

# This build script takes a single argument: the number of tiles.
# This is important for parts of this package that depend on xGPU, for which
# the number of stations ("tiles", in the MWA's case) must be set at compile
# time.
# For more information, see
# 	https://wiki.mwatelescope.org/display/MP/MWAX+Offline+Correlator
if [ -z "$1" ]; then
echo "usage: $0 NUMTILES"
echo "where NUMTILES is the number of MWA tiles to be correlated"
exit
fi

NUMTILES=$1
echo "Compiling for ${NUMTILES}T configuration..."

set -eu

source /pawsey/mwa/software/python3/build_base.sh

# PACKAGE is the name of the software (e.g. ruby, casacore etc.)
PACKAGE=mwax_offline_correlator
PREFIX="${ROOT}/${PACKAGE}"

for DEPENDENCY in mwax_common mwax_cbf mwax_db2fits mwax_user_tools mwax-xGPU
do
	cd "${PREFIX}"
	[ ! -r "${DEPENDENCY}-git" ] && git clone git@github.com:MWATelescope/${DEPENDENCY}.git ${DEPENDENCY}-git

	cd "${PREFIX}/${DEPENDENCY}"-git

	# Checkout the desired branch
	#git checkout master # Don't do this for mwax_db2fits until garrawarla_installion branch has been merged into master
	#git pull
done

# Set the target module version here
VERSION=${NUMTILES}T

echo "Package: ${PACKAGE}"
echo "Version: ${VERSION}"
echo "Prefix: ${PREFIX}"

# Dependencies for mwax_offline_correlator
#module purge

module load gcc/8.3.0
module load cuda
module use /pawsey/mwa/software/python3/modulefiles
module load cfitsio_reentrant
module load psrdada/mwax

module list

set -x

# Compile!

# Root directories for three modules
MWAX_COMMON_ROOT=${PREFIX}/mwax_common-git
MWAX_CBF_ROOT=${PREFIX}/mwax_cbf-git
MWAX_DB2FITS_ROOT=${PREFIX}/mwax_db2fits-git
MWAX_USER_TOOLS_ROOT=${PREFIX}/mwax_user_tools-git
MWAX_XGPU_ROOT=${PREFIX}/mwax-xGPU-git

# Where to put final products
INSTALL_PREFIX=${ROOT}/${PACKAGE}/${VERSION}/${PAWSEY_CLUSTER}

# Don't use cray compilers 
export CC=$(which gcc)
export CXX=$(which g++)


# ----- Compile mwax_xGPU -----

# Original build instructions for mwax_cbf at
# 	https://github.com/MWATelescope/mwax-xGPU
# and
# 	https://wiki.mwatelescope.org/display/MP/MWAX+Offline+Correlator
# The following has been adapted for Garrawarla.

CUDA_DIR=${NVIDIA_CUDA_HOME}
CUDA_ARCH="sm_70"

# Build and install
NFREQUENCY=6400
if [ "${NUMTILES}" -le 192 ]; then
	NTIME=52
	NTIME_PIPE=52
else
	NTIME=56
	NTIME_PIPE=28
fi

cd ${MWAX_XGPU_ROOT}/src
make CUDA_DIR=${CUDA_HOME} clean NSTATION=${NUMTILES}
make CUDA_ARCH=${CUDA_ARCH} CUDA_DIR=${CUDA_HOME} NSTATION=${NUMTILES} NFREQUENCY=${NFREQUENCY} NTIME=${NTIME} NTIME_PIPE=${NTIME_PIPE}
make CUDA_DIR=${CUDA_HOME} install prefix=${INSTALL_PREFIX} NSTATION=${NUMTILES} NFREQUENCY=${NFREQUENCY}


# ----- Compile mwax_cbf -----

# Original build instructions for mwax_cbf at
# 	https://github.com/MWATelescope/mwax_cbf
# The following has been adapted for Garrawarla.

cd "${MWAX_CBF_ROOT}"

nvcc -c -Wno-deprecated-gpu-targets -O3 -gencode arch=compute_70,code=sm_70 \
     -I${MWAX_CBF_ROOT}/mwax-cbf-cuda \
     ${MWAX_COMMON_ROOT}/mwax-cbf-cuda/mwax_cbf_cuda.cu

${CC} -c -O3 \
	  -I${INSTALL_PREFIX}/include -I${PSRDADA_INCLUDE_DIR} -I${MWAX_COMMON_ROOT} -I${MWAX_COMMON_ROOT}/mwax-cbf-cuda \
      ${MWAX_COMMON_ROOT}/mwax_global_defs.c \
      ${MWAX_CBF_ROOT}/mwax-fullcorrelator/src/mwax_db2correlator2db/health.c \
      ${MWAX_CBF_ROOT}/mwax-fullcorrelator/src/mwax_db2correlator2db/quit.c \
      ${MWAX_CBF_ROOT}/mwax-fullcorrelator/src/mwax_db2correlator2db/mwax_db2correlate2db.c

${CC} -o mwax_db2correlate2db -fopenmp \
      -L${PSRDADA_LIB} -L${INSTALL_PREFIX}/lib -lpsrdada -lxgpumwax${NUMTILES}t_6400 -lcudart -lcufft -lm -lstdc++ \
      mwax_global_defs.o health.o quit.o mwax_db2correlate2db.o mwax_cbf_cuda.o
	
install -D -t ${INSTALL_PREFIX}/bin mwax_db2correlate2db


# ----- Compile mwax_db2fits -----

# Original build instructions for mwax_db2fits at
# https://github.com/MWATelescope/mwax_db2fits

cd ${MWAX_DB2FITS_ROOT}

rm -rf build
mkdir -p build && cd build

cmake -DCMAKE_C_COMPILER=${CC} -DCMAKE_CXX_COMPILER=${CXX} \
      -DPSRDADA_ROOT_DIR=${PSRDADA_ROOT} \
      -DMWAX_COMMON_ROOT=../mwax_common-git \
      -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \
      ..

make
make install


# ----- Compile mwax_user_tools -----

# Original build instructions for mwax_db2fits at
# 	https://github.com/MWATelescope/mwax_user_tools

cd ${MWAX_USER_TOOLS_ROOT}/utils/mwax_update_subfile_header

${CC} -c -Wall -Wextra -Ofast \
      -I${PSRDADA_INCLUDE_DIR} \
      -o mwax_update_subfile_header.o \
      mwax_update_subfile_header.c

${CC} -o mwax_update_subfile_header \
      -L${PSRDADA_LIB} -lpsrdada \
      mwax_update_subfile_header.o 

install -D -t ${INSTALL_PREFIX}/bin mwax_update_subfile_header