"""
Statespace Tools

Author: Chad Fulton
License: Simplified-BSD
"""
import numpy as np
from scipy.linalg import solve_sylvester
import pandas as pd

from statsmodels.compat.pandas import Appender
from statsmodels.tools.data import _is_using_pandas
from scipy.linalg.blas import find_best_blas_type
from . import (_initialization, _representation, _kalman_filter,
               _kalman_smoother, _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_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
}


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) or isinstance(polynomial, 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.
    diff : int, optional
        The number of simple differences to perform. Default is 1.
    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]]

    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 s in series[1:]:
            if isinstance(s, pd.DataFrame):
                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