mwe.jl

    using BasisMatrices, Optim, QuantEcon, Parameters
    using BasisMatrices: Degree, Derivative
    using Printf, Random
    using LinearAlgebra

    """
    The stochastic Neoclassical growth model type contains parameters
    which define the model

    * α: Capital share in output
    * β: Discount factor
    * δ: Depreciation rate
    * γ: Risk aversion
    * ρ: Persistence of the log of the productivity level
    * σ: Standard deviation of shocks to log productivity level
    * A: Coefficient on C-D production function

    * kgrid: Grid over capital
    * zgrid: Grid over productivity
    * grid: Grid of (k, z) pairs
    * eps_nodes: Nodes used to integrate
    * weights: Weights used to integrate
    * z1: A grid of the possible z1s tomorrow given eps_nodes and zgrid
    """
    @with_kw struct NeoclassicalGrowth
        # Parameters
        α::Float64 = 0.36
        β::Float64 = 0.99
        δ::Float64 = 0.02
        γ::Float64 = 2.0
        ρ::Float64 = 0.95
        σ::Float64 = 0.01
        A::Float64 = (1.0/β - (1 - δ)) / α

        # Grids
        kgrid::Vector{Float64} = collect(range(0.9, stop = 1.1,length = 10))
        zgrid::Vector{Float64} = collect(range(0.9, stop = 1.1,length = 10))
        grid::Matrix{Float64} = gridmake(kgrid, zgrid)
        eps_nodes::Vector{Float64} = qnwnorm(5, 0.0, σ^2)[1]
        weights::Vector{Float64} = qnwnorm(5, 0.0, σ^2)[2]
        z1::Matrix{Float64} = (zgrid.^(ρ))' .* exp.(eps_nodes)
    end

    # Helper functions
    f(ncgm::NeoclassicalGrowth, k, z) = z .* (ncgm.A*k.^ncgm.α)
    df(ncgm::NeoclassicalGrowth, k, z) = ncgm.α * z .* (ncgm.A * k.^(ncgm.α-1.0))

    u(ncgm::NeoclassicalGrowth, c) = c > 1e-10 ? (c^(1-ncgm.γ)-1)/(1-ncgm.γ) : -1e10
    du(ncgm::NeoclassicalGrowth, c) = c > 1e-10 ? c^(-ncgm.γ) : 1e10
    duinv(ncgm::NeoclassicalGrowth, u) = u .^ (-1 / ncgm.γ)

    expendables_t(ncgm::NeoclassicalGrowth, k, z) = (1-ncgm.δ)*k + f(ncgm, k, z)


    # Types for solution methods
    abstract type SolutionMethod end

    struct IterateOnPolicy <: SolutionMethod end
    struct VFI_ECM <: SolutionMethod end
    struct VFI_EGM <: SolutionMethod end
    struct VFI <: SolutionMethod end
    struct PFI_ECM <: SolutionMethod end
    struct PFI <: SolutionMethod end
    struct dVFI_ECM <: SolutionMethod end
    struct EulEq <: SolutionMethod end

    #
    # Type for Approximating Value and Policy
    #
    mutable struct ValueCoeffs{T<:SolutionMethod,D<:Degree}
        d::D
        v_coeffs::Vector{Float64}
        k_coeffs::Vector{Float64}
    end


    function ValueCoeffs(::Type{Val{d}}, method::T) where {T <: SolutionMethod, d}
        # Initialize two vectors of zeros
        deg = Degree{d}()
        n = n_complete(2, deg)
        v_coeffs = zeros(n)
        k_coeffs = zeros(n)

        return ValueCoeffs{T,Degree{d}}(deg, v_coeffs, k_coeffs)
    end

    function ValueCoeffs(ncgm::NeoclassicalGrowth, ::Type{Val{d}}, method::T) where {T <: SolutionMethod, d}
        # Initialize with vector of zeros
        deg = Degree{d}()
        n = n_complete(2, deg)
        v_coeffs = zeros(n)

        # Policy guesses based on k and z
        k, z = ncgm.grid[:, 1], ncgm.grid[:, 2]
        css = ncgm.A - ncgm.δ
        yss = ncgm.A
        c_pol = f(ncgm, k, z) * (css/yss)

        # Figure out what kp is
        k_pol = expendables_t(ncgm, k, z) - c_pol
        k_coeffs = complete_polynomial(ncgm.grid, d) \ k_pol

        return ValueCoeffs{T,Degree{d}}(deg, v_coeffs, k_coeffs)
    end

    function solutionmethod(::ValueCoeffs{T}) where T <: SolutionMethod
        return T
    end
    # solutionmethod{T<:SolutionMethod}(::ValueCoeffs{T}) = T

    # A few copy methods to make life easier
    Base.copy(vp::ValueCoeffs{T, D}) where {T, D} =
        ValueCoeffs{T,D}(vp.d, vp.v_coeffs, vp.k_coeffs)

    Base.copy(vp::ValueCoeffs{T1, D}, ::T2) where {T1, D, T2 <: SolutionMethod} =
        ValueCoeffs{T2,D}(vp.d, vp.v_coeffs, vp.k_coeffs)

    function Base.copy(ncgm::NeoclassicalGrowth, vp::ValueCoeffs{T}, ::Type{Val{new_degree}}) where {T, new_degree}
        # Build Value and policy matrix
        deg = Degree{new_degree}()
        V = build_V(ncgm, vp)
        k = build_k(ncgm, vp)

        # Build new Phi
        Phi = complete_polynomial(ncgm.grid, deg)
        v_coeffs = Phi \ V
        k_coeffs = Phi \ k

        return ValueCoeffs{T,Degree{new_degree}}(deg, v_coeffs, k_coeffs)
    end


    """
    Updates the coefficients for the value function inplace in `vp`
    """
    function update_v!(vp::ValueCoeffs, new_coeffs::Vector{Float64}, dampen::Float64)
        vp.v_coeffs = (1-dampen)*vp.v_coeffs + dampen*new_coeffs
    end

    """
    Updates the coefficients for the policy function inplace in `vp`
    """
    function update_k!(vp::ValueCoeffs, new_coeffs::Vector{Float64}, dampen::Float64)
        vp.k_coeffs = (1-dampen)*vp.k_coeffs + dampen*new_coeffs
    end

    """
    Builds either V or dV depending on the solution method that is given. If it
    is a solution method that iterates on the derivative of the value function
    then it will return derivative of the value function, otherwise the
    value function itself
    """
    build_V_or_dV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs) =
        build_V_or_dV(ncgm, vp, solutionmethod(vp)())

    build_V_or_dV(ncgm, vp::ValueCoeffs, ::SolutionMethod) = build_V(ncgm, vp)
    build_V_or_dV(ncgm, vp::ValueCoeffs, T::dVFI_ECM) = build_dV(ncgm, vp)

    function build_dV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
        Φ = complete_polynomial(ncgm.grid, vp.d, Derivative{1}())
        Φ*vp.v_coeffs
    end

    function build_V(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
        Φ = complete_polynomial(ncgm.grid, vp.d)
        Φ*vp.v_coeffs
    end

    function build_k(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
        Φ = complete_polynomial(ncgm.grid, vp.d)
        Φ*vp.k_coeffs
    end


    function compute_EV!(cp_kpzp::Vector{Float64}, ncgm::NeoclassicalGrowth,
                         vp::ValueCoeffs, kp, iz)
        # Pull out information from types
        z1, weightsz = ncgm.z1, ncgm.weights

        # Get number nodes
        nzp = length(weightsz)

        EV = 0.0
        for izp in 1:nzp
            zp = z1[izp, iz]
            complete_polynomial!(cp_kpzp, [kp, zp], vp.d)
            EV += weightsz[izp] * dot(vp.v_coeffs, cp_kpzp)
        end

        return EV
    end

    function compute_EV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs, kp, iz)
        cp_kpzp = Array{Float64}(undef, n_complete(2, vp.d))

        return compute_EV!(cp_kpzp, ncgm, vp, kp, iz)
    end

    function compute_EV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
        # Get length of k and z grids
        kgrid, zgrid = ncgm.kgrid, ncgm.zgrid
        nk, nz = length(kgrid), length(zgrid)
        temp = Array{Float64}(undef, n_complete(2, vp.d))

        # Allocate space to store EV
        EV = Array{Float64}(undef, nk*nz)

        for ik in 1:nk, iz in 1:nz
            # Pull out states
            k = kgrid[ik]
            z = zgrid[iz]
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]

            # Pass to scalar EV
            complete_polynomial!(temp, [k, z], vp.d)
            kp = dot(vp.k_coeffs, temp)
            EV[ikiz_index] = compute_EV!(temp, ncgm, vp, kp, iz)
        end

        return EV
    end


    function compute_dEV!(cp_dkpzp::Vector,  ncgm::NeoclassicalGrowth,
                          vp::ValueCoeffs, kp, iz)
        # Pull out information from types
        z1, weightsz = ncgm.z1, ncgm.weights

        # Get number nodes
        nzp = length(weightsz)

        dEV = 0.0
        for izp in 1:nzp
            zp = z1[izp, iz]
            complete_polynomial!(cp_dkpzp, [kp, zp], vp.d, Derivative{1}())
            dEV += weightsz[izp] * dot(vp.v_coeffs, cp_dkpzp)
        end

        return dEV
    end

    function compute_dEV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs, kp, iz)
        compute_dEV!(Array{Float64}(undef, n_complete(2, vp.d)), ncgm, vp, kp, iz)
    end

## solver

    function solve(ncgm::NeoclassicalGrowth, vp::ValueCoeffs{T};
        tol::Float64 = 1e-06, maxiter::Int = 5000, dampen::Float64 = 1, 
        nskipprint::Int = 1, verbose::Bool = true) where {T <: SolutionMethod}

        # Get number of k and z on grid
        nk, nz = length(ncgm.kgrid), length(ncgm.zgrid)

        # Build basis matrix and value function
        dPhi = complete_polynomial(ncgm.grid, vp.d, Derivative{1}())
        Phi = complete_polynomial(ncgm.grid, vp.d)
        V = build_V_or_dV(ncgm, vp)
        k = build_k(ncgm, vp)
        Vnew = copy(V)
        knew = copy(k)

        # Print column names
        if verbose
            @printf("| Iteration | Distance V | Distance K |\n")
        end

        # Iterate to convergence
        dist, iter = 10.0, 0
        while (tol < dist) & (iter < maxiter)
            # Update the value function using appropriate update methods
            update!(Vnew, knew, ncgm, vp, Phi, dPhi)

            # Compute distance and update all relevant elements
            iter += 1
            dist_v = maximum(abs, 1.0 .- Vnew./V)
            dist_k = maximum(abs, 1.0 .- knew./k)
            copyto!(V, Vnew)
            copyto!(k, knew)

            # If we are iterating on a policy, use the difference of values
            # otherwise use the distance on policy
            dist = ifelse(solutionmethod(vp) == IterateOnPolicy, dist_v, dist_k)

            # Print status update
            if verbose && (iter%nskipprint == 0)
                @printf("|%-11d|%-12e|%-12e|\n", iter, dist_v, dist_k)
            end
        end

        # Update value and policy functions one last time as long as the
        # solution method isn't IterateOnPolicy
        if ~(solutionmethod(vp) == IterateOnPolicy)
            # Update capital policy after finished
            kp = env_condition_kp(ncgm, vp)
            update_k!(vp, complete_polynomial(ncgm.grid, vp.d) \ kp, 1.0)

            # Update value function according to specified policy
            vp_igp = copy(vp, IterateOnPolicy())
            solve(ncgm, vp_igp; tol=1e-10, maxiter=5000, verbose=false)
            update_v!(vp, vp_igp.v_coeffs, 1.0)

        end

        return vp
    end

## updaters

    function update!(V::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{IterateOnPolicy},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Get sizes and allocate for complete_polynomial
        kgrid = ncgm.kgrid; zgrid = ncgm.zgrid;
        nk, nz = length(kgrid), length(zgrid)

        # Iterate over all states
        for ik in 1:nk, iz in 1:nz
            # Pull out states
            k = kgrid[ik]
            z = zgrid[iz]

            # Pull out policy and evaluate consumption
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            k1 = kpol[ikiz_index]
            c = expendables_t(ncgm, k, z) - k1

            # New value
            EV = compute_EV(ncgm, vp, k1, iz)
            V[ikiz_index] = u(ncgm, c) + ncgm.β*EV
        end

        # Update coefficients
        update_v!(vp, Φ \ V, 1.0)
        update_k!(vp, Φ \ kpol, 1.0)

        return V
    end


    function update!(V::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{VFI},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Get sizes and allocate for complete_polynomial
        kgrid = ncgm.kgrid; zgrid = ncgm.zgrid
        nk, nz = length(kgrid), length(zgrid)

        # Iterate over all states
        temp = Array{Float64}(undef,n_complete(2, vp.d))
        for iz=1:nz, ik=1:nk
            k = kgrid[ik]; z = zgrid[iz]

            # Define an objective function (negative for minimization)
            y = expendables_t(ncgm, k, z)
            solme(kp) = du(ncgm, y - kp) - ncgm.β*compute_dEV!(temp, ncgm, vp, kp, iz)

            # Find sol to foc
            kp = brent(solme, 1e-8, y-1e-8; rtol=1e-12)
            c = expendables_t(ncgm, k, z) - kp

            # New value
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            # ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            EV = compute_EV!(temp, ncgm, vp, kp, iz)
            V[ikiz_index] = u(ncgm, c) + ncgm.β*EV
            kpol[ikiz_index] = kp
        end

        # Update coefficients
        update_v!(vp, Φ \ V, 1.0)
        update_k!(vp, Φ \ kpol, 1.0)

        return V
    end


    function update!(V::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{VFI_EGM},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Get sizes and allocate for complete_polynomial
        kgrid = ncgm.kgrid; zgrid = ncgm.zgrid; grid = ncgm.grid;
        nk, nz = length(kgrid), length(zgrid)

        # Iterate
        temp = Array{Float64}(undef, n_complete(2, vp.d))
        for iz=1:nz, ik=1:nk

            # In EGM we use the grid points as if they were our
            # policy for yesterday and find implied kt
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            k1 = kgrid[ik];z = zgrid[iz];

            # Compute the derivative of expected values
            dEV = compute_dEV!(temp, ncgm, vp, k1, iz)

            # Compute optimal consumption
            c = duinv(ncgm, ncgm.β*dEV)

            # Need to find corresponding kt for optimal c
            obj(kt) = expendables_t(ncgm, kt, z) - c - k1
            kt_star = brent(obj, 0.0, 2.0, xtol=1e-10)

            # New value
            EV = compute_EV!(temp, ncgm, vp, k1, iz)
            V[ikiz_index] = u(ncgm, c) + ncgm.β*EV
            kpol[ikiz_index] = kt_star
        end

        # New Φ (has our new "kt_star" and z points)
        Φ_egm = complete_polynomial([kpol grid[:, 2]], vp.d)

        # Update coefficients
        update_v!(vp, Φ_egm \ V, 1.0)
        update_k!(vp, Φ_egm \ grid[:, 1], 1.0)

        # Update V and kpol to be value and policy corresponding
        # to our grid again
        copyto!(V, Φ*vp.v_coeffs)
        copyto!(kpol, Φ*vp.k_coeffs)

        return V
    end

    function env_condition_kp!(cp_out::Vector{Float64}, ncgm::NeoclassicalGrowth,
                               vp::ValueCoeffs, k::Float64, z::Float64)
        # Compute derivative of VF
        dV = dot(vp.v_coeffs, complete_polynomial!(cp_out, [k, z], vp.d, Derivative{1}()))

        # Consumption is then computed as
        c = duinv(ncgm, dV / (1 - ncgm.δ + df(ncgm, k, z)))

        return expendables_t(ncgm, k, z) - c
    end

    function env_condition_kp(ncgm::NeoclassicalGrowth, vp::ValueCoeffs,
                              k::Float64, z::Float64)
        cp_out = Array{Float64}(undef, n_complete(2, vp.d))
        env_condition_kp!(cp_out, ncgm, vp, k, z)
    end

    function env_condition_kp(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)
        # Pull out k and z from grid
        k = ncgm.grid[:, 1]
        z = ncgm.grid[:, 2]

        # Create basis matrix for entire grid
        dPhi = complete_polynomial(ncgm.grid, vp.d, Derivative{1}())

        # Compute consumption
        c = duinv(ncgm, (dPhi*vp.v_coeffs) ./ (1 .- ncgm.δ .+ df(ncgm, k, z)))

        return expendables_t(ncgm, k, z) .- c
    end

    function update!(V::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{VFI_ECM},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Get sizes and allocate for complete_polynomial
        kgrid = ncgm.kgrid; zgrid = ncgm.zgrid;
        nk, nz = length(kgrid), length(zgrid)

        # Iterate over all states
        temp = Array{Float64}(undef, n_complete(2, vp.d))
        for ik in 1:nk, iz in 1:nz
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            k = kgrid[ik]
            z = zgrid[iz]

            # Policy from envelope condition
            kp = env_condition_kp!(temp, ncgm, vp, k, z)
            c = expendables_t(ncgm, k, z) - kp
            kpol[ikiz_index] = kp

            # New value
            EV = compute_EV!(temp, ncgm, vp, kp, iz)
            V[ikiz_index] = u(ncgm, c) + ncgm.β*EV
        end

        # Update coefficients
        update_v!(vp, Φ \ V, 1.0)
        update_k!(vp, Φ \ kpol, 1.0)

        return V
    end



    function update!(V::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{PFI},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Get sizes and allocate for complete_polynomial
        kgrid = ncgm.kgrid; zgrid = ncgm.zgrid; grid = ncgm.grid;
        nk, nz = length(kgrid), length(zgrid)

        # Copy valuecoeffs object and use to iterate to
        # convergence given a policy
        vp_igp = copy(vp, IterateOnPolicy())
        solve(ncgm, vp_igp; nskipprint=1000, maxiter=5000, verbose=false)

        # Update the policy and values
        temp = Array{Float64}(undef, n_complete(2, vp.d))
        for ik in 1:nk, iz in 1:nz
            k = kgrid[ik]; z = zgrid[iz];

            # Define an objective function (negative for minimization)
            y = expendables_t(ncgm, k, z)
            solme(kp) = du(ncgm, y - kp) - ncgm.β*compute_dEV!(temp, ncgm, vp, kp, iz)

            # Find minimum of objective
            kp = brent(solme, 1e-8, y-1e-8; rtol=1e-12)

            # Update policy function
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            kpol[ikiz_index] = kp
        end

        # Get new coeffs
        update_k!(vp, Φ \ kpol, 1.0)
        update_v!(vp, vp_igp.v_coeffs, 1.0)

        # Update all elements of value
        copyto!(V, Φ*vp.v_coeffs)

        return V
    end


    function update!(V::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{PFI_ECM},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Copy valuecoeffs object and use to iterate to
        # convergence given a policy
        vp_igp = copy(vp, IterateOnPolicy())
        solve(ncgm, vp_igp; nskipprint=1000, maxiter=5000, verbose=false)

        # Update the policy and values
        kp = env_condition_kp(ncgm, vp)
        update_k!(vp, Φ \ kp, 1.0)
        update_v!(vp, vp_igp.v_coeffs, 1.0)

        # Update all elements of value
        copyto!(V, Φ*vp.v_coeffs)
        copyto!(kpol, kp)

        return V
    end


    function update!(dV::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{dVFI_ECM},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Get sizes and allocate for complete_polynomial
        kgrid = ncgm.kgrid; zgrid = ncgm.zgrid; grid = ncgm.grid;
        nk, nz, ns = length(kgrid), length(zgrid), size(grid, 1)

        # Iterate over all states
        temp = Array{Float64}(undef, n_complete(2, vp.d))
        for iz=1:nz, ik=1:nk
            k = kgrid[ik]; z = zgrid[iz];

            # Envelope condition implies optimal kp
            kp = env_condition_kp!(temp, ncgm, vp, k, z)
            c = expendables_t(ncgm, k, z) - kp

            # New value
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            dEV = compute_dEV!(temp, ncgm, vp, kp, iz)
            dV[ikiz_index] = (1-ncgm.δ+df(ncgm, k, z))*ncgm.β*dEV
            kpol[ikiz_index] = kp
        end

        # Get new coeffs
        update_k!(vp, Φ \ kpol, 1.0)
        update_v!(vp, dΦ \ dV, 1.0)

        return dV
    end

    # Conventional Euler equation method
    function update!(V::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{EulEq},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Get sizes and allocate for complete_polynomial
        @unpack kgrid, zgrid, weights, z1 = ncgm
        nz1, nz = size(z1)
        nk = length(kgrid)

        # Iterate over all states
        temp = Array{Float64}(undef, n_complete(2, vp.d))
        for iz in 1:nz, ik in 1:nk
            k = kgrid[ik]; z = zgrid[iz];

            # Create current polynomial
            complete_polynomial!(temp, [k, z], vp.d)

            # Compute what capital will be tomorrow according to policy
            kp = dot(temp, vp.k_coeffs)

            # Compute RHS of EE
            rhs_ee = 0.0
            for iz1 in 1:nz1
                # Possible z in t+1
                zp = z1[iz1, iz]

                # Policy for k_{t+2}
                complete_polynomial!(temp, [kp, zp], vp.d)
                kpp = dot(temp, vp.k_coeffs)

                # Implied t+1 consumption
                cp = expendables_t(ncgm, kp, zp) - kpp

                # Add to running expectation
                rhs_ee += ncgm.β*weights[iz1]*du(ncgm, cp)*(1-ncgm.δ+df(ncgm, kp, zp))
            end

            # The rhs of EE implies consumption and investment in t
            c = duinv(ncgm, rhs_ee)
            kp_star = expendables_t(ncgm, k, z) - c

            # New value
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            EV = compute_EV!(temp, ncgm, vp, kp_star, iz)
            V[ikiz_index] = u(ncgm, c) + ncgm.β*EV
            kpol[ikiz_index] = kp_star
        end

        # Update coefficients
        update_v!(vp, Φ \ V, 1.0)
        update_k!(vp, Φ \ kpol, 1.0)

        return V
    end

    function update!(dV::Vector{Float64}, kpol::Vector{Float64},
                     ncgm::NeoclassicalGrowth, vp::ValueCoeffs{dVFI_ECM},
                     Φ::Matrix{Float64}, dΦ::Matrix{Float64})
        # Get sizes and allocate for complete_polynomial
        kgrid = ncgm.kgrid; zgrid = ncgm.zgrid; grid = ncgm.grid;
        nk, nz, ns = length(kgrid), length(zgrid), size(grid, 1)

        # Iterate over all states
        temp = Array{Float64}(undef, n_complete(2, vp.d))
        for iz=1:nz, ik=1:nk
            k = kgrid[ik]; z = zgrid[iz];

            # Envelope condition implies optimal kp
            kp = env_condition_kp!(temp, ncgm, vp, k, z)
            c = expendables_t(ncgm, k, z) - kp

            # New value
            ikiz_index = LinearIndices((nk,nz))[(ik, iz)...]
            dEV = compute_dEV!(temp, ncgm, vp, kp, iz)
            dV[ikiz_index] = (1-ncgm.δ+df(ncgm, k, z))*ncgm.β*dEV
            kpol[ikiz_index] = kp
        end

        # Get new coeffs
        update_k!(vp, Φ \ kpol, 1.0)
        update_v!(vp, dΦ \ dV, 1.0)

        return dV
    end


    """
    Simulates the neoclassical growth model for a given set of solution
    coefficients. It simulates for `capT` periods and discards first
    `nburn` observations.
    """
    function simulate(ncgm::NeoclassicalGrowth, vp::ValueCoeffs,
                      shocks::Vector{Float64}; capT::Int=10_000,
                      nburn::Int=200)
        # Unpack parameters
        kp = 0.0  # Policy holder
        temp = Array{Float64}(undef, n_complete(2, vp.d))

        # Allocate space for k and z
        ksim = Array{Float64}(undef, capT+nburn)
        zsim = Array{Float64}(undef, capT+nburn)

        # Initialize both k and z at 1
        ksim[1] = 1.0
        zsim[1] = 1.0

        # Simulate
        temp = Array{Float64}(undef, n_complete(2, vp.d))
        for t in 2:capT+nburn
            # Evaluate k_t given yesterday's (k_{t-1}, z_{t-1})
            kp = env_condition_kp!(temp, ncgm, vp, ksim[t-1], zsim[t-1])

            # Draw new z and update k using policy above
            zsim[t] = zsim[t-1]^ncgm.ρ * exp(ncgm.σ*shocks[t])
            ksim[t] = kp
        end

        return ksim[nburn+1:end], zsim[nburn+1:end]
    end

    function simulate(ncgm::NeoclassicalGrowth, vp::ValueCoeffs;
                      capT::Int=10_000, nburn::Int=200, seed=42)
        srand(seed)  # Set specific seed
        shocks = randn(capT + nburn)

        return simulate(ncgm, vp, shocks; capT=capT, nburn=nburn)
    end

    """
    This function evaluates the Euler Equation residual for a single point (k, z)
    """
    function EulerEquation!(out::Vector{Float64}, ncgm::NeoclassicalGrowth,
                            vp::ValueCoeffs, k::Float64, z::Float64,
                            nodes::Vector{Float64}, weights::Vector{Float64})
        # Evaluate consumption today
        k1 = env_condition_kp!(out, ncgm, vp, k, z)
        c = expendables_t(ncgm, k, z) - k1
        LHS = du(ncgm, c)

        # For each of realizations tomorrow, evaluate expectation on RHS
        RHS = 0.0
        for (eps, w) in zip(nodes, weights)
            # Compute ztp1
            z1 = z^ncgm.ρ * exp(eps)

            # Evaluate the ktp2
            ktp2 = env_condition_kp!(out, ncgm, vp, k1, z1)

            # Get c1
            c1 = expendables_t(ncgm, k1, z1) - ktp2

            # Update RHS of equation
            RHS = RHS + w*du(ncgm, c1)*(1 - ncgm.δ + df(ncgm, k1, z1))
        end

        return abs(ncgm.β*RHS/LHS - 1.0)
    end

    """
    Given simulations for k and z, it computes the euler equation residuals
    along the entire simulation. It reports the mean and max values in
    log10.
    """
    function ee_residuals(ncgm::NeoclassicalGrowth, vp::ValueCoeffs,
                          ksim::Vector{Float64}, zsim::Vector{Float64}; Qn::Int=10)
        # Figure out how many periods we simulated for and make sure k and z
        # are same length
        capT = length(ksim)
        @assert length(zsim) == capT

        # Finer integration nodes
        eps_nodes, weight_nodes = qnwnorm(Qn, 0.0, ncgm.σ^2)
        temp = Array{Float64}(undef, n_complete(2, vp.d))

        # Compute EE for each period
        EE_resid = Array{Float64}(undef, capT)
        for t=1:capT
            # Pull out current state
            k, z = ksim[t], zsim[t]

            # Compute residual of Euler Equation
            EE_resid[t] = EulerEquation!(temp, ncgm, vp, k, z, eps_nodes, weight_nodes)
        end

        return EE_resid
    end

    function ee_residuals(ncgm::NeoclassicalGrowth, vp::ValueCoeffs; Qn::Int=10)
        # Simulate and then call other ee_residuals method
        ksim, zsim = simulate(ncgm, vp)

        return ee_residuals(ncgm, vp, ksim, zsim; Qn=Qn)
    end

## main
    function main(sm::SolutionMethod, nd::Int=5, shocks=randn(capT+nburn);
                  capT=10_000, nburn=200, tol=1e-9, maxiter=2500,
                  nskipprint=25, verbose=true)
        # Create model
        ncgm = NeoclassicalGrowth()

        # Create initial quadratic guess
        vp = ValueCoeffs(ncgm, Val{2}, IterateOnPolicy())
        solve(ncgm, vp, tol=1e-6, verbose=true)
        # solve(ncgm, 1.0)

        # Allocate memory for timings
        times = Array{Float64}(undef,nd-1)
        sols = Array{ValueCoeffs}(undef,nd-1)
        mean_ees = Array{Float64}(undef,nd-1)
        max_ees = Array{Float64}(undef,nd-1)

        # Solve using the solution method for degree 2 to 5
        vp = copy(vp, sm)
        for d in 2:nd
            # Change degree of solution method
            vp = copy(ncgm, vp, Val{d})

            # Time the current method
            start_time = time()
            solve(ncgm, vp; tol=tol, maxiter=maxiter, nskipprint=nskipprint,
                  verbose=verbose)
            end_time = time()

            # Save the time and solution
            times[d-1] = end_time - start_time
            sols[d-1] = vp

            # Simulate and compute EE
            ks, zs = simulate(ncgm, vp, shocks; capT=capT, nburn=nburn)
            resids = ee_residuals(ncgm, vp, ks, zs; Qn=10)
            mean_ees[d-1] = log10.(mean(abs.(resids)))
            max_ees[d-1] = log10.(maximum(abs, resids))
        end

        return sols, times, mean_ees, max_ees
    end

    
    Random.seed!(52)
    shocks = randn(10200)
    
    for sol_method in [VFI(), VFI_EGM(), VFI_ECM(), dVFI_ECM(),
                       PFI(), PFI_ECM(), EulEq()]
        # Make sure everything is compiled
        main(sol_method, 5, shocks; maxiter=2, verbose=false)
    
        # Run for real
        s_sm, t_sm, mean_eem, max_eem = main(sol_method, 5, shocks;
                                             tol=1e-8, verbose=false)
    
        println("Solution Method: $sol_method")
        for (d, t) in zip([2, 3, 4, 5], t_sm)
            println("\tDegree $d took time $t")
            println("\tMean & Max EE are" *
                    "$(round(mean_eem[d-1], digits = 3)) & $(round(max_eem[d-1], digits =  3))")
        end
    end