# Estimating DSGE by Maximum Likelihood in Python
# Author: Michal Miktus
# Date: 20.03.2019

# Import libraries

from LinApp_Solve import LinApp_Solve
import matplotlib.pyplot as plt
import seaborn as sn
import pandas as pd
import numpy as np
import pickle
from numpy import vstack, hstack, zeros, eye
from Linear_Time_Iteration import Linear_Time_Iteration

# %matplotlib inline

# Set printing options

np.set_printoptions(precision=3, suppress=True, linewidth=120)

pd.set_option('float_format', lambda x: '%.3g' % x, )

# --------------------------------------------------------------------------
# -- Define the model
# --------------------------------------------------------------------------

#  VARIABLES   [M,3] cell array: Variable name (case sensitive) ~ Variable type ~ Description
#              Variable type: 'X' ... Endogenous state variable
#                             'Y' ... Endogenous other (jump) variable
#                             'Z' ... Exogenous state variable
#                             'U' ... Innovation to exogenous state variable
#                             ''  ... Skip variable

variable_symbols = [r'$e_a$', r'$e_I$', r'$e_b$', r'$e_L$',
                    r'$e_G$', r'$e_pi$', r'$pi$', r'$w$', r'$K$',
                    r'$Q$', r'$I$', r'$C$', r'$R$', r'$r$', r'$L$',
                    r'$Y$', r'$pi_f$', r'$w_f$', r'$K_f$', r'$Q_f$',
                    r'$I_f$', r'$C_f$', r'$R_f$', r'$r_f$', r'$L_f$',
                    r'$Y_f$', r'$ee_a$', r'$ee_I$', r'$ee_b$', r'$ee_L$', r'$ee_G$', r'$ee_pi$', r'$n_p$', r'$n_Q$', r'$n_R$', r'$n_w$', r'$one$']

variable_types = ['Z',    'Z',    'Z',    'Z',    'Z',    'Z',
                  'X',    'X',    'X',    'X',    'X',    'X',
                  'X',    'Y',    'Y',    'X',    'X',    'X',
                  'X',    'X',    'X',    'X',    'X',    'Y',
                  'Y',    'X',    'U',    'U',    'U',    'U',
                  'U',    'U',    'Z',    'Z',    'Z',    'Z',
                  'X', ]

variable_names = ['Aggregate productivity shock', 'Adjustment cost shock',
                  'Preference shock', 'Labor supply shock',
                  'Public spending shock', 'Inflation target', 'Inflation rate', 'Real wages', 'Capital', 'Q-Tobin ratio', 'Investment',
                  'Consumption', 'Monetary policy interest rate', 'Capital rental rate', 'Labor', 'Output',
                  'Flexible inflation rate', 'Flexible wages', 'Flexible capital', 'Flexible Q-Tobin ratio',
                  'Flexible investment', 'Flexible consumption', 'Flexible monetary policy interest rate',
                  'Flexible capital rental rate', 'Flexible labor', 'Flexible output', 'Aggregate productivity shock error',
                  'Adjustment cost shock error', 'Preference shock error', 'Labor supply shock error', 'Public spending shock error',
                  'Inflation target error', 'Prices error', 'Tobin Q-ratio error', 'Monetary policy interest rate error', 'Wages error',
                  'Temporary variable']


variables = pd.DataFrame({
        'names': variable_names,
        'types': variable_types,
        'symbols': variable_symbols
    })

var_endo_states = variables.loc[variables['types'] == 'X']
var_endo_controls = variables.loc[variables['types'] == 'Y']
var_exo  = variables.loc[variables['types'] == 'Z']

#    EQUATIONS   [N,3] cell array: Equation type ~ Equation name ~ equation
#                Equation type: 'D' ... Deterministic equation
#                               'E' ... Expectational equation
#                               'S' ... Shock equation
#                               ''  ... Skip equation

equation_formulas = [
    '0 = - C(t) + (h/(1 + h))*C(t-1) + (1/(1 + h))*C(t+1) - ((1 - h)/((1 + h)*sigma_c))*(R(t) - pi(t+1)) + ((1 - h)/((1 + h)*sigma_c))*(e_b(t)-e_b(t+1))',
    '0 = - I(t) + (1/(1 + betta))*I(t-1) + (betta/(1 + betta))*I(t+1) + (adj/(1 + betta))*Q(t) + (betta*e_I(t+1) - e_I(t))/(1 + betta)',
    '0 = - Q(t)  -R(t) + pi(t+1) + ((1 - tau)/(1 - tau + r_k))*Q(t+1) + (r_k/(1 - tau + r_k))*r(t+1) + n_Q(t)',
    '0 = - K(t) + (1 - tau)*K(t-1) + tau*I(t-1)', '0 = - pi(t) + (betta/(1 + betta*gamma_p))*pi(t+1) + (gamma_p/(1 + betta*gamma_p))*pi(t-1) + (((1 - betta*xi_p)*(1 - xi_p))/((1 + betta*gamma_p)*xi_p))*(alphaa*r(t) + (1 - alphaa)*w(t) - e_a(t) + n_p(t))',
    '0 = (-1 - ((1/(1 + betta))*((1 - betta*xi_w)*(1 - xi_w))/((1 + (1/lambda_w)*((1 + lambda_w)*sigma_L))*xi_w)))*w(t) + (betta/(1 + betta))*w(t+1) + (1/(1 + betta))*w(t-1)+ (betta/(1 + betta))*pi(t+1) - ((1 + betta*gamma_w)/(1 + betta))*pi(t) + (gamma_w/( 1 + betta))*pi(t-1) - ((1/(1 + betta))*((1 - betta*xi_w)*(1 - xi_w))/((1 + (1/lambda_w)*((1 + lambda_w)*sigma_L))*xi_w))*(- sigma_L*L(t) - (sigma_c/(1 - h))*(C(t) - h*C(t-1)) - e_L(t) - n_w(t))', '0 = - L(t) + -w(t) + (1 + psi)*r(t) + K(t-1)',
    '0 = - Y(t) + phi*e_a(t) + phi*alphaa*K(t-1) + phi*alphaa*psi*r(t) + phi*(1 - alphaa)*L(t)',
    '0 = - Y(t) + (1 - tau*k - g)*C(t) + tau*k*I(t) + e_G(t)',
    '0 = - R(t) + rho*R(t-1) + ((1 - rho)*r_pi - r_dpi)*pi(t-1) + (1 - rho)*(1 - r_pi)*e_pi(t) + Y(t)*((1-rho)*r_Y+r_dY) - Y_f(t)*((1-rho)*r_Y+r_dY) + r_dpi*pi(t) - r_dY*Y(t-1) - r_dY*Y_f(t-1) + n_R(t)',
    '0 = - C_f(t) + (h/(1 + h))*C_f(t-1) + (1/(1 + h))*C_f(t+1) - ((1 - h)/((1 + h)*sigma_c))*(R_f(t) - pi_f(t+1)) + ((1 - h)/((1 + h)*sigma_c))*(e_b(t)-e_b(t+1))',
    '0 = - I_f(t) + (1/(1 + betta))*I_f(t-1) + (betta/(1 + betta))*I_f(t+1) + (adj/(1 + betta))*Q_f(t) + (betta*e_I(t+1) - e_I(t))/(1 + betta)',
    '0 = - Q_f(t) - R_f(t) + pi_f(t+1) + ((1 - tau)/(1 - tau + r_k))*Q_f(t+1) + (r_k/(1 - tau + r_k))*r_f(t+1)',
    '0 = - K_f(t) + (1 - tau)*K_f(t-1) + tau*I_f(t-1)',
    '0 = alphaa*r_f(t)+(1-alphaa)*w_f(t) - e_a(t)',
    '0 = - w_f(t) + sigma_L*L_f(t) + (sigma_c/(1 - h))*(C_f(t) - h*C_f(t-1)) - e_L(t)',
    '0 = - L_f(t) + -w_f(t) + (1 + psi)*r_f(t) + K_f(t-1)',
    '0 = - Y_f(t) + phi*e_a(t) + phi*alphaa*K_f(t-1) + phi*alphaa*psi*r_f(t) + phi*(1 - alphaa)*L_f(t)',
    '0 = - Y_f(t) + (1 - tau*k - g)*C_f(t) + tau*k*I_f(t) + e_G(t)',
    'e_I(t+1) = rho_I*e_I(t) + ee_I(t+1)',
    'e_b(t+1) = rho_b*e_b(t) + ee_b(t+1)',
    'e_L(t+1) = rho_L*e_L(t) + ee_L(t+1)',
    'e_G(t+1) = rho_G*e_G(t) + ee_G(t+1)',
    'e_a(t+1) = rho_a*e_a(t) + ee_a(t+1)',
    'e_pi(t+1) = rho_pi*e_pi(t)+ ee_pi(t+1)',
    'pi_f(t) = 0*one(t)',
    'one(t) = 0*one(t-1)',
    'n_p(t) = 0',
    'n_Q(t) = 0',
    'n_R(t) = 0',
    'n_w(t) = 0'
    ]

equation_type = ['E',    'E',    'E',    'E',    'E',    'E',
                 'D',    'D',    'E',    'E',    'E',    'E',
                 'E',    'E',    'E',    'E',    'D',    'D',
                 'E',    'S',    'S',    'S',    'S',    'S',
                 'S',    'E',    'E',    'S',    'S',    'S',
                 'S']

equation_names = ['Consumption', 'Investment', 'Q-equation', 'Capital', 'Prices', 'Wages', 'Labor demand', 'Production',
                  'Goods market', 'Monetary policy', 'Flexible Consumption', 'Flexible Investment', 'Flexible Q-equation',
                  'Flexible Capital', 'Flexible Prices', 'Flexible Wages', 'Flexible Labor demand', 'Flexible Production',
                  'Flexible Goods market', 'Aggregate productivity shock', 'Adjustment cost shock',
                  'Preference shock', 'Labor supply shock', 'Public spending shock', 'Inflation target', 'Temporary first equation', 'Temporary second equation', 'Prices error expectation', 'Tobin Q-ratio error expectation',
                  'Monetary policy interest rate error expectation', 'Wages error expectation']

equations = pd.DataFrame({
        'names': equation_names,
        'types': equation_type,
        'formulas': equation_formulas
    })

# --------------------------------------------------------------------------
# --- Setting parameters
# --------------------------------------------------------------------------

betta = 0.99
tau = 0.025
alphaa = 0.3
psi = 1/0.169
gamma_p = 0.469
gamma_w = 0.763
lambda_w = 0.5
xi_p = 0.908
xi_w = 0.737
sigma_L = 2.4
sigma_c = 1.353
h = 0.573
phi = 1.408
adj = 1/6.771
r_k = (1/betta)-1+tau
k = 8.8
g = 0.18
r_dpi = 0.14
r_Y = 0.099
r_dY = 0.159
rho = 0.961
r_pi = 1.684

rho_L = 0.889
rho_a = 0.823
rho_b = 0.855
rho_G = 0.949
rho_pi = 0.924
rho_I = 0.927

sd_R = 0.081
sd_p = 0.16
sd_w = 0.289
sd_Q = 0.604

# --------------------------------------------------------------------------
# --- Matrix form in spirit of Uhlig
# --------------------------------------------------------------------------

# Structure of the model in Uhlig's notation:
#   0 = AA x(t) + BB x(t-1) + CC y(t) + DD z(t)
#   0 = E_t [ FF x(t+1) + GG x(t) + HH x(t-1) + JJ y(t+1) + KK y(t) + LL z(t+1) + MM z(t)]
#   z(t+1) = NN z(t) + epsilon(t+1),    with E_t [ epsilon(t+1) ] = 0
#
#   x(t)      : Endogenous state variables
#   y(t)      : Endogenous other variables
#   z(t)      : Exogenous state variables
#   epsilon(t): Innovation to exogenous state variable

AA = np.empty((4, 17), dtype=object)
BB = np.empty((4, 17), dtype=object)
CC = np.empty((4, 4), dtype=object)
DD = np.empty((4, 10), dtype=object)
FF = np.empty((17, 17), dtype=object)
GG = np.empty((17, 17), dtype=object)
HH = np.empty((17, 17), dtype=object)
JJ = np.empty((17, 4), dtype=object)
KK = np.empty((17, 4), dtype=object)
LL = np.empty((17, 10), dtype=object)
MM = np.empty((17, 10), dtype=object)
NN = np.empty((10, 10), dtype=object)

matrix_names = [AA, BB, CC, DD, FF, GG, HH, JJ, KK, LL, MM, NN]

nx = AA.shape[1]
ny = AA.shape[0]
nz = DD.shape[1]

# --- 1: Consumption

GG[0, 5] = '-1'                               # C(t)
HH[0, 5] = '(h/(1+h))'                        # C(t-1)
FF[0, 5] = '(1/(1+h))'                        # C(t+1)
GG[0, 6] = '-((1-h)/((1+h)*sigma_c))'         # R(t)
FF[0, 0] = '((1-h)/((1+h)*sigma_c))'          # pi(t+1)
MM[0, 2] = '((1-h)/((1+h)*sigma_c))'          # e_b(t)
LL[0, 2] = '-((1-h)/((1+h)*sigma_c)) '        # e_b(t+1)


# --- 2:  Investment

GG[1, 4] = '-1'                               # I(t)
HH[1, 4] = '(1/(1+betta))'                    # I(t-1)
FF[1, 4] = '(betta/(1+betta))'                # I(t+1)
GG[1, 3] = '(adj/(1+betta))'                  # Q(t)
LL[1, 1] = '(1+betta)*betta'                  # e_I(t+1)
MM[1, 1] = '-(1+betta)'                       # e_I(t)


# --- 3:  Q-equation

GG[2, 3] = '-1'                               # Q(t)
GG[2, 6] = '-1'                               # R(t)
FF[2, 0] = '1'                                # pi(t+1)
FF[2, 3] = '((1-tau)/(1-tau+r_k))'            # Q(t+1)
JJ[2, 0] = '(r_k/(1-tau+r_k))'                # r(t+1)
MM[2, 7] = '1'                                # n_Q(t)


# --- 4:  Capital

GG[3, 2] = '-1'                               # K(t)
HH[3, 2] = '(1-tau)'                          # K(t-1)
HH[3, 4] = 'tau'                              # I(t-1)


# --- 5:  Prices

GG[4, 0] = '-1'                               # pi(t)
FF[4, 0] = '(betta/(1+betta*gamma_p))'        # pi(t+1)
HH[4, 0] = '(gamma_p/(1+betta*gamma_p))'      # pi(t-1)
KK[4, 0] = '(((1-betta*xi_p)*(1-xi_p))/((1+betta*gamma_p)*xi_p))*alphaa'         # r(t)
GG[4, 1] = '(((1-betta*xi_p)*(1-xi_p))/((1+betta*gamma_p)*xi_p))*(1-alphaa)'     # w(t)
MM[4, 0] = '-(((1-betta*xi_p)*(1-xi_p))/((1+betta*gamma_p)*xi_p))'               # e_a(t)
MM[4, 6] = '(((1-betta*xi_p)*(1-xi_p))/((1+betta*gamma_p)*xi_p))'                # n_p(t)


# --- 6:  Wages

GG[5, 1] = '(-1-((1/(1+betta))*((1-betta*xi_w)*(1-xi_w))/((1+(1/lambda_w)*((1+lambda_w)*sigma_L))*xi_w)))'  # w(t)
FF[5, 1] = '(betta/(1+betta))'                # w(t+1)
HH[5, 1] = '(1/(1+betta))'                    # w(t-1)
FF[5, 0] = '(betta/(1+betta))'                # pi(t+1)
GG[5, 0] = '-((1+betta*gamma_w)/(1+betta))'   # pi(t)
HH[5, 0] = '(gamma_w/(1+betta)) '             # pi(t-1)
KK[5, 1] = '((1/(1+betta))*((1-betta*xi_w)*(1-xi_w))/((1+(1/lambda_w)*((1+lambda_w)*sigma_L))*xi_w))*sigma_L'  # L(t)
HH[5, 5] = '-((1/(1+betta))*((1-betta*xi_w)*(1-xi_w))/((1+(1/lambda_w)*((1+lambda_w)*sigma_L))*xi_w))*(sigma_c/(1-h))*h'  # C(t-1)
GG[5, 5] = '((1/(1+betta))*((1-betta*xi_w)*(1-xi_w))/((1+(1/lambda_w)*((1+lambda_w)*sigma_L))*xi_w))*(sigma_c/(1-h))'  # C(t)
MM[5, 3] = '((1/(1+betta))*((1-betta*xi_w)*(1-xi_w))/((1+(1/lambda_w)*((1+lambda_w)*sigma_L))*xi_w))'  # e_L(t)
MM[5, 9] = '((1/(1+betta))*((1-betta*xi_w)*(1-xi_w))/((1+(1/lambda_w)*((1+lambda_w)*sigma_L))*xi_w))'  # n_w(t)


# --- 7:  Labor demand

CC[0, 1] = '-1'                               # L(t)
AA[0, 1] = '-1'                               # w(t)
CC[0, 0] = '(1+psi)'                          # r(t)
BB[0, 2] = '1'                                # K(t-1)


# --- 8:  Production

AA[1, 7] = '-1'                               # Y(t)
DD[1, 0] = 'phi'                              # e_a(t)
BB[1, 2] = 'phi*alphaa'                       # K(t-1)
CC[1, 0] = 'phi*alphaa*psi'                   # r(t)
CC[1, 1] = 'phi*(1-alphaa)'                   # L(t)


# --- 9:     Goods market

GG[6, 7] = '-1'                               # Y(t)
GG[6, 5] = '(1-tau*k-g)'                      # C(t)
GG[6, 4] = 'tau*k'                            # I(t)
MM[6, 4] = '1'                                # e_G(t)


# --- 10: Monetary policy

GG[7, 6] = '-1'                               # R(t)
HH[7, 6] = 'rho'                              # R(t-1)
HH[7, 0] = '((1-rho)*r_pi-r_dpi)'             # pi(t-1)
MM[7, 5] = '(1-rho)*(1-r_pi)'                 # e_pi(t)
GG[7, 7] = '((1-rho)*r_Y+r_dY)'               # Y(t)
GG[7, 15] = '-((1-rho)*r_Y+r_dY)'             # Y_f(t)
GG[7, 0] = 'r_dpi'                            # pi(t)
HH[7, 7] = '-r_dY'                            # Y(t-1)
HH[7, 15] = '-r_dY'                           # Y_f(t-1)
MM[7, 8] = '1'                                # n_R(t)


# --- 11:  Flexible Consumption

GG[8, 13] = '-1'                               # C_f(t)
HH[8, 13] = '(h/(1+h))'                        # C_f(t-1)
FF[8, 13] = '(1/(1+h))'                        # C_f(t+1)
GG[8, 14] = '-((1-h)/((1+h)*sigma_c))'         # R_f(t)
FF[8, 8] = '((1-h)/((1+h)*sigma_c))'           # pi_f(t+1)
MM[8, 2] = '((1-h)/((1+h)*sigma_c))'           # e_b(t)
LL[8, 2] = '-((1-h)/((1+h)*sigma_c))'          # e_b(t+1)


# --- 12:  Flexible Investment

GG[9, 12] = '-1'                               # I_f(t)
HH[9, 12] = '(1/(1+betta))'                    # I_f(t-1)
FF[9, 12] = '(betta/(1+betta))'                # I_f(t+1)
GG[9, 11] = '(adj/(1+betta))'                  # Q_f(t)
LL[9, 1] = '(1+betta)*betta'                   # e_I(t+1)
MM[9, 1] = '-(1+betta)'                        # e_I(t)


# --- 13:  Flexible Q-equation

GG[10, 11] = '-1'                               # Q_f(t)
GG[10, 14] = '-1'                               # R_f(t)
FF[10, 8] = '1'                                 # pi_f(t+1)
FF[10, 11] = '((1-tau)/(1-tau+r_k))'            # Q_f(t+1)
JJ[10, 2] = '(r_k/(1-tau+r_k))'                 # r_f(t+1)


# --- 14:  Flexible Capital

GG[11, 10] = '-1'                               # K_f(t)
HH[11, 10] = '(1-tau)'                          # K_f(t-1)
HH[11, 12] = 'tau'                              # I_f(t-1)


# --- 15:  Flexible Prices

KK[12, 2] = 'alphaa'                           # r_f(t)
GG[12, 9] = '(1-alphaa)'                       # w_f(t)
MM[12, 0] = '-1'                               # e_a(t)


# --- 16:  Flexible Wages

GG[13, 9] = '-1'                               # w_f(t)
KK[13, 3] = 'sigma_L'                          # L_f(t)
HH[13, 13] = '-(sigma_c/(1-h))*h'              # C_f(t-1)
GG[13, 13] = '(sigma_c/(1-h))'                 # C_f(t)
MM[13, 3] = '-1'                               # e_L(t)


# --- 17:  Flexible Labor demand

CC[2, 3] = '-1'                               # L_f(t)
AA[2, 9] = '-1'                               # w_f(t)
CC[2, 2] = '(1+psi)'                          # r_f(t)
BB[2, 10] = '1'                               # K_f(t-1)


# --- 18:  Flexible Production

AA[3, 15] = '-1'                              # Y_f(t)
DD[3, 0] = 'phi'                              # e_a(t)
BB[3, 10] = 'phi*alphaa'                      # K_f(t-1)
CC[3, 2] = 'phi*alphaa*psi'                   # r_f(t)
CC[3, 3] = 'phi*(1-alphaa)'                   # L_f(t)


# --- 19:  Flexible Goods market

GG[14, 15] = '-1'                              # Y_f(t)
GG[14, 13] = '(1-tau*k-g)'                     # C_f(t)
GG[14, 12] = 'tau*k'                           # I_f(t)
MM[14, 4] = '1'                                # e_G(t)


# --- 20:   Aggregate productivity shock

NN[0, 1] = 'rho_I'                            # e_I(t)


# --- 21:   Adjustment cost shock

NN[1, 2] = 'rho_b'                            # e_b(t)


# --- 22:   Preference shock

NN[2, 3] = 'rho_L'                            # e_L(t)


# --- 23:   Labor supply shock

NN[3, 4] = 'rho_G'                            # e_G(t)


# --- 24:   Public spending shock

NN[4, 0] = 'rho_a'                            # e_a(t)


# --- 25:   Inflation target

NN[5, 5] = 'rho_pi'                           # e_pi(t)


# --- 26: Temporary first equation

GG[15, 8] = '-1'                               # pi_f(t)
GG[15, 16] = '0'                                # one(t)


# --- 27: Temporary second equation

GG[16, 16] = '-1'                               # one(t)
HH[16, 16] = '0'                                # one(t-1)


# --- 26: Prices error expectation

NN[6, 6] = '0'                               # n_p(t)


# --- 27: Tobin Q-ratio error expectation

NN[7, 7] = '0'                               # n_Q(t)


# --- 28: Monetary policy interest rate error expectation

NN[8, 8] = '0'                               # n_R(t)


# --- 29: Wages error expectation

NN[9, 9] = '0'                               # n_w(t)


# --------------------------------------------------------------------------
# --- Evaluate
# --------------------------------------------------------------------------

# A = vstack((hstack((matrix_names[1], zeros((ny,ny)), zeros((ny,nz)))), hstack((matrix_names[6], zeros((nx, ny)), zeros((nx, nz)))),
#             hstack((zeros((nz,nx)), zeros((nz,ny)), matrix_names[11]))))
#
# B = vstack((hstack((matrix_names[0], matrix_names[2], matrix_names[3])), hstack((matrix_names[5], matrix_names[8], matrix_names[10])),
#             hstack((zeros((nz, nx)), zeros((nz, ny)), -eye((nz))))))
#
# C = vstack((zeros((ny, nx+ny+nz)), hstack((matrix_names[4], matrix_names[7], matrix_names[9])), zeros((nz, nx+ny+nz))))
#
# with open('MatricesABC.pickle', 'wb') as f:
#     pickle.dump([A, B, C], f)
#
# np.save('A', A)
# np.save('B', B)
# np.save('C', C)

for name in matrix_names:
    for (x, y), i in np.ndenumerate(name):
        name[x, y] = float(0 if i is None else eval(i))

matrix_names = [np.asarray(x, dtype=np.float64, order='C') for x in matrix_names]

# --------------------------------------------------------------------------
# --- Matrix form in spirit of Uhlig
# --------------------------------------------------------------------------

# Sticking to Uhlig notation
# We know that:
# 1) 0 = Ax + Bx(-1) + Cy + Dz
# 2) 0 = E[Fx(+1) + Gx + Hx(-1) + Jy(+1) + Ky + Lz(+1) + Mz]
# 3) z(+1) = Nz + epsilon(+1)
# 4) E[epsilon(+1)] = 0
# Solution in the form:
# 1) x(+1) = Px + Qz
# 2) y = Rx + Sz

WWW = np.zeros((matrix_names[0].shape[0], 1))
TT = np.zeros((matrix_names[0].shape[1], 1))
Z0 = np.zeros((matrix_names[3].shape[1], 1))

P, Q, U, R, S, V = LinApp_Solve(matrix_names[0], matrix_names[1], matrix_names[2], matrix_names[3], matrix_names[4],
                                matrix_names[5], matrix_names[6], matrix_names[7], matrix_names[8], matrix_names[9],
                                matrix_names[10], WWW, TT, matrix_names[11], Z0, 1)

# Form to compare with Linear Time Iteration

F_Uhlig = vstack((P, R))
Q_Uhlig = vstack((Q, S))

# --------------------------------------------------------------------------
# --- Matrix form in spirit of Linear Time Iteration
# --------------------------------------------------------------------------

# System to solve: Ax(-1) + Bx + CE[x(+1)] + epsilon = 0
# Solution in the form: x = Fx(-1) + Q*epsilon
# Sticking to the notation from the paper: x is (nx+ny+nz) x 1 vector

A = vstack((hstack((matrix_names[1], zeros((ny,ny)), zeros((ny,nz)))), hstack((matrix_names[6], zeros((nx, ny)), zeros((nx, nz)))),
            hstack((zeros((nz,nx)), zeros((nz,ny)), matrix_names[11]))))

B = vstack((hstack((matrix_names[0], matrix_names[2], matrix_names[3])), hstack((matrix_names[5], matrix_names[8], matrix_names[10])),
            hstack((zeros((nz, nx)), zeros((nz, ny)), -eye((nz))))))

C = vstack((zeros((ny, nx+ny+nz)), hstack((matrix_names[4], matrix_names[7], matrix_names[9])), zeros((nz, nx+ny+nz))))

# np.save('A_value', A, allow_pickle=False)
# np.save('B_value', B, allow_pickle=False)
# np.save('C_value', C, allow_pickle=False)

F_iter, Q_iter = Linear_Time_Iteration(A, B, C, 1e-16, 1e-4)

# Check whether the results are identical with Uhlig

print(np.amax(F_iter[:nx+ny, :nx] - F_Uhlig))
# print(Q_iter)
# print('\n')
# print(Q_Uhlig)
# print(np.amax(Q_iter[:nx+ny, nx+ny:] - Q_Uhlig))

# --------------------------------------------------------------------------
# --- Simulation
# --------------------------------------------------------------------------

T = 1000  # Number of periods to simulate

epsilon = np.array([np.random.normal(loc = 0, scale = sd_Q),
                    np.random.normal(loc = 0, scale = sd_p), np.random.normal(loc = 0, scale = sd_w),
                    np.random.normal(loc = 0, scale = sd_R)])
epsilon = np.append(np.random.randn(6), epsilon)
epsilon = np.append(np.zeros((nx+ny)), epsilon)

X_sim = np.zeros((nx+ny+nz, T))
X_sim[:,0] = Q_iter.dot(epsilon)

for t in range(1, T):
    epsilon = np.array([np.random.normal(loc=0, scale=sd_Q),
                        np.random.normal(loc=0, scale=sd_p), np.random.normal(loc=0, scale=sd_w),
                        np.random.normal(loc=0, scale=sd_R)])
    epsilon = np.append(np.random.randn(6), epsilon)
    epsilon = np.append(np.zeros((nx + ny)), epsilon)
    X_sim[:, t] = F_iter.dot(X_sim[:, t-1]) + Q_iter.dot(epsilon)

# Plot for consumption

# plt.plot(X_sim[11,:])
# plt.show()

# --------------------------------------------------------------------------
# --- Estimation
# --------------------------------------------------------------------------

# Observables: labor, consumption, investment
# Indexes: 14, 11, 10

no = 3  # Number of observables

observation_matrix = zeros((no, nx+ny+nz))
observation_matrix[0,10] = observation_matrix[1,11] = observation_matrix[2,14] = 1

observables = X_sim[(10,11,14), :]

# Kalman filter

from Kalman import Kalman

kalman = Kalman(A = F_iter, C = Q_iter, G = observation_matrix, x_hat = X_sim[:, 0])
log_like = kalman.compute_loglikelihood(observables)



