!pip install --upgrade quantecon
!pip install interpolation

import numpy as np
from quantecon.optimize import brent_max, brentq
from interpolation import interp
from numba import njit
import matplotlib.pyplot as plt
%matplotlib inline
from quantecon import MarkovChain

class ConsumerProblem:
    """
    A class that stores primitives for the income fluctuation problem. The
    income process is assumed to be a finite state Markov chain.
    """
    def __init__(self,
                 r=0.01,                        # Interest rate
                 β=0.96,                        # Discount factor
                 Π=((0.6, 0.4),
                    (0.05, 0.95)),              # Markov matrix for z_t
                 z_vals=(0.5, 1.0),             # State space of z_t
                 b=0,                           # Borrowing constraint
                 grid_max=16,
                 grid_size=50,
                 u=np.log,                      # Utility function
                 du=njit(lambda x: 1/x)):       # Derivative of utility

        self.u, self.du = u, du
        self.r, self.R = r, 1 + r
        self.β, self.b = β, b
        self.Π, self.z_vals = np.array(Π), tuple(z_vals)
        self.asset_grid = np.linspace(-b, grid_max, grid_size)

def operator_factory(cp):
    """
    A function factory for building operator K.

    Here cp is an instance of ConsumerProblem.
    """
    # Simplify names, set up arrays
    R, Π, β, u, b, du = cp.R, cp.Π, cp.β, cp.u, cp.b, cp.du
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    γ = R * β


    @njit
    def euler_diff(c, a, z, i_z, σ):
        """
        The difference of the left-hand side and the right-hand side
        of the Euler Equation.
        """
        lhs = du(c)
        expectation = 0
        for i in range(len(z_vals)):
            expectation += du(interp(asset_grid, σ[:, i], R * a + z - c)) \
                * Π[i_z, i]
        rhs = max(γ * expectation, du(R * a + z + b))

        return lhs - rhs

    @njit
    def K(σ):
        """
        The operator K.

        Iteration with this operator corresponds to time iteration on the
        Euler equation.  Computes and returns the updated consumption policy
        σ.  The array σ is replaced with a function cf that implements
        univariate linear interpolation over the asset grid for each
        possible value of z.
        """
        σ_new = np.empty_like(σ)
        for i_a in range(len(asset_grid)):
            a = asset_grid[i_a]
            for i_z in range(len(z_vals)):
                z = z_vals[i_z]
                c_star = brentq(euler_diff, 1e-8, R * a + z + b, \
                    args=(a, z, i_z, σ)).root
                σ_new[i_a, i_z] = c_star

        return σ_new

    return K

def solve_model(cp,
                tol=1e-4,
                max_iter=1000,
                verbose=True,
                print_skip=25):

    """
    Solves for the optimal policy using time iteration

    * cp is an instance of ConsumerProblem
    """

    u, β, b, R = cp.u, cp.β, cp.b, cp.R
    asset_grid, z_vals = cp.asset_grid, cp.z_vals

    # Initial guess of σ
    σ = np.empty((len(asset_grid), len(z_vals)))
    for i_a, a in enumerate(asset_grid):
        for i_z, z in enumerate(z_vals):
            c_max = R * a + z + b
            σ[i_a, i_z] = c_max

    K = operator_factory(cp)

    i = 0
    error = tol + 1

    while i < max_iter and error > tol:
        σ_new = K(σ)
        error = np.max(np.abs(σ - σ_new))
        i += 1
        if verbose and i % print_skip == 0:
            print(f"Error at iteration {i} is {error}.")
        σ = σ_new

    if i == max_iter:
        print("Failed to converge!")

    if verbose and i < max_iter:
        print(f"\nConverged in {i} iterations.")

    return σ_new

cp = ConsumerProblem()
σ_star = solve_model(cp)

fig, ax = plt.subplots(figsize=(10, 6))
ax.plot(cp.asset_grid, σ_star[:, 0], label='$\sigma^*$')
ax.set(xlabel='asset level', ylabel='optimal consumption')
ax.legend()
plt.show()

m = ConsumerProblem(r=0.03, grid_max=4)
K = operator_factory(m)

σ_star = solve_model(m, verbose=False)
a = m.asset_grid
R, z_vals = m.R, m.z_vals

fig, ax = plt.subplots(figsize=(10, 8))
ax.plot(a, R * a + z_vals[0] - σ_star[:, 0], label='Low income')
ax.plot(a, R * a + z_vals[1] - σ_star[:, 1], label='High income')
ax.plot(a, a, 'k--')
ax.set(xlabel='Current assets',
       ylabel='Next period assets',
       xlim=(0, 4), ylim=(0, 4))
ax.legend()
plt.show()

r_vals = np.linspace(0, 0.04, 4)

fig, ax = plt.subplots(figsize=(10, 8))
for r_val in r_vals:
    cp = ConsumerProblem(r=r_val)
    σ_star = solve_model(cp, verbose=False)
    ax.plot(cp.asset_grid, σ_star[:, 0], label=f'$r = {r_val:.3f}$')

ax.set(xlabel='asset level', ylabel='consumption (low income)')
ax.legend()
plt.show()

def compute_asset_series(cp, T=500000, verbose=False):
    """
    Simulates a time series of length T for assets, given optimal
    savings behavior.

    cp is an instance of ConsumerProblem
    """
    Π, z_vals, R = cp.Π, cp.z_vals, cp.R  # Simplify names
    mc = MarkovChain(Π)
    σ_star = solve_model(cp, verbose=False)
    cf = lambda a, i_z: interp(cp.asset_grid, σ_star[:, i_z], a)
    a = np.zeros(T+1)
    z_seq = mc.simulate(T)
    for t in range(T):
        i_z = z_seq[t]
        a[t+1] = R * a[t] + z_vals[i_z] - cf(a[t], i_z)
    return a

cp = ConsumerProblem(r=0.03, grid_max=4)
a = compute_asset_series(cp)

fig, ax = plt.subplots(figsize=(10, 8))
ax.hist(a, bins=20, alpha=0.5, density=True)
ax.set(xlabel='assets', xlim=(-0.05, 0.75))
plt.show()

M = 25
r_vals = np.linspace(0, 0.04, M)
fig, ax = plt.subplots(figsize=(10, 8))

for b in (1, 3):
    asset_mean = []
    for r_val in r_vals:
        cp = ConsumerProblem(r=r_val, b=b)
        mean = np.mean(compute_asset_series(cp, T=250000))
        asset_mean.append(mean)
    ax.plot(asset_mean, r_vals, label=f'$b = {b:d}$')
    print(f"Finished iteration b = {b:d}")

ax.set(xlabel='capital', ylabel='interest rate')
ax.grid()
ax.legend()
plt.show()