# 1. Input model parameters and print
rho = .75
sigma = 0.006
alpha = 0.35
delta = 0.025
beta = 0.99
phi = 1.7317

parameters = pd.Series()
parameters['rho'] = rho
parameters['sigma'] = sigma
parameters['alpha'] = alpha
parameters['delta'] = delta
parameters['beta'] = beta
parameters['phi'] = phi

# 2. Compute the steady state equilibrium manually
A = 1
L = 1 / ( 1 + phi/(1-alpha)*((beta**-1 + (1-alpha)*delta -1)/(beta**-1+delta-1)))
K = L*(alpha*A/(beta**-1+delta-1))**(1/(1-alpha))
C = A*K**alpha*L**(1-alpha) - delta*K
Y = A*K**alpha*L**(1-alpha)
I = delta*K

print('A:',round(A,5))
print('L:',round(L,5))
print('K:',round(K,5))
print('Y:',round(Y,5))
print('I:',round(I,5))
print('C:',round(C,5))

# 3. 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
        
    # Resource constraint
    resource =  cur.a*cur.k**p.alpha*cur.l**(1-p.alpha) + (1-p.delta)*cur.k - fwd.k - cur.c
    
    # Exogenous tfp
    tfp_process = p.rho*np.log(cur.a) - np.log(fwd.a)
    
    # Euler equation
    euler = p.beta*(p.alpha*fwd.a*fwd.k**(p.alpha-1)*fwd.l**(1-p.alpha) + 1 - p.delta)/fwd.c - 1/cur.c
    
    # Labor supply
    labor_supply = p.phi/(1-cur.l) - (1-p.alpha)*cur.a*cur.k**alpha *cur.l**(-p.alpha)/cur.c
    
    # Production
    production = cur.a*cur.k**p.alpha*cur.l**(1-p.alpha) - cur.y
    
    # Capital evoluation
    capital_evolution = cur.i + (1-p.delta)*cur.k - fwd.k
    
    # Stack equilibrium conditions into a numpy array
    return np.array([
        resource,
        tfp_process,
        euler,
        labor_supply,
        production,
        capital_evolution
        ])

# 4. Initialize the model
model = ls.model(equations = equilibrium_equations,
                 nstates=2,
                 varNames=['a','k','c','y','i','l'],                   # Any order as long as the state variables are named first
                 shockNames=['eA','eK'],                 # Name a shock for each state variable *even if there is no corresponding shock in the model*
                 parameters = parameters)

guess = [A,K,C,Y,I,L]
model.compute_ss(guess)
model.approximate_and_solve()

print(model.ss)

# 5(a) Compute impulse responses and print the computed impulse responses
model.impulse(T=41,t0=5,shock=None,percent=True)

# 5(b) Plot the computed impulse responses to a TFP shock
fig = plt.figure(figsize=(12,12))

ax1 = fig.add_subplot(3,2,1)
model.irs['eA'][['y','i','k']].plot(lw=5,alpha=0.5,grid=True,ax = ax1).legend(loc='upper right',ncol=4)
ax1.set_title('Output, investment, capital')
ax1.set_ylabel('% dev')
ax1.set_xlabel('quarters')

ax2 = fig.add_subplot(3,2,2)
model.irs['eA'][['a','eA']].plot(lw=5,alpha=0.5,grid=True,ax = ax2).legend(loc='upper right',ncol=2)
ax2.set_title('TFP and TFP shock')
ax2.set_ylabel('% dev')
ax2.set_xlabel('quarters')

ax3 = fig.add_subplot(3,2,3)
model.irs['eA'][['y','c','l']].plot(lw=5,alpha=0.5,grid=True,ax = ax3).legend(loc='upper right',ncol=4)
ax3.set_title('Output, consumption, and labor')
ax3.set_ylabel('% dev')
ax3.set_xlabel('quarters')

plt.tight_layout()

# 6(a) Compute stochastic simulation and print the simulated values
model.stoch_sim(seed=192,covMat= [[parameters['sigma']**2,0],[0,0]])

# 6(b) Plot the computed stochastic simulation
fig = plt.figure(figsize=(12,4))

ax1 = fig.add_subplot(1,2,1)
model.simulated[['k','c','y','i','l']].plot(lw=5,alpha=0.5,grid=True,ax = ax1).legend(ncol=5,bbox_to_anchor=(0., 1.02, 1., .102), loc=3, mode="expand")

ax2 = fig.add_subplot(1,2,2)
model.simulated[['eA','a']].plot(lw=5,alpha=0.5,grid=True,ax = ax2).legend(ncol=2,bbox_to_anchor=(0., 1.02, 1., .102), loc=3)

"""
Explanation: Class 18: A Centralized Real Business Cycle Model with Labor
The Model
Setup
A representative household lives for an infinite number of periods. The expected present value of lifetime utility to the household from consuming $C_0, C_1, C_2, \ldots $ and working $L_0, L_1, L_2, \ldots $ is denoted by $U_0$:
\begin{align}
U_0 & =  \log (C_0) + \varphi \log (1-L_0) + \beta E_0 \left[\log (C_1) + \varphi \log(1-L_1)\right] + \beta^2 E_0 \left[\log (C_2) + \varphi \log(1-L_2)\right] + \cdots\
& =  E_0\sum_{t = 0}^{\infty} \beta^t\left[ \log (C_t) + \varphi \log (1-L_t) \right],
\end{align}
where $0<\beta<1$ is the household's subjective discount factor. $E_0$ denotes the expectation with respect to all information available as of date 0.
The household enters period 0 with capital $K_0>0$. Production in period $t$:
\begin{align}
F(A_t,K_t,L_t) & = A_t K_t^{\alpha}L_t^{1-\alpha}
\end{align}
where TFP $A_t$ is stochastic:
\begin{align}
\log A_{t+1} & = \log A_t + \epsilon_{t+1}
\end{align}
Capital depreciates at the constant rate $\delta$ per period and so the household's resource constraint in each period $t$ is:
\begin{align}
C_t + K_{t+1} & = A_t K_{t}^{\alpha}L_t^{1-\alpha}  + (1-\delta)K_t
\end{align}
Define output and investment:
\begin{align}
Y_t & =  A_t K_{t}^{\alpha}L_t^{1-\alpha} \
I_t & = K_{t+1} - (1-\delta)K_t
\end{align}
Optimization problem
In period 0, the household solves:
\begin{align}
& \max_{C_0,L_0,K_1} \; E_0\sum_{t=0}^{\infty}\beta^t\left[ \log (C_t) + \varphi \log (1-L_t) \right] \
& \; \; \; \;  \; \; \; \; \text{s.t.} \; \; \; \; C_t + K_{t+1} = A_t K_{t}^{\alpha}L_t^{1-\alpha}  + (1-\delta)K_t
\end{align}
which can be written as a choice of $L_0$ and $K_1$ only:
\begin{align}
\max_{L_0,K_1} \; E_0\sum_{t=0}^{\infty}\beta^t\left[\log ( A_t K_{t}^{\alpha}L_t^{1-\alpha}  + (1-\delta)K_t - K_{t+1}) + \varphi \log (1-L_t)\right]
\end{align}
Equilibrium
So given $K_0>0$ and $A_0$, the equilibrium paths for consumption, labor, capital, output, investment, and TFP are described described by:
\begin{align}
\frac{1}{C_t} & = \beta E_t \left[\frac{\alpha A_{t+1}K_{t+1}^{\alpha - 1}L_{t+1}^{1-\alpha} + 1 - \delta}{C_{t+1}}\right]\
\frac{\varphi}{1-L_t} & = \frac{(1-\alpha) A_t K_t^{\alpha} L_t^{-\alpha}}{C_t}\
C_t + K_{t+1} & = A_{t} K_t^{\alpha}L_t^{1-\alpha} + (1-\delta) K_t\
Y_t & = A_t K_t^{\alpha} L_t^{1-\alpha}\
I_t & = K_{t+1} - (1-\delta) K_t\
\log A_{t+1} & = \rho \log A_t + \epsilon_{t+1}
\end{align}
Calibration
For computation purposes, assume the following values for the parameters of the model:
\begin{align}
\beta & = 0.99\
\rho & = .75\
\sigma & = 0.006\
\alpha & = 0.35\
\delta & = 0.025\
\varphi & = 1.7317
\end{align}
The value for $\beta$ implies a steady state (annualized) real interest rate of about 4 percent: 
\begin{align}
4 \cdot \left(\beta^{-1} - 1\right)  & \approx 0.04040
\end{align}
$\rho = 0.75$ and $\sigma = 0.006$ are consistent with the statistical properties of the cyclical component of TFP in the US. $\alpha$ is set so that, consistent with the long-run average of the US, the labor share of income is about 65 percent of GDP. The deprecation rate of capital is calibrated to be about 10 percent annually. Finally, $\varphi$ was chosen last to ensure that in the steady state households allocate about 33 percent of their available time to labor.
Steady State
The steady state:
\begin{align}
A & = 1\
L & = \frac{1}{1 + \left(\frac{\varphi}{1-\alpha}\right)\left[ \frac{\beta^{-1} + (1-\alpha) \delta - 1}{\beta^{-1} + \delta - 1}\right]}\
K & = \left(\frac{\alpha A}{\beta^{-1} - 1 + \delta} \right)^{\frac{1}{1-\alpha}}L\
C & = AK^{\alpha}L^{1-\alpha} - \delta K\
Y & = AK^{\alpha} L^{1-\alpha}\
I & = \delta K
\end{align}
End of explanation
"""


# 7. Compute the standard deviations of Y, C, and I in model.simulated
print(model.simulated[['y','c','i']].std())

# 8. Compute the coefficients of correlation for Y, C, and I
print(model.simulated[['y','c','i']].corr())

"""
Explanation: Evaluation
Now we want to examine the statistical properties of the simulated model
End of explanation
"""


