"""VAR and VARMA process

this does not actually do much, trying out a version for a time loop

alternative representation:
* textbook, different blocks in matrices
* Kalman filter
* VAR, VARX and ARX could be calculated with signal.lfilter
  only tried some examples, not implemented

TODO: try minimizing sum of squares of (Y-Yhat)

Note: filter has smallest lag at end of array and largest lag at beginning,
    be careful for asymmetric lags coefficients
    check this again if it is consistently used


changes
2009-09-08 : separated from movstat.py

Author : josefpkt
License : BSD
"""

import numpy as np
from scipy import signal


# NOTE: this just returns that predicted values given the
# B matrix in polynomial form.
# TODO: make sure VAR class returns B/params in this form.
def VAR(x, B, const=0):
    """multivariate linear filter

    Parameters
    ----------
    x: (TxK) array
        columns are variables, rows are observations for time period
    B: (PxKxK) array
        b_t-1 is bottom "row", b_t-P is top "row" when printing
        B(:,:,0) is lag polynomial matrix for variable 1
        B(:,:,k) is lag polynomial matrix for variable k
        B(p,:,k) is pth lag for variable k
        B[p,:,:].T corresponds to A_p in Wikipedia
    const : float or array (not tested)
        constant added to autoregression

    Returns
    -------
    xhat: (TxK) array
        filtered, predicted values of x array

    Notes
    -----
    xhat(t,i) = sum{_p}sum{_k} { x(t-P:t,:) .* B(:,:,i) }  for all i = 0,K-1, for all t=p..T

    xhat does not include the forecasting observation, xhat(T+1),
    xhat is 1 row shorter than signal.correlate

    References
    ----------
    https://en.wikipedia.org/wiki/Vector_Autoregression
    https://en.wikipedia.org/wiki/General_matrix_notation_of_a_VAR(p)
    """
    p = B.shape[0]
    T = x.shape[0]
    xhat = np.zeros(x.shape)
    for t in range(p, T):  # [p+2]:#
        xhat[t, :] = const + (x[t - p : t, :, np.newaxis] * B).sum(axis=1).sum(axis=0)
    return xhat


def VARMA(x, B, C, const=0):
    """multivariate linear filter

    x (TxK)
    B (PxKxK)

    xhat(t,i) = sum{_p}sum{_k} { x(t-P:t,:) .* B(:,:,i) } +
                sum{_q}sum{_k} { e(t-Q:t,:) .* C(:,:,i) }for all i = 0,K-1

    """
    P = B.shape[0]
    Q = C.shape[0]
    T = x.shape[0]
    xhat = np.zeros(x.shape)
    e = np.zeros(x.shape)
    start = max(P, Q)
    for t in range(start, T):  # [p+2]:#
        xhat[t, :] = (
            const
            + (x[t - P : t, :, np.newaxis] * B).sum(axis=1).sum(axis=0)
            + (e[t - Q : t, :, np.newaxis] * C).sum(axis=1).sum(axis=0)
        )
        e[t, :] = x[t, :] - xhat[t, :]
    return xhat, e


if __name__ == "__main__":

    T = 20
    K = 2
    P = 3
    # x = np.arange(10).reshape(5,2)
    x = np.column_stack([np.arange(T)] * K)
    B = np.ones((P, K, K))
    # B[:,:,1] = 2
    B[:, :, 1] = [[0, 0], [0, 0], [0, 1]]
    xhat = VAR(x, B)
    print(np.all(xhat[P:, 0] == np.correlate(x[:-1, 0], np.ones(P)) * 2))
    # print(xhat)

    T = 20
    K = 2
    Q = 2
    P = 3
    const = 1
    # x = np.arange(10).reshape(5,2)
    x = np.column_stack([np.arange(T)] * K)
    B = np.ones((P, K, K))
    # B[:,:,1] = 2
    B[:, :, 1] = [[0, 0], [0, 0], [0, 1]]
    C = np.zeros((Q, K, K))
    xhat1 = VAR(x, B, const=const)
    xhat2, err2 = VARMA(x, B, C, const=const)
    print(np.all(xhat2 == xhat1))
    print(np.all(xhat2[P:, 0] == np.correlate(x[:-1, 0], np.ones(P)) * 2 + const))

    C[1, 1, 1] = 0.5
    xhat3, err3 = VARMA(x, B, C)

    x = np.r_[np.zeros((P, K)), x]  # prepend initial conditions
    xhat4, err4 = VARMA(x, B, C)

    C[1, 1, 1] = 1
    B[:, :, 1] = [[0, 0], [0, 0], [0, 1]]
    xhat5, err5 = VARMA(x, B, C)
    # print(err5)

    # in differences
    # VARMA(np.diff(x,axis=0),B,C)

    # Note:
    # * signal correlate applies same filter to all columns if kernel.shape[1]<K
    #   e.g. signal.correlate(x0,np.ones((3,1)),'valid')
    # * if kernel.shape[1]==K, then `valid` produces a single column
    #   -> possible to run signal.correlate K times with different filters,
    #      see the following example, which replicates VAR filter
    x0 = np.column_stack([np.arange(T), 2 * np.arange(T)])
    B[:, :, 0] = np.ones((P, K))
    B[:, :, 1] = np.ones((P, K))
    B[1, 1, 1] = 0
    xhat0 = VAR(x0, B)
    xcorr00 = signal.correlate(x0, B[:, :, 0])  # [:,0]
    xcorr01 = signal.correlate(x0, B[:, :, 1])
    print(np.all(signal.correlate(x0, B[:, :, 0], "valid")[:-1, 0] == xhat0[P:, 0]))
    print(np.all(signal.correlate(x0, B[:, :, 1], "valid")[:-1, 0] == xhat0[P:, 1]))

    # import error
    # from movstat import acovf, acf
    from statsmodels.tsa.stattools import acf, acovf

    aav = acovf(x[:, 0])
    print(aav[0] == np.var(x[:, 0]))
    aac = acf(x[:, 0])
