"""
Statespace Tools

Author: Chad Fulton
License: Simplified-BSD
"""
from statsmodels.compat.pandas import Appender

import numpy as np
import pandas as pd
from scipy.linalg import solve_sylvester
from scipy.linalg.blas import find_best_blas_type

from statsmodels.tools.data import _is_using_pandas

from . import (
    _cfa_simulation_smoother,
    _initialization,
    _kalman_filter,
    _kalman_smoother,
    _representation,
    _simulation_smoother,
    _tools,
)

compatibility_mode = False
has_trmm = True
prefix_dtype_map = {
    "s": np.float32, "d": np.float64, "c": np.complex64, "z": np.complex128
}
prefix_initialization_map = {
    "s": _initialization.sInitialization,
    "d": _initialization.dInitialization,
    "c": _initialization.cInitialization,
    "z": _initialization.zInitialization
}
prefix_statespace_map = {
    "s": _representation.sStatespace, "d": _representation.dStatespace,
    "c": _representation.cStatespace, "z": _representation.zStatespace
}
prefix_kalman_filter_map = {
    "s": _kalman_filter.sKalmanFilter,
    "d": _kalman_filter.dKalmanFilter,
    "c": _kalman_filter.cKalmanFilter,
    "z": _kalman_filter.zKalmanFilter
}
prefix_kalman_smoother_map = {
    "s": _kalman_smoother.sKalmanSmoother,
    "d": _kalman_smoother.dKalmanSmoother,
    "c": _kalman_smoother.cKalmanSmoother,
    "z": _kalman_smoother.zKalmanSmoother
}
prefix_simulation_smoother_map = {
    "s": _simulation_smoother.sSimulationSmoother,
    "d": _simulation_smoother.dSimulationSmoother,
    "c": _simulation_smoother.cSimulationSmoother,
    "z": _simulation_smoother.zSimulationSmoother
}
prefix_cfa_simulation_smoother_map = {
    "s": _cfa_simulation_smoother.sCFASimulationSmoother,
    "d": _cfa_simulation_smoother.dCFASimulationSmoother,
    "c": _cfa_simulation_smoother.cCFASimulationSmoother,
    "z": _cfa_simulation_smoother.zCFASimulationSmoother
}
prefix_pacf_map = {
    "s": _tools._scompute_coefficients_from_multivariate_pacf,
    "d": _tools._dcompute_coefficients_from_multivariate_pacf,
    "c": _tools._ccompute_coefficients_from_multivariate_pacf,
    "z": _tools._zcompute_coefficients_from_multivariate_pacf
}
prefix_sv_map = {
    "s": _tools._sconstrain_sv_less_than_one,
    "d": _tools._dconstrain_sv_less_than_one,
    "c": _tools._cconstrain_sv_less_than_one,
    "z": _tools._zconstrain_sv_less_than_one
}
prefix_reorder_missing_matrix_map = {
    "s": _tools.sreorder_missing_matrix,
    "d": _tools.dreorder_missing_matrix,
    "c": _tools.creorder_missing_matrix,
    "z": _tools.zreorder_missing_matrix
}
prefix_reorder_missing_vector_map = {
    "s": _tools.sreorder_missing_vector,
    "d": _tools.dreorder_missing_vector,
    "c": _tools.creorder_missing_vector,
    "z": _tools.zreorder_missing_vector
}
prefix_copy_missing_matrix_map = {
    "s": _tools.scopy_missing_matrix,
    "d": _tools.dcopy_missing_matrix,
    "c": _tools.ccopy_missing_matrix,
    "z": _tools.zcopy_missing_matrix
}
prefix_copy_missing_vector_map = {
    "s": _tools.scopy_missing_vector,
    "d": _tools.dcopy_missing_vector,
    "c": _tools.ccopy_missing_vector,
    "z": _tools.zcopy_missing_vector
}
prefix_copy_index_matrix_map = {
    "s": _tools.scopy_index_matrix,
    "d": _tools.dcopy_index_matrix,
    "c": _tools.ccopy_index_matrix,
    "z": _tools.zcopy_index_matrix
}
prefix_copy_index_vector_map = {
    "s": _tools.scopy_index_vector,
    "d": _tools.dcopy_index_vector,
    "c": _tools.ccopy_index_vector,
    "z": _tools.zcopy_index_vector
}
prefix_compute_smoothed_state_weights_map = {
    "s": _tools._scompute_smoothed_state_weights,
    "d": _tools._dcompute_smoothed_state_weights,
    "c": _tools._ccompute_smoothed_state_weights,
    "z": _tools._zcompute_smoothed_state_weights
}


def set_mode(compatibility=None):
    if compatibility:
        raise NotImplementedError("Compatibility mode is only available in"
                                  " statsmodels <= 0.9")


def companion_matrix(polynomial):
    r"""
    Create a companion matrix

    Parameters
    ----------
    polynomial : array_like or list
        If an iterable, interpreted as the coefficients of the polynomial from
        which to form the companion matrix. Polynomial coefficients are in
        order of increasing degree, and may be either scalars (as in an AR(p)
        model) or coefficient matrices (as in a VAR(p) model). If an integer,
        it is interpreted as the size of a companion matrix of a scalar
        polynomial, where the polynomial coefficients are initialized to zeros.
        If a matrix polynomial is passed, :math:`C_0` may be set to the scalar
        value 1 to indicate an identity matrix (doing so will improve the speed
        of the companion matrix creation).

    Returns
    -------
    companion_matrix : ndarray

    Notes
    -----
    Given coefficients of a lag polynomial of the form:

    .. math::

        c(L) = c_0 + c_1 L + \dots + c_p L^p

    returns a matrix of the form

    .. math::
        \begin{bmatrix}
            \phi_1 & 1      & 0 & \cdots & 0 \\
            \phi_2 & 0      & 1 &        & 0 \\
            \vdots &        &   & \ddots & 0 \\
                   &        &   &        & 1 \\
            \phi_n & 0      & 0 & \cdots & 0 \\
        \end{bmatrix}

    where some or all of the :math:`\phi_i` may be non-zero (if `polynomial` is
    None, then all are equal to zero).

    If the coefficients provided are scalars :math:`(c_0, c_1, \dots, c_p)`,
    then the companion matrix is an :math:`n \times n` matrix formed with the
    elements in the first column defined as
    :math:`\phi_i = -\frac{c_i}{c_0}, i \in 1, \dots, p`.

    If the coefficients provided are matrices :math:`(C_0, C_1, \dots, C_p)`,
    each of shape :math:`(m, m)`, then the companion matrix is an
    :math:`nm \times nm` matrix formed with the elements in the first column
    defined as :math:`\phi_i = -C_0^{-1} C_i', i \in 1, \dots, p`.

    It is important to understand the expected signs of the coefficients. A
    typical AR(p) model is written as:

    .. math::
        y_t = a_1 y_{t-1} + \dots + a_p y_{t-p} + \varepsilon_t

    This can be rewritten as:

    .. math::
        (1 - a_1 L - \dots - a_p L^p )y_t = \varepsilon_t \\
        (1 + c_1 L + \dots + c_p L^p )y_t = \varepsilon_t \\
        c(L) y_t = \varepsilon_t

    The coefficients from this form are defined to be :math:`c_i = - a_i`, and
    it is the :math:`c_i` coefficients that this function expects to be
    provided.
    """
    identity_matrix = False
    if isinstance(polynomial, (int, np.integer)):
        # GH 5570, allow numpy integer types, but coerce to python int
        n = int(polynomial)
        m = 1
        polynomial = None
    else:
        n = len(polynomial) - 1

        if n < 1:
            raise ValueError("Companion matrix polynomials must include at"
                             " least two terms.")

        if isinstance(polynomial, (list, tuple)):
            try:
                # Note: cannot use polynomial[0] because of the special
                # behavior associated with matrix polynomials and the constant
                # 1, see below.
                m = len(polynomial[1])
            except TypeError:
                m = 1

            # Check if we just have a scalar polynomial
            if m == 1:
                polynomial = np.asanyarray(polynomial)
            # Check if 1 was passed as the first argument (indicating an
            # identity matrix)
            elif polynomial[0] == 1:
                polynomial[0] = np.eye(m)
                identity_matrix = True
        else:
            m = 1
            polynomial = np.asanyarray(polynomial)

    matrix = np.zeros((n * m, n * m), dtype=np.asanyarray(polynomial).dtype)
    idx = np.diag_indices((n - 1) * m)
    idx = (idx[0], idx[1] + m)
    matrix[idx] = 1
    if polynomial is not None and n > 0:
        if m == 1:
            matrix[:, 0] = -polynomial[1:] / polynomial[0]
        elif identity_matrix:
            for i in range(n):
                matrix[i * m:(i + 1) * m, :m] = -polynomial[i+1].T
        else:
            inv = np.linalg.inv(polynomial[0])
            for i in range(n):
                matrix[i * m:(i + 1) * m, :m] = -np.dot(inv, polynomial[i+1]).T
    return matrix


def diff(series, k_diff=1, k_seasonal_diff=None, seasonal_periods=1):
    r"""
    Difference a series simply and/or seasonally along the zero-th axis.

    Given a series (denoted :math:`y_t`), performs the differencing operation

    .. math::

        \Delta^d \Delta_s^D y_t

    where :math:`d =` `diff`, :math:`s =` `seasonal_periods`,
    :math:`D =` `seasonal\_diff`, and :math:`\Delta` is the difference
    operator.

    Parameters
    ----------
    series : array_like
        The series to be differenced.
    k_diff : int, optional
        The number of simple differences to perform. Default is 1.
    k_seasonal_diff : int or None, optional
        The number of seasonal differences to perform. Default is no seasonal
        differencing.
    seasonal_periods : int, optional
        The seasonal lag. Default is 1. Unused if there is no seasonal
        differencing.

    Returns
    -------
    differenced : ndarray
        The differenced array.
    """
    pandas = _is_using_pandas(series, None)
    differenced = np.asanyarray(series) if not pandas else series

    # Seasonal differencing
    if k_seasonal_diff is not None:
        while k_seasonal_diff > 0:
            if not pandas:
                differenced = (differenced[seasonal_periods:] -
                               differenced[:-seasonal_periods])
            else:
                sdiffed = differenced.diff(seasonal_periods)
                differenced = sdiffed[seasonal_periods:]
            k_seasonal_diff -= 1

    # Simple differencing
    if not pandas:
        differenced = np.diff(differenced, k_diff, axis=0)
    else:
        while k_diff > 0:
            differenced = differenced.diff()[1:]
            k_diff -= 1
    return differenced


def concat(series, axis=0, allow_mix=False):
    """
    Concatenate a set of series.

    Parameters
    ----------
    series : iterable
        An iterable of series to be concatenated
    axis : int, optional
        The axis along which to concatenate. Default is 1 (columns).
    allow_mix : bool
        Whether or not to allow a mix of pandas and non-pandas objects. Default
        is False. If true, the returned object is an ndarray, and additional
        pandas metadata (e.g. column names, indices, etc) is lost.

    Returns
    -------
    concatenated : array or pd.DataFrame
        The concatenated array. Will be a DataFrame if series are pandas
        objects.
    """
    is_pandas = np.r_[[_is_using_pandas(s, None) for s in series]]
    ndim = np.r_[[np.ndim(s) for s in series]]
    max_ndim = np.max(ndim)

    if max_ndim > 2:
        raise ValueError("`tools.concat` does not support arrays with 3 or"
                         " more dimensions.")

    # Make sure the iterable is mutable
    if isinstance(series, tuple):
        series = list(series)

    # Standardize ndim
    for i in range(len(series)):
        if ndim[i] == 0 and max_ndim == 1:
            series[i] = np.atleast_1d(series[i])
        elif ndim[i] == 0 and max_ndim == 2:
            series[i] = np.atleast_2d(series[i])
        elif ndim[i] == 1 and max_ndim == 2 and is_pandas[i]:
            name = series[i].name
            series[i] = series[i].to_frame()
            series[i].columns = [name]
        elif ndim[i] == 1 and max_ndim == 2 and not is_pandas[i]:
            series[i] = np.atleast_2d(series[i]).T

    if np.all(is_pandas):
        if isinstance(series[0], pd.DataFrame):
            base_columns = series[0].columns
        else:
            base_columns = pd.Index([series[0].name])
        for i in range(1, len(series)):
            s = series[i]

            if isinstance(s, pd.DataFrame):
                # Handle case where we were passed a dataframe and a series
                # to concatenate, and the series did not have a name.
                if s.columns.equals(pd.Index([None])):
                    s.columns = base_columns[:1]
                s_columns = s.columns
            else:
                s_columns = pd.Index([s.name])

            if axis == 0 and not base_columns.equals(s_columns):
                raise ValueError("Columns must match to concatenate along"
                                 " rows.")
            elif axis == 1 and not series[0].index.equals(s.index):
                raise ValueError("Index must match to concatenate along"
                                 " columns.")
        concatenated = pd.concat(series, axis=axis)
    elif np.all(~is_pandas) or allow_mix:
        concatenated = np.concatenate(series, axis=axis)
    else:
        raise ValueError("Attempted to concatenate Pandas objects with"
                         " non-Pandas objects with `allow_mix=False`.")

    return concatenated


def is_invertible(polynomial, threshold=1 - 1e-10):
    r"""
    Determine if a polynomial is invertible.

    Requires all roots of the polynomial lie inside the unit circle.

    Parameters
    ----------
    polynomial : array_like or tuple, list
        Coefficients of a polynomial, in order of increasing degree.
        For example, `polynomial=[1, -0.5]` corresponds to the polynomial
        :math:`1 - 0.5x` which has root :math:`2`. If it is a matrix
        polynomial (in which case the coefficients are coefficient matrices),
        a tuple or list of matrices should be passed.
    threshold : number
        Allowed threshold for `is_invertible` to return True. Default is 1.

    See Also
    --------
    companion_matrix

    Notes
    -----

    If the coefficients provided are scalars :math:`(c_0, c_1, \dots, c_n)`,
    then the corresponding polynomial is :math:`c_0 + c_1 L + \dots + c_n L^n`.


    If the coefficients provided are matrices :math:`(C_0, C_1, \dots, C_n)`,
    then the corresponding polynomial is :math:`C_0 + C_1 L + \dots + C_n L^n`.

    There are three equivalent methods of determining if the polynomial
    represented by the coefficients is invertible:

    The first method factorizes the polynomial into:

    .. math::

        C(L) & = c_0 + c_1 L + \dots + c_n L^n \\
             & = constant (1 - \lambda_1 L)
                 (1 - \lambda_2 L) \dots (1 - \lambda_n L)

    In order for :math:`C(L)` to be invertible, it must be that each factor
    :math:`(1 - \lambda_i L)` is invertible; the condition is then that
    :math:`|\lambda_i| < 1`, where :math:`\lambda_i` is a root of the
    polynomial.

    The second method factorizes the polynomial into:

    .. math::

        C(L) & = c_0 + c_1 L + \dots + c_n L^n \\
             & = constant (L - \zeta_1) (L - \zeta_2) \dots (L - \zeta_3)

    The condition is now :math:`|\zeta_i| > 1`, where :math:`\zeta_i` is a root
    of the polynomial with reversed coefficients and
    :math:`\lambda_i = \frac{1}{\zeta_i}`.

    Finally, a companion matrix can be formed using the coefficients of the
    polynomial. Then the eigenvalues of that matrix give the roots of the
    polynomial. This last method is the one actually used.
    """
    # First method:
    # np.all(np.abs(np.roots(np.r_[1, params])) < 1)
    # Second method:
    # np.all(np.abs(np.roots(np.r_[1, params][::-1])) > 1)
    # Final method:
    eigvals = np.linalg.eigvals(companion_matrix(polynomial))
    return np.all(np.abs(eigvals) < threshold)


def solve_discrete_lyapunov(a, q, complex_step=False):
    r"""
    Solves the discrete Lyapunov equation using a bilinear transformation.

    Notes
    -----
    This is a modification of the version in Scipy (see
    https://github.com/scipy/scipy/blob/master/scipy/linalg/_solvers.py)
    which allows passing through the complex numbers in the matrix a
    (usually the transition matrix) in order to allow complex step
    differentiation.
    """
    eye = np.eye(a.shape[0], dtype=a.dtype)
    if not complex_step:
        aH = a.conj().transpose()
        aHI_inv = np.linalg.inv(aH + eye)
        b = np.dot(aH - eye, aHI_inv)
        c = 2*np.dot(np.dot(np.linalg.inv(a + eye), q), aHI_inv)
        return solve_sylvester(b.conj().transpose(), b, -c)
    else:
        aH = a.transpose()
        aHI_inv = np.linalg.inv(aH + eye)
        b = np.dot(aH - eye, aHI_inv)
        c = 2*np.dot(np.dot(np.linalg.inv(a + eye), q), aHI_inv)
        return solve_sylvester(b.transpose(), b, -c)


def constrain_stationary_univariate(unconstrained):
    """
    Transform unconstrained parameters used by the optimizer to constrained
    parameters used in likelihood evaluation

    Parameters
    ----------
    unconstrained : ndarray
        Unconstrained parameters used by the optimizer, to be transformed to
        stationary coefficients of, e.g., an autoregressive or moving average
        component.

    Returns
    -------
    constrained : ndarray
        Constrained parameters of, e.g., an autoregressive or moving average
        component, to be transformed to arbitrary parameters used by the
        optimizer.

    References
    ----------
    .. [*] Monahan, John F. 1984.
       "A Note on Enforcing Stationarity in
       Autoregressive-moving Average Models."
       Biometrika 71 (2) (August 1): 403-404.
    """

    n = unconstrained.shape[0]
    y = np.zeros((n, n), dtype=unconstrained.dtype)
    r = unconstrained/((1 + unconstrained**2)**0.5)
    for k in range(n):
        for i in range(k):
            y[k, i] = y[k - 1, i] + r[k] * y[k - 1, k - i - 1]
        y[k, k] = r[k]
    return -y[n - 1, :]


def unconstrain_stationary_univariate(constrained):
    """
    Transform constrained parameters used in likelihood evaluation
    to unconstrained parameters used by the optimizer

    Parameters
    ----------
    constrained : ndarray
        Constrained parameters of, e.g., an autoregressive or moving average
        component, to be transformed to arbitrary parameters used by the
        optimizer.

    Returns
    -------
    unconstrained : ndarray
        Unconstrained parameters used by the optimizer, to be transformed to
        stationary coefficients of, e.g., an autoregressive or moving average
        component.

    References
    ----------
    .. [*] Monahan, John F. 1984.
       "A Note on Enforcing Stationarity in
       Autoregressive-moving Average Models."
       Biometrika 71 (2) (August 1): 403-404.
    """
    n = constrained.shape[0]
    y = np.zeros((n, n), dtype=constrained.dtype)
    y[n-1:] = -constrained
    for k in range(n-1, 0, -1):
        for i in range(k):
            y[k-1, i] = (y[k, i] - y[k, k]*y[k, k-i-1]) / (1 - y[k, k]**2)
    r = y.diagonal()
    x = r / ((1 - r**2)**0.5)
    return x


def _constrain_sv_less_than_one_python(unconstrained, order=None,
                                       k_endog=None):
    """
    Transform arbitrary matrices to matrices with singular values less than
    one.

    Parameters
    ----------
    unconstrained : list
        Arbitrary matrices. Should be a list of length `order`, where each
        element is an array sized `k_endog` x `k_endog`.
    order : int, optional
        The order of the autoregression.
    k_endog : int, optional
        The dimension of the data vector.

    Returns
    -------
    constrained : list
        Partial autocorrelation matrices. Should be a list of length
        `order`, where each element is an array sized `k_endog` x `k_endog`.

    See Also
    --------
    constrain_stationary_multivariate

    Notes
    -----
    Corresponds to Lemma 2.2 in Ansley and Kohn (1986). See
    `constrain_stationary_multivariate` for more details.
    """

    from scipy import linalg

    constrained = []  # P_s,  s = 1, ..., p
    if order is None:
        order = len(unconstrained)
    if k_endog is None:
        k_endog = unconstrained[0].shape[0]

    eye = np.eye(k_endog)
    for i in range(order):
        A = unconstrained[i]
        B, lower = linalg.cho_factor(eye + np.dot(A, A.T), lower=True)
        constrained.append(linalg.solve_triangular(B, A, lower=lower))
    return constrained


def _compute_coefficients_from_multivariate_pacf_python(
        partial_autocorrelations, error_variance, transform_variance=False,
        order=None, k_endog=None):
    """
    Transform matrices with singular values less than one to matrices
    corresponding to a stationary (or invertible) process.

    Parameters
    ----------
    partial_autocorrelations : list
        Partial autocorrelation matrices. Should be a list of length `order`,
        where each element is an array sized `k_endog` x `k_endog`.
    error_variance : ndarray
        The variance / covariance matrix of the error term. Should be sized
        `k_endog` x `k_endog`. This is used as input in the algorithm even if
        is not transformed by it (when `transform_variance` is False). The
        error term variance is required input when transformation is used
        either to force an autoregressive component to be stationary or to
        force a moving average component to be invertible.
    transform_variance : bool, optional
        Whether or not to transform the error variance term. This option is
        not typically used, and the default is False.
    order : int, optional
        The order of the autoregression.
    k_endog : int, optional
        The dimension of the data vector.

    Returns
    -------
    coefficient_matrices : list
        Transformed coefficient matrices leading to a stationary VAR
        representation.

    See Also
    --------
    constrain_stationary_multivariate

    Notes
    -----
    Corresponds to Lemma 2.1 in Ansley and Kohn (1986). See
    `constrain_stationary_multivariate` for more details.
    """
    from scipy import linalg

    if order is None:
        order = len(partial_autocorrelations)
    if k_endog is None:
        k_endog = partial_autocorrelations[0].shape[0]

    # If we want to keep the provided variance but with the constrained
    # coefficient matrices, we need to make a copy here, and then after the
    # main loop we will transform the coefficients to match the passed variance
    if not transform_variance:
        initial_variance = error_variance
        # Need to make the input variance large enough that the recursions
        # do not lead to zero-matrices due to roundoff error, which would case
        # exceptions from the Cholesky decompositions.
        # Note that this will still not always ensure positive definiteness,
        # and for k_endog, order large enough an exception may still be raised
        error_variance = np.eye(k_endog) * (order + k_endog)**10

    forward_variances = [error_variance]   # \Sigma_s
    backward_variances = [error_variance]  # \Sigma_s^*,  s = 0, ..., p
    autocovariances = [error_variance]     # \Gamma_s
    # \phi_{s,k}, s = 1, ..., p
    #             k = 1, ..., s+1
    forwards = []
    # \phi_{s,k}^*
    backwards = []

    error_variance_factor = linalg.cholesky(error_variance, lower=True)

    forward_factors = [error_variance_factor]
    backward_factors = [error_variance_factor]

    # We fill in the entries as follows:
    # [1,1]
    # [2,2], [2,1]
    # [3,3], [3,1], [3,2]
    # ...
    # [p,p], [p,1], ..., [p,p-1]
    # the last row, correctly ordered, is then used as the coefficients
    for s in range(order):  # s = 0, ..., p-1
        prev_forwards = forwards
        prev_backwards = backwards
        forwards = []
        backwards = []

        # Create the "last" (k = s+1) matrix
        # Note: this is for k = s+1. However, below we then have to fill
        # in for k = 1, ..., s in order.
        # P L*^{-1} = x
        # x L* = P
        # L*' x' = P'
        forwards.append(
            linalg.solve_triangular(
                backward_factors[s], partial_autocorrelations[s].T,
                lower=True, trans="T"))
        forwards[0] = np.dot(forward_factors[s], forwards[0].T)

        # P' L^{-1} = x
        # x L = P'
        # L' x' = P
        backwards.append(
            linalg.solve_triangular(
                forward_factors[s], partial_autocorrelations[s],
                lower=True, trans="T"))
        backwards[0] = np.dot(backward_factors[s], backwards[0].T)

        # Update the variance
        # Note: if s >= 1, this will be further updated in the for loop
        # below
        # Also, this calculation will be re-used in the forward variance
        tmp = np.dot(forwards[0], backward_variances[s])
        autocovariances.append(tmp.copy().T)

        # Create the remaining k = 1, ..., s matrices,
        # only has an effect if s >= 1
        for k in range(s):
            forwards.insert(k, prev_forwards[k] - np.dot(
                forwards[-1], prev_backwards[s-(k+1)]))

            backwards.insert(k, prev_backwards[k] - np.dot(
                backwards[-1], prev_forwards[s-(k+1)]))

            autocovariances[s+1] += np.dot(autocovariances[k+1],
                                           prev_forwards[s-(k+1)].T)

        # Create forward and backwards variances
        forward_variances.append(
            forward_variances[s] - np.dot(tmp, forwards[s].T)
        )
        backward_variances.append(
            backward_variances[s] -
            np.dot(
                np.dot(backwards[s], forward_variances[s]),
                backwards[s].T
            )
        )

        # Cholesky factors
        forward_factors.append(
            linalg.cholesky(forward_variances[s+1], lower=True)
        )
        backward_factors.append(
            linalg.cholesky(backward_variances[s+1], lower=True)
        )

    # If we do not want to use the transformed variance, we need to
    # adjust the constrained matrices, as presented in Lemma 2.3, see above
    variance = forward_variances[-1]
    if not transform_variance:
        # Here, we need to construct T such that:
        # variance = T * initial_variance * T'
        # To do that, consider the Cholesky of variance (L) and
        # input_variance (M) to get:
        # L L' = T M M' T' = (TM) (TM)'
        # => L = T M
        # => L M^{-1} = T
        initial_variance_factor = np.linalg.cholesky(initial_variance)
        transformed_variance_factor = np.linalg.cholesky(variance)
        transform = np.dot(initial_variance_factor,
                           np.linalg.inv(transformed_variance_factor))
        inv_transform = np.linalg.inv(transform)

        for i in range(order):
            forwards[i] = (
                np.dot(np.dot(transform, forwards[i]), inv_transform)
            )

    return forwards, variance


def constrain_stationary_multivariate_python(unconstrained, error_variance,
                                             transform_variance=False,
                                             prefix=None):
    r"""
    Transform unconstrained parameters used by the optimizer to constrained
    parameters used in likelihood evaluation for a vector autoregression.

    Parameters
    ----------
    unconstrained : array or list
        Arbitrary matrices to be transformed to stationary coefficient matrices
        of the VAR. If a list, should be a list of length `order`, where each
        element is an array sized `k_endog` x `k_endog`. If an array, should be
        the matrices horizontally concatenated and sized
        `k_endog` x `k_endog * order`.
    error_variance : ndarray
        The variance / covariance matrix of the error term. Should be sized
        `k_endog` x `k_endog`. This is used as input in the algorithm even if
        is not transformed by it (when `transform_variance` is False). The
        error term variance is required input when transformation is used
        either to force an autoregressive component to be stationary or to
        force a moving average component to be invertible.
    transform_variance : bool, optional
        Whether or not to transform the error variance term. This option is
        not typically used, and the default is False.
    prefix : {'s','d','c','z'}, optional
        The appropriate BLAS prefix to use for the passed datatypes. Only
        use if absolutely sure that the prefix is correct or an error will
        result.

    Returns
    -------
    constrained : array or list
        Transformed coefficient matrices leading to a stationary VAR
        representation. Will match the type of the passed `unconstrained`
        variable (so if a list was passed, a list will be returned).

    Notes
    -----
    In the notation of [1]_, the arguments `(variance, unconstrained)` are
    written as :math:`(\Sigma, A_1, \dots, A_p)`, where :math:`p` is the order
    of the vector autoregression, and is here determined by the length of
    the `unconstrained` argument.

    There are two steps in the constraining algorithm.

    First, :math:`(A_1, \dots, A_p)` are transformed into
    :math:`(P_1, \dots, P_p)` via Lemma 2.2 of [1]_.

    Second, :math:`(\Sigma, P_1, \dots, P_p)` are transformed into
    :math:`(\Sigma, \phi_1, \dots, \phi_p)` via Lemmas 2.1 and 2.3 of [1]_.

    If `transform_variance=True`, then only Lemma 2.1 is applied in the second
    step.

    While this function can be used even in the univariate case, it is much
    slower, so in that case `constrain_stationary_univariate` is preferred.

    References
    ----------
    .. [1] Ansley, Craig F., and Robert Kohn. 1986.
       "A Note on Reparameterizing a Vector Autoregressive Moving Average Model
       to Enforce Stationarity."
       Journal of Statistical Computation and Simulation 24 (2): 99-106.
    .. [*] Ansley, Craig F, and Paul Newbold. 1979.
       "Multivariate Partial Autocorrelations."
       In Proceedings of the Business and Economic Statistics Section, 349-53.
       American Statistical Association
    """

    use_list = type(unconstrained) is list
    if not use_list:
        k_endog, order = unconstrained.shape
        order //= k_endog

        unconstrained = [
            unconstrained[:k_endog, i*k_endog:(i+1)*k_endog]
            for i in range(order)
        ]

    order = len(unconstrained)
    k_endog = unconstrained[0].shape[0]

    # Step 1: convert from arbitrary matrices to those with singular values
    # less than one.
    sv_constrained = _constrain_sv_less_than_one_python(
        unconstrained, order, k_endog)

    # Step 2: convert matrices from our "partial autocorrelation matrix" space
    # (matrices with singular values less than one) to the space of stationary
    # coefficient matrices
    constrained, var = _compute_coefficients_from_multivariate_pacf_python(
        sv_constrained, error_variance, transform_variance, order, k_endog)

    if not use_list:
        constrained = np.concatenate(constrained, axis=1).reshape(
            k_endog, k_endog * order)

    return constrained, var


@Appender(constrain_stationary_multivariate_python.__doc__)
def constrain_stationary_multivariate(unconstrained, variance,
                                      transform_variance=False,
                                      prefix=None):

    use_list = type(unconstrained) is list
    if use_list:
        unconstrained = np.concatenate(unconstrained, axis=1)

    k_endog, order = unconstrained.shape
    order //= k_endog

    if order < 1:
        raise ValueError("Must have order at least 1")
    if k_endog < 1:
        raise ValueError("Must have at least 1 endogenous variable")

    if prefix is None:
        prefix, dtype, _ = find_best_blas_type(
            [unconstrained, variance])
    dtype = prefix_dtype_map[prefix]

    unconstrained = np.asfortranarray(unconstrained, dtype=dtype)
    variance = np.asfortranarray(variance, dtype=dtype)

    # Step 1: convert from arbitrary matrices to those with singular values
    # less than one.
    # sv_constrained = _constrain_sv_less_than_one(unconstrained, order,
    #                                              k_endog, prefix)
    sv_constrained = prefix_sv_map[prefix](unconstrained, order, k_endog)

    # Step 2: convert matrices from our "partial autocorrelation matrix"
    # space (matrices with singular values less than one) to the space of
    # stationary coefficient matrices
    constrained, variance = prefix_pacf_map[prefix](
        sv_constrained, variance, transform_variance, order, k_endog)

    constrained = np.array(constrained, dtype=dtype)
    variance = np.array(variance, dtype=dtype)

    if use_list:
        constrained = [
            constrained[:k_endog, i*k_endog:(i+1)*k_endog]
            for i in range(order)
        ]

