"""
Module providing the base `CodeObject` and related functions.
"""

import collections
import copy
import platform
from abc import ABC, abstractmethod

from brian2.core.base import weakproxy_with_fallback
from brian2.core.functions import DEFAULT_FUNCTIONS, Function
from brian2.core.names import Nameable
from brian2.equations.unitcheck import check_units_statements
from brian2.utils.logger import get_logger
from brian2.utils.stringtools import code_representation, indent

from .translation import analyse_identifiers

__all__ = ["CodeObject", "constant_or_scalar"]

logger = get_logger(__name__)


#: Dictionary with basic information about the current system (OS, etc.)
sys_info = {
    "system": platform.system(),
    "architecture": platform.architecture(),
    "machine": platform.machine(),
}


def constant_or_scalar(varname, variable):
    """
    Convenience function to generate code to access the value of a variable.
    Will return ``'varname'`` if the ``variable`` is a constant, and
    ``array_name[0]`` if it is a scalar array.
    """
    from brian2.devices.device import get_device  # avoid circular import

    if variable.array:
        return f"{get_device().get_array_name(variable)}[0]"
    else:
        return f"{varname}"


class CodeObject(Nameable, ABC):
    """
    Executable code object.

    The ``code`` can either be a string or a
    `brian2.codegen.templates.MultiTemplate`.

    After initialisation, the code is compiled with the given namespace
    using ``code.compile(namespace)``.

    Calling ``code(key1=val1, key2=val2)`` executes the code with the given
    variables inserted into the namespace.
    """

    #: The `CodeGenerator` class used by this `CodeObject`
    generator_class = None
    #: A short name for this type of `CodeObject`
    class_name = None

    def __init__(
        self,
        owner,
        code,
        variables,
        variable_indices,
        template_name,
        template_source,
        compiler_kwds,
        name="codeobject*",
    ):
        Nameable.__init__(self, name=name)
        self.owner = weakproxy_with_fallback(owner)
        self.code = code
        self.compiled_code = {}
        self.variables = variables
        self.variable_indices = variable_indices
        self.template_name = template_name
        self.template_source = template_source
        self.compiler_kwds = compiler_kwds

    def __getstate__(self):
        state = self.__dict__.copy()
        state["owner"] = self.owner.__repr__.__self__
        # Replace Function objects for standard functions by their name
        state["variables"] = self.variables.copy()
        for k, v in state["variables"].items():
            if isinstance(v, Function) and v is DEFAULT_FUNCTIONS[k]:
                state["variables"][k] = k
        return state

    def __setstate__(self, state):
        state["owner"] = weakproxy_with_fallback(state["owner"])
        for k, v in state["variables"].items():
            if isinstance(v, str):
                state["variables"][k] = DEFAULT_FUNCTIONS[k]
        self.__dict__ = state

    @classmethod
    def is_available(cls):
        """
        Whether this target for code generation is available. Should use a
        minimal example to check whether code generation works in general.
        """
        return True

    def update_namespace(self):
        """
        Update the namespace for this timestep. Should only deal with variables
        where *the reference* changes every timestep, i.e. where the current
        reference in `namespace` is not correct.
        """
        pass

    @abstractmethod
    def compile_block(self, block): ...

    def compile(self):
        for block in ["before_run", "run", "after_run"]:
            self.compiled_code[block] = self.compile_block(block)

    def __call__(self, **kwds):
        self.update_namespace()
        self.namespace.update(**kwds)

        return self.run()

    @abstractmethod
    def run_block(self, block): ...

    def before_run(self):
        """
        Runs the preparation code in the namespace. This code will only be
        executed once per run.

        Returns
        -------
        return_value : dict
            A dictionary with the keys corresponding to the `output_variables`
            defined during the call of `CodeGenerator.code_object`.
        """
        return self.run_block("before_run")

    def run(self):
        """
        Runs the main code in the namespace.

        Returns
        -------
        return_value : dict
            A dictionary with the keys corresponding to the `output_variables`
            defined during the call of `CodeGenerator.code_object`.
        """
        return self.run_block("run")

    def after_run(self):
        """
        Runs the finalizing code in the namespace. This code will only be
        executed once per run.

        Returns
        -------
        return_value : dict
            A dictionary with the keys corresponding to the `output_variables`
            defined during the call of `CodeGenerator.code_object`.
        """
        return self.run_block("after_run")


def _error_msg(code, name):
    """
    Little helper function for error messages.
    """
    error_msg = f"Error generating code for code object '{name}' "
    code_lines = [line for line in code.split("\n") if len(line.strip())]
    # If the abstract code is only one line, display it in full
    if len(code_lines) <= 1:
        error_msg += f"from this abstract code: '{code_lines[0]}'\n"
    else:
        error_msg += (
            f"from {len(code_lines)} lines of abstract code, first line is: "
            "'code_lines[0]'\n"
        )
    return error_msg


def check_compiler_kwds(compiler_kwds, accepted_kwds, target):
    """
    Internal function to check the provided compiler keywords against the list
    of understood keywords.

    Parameters
    ----------
    compiler_kwds : dict
        Dictionary of compiler keywords and respective list of values.
    accepted_kwds : list of str
        The compiler keywords understood by the code generation target
    target : str
        The name of the code generation target (used for the error message).

    Raises
    ------
    ValueError
        If a compiler keyword is not understood
    """
    for key in compiler_kwds:
        if key not in accepted_kwds:
            formatted_kwds = ", ".join(f"'{kw}'" for kw in accepted_kwds)
            raise ValueError(
                f"The keyword argument '{key}' is not understood by "
                f"the code generation target '{target}'. The valid "
                f"arguments are: {formatted_kwds}."
            )


def _merge_compiler_kwds(list_of_kwds):
    """
    Merges a list of keyword dictionaries. Values in these dictionaries are
    lists of values, the merged dictionaries will contain the concatenations
    of lists specified for the same key.

    Parameters
    ----------
    list_of_kwds : list of dict
        A list of compiler keyword dictionaries that should be merged.

    Returns
    -------
    merged_kwds : dict
        The merged dictionary
    """
    merged_kwds = collections.defaultdict(list)
    for kwds in list_of_kwds:
        for key, values in kwds.items():
            if not isinstance(values, list):
                raise TypeError(
                    f"Compiler keyword argument '{key}' requires a list of values."
                )
            merged_kwds[key].extend(values)
    return merged_kwds


def _gather_compiler_kwds(function, codeobj_class):
    """
    Gather all the compiler keywords for a function and its dependencies.

    Parameters
    ----------
    function : `Function`
        The function for which the compiler keywords should be gathered
    codeobj_class : type
        The class of `CodeObject` to use

    Returns
    -------
    kwds : dict
        A dictionary with the compiler arguments, a list of values for each
        key.
    """
    implementation = function.implementations[codeobj_class]
    all_kwds = [implementation.compiler_kwds]
    if implementation.dependencies is not None:
        for dependency in implementation.dependencies.values():
            all_kwds.append(_gather_compiler_kwds(dependency, codeobj_class))
    return _merge_compiler_kwds(all_kwds)


def create_runner_codeobj(
    group,
    code,
    template_name,
    run_namespace,
    user_code=None,
    variable_indices=None,
    name=None,
    check_units=True,
    needed_variables=None,
    additional_variables=None,
    template_kwds=None,
    override_conditional_write=None,
    codeobj_class=None,
):
    """Create a `CodeObject` for the execution of code in the context of a
    `Group`.

    Parameters
    ----------
    group : `Group`
        The group where the code is to be run
    code : str or dict of str
        The code to be executed.
    template_name : str
        The name of the template to use for the code.
    run_namespace : dict-like
        An additional namespace that is used for variable lookup (either
        an explicitly defined namespace or one taken from the local
        context).
    user_code : str, optional
        The code that had been specified by the user before other code was
        added automatically. If not specified, will be assumed to be identical
        to ``code``.
    variable_indices : dict-like, optional
        A mapping from `Variable` objects to index names (strings).  If none is
        given, uses the corresponding attribute of `group`.
    name : str, optional
        A name for this code object, will use ``group + '_codeobject*'`` if
        none is given.
    check_units : bool, optional
        Whether to check units in the statement. Defaults to ``True``.
    needed_variables: list of str, optional
        A list of variables that are neither present in the abstract code, nor
        in the ``USES_VARIABLES`` statement in the template. This is only
        rarely necessary, an example being a `StateMonitor` where the
        names of the variables are neither known to the template nor included
        in the abstract code statements.
    additional_variables : dict-like, optional
        A mapping of names to `Variable` objects, used in addition to the
        variables saved in `group`.
    template_kwds : dict, optional
        A dictionary of additional information that is passed to the template.
    override_conditional_write: list of str, optional
        A list of variable names which are used as conditions (e.g. for
        refractoriness) which should be ignored.
    codeobj_class : class, optional
        The `CodeObject` class to run code with. If not specified, defaults to
        the `group`'s ``codeobj_class`` attribute.
    """

    if name is None:
        if group is not None:
            name = f"{group.name}_{template_name}_codeobject*"
        else:
            name = f"{template_name}_codeobject*"

    if user_code is None:
        user_code = code

    if isinstance(code, str):
        code = {None: code}
        user_code = {None: user_code}

    msg = (
        f"Creating code object (group={group.name}, template name={template_name}) for"
        " abstract code:\n"
    )
    msg += indent(code_representation(code))
    logger.diagnostic(msg)
    from brian2.devices import get_device

    device = get_device()

    if override_conditional_write is None:
        override_conditional_write = set()
    else:
        override_conditional_write = set(override_conditional_write)

    if codeobj_class is None:
        codeobj_class = device.code_object_class(group.codeobj_class)
    else:
        codeobj_class = device.code_object_class(codeobj_class)

    template = getattr(codeobj_class.templater, template_name)
    template_variables = getattr(template, "variables", None)

    all_variables = dict(group.variables)
    if additional_variables is not None:
        all_variables.update(additional_variables)

    # Determine the identifiers that were used
    identifiers = set()
    user_identifiers = set()
    for v, u_v in zip(code.values(), user_code.values(), strict=True):
        _, uk, u = analyse_identifiers(v, all_variables, recursive=True)
        identifiers |= uk | u
        _, uk, u = analyse_identifiers(u_v, all_variables, recursive=True)
        user_identifiers |= uk | u

    # Add variables that are not in the abstract code, nor specified in the
    # template but nevertheless necessary
    if needed_variables is None:
        needed_variables = []
    # Resolve all variables (variables used in the code and variables needed by
    # the template)
    variables = group.resolve_all(
        sorted(identifiers | set(needed_variables) | set(template_variables)),
        # template variables are not known to the user:
        user_identifiers=user_identifiers,
        additional_variables=additional_variables,
        run_namespace=run_namespace,
    )
    # We raise this error only now, because there is some non-obvious code path
    # where Jinja tries to get a Synapse's "name" attribute via syn['name'],
    # which then triggers the use of the `group_get_indices` template which does
    # not exist for standalone. Putting the check for template == None here
    # means we will first raise an error about the unknown identifier which will
    # then make Jinja try syn.name
    if template is None:
        codeobj_class_name = codeobj_class.class_name or codeobj_class.__name__
        raise AttributeError(
            f"'{codeobj_class_name}' does not provide a code "
            f"generation template '{template_name}'"
        )

    conditional_write_variables = {}
    # Add all the "conditional write" variables
    for var in variables.values():
        cond_write_var = getattr(var, "conditional_write", None)
        if cond_write_var in override_conditional_write:
            continue
        if cond_write_var is not None:
            if (
                cond_write_var.name in variables
                and variables[cond_write_var.name] is not cond_write_var
            ):
                logger.diagnostic(
                    f"Variable '{cond_write_var.name}' is needed for the "
                    "conditional write mechanism of variable "
                    f"'{var.name}'. Its name is already used for "
                    f"{variables[cond_write_var.name]!r}."
                )
            else:
                conditional_write_variables[cond_write_var.name] = cond_write_var

    variables.update(conditional_write_variables)

    if check_units:
        for c in code.values():
            # This is the first time that the code is parsed, catch errors
            try:
                check_units_statements(c, variables)
            except (SyntaxError, ValueError) as ex:
                error_msg = _error_msg(c, name)
                raise ValueError(error_msg) from ex

    all_variable_indices = copy.copy(group.variables.indices)
    if additional_variables is not None:
        all_variable_indices.update(additional_variables.indices)
    if variable_indices is not None:
        all_variable_indices.update(variable_indices)

    # Make "conditional write" variables use the same index as the variable
    # that depends on them
    for varname, var in variables.items():
        cond_write_var = getattr(var, "conditional_write", None)
        if cond_write_var is not None:
            all_variable_indices[cond_write_var.name] = all_variable_indices[varname]

    # Check that all functions are available
    for varname, value in variables.items():
        if isinstance(value, Function):
            try:
                value.implementations[codeobj_class]
            except KeyError as ex:
                # if we are dealing with numpy, add the default implementation
                from brian2.codegen.runtime.numpy_rt import NumpyCodeObject

                if codeobj_class is NumpyCodeObject:
                    value.implementations.add_numpy_implementation(value.pyfunc)
                else:
                    raise NotImplementedError(
                        f"Cannot use function '{varname}': {ex}"
                    ) from ex

    # Gather the additional compiler arguments declared by function
    # implementations
    all_keywords = [
        _gather_compiler_kwds(var, codeobj_class)
        for var in variables.values()
        if isinstance(var, Function)
    ]
    compiler_kwds = _merge_compiler_kwds(all_keywords)

    # Add the indices needed by the variables
    for varname in list(variables):
        var_index = all_variable_indices[varname]
        if var_index not in ("_idx", "0"):
            variables[var_index] = all_variables[var_index]

    return device.code_object(
        owner=group,
        name=name,
        abstract_code=code,
        variables=variables,
        template_name=template_name,
        variable_indices=all_variable_indices,
        template_kwds=template_kwds,
        codeobj_class=codeobj_class,
        override_conditional_write=override_conditional_write,
        compiler_kwds=compiler_kwds,
    )
