# Author: Nathan Trouvain at 14/10/2021 <nathan.trouvain@inria.fr>
# Licence: MIT License
# Copyright: Xavier Hinaut (2018) <xavier.hinaut@inria.fr>
import itertools as it
from math import comb
from typing import Optional, Sequence, Union

import numpy as np

from ..node import Node
from ..type import NodeInput, State, Timestep


class NVAR(Node):
    """Non-linear Vector AutoRegressive machine.

    NVAR is implemented as described in [1]_.

    The state :math:`\\mathbb{O}_{total}` of the NVAR first contains a series of linear
    features :math:`\\mathbb{O}_{lin}` made of input data concatenated
    with delayed inputs:

    .. math::

        \\mathbb{O}_{lin}[t] = \\mathbf{X}[t] \\oplus \\mathbf{X}[t - s] \\oplus
        \\mathbf{X}[t - 2s] \\oplus \\dots \\oplus \\mathbf{X}[t - (k-1)s]

    where :math:`\\mathbf{X}[t]` are the inputs at time :math:`t`, :math:`k` is the
    delay and :math:`s` is the strides (only one input every :math:`s`
    inputs within the delayed inputs is used).
    The operator :math:`\\oplus` denotes the concatenation.

    In addition to these linear features, nonlinear representations
    :math:`\\mathbb{O}_{nonlin}^n` of the inputs are constructed using all unique
    monomials of order :math:`n` of these inputs:

    .. math::

        \\mathbb{O}_{nonlin}^n[t] = \\mathbb{O}_{lin}[t] \\otimes \\mathbb{O}_{lin}[t]
        \\overbrace{\\otimes \\dots \\otimes}^{n-1~\\mathrm{times}} \\mathbb{O}_{lin}[t]

    where :math:`\\otimes` is the operator denoting an outer product followed by the
    selection of all unique monomials generated by this outer product.

    Note
    ----

        Under the hood,
        this product is computed by finding all unique combinations
        of input features and multiplying each combination of terms.

    Finally, all representations are gathered to form the final feature
    vector :math:`\\mathbb{O}_{total}`:

    .. math::

        \\mathbb{O}_{total} = \\mathbb{O}_{lin}[t] \\oplus \\mathbb{O}_{nonlin}^n[t]


    Parameters
    ----------
    delay : int
        Maximum delay of inputs.
    order : int
        Order of the non-linear monomials.
    strides : int, default to 1
        Strides between delayed inputs.
    input_dim : int, optional
        Input dimension. Can be inferred at first call.
    name : str, optional
        Node name.

    References
    ----------
    .. [1] Gauthier, D. J., Bollt, E., Griffith, A., & Barbosa, W. A. S. (2021).
           Next generation reservoir computing. Nature Communications, 12(1), 5564.
           https://doi.org/10.1038/s41467-021-25801-2

    Example
    -------

    >>> import numpy as np
    >>> from reservoirpy.nodes import NVAR, Ridge
    >>> nvar = NVAR(delay=2, order=2, strides=1)
    >>> readout = Ridge(ridge=2.5e-6, output_dim=3)
    >>> model = nvar >> readout

    Using the :py:func:`~reservoirpy.datasets.lorenz` timeseries and learning to
    predict the next difference:

    >>> from reservoirpy.datasets import lorenz
    >>> X = lorenz(5400, x0=[17.677, 12.931, 43.914], h=0.025, method="RK23")
    >>> Xi  = X[:600]
    >>> dXi = X[1:601] - X[:600]  # difference u[t+1] - u[t]
    >>> Y_test = X[600:]  # testing data
    >>> _ = model.fit(Xi, dXi, warmup=200)

    We can now predict the differences and integrate these predictions:

    >>> u = X[600]
    >>> res = np.zeros((5400-600, readout.output_dim))
    >>> for i in range(5400-600):
    ...     u = u + model(u)
    ...     res[i, :] = u
    ...

    .. plot:: ./api/nvar_example.py

    """

    #: Time window over the inputs (of shape (delay * strides, features)).
    store: np.ndarray
    #: Maximum delay of inputs (:math:`k`).
    delay: int
    #: Order of the non-linear monomials (:math:`n`).
    order: int
    #: Strides between delayed inputs, by default 1 (:math:`s`).
    strides: int

    def __init__(
        self,
        delay: int,
        order: int,
        strides: int = 1,
        input_dim: Optional[int] = None,
        name: Optional[str] = None,
    ):
        self.store = None
        self._monomial_idx = None
        self.delay = delay
        self.order = order
        self.strides = strides
        self.name = name
        self.initialized = False
        self.state = {}
        self.input_dim = input_dim
        self.output_dim = None

    def initialize(self, x: Union[NodeInput, Timestep]):

        # set input_dim
        self._set_input_dim(x)

        order = self.order
        delay = self.delay
        strides = self.strides

        linear_dim = delay * self.input_dim
        # number of non linear components is (d + n - 1)! / (d - 1)! n!
        # i.e. number of all unique monomials of order n made from the
        # linear components.
        nonlinear_dim = comb(linear_dim + order - 1, order)

        output_dim = int(linear_dim + nonlinear_dim)

        self.output_dim = output_dim

        # for each monomial created in the non linear part, indices
        # of the n components involved, n being the order of the
        # monomials. Pre-compute them to improve efficiency.
        idx = np.array(
            list(it.combinations_with_replacement(np.arange(linear_dim), order))
        )

        self._monomial_idx = idx

        # to store the k*s last inputs, k being the delay and s the strides
        self.store = np.zeros((delay * strides, self.input_dim))
        self.initialized = True

    def _step(self, state: tuple, x: Timestep) -> State:
        store = self.store
        strides = self.strides
        idxs = self._monomial_idx
        output_dim = self.output_dim

        # store the current input
        new_store = np.roll(store, 1, axis=0)
        new_store[0] = x
        self.store = new_store

        output = np.zeros((output_dim,))

        # select all previous inputs, including the current, with strides
        linear_feats = np.ravel(new_store[::strides, :])
        linear_len = linear_feats.shape[0]

        output[:linear_len] = linear_feats

        # select monomial terms and compute them
        output[linear_len:] = np.prod(linear_feats[idxs], axis=1)

        return {"out": output}
