##
# Copyright 2009-2025 Ghent University
#
# This file is part of EasyBuild,
# originally created by the HPC team of Ghent University (http://ugent.be/hpc/en),
# with support of Ghent University (http://ugent.be/hpc),
# the Flemish Supercomputer Centre (VSC) (https://www.vscentrum.be),
# Flemish Research Foundation (FWO) (http://www.fwo.be/en)
# and the Department of Economy, Science and Innovation (EWI) (http://www.ewi-vlaanderen.be/en).
#
# https://github.com/easybuilders/easybuild
#
# EasyBuild is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation v2.
#
# EasyBuild is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with EasyBuild.  If not, see <http://www.gnu.org/licenses/>.
##
"""
EasyBuild support for building and installing ScaLAPACK, implemented as an easyblock

@author: Stijn De Weirdt (Ghent University)
@author: Dries Verdegem (Ghent University)
@author: Kenneth Hoste (Ghent University)
@author: Pieter De Baets (Ghent University)
@author: Jens Timmerman (Ghent University)
"""

import glob
import os
from easybuild.tools import LooseVersion

import easybuild.tools.toolchain as toolchain
from easybuild.easyblocks.generic.cmakemake import CMakeMake
from easybuild.toolchains.linalg.acml import Acml
from easybuild.toolchains.linalg.atlas import Atlas
from easybuild.toolchains.linalg.blis import Blis
from easybuild.toolchains.linalg.flexiblas import FlexiBLAS, det_flexiblas_backend_libs
from easybuild.toolchains.linalg.gotoblas import GotoBLAS
from easybuild.toolchains.linalg.lapack import Lapack
from easybuild.toolchains.linalg.openblas import OpenBLAS
from easybuild.toolchains.linalg.intelmkl import IntelMKL
from easybuild.tools.build_log import EasyBuildError
from easybuild.tools.filetools import copy_file, remove_file
from easybuild.tools.modules import get_software_root
from easybuild.tools.run import run_shell_cmd


class EB_ScaLAPACK(CMakeMake):
    """
    Support for building and installing ScaLAPACK, both versions 1.x and 2.x
    """

    def __init__(self, *args, **kwargs):
        """Constructor of ScaLAPACK easyblock."""
        super().__init__(*args, **kwargs)

        self.loosever = LooseVersion(self.version)

        # use CMake for recent versions, but only if CMake is listed as a build dep
        build_deps_names = [dep['name'].lower() for dep in self.cfg.builddependencies()]
        self.use_cmake = self.loosever >= LooseVersion('2.1.0') and 'cmake' in build_deps_names

    def configure_step(self):
        """Configure ScaLAPACK build by copying SLmake.inc.example to SLmake.inc and checking dependencies."""

        # use CMake for recent versions, but only if CMake is listed as a build dep
        if self.use_cmake:
            super().configure_step()
        else:
            src = os.path.join(self.cfg['start_dir'], 'SLmake.inc.example')
            dest = os.path.join(self.cfg['start_dir'], 'SLmake.inc')

            if os.path.exists(dest):
                raise EasyBuildError("Destination file %s exists", dest)
            else:
                copy_file(src, dest)

    def build_libscalapack_make(self):
        """Build libscalapack using 'make -j', after determining the options to pass to make."""
        # MPI compiler commands
        known_mpi_libs = [toolchain.MPICH, toolchain.MPICH2, toolchain.MVAPICH2]  # @UndefinedVariable
        known_mpi_libs += [toolchain.OPENMPI, toolchain.QLOGICMPI]  # @UndefinedVariable
        known_mpi_libs += [toolchain.INTELMPI]  # @UndefinedVariable
        if os.getenv('MPICC') and os.getenv('MPIF90'):
            mpicc = os.getenv('MPICC')
            mpif90 = os.getenv('MPIF90')
        elif self.toolchain.mpi_family() in known_mpi_libs:
            mpicc = 'mpicc'
            mpif90 = 'mpif90'
        else:
            raise EasyBuildError("Don't know which compiler commands to use.")

        # determine build options BLAS and LAPACK libs
        extra_makeopts = []

        acml = get_software_root(Acml.LAPACK_MODULE_NAME[0])
        flexiblas = get_software_root(FlexiBLAS.LAPACK_MODULE_NAME[0])
        intelmkl = get_software_root(IntelMKL.LAPACK_MODULE_NAME[0])
        lapack = get_software_root(Lapack.LAPACK_MODULE_NAME[0])
        openblas = get_software_root(OpenBLAS.LAPACK_MODULE_NAME[0])

        if flexiblas:
            libdir = os.path.join(flexiblas, 'lib')
            blas_libs = ' '.join(['-l%s' % lib for lib in FlexiBLAS.BLAS_LIB])
            extra_makeopts.extend([
                'BLASLIB="-L%s %s -lpthread"' % (libdir, blas_libs),
                'LAPACKLIB="-L%s %s"' % (libdir, blas_libs),
            ])
        elif lapack:
            extra_makeopts.append('LAPACKLIB=%s' % os.path.join(lapack, 'lib', 'liblapack.a'))

            for blas in [Atlas, Blis, GotoBLAS]:
                blas_root = get_software_root(blas.BLAS_MODULE_NAME[0])
                if blas_root:
                    blas_libs = ' '.join(['-l%s' % lib for lib in blas.BLAS_LIB])
                    blas_libdir = os.path.join(blas_root, 'lib')
                    extra_makeopts.append('BLASLIB="-L%s %s -lpthread"' % (blas_libdir, blas_libs))
                    break

            if not blas_root:
                raise EasyBuildError("Failed to find a known BLAS library, don't know how to define 'BLASLIB'")

        elif acml:
            acml_base_dir = os.getenv('ACML_BASEDIR', 'NO_ACML_BASEDIR')
            acml_static_lib = os.path.join(acml, acml_base_dir, 'lib', 'libacml.a')
            extra_makeopts.extend([
                'BLASLIB="%s -lpthread"' % acml_static_lib,
                'LAPACKLIB=%s' % acml_static_lib
            ])
        elif openblas:
            libdir = os.path.join(openblas, 'lib')
            blas_libs = ' '.join(['-l%s' % lib for lib in OpenBLAS.BLAS_LIB])
            extra_makeopts.extend([
                'BLASLIB="-L%s %s -lpthread"' % (libdir, blas_libs),
                'LAPACKLIB="-L%s %s"' % (libdir, blas_libs),
            ])
        elif intelmkl:
            libdir = os.path.join(intelmkl, 'mkl', 'lib', 'intel64')
            blas_libs = os.environ['LIBLAPACK']
            extra_makeopts.extend([
                'BLASLIB="-L%s %s -lpthread"' % (libdir, blas_libs),
                'LAPACKLIB="-L%s %s"' % (libdir, blas_libs),
            ])
        else:
            raise EasyBuildError("Unknown LAPACK library used, no idea how to set BLASLIB/LAPACKLIB make options")

        # determine interface
        if self.toolchain.mpi_family() in known_mpi_libs:
            interface = 'Add_'
        else:
            raise EasyBuildError("Don't know which interface to pick for the MPI library being used.")

        # set compilers and options
        extra_makeopts += [
            'FC="%s"' % mpif90,
            'CC="%s"' % mpicc,
            'CCFLAGS="%s"' % os.getenv('CFLAGS'),
            'FCFLAGS="%s"' % os.getenv('FFLAGS'),
        ]

        # set interface
        extra_makeopts.append('CDEFS="-D%s"' % interface)

        # update make opts, and build_step
        saved_buildopts = self.cfg['buildopts']

        # Only build the library first, that can be done in parallel.
        # Creating libscalapack.a may fail in parallel, but should work
        # fine with non-parallel make afterwards
        self.cfg.update('buildopts', 'lib')
        self.cfg.update('buildopts', ' '.join(extra_makeopts))

        cmd = "%s make %s %s" % (self.cfg['prebuildopts'], self.parallel_flag, self.cfg['buildopts'])

        # Ignore exit code for parallel run
        run_shell_cmd(cmd, fail_on_error=False)

        # Now prepare to remake libscalapack.a serially and the tests.
        self.cfg['buildopts'] = saved_buildopts
        self.cfg.update('buildopts', ' '.join(extra_makeopts))

        remove_file('libscalapack.a')
        self.cfg.parallel = 1

    def build_step(self):
        """Build ScaLAPACK using make after setting make options."""

        # only do a parallel pre-build of libscalapack and set up build options if we're not using CMake
        if not self.use_cmake:
            self.build_libscalapack_make()

        super().build_step()

    def install_step(self):
        """Install by copying files to install dir."""

        if self.use_cmake:
            super().install_step()
        else:
            # 'manually' install ScaLAPACK by copying headers and libraries if we're not using CMake
            path_info = [
                ('SRC', 'include', '.h'),  # include files
                ('', 'lib', '.a'),  # libraries
            ]
            for (srcdir, destdir, ext) in path_info:

                src = os.path.join(self.cfg['start_dir'], srcdir)
                dest = os.path.join(self.installdir, destdir)

                for lib in glob.glob(os.path.join(src, '*%s' % ext)):
                    copy_file(lib, os.path.join(dest, os.path.basename(lib)))
                    self.log.debug("Copied %s to %s", lib, dest)

    def banned_linked_shared_libs(self):
        """
        List of shared libraries which are not allowed to be linked in any installed binary/library.
        """
        res = super().banned_linked_shared_libs()

        # register FlexiBLAS backends as banned libraries,
        # ScaLAPACK should not be linking to those directly
        if get_software_root(FlexiBLAS.LAPACK_MODULE_NAME[0]):
            res.extend(det_flexiblas_backend_libs())

        return res

    def sanity_check_step(self):
        """Custom sanity check for ScaLAPACK."""

        custom_paths = {
            'files': [os.path.join('lib', 'libscalapack.a')],
            'dirs': []
        }

        super().sanity_check_step(custom_paths=custom_paths)
