"""
Implementation of `TimedArray`.
"""

import numpy as np

from rockpool.utilities.backend_management import backend_available

if not backend_available("brian2"):
    raise ModuleNotFoundError(
        "'Brian2' backend not found. Modules that rely on Brian will not be available."
    )

from brian2.core.clocks import defaultclock
from brian2.core.functions import Function
from brian2.units.allunits import second
from brian2.units.fundamentalunits import (
    check_units,
    get_dimensions,
    Quantity,
    get_unit,
)
from brian2.core.names import Nameable
from brian2.utils.logger import get_logger
from brian2.utils.stringtools import replace


__all__ = ["TimedArray"]


logger = get_logger(__name__)


def _find_K(group_dt, dt):
    dt_ratio = dt / group_dt
    if dt_ratio > 1 and np.floor(dt_ratio) != dt_ratio:
        logger.warn(
            ("Group uses a dt of %s while TimedArray uses dt " "of %s")
            % (group_dt * second, dt * second),
            once=True,
        )
    # Find an upsampling factor that should avoid rounding issues even
    # for multistep methods
    K = max(int(2 ** np.ceil(np.log2(8 / group_dt * dt))), 1)
    return K


class TimedArray(Function, Nameable):
    """
    TimedArray(values, dt, name=None, tOffset=0*second)

    A function of time built from an array of values. The returned object can
    be used as a function, including in model equations etc. The resulting
    function has to be called as `funcion_name(t)` if the provided value array
    is one-dimensional and as `function_name(t, i)` if it is two-dimensional.

    Parameters
    ----------
    values : ndarray or `Quantity`
        An array of values providing the values at various points in time. This
        array can either be one- or two-dimensional. If it is two-dimensional
        it's first dimension should be the time.
    dt : `Quantity`
        The time distance between values in the `values` array.
    name : str, optional
        A unique name for this object, see `Nameable` for details. Defaults
        to ``'_timedarray*'``.
    tOffset : float, optional
        A temporal offset at which the data in the TimedArray should begin. Prior to
        t = tOffset, the first element of the array will be used

    Notes
    -----
    For time values corresponding to elements outside of the range of `values`
    provided, the first respectively last element is returned.

    Examples
    --------
    >>> from brian2 import *
    >>> ta = TimedArray([1, 2, 3, 4] * mV, dt=0.1*ms)
    >>> print((ta(0.3*ms)))
    4. mV
    >>> G = NeuronGroup(1, 'v = ta(t) : volt')
    >>> mon = StateMonitor(G, 'v', record=True)
    >>> net = Network(G, mon)
    >>> net.run(1*ms)  # doctest: +ELLIPSIS
    ...
    >>> print((mon[0].v))
    [ 1.  2.  3.  4.  4.  4.  4.  4.  4.  4.] mV
    >>> ta2d = TimedArray([[1, 2], [3, 4], [5, 6]]*mV, dt=0.1*ms)
    >>> G = NeuronGroup(4, 'v = ta2d(t, i%2) : volt')
    >>> mon = StateMonitor(G, 'v', record=True)
    >>> net = Network(G, mon)
    >>> net.run(0.2*ms)  # doctest: +ELLIPSIS
    ...
    >>> print(mon.v[:])
    [[ 1.  3.]
     [ 2.  4.]
     [ 1.  3.]
     [ 2.  4.]] mV
    """

    @check_units(dt=second, tOffset=second)
    def __init__(self, values, dt, name=None, tOffset=0 * second):
        if name is None:
            name = "_timedarray*"
        Nameable.__init__(self, name)
        dimensions = get_dimensions(values)
        self.dim = dimensions
        values = np.asarray(values, dtype=np.double)
        self.values = values
        dt = float(dt)
        self.dt = dt
        self.tOffset = tOffset
        if values.ndim == 1:
            self._init_1d()
        elif values.ndim == 2:
            self._init_2d()
        else:
            raise NotImplementedError(
                ("Only 1d and 2d arrays are supported " "for TimedArray")
            )

    def _init_1d(self):
        dimensions = self.dim
        unit = get_unit(dimensions)
        values = self.values
        dt = self.dt
        tOffset = self.tOffset

        # Python implementation (with units), used when calling the TimedArray
        # directly, outside of a simulation
        @check_units(t=second, result=unit)
        def timed_array_func(t):
            t -= tOffset
            # We round according to the current defaultclock.dt
            K = _find_K(float(defaultclock.dt), dt)
            epsilon = dt / K
            i = np.clip(
                np.int_(np.round(np.asarray(t / epsilon)) / K), 0, len(values) - 1
            )
            return Quantity(values[i], dim=dimensions)

        Function.__init__(self, pyfunc=timed_array_func)

        # we use dynamic implementations because we want to do upsampling
        # in a way that avoids rounding problems with the group's dt
        def create_numpy_implementation(owner):
            group_dt = owner.clock.dt_

            K = _find_K(group_dt, dt)
            n_values = len(values)
            epsilon = dt / K

            def unitless_timed_array_func(t):
                t = t - tOffset
                _timestep = np.clip(np.int_(np.round(t / epsilon) / K), 0, n_values - 1)
                return values[_timestep]

            unitless_timed_array_func._arg_units = [second]
            unitless_timed_array_func._return_unit = unit

            return unitless_timed_array_func

        self.implementations.add_dynamic_implementation(
            "numpy", create_numpy_implementation
        )

        def create_cpp_implementation(owner):
            group_dt = owner.clock.dt_
            K = _find_K(group_dt, dt)
            support_code = (
                """
            inline double %NAME%(const double t)
            {
                const double epsilon = %DT% / %K%;
                int i = (int)(((t-%OFFSET%)/epsilon + 0.5)/%K%);
                if(i < 0)
                   i = 0;
                if(i >= %NUM_VALUES%)
                    i = %NUM_VALUES%-1;
                return _namespace%NAME%_values[i];
            }
            """.replace(
                    "%NAME%", self.name
                )
                .replace("%DT%", "%.18f" % dt)
                .replace("%K%", str(K))
                .replace("%NUM_VALUES%", str(len(self.values)))
                .replace("%OFFSET%", str(np.asarray(tOffset)))
            )
            cpp_code = {"support_code": support_code}

            return cpp_code

        def create_cpp_namespace(owner):
            return {"%s_values" % self.name: self.values}

        self.implementations.add_dynamic_implementation(
            "cpp",
            code=create_cpp_implementation,
            namespace=create_cpp_namespace,
            name=self.name,
        )

        def create_cython_implementation(owner):
            group_dt = owner.clock.dt_
            K = _find_K(group_dt, dt)
            code = (
                """
            cdef double %NAME%(const double t):
                global _namespace%NAME%_values
                cdef double epsilon = %DT% / %K%
                cdef int i = (int)(((t-%OFFSET%)/epsilon + 0.5)/%K%)
                if i < 0:
                   i = 0
                if i >= %NUM_VALUES%:
                    i = %NUM_VALUES% - 1
                return _namespace%NAME%_values[i]
            """.replace(
                    "%NAME%", self.name
                )
                .replace("%DT%", "%.18f" % dt)
                .replace("%K%", str(K))
                .replace("%NUM_VALUES%", str(len(self.values)))
                .replace("%OFFSET%", str(np.asarray(tOffset)))
            )

            return code

        def create_cython_namespace(owner):
            return {"%s_values" % self.name: self.values}

        self.implementations.add_dynamic_implementation(
            "cython",
            code=create_cython_implementation,
            namespace=create_cython_namespace,
            name=self.name,
        )

    def _init_2d(self):
        dimensions = self.dim
        unit = get_unit(dimensions)
        values = self.values
        dt = self.dt
        tOffset = self.tOffset

        # Python implementation (with units), used when calling the TimedArray
        # directly, outside of a simulation
        @check_units(i=1, t=second, result=unit)
        def timed_array_func(t, i):
            t -= tOffset
            # We round according to the current defaultclock.dt
            K = _find_K(float(defaultclock.dt), dt)
            epsilon = dt / K
            time_step = np.clip(
                np.int_(np.round(np.asarray(t / epsilon)) / K), 0, len(values) - 1
            )
            return Quantity(values[time_step, i], dim=dimensions)

        Function.__init__(self, pyfunc=timed_array_func)

        # we use dynamic implementations because we want to do upsampling
        # in a way that avoids rounding problems with the group's dt
        def create_numpy_implementation(owner):
            group_dt = owner.clock.dt_

            K = _find_K(group_dt, dt)
            n_values = len(values)
            epsilon = dt / K

            def unitless_timed_array_func(t, i):
                t -= np.asarray(tOffset)
                _timestep = np.clip(np.int_(np.round(t / epsilon) / K), 0, n_values - 1)
                return values[_timestep, i]

            unitless_timed_array_func._arg_units = [second]
            unitless_timed_array_func._return_unit = unit

            return unitless_timed_array_func

        self.implementations.add_dynamic_implementation(
            "numpy", create_numpy_implementation
        )

        def create_cpp_implementation(owner):
            group_dt = owner.clock.dt_
            K = _find_K(group_dt, dt)
            support_code = """
            inline double %NAME%(const double t, const int i)
            {
                const double epsilon = %DT% / %K%;
                if (i < 0 || i >= %COLS%)
                    return NAN;
                int timestep = (int)(((t-%OFFSET%)/epsilon + 0.5)/%K%);
                if(timestep < 0)
                   timestep = 0;
                else if(timestep >= %ROWS%)
                    timestep = %ROWS%-1;
                return _namespace%NAME%_values[timestep*%COLS% + i];
            }
            """
            support_code = replace(
                support_code,
                {
                    "%NAME%": self.name,
                    "%DT%": "%.18f" % dt,
                    "%K%": str(K),
                    "%COLS%": str(self.values.shape[1]),
                    "%ROWS%": str(self.values.shape[0]),
                    "%OFFSET%": str(np.asarray(tOffset)),
                },
            )
            cpp_code = {"support_code": support_code}

            return cpp_code

        def create_cpp_namespace(owner):
            return {
                "%s_values"
                % self.name: self.values.astype(
                    np.double, order="C", copy=False
                ).ravel()
            }

        self.implementations.add_dynamic_implementation(
            "cpp",
            code=create_cpp_implementation,
            namespace=create_cpp_namespace,
            name=self.name,
        )

        def create_cython_implementation(owner):
            group_dt = owner.clock.dt_
            K = _find_K(group_dt, dt)
            code = """
            cdef double %NAME%(const double t, const int i):
                global _namespace%NAME%_values
                cdef double epsilon = %DT% / %K%
                if i < 0 or i >= %COLS%:
                    return _numpy.nan
                cdef int timestep = (int)(((t-%OFFSET%)/epsilon + 0.5)/%K%)
                if timestep < 0:
                   timestep = 0
                elif timestep >= %ROWS%:
                    timestep = %ROWS%-1
                return _namespace%NAME%_values[timestep*%COLS% + i]
            """
            code = replace(
                code,
                {
                    "%NAME%": self.name,
                    "%DT%": "%.18f" % dt,
                    "%K%": str(K),
                    "%COLS%": str(self.values.shape[1]),
                    "%ROWS%": str(self.values.shape[0]),
                    "%OFFSET%": str(np.asarray(tOffset)),
                },
            )

            return code

        def create_cython_namespace(owner):
            return {
                "%s_values"
                % self.name: self.values.astype(
                    np.double, order="C", copy=False
                ).ravel()
            }

        self.implementations.add_dynamic_implementation(
            "cython",
            code=create_cython_implementation,
            namespace=create_cython_namespace,
            name=self.name,
        )

    def is_locally_constant(self, dt):
        if dt > self.dt:
            return False
        dt_ratio = self.dt / float(dt)
        if np.floor(dt_ratio) != dt_ratio:
            logger.info(
                (
                    "dt of the TimedArray is not an integer multiple of "
                    "the group's dt, the TimedArray's return value can "
                    "therefore not be considered constant over one "
                    "timestep, making linear integration impossible."
                ),
                once=True,
            )
            return False
        return True
