using InstantiateFromURL
github_project("QuantEcon/quantecon-notebooks-julia", version = "0.2.0")

using LinearAlgebra, Statistics, Compat
using Distributions, Interpolations, Parameters, Plots, QuantEcon, Random
gr(fmt = :png);

# model
function LucasTree(;γ = 2.0,
                    β = 0.95,
                    α = 0.9,
                    σ = 0.1,
                    grid_size = 100)

    ϕ = LogNormal(0.0, σ)
    shocks = rand(ϕ, 500)

    # build a grid with mass around stationary distribution
    ssd = σ / sqrt(1 - α^2)
    grid_min, grid_max = exp(-4ssd), exp(4ssd)
    grid = range(grid_min, grid_max, length = grid_size)

    # set h(y) = β * int u'(G(y,z)) G(y,z) ϕ(dz)
    h = similar(grid)
    for (i, y) in enumerate(grid)
        h[i] = β * mean((y^α .* shocks).^(1 - γ))
    end

    return (γ = γ, β = β, α = α, σ = σ, ϕ = ϕ, grid = grid, shocks = shocks, h = h)
end

# approximate Lucas operator, which returns the updated function Tf on the grid
function lucas_operator(lt, f)

    # unpack input
    @unpack grid, α, β, h = lt
    z = lt.shocks

    Af = LinearInterpolation(grid, f, extrapolation_bc=Line())

    Tf = [ h[i] + β * mean(Af.(grid[i]^α .* z)) for i in 1:length(grid) ]
    return Tf
end

# get equilibrium price for Lucas tree
function solve_lucas_model(lt;
                           tol = 1e-6,
                           max_iter = 500)

    @unpack grid, γ = lt

    i = 0
    f = zero(grid)  # Initial guess of f
    error = tol + 1

    while (error > tol) && (i < max_iter)
        f_new = lucas_operator(lt, f)
        error = maximum(abs, f_new - f)
        f = f_new
        i += 1
    end

    # p(y) = f(y) * y ^ γ
    price = f .* grid.^γ

    return price
end

Random.seed!(42) # For reproducible results.

tree = LucasTree(γ = 2.0, β = 0.95, α = 0.90, σ = 0.1)
price_vals = solve_lucas_model(tree);

plot(tree.grid, price_vals, lw = 2, label = "p*(y)")
plot!(xlabel = "y", ylabel = "price", legend = :topleft)

plot()
for β in (.95, 0.98)
    tree = LucasTree(;β = β)
    grid = tree.grid
    price_vals = solve_lucas_model(tree)
    plot!(grid, price_vals, lw = 2, label = "beta = beta_var")
end

plot!(xlabel = "y", ylabel = "price", legend = :topleft)