#!/bin/bash
#
# install_netcdf.sh – build NetCDF from source using GCC + 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
#
# The script builds zlib, HDF5, NetCDF-C, and NetCDF-Fortran in sequence,
# installs everything under ${WORK_DIR}/install, and writes a small
# environment file (netcdf_env.sh) that get_siesta.sh sources later.
#

set -e

###############################################################################
# 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

echo "=========================================="
echo "NetCDF installation – GCC + Fujitsu MPI"
echo "=========================================="

###############################################################################
# Paths
###############################################################################
BASE_DIR="/${MY_VOLUME}/mdt1/data/${MY_PROJECT}/${MY_USER}"
WORK_DIR="${BASE_DIR}/netcdf_gcc"
SRC_DIR="${WORK_DIR}/src"
INSTALL_DIR="${WORK_DIR}/install"

MPI_GCC_DIR="/vol0004/apps/oss/mpigcc/fjmpi-gcc12"
FUJITSU_MPI_DIR="/opt/FJSVxtclanga/tcsds-1.2.42"

###############################################################################
# Library versions
###############################################################################
ZLIB_VERSION="1.3.1"
HDF5_VERSION="1.14.3"
NETCDF_C_VERSION="4.9.2"
NETCDF_FORTRAN_VERSION="4.6.1"

###############################################################################
# Download URLs
###############################################################################
ZLIB_URL="https://zlib.net/fossils/zlib-${ZLIB_VERSION}.tar.gz"
HDF5_URL="https://support.hdfgroup.org/ftp/HDF5/releases/hdf5-1.14/hdf5-${HDF5_VERSION}/src/hdf5-${HDF5_VERSION}.tar.gz"
NETCDF_C_URL="https://downloads.unidata.ucar.edu/netcdf-c/${NETCDF_C_VERSION}/netcdf-c-${NETCDF_C_VERSION}.tar.gz"
NETCDF_FORTRAN_URL="https://downloads.unidata.ucar.edu/netcdf-fortran/${NETCDF_FORTRAN_VERSION}/netcdf-fortran-${NETCDF_FORTRAN_VERSION}.tar.gz"

###############################################################################
# Environment
###############################################################################
export PATH="${MPI_GCC_DIR}/bin:${INSTALL_DIR}/bin:${PATH}"
export LD_LIBRARY_PATH="${INSTALL_DIR}/lib:${FUJITSU_MPI_DIR}/lib64:${LD_LIBRARY_PATH:-}"

mkdir -p "${SRC_DIR}" "${INSTALL_DIR}"

###############################################################################
# Helper: resolve Fortran MPI wrapper (mpif90 or mpifort)
###############################################################################
get_fc_wrapper() {
    if [ -f "${MPI_GCC_DIR}/bin/mpif90" ]; then
        echo "${MPI_GCC_DIR}/bin/mpif90"
    elif [ -f "${MPI_GCC_DIR}/bin/mpifort" ]; then
        echo "${MPI_GCC_DIR}/bin/mpifort"
    else
        echo "Error: no Fortran MPI wrapper found in ${MPI_GCC_DIR}/bin." >&2
        exit 1
    fi
}

###############################################################################
# Helper: download and extract a tarball
###############################################################################
download_and_extract() {
    local url=$1
    local filename
    filename=$(basename "${url}")
    cd "${SRC_DIR}"
    if [ ! -f "${filename}" ]; then
        echo "Downloading ${filename} ..."
        wget "${url}" || curl -L -O "${url}"
    else
        echo "${filename} already downloaded."
    fi
    echo "Extracting ${filename} ..."
    tar -xzf "${filename}"
}

###############################################################################
# Verify MPI wrappers before starting
###############################################################################
verify_mpi_wrappers() {
    echo ""
    echo "=========================================="
    echo "Verifying MPI wrappers"
    echo "=========================================="

    if [ ! -d "${MPI_GCC_DIR}" ]; then
        echo "Error: MPI directory not found: ${MPI_GCC_DIR}" >&2
        exit 1
    fi

    if [ -f "${MPI_GCC_DIR}/bin/mpicc" ]; then
        echo "  ✓ mpicc found"
    else
        echo "  ✗ mpicc NOT found in ${MPI_GCC_DIR}/bin" >&2
        exit 1
    fi

    local fc_wrapper
    fc_wrapper=$(get_fc_wrapper)
    echo "  ✓ Fortran wrapper found: $(basename "${fc_wrapper}")"

    # Compile a minimal MPI test to confirm the wrappers work
    local test_dir="${WORK_DIR}/mpi_test"
    mkdir -p "${test_dir}"
    cd "${test_dir}"

    cat > test_mpi.f90 << 'TESTMPI'
program test_mpi
    include 'mpif.h'
    integer :: ierr, rank
    call MPI_Init(ierr)
    call MPI_Comm_rank(MPI_COMM_WORLD, rank, ierr)
    print *, 'MPI OK'
    call MPI_Finalize(ierr)
end program test_mpi
TESTMPI

    echo "Compiling MPI test ..."
    "${fc_wrapper}" -o test_mpi test_mpi.f90

    if [ $? -eq 0 ]; then
        echo "  ✓ MPI test compiled successfully (not executed – requires batch system)"
    else
        echo "  ✗ MPI test compilation failed" >&2
        exit 1
    fi
}

###############################################################################
# zlib
###############################################################################
install_zlib() {
    echo ""
    echo "=========================================="
    echo "Installing zlib ${ZLIB_VERSION}"
    echo "=========================================="

    download_and_extract "${ZLIB_URL}"
    cd "${SRC_DIR}/zlib-${ZLIB_VERSION}"

    ./configure --prefix="${INSTALL_DIR}"
    make -j 4
    make install
    echo "  ✓ zlib installed"
}

###############################################################################
# HDF5
###############################################################################
install_hdf5() {
    echo ""
    echo "=========================================="
    echo "Installing HDF5 ${HDF5_VERSION}"
    echo "=========================================="

    download_and_extract "${HDF5_URL}"
    cd "${SRC_DIR}/hdf5-${HDF5_VERSION}"

    local fc_wrapper
    fc_wrapper=$(get_fc_wrapper)

    echo "Configuring HDF5 ..."
    echo "  CC: ${MPI_GCC_DIR}/bin/mpicc"
    echo "  FC: ${fc_wrapper}"

    ./configure --prefix="${INSTALL_DIR}" \
                --enable-fortran \
                --enable-parallel \
                --enable-shared \
                --enable-static \
                --with-zlib="${INSTALL_DIR}" \
                CC="${MPI_GCC_DIR}/bin/mpicc" \
                FC="${fc_wrapper}" \
                CFLAGS="-O2 -fPIC" \
                FCFLAGS="-O2 -fPIC" \
                LDFLAGS="-L${INSTALL_DIR}/lib -Wl,-rpath,${INSTALL_DIR}/lib"

    if [ $? -ne 0 ]; then
        echo "  ✗ HDF5 configure failed. Last 100 lines of config.log:"
        tail -100 config.log
        exit 1
    fi

    echo "Building HDF5 ..."
    make -j 4

    echo "Installing HDF5 ..."
    make install
    echo "  ✓ HDF5 installed"
}

###############################################################################
# NetCDF-C
###############################################################################
install_netcdf_c() {
    echo ""
    echo "=========================================="
    echo "Installing NetCDF-C ${NETCDF_C_VERSION}"
    echo "=========================================="

    download_and_extract "${NETCDF_C_URL}"
    cd "${SRC_DIR}/netcdf-c-${NETCDF_C_VERSION}"

    echo "Configuring NetCDF-C ..."

    ./configure --prefix="${INSTALL_DIR}" \
                --enable-netcdf-4 \
                --enable-shared \
                --enable-static \
                --disable-dap \
                CC="${MPI_GCC_DIR}/bin/mpicc" \
                CFLAGS="-O2 -fPIC" \
                CPPFLAGS="-I${INSTALL_DIR}/include" \
                LDFLAGS="-L${INSTALL_DIR}/lib -Wl,-rpath,${INSTALL_DIR}/lib"

    if [ $? -ne 0 ]; then
        echo "  ✗ NetCDF-C configure failed. Last 100 lines of config.log:"
        tail -100 config.log
        exit 1
    fi

    echo "Building NetCDF-C ..."
    make -j 4

    echo "Installing NetCDF-C ..."
    make install
    echo "  ✓ NetCDF-C installed"
}

###############################################################################
# NetCDF-Fortran
###############################################################################
install_netcdf_fortran() {
    echo ""
    echo "=========================================="
    echo "Installing NetCDF-Fortran ${NETCDF_FORTRAN_VERSION}"
    echo "=========================================="

    download_and_extract "${NETCDF_FORTRAN_URL}"
    cd "${SRC_DIR}/netcdf-fortran-${NETCDF_FORTRAN_VERSION}"

    local fc_wrapper
    fc_wrapper=$(get_fc_wrapper)

    echo "Configuring NetCDF-Fortran ..."

    ./configure --prefix="${INSTALL_DIR}" \
                --enable-shared \
                --enable-static \
                CC="${MPI_GCC_DIR}/bin/mpicc" \
                FC="${fc_wrapper}" \
                F77="${fc_wrapper}" \
                CFLAGS="-O2 -fPIC" \
                FCFLAGS="-O2 -fPIC" \
                FFLAGS="-O2 -fPIC" \
                CPPFLAGS="-I${INSTALL_DIR}/include" \
                LDFLAGS="-L${INSTALL_DIR}/lib -Wl,-rpath,${INSTALL_DIR}/lib"

    if [ $? -ne 0 ]; then
        echo "  ✗ NetCDF-Fortran configure failed. Last 100 lines of config.log:"
        tail -100 config.log
        exit 1
    fi

    echo "Building NetCDF-Fortran ..."
    make -j 4

    echo "Installing NetCDF-Fortran ..."
    make install
    echo "  ✓ NetCDF-Fortran installed"
}

###############################################################################
# Write environment file sourced by get_siesta.sh
#
# INSTALL_DIR is expanded now (at write time) so the generated file contains
# the resolved absolute path and works correctly when sourced in a fresh
# session where MY_* variables may not be set.
###############################################################################
create_env_file() {
    echo ""
    echo "Writing environment file ..."
    ENV_FILE="${WORK_DIR}/netcdf_env.sh"

    cat > "${ENV_FILE}" << ENVEOF
#!/bin/bash
# netcdf_env.sh – sourced by get_siesta.sh to expose NetCDF paths.
# Generated by install_netcdf.sh on $(date).

export NETCDF_DIR="${INSTALL_DIR}"
export MPI_GCC_DIR="${MPI_GCC_DIR}"
export FUJITSU_MPI_DIR="${FUJITSU_MPI_DIR}"

export PATH="\${MPI_GCC_DIR}/bin:\${NETCDF_DIR}/bin:\${PATH}"
export LD_LIBRARY_PATH="\${NETCDF_DIR}/lib:\${FUJITSU_MPI_DIR}/lib64:\${LD_LIBRARY_PATH:-}"

export NETCDF_INC="-I\${NETCDF_DIR}/include"
export NETCDF_LIB="-L\${NETCDF_DIR}/lib -lnetcdff -lnetcdf -lhdf5_hl -lhdf5 -lz"

echo "NetCDF environment configured."
echo "  NETCDF_DIR  : \${NETCDF_DIR}"
echo "  MPI wrappers: \${MPI_GCC_DIR}/bin"
ENVEOF

    chmod +x "${ENV_FILE}"
    echo "  ✓ Environment file written: ${ENV_FILE}"
}

###############################################################################
# Verify critical libraries
###############################################################################
verify_installation() {
    echo ""
    echo "=========================================="
    echo "Verifying installation"
    echo "=========================================="

    source "${WORK_DIR}/netcdf_env.sh"

    if [ -f "${INSTALL_DIR}/bin/nc-config" ]; then
        echo "nc-config --version: $(nc-config --version)"
        echo "nc-config --libs:    $(nc-config --libs)"
    fi

    if [ -f "${INSTALL_DIR}/bin/nf-config" ]; then
        echo "nf-config --version: $(nf-config --version)"
        echo "nf-config --flibs:   $(nf-config --flibs)"
    fi

    echo ""
    echo "Installed libraries:"
    ls -lh "${INSTALL_DIR}/lib/"*.{a,so} 2>/dev/null | grep -E "libz|libhdf5|libnetcdf" || true

    echo ""
    CRITICAL_FILES=(
        "${INSTALL_DIR}/lib/libz.a"
        "${INSTALL_DIR}/lib/libhdf5.a"
        "${INSTALL_DIR}/lib/libhdf5_hl.a"
        "${INSTALL_DIR}/lib/libhdf5_fortran.a"
        "${INSTALL_DIR}/lib/libnetcdf.a"
        "${INSTALL_DIR}/lib/libnetcdff.a"
    )

    local all_ok=true
    for f in "${CRITICAL_FILES[@]}"; do
        if [ -f "$f" ]; then
            echo "  ✓ $(basename "$f")"
        else
            echo "  ✗ Missing: $(basename "$f")"
            all_ok=false
        fi
    done

    if [ "${all_ok}" = true ]; then
        echo ""
        echo "  ✓ All critical libraries present."
    fi
}

###############################################################################
# Compile a NetCDF+MPI test program
###############################################################################
create_test_program() {
    echo ""
    echo "=========================================="
    echo "Compiling NetCDF test program"
    echo "=========================================="

    local test_dir="${WORK_DIR}/test"
    mkdir -p "${test_dir}"
    cd "${test_dir}"

    cat > test_netcdf.f90 << 'EOF'
program test_netcdf
  use netcdf
  implicit none
  include 'mpif.h'
  integer :: ncid, status, ierr, rank
  character(len=256) :: filename

  call MPI_Init(ierr)
  call MPI_Comm_rank(MPI_COMM_WORLD, rank, ierr)

  if (rank == 0) then
    print *, '=========================================='
    print *, 'NetCDF-Fortran + MPI test'
    print *, 'NetCDF version: ', trim(nf90_inq_libvers())
    print *, '=========================================='

    filename = 'test.nc'
    status = nf90_create(filename, NF90_CLOBBER, ncid)

    if (status == NF90_NOERR) then
      status = nf90_close(ncid)
      print *, 'TEST SUCCESSFUL'
    else
      print *, 'TEST FAILED'
    end if
  end if

  call MPI_Finalize(ierr)
end program test_netcdf
EOF

    local fc_wrapper
    fc_wrapper=$(get_fc_wrapper)

    echo "Compiling test ..."
    "${fc_wrapper}" -o test_netcdf test_netcdf.f90 \
        ${NETCDF_INC} ${NETCDF_LIB}

    if [ $? -eq 0 ]; then
        echo "  ✓ Test compiled successfully."
        echo ""
        echo "  To run:  mpiexec -n 1 ./test_netcdf"
        echo "  Or submit via your batch system."
    else
        echo "  ✗ Test compilation failed."
    fi
}

###############################################################################
# Main
###############################################################################
main() {
    echo "Start: $(date)"
    echo ""

    verify_mpi_wrappers
    install_zlib
    install_hdf5
    install_netcdf_c
    install_netcdf_fortran
    create_env_file
    verify_installation
    create_test_program

    echo ""
    echo "=========================================="
    echo "✓ Installation complete"
    echo "=========================================="
    echo ""
    echo "Elapsed: ${SECONDS} seconds"
    echo ""
    echo "To use:  source ${WORK_DIR}/netcdf_env.sh"
    echo ""
    echo "Example compilation:"
    echo "  mpif90 program.f90 \${NETCDF_INC} \${NETCDF_LIB} -o program"
}

LOG_FILE="${WORK_DIR}/install_log_$(date +%Y%m%d_%H%M%S).txt"
mkdir -p "${WORK_DIR}"
main 2>&1 | tee "${LOG_FILE}"