# 1. Input model parameters
parameters = pd.Series()
parameters['rhoa'] = .9
parameters['sigma'] = 0.001

print(parameters)

# 2. Define a function that evaluates the equilibrium conditions
def equilibrium_equations(variables_forward,variables_current,parameters):
    
    # Parameters 
    p = parameters
    
    # Variables
    fwd = variables_forward
    cur = variables_current
        
    # Exogenous tfp
    tfp_proc = p.rhoa*np.log(cur.a) - np.log(fwd.a)
    
    # Stack equilibrium conditions into a numpy array
    return np.array([
        tfp_proc
        ])

# 3. Initialize the model
model = ls.model(equations = equilibrium_equations,
                 nstates=1,
                 varNames=['a'],
                 shockNames=['eA'],
                 parameters = parameters)

# 4. Have linearsolve compute the steady state numerically
guess = [1]
model.compute_ss(guess)

print(model.ss)

# 5. Find the log-linear approximation around the non-stochastic steady state and solve
model.approximate_and_solve()

# 6 (a) Compute impulse responses
model.impulse(T=41,t0=5,shock=None)
print(model.irs['eA'].head(10))

# 6 (b) Plot impulse responses
model.irs['eA'][['eA','a']].plot(lw='5',alpha=0.5,grid=True).legend(loc='upper right',ncol=2)

# 6(c) Compute stochastic simulation
model.stoch_sim(seed=192,covMat= [[parameters['sigma']]])
print(model.simulated.head(10))

# 6(d) Plot stochastic simulation
model.simulated[['eA','a']].plot(lw='5',alpha=0.5,grid=True).legend(loc='upper right',ncol=2)

"""
Explanation: Class 15: Introduction to the linearsolve package
In general, dynamic stochastic general equilibrium (DSGE) models are time-consuming to work with. The linearsolve Python package approximates, solves, and simulates DSGE models.
Example 1: A one-equation model of TFP
Consider the following AR(1) specification for $\log$ TFP:
\begin{align}
\log A_{t+1} & = \rho \log A_t + \epsilon_{t+1} \tag{1}
\end{align}
where $\epsilon_{t+1} \sim \mathcal{N}(0,\sigma^2)$. Let's simulate the model with linearsolve. To do this we need to do several things:

Create a Pandas series that stores the names of the parameters of the model.
Define a function that returns the equilibrium conditions of the model solved for zero.
Initialize an instance of the linearsolve.model class
Compute and input the steady state of the model.
Approximate and solve the model.
Compute simulations of the model.
End of explanation
"""


# 1. Input model parameters


# 2. Define a function that evaluates the equilibrium conditions
def equilibrium_equations(variables_forward,variables_current,parameters):
    
    # Parameters 
    p = parameters
    
    # Variables
    fwd = variables_forward
    cur = variables_current
        
    # Production function
    prod_fn = cur.a*cur.k**p.alpha - cur.y
    
    # Capital evolution
    
    
    # Exogenous tfp
    tfp_proc = p.rhoa*np.log(cur.a) - np.log(fwd.a)
    
    
    # Stack equilibrium conditions into a numpy array
    return np.array([
        tfp_proc
        
        ])

# 3. Initialize the model
model = ls.model(equations = equilibrium_equations,
                 nstates=,
                 varNames=[],                   # Any order as long as the state variables are named first
                 shockNames=[],                 # Name a shock for each state variable *even if there is no corresponding shock in the model*
                 parameters = parameters)

# 4. Have linearsolve compute the steady state numerically


# 5. Find the log-linear approximation around the non-stochastic steady state and solve


"""
Explanation: Example 2: A three-equation model business cycle model
Now consider the following system of equations:
\begin{align}
Y_t & = A_t K_t^{\alpha} \tag{2}\
K_{t+1} & = sY_t + (1-\delta) K_t \tag{3}\
\log A_{t+1} & = \rho \log A_t + \epsilon_{t+1} \tag{4}
\end{align}
where $\epsilon_{t+1} \sim \mathcal{N}(0,\sigma^2)$. Let's simulate the model with linearsolve. Before proceding, let's also go ahead and rewrite the model with all variables moved to the lefthand side of the equations:
\begin{align}
0 & = A_t K_t^{\alpha} - Y_t \tag{5}\
0 & = sY_t + (1-\delta) K_t - K_{t+1} \tag{6}\
0 & = \rho \log A_t + \epsilon_{t+1} - \log A_{t+1} \tag{7}
\end{align}
Capital and TFP are called state variables because they're $t+1$ values are predetermined. Output is called a costate variable. Note that the model as 3 endogenous variables with 2 state variables.
End of explanation
"""


# Print the coeficient matrix P


# Print the coeficient matrix F


# 6 (a) Compute impulse responses and print the computed impulse responses


# 6(b) Plot the computed impulse responses to a TFP shock


# 6(c) Compute stochastic simulation and print the simulated values


# 6(d) Plot the computed stochastic simulation


"""
Explanation: The previous step constructs a log-linear approximation of the model and then solves for the endogenous variables as functions of the state variables and exogenous shocks only:
\begin{align}
\left[ \hat{y}t\right]  &  = F \left[\begin{array}{c} \hat{a}_t \ \hat{k}_t \end{array}\right]\
\left[\begin{array}{c} \hat{a}{t+1} \ \hat{k}{t+1} \end{array}\right] & = P \left[\begin{array}{c} \hat{a}_t \ \hat{k}_t \end{array}\right] + \left[\begin{array}{c} \epsilon{t+1}^a \ \epsilon^k_{t+1} \end{array}\right].
\end{align}
where $F$ and $P$ are coefficient matrices computed by the program.
End of explanation
"""


