"""Helper and filter functions for VAR and VARMA, and basic VAR class

Created on Mon Jan 11 11:04:23 2010
Author: josef-pktd
License: BSD

This is a new version, I did not look at the old version again, but similar
ideas.

not copied/cleaned yet:
 * fftn based filtering, creating samples with fft
 * Tests: I ran examples but did not convert them to tests
   examples look good for parameter estimate and forecast, and filter functions

main TODOs:
* result statistics
* see whether Bayesian dummy observation can be included without changing
  the single call to linalg.lstsq
* impulse response function does not treat correlation, see Hamilton and jplv

Extensions
* constraints, Bayesian priors/penalization
* Error Correction Form and Cointegration
* Factor Models Stock-Watson,  ???


see also VAR section in Notes.txt

"""

import numpy as np
from scipy import signal

from statsmodels.tsa.tsatools import lagmat


def varfilter(x, a):
    """apply an autoregressive filter to a series x

    Warning: I just found out that convolve does not work as I
       thought, this likely does not work correctly for
       nvars>3


    x can be 2d, a can be 1d, 2d, or 3d

    Parameters
    ----------
    x : array_like
        data array, 1d or 2d, if 2d then observations in rows
    a : array_like
        autoregressive filter coefficients, ar lag polynomial
        see Notes

    Returns
    -------
    y : ndarray, 2d
        filtered array, number of columns determined by x and a

    Notes
    -----

    In general form this uses the linear filter ::

        y = a(L)x

    where
    x : nobs, nvars
    a : nlags, nvars, npoly

    Depending on the shape and dimension of a this uses different
    Lag polynomial arrays

    case 1 : a is 1d or (nlags,1)
        one lag polynomial is applied to all variables (columns of x)
    case 2 : a is 2d, (nlags, nvars)
        each series is independently filtered with its own
        lag polynomial, uses loop over nvar
    case 3 : a is 3d, (nlags, nvars, npoly)
        the ith column of the output array is given by the linear filter
        defined by the 2d array a[:,:,i], i.e. ::

            y[:,i] = a(.,.,i)(L) * x
            y[t,i] = sum_p sum_j a(p,j,i)*x(t-p,j)
                     for p = 0,...nlags-1, j = 0,...nvars-1,
                     for all t >= nlags


    Note: maybe convert to axis=1, Not

    TODO: initial conditions

    """
    x = np.asarray(x)
    a = np.asarray(a)
    if x.ndim == 1:
        x = x[:, None]
    if x.ndim > 2:
        raise ValueError("x array has to be 1d or 2d")
    nvar = x.shape[1]
    nlags = a.shape[0]
    ntrim = nlags // 2
    # for x is 2d with ncols >1

    if a.ndim == 1:
        # case: identical ar filter (lag polynomial)
        return signal.convolve(x, a[:, None], mode="valid")
        # alternative:
        # return signal.lfilter(a,[1],x.astype(float),axis=0)
    elif a.ndim == 2:
        if min(a.shape) == 1:
            # case: identical ar filter (lag polynomial)
            return signal.convolve(x, a, mode="valid")

        # case: independent ar
        # (a bit like recserar in gauss, but no x yet)
        # (no, reserar is inverse filter)
        result = np.zeros((x.shape[0] - nlags + 1, nvar))
        for i in range(nvar):
            # could also use np.convolve, but easier for swiching to fft
            result[:, i] = signal.convolve(x[:, i], a[:, i], mode="valid")
        return result

    elif a.ndim == 3:
        # case: vector autoregressive with lag matrices
        # Note: we must have shape[1] == shape[2] == nvar
        yf = signal.convolve(x[:, :, None], a)
        yvalid = yf[ntrim:-ntrim, yf.shape[1] // 2, :]
        return yvalid


def varinversefilter(ar, nobs, version=1):
    """creates inverse ar filter (MA representation) recursively

    The VAR lag polynomial is defined by ::

        ar(L) y_t = u_t  or
        y_t = -ar_{-1}(L) y_{t-1} + u_t

    the returned lagpolynomial is arinv(L)=ar^{-1}(L) in ::

        y_t = arinv(L) u_t



    Parameters
    ----------
    ar : ndarray, (nlags,nvars,nvars)
        matrix lagpolynomial, currently no exog
        first row should be identity

    Returns
    -------
    arinv : ndarray, (nobs,nvars,nvars)


    Notes
    -----

    """
    nlags, nvars, nvarsex = ar.shape
    if nvars != nvarsex:
        print("exogenous variables not implemented not tested")
    arinv = np.zeros((nobs + 1, nvarsex, nvars))
    arinv[0, :, :] = ar[0]
    arinv[1:nlags, :, :] = -ar[1:]
    if version == 1:
        for i in range(2, nobs + 1):
            tmp = np.zeros((nvars, nvars))
            for p in range(1, nlags):
                tmp += np.dot(-ar[p], arinv[i - p, :, :])
            arinv[i, :, :] = tmp
    if version == 0:
        for i in range(nlags + 1, nobs + 1):
            print(ar[1:].shape, arinv[i - 1 : i - nlags : -1, :, :].shape)
            # arinv[i,:,:] = np.dot(-ar[1:],arinv[i-1:i-nlags:-1,:,:])
            # print(np.tensordot(-ar[1:],arinv[i-1:i-nlags:-1,:,:],axes=([2],[1])).shape
            # arinv[i,:,:] = np.tensordot(-ar[1:],arinv[i-1:i-nlags:-1,:,:],axes=([2],[1]))
            raise NotImplementedError("waiting for generalized ufuncs or something")

    return arinv


def vargenerate(ar, u, initvalues=None):
    """generate an VAR process with errors u

    similar to gauss
    uses loop

    Parameters
    ----------
    ar : array (nlags,nvars,nvars)
        matrix lagpolynomial
    u : array (nobs,nvars)
        exogenous variable, error term for VAR

    Returns
    -------
    sar : array (1+nobs,nvars)
        sample of var process, inverse filtered u
        does not trim initial condition y_0 = 0

    Examples
    --------
    # generate random sample of VAR
    nobs, nvars = 10, 2
    u = numpy.random.randn(nobs,nvars)
    a21 = np.array([[[ 1. ,  0. ],
                     [ 0. ,  1. ]],

                    [[-0.8,  0. ],
                     [ 0.,  -0.6]]])
    vargenerate(a21,u)

    # Impulse Response to an initial shock to the first variable
    imp = np.zeros((nobs, nvars))
    imp[0,0] = 1
    vargenerate(a21,imp)

    """
    nlags, nvars, nvarsex = ar.shape
    nlagsm1 = nlags - 1
    nobs = u.shape[0]
    if nvars != nvarsex:
        print("exogenous variables not implemented not tested")
    if u.shape[1] != nvars:
        raise ValueError("u needs to have nvars columns")
    if initvalues is None:
        sar = np.zeros((nobs + nlagsm1, nvars))
        start = nlagsm1
    else:
        start = max(nlagsm1, initvalues.shape[0])
        sar = np.zeros((nobs + start, nvars))
        sar[start - initvalues.shape[0] : start] = initvalues
    # sar[nlagsm1:] = u
    sar[start:] = u
    # if version == 1:
    for i in range(start, start + nobs):
        for p in range(1, nlags):
            sar[i] += np.dot(sar[i - p, :], -ar[p])

    return sar


def padone(x, front=0, back=0, axis=0, fillvalue=0):
    """pad with zeros along one axis, currently only axis=0


    can be used sequentially to pad several axis

    Examples
    --------
    >>> padone(np.ones((2,3)),1,3,axis=1)
    array([[ 0.,  1.,  1.,  1.,  0.,  0.,  0.],
           [ 0.,  1.,  1.,  1.,  0.,  0.,  0.]])

    >>> padone(np.ones((2,3)),1,1, fillvalue=np.nan)
    array([[ NaN,  NaN,  NaN],
           [  1.,   1.,   1.],
           [  1.,   1.,   1.],
           [ NaN,  NaN,  NaN]])
    """
    # primitive version
    shape = np.array(x.shape)
    shape[axis] += front + back
    shapearr = np.array(x.shape)
    out = np.empty(shape)
    out.fill(fillvalue)
    startind = np.zeros(x.ndim)
    startind[axis] = front
    endind = startind + shapearr
    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]
    # print(myslice
    # print(out.shape
    # print(out[tuple(myslice)].shape
    out[tuple(myslice)] = x
    return out


def trimone(x, front=0, back=0, axis=0):
    """trim number of array elements along one axis


    Examples
    --------
    >>> xp = padone(np.ones((2,3)),1,3,axis=1)
    >>> xp
    array([[ 0.,  1.,  1.,  1.,  0.,  0.,  0.],
           [ 0.,  1.,  1.,  1.,  0.,  0.,  0.]])
    >>> trimone(xp,1,3,1)
    array([[ 1.,  1.,  1.],
           [ 1.,  1.,  1.]])
    """
    shape = np.array(x.shape)
    shape[axis] -= front + back
    # print(shape, front, back
    startind = np.zeros(x.ndim)
    startind[axis] = front
    endind = startind + shape
    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]
    # print(myslice
    # print(shape, endind
    # print(x[tuple(myslice)].shape
    return x[tuple(myslice)]


def ar2full(ar):
    """make reduced lagpolynomial into a right side lagpoly array"""
    nlags, nvar, nvarex = ar.shape
    return np.r_[np.eye(nvar, nvarex)[None, :, :], -ar]


def ar2lhs(ar):
    """convert full (rhs) lagpolynomial into a reduced, left side lagpoly array

    this is mainly a reminder about the definition
    """
    return -ar[1:]


class _Var:
    """obsolete VAR class, use tsa.VAR instead, for internal use only


    Examples
    --------

    >>> v = Var(ar2s)
    >>> v.fit(1)
    >>> v.arhat
    array([[[ 1.        ,  0.        ],
            [ 0.        ,  1.        ]],

           [[-0.77784898,  0.01726193],
            [ 0.10733009, -0.78665335]]])

    """

    def __init__(self, y):
        self.y = y
        self.nobs, self.nvars = y.shape

    def fit(self, nlags):
        """estimate parameters using ols

        Parameters
        ----------
        nlags : int
            number of lags to include in regression, same for all variables

        Returns
        -------
        None, but attaches

        arhat : array (nlags, nvar, nvar)
            full lag polynomial array
        arlhs : array (nlags-1, nvar, nvar)
            reduced lag polynomial for left hand side
        other statistics as returned by linalg.lstsq : need to be completed



        This currently assumes all parameters are estimated without restrictions.
        In this case SUR is identical to OLS

        estimation results are attached to the class instance


        """
        self.nlags = nlags  # without current period
        nvars = self.nvars
        # TODO: ar2s looks like a module variable, bug?
        # lmat = lagmat(ar2s, nlags, trim='both', original='in')
        lmat = lagmat(self.y, nlags, trim="both", original="in")
        self.yred = lmat[:, :nvars]
        self.xred = lmat[:, nvars:]
        res = np.linalg.lstsq(self.xred, self.yred, rcond=-1)
        self.estresults = res
        self.arlhs = res[0].reshape(nlags, nvars, nvars)
        self.arhat = ar2full(self.arlhs)
        self.rss = res[1]
        self.xredrank = res[2]

    def predict(self):
        """calculate estimated timeseries (yhat) for sample"""

        if not hasattr(self, "yhat"):
            self.yhat = varfilter(self.y, self.arhat)
        return self.yhat

    def covmat(self):
        """covariance matrix of estimate
        # not sure it's correct, need to check orientation everywhere
        # looks ok, display needs getting used to
        >>> v.rss[None,None,:]*np.linalg.inv(np.dot(v.xred.T,v.xred))[:,:,None]
        array([[[ 0.37247445,  0.32210609],
                [ 0.1002642 ,  0.08670584]],

               [[ 0.1002642 ,  0.08670584],
                [ 0.45903637,  0.39696255]]])
        >>>
        >>> v.rss[0]*np.linalg.inv(np.dot(v.xred.T,v.xred))
        array([[ 0.37247445,  0.1002642 ],
               [ 0.1002642 ,  0.45903637]])
        >>> v.rss[1]*np.linalg.inv(np.dot(v.xred.T,v.xred))
        array([[ 0.32210609,  0.08670584],
               [ 0.08670584,  0.39696255]])
        """

        # check if orientation is same as self.arhat
        self.paramcov = (
            self.rss[None, None, :]
            * np.linalg.inv(np.dot(self.xred.T, self.xred))[:, :, None]
        )

    def forecast(self, horiz=1, u=None):
        """calculates forcast for horiz number of periods at end of sample

        Parameters
        ----------
        horiz : int (optional, default=1)
            forecast horizon
        u : array (horiz, nvars)
            error term for forecast periods. If None, then u is zero.

        Returns
        -------
        yforecast : array (nobs+horiz, nvars)
            this includes the sample and the forecasts
        """
        if u is None:
            u = np.zeros((horiz, self.nvars))
        return vargenerate(self.arhat, u, initvalues=self.y)


class VarmaPoly:
    """class to keep track of Varma polynomial format


    Examples
    --------

    ar23 = np.array([[[ 1. ,  0. ],
                     [ 0. ,  1. ]],

                    [[-0.6,  0. ],
                     [ 0.2, -0.6]],

                    [[-0.1,  0. ],
                     [ 0.1, -0.1]]])

    ma22 = np.array([[[ 1. ,  0. ],
                     [ 0. ,  1. ]],

                    [[ 0.4,  0. ],
                     [ 0.2, 0.3]]])


    """

    def __init__(self, ar, ma=None):
        self.ar = ar
        self.ma = ma
        nlags, nvarall, nvars = ar.shape
        self.nlags, self.nvarall, self.nvars = nlags, nvarall, nvars
        self.isstructured = not (ar[0, :nvars] == np.eye(nvars)).all()
        if self.ma is None:
            self.ma = np.eye(nvars)[None, ...]
            self.isindependent = True
        else:
            self.isindependent = not (ma[0] == np.eye(nvars)).all()
        self.malags = ar.shape[0]
        self.hasexog = nvarall > nvars
        self.arm1 = -ar[1:]

    # @property
    def vstack(self, a=None, name="ar"):
        """stack lagpolynomial vertically in 2d array"""
        if a is not None:
            _a = a
        elif name == "ar":
            _a = self.ar
        elif name == "ma":
            _a = self.ma
        else:
            raise ValueError("no array or name given")
        return _a.reshape(-1, self.nvarall)

    # @property
    def hstack(self, a=None, name="ar"):
        """stack lagpolynomial horizontally in 2d array"""
        if a is not None:
            _a = a
        elif name == "ar":
            _a = self.ar
        elif name == "ma":
            _a = self.ma
        else:
            raise ValueError("no array or name given")
        return _a.swapaxes(1, 2).reshape(-1, self.nvarall).T

    # @property
    def stacksquare(self, a=None, name="ar", orientation="vertical"):
        """stack lagpolynomial vertically in 2d square array with eye"""
        if a is not None:
            _a = a
        elif name == "ar":
            _a = self.ar
        elif name == "ma":
            _a = self.ma
        else:
            raise ValueError("no array or name given")
        astacked = _a.reshape(-1, self.nvarall)
        lenpk, nvars = astacked.shape  # [0]
        amat = np.eye(lenpk, k=nvars)
        amat[:, :nvars] = astacked
        return amat

    # @property
    def vstackarma_minus1(self):
        """stack ar and lagpolynomial vertically in 2d array"""
        a = np.concatenate((self.ar[1:], self.ma[1:]), 0)
        return a.reshape(-1, self.nvarall)

    # @property
    def hstackarma_minus1(self):
        """stack ar and lagpolynomial vertically in 2d array

        this is the Kalman Filter representation, I think
        """
        a = np.concatenate((self.ar[1:], self.ma[1:]), 0)
        return a.swapaxes(1, 2).reshape(-1, self.nvarall)

    def getisstationary(self, a=None):
        """check whether the auto-regressive lag-polynomial is stationary

        Returns
        -------
        isstationary : bool

        *attaches*

        areigenvalues : complex array
            eigenvalues sorted by absolute value

        References
        ----------
        formula taken from NAG manual

        """
        if a is not None:
            _a = a
        elif self.isstructured:
            _a = -self.reduceform(self.ar)[1:]
        else:
            _a = -self.ar[1:]
        amat = self.stacksquare(_a)
        ev = np.sort(np.linalg.eigvals(amat))[::-1]
        self.areigenvalues = ev
        return (np.abs(ev) < 1).all()

    def getisinvertible(self, a=None):
        """check whether the auto-regressive lag-polynomial is stationary

        Returns
        -------
        isinvertible : bool

        *attaches*

        maeigenvalues : complex array
            eigenvalues sorted by absolute value

        References
        ----------
        formula taken from NAG manual

        """
        if a is not None:
            _a = a
        elif self.isindependent:
            _a = self.reduceform(self.ma)[1:]
        else:
            _a = self.ma[1:]
        if _a.shape[0] == 0:
            # no ma lags
            self.maeigenvalues = np.array([], np.complex)
            return True

        amat = self.stacksquare(_a)
        ev = np.sort(np.linalg.eigvals(amat))[::-1]
        self.maeigenvalues = ev
        return (np.abs(ev) < 1).all()

    def reduceform(self, apoly):
        """

        this assumes no exog, todo

        """
        if apoly.ndim != 3:
            raise ValueError("apoly needs to be 3d")
        nlags, nvarsex, nvars = apoly.shape

        a = np.empty_like(apoly)
        try:
            a0inv = np.linalg.inv(a[0, :nvars, :])
        except np.linalg.LinAlgError as la_err:
            raise ValueError(
                "matrix not invertible, ask for implementation of pinv"
            ) from la_err

        for lag in range(nlags):
            a[lag] = np.dot(a0inv, apoly[lag])

        return a


if __name__ == "__main__":
    # some example lag polynomials
    a21 = np.array([[[1.0, 0.0], [0.0, 1.0]], [[-0.8, 0.0], [0.0, -0.6]]])

    a22 = np.array([[[1.0, 0.0], [0.0, 1.0]], [[-0.8, 0.0], [0.1, -0.8]]])

    a23 = np.array([[[1.0, 0.0], [0.0, 1.0]], [[-0.8, 0.2], [0.1, -0.6]]])

    a24 = np.array(
        [
            [[1.0, 0.0], [0.0, 1.0]],
            [[-0.6, 0.0], [0.2, -0.6]],
            [[-0.1, 0.0], [0.1, -0.1]],
        ]
    )

    a31 = np.r_[np.eye(3)[None, :, :], 0.8 * np.eye(3)[None, :, :]]
    a32 = np.array(
        [
            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
            [[0.8, 0.0, 0.0], [0.1, 0.6, 0.0], [0.0, 0.0, 0.9]],
        ]
    )

    ########
    ut = np.random.randn(1000, 2)
    ar2s = vargenerate(a22, ut)
    # res = np.linalg.lstsq(lagmat(ar2s,1)[:,1:], ar2s)
    res = np.linalg.lstsq(lagmat(ar2s, 1), ar2s, rcond=-1)
    bhat = res[0].reshape(1, 2, 2)
    arhat = ar2full(bhat)
    # print(maxabs(arhat - a22)

    v = _Var(ar2s)
    v.fit(1)
    v.forecast()
    v.forecast(25)[-30:]

    ar23 = np.array(
        [
            [[1.0, 0.0], [0.0, 1.0]],
            [[-0.6, 0.0], [0.2, -0.6]],
            [[-0.1, 0.0], [0.1, -0.1]],
        ]
    )

    ma22 = np.array([[[1.0, 0.0], [0.0, 1.0]], [[0.4, 0.0], [0.2, 0.3]]])

    ar23ns = np.array(
        [
            [[1.0, 0.0], [0.0, 1.0]],
            [[-1.9, 0.0], [0.4, -0.6]],
            [[0.3, 0.0], [0.1, -0.1]],
        ]
    )

    vp = VarmaPoly(ar23, ma22)
    print(vars(vp))
    print(vp.vstack())
    print(vp.vstack(a24))
    print(vp.hstackarma_minus1())
    print(vp.getisstationary())
    print(vp.getisinvertible())

    vp2 = VarmaPoly(ar23ns)
    print(vp2.getisstationary())
    print(vp2.getisinvertible())  # no ma lags
