import pandas as pd
import matplotlib.pyplot as plt
from dolo import yaml_import, approximate_controls, display
from dolo.algos import perfect_foresight as pf

# Declare plotting style
pd.options.display.mpl_style='default'

model = yaml_import("dolo/Figv4_1191.yaml")
T = 101  # simluation length
p_T = 40  # Periods to plot
g = np.ones(T + 1) * .4
g[:9] = 0.2
g_shock = np.atleast_2d(g[:100])
sol = pf.deterministic_solve(model, shocks=g_shock, T=T, ignore_constraints=True)
sol['g'] = g

# Compute initial steady states
kbar, cbar = pf.find_steady_state(model, np.array([0.2]))
a = model.functions['auxiliary']
etabar, wbar, Rbar = a(kbar, cbar, model.calibration['parameters'])

# Set up plotting materials 
o = np.ones(p_T)
names = [['k', 'c', 'R'], ['eta', 'g', 'w']]
titles = [['k', 'c', r'$\bar{R}$'], [r'$\eta$', 'g', 'w']]
ss = [[kbar, cbar, Rbar], [etabar, g[0], wbar]]

# Generate plots
psol = sol.ix[:p_T]
fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(8, 8))
for ind_i, i in enumerate(names):
    for ind_k, k in enumerate(i):
        ax_ik = axes[ind_i, ind_k]
        psol[k].plot(ax=ax_ik, linewidth=2, title=titles[ind_i][ind_k])
        ax_ik.plot(o * ss[ind_i][ind_k], 'k--')
axes[1,1].set_ybound(.18, .42)
fig.suptitle('RMT4 Figure 11.9.1', fontsize=18)

# Change gamma and compute solution again
model.set_calibration('gamma', 0.2)
sol2 = pf.deterministic_solve(model, shocks=g_shock, T=T, ignore_constraints=True)
sol2['g'] = g

# Generate plot
psol2 = sol2.ix[:p_T]
fig2, axes2 = plt.subplots(nrows=2, ncols=3, figsize=(8, 8))
for ind_i, i in enumerate(names):
    for ind_k, k in enumerate(i):
        ax_ik = axes2[ind_i, ind_k]
        psol[k].plot(ax=ax_ik, linewidth=2, title=titles[ind_i][ind_k])
        psol2[k].plot(ax=ax_ik, linewidth=3, style='-.')
        ax_ik.plot(o * ss[ind_i][ind_k], 'k--')
axes2[1,1].set_ybound(.18, .42)
fig2.suptitle('RMT4 Figure 11.9.2', fontsize=18)

model.set_calibration('gamma', 2.0)
beta, gamma = model.get_calibration('beta'), 2.0
n = np.arange(1, sol.shape[0] + 1)

# Compute q and qbar
q = beta ** n  * sol['c'] ** (-gamma)  / (beta ** 2 * sol['c'][1] ** (-gamma))
qbar = beta ** (n - 1)
sol['q'] = q
sol['qbar'] = qbar

# Compute r and rbar
sol['r'] = sol['R'] - 1
rbar = sol['R'][0] - 1

# comupte r_{t, s+t}
rst = np.zeros((40, 62))
for t in range(62):
    for i in range(40):
        rst[i, t] = np.log(q[i+t]/  q[t]) / (-i)      

fig3, axes3 = plt.subplots(nrows=2, ncols=3, figsize=(8, 8))
psol3 = sol.ix[:p_T]
psol3['c'].plot(ax=axes3[0,0], linewidth=2, title='c'); axes3[0,0].plot(o * cbar, 'k--')
psol3[['q', 'qbar']].plot(ax=axes3[0,1], linewidth=2, title='q', legend=None, style=['-', 'k--'])
psol3['r'].plot(ax=axes3[0,2], linewidth=2, title=r'$r_{t, t+1}$'); axes3[0,2].plot(o * rbar, 'k--')
axes3[1,0].plot(range(p_T), rst[:, 1], range(p_T), rst[:, 10], '-.', range(p_T), rst[:, -1], 'k--', linewidth=2)
axes3[1, 0].set_title(r'$r_{t, t+t}$')
psol3['g'].plot(ax=axes3[1,1], linewidth=2, title='g'); axes3[1,1].plot(o * g[0], 'k--'); axes3[1,1].set_ybound(.18, .42)
fig3.suptitle('RMT4 Figure 11.9.3', fontsize=18)