# 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
import sys

if sys.version_info < (3, 8):
    from scipy.special import comb
else:
    from math import comb

import numpy as np

from ...node import Node


def forward(node, x):
    store = node.store
    strides = node.strides
    idxs = node._monomial_idx

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

    output = np.zeros((node.output_dim, 1))

    # select all previous inputs, including the current, with strides
    linear_feats = np.ravel(new_store[::strides, :]).reshape(-1, 1)
    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.astype(int)], axis=1)

    return output.reshape(1, -1)


def initialize(node, x=None, *args, **kwargs):
    if x is not None:
        input_dim = x.shape[1]

        order = node.order
        delay = node.delay
        strides = node.strides

        linear_dim = delay * 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)

        node.set_output_dim(output_dim)
        node.set_input_dim(input_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))
        )

        node.set_param("_monomial_idx", idx)

        # to store the k*s last inputs, k being the delay and s the strides
        node.set_param("store", np.zeros((delay * strides, node.input_dim)))


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]


    :py:attr:`NVAR.params` **list:**

    ================== ===================================================================
    ``store``          Time window over the inputs (of shape (delay * strides, features)).
    ================== ===================================================================

    :py:attr:`NVAR.hypers` **list:**

    ================== =================================================================
    ``delay``          Maximum delay of inputs (:math:`k`).
    ``order``          Order of the non-linear monomials (:math:`n`).
    ``strides``        Strides between delayed inputs, by default 1 (:math:`s`).
    ================== =================================================================

    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(3, ridge=2.5e-6)
    >>> 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

    """

    def __init__(self, delay: int, order: int, strides: int = 1, **kwargs):
        super(NVAR, self).__init__(
            params={"store": None, "_monomial_idx": None},
            hypers={"delay": delay, "order": order, "strides": strides},
            forward=forward,
            initializer=initialize,
            **kwargs,
        )
