# Copyright 2013-2022 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
from spack import *


class KokkosKernels(CMakePackage, CudaPackage):
    """Kokkos Kernels provides math kernels, often BLAS or LAPACK
    for small matrices, that can be used in larger Kokkos parallel routines"""

    homepage = "https://github.com/kokkos/kokkos-kernels"
    git      = "https://github.com/kokkos/kokkos-kernels.git"
    url      = "https://github.com/kokkos/kokkos-kernels/archive/3.6.00.tar.gz"

    tags = ['e4s']

    test_requires_compiler = True

    maintainers = ['lucbv', 'srajama1', 'brian-kelley']

    # generate checksum for each release tarball with the following command
    # openssl sha256 kokkos-kernels-x.y.z.tar.gz
    version('develop', branch='develop')
    version('master',  branch='master')
    version('3.6.00', sha256="2753643fd643b9eed9f7d370e0ff5fa957211d08a91aa75398e31cbc9e5eb0a5")
    version('3.5.00', sha256="a03a41a047d95f9f07cd1e1d30692afdb75b5c705ef524e19c1d02fe60ccf8d1")
    version('3.4.01', sha256="f504aa4afbffb58fa7c4430d0fdb8fd5690a268823fa15eb0b7d58dab9d351e6")
    version('3.4.00', sha256="07ba11869e686cb0d47272d1ef494ccfbcdef3f93ff1c8b64ab9e136a53a227a")
    version('3.3.01', sha256="0f21fe6b5a8b6ae7738290e293aa990719aefe88b32f84617436bfd6074a8f77")
    version('3.3.00', sha256="8d7f78815301afb90ddba7914dce5b718cea792ac0c7350d2f8d00bd2ef1cece")
    version('3.2.01', sha256="c486e5cac19e354a517498c362838619435734d64b44f44ce909b0531c21d95c")
    version('3.2.00', sha256="8ac20ee28ae7813ce1bda461918800ad57fdbac2af86ef5d1ba74e83e10956de")
    version('3.1.00', sha256="27fea241ae92f41bd5b070b1a590ba3a56a06aca750207a98bea2f64a4a40c89")
    version('3.0.00', sha256="e4b832aed3f8e785de24298f312af71217a26067aea2de51531e8c1e597ef0e6")

    depends_on("kokkos")
    depends_on("kokkos@master", when="@master")
    depends_on("kokkos@develop", when="@develop")
    depends_on("kokkos@3.6.00", when="@3.6.00")
    depends_on("kokkos@3.5.00", when="@3.5.00")
    depends_on("kokkos@3.4.01", when="@3.4.01")
    depends_on("kokkos@3.4.00", when="@3.4.00")
    depends_on("kokkos@3.3.01", when="@3.3.01")
    depends_on("kokkos@3.3.00", when="@3.3.00")
    depends_on("kokkos@3.2.01", when="@3.2.01")
    depends_on("kokkos@3.2.00", when="@3.2.00")
    depends_on("kokkos@3.1.00", when="@3.1.00")
    depends_on("kokkos@3.0.00", when="@3.0.00")
    depends_on("cmake@3.16:", type='build')

    backends = {
        'serial': (False,  "enable Serial backend (default)"),
        'cuda': (False,  "enable Cuda backend"),
        'openmp': (False,  "enable OpenMP backend"),
        'pthread': (False,  "enable Pthread backend"),
    }

    for backend in backends:
        deflt, descr = backends[backend]
        variant(backend.lower(), default=deflt, description=descr)
        depends_on("kokkos+%s" % backend.lower(), when="+%s" % backend.lower())

    space_etis = {
        "execspace_cuda": ('auto', "", "cuda"),
        "execspace_openmp": ('auto', "", "openmp"),
        "execspace_threads": ('auto', "", "pthread"),
        "execspace_serial": ('auto', "", "serial"),
        "memspace_cudauvmspace": ('auto', "", "cuda"),
        "memspace_cudaspace": ('auto', "", "cuda"),
    }
    for eti in space_etis:
        deflt, descr, backend_required = space_etis[eti]
        variant(eti, default=deflt, description=descr)
        depends_on("kokkos+%s" % backend_required, when="+%s" % eti)

    numeric_etis = {
        "ordinals": ("int",        "ORDINAL_",  # default, cmake name
                     ["int", "int64_t"]),  # allowed values
        "offsets": ("int,size_t", "OFFSET_",
                    ["int", "size_t"]),
        "layouts": ("left",       "LAYOUT",
                    ["left", "right"]),
        "scalars": ("double",     "",
                    ["float", "double", "complex_float", "complex_double"])
    }
    for eti in numeric_etis:
        deflt, cmake_name, vals = numeric_etis[eti]
        variant(eti, default=deflt, values=vals, multi=True)

    tpls = {
        # variant name   #deflt   #spack name   #root var name #docstring
        "blas": (False, "blas", "BLAS", "Link to system BLAS"),
        "lapack": (False, "lapack", "LAPACK", "Link to system LAPACK"),
        "mkl": (False, "mkl", "MKL", "Link to system MKL"),
        "cublas": (False, "cuda", None, "Link to CUDA BLAS library"),
        "cusparse": (False, "cuda", None, "Link to CUDA sparse library"),
        "superlu": (False, "superlu", "SUPERLU", "Link to SuperLU library"),
        "cblas": (False, "cblas", "CBLAS", "Link to CBLAS library"),
        "lapacke": (False, "clapack", "LAPACKE", "Link to LAPACKE library"),
    }

    for tpl in tpls:
        deflt, spackname, rootname, descr = tpls[tpl]
        variant(tpl, default=deflt, description=descr)
        depends_on(spackname, when="+%s" % tpl)

    variant('shared', default=True, description='Build shared libraries')

    def cmake_args(self):
        spec = self.spec
        options = []

        isdiy = "+diy" in spec
        if isdiy:
            options.append("-DSpack_WORKAROUND=On")

        options.append("-DKokkos_ROOT=%s" % spec["kokkos"].prefix)
        if spec.satisfies('^kokkos+rocm'):
            options.append("-DCMAKE_CXX_COMPILER=%s" % spec['hip'].hipcc)
        else:
            # Compiler weirdness due to nvcc_wrapper
            options.append("-DCMAKE_CXX_COMPILER=%s" % spec["kokkos"].kokkos_cxx)

        if self.run_tests:
            options.append("-DKokkosKernels_ENABLE_TESTS=ON")

        for tpl in self.tpls:
            on_flag = "+%s" % tpl
            off_flag = "~%s" % tpl
            dflt, spackname, rootname, descr = self.tpls[tpl]
            if on_flag in self.spec:
                options.append("-DKokkosKernels_ENABLE_TPL_%s=ON" %
                               tpl.upper())
                if rootname:
                    options.append("-D%s_ROOT=%s" %
                                   (rootname, spec[spackname].prefix))
                else:
                    pass  # this should get picked up automatically, we hope
            elif off_flag in self.spec:
                options.append(
                    "-DKokkosKernels_ENABLE_TPL_%s=OFF" % tpl.upper())

        for eti in self.numeric_etis:
            deflt, cmake_name, vals = self.numeric_etis[eti]
            for val in vals:
                keyval = "%s=%s" % (eti, val)
                cmake_option = "KokkosKernels_INST_%s%s" % (
                    cmake_name.upper(), val.upper())
                if keyval in spec:
                    options.append("-D%s=ON" % cmake_option)
                else:
                    options.append("-D%s=OFF" % cmake_option)

        for eti in self.space_etis:
            deflt, descr, _ = self.space_etis[eti]
            if deflt == "auto":
                value = spec.variants[eti].value
                # spack does these as strings, not reg booleans
                if str(value) == "True":
                    options.append("-DKokkosKernels_INST_%s=ON" % eti.upper())
                elif str(value) == "False":
                    options.append("-DKokkosKernels_INST_%s=OFF" % eti.upper())
                else:
                    pass  # don't pass anything, let CMake decide
            else:  # simple option
                on_flag = "+%s" % eti
                off_flag = "~%s" % eti
                if on_flag in self.spec:
                    options.append("-DKokkosKernels_INST_%s=ON" % eti.upper())
                elif off_flag in self.spec:
                    options.append("-DKokkosKernels_INST_%s=OFF" % eti.upper())

        options.append(self.define_from_variant('BUILD_SHARED_LIBS', 'shared'))

        return options
