#cython: boundscheck=False
#cython: wraparound=False
#cython: cdivision=False
"""
State Space Models

Author: Chad Fulton  
License: Simplified-BSD
"""

{{py:

TYPES = {
    "s": ("np.float32_t", "np.float32", "np.NPY_FLOAT32"),
    "d": ("np.float64_t", "float", "np.NPY_FLOAT64"),
    "c": ("np.complex64_t", "np.complex64", "np.NPY_COMPLEX64"),
    "z": ("np.complex128_t", "complex", "np.NPY_COMPLEX128"),
}

}}

cimport numpy as np
import numpy as np
from statsmodels.src.math cimport *
cimport scipy.linalg.cython_blas as blas
cimport scipy.linalg.cython_lapack as lapack

from statsmodels.tsa.statespace._kalman_filter cimport (
    MEMORY_NO_SMOOTHING, MEMORY_NO_STD_FORECAST)

{{for prefix, types in TYPES.items()}}
{{py:cython_type, dtype, typenum = types}}
{{py:
combined_prefix = prefix
combined_cython_type = cython_type
if prefix == 'c':
    combined_prefix = 'z'
    combined_cython_type = 'np.complex128_t'
if prefix == 's':
    combined_prefix = 'd'
    combined_cython_type = 'np.float64_t'
}}

# ## Forecast error covariance inversion
#
# The following are routines that can calculate the inverse of the forecast
# error covariance matrix (defined in `forecast_<filter type>`).
#
# These routines are aware of the possibility that the Kalman filter may have
# converged to a steady state, in which case they do not need to perform the
# inversion or calculate the determinant.

cdef {{cython_type}} {{prefix}}inverse_univariate({{prefix}}KalmanFilter kfilter, {{prefix}}Statespace model, {{cython_type}} determinant) except *:
    """
    Factorize the forecast error covariance matrix using simple division
    in the case that the observations are univariate.

    If the model has converged to a steady-state, this is a NOOP and simply
    returns the determinant that was passed in.
    """
    # TODO it's not a noop if converged, but the docstring says it is...

    # #### Intermediate values
    cdef:
        int inc = 1
        {{cython_type}} scalar

    # Take the inverse of the forecast error covariance matrix
    if not kfilter.converged:
        determinant = {{combined_prefix}}log(kfilter._forecast_error_cov[0])
    try:
        # If we're close to singular, switch to univariate version even if the
        # following doesn't trigger a ZeroDivisionError
        if kfilter._forecast_error_cov[0]{{if combined_prefix == 'z'}}.real{{endif}} < 1e-12:
            raise Exception

        scalar = 1.0 / kfilter._forecast_error_cov[0]
    except:
        raise np.linalg.LinAlgError('Non-positive-definite forecast error'
                                    ' covariance matrix encountered at'
                                    ' period %d' % kfilter.t)
    kfilter._tmp2[0] = scalar * kfilter._forecast_error[0]
    blas.{{prefix}}copy(&model._k_endogstates, model._design, &inc, kfilter._tmp3, &inc)
    blas.{{prefix}}scal(&model._k_endogstates, &scalar, kfilter._tmp3, &inc)

    if not (kfilter.conserve_memory & MEMORY_NO_STD_FORECAST > 0):
        kfilter._standardized_forecast_error[0] = kfilter._forecast_error[0] * scalar**0.5

    if not (kfilter.conserve_memory & MEMORY_NO_SMOOTHING > 0):
        kfilter._tmp4[0] = scalar * model._obs_cov[0]

    return determinant

cdef {{cython_type}} {{prefix}}factorize_cholesky({{prefix}}KalmanFilter kfilter, {{prefix}}Statespace model, {{cython_type}} determinant) except *:
    """
    Factorize the forecast error covariance matrix using a Cholesky
    decomposition. Called by either of the `solve_cholesky` or
    `invert_cholesky` routines.

    Requires a positive definite matrix, but is faster than an LU
    decomposition.

    If the model has converged to a steady-state, this is a NOOP and simply
    returns the determinant that was passed in.
    """
    cdef:
        int inc = 1
        int info
        int i

    if not kfilter.converged or not (kfilter.conserve_memory & MEMORY_NO_STD_FORECAST > 0):
        blas.{{prefix}}copy(&kfilter.k_endog2, kfilter._forecast_error_cov, &inc, kfilter._forecast_error_fac, &inc)
        lapack.{{prefix}}potrf("U", &model._k_endog, kfilter._forecast_error_fac, &kfilter.k_endog, &info)

        if info < 0:
            raise np.linalg.LinAlgError('Illegal value in forecast error'
                                        ' covariance matrix encountered at'
                                        ' period %d' % kfilter.t)
        if info > 0:
            raise np.linalg.LinAlgError('Non-positive-definite forecast error'
                                       ' covariance matrix encountered at'
                                       ' period %d' % kfilter.t)

        # Calculate the determinant (just the squared product of the
        # diagonals, in the Cholesky decomposition case)
        determinant = 0.0
        for i in range(model._k_endog):
            determinant = determinant + {{combined_prefix}}log(kfilter.forecast_error_fac[i, i])
        determinant = 2 * determinant

    return determinant

cdef {{cython_type}} {{prefix}}factorize_lu({{prefix}}KalmanFilter kfilter, {{prefix}}Statespace model, {{cython_type}} determinant) except *:
    """
    Factorize the forecast error covariance matrix using an LU
    decomposition. Called by either of the `solve_lu` or `invert_lu`
    routines.

    Is slower than a Cholesky decomposition, but does not require a
    positive definite matrix.

    If the model has converged to a steady-state, this is a NOOP and simply
    returns the determinant that was passed in.
    """
    cdef:
        int inc = 1
        int info
        int i

    if not kfilter.converged:
        # Perform LU decomposition into `forecast_error_fac`
        blas.{{prefix}}copy(&kfilter.k_endog2, kfilter._forecast_error_cov, &inc, kfilter._forecast_error_fac, &inc)
        
        lapack.{{prefix}}getrf(&model._k_endog, &model._k_endog,
                        kfilter._forecast_error_fac, &kfilter.k_endog,
                        kfilter._forecast_error_ipiv, &info)

        if info < 0:
            raise np.linalg.LinAlgError('Illegal value in forecast error'
                                        ' covariance matrix encountered at'
                                        ' period %d' % kfilter.t)
        if info > 0:
            raise np.linalg.LinAlgError('Singular forecast error covariance'
                                        ' matrix encountered at period %d' %
                                        kfilter.t)

        # Calculate the determinant (product of the diagonals, but with
        # sign modifications according to the permutation matrix)    
        determinant = 1
        for i in range(model._k_endog):
            if not kfilter._forecast_error_ipiv[i] == i+1:
                determinant *= -1*kfilter.forecast_error_fac[i, i]
            else:
                determinant *= kfilter.forecast_error_fac[i, i]
        determinant = {{combined_prefix}}log(determinant)

    return determinant

cdef {{cython_type}} {{prefix}}inverse_cholesky({{prefix}}KalmanFilter kfilter, {{prefix}}Statespace model, {{cython_type}} determinant) except *:
    """
    inverse_cholesky(self, determinant)

    If the model has converged to a steady-state, this is a NOOP and simply
    returns the determinant that was passed in.
    """
    cdef:
        int info
        int inc = 1
        int i, j
        {{cython_type}} alpha = 1.0
        {{cython_type}} beta = 0.0

    if not kfilter.converged or not (kfilter.conserve_memory & MEMORY_NO_STD_FORECAST > 0):
        # Perform the Cholesky decomposition and get the determinant
        determinant = {{prefix}}factorize_cholesky(kfilter, model, determinant)

        # Use the Cholesky factor to get standardized forecast errors, if desired
        # (in the notation of DK, kfilter._forecast_error_fac holds L_t', and
        # we want to solve the equation L_t' v_t^s = v_t)
        if not (kfilter.conserve_memory & MEMORY_NO_STD_FORECAST > 0):
            blas.{{prefix}}copy(&kfilter.k_endog, kfilter._forecast_error, &inc, kfilter._standardized_forecast_error, &inc)
            lapack.{{prefix}}trtrs("U", "T", "N", &model._k_endog, &inc,
                                   kfilter._forecast_error_fac, &kfilter.k_endog,
                                   kfilter._standardized_forecast_error, &kfilter.k_endog, &info)

            if info != 0:
                raise np.linalg.LinAlgError('Error computing standardized'
                                            ' forecast error at period %d'
                                            % kfilter.t)

        # Continue taking the inverse
        lapack.{{prefix}}potri("U", &model._k_endog, kfilter._forecast_error_fac, &kfilter.k_endog, &info)

        # ?potri only fills in the upper triangle of the symmetric array, and
        # since the ?symm and ?symv routines are not available as of scipy
        # 0.11.0, we cannot use them, so we must fill in the lower triangle
        # by hand
        for i in range(model._k_endog): # columns
            for j in range(i): # rows
                kfilter.forecast_error_fac[i,j] = kfilter.forecast_error_fac[j,i]


    # Get `tmp2` and `tmp3` via matrix multiplications

    # `tmp2` array used here, dimension $(p \times 1)$  
    # $\\#_2 = F_t^{-1} v_t$  
    #blas.{{prefix}}symv("U", &model._k_endog, &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
    #               kfilter._forecast_error, &inc, &beta, kfilter._tmp2, &inc)
    blas.{{prefix}}gemv("N", &model._k_endog, &model._k_endog,
                   &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
                           kfilter._forecast_error, &inc,
                   &beta, kfilter._tmp2, &inc)

    # `tmp3` array used here, dimension $(p \times m)$  
    # $\\#_3 = F_t^{-1} Z_t$
    #blas.{{prefix}}symm("L", "U", &kfilter.k_endog, &kfilter.k_states,
    #               &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
    #                       kfilter._design, &kfilter.k_endog,
    #               &beta, kfilter._tmp3, &kfilter.k_endog)
    blas.{{prefix}}gemm("N", "N", &model._k_endog, &model._k_states, &model._k_endog,
                   &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
                           model._design, &model._k_endog,
                   &beta, kfilter._tmp3, &kfilter.k_endog)

    if not (kfilter.conserve_memory & MEMORY_NO_SMOOTHING > 0):
        # `tmp4` array used here, dimension $(p \times p)$  
        # $\\#_4 = F_t^{-1} H_t$
        #blas.{{prefix}}symm("L", "U", &kfilter.k_endog, &kfilter.k_endog,
        #               &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
        #                       kfilter._obs_cov, &kfilter.k_endog,
        #               &beta, kfilter._tmp4, &kfilter.k_endog)
        blas.{{prefix}}gemm("N", "N", &model._k_endog, &model._k_endog, &model._k_endog,
                       &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
                               model._obs_cov, &model._k_endog,
                       &beta, kfilter._tmp4, &kfilter.k_endog)

    return determinant

cdef {{cython_type}} {{prefix}}inverse_lu({{prefix}}KalmanFilter kfilter, {{prefix}}Statespace model, {{cython_type}} determinant) except *:
    """
    inverse_cholesky(self, determinant)

    If the model has converged to a steady-state, this is a NOOP and simply
    returns the determinant that was passed in.
    """
    cdef:
        int info
        int inc = 1
        {{cython_type}} alpha = 1.0
        {{cython_type}} beta = 0.0

    if not kfilter.converged:
        # Perform the Cholesky decomposition and get the determinant
        determinant = {{prefix}}factorize_lu(kfilter, model, determinant)

        # Continue taking the inverse
        lapack.{{prefix}}getri(&model._k_endog, kfilter._forecast_error_fac, &kfilter.k_endog,
               kfilter._forecast_error_ipiv, kfilter._forecast_error_work, &kfilter.ldwork, &info)

    # Get `tmp2` and `tmp3` via matrix multiplications

    # `tmp2` array used here, dimension $(p \times 1)$  
    # $\\#_2 = F_t^{-1} v_t$  
    blas.{{prefix}}gemv("N", &model._k_endog, &model._k_endog,
                   &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
                           kfilter._forecast_error, &inc,
                   &beta, kfilter._tmp2, &inc)

    # `tmp3` array used here, dimension $(p \times m)$  
    # $\\#_3 = F_t^{-1} Z_t$
    blas.{{prefix}}gemm("N", "N", &model._k_endog, &model._k_states, &model._k_endog,
                   &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
                           model._design, &model._k_endog,
                   &beta, kfilter._tmp3, &kfilter.k_endog)

    if not (kfilter.conserve_memory & MEMORY_NO_SMOOTHING > 0):
        # `tmp4` array used here, dimension $(p \times p)$  
        # $\\#_4 = F_t^{-1} H_t$
        blas.{{prefix}}gemm("N", "N", &model._k_endog, &model._k_endog, &model._k_endog,
                       &alpha, kfilter._forecast_error_fac, &kfilter.k_endog,
                               model._obs_cov, &model._k_endog,
                       &beta, kfilter._tmp4, &kfilter.k_endog)

    return determinant

cdef {{cython_type}} {{prefix}}solve_cholesky({{prefix}}KalmanFilter kfilter, {{prefix}}Statespace model, {{cython_type}} determinant) except *:
    """
    solve_cholesky(self, determinant)

    If the model has converged to a steady-state, this is a NOOP and simply
    returns the determinant that was passed in.
    """
    cdef:
        int info, i, j
        int inc = 1

    if not kfilter.converged:
        # Perform the Cholesky decomposition and get the determinant
        determinant = {{prefix}}factorize_cholesky(kfilter, model, determinant)

    # Use the Cholesky factor to get standardized forecast errors, if desired
    # (in the notation of DK, kfilter._forecast_error_fac holds L_t', and
    # we want to solve the equation L_t' v_t^s = v_t)
    if not (kfilter.conserve_memory & MEMORY_NO_STD_FORECAST > 0):
        blas.{{prefix}}copy(&kfilter.k_endog, kfilter._forecast_error, &inc, kfilter._standardized_forecast_error, &inc)
        lapack.{{prefix}}trtrs("U", "T", "N", &model._k_endog, &inc,
                               kfilter._forecast_error_fac, &kfilter.k_endog,
                               kfilter._standardized_forecast_error, &kfilter.k_endog, &info)

        if info != 0:
                raise np.linalg.LinAlgError('Error computing standardized'
                                            ' forecast error at period %d'
                                            % kfilter.t)

    # Solve the linear systems  
    # `tmp2` array used here, dimension $(p \times 1)$  
    # $F_t \\#_2 = v_t$  
    blas.{{prefix}}copy(&kfilter.k_endog, kfilter._forecast_error, &inc, kfilter._tmp2, &inc)
    lapack.{{prefix}}potrs("U", &model._k_endog, &inc, kfilter._forecast_error_fac, &kfilter.k_endog, kfilter._tmp2, &kfilter.k_endog, &info)

    # `tmp3` array used here, dimension $(p \times m)$  
    # $F_t \\#_3 = Z_t$
    if model._k_states == model.k_states and model._k_endog == model.k_endog:
        blas.{{prefix}}copy(&kfilter.k_endogstates, model._design, &inc, kfilter._tmp3, &inc)
    else:
        for i in range(model._k_states): # columns
            for j in range(model._k_endog): # rows
                kfilter._tmp3[j + i*kfilter.k_endog] = model._design[j + i*model._k_endog]
    lapack.{{prefix}}potrs("U", &model._k_endog, &model._k_states, kfilter._forecast_error_fac, &kfilter.k_endog, kfilter._tmp3, &kfilter.k_endog, &info)

    if not (kfilter.conserve_memory & MEMORY_NO_SMOOTHING > 0):
        # `tmp4` array used here, dimension $(p \times p)$  
        # $F_t \\#_4 = H_t$
        if model._k_states == model.k_states and model._k_endog == model.k_endog:
            blas.{{prefix}}copy(&kfilter.k_endog2, model._obs_cov, &inc, kfilter._tmp4, &inc)
        else:
            for i in range(model._k_endog): # columns
                for j in range(model._k_endog): # rows
                    kfilter._tmp4[j + i*kfilter.k_endog] = model._obs_cov[j + i*model._k_endog]
        lapack.{{prefix}}potrs("U", &model._k_endog, &model._k_endog, kfilter._forecast_error_fac, &kfilter.k_endog, kfilter._tmp4, &kfilter.k_endog, &info)

    return determinant

cdef {{cython_type}} {{prefix}}solve_lu({{prefix}}KalmanFilter kfilter, {{prefix}}Statespace model, {{cython_type}} determinant) except *:
    """
    inverse_cholesky(self, determinant)

    If the model has converged to a steady-state, this is a NOOP and simply
    returns the determinant that was passed in.
    """
    cdef:
        int info
        int inc = 1
        {{cython_type}} alpha = 1.0
        {{cython_type}} beta = 0.0

    if not kfilter.converged:
        # Perform the Cholesky decomposition and get the determinant
        determinant = {{prefix}}factorize_lu(kfilter, model, determinant)

    # Solve the linear systems  
    # `tmp2` array used here, dimension $(p \times 1)$  
    # $F_t \\#_2 = v_t$  
    blas.{{prefix}}copy(&kfilter.k_endog, kfilter._forecast_error, &inc, kfilter._tmp2, &inc)
    lapack.{{prefix}}getrs("N", &model._k_endog, &inc, kfilter._forecast_error_fac, &kfilter.k_endog,
                    kfilter._forecast_error_ipiv, kfilter._tmp2, &kfilter.k_endog, &info)

    # `tmp3` array used here, dimension $(p \times m)$  
    # $F_t \\#_3 = Z_t$
    if model._k_states == model.k_states and model._k_endog == model.k_endog:
        blas.{{prefix}}copy(&kfilter.k_endogstates, model._design, &inc, kfilter._tmp3, &inc)
    else:
        for i in range(model._k_states): # columns
            for j in range(model._k_endog): # rows
                kfilter._tmp3[j + i*kfilter.k_endog] = model._design[j + i*model._k_endog]
    lapack.{{prefix}}getrs("N", &model._k_endog, &model._k_states, kfilter._forecast_error_fac, &kfilter.k_endog,
                    kfilter._forecast_error_ipiv, kfilter._tmp3, &kfilter.k_endog, &info)

    if not (kfilter.conserve_memory & MEMORY_NO_SMOOTHING > 0):
        # `tmp4` array used here, dimension $(p \times p)$  
        # $F_t \\#_4 = H_t$
        if model._k_states == model.k_states and model._k_endog == model.k_endog:
            blas.{{prefix}}copy(&kfilter.k_endog2, model._obs_cov, &inc, kfilter._tmp4, &inc)
        else:
            for i in range(model._k_endog): # columns
                for j in range(model._k_endog): # rows
                    kfilter._tmp4[j + i*kfilter.k_endog] = model._obs_cov[j + i*model._k_endog]
        lapack.{{prefix}}getrs("N", &model._k_endog, &model._k_endog, kfilter._forecast_error_fac, &kfilter.k_endog,
                        kfilter._forecast_error_ipiv, kfilter._tmp4, &kfilter.k_endog, &info)

    return determinant

{{endfor}}
