"""
Module providing `NumpyCodeObject`.
"""

import sys
from collections.abc import Iterable

import numpy as np

from brian2.core.base import BrianObjectException
from brian2.core.functions import Function
from brian2.core.preferences import BrianPreference, prefs
from brian2.core.variables import (
    ArrayVariable,
    AuxiliaryVariable,
    DynamicArrayVariable,
    Subexpression,
)

from ...codeobject import CodeObject, check_compiler_kwds, constant_or_scalar
from ...generators.numpy_generator import NumpyCodeGenerator
from ...targets import codegen_targets
from ...templates import Templater

__all__ = ["NumpyCodeObject"]

# Preferences
prefs.register_preferences(
    "codegen.runtime.numpy",
    "Numpy runtime codegen preferences",
    discard_units=BrianPreference(
        default=False,
        docs="""
        Whether to change the namespace of user-specifed functions to remove
        units.
        """,
    ),
)


class LazyArange(Iterable):
    """
    A class that can be used as a `~numpy.arange` replacement (with an implied
    step size of 1) but does not actually create an array of values until
    necessary. It is somewhat similar to the ``range()`` function in Python 3,
    but does not use a generator. It is tailored to a special use case, the
    ``_vectorisation_idx`` variable in numpy templates, and not meant for
    general use. The ``_vectorisation_idx`` is used for stateless function
    calls such as ``rand()`` and for the numpy codegen target determines the
    number of values produced by such a call. This will often be the number of
    neurons or synapses, and this class avoids creating a new array of that size
    at every code object call when all that is needed is the *length* of the
    array.

    Examples
    --------
    >>> from brian2.codegen.runtime.numpy_rt.numpy_rt import LazyArange
    >>> ar = LazyArange(10)
    >>> len(ar)
    10
    >>> len(ar[:5])
    5
    >>> type(ar[:5])
    <class 'brian2.codegen.runtime.numpy_rt.numpy_rt.LazyArange'>
    >>> ar[5]
    5
    >>> for value in ar[3:7]:
    ...     print(value)
    ...
    3
    4
    5
    6
    >>> len(ar[np.array([1, 2, 3])])
    3
    """

    def __init__(self, stop, start=0, indices=None):
        self.start = start
        self.stop = stop
        self.indices = indices

    def __len__(self):
        if self.indices is None:
            return self.stop - self.start
        else:
            return len(self.indices)

    def __getitem__(self, item):
        if isinstance(item, slice):
            if self.indices is None:
                start, stop, step = item.start, item.stop, item.step
                if step not in [None, 1]:
                    raise NotImplementedError("Step should be 1")
                if start is None:
                    start = 0
                if stop is None:
                    stop = len(self)
                return LazyArange(
                    start=self.start + start, stop=min([self.start + stop, self.stop])
                )
            else:
                raise NotImplementedError("Cannot slice LazyArange with indices")
        elif isinstance(item, np.ndarray):
            if item.dtype == np.dtype(bool):
                item = np.nonzero(item)[0]  # convert boolean array into integers
            if len(item) == 0:
                return np.array([], dtype=np.int32)
            if np.min(item) < 0 or np.max(item) > len(self):
                raise IndexError("Indexing array contains out-of-bounds values")
            return LazyArange(start=self.start, stop=self.stop, indices=item)
        elif isinstance(item, int):
            if self.indices is None:
                index = self.start + item
                if index >= self.stop:
                    raise IndexError(index)
                return index
            else:
                return self.indices[item]
        else:
            raise TypeError("Can only index with integer, numpy array, or slice.")

    def __iter__(self):
        if self.indices is None:
            return iter(np.arange(self.start, self.stop))
        else:
            return iter(self.indices)

    # Allow conversion to a proper array with np.array(...)
    def __array__(self, dtype=None, copy=None):
        if copy is False:
            raise ValueError("LazyArange does not support copy=False")
        if self.indices is None:
            return np.arange(self.start, self.stop, dtype=dtype)
        else:
            return (self.indices + self.start).astype(dtype)

    # Allow basic arithmetics (used when shifting stuff for subgroups)
    def __add__(self, other):
        if isinstance(other, int):
            return LazyArange(start=self.start + other, stop=self.stop + other)
        else:
            return NotImplemented

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        if isinstance(other, int):
            return LazyArange(start=self.start - other, stop=self.stop - other)
        else:
            return NotImplemented


class NumpyCodeObject(CodeObject):
    """
    Execute code using Numpy

    Default for Brian because it works on all platforms.
    """

    templater = Templater(
        "brian2.codegen.runtime.numpy_rt",
        ".py_",
        env_globals={"constant_or_scalar": constant_or_scalar},
    )
    generator_class = NumpyCodeGenerator
    class_name = "numpy"

    def __init__(
        self,
        owner,
        code,
        variables,
        variable_indices,
        template_name,
        template_source,
        compiler_kwds,
        name="numpy_code_object*",
    ):
        check_compiler_kwds(compiler_kwds, [], "numpy")
        from brian2.devices.device import get_device

        self.device = get_device()
        self.namespace = {
            "_owner": owner,
            # TODO: This should maybe go somewhere else
            "logical_not": np.logical_not,
        }
        CodeObject.__init__(
            self,
            owner,
            code,
            variables,
            variable_indices,
            template_name,
            template_source,
            compiler_kwds=compiler_kwds,
            name=name,
        )
        self.variables_to_namespace()

    @classmethod
    def is_available(cls):
        # no test necessary for numpy
        return True

    def variables_to_namespace(self):
        # Variables can refer to values that are either constant (e.g. dt)
        # or change every timestep (e.g. t). We add the values of the
        # constant variables here and add the names of non-constant variables
        # to a list

        # A list containing tuples of name and a function giving the value
        self.nonconstant_values = []

        for name, var in self.variables.items():
            if isinstance(var, (AuxiliaryVariable, Subexpression)):
                continue

            try:
                if not hasattr(var, "get_value"):
                    raise TypeError()
                value = var.get_value()
            except TypeError:
                # Either a dummy Variable without a value or a Function object
                if isinstance(var, Function):
                    impl = var.implementations[self.__class__].get_code(self.owner)
                    self.namespace[name] = impl
                else:
                    self.namespace[name] = var
                continue

            if isinstance(var, ArrayVariable):
                self.namespace[self.generator_class.get_array_name(var)] = value
                if var.scalar and var.constant:
                    self.namespace[name] = value[0]
            else:
                self.namespace[name] = value

            if isinstance(var, DynamicArrayVariable):
                dyn_array_name = self.generator_class.get_array_name(
                    var, access_data=False
                )
                self.namespace[dyn_array_name] = self.device.get_value(
                    var, access_data=False
                )

            # Also provide the Variable object itself in the namespace (can be
            # necessary for resize operations, for example)
            self.namespace[f"_var_{name}"] = var

            # There is one type of objects that we have to inject into the
            # namespace with their current value at each time step: dynamic
            # arrays that change in size during runs (i.e. not synapses but
            # e.g. the structures used in monitors)
            if isinstance(var, DynamicArrayVariable) and var.needs_reference_update:
                self.nonconstant_values.append(
                    (
                        self.generator_class.get_array_name(var, self.variables),
                        var.get_value,
                    )
                )

    def update_namespace(self):
        # update the values of the non-constant values in the namespace
        for name, func in self.nonconstant_values:
            self.namespace[name] = func()

    def compile_block(self, block):
        code = getattr(self.code, block, "").strip()
        if not code or "EMPTY_CODE_BLOCK" in code:
            return None
        return compile(code, "(string)", "exec")

    def run_block(self, block):
        compiled_code = self.compiled_code[block]
        if not compiled_code:
            return
        try:
            exec(compiled_code, self.namespace)
        except Exception as exc:
            code = getattr(self.code, block)
            message = (
                "An exception occured during the execution of the "
                f"'{block}' block of code object {self.name}.\n"
            )
            lines = code.split("\n")
            message += "The error was raised in the following line:\n"
            _, _, tb = sys.exc_info()
            tb = tb.tb_next  # Line in the code object's code
            message += f"{lines[tb.tb_lineno - 1]}\n"
            raise BrianObjectException(message, self.owner) from exc
        # output variables should land in the variable name _return_values
        if "_return_values" in self.namespace:
            return self.namespace["_return_values"]


codegen_targets.add(NumpyCodeObject)