#!/usr/bin/env python
"""Script for generating mathlib.cmake."""
import os
import sys
import argparse
from dataclasses import dataclass

from cmake import CMakeFile
from gen_config import BuildConfig


@dataclass
class TBPLaSOptions:
    """Collection of internal TBPLaS options."""
    use_mkl = False
    use_aocl = False
    use_openblas = False
    use_netlib = False
    ilp64 = False


CONFIG_MKL = r"""
find_package(MKL CONFIG REQUIRED)
add_compile_options($<TARGET_PROPERTY:MKL::MKL,INTERFACE_COMPILE_OPTIONS>)
include_directories($<TARGET_PROPERTY:MKL::MKL,INTERFACE_INCLUDE_DIRECTORIES>)
link_libraries($<LINK_ONLY:MKL::MKL>)

# For Intel compiler, add proper mkl options.
if(TBPLAS_USE_INTEL)
    if(WITH_OPENMP)
        set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -qmkl=parallel")
        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -qmkl=parallel")
    else()
        set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -qmkl=sequential")
        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -qmkl=sequential")
    endif()
endif()
"""

CONFIG_FEAST = r"""
add_compile_definitions(WITH_FEAST)
include_directories("$ENV{FEASTROOT}/include")
link_directories("$ENV{FEASTROOT}/lib/x64")

# Link compiler-dependent fortran library
# FEAST Eigenvalue Solver v4.0 User Guide, pp.11
# FEAST does not support flang, so we need to compile it using gfortran and link
# to its runtime library.
if(TBPLAS_USE_INTEL)
    link_libraries(feast ifcore ifcoremt)
elseif(TBPLAS_USE_CLANG)
    link_libraries(feast gfortran)
elseif(TBPLAS_USE_GCC)
    link_libraries(feast gfortran)
else()
    message(FATAL_ERROR "Fortran library for FEAST not set")
endif()
"""


def check_math_vendor(config: BuildConfig) -> None:
    # Allowed math library vendors for each component
    eigen_vendors = ("mkl", "aocl", "openblas", "netlib", "default")
    diag_vendors = eigen_vendors
    tbpm_vendors = ("mkl", "aocl", "default")

    # Check for illegal math library vendor
    eigen_backend = config.eigen_backend
    diag_backend = config.diag_backend
    tbpm_backend = config.tbpm_backend
    if eigen_backend not in eigen_vendors:
        raise ValueError(f"Unknown eigen_backend: {eigen_backend}")
    if diag_backend not in diag_vendors:
        raise ValueError(f"Unknown diag_backend: {diag_backend}")
    if tbpm_backend not in tbpm_vendors:
        raise ValueError(f"Unknown tbpm_backend: {tbpm_backend}")

    # Check for incompatible combinations
    if eigen_backend != "default":
        # If eigen_backend is non-default, diag_backend should be either default or
        # equal to eigen_backend. And tbpm_backend must be default.
        if diag_backend not in (eigen_backend, "default"):
            raise ValueError(f"Mixed usage of {eigen_backend} and {diag_backend}")
        if tbpm_backend != "default":
            raise ValueError(f"TBPM_BACKEND must be default when EIGEN_BACKEND is {eigen_backend}")
    elif diag_backend != "default":
        # If eigen_backend is default, then the choice of diag_backend is free.
        # But tbpm_backend cannot conflict with diag_backend.
        if tbpm_backend != "default":
            raise ValueError(f"TBPM_BACKEND must be default when DIAG_BACKEND is {diag_backend}")


def check_aocl(opt: TBPLaSOptions) -> None:
    # Check if AOCL is properly installed
    try:
        aocl_root = os.environ["AOCL_ROOT"]
    except KeyError:
        raise RuntimeError("AOCL_ROOT not set")

    # Check targets of symbolic links under AOCL_ROOT
    target = os.popen(f"readlink {aocl_root}/lib").readlines()[0].rstrip()
    link_ilp64 = True if target.find("ILP64") != -1 else False
    if link_ilp64 != opt.ilp64:
        opt_type = "ilp64" if opt.ilp64 else "lp64"
        raise RuntimeError(f"AOCL links not set to {opt_type}")


def check_feast(opt: TBPLaSOptions) -> None:
    # Check if LAPACK has been enabled
    if not (opt.use_mkl or opt.use_aocl or opt.use_openblas or opt.use_netlib):
        raise ValueError("FEAST requires a working LAPACK library")

    # Check if using the LP64 interface of MKL/AOCL
    if opt.ilp64:
        raise ValueError("FEAST requires MKL/AOCL LP64 interface")

    # Check if FEAST is properly installed
    try:
        os.environ["FEASTROOT"]
    except KeyError:
        raise RuntimeError("FEASTROOT not set")


def gen_cmake(config: BuildConfig, output: str) -> None:
    opt = TBPLaSOptions()
    cm = CMakeFile(output)
    check_math_vendor(config)

    # Backend for Eigen
    if config.eigen_backend == "mkl":
        opt.use_mkl = True
        cm.add_compile_definitions("EIGEN_USE_MKL_ALL")
    elif config.eigen_backend == "aocl":
        opt.use_aocl = True
        cm.add_compile_definitions("EIGEN_USE_BLAS")
        cm.add_compile_definitions("EIGEN_USE_LAPACKE")
    elif config.eigen_backend == "openblas":
        opt.use_openblas = True
        cm.add_compile_definitions("EIGEN_USE_BLAS")
        cm.add_compile_definitions("EIGEN_USE_LAPACKE")
    elif config.eigen_backend == "netlib":
        opt.use_netlib = True
        cm.add_compile_definitions("EIGEN_USE_BLAS")
        cm.add_compile_definitions("EIGEN_USE_LAPACKE")

    # Backend for Diag
    if config.diag_backend == "mkl":
        opt.use_mkl = True
        cm.add_compile_definitions("DIAG_USE_MKL")
        cm.add_compile_definitions("WITH_LAPACKE")
    elif config.diag_backend == "aocl":
        opt.use_aocl = True
        cm.add_compile_definitions("WITH_LAPACKE")
    elif config.diag_backend == "openblas":
        opt.use_openblas = True
        cm.add_compile_definitions("WITH_LAPACKE")
    elif config.diag_backend == "netlib":
        opt.use_netlib = True
        cm.add_compile_definitions("WITH_LAPACKE")

    # Backend for TBPM
    if config.tbpm_backend == "mkl":
        opt.use_mkl = True
        opt.ilp64 = True
        cm.add_compile_definitions("TBPM_USE_MKL")
    elif config.tbpm_backend == "aocl":
        opt.use_aocl = True
        opt.ilp64 = True
        cm.add_compile_definitions("TBPM_USE_AOCL")

    # Write and clear current configurations
    # MKL_INTERFACE must go first
    if opt.use_mkl:
        if opt.ilp64:
            cm.write_command(f"set(MKL_INTERFACE ilp64)")
        else:
            cm.write_command(f"set(MKL_INTERFACE lp64)")
    cm.write_contents()

    # Vendor-specific configuration
    if opt.use_mkl:
        cm.write_command(CONFIG_MKL)
    elif opt.use_aocl:
        check_aocl(opt)
        cm.add_link_options("-fopenmp")
        cm.link_libraries("blis-mt m pthread rt")
        cm.link_libraries("flame pthread aoclutils")
        if opt.ilp64:
            cm.add_compile_definitions("AOCLDA_ILP64")
            cm.add_compile_definitions("aoclsparse_ILP64")
            cm.link_libraries("aoclsparse")
    elif opt.use_openblas:
        cm.find_package("OpenBLAS")
        cm.include_directories(r"${OpenBLAS_INCLUDE_DIRS}")
        cm.link_libraries(r"${OpenBLAS_LIBRARIES}")
    elif opt.use_netlib:
        cm.find_package("cblas")
        cm.find_package("lapacke")
        cm.include_directories(r"${CBLAS_INCLUDE_DIRS}")
        cm.include_directories(r"${LAPACKE_INCLUDE_DIRS}")
        cm.link_libraries(r"${CBLAS_LIBRARIES}")
        cm.link_libraries(r"${LAPACKE_LIBRARIES}")

    # Write and clear current configurations
    cm.write_contents()

    # Write feast
    if config.with_feast:
        check_feast(opt)
        cm.write_command(CONFIG_FEAST)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--eigen-backend", type=str, action="store", default="default")
    parser.add_argument("--diag-backend", type=str, action="store", default="default")
    parser.add_argument("--tbpm-backend", type=str, action="store", default="default")
    parser.add_argument("--with-feast", type=str, action="store", default="off")
    parser.add_argument("--output", type=str, action="store", default="math_lib.cmake")
    args = parser.parse_args()

    config = BuildConfig()
    config.eigen_backend = args.eigen_backend
    config.diag_backend = args.diag_backend
    config.tbpm_backend = args.tbpm_backend
    config.with_feast = True if args.with_feast == "on" else False
    gen_cmake(config, args.output)


if __name__ == "__main__":
    try:
        main()
    except (ValueError, RuntimeError) as err:
        print(f"ERROR: {err}")
        sys.exit(-1)
