"""
Module implementing the C++ "standalone" `CodeObject`
"""

from brian2.codegen.codeobject import CodeObject, constant_or_scalar
from brian2.codegen.generators.cpp_generator import CPPCodeGenerator, c_data_type
from brian2.codegen.targets import codegen_targets
from brian2.codegen.templates import Templater
from brian2.core.functions import DEFAULT_FUNCTIONS
from brian2.core.preferences import prefs
from brian2.devices.device import get_device
from brian2.utils.stringtools import replace

__all__ = ["CPPStandaloneCodeObject"]


def openmp_pragma(pragma_type):
    nb_threads = prefs.devices.cpp_standalone.openmp_threads
    openmp_on = not (nb_threads == 0)

    ## First we need to deal with some special cases that have to be handle in case
    ## openmp is not activated
    if pragma_type == "set_num_threads":
        if not openmp_on:
            return ""
        elif nb_threads > 0:
            # We have to fix the exact number of threads in all parallel sections
            return f"omp_set_dynamic(0);\nomp_set_num_threads({int(nb_threads)});"
    elif pragma_type == "get_thread_num":
        if not openmp_on:
            return "0"
        else:
            return "omp_get_thread_num()"
    elif pragma_type == "get_num_threads":
        if not openmp_on:
            return "1"
        else:
            return f"{int(nb_threads)}"
    elif pragma_type == "with_openmp":
        # The returned value is a proper Python boolean, i.e. not something
        # that should be included in the generated code but rather for use
        # in {% if ... %} statements in the template
        return openmp_on

    ## Then by default, if openmp is off, we do not return any pragma statement in the templates
    elif not openmp_on:
        return ""
    ## Otherwise, we return the correct pragma statement
    elif pragma_type == "static":
        return "#pragma omp for schedule(static)"
    elif pragma_type == "single":
        return "#pragma omp single"
    elif pragma_type == "single-nowait":
        return "#pragma omp single nowait"
    elif pragma_type == "critical":
        return "#pragma omp critical"
    elif pragma_type == "atomic":
        return "#pragma omp atomic"
    elif pragma_type == "once":
        return "#pragma once"
    elif pragma_type == "parallel-static":
        return "#pragma omp parallel for schedule(static)"
    elif pragma_type == "static-ordered":
        return "#pragma omp for schedule(static) ordered"
    elif pragma_type == "ordered":
        return "#pragma omp ordered"
    elif pragma_type == "include":
        return "#include <omp.h>"
    elif pragma_type == "parallel":
        return "#pragma omp parallel"
    elif pragma_type == "master":
        return "#pragma omp master"
    elif pragma_type == "barrier":
        return "#pragma omp barrier"
    elif pragma_type == "compilation":
        return "-fopenmp"
    elif pragma_type == "sections":
        return "#pragma omp sections"
    elif pragma_type == "section":
        return "#pragma omp section"
    else:
        raise ValueError(f'Unknown OpenMP pragma "{pragma_type}"')


class CPPStandaloneCodeObject(CodeObject):
    """
    C++ standalone code object

    The ``code`` should be a `~brian2.codegen.templates.MultiTemplate`
    object with two macros defined, ``main`` (for the main loop code) and
    ``support_code`` for any support code (e.g. function definitions).
    """

    templater = Templater(
        "brian2.devices.cpp_standalone",
        ".cpp",
        env_globals={
            "c_data_type": c_data_type,
            "openmp_pragma": openmp_pragma,
            "constant_or_scalar": constant_or_scalar,
            "prefs": prefs,
            "zip": zip,
        },
    )
    generator_class = CPPCodeGenerator

    def __init__(self, *args, **kwds):
        super().__init__(*args, **kwds)
        #: Store whether this code object defines before/after blocks
        self.before_after_blocks = []

    def __call__(self, **kwds):
        return self.run()

    def compile_block(self, block):
        pass  # Compilation will be handled in device

    def run_block(self, block):
        if block == "run":
            get_device().main_queue.append((f"{block}_code_object", (self,)))
        else:
            # Check the C++ code whether there is anything to run
            cpp_code = getattr(self.code, f"{block}_cpp_file")
            if len(cpp_code) and "EMPTY_CODE_BLOCK" not in cpp_code:
                get_device().main_queue.append((f"{block}_code_object", (self,)))
                self.before_after_blocks.append(block)


codegen_targets.add(CPPStandaloneCodeObject)


# At module initialization time, we do not yet know whether the code will be
# run with OpenMP or not. We therefore use a "dynamic implementation" which
# generates the rand/randn implementation during code generation.
def generate_rand_code(rand_func, owner):
    nb_threads = prefs.devices.cpp_standalone.openmp_threads
    if nb_threads == 0:  # no OpenMP
        thread_number = "0"
    else:
        thread_number = "omp_get_thread_num()"
    if rand_func not in ["rand", "randn"]:
        raise AssertionError(rand_func)
    code = """
           inline double _%RAND_FUNC%(const int _vectorisation_idx) {
               return brian::_random_generators[%THREAD_NUMBER%].%RAND_FUNC%();
           }
           """
    code = replace(
        code,
        {
            "%THREAD_NUMBER%": thread_number,
            "%RAND_FUNC%": rand_func,
        },
    )
    return {"support_code": code}


rand_impls = DEFAULT_FUNCTIONS["rand"].implementations
rand_impls.add_dynamic_implementation(
    CPPStandaloneCodeObject,
    code=lambda owner: generate_rand_code("rand", owner),
    namespace=lambda owner: {},
    name="_rand",
)

randn_impls = DEFAULT_FUNCTIONS["randn"].implementations
randn_impls.add_dynamic_implementation(
    CPPStandaloneCodeObject,
    code=lambda owner: generate_rand_code("randn", owner),
    namespace=lambda owner: {},
    name="_randn",
)