text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

module DiskResourceProviderClient

using Random
using Dates
using Swagger
import Swagger: field_name, property_type, hasproperty, validate_property, SwaggerApi, SwaggerModel
import Base: convert, propertynames

include("modelincludes.jl")

include("api_DiskAccessesApi.jl")
include("api_DiskEncryptionSetsApi.jl")
include("api_DisksApi.jl")
include("api_SnapshotsApi.jl")

# export models
export convert, AccessUri
export convert, ApiError
export convert, ApiErrorBase
export convert, CloudError
export convert, CreationData
export convert, Disk
export convert, DiskAccess
export convert, DiskAccessList
export convert, DiskAccessProperties
export convert, DiskAccessUpdate
export convert, DiskEncryptionSet
export convert, DiskEncryptionSetList
export convert, DiskEncryptionSetUpdate
export convert, DiskEncryptionSetUpdateProperties
export convert, DiskList
export convert, DiskProperties
export convert, DiskSku
export convert, DiskUpdate
export convert, DiskUpdateProperties
export convert, Encryption
export convert, EncryptionSetIdentity
export convert, EncryptionSetProperties
export convert, EncryptionSettingsCollection
export convert, EncryptionSettingsElement
export convert, EncryptionType
export convert, GrantAccessData
export convert, ImageDiskReference
export convert, InnerError
export convert, KeyVaultAndKeyReference
export convert, KeyVaultAndSecretReference
export convert, NetworkAccessPolicy
export convert, PrivateEndpoint
export convert, PrivateEndpointConnection
export convert, PrivateEndpointConnectionProperties
export convert, PrivateEndpointConnectionProvisioningState
export convert, PrivateEndpointServiceConnectionStatus
export convert, PrivateLinkResource
export convert, PrivateLinkResourceListResult
export convert, PrivateLinkResourceProperties
export convert, PrivateLinkServiceConnectionState
export convert, Resource
export convert, ShareInfoElement
export convert, Snapshot
export convert, SnapshotList
export convert, SnapshotProperties
export convert, SnapshotSku
export convert, SnapshotUpdate
export convert, SnapshotUpdateProperties
export convert, SourceVault

# export operations
export convert, DiskAccessesApi, DiskEncryptionSetsApi, DisksApi, SnapshotsApi

export check_required, field_name, property_type, hasproperty, propertynames, validate_property, convert

end


text:
using Base.Test
import SensitiveFloats: zero_sensitive_float_array

function test_sky_noise_estimates()
    images_vec = Array(Vector{Image}, 2)
    images_vec[1], ea, three_bodies = gen_three_body_dataset()  # synthetic
    images_vec[2] = SampleData.load_stamp_blob(datadir, "164.4311-39.0359_2kpsf")  # real

    for images in images_vec
        for b in 1:5
            sdss_sky_estimate = median(images[b].epsilon_mat) * median(images[b].iota_vec)
            crude_estimate = median(images[b].pixels)
            @test_approx_eq_eps sdss_sky_estimate / crude_estimate 1. .3
        end
    end
end

function test_zero_sensitive_float_array()
    S = 3
    sf_vec = zero_sensitive_float_array(Float64, length(ids), S, 3, 5)
    @test size(sf_vec) == (3, 5)
    for sf in sf_vec
        @test sf.v[] == 0
        @test all([ all(sf_d .== 0) for sf_d in sf.d ])
        @test all(sf.h .== 0)
    end

    # Test that the sensitive floats are distinct and not pointers to the
    # same sensitive float.
    sf_vec[1].v[] = rand()
    sf_vec[1].d[:, 1] = rand(length(CanonicalParams))
    sf_vec[1].h[:, :] = 2 * diagm(ones(size(sf_vec[1].h, 1)))

    for ind in 2:length(sf_vec)
        sf = sf_vec[ind]
        @test sf.v[] == 0
        @test all([ all(sf_d .== 0) for sf_d in sf.d ])
        @test all(sf.h .== 0)
    end
end

####################################################

test_sky_noise_estimates()
test_zero_sensitive_float_array()

ext: jl


text:
using Fractal
using Test

@testset "Fractal.jl" begin
    # set the parameters
    linearIFS = [[0  0; 0 0.16],
    [0.85  0.04; -0.04 0.85 ],
    [0.2  -0.26; 0.23 0.22],
    [-0.15  0.28; 0.26 0.24]]
    transIFS = [[0, 0],
    [0, 1.6],
    [0, 1.6],
    [0, 0.44]]
    weights = [0.01, 0.84, 0.08, 0.07]
    myWIFS = WIFS(linearIFS, transIFS, weights)

    @test typeof(itrPtsProb(myWIFS)) == Vector{Vector{Float64}}
end


text:
@info "$("-"^70)\nStarting sandwich shop demo\n$("-"^70)"

using Dates              # for actual time & date
using Distributions      # for distributions and random behaviour
using HypothesisTests    # for more statistical analysis
using Logging            # for debugging
using Plots              # for figures
using SimJulia           # for DES
using StatsPlots         # for nicer histograms
using Statistics         # for statistics

# import Base.show in order to use it for our own types
import Base.show

# logging settings
logger = Logging.ConsoleLogger(stdout,LogLevel(-3000))
Logging.global_logger(logger)
Logging.disable_logging(LogLevel(-1000))
@info "$("-"^70)\nLoaded dependencies\n$("-"^70)"

# ------------------------------------------------ #
#               Simulation core code               #
# ------------------------------------------------ #

# Available items in the shop
const menulist = ["sandwich, cold","sandwich, hot","pasta","soup"]
# Historical data product demand
const menuprob = Distributions.Categorical([0.6,0.2,0.1,0.1])
# Preparation times [s]
const preptimes = Dict( menulist[1]=>Distributions.Uniform(60,90),
                        menulist[2]=>Distributions.Uniform(60,120),
                        menulist[3]=>Distributions.Uniform(60,90),
                        menulist[4]=>Distributions.Uniform(30,45))
# Client arrival rates [s]
const arrivals = Dict("[08-11[" => Distributions.Exponential(25*60),
                        "[11-14[" => Distributions.Exponential(1*60),
                        "[14-17[" => Distributions.Exponential(10*60),
                        "[17-19[" => Distributions.Exponential(2*60),
                        "[19-20[" => Distributions.Exponential(5*60))
# Time required to choose something [s]
const choicetime = Distributions.Uniform(10,30)
# Arrival time function [s]
function nextarrival(t::DateTime; arrivals::Dict=arrivals)
    if (hour(t) >= 8) & (hour(t) < 11)
        return Second(round(rand(arrivals["[08-11["])))
    elseif (hour(t) >= 11) & (hour(t) < 14)
        return Second(round(rand(arrivals["[11-14["])))
    elseif (hour(t) >= 14) & (hour(t) < 17)
        return Second(round(rand(arrivals["[14-17["])))
    elseif (hour(t) >= 17) & (hour(t) < 19)
        return Second(round(rand(arrivals["[17-19["])))
    elseif (hour(t) >= 19) & (hour(t) < 20)
        return Second(round(rand(arrivals["[19-20["])))
    else
        return nothing
    end
end

# define a shop
struct Shop
    staff::Resource
    queuelength::Array{Tuple{DateTime,Int64},1}
    clients::Dict
    waitingtimes::Array{Millisecond, 1}
    renegtimes::Array{DateTime,1}
    function Shop(env::Environment, nstaff::Int=1)
        # add the crew
        staff = Resource(env,nstaff)
        # no queue at the start of the simulation
        queuelength = [(nowDatetime(env),0)]
        # keep record of each client, for possible extensions such as:
        #   - a free sandwich after having bought a specific number
        #   - including a feeling of disconent after leaving the queue
        clients = Dict()
        # client waiting times
        waitingtimes = Array{Millisecond,1}()
        # moments when a client leaves due to losing his patience
        renegtimes = Array{DateTime,1}()
        return new(staff,queuelength,clients,waitingtimes,renegtimes)
    end
end

function show(io::IO,s::Shop)
    print(io,"Shop with $(s.staff.capacity) employees, currently $(length(s.staff.put_queue)) people waiting")
end

# define a client
mutable struct Client
    id::Int
    patience::Second
    orderhist::Dict
    proc::Process
    function Client(env::Environment,shop::Shop)
        client = new()
        client.id = length(shop.clients) + 1
        client.patience = Second(round(rand(Distributions.Uniform(5*60,10*60))) )
        client.orderhist = Dict()
        # add client to shop registry
        shop.clients[client.id]= client
        # start the client process
        client.proc = @process clientbehavior(env, shop, client)
        return client
    end
end

function show(io::IO,c::Client)
    print(io,"Client N° $(c.id) with $(round(c.patience.value / 60,digits=2)) minutes of patience")
end

# Generating function for clients.
@resumable function clientgenerator(env::Environment, shop::Shop, topen::Int=8, tclose::Int=20)
    while true
        # before opening hours => wait
        if hour(nowDatetime(env)) < topen
            delta = floor(nowDatetime(env), Day) + Hour(topen) - nowDatetime(env)
            @logmsg LogLevel(-3000) "$(nowDatetime(env)) - Not open yet, waiting for $(delta)"
            @yield timeout(env, delta)
        end

        # after opening hours => wait
        if hour(nowDatetime(env)) >= tclose
            delta = floor(nowDatetime(env), Day) + Day(1) + Hour(topen) - nowDatetime(env)
            @logmsg LogLevel(-3000) "$(nowDatetime(env)) - Closing shop, waiting for $(delta)"
            @yield timeout(env, delta)
        end

        # during opening hours => generate clients
        tnext = nextarrival(nowDatetime(env))
        @yield timeout(env, tnext)
        c = Client(env,shop)
    end
end

@resumable function clientbehavior(env::Environment, s::Shop, c::Client)
    @debug "$(nowDatetime(env)) - Client (N° $(c.id) arrives (Client patience: $(c.patience))"
    # choose what to get
    @yield timeout(env, Second(round(rand(choicetime))))
    choice = menulist[rand(menuprob)]
    # ready to order
    tin = nowDatetime(env)
    # try to obtain a staff member
    req = request(s.staff)
    res  = @yield req | timeout(env, c.patience)
    # log of current queuelength
    push!(s.queuelength, (nowDatetime(env), length(s.staff.put_queue) ) )

    if res[req].state == SimJulia.processed
        @debug "$(nowDatetime(env)) - Client N° $(c.id) is being served and orders a $(choice)"
        tserved = nowDatetime(env)
        twait = tserved - tin # in milliseconds
        push!(s.waitingtimes, twait)
        @yield timeout(env, Second(round(rand(preptimes[choice])))) # preparing order
        @debug "$(nowDatetime(env)) - Client N° $(c.id) receives order"
        @yield release(s.staff) # release staff
    else
        @debug "$(nowDatetime(env)) - Client N° $(c.id) ran out of patience"
        if haskey(s.staff.put_queue, req)
            cancel(s.staff, req)                   # cancel staff request
        end
        push!(s.renegtimes, nowDatetime(env))  # store renegtimes
    end
end

@info "$("-"^70)\nLoaded functions\n$("-"^70)"

# ------------------------------------------------ #
#               helper functions                   #
# ------------------------------------------------ #

"""
    plotqueue(s::Shop)

Store an illustration of the queuelength over time from a `::Store`
"""
function plotqueue(s::Shop)
    # some makeup
    tstart = floor(s.queuelength[1][1], Day) + Hour(8)
    tstop  = floor(s.queuelength[1][1], Day) + Hour(23)
    daterange =  tstart : Minute(60) : tstop
    datexticks = [Dates.value(mom) for mom in daterange]
    datexticklabels = Dates.format.(daterange,"HH:MM")
    # queue length
    x::Array{DateTime,1} = map(v -> v[1], s.queuelength)
    y::Array{Int,1} = map(v -> v[2], s.queuelength)
    p1 = plot(x, y, linetype=:steppost, label="Queue length")
    xticks!(datexticks, datexticklabels,rotation=0)
    xlims!(Dates.value(tstart),Dates.value(tstop))
    yticks!(0:1:maximum(y)+1)
    # reneg moments
    nreneg = 1:length(s.renegtimes)
    p2 = plot(s.renegtimes, collect(nreneg), linetype=:steppost, label="Number of renegs")
    xticks!(datexticks, datexticklabels,rotation=0)
    xlims!(Dates.value(tstart),Dates.value(tstop))
    yticks!(0:5:length(s.renegtimes)+1)
    # global figure
    p = plot(p1,p2, layout=(2,1), size=(1000,800))
    savefig(p, "queuelength.png")
end

"""
    overviewplot(Nreneg::Array{Int64,1}, MWT::Array{Float64,1},
MQL::Array{Float64,1}, fname::String)

function to make an overview plot that includes a histogram of
the amount of renegs, the mean waiting time and the mean queue
length.

"""
function overviewplot(Nreneg::Array{Int64,1}, MWT::Array{Float64,1},
                      MQL::Array{Float64,1}, fname::String)
    @info "making overviewplot $(fname)"
    p1 = StatsPlots.histogram(Nreneg,bins=0:2:80,normalize=:probability,
                              title="Reneg distribution", label="",
                              xlabel="Amount of renegs",
                              ylabel="Experimental probability")
    p2 = StatsPlots.histogram(MWT,bins=20,normalize=:probability,
                              title="MWT distribution", label="",
                              xlabel="MWT [s]",
                              ylabel="Experimental probability")
    p3 = StatsPlots.histogram(MQL,bins=20,normalize=:probability,
                              title="MWT distribution", label="",
                              xlabel="MQL",
                              ylabel="Experimental probability")
    p = plot(p1, p2, p3, layout=(1,3), size=(1200, 500))
    savefig(p, fname)
    return
end

# ------------------------------------------------ #
#            Simulation configuration              #
# ------------------------------------------------ #

"""
    test()

Run a simple simulation for one day without staff. Conclusions:
    - a client can arrive after closing time
    - when a client arrives after closing time, he still waits until he runs out of patience
"""
function test()
    @info "$("-"^70)\nStarting a store without staff\n$("-"^70)"
    # Start a simulation on today 00Hr00
    sim = Simulation(floor(Dates.now(),Day))
    s = Shop(sim, 0)
    @process clientgenerator(sim, s)
    # Run the sim for one day
    run(sim, floor(Dates.now(),Day) + Day(1))
end

test()

"""
    runsim()

From our first simulation and the generated figure we can clearly see that there is a problem
serving the clients fast enough during the midday peak hours.

"""
function runsim()
    @info "$("-"^70)\nStarting a complete simulation\n$("-"^70)"
    # Start a simulation on today 00Hr00
    sim = Simulation(floor(Dates.now(),Day))
    s = Shop(sim, 1)
    @process clientgenerator(sim, s)
    # Run the sim for one day
    run(sim, floor(Dates.now(),Day) + Day(1))
    # Make an illustration of the queue length
    @info "Making queue length figure"
    plotqueue(s)
end

runsim()

"""
    multisim(;n::Int=100, staff::Int=1,tstart=floor(now(),Day), duration::Period=Day(1))

Run `n` simulations with shop staffed with `staff` persons. By default start on current day and
runs for one day.

Returns a vector of the number of renegs, the Mean Waiting Times (MWT) and Mean Queue Length (MQL)
"""
function multisim(;n::Int=100, staff::Int=1,
                  tstart::DateTime=floor(now(),Day),
                  duration::Period=Day(1), plt::Bool=false)
    @info "<multisim>: Running a multisim simulation on $(Threads.nthreads()) threads"
    # pre-allocate results
    Nreneg = Array{Int64,1}(undef, n)
    MWT = Array{Float64,1}(undef, n)
    MQL = Array{Float64,1}(undef, n)
    # run all simulations (in parallel where available)
    Threads.@threads for i in 1:n
        sim = Simulation(tstart)
        shop = Shop(sim, staff)
        @process clientgenerator(sim,shop)
        run(sim, tstart + duration)
        Nreneg[i] = length(shop.renegtimes)
        MWT[i] = mean(Dates.value.(shop.waitingtimes)/1e3) # in seconds
        MQL[i] = mean([x[2] for x in shop.queuelength])
    end
    # generate a nice illustration
    if plt
        #@info Nreneg, MWT, MQL
        fname = "multisim - $(n) iterations - $(staff) staff.png"
        overviewplot(Nreneg, MWT, MQL, fname)
    end

    return Nreneg, MWT, MQL
end

# run a single multisim in different configurations
multisim(plt=true)
multisim(staff=2,plt=true)

# Determine when you have enough sample data
"""
    determinesamplesize()

Make a boxplot of Nreneg, MWT, MQL for different sample sizes.

When looking at the boxplot, you see the boxes stay fixed starting from
a sample length of around 1000. Of course the amount of outlier will increase
for an increasing sample length.
"""
function determinesamplesize(;maxpow::Int=3, fname::String="determinesamplesize.png")
    @info "<determinesamplesize>: Running determinesamplesize $(fname)"
    r_reneg = []
    r_MWT = []
    r_MQL = []
    labels = collect(10 .^ (1:maxpow))

    for n in labels
        @info "\t<determinesamplesize>: working on sample length $(n)"
        (Nreneg, MWT, MQL) = multisim(n, 1);
        push!(r_reneg, Nreneg)
        push!(r_MWT, MWT)
        push!(r_MQL, MQL)
    end

    @info "\t<determinesamplesize>: rendering plot"
    p = plot(
        boxplot(r_reneg, label="", title="Renegs"),
        boxplot(r_MWT, label="", title="MWT"),
        boxplot(r_MQL, label="", title="MQL"),
        layout=(1,3), size=(1200, 500)
    )
    xticks!(collect(1:maxpow), ["$(n)" for n in labels])
    xlabel!("Sample length")
    savefig(p, fname)

end

#determinesamplesize(maxpow=4)

# Evaluate normality and required sample lengths
function evalnorm(;α::Float64=0.05)
    @info "running normality test"
    # use MWT for this example
    _, MWT, _ = multisim()
    # run K-S test
    res = ExactOneSampleKSTest(MWT, Normal(mean(MWT), std(MWT)))
    if pvalue(res) > α
        # make plot for CI
        @info "normality respected, making sample size plot"
        # basis stat formulas
        E = range(0.25, stop=20, length=50)
        n = (std(MWT)*quantile(Normal(),1 - 0.05/2) ./ E ) .^2
        p = plot(E,n,marker=:circle,label="", yscale=:log10)
        xlabel!("precision on MWT [s]")
        ylabel!("Required sample size")
        savefig(p,"samplesize.png")
    else
        @info "data not normal"
    end
    return res
end

evalnorm()

# verify statistical difference between two staffings
function comparestaffing(;α::Float64=0.05)
    @info "comparing two populations"
    _, MWT_1, _ = multisim()
    _, MWT_2, _ = multisim(staff=2)
    @info "  checking equality of variances"
    eqvar = VarianceFTest(MWT_1, MWT_2)
    if pvalue(eqvar) > α
        @info "  assuming equal variance"
        r = EqualVarianceTTest(MWT_1, MWT_2)
    else
        @info "  assuming unequal variance"
        r = UnequalVarianceTTest(MWT_1, MWT_2)
    end
    @info """  Conclusion: $(ifelse(pvalue(r) > α, "Equal population mean",
                            "Unequal population mean"))"""
end

comparestaffing()


text:
@testset "add_agent! (discrete)" begin
    properties = Dict(:x1 => 1)
    space = GraphSpace(complete_digraph(10))
    model = ABM(Agent7, space; properties = properties)
    attributes = (f1 = true, f2 = 1)
    add_agent!(1, model, attributes...)
    attributes = (f2 = 1, f1 = true)
    add_agent!(1, model; attributes...)
    @test model.agents[1].id != model.agents[2].id
    @test model.agents[1].pos == model.agents[2].pos
    @test model.agents[1].f1 == model.agents[2].f1
    @test model.agents[1].f2 == model.agents[2].f2
    @test add_agent_single!(model, attributes...).pos ∈ 1:10
    for id in 4:11
        agent = Agent7(id, 2, attributes...)
        add_agent_single!(agent, model)
    end
    @test !has_empty_positions(model)
    agent = Agent7(12, 5, attributes...)
    add_agent_single!(agent, model)
    @test_throws KeyError model[12]
    add_agent!(agent, model)
    @test model[12].pos ∈ 1:10

    agent = Agent7(13, 5, attributes...)
    @test add_agent!(agent, 3, model).pos == 3

    model = ABM(Agent1, GridSpace((10, 10)))
    agent = Agent1(1, (3, 6))
    @test add_agent!(agent, (7, 8), model).pos == (7, 8)
end

@testset "add_agent! (continuous)" begin
    properties = Dict(:x1 => 1)
    space2d = ContinuousSpace((1, 1), 0.1; periodic = true)
    model = ABM(Agent8, space2d; properties = properties)
    attributes = (f1 = true, f2 = 1)
    add_agent!(model, attributes...)
    attributes = (f2 = 1, f1 = true)
    add_agent!(model; attributes...)
    @test model.agents[1].id != model.agents[2].id
    @test model.agents[1].f1 == model.agents[2].f1
    @test model.agents[1].f2 == model.agents[2].f2
    agent = Agent8(3, (0, 0), false, 6)
    @test 0 <= add_agent!(agent, model).pos[1] <= 1
    agent.id = 4
    @test add_agent!(agent, (0.5, 0.5), model).pos[1] ≈ 0.5 atol = 1e-3
end

@testset "move_agent!" begin
    # GraphSpace
    model = ABM(Agent5, GraphSpace(path_graph(6)))
    agent = add_agent!(model, 5.3)
    init_pos = agent.pos
    # Checking specific indexing
    move_agent!(agent, rand(model.rng, [i for i in 1:6 if i != init_pos]), model)
    new_pos = agent.pos
    @test new_pos != init_pos
    # Checking a random move
    ni = 0
    init_pos = agent.pos
    while agent.pos == init_pos
        move_agent!(agent, model)
    end
    @test ni < Inf

    # GridSpace
    model = ABM(Agent1, GridSpace((5, 5)))
    agent = add_agent!((2, 4), model)
    move_agent!(agent, (1, 3), model)
    @test agent.pos == (1, 3)
    ni = 0
    init_pos = agent.pos
    while agent.pos == init_pos
        move_agent!(agent, model)
    end
    @test ni < Inf

    model = ABM(Agent1, GridSpace((2, 1)))
    agent = add_agent!((1, 1), model)
    move_agent_single!(agent, model)
    @test agent.pos == (2, 1)
    agent2 = add_agent!((1, 1), model)
    move_agent_single!(agent2, model)
    # Agent shouldn't move since the grid is saturated
    @test agent2.pos == (1, 1)

    # ContinuousSpace
    model = ABM(Agent6, ContinuousSpace((1, 1), 0.1))
    agent = add_agent!((0.0, 0.0), model, (0.5, 0.0), 1.0)
    move_agent!(agent, model)
    @test agent.pos == (0.5, 0.0)
end

@testset "kill_agent!" begin
    # No Space
    model = ABM(Agent0)
    add_agent!(model)
    agent = add_agent!(model)
    @test nagents(model) == 2
    kill_agent!(agent, model)
    @test nagents(model) == 1
    add_agent!(model)
    genocide!(model, [1, 3])
    @test nagents(model) == 0
    # GraphSpace
    model = ABM(Agent5, GraphSpace(path_graph(6)))
    add_agent!(model, 5.3)
    add_agent!(model, 2.7)
    @test nagents(model) == 2
    kill_agent!(model.agents[1], model)
    @test nagents(model) == 1
    kill_agent!(2, model)
    @test nagents(model) == 0
    # GridSpace
    model = ABM(Agent1, GridSpace((5, 5)))
    add_agent!((1, 3), model)
    add_agent!((1, 3), model)
    add_agent!((5, 2), model)
    @test nagents(model) == 3
    for id in copy(ids_in_position((1, 3), model))
        kill_agent!(id, model)
    end
    @test nagents(model) == 1
    # ContinuousSpace
    model = ABM(Agent6, ContinuousSpace((1, 1), 0.1))
    add_agent!((0.7, 0.1), model, (15, 20), 5.0)
    add_agent!((0.2, 0.9), model, (8, 35), 1.7)
    @test nagents(model) == 2
    kill_agent!(model[1], model)
    @test nagents(model) == 1
end

@testset "genocide!" begin
    # Testing no space
    model = ABM(Agent0)
    for i in 1:10
        a = Agent0(i)
        add_agent!(a, model)
    end
    genocide!(model)
    @test nagents(model) == 0
    for i in 1:10
        a = Agent0(i)
        add_agent!(a, model)
    end
    genocide!(model, 5)
    @test nagents(model) == 5
    genocide!(model, a -> a.id < 3)
    @test nagents(model) == 3

    model = ABM(Agent3, GridSpace((10, 10)))

    # Testing genocide!(model::ABM)
    for i in 1:20
        agent = Agent3(i, (1, 1), rand(model.rng))
        add_agent_single!(agent, model)
    end
    genocide!(model)
    @test nagents(model) == 0

    # Testing genocide!(model::ABM, n::Int)
    for i in 1:20
        # Explicitly override agents each time we replenish the population,
        # so we always start the genocide with 20 agents.
        agent = Agent3(i, (1, 1), rand(model.rng))
        add_agent_single!(agent, model)
    end
    genocide!(model, 10)
    @test nagents(model) == 10

    # Testing genocide!(model::ABM, f::Function) with an anonymous function
    for i in 1:20
        agent = Agent3(i, (1, 1), rand(model.rng))
        add_agent_single!(agent, model)
    end
    @test nagents(model) == 20
    genocide!(model, a -> a.id > 5)
    @test nagents(model) == 5

    Random.seed!(6465)
    # Testing genocide!(model::ABM, f::Function) when the function is invalid
    # (i.e. does not return a bool)
    for i in 1:20
        agent = Agent3(i, (rand(model.rng, 1:10), rand(model.rng, 1:10)), i * 2)
        add_agent_pos!(agent, model)
    end
    @test_throws TypeError genocide!(model, a -> a.id)
    N = nagents(model)

    # Testing genocide!(model::ABM, f::Function) with a named function
    # No need to replenish population since the last test fails
    function complex_logic(agent::A) where {A<:AbstractAgent}
        if agent.pos[1] <= 5 && agent.weight > 25
            true
        else
            false
        end
    end
    genocide!(model, complex_logic)
    @test nagents(model) < N

    space2d = ContinuousSpace((1, 1), 0.1; periodic = true)
    model = ABM(Agent8, space2d)
    attributes = (f1 = true, f2 = 1)
    for _ in 1:10
        add_agent!(model, attributes...)
    end
    genocide!(model)
    @test nagents(model) == 0
    for _ in 1:10
        add_agent!(model, attributes...)
    end
    genocide!(model, 5)
    @test nagents(model) == 5
    genocide!(model, a -> a.id < 3)
    @test nagents(model) == 3
end

mutable struct Daisy <: AbstractAgent
    id::Int
    pos::Dims{2}
    breed::String
end
mutable struct Land <: AbstractAgent
    id::Int
    pos::Dims{2}
    temperature::Float64
end
@testset "fill space" begin
    space = GridSpace((10, 10))
    model = ABM(Land, space)
    fill_space!(model, 15)
    @test nagents(model) == 100
    for a in allagents(model)
        @test a isa Land
        @test a.temperature == 15
    end

    space = GridSpace((10, 10))
    model = ABM(Union{Daisy,Land}, space; warn = false)
    fill_space!(Daisy, model, "black")
    @test nagents(model) == 100
    for a in values(model.agents)
        @test a isa Daisy
        @test a.breed == "black"
    end

    space = GridSpace((10, 10), periodic = true)
    model = ABM(Union{Daisy,Land}, space; warn = false)
    temperature(pos) = (pos[1] / 10,) # make it Tuple!
    fill_space!(Land, model, temperature)
    @test nagents(model) == 100
    for a in values(model.agents)
        @test a.temperature == a.pos[1] / 10
    end

end

@testset "model step order" begin
    function model_step!(model)
        for a in allagents(model)
            if a.weight > 1.0
                model.count += 1
            end
        end
    end
    function agent_step!(a, model)
        a.weight += 1
    end

    for bool in (true, false)
        model = ABM(Agent2; properties = Dict(:count => 0))
        for i in 1:100
            add_agent!(model, rand(model.rng))
        end
        step!(model, agent_step!, model_step!, 1, bool)
        if bool
            @test model.count == 100
        else
            @test model.count == 0
        end
    end
end

@testset "Higher order groups" begin
    model = ABM(Agent3, GridSpace((10, 10)))
    for i in 1:10
        add_agent!(model, i)
    end

    iter_second_ids = map(x -> (x[1].id, x[2].id), iter_agent_groups(2, model))
    @test size(iter_second_ids) == (10, 10)
    @test iter_second_ids[1] == (1, 1)
    @test iter_second_ids[15] == (5, 2)
    @test iter_second_ids[end] == (10, 10)

    second = collect(map_agent_groups(2, x -> x[1].weight + x[2].weight, model))
    @test size(second) == (10, 10)
    @test second[1] == 2.0
    @test second[15] == 7.0
    @test second[end] == 20.0

    third =
        collect(map_agent_groups(3, x -> x[1].weight + x[2].weight + x[3].weight, model))
    @test size(third) == (10, 10, 10)
    @test third[1] == 3.0
    @test third[15] == 8.0
    @test third[end] == 30.0

    second_filtered =
        collect(map_agent_groups(2, x -> x[1].weight + x[2].weight, model, allunique))
    @test size(second_filtered) == (90,)
    @test second_filtered[1] == 3.0
    @test second_filtered[15] == 9.0
    @test second_filtered[end] == 19.0

    idx_second_filtered = collect(index_mapped_groups(2, model, allunique))
    @test size(idx_second_filtered) == (90,)
    @test idx_second_filtered[1] == (2, 1)
    @test idx_second_filtered[15] == (7, 2)
    @test idx_second_filtered[end] == (9, 10)
end


text:
using MPI
using CLIMA
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.DGBalanceLawDiscretizations
using Printf
using LinearAlgebra
using Logging

@inline function auxiliary_state_initialization!(aux, x, y, z, dim)
    @inbounds begin
        if dim == 2
            aux[1] = x^2 + y^3 - x * y
            aux[5] = 2 * x - y
            aux[6] = 3 * y^2 - x
            aux[7] = 0
        else
            aux[1] = x^2 + y^3 + z^2 * y^2 - x * y * z
            aux[5] = 2 * x - y * z
            aux[6] = 3 * y^2 + 2 * z^2 * y - x * z
            aux[7] = 2 * z * y^2 - x * y
        end
    end
end

using Test
function run(mpicomm, dim, ArrayType, Ne, N, FT)

    brickrange = ntuple(j -> range(FT(-1); length = Ne[j] + 1, stop = 1), dim)
    topl =
        BrickTopology(mpicomm, brickrange, periodicity = ntuple(j -> true, dim))

    grid = DiscontinuousSpectralElementGrid(
        topl,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = N,
    )

    spacedisc = DGBalanceLaw(
        grid = grid,
        length_state_vector = 0,
        flux! = (x...) -> (),
        numerical_flux! = (x...) -> (),
        auxiliary_state_length = 7,
        auxiliary_state_initialization! = (x...) ->
            auxiliary_state_initialization!(x..., dim),
    )

    DGBalanceLawDiscretizations.grad_auxiliary_state!(spacedisc, 1, (2, 3, 4))

    # Wrapping in Array ensure both GPU and CPU code use same approx
    @test Array(spacedisc.auxstate.data[:, 2, :]) ≈
          Array(spacedisc.auxstate.data[:, 5, :])
    @test Array(spacedisc.auxstate.data[:, 3, :]) ≈
          Array(spacedisc.auxstate.data[:, 6, :])
    @test Array(spacedisc.auxstate.data[:, 4, :]) ≈
          Array(spacedisc.auxstate.data[:, 7, :])
end

let
    CLIMA.init()
    ArrayTypes = (CLIMA.array_type(),)

    mpicomm = MPI.COMM_WORLD

    numelem = (5, 5, 1)
    lvls = 1

    polynomialorder = 4

    @testset "$(@__FILE__)" for ArrayType in ArrayTypes
        for FT in (Float64,) #Float32)
            for dim in 2:3
                err = zeros(FT, lvls)
                for l in 1:lvls
                    @info (ArrayType, FT, dim)
                    run(
                        mpicomm,
                        dim,
                        ArrayType,
                        ntuple(j -> 2^(l - 1) * numelem[j], dim),
                        polynomialorder,
                        FT,
                    )
                end
            end
        end
    end
end

nothing

max_stars_repo_name: yairchn/CLIMA

max_issues_repo_name: yairchn/CLIMA

max_forks_repo_name: yairchn/CLIMA


text:
# The contents of this file are slightly modified from InvertedIndices.jl: https://github.com/mbauman/InvertedIndices.jl/blob/7c97b19e6c75043f231f9d8e6661886b688259f1/src/InvertedIndices.jl
# Available under the following MIT "Expat" License:
#
# > Copyright (c) 2017: Matt Bauman.
# >
# > Permission is hereby granted, free of charge, to any person obtaining a copy
# > of this software and associated documentation files (the "Software"), to deal
# > in the Software without restriction, including without limitation the rights
# > to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# > copies of the Software, and to permit persons to whom the Software is
# > furnished to do so, subject to the following conditions:
# >
# > The above copyright notice and this permission notice shall be included in all
# > copies or substantial portions of the Software.
# >
# > THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# > IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# > FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# > AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# > LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# > OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# > SOFTWARE.
# >

# A very simple and primitive static array to avoid allocations while fulfilling the indexing API
struct TupleVector{T<:Tuple{Vararg{Int,N}} where N} <: AbstractVector{Int}
    data::T
end
Base.size(::TupleVector{<:NTuple{N}}) where {N} = (N,)
@inline function Base.getindex(t::TupleVector, i::Int)
    @boundscheck checkbounds(t, i)
    return @inbounds t.data[i]
end


text:
## rademacher
rademacher(x) = x .* sign.(randn(nobs(x)))

const _mammen_dist = Binomial(1, (sqrt(5) + 1) / (2 * sqrt(5)))
const _mammen_val1 = -(sqrt(5) - 1) / 2
const _mammen_val2 = (sqrt(5) + 1) / 2

function mammen(x)
    r = rand(_mammen_dist, nobs(x))
    return ifelse.(r .== 1, _mammen_val1, _mammen_val2)
end

## Number of exact bootstrap runs
nrun_exact(n::Integer) = binomial(2n - 1, n)

## Jack-Knife estimate
function jack_knife_estimate(data, statistic::Function, j::Int = 1, typ::Type = typeof(statistic(data)[j]))
    n = nobs(data)
    y = zeros(typ, n)
    idx = trues(n)
    for i in 1:n
        idx[i] = false
        if i > 1
            idx[i - 1] = true
        end
        y[i] = tx(statistic(pick(data, idx)))[j]
    end
    return y
end

## Sample quantiles with Gaussian interpolation
## Davison and Hinkley, equation 5.6

function iquantile(x::AbstractVector, alpha::AbstractVector)
    x = sort(x)
    qn = Float64[iquantile(x, a, true) for a in alpha]
    return qn
end

function iquantile(x::AbstractVector, alpha::AbstractFloat, is_sorted::Bool = false)
    if !is_sorted
        x = sort(x)
    end
    n = length(x)
    k = trunc(Int, (n + 1) * alpha)
    ## infinity outside of data range
    if k == 0
        return -Inf
    elseif k > (n - 1)
        return Inf
    end
    tx = [quantile(Normal(), a) for a in [alpha, k / (n + 1), (k + 1) / (n + 1)]]
    qn = (tx[1] - tx[2]) / (tx[3] - tx[2]) * (x[k + 1] - x[k]) + x[k]
    return qn
end


text:
export PPR3, PKatz3

function full_solve(M::SpFltMat)
    K = inv(Matrix(M))
    S = copy(M)
    for j in 1:size(M, 2), i in nz_row_inds(M, j); S[i, j] = K[i, j]; end
    return S
end

function iterative_solve(M::SpFltMat, triangles::Vector{NTuple{3,Int64}})
    n = size(M, 2)
    # only compute for indices that appear in at least one triangle
    inds = zeros(Int64, n)
    for (i, j, k) in triangles; inds[[i, j, k]] .= 1; end
    shuffled_inds = shuffle(findall(inds .> 0))

    nthreads = Threads.nthreads()
    I = Vector{Vector{Int64}}(undef, nthreads)
    J = Vector{Vector{Int64}}(undef, nthreads)
    V = Vector{Vector{Float64}}(undef, nthreads)
    Threads.@threads for t = 1:nthreads
        I[t] = Vector{Int64}()
        J[t] = Vector{Int64}()
        V[t] = Vector{Float64}()
    end
    Threads.@threads for ind = 1:length(shuffled_inds)
        tid = Threads.threadid()
        if tid == 1
            print(stdout, "$(ind) of $(length(shuffled_inds)) \r")
            flush(stdout)
        end
        node = shuffled_inds[ind]
        b = zeros(n)
        b[node] = 1
        sol = dqgmres(M, b, atol=1e-4, rtol=1e-4)[1]
        for i in nz_row_inds(M, node)
            push!(I[tid], i)
            push!(J[tid], node)
            push!(V[tid], sol[i])
        end
    end

    total = sum([length(It) for It in I])
    cI = Vector{Int64}(undef, total)
    cJ = Vector{Int64}(undef, total)
    cV = Vector{Float64}(undef, total)
    curr_ind = 1
    for t in 1:nthreads
        size = length(I[t])
        curr_range = collect(curr_ind:(curr_ind + size - 1))
        cI[curr_range] .= I[t]
        cJ[curr_range] .= J[t]
        cV[curr_range] .= V[t]
        curr_ind += size
    end
    return sparse(cI, cJ, cV, n, n)
end

"""
PKatz3
----

Compute 3-way personalized Katz scores for triangles.

PKatz3(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat,
       unweighted::Bool, dense_solve::Bool=false)

Example usage:
```
(scores, S) = PPR3(T, B, false, false)
```

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- B::SpIntMat: Projected graph as a Sparse integer matrix, where B[i, j] is the number of times that i and j co-appear in a simplex.
- unweighted::Bool: Whether or not to use the unweighted version of the matrix.
- dense_solve::Bool=false: whether or not to use a dense solver. If B is small, then it is worth setting this option to true.

returns a tuple (scores, S):
- scores::Vector{Float64}: a vector of 3-way PPR scores for the triangles
- S::SpFltMat: a sparse matrix with the same sparsity pattern as B, where S[i, j] is the personalized Katz score of node i with respect to the seed node j.
"""
function PKatz3(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat,
                unweighted::Bool, dense_solve::Bool=false)
    W = copy(B)
    if unweighted; W = make_sparse_ones(W); end
    σ_1 = svds(W,  nsv=1)[1].S[1]
    β = min(0.25 / σ_1, 0.5)
    M = I - β * W
    S = (dense_solve ? full_solve(M) : iterative_solve(M, triangles)) - I
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        scores[ind] = S[i, j] + S[j, i] + S[i, k] + S[k, i] + S[j, k] + S[k, j]
    end
    return scores, S
end

"""
PPR3
----

Compute 3-way personalized PageRank scores for triangles.

PPR3(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat,
     unweighted::Bool, dense_solve::Bool=false, α::Float64=0.85)

Example usage:
```
(scores, S) = PPR3(T, B, false, false, 0.85)
```

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- B::SpIntMat: Projected graph as a Sparse integer matrix, where B[i, j] is the number of times that i and j co-appear in a simplex.
- unweighted::Bool: Whether or not to use the unweighted version of the matrix.
- dense_solve::Bool=false: whether or not to use a dense solver. If B is small, then it is worth setting this option to true.
- α::Float64=0.85: teleportation parameter for PageRank

returns a tuple (scores, S):
- scores::Vector{Float64}: a vector of 3-way PPR scores for the triangles
- S::SpFltMat: a sparse matrix with the same sparsity pattern as B, where S[i, j] is the personalized PageRank score of node i with respect to the seed node j.
"""
function PPR3(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat,
              unweighted::Bool, dense_solve::Bool=false, α::Float64=0.85)
    W = copy(B)
    if unweighted; W = make_sparse_ones(W); end
    W = convert(SpFltMat, W)
    d = vec(sum(W, dims=1))
    nonzeros = findall(d .> 0)
    d[nonzeros] = 1.0 ./ d[nonzeros]
    M = I - α * W * Diagonal(d)
    S = (dense_solve ? full_solve(M) : iterative_solve(M, triangles)) * (1 - α)
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        scores[ind] = S[i, j] + S[j, i] + S[i, k] + S[k, i] + S[j, k] + S[k, j]
    end
    return scores, S
end


text:
"""
    as_to_bs(a1, a2)

Calculate reciprocal lattice vectors, `b1` and `b2`, from the real space lattice
vectors, `a1` and `a2`.
"""
function as_to_bs(a1, a2)
    b1 = 2pi * [+a2[2],-a2[1]] / (+a1[1]*a2[2]-a1[2]*a2[1])
    b2 = 2pi * [-a1[2],+a1[1]] / (-a2[1]*a1[2]+a2[2]*a1[1])
    return b1, b2
end

"""
    bs_to_as(b1, b2)

Calculate real space lattice vectors, `a1` and `a2`, from reciprocal lattice
vectors, `b1` and `b2`.

This is actually the same as `as_to_bs`, but I think having both functions makes
the intention of the code more obvious.
"""
function bs_to_as(b1, b2)
    return as_to_bs(b1, b2)
end

# """
#     DiagonalMatrix(diag::AbstractVector{ComplexF64})

# A sparse diagonal matrix that can be used in left division (D \\ X)
# """
# struct DiagonalMatrix <: AbstractMatrix{ComplexF64}
#     diag::AbstractVector{ComplexF64}
# end
# Base.size(A::DiagonalMatrix) = (length(A.diag), length(A.diag))
# Base.getindex(A::DiagonalMatrix, I::Vararg{Int,2}) = I[1]==I[2] ? A.diag[I[1]] : 0


text:
include("../src/hypergraph-affinity-functions.jl")

names = [
"Clothing, Shoes & Accessories"
"Electronics and Office"
"Home, Furniture & Appliances"
"Home Improvement & Patio"
"Baby"
"Toys, Games, and Video Games"
"Food, Household & Pets"
"Pharmacy, Health & Beauty"
"Sports, Fitness & Outdoors"
"Auto, Tires & Industrial"
"Other"
]
dataset = "walmart-trips"
r = 25
H, labels = read_hypergraph_data(dataset,r)

class1 = 10
namelabel = names[class1]

n = size(H,2)
classes = zeros(Int64,n)
C1 = findall(x->x==class1,labels)
classes[C1] .= 1

# Turn labels to two classes

alpha, B1, B2, H1, H2, R1, R2, N = Hypergraph_to_Scores(H,classes,25)
hr, hd = clique_expansion_homophily(N,r)
@show alpha
## Plot
lw = 2.5
ms = 9
gfs = 15
tfs = 15
titlesize = 14
s1 = 350
s2 = 300
k = 5
color1 = :blue
color2 = :gray
title = "$k-size hyperedge, $namelabel"
title = ""
leg = false
xlab = "Affinity type t"
ylab = "Affinity"
p = Plots.plot(fontfamily = "helvetica",linewidth = 2,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = k/5:k/5:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs)

plot!(p,1:k,H2[k,1:k],color = color2, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color2, label = class2)
plot!(p,1:k,H1[k,1:k],color = color1, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color1, label = class1)

savefig("Figures/Label$class1/Affinities_$(k)_Label_$class1.pdf")

## Next
for k = 2:15
s1 = 400
s2 = 300
title = "$k-size hyperedge"
title = ""
leg = false
xlab = "Affinity type t"
ylab = "Affinity / Baseline"
p = Plots.plot([1; k],[1; 1],yscale = :log10,fontfamily = "helvetica",linewidth = 2,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = k/5:k/5:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs)
plot!(p,1:k,R1[k,1:k],color = color1, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color1, label = class1)
plot!(p,1:k,R2[k,1:k],color = color2, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color2, label = class2)

@show k
savefig("Figures/Label$class1/Ratios_$(k)_Label_$class1.pdf")
end

## Group Homophily Index, print usable LaTeX code for building table
include("../src/hypergraph-affinity-functions.jl")
S1,S2  = GHI(H,classes,k)

max_stars_repo_name: nveldt/HypergraphHomophily

max_issues_repo_name: nveldt/HypergraphHomophily

max_forks_repo_name: nveldt/HypergraphHomophily


text:
function internalsolve(mdp::MDP, svi::SerialValueIteration)
  if length(mdp.transition.argnames) == length(mdp.statemap) + length(mdp.actionmap)
    return solveset(mdp, svi)
  elseif length(mdp.transition.argnames) > length(mdp.statemap) + length(mdp.actionmap)
    return solveprob(mdp, svi)
  else
    error(string(
      "unkown transition function of neither T(s,a,s') nor T(s,a) type\n",
      "number of transition arguments: ",
      length(mdp.transition.argnames), "\n",
      "number of state and action variables",
      length(mdp.statemap) + length(mdp.actionmap)))
  end
end

function solveset(mdp::MDP, svi::SerialValueIteration)

  statedim = length(mdp.statemap)
  stateargs = mdp.reward.argnames[1:statedim]
  actionargs = mdp.reward.argnames[1 + statedim:end]

  nstates = length(svi.stategrid)
  nactions = length(svi.actiongrid)

  vold = zeros(nstates)
  vnew = zeros(nstates)
  qval = zeros(nactions, nstates)
  resid = 0.0

  iter = 0
  itertime = 0.0
  cputime = 0.0

  for iter in 1:svi.maxiter

    tic()
    resid = -Inf

    for istate in 1:nstates

      state = getvar(svi.stategrid, mdp.statemap, stateargs, istate)
      qhi = -Inf

      for iaction in 1:nactions

        action = getvar(svi.actiongrid, mdp.actionmap, actionargs, iaction)

        statepIdxs, probs = transition(mdp, svi, state, action, stateargs)
        qnow = 0.0

        for istatep in 1:length(statepIdxs)
          qnow += probs[istatep] * vold[statepIdxs[istatep]]
        end

        qnow *= svi.discount
        qnow += reward(mdp, state, action)

        qval[iaction, istate] = qnow

        if qnow > qhi
          qhi = qnow
          vnew[istate] = qhi
        end

      end

      # use infinity-norm
      newresid = (vold[istate] - vnew[istate])^2
      newresid > resid ? resid = newresid : nothing

    end

    itertime = toq()
    cputime += itertime

    if svi.verbose
      println(string(
        "iter $iter, resid: $resid, itertime: $itertime, cputime: $cputime"))
    end

    resid < svi.tol ? break : nothing

    vtmp = vold
    vold = vnew
    vnew = vtmp

  end

  if iter == svi.maxiter
    warn(string(
      "maximum number of iterations reached; check accuracy of solutions"))
  end

  info(string(
    "value iteration solution generated\n",
    "cputime [s] = ", cputime, "\n",
    "number of iterations = ", iter, "\n",
    "residual = ", resid))

  return ValueIterationSolution(
    qval,
    svi.stategrid,
    svi.actiongrid,
    cputime,
    iter,
    resid)

end

function solveprob(mdp::MDP, svi::SerialValueIteration)

  statedim = length(mdp.statemap)
  stateargs = mdp.reward.argnames[1:statedim]
  actionargs = mdp.reward.argnames[1 + statedim:end]

  nstates = length(svi.stategrid)
  nactions = length(svi.actiongrid)

  vold = zeros(nstates)
  vnew = zeros(nstates)
  qval = zeros(nactions, nstates)
  resid = 0.0

  iter = 0
  itertime = 0.0
  cputime = 0.0

  for iter in 1:svi.maxiter

    tic()
    resid = -Inf

    for istate in 1:nstates

      state = getvar(svi.stategrid, mdp.statemap, stateargs, istate)
      qhi = -Inf

      for iaction in 1:nactions

        action = getvar(svi.actiongrid, mdp.actionmap, actionargs, iaction)
        qnow = 0.0

        for istatep in 1:nstates
          statep = getvar(svi.stategrid, mdp.statemap, stateargs, istatep)
          prob = mdp.transition.fn(state..., action..., statep...)
          qnow += prob * vold[istatep]
        end

        qnow *= svi.discount
        qnow += reward(mdp, state, action)

        qval[iaction, istate] = qnow

        if qnow > qhi
          qhi = qnow
          vnew[istate] = qhi
        end

      end

      # use infinity-norm
      newresid = (vold[istate] - vnew[istate])^2
      newresid > resid ? resid = newresid : nothing

    end

    itertime = toq()
    cputime += itertime

    if svi.verbose
      println(string(
        "iter $iter, resid: $resid, itertime: $itertime, cputime: $cputime"))
    end

    resid < svi.tol ? break : nothing

    vtmp = vold
    vold = vnew
    vnew = vtmp

  end

  if iter == svi.maxiter
    warn(string(
      "maximum number of iterations reached; check accuracy of solutions"))
  end

  info(string(
    "value iteration solution generated\n",
    "cputime [s] = ", cputime, "\n",
    "number of iterations = ", iter, "\n",
    "residual = ", resid))

  return ValueIterationSolution(
    qval,
    svi.stategrid,
    svi.actiongrid,
    cputime,
    iter,
    resid)

end


text:
# A representation of HloIf that takes two IRCodes to represent the computations
struct _HloIf <: HloOp{:conditional}
    if_func::IRCode
    else_func::IRCode
end

struct _HloWhile <: HloOp{:while}
    condition_func::IRCode
    body_func::IRCode
end

Base.show(io::IO, hloif::_HloIf) = print(io, "_HloIf(<omitted>)")
Base.show(io::IO, hloif::_HloWhile) = print(io, "_HloWhile(<omitted>)")

function outline_if_bb(ir, bb_to_outline, types, used_from_outside, merge_phi)
    block = ir.cfg.blocks[bb_to_outline]
    block = ir.cfg.blocks[bb_to_outline]
    outlined_ir = IRCode(Any[], Any[], Int32[], UInt8[], Compiler.CFG(BasicBlock[
        BasicBlock(Compiler.StmtRange(1, 0))
    ], Int64[]), LineInfoNode[], Any[], Any[], Core.svec())
    for i in 1:length(used_from_outside)
        Compiler.append_node!(outlined_ir, types[i], Expr(:call, Core.getfield, Compiler.Argument(2), i), 0)
    end
    for idx in block.stmts
        stmt = ir.stmts[idx]
        isa(stmt, GotoNode) && continue
        stmt′ = isa(stmt, Expr) ? copy(stmt) : stmt
        urs = userefs(stmt′)
        for op in urs
            val = op[]
            if isa(val, SSAValue)
                if val.id in block.stmts
                    op[] = SSAValue(val.id - first(block.stmts) + length(used_from_outside) + 1)
                else
                    op[] = SSAValue(findfirst(==(val), used_from_outside))
                end
            elseif isa(val, Compiler.Argument)
                op[] = SSAValue(findfirst(==(val), used_from_outside))
            end
        end
        stmt′ = urs[]
        Compiler.append_node!(outlined_ir, ir.types[idx], stmt′, 0)
    end
    edge_idx = findfirst(==(bb_to_outline), merge_phi.edges)
    retval_val = (merge_phi.values[edge_idx]::SSAValue)
    val_id = retval_val.id
    retval_outside = findfirst(==(retval_val), used_from_outside)
    retval_val = retval_outside === nothing ?
        SSAValue(val_id - first(block.stmts) + length(used_from_outside) + 1) :
        SSAValue(retval_outside)
    Compiler.append_node!(outlined_ir, ir.types[val_id], ReturnNode(retval_val), 0)
    return outlined_ir
end

function outline_loop_blocks(ir, domtree, bbs_to_outline, types, phi_nodes, used_from_outside, is_header=false)
    bbs = ir.cfg.blocks
    outlined_ir = IRCode(Any[], Any[], Int32[], UInt8[], Compiler.CFG(BasicBlock[
        BasicBlock(Compiler.StmtRange(1, 0))
    ], Int64[]), LineInfoNode[], Any[], Any[], Core.svec())
    ncarried_vals = length(phi_nodes)+length(used_from_outside)
    for i in 1:ncarried_vals
        Compiler.append_node!(outlined_ir, types[i], Expr(:call, Core.getfield, Compiler.Argument(2), i), 0)
    end
    bb_starts = Int[ncarried_vals+1]
    for bb_to_outline in bbs_to_outline
        block = bbs[bb_to_outline]
        for idx in block.stmts
            stmt = ir.stmts[idx]
            isa(stmt, GotoNode) && continue
            if isa(stmt, GotoIfNot)
                @assert is_header && bb_to_outline == bbs_to_outline[end] && idx == block.stmts[end]
            end
            if isa(stmt, PhiNode)
                continue
            end
            stmt′ = isa(stmt, Expr) ? copy(stmt) : stmt
            urs = userefs(stmt′)
            for op in urs
                val = op[]
                if isa(val, SSAValue)
                    val_bb = block_for_inst(ir.cfg, val.id)
                    bb_idx = findfirst(==(val_bb), bbs_to_outline)
                    if bb_idx !== nothing && !isa(ir.stmts[val.id], PhiNode)
                        def_block = bbs[val_bb]
                        block_local_id = val.id - first(def_block.stmts)
                        if bb_idx == 1 && is_header
                            block_local_id -= length(phi_nodes)
                        end
                        op[] = SSAValue(bb_starts[bb_idx] + block_local_id)
                    else
                        found = findfirst(e->e.first==val, phi_nodes)
                        if found == nothing
                            found = findfirst(==(val), used_from_outside) + length(phi_nodes)
                            @assert found !== nothing
                        end
                        op[] = SSAValue(found)
                    end
                elseif isa(val, Compiler.Argument)
                    op[] = SSAValue(findfirst(==(val), used_from_outside) + length(phi_nodes))
                end
            end
            stmt′ = urs[]
            Compiler.append_node!(outlined_ir, ir.types[idx], stmt′, 0)
        end
        push!(bb_starts, length(outlined_ir.stmts) + 1)
    end
    return (bb_starts, outlined_ir)
end

function outline_loop_header(ir, domtree, bb_to_outline, types, phi_nodes, used_from_outside)
    _, outlined_ir = outline_loop_blocks(ir, domtree, [bb_to_outline], types, phi_nodes, used_from_outside, true)
    branch = outlined_ir.stmts[end]
    outlined_ir.stmts[end] = ReturnNode(branch.cond)
    outlined_ir.types[end] = Bool
    return outlined_ir
end

function outline_loop_body(ir, domtree, bbs_to_outline, types, phi_nodes, used_from_outside)
    bbs = ir.cfg.blocks
    ncarried_vals = length(phi_nodes)+length(used_from_outside)
    # This is significant overkill, but I don't care for now. If all you have is a hammer...
    bbs_to_outline = sort(bbs_to_outline, lt=(a, b)->dominates(domtree, a, b))
    bb_starts, outlined_ir = outline_loop_blocks(ir, domtree, bbs_to_outline, types, phi_nodes, used_from_outside)
    # Reconstruct the loop carried tuple
    args = Any[]
    for (_, phi) in phi_nodes
        idx = findfirst(==(bbs_to_outline[end]), phi.edges)
        val = phi.values[idx]
        # It's possible for a phi node to reference a previous phi
        # node, in which case we can handle that
        carried_phi = findfirst(x->x[1] == val, phi_nodes)
        if carried_phi !== nothing
            push!(args, SSAValue(carried_phi))
            continue
        end
        val_bb = block_for_inst(ir.cfg, val.id)
        bb_idx = findfirst(==(val_bb), bbs_to_outline)
        @assert bb_idx !== nothing
        def_block = bbs[val_bb]
        block_local_id = val.id - first(def_block.stmts)
        push!(args, SSAValue(bb_starts[bb_idx] + block_local_id))
    end
    # The rest just gets carried through
    for i in 1:length(used_from_outside)
        push!(args, SSAValue(length(phi_nodes) + i))
    end
    tuple_type = Compiler.tuple_tfunc(types)
    Compiler.append_node!(outlined_ir, tuple_type, Expr(:call, Core.tuple, args...), 0)
    Compiler.append_node!(outlined_ir, tuple_type, ReturnNode(SSAValue(length(outlined_ir.stmts))), 0)
    return outlined_ir
end

function erase_block!(ir, bbidx)
    block = ir.cfg.blocks[bbidx]
    ir.stmts[block.stmts[1:end-1]] .= nothing
    ir.stmts[block.stmts[end]] = ReturnNode()
    empty!(ir.cfg.blocks[bbidx].preds)
    empty!(ir.cfg.blocks[bbidx].succs)
end

function NCA(domtree, nodes)
    nodes = copy(nodes)
    min_level = minimum(n->domtree.nodes[n].level, nodes)
    # Ascend all nodes that are not at the same level
    for i = 1:length(nodes)
        if domtree.nodes[nodes[i]].level != min_level
            nodes[i] = domtree.idoms[nodes[i]]
        end
    end
    # Ascend levels together until all the nodes are the same
    while !all(==(nodes[1]), nodes)
        for i = 1:length(nodes)
            nodes[i] = domtree.idoms[nodes[i]]
        end
    end
    nodes[1]
end

struct OutlineLoop
    header::Int
end

struct OutlineIf
    head::Int
    join::Int
end
entry(l::OutlineLoop) = l.header
entry(l::OutlineIf) = l.head

# Traverse the domtree and find all outlining regions
# TODO: There's much better ways to do this. For now
# I just want to get it to work at all
function analyze_regions(domtree, cfg)
    bbs = cfg.blocks
    todo = Any[]
    for i = 1:length(bbs)
        if length(bbs[i].preds) > 1
            # If we dominate one of the the predecessors, this is a backege,
            # and we outline it as a loop.
            # Else this is a conditional and we need to find the nearest common
            # ancestor of all predecessors in the dominator tree for outlining
            # as a conditional region.
            dominate_any = false
            backedges = Pair{Int, Int}[]
            for pred in bbs[i].preds
                if dominates(domtree, i, pred)
                    dominate_any = true
                    push!(backedges, pred=>i)
                    break
                end
            end
            if dominate_any
                push!(todo, OutlineLoop(i))
            else
                push!(todo, OutlineIf(NCA(domtree, bbs[i].preds), i))
            end
        end
    end
    sort(todo, lt=(a, b)->dominates(domtree, entry(b), entry(a)))
end

function outline_loop!(ir, sv, domtree, loop_header)
    bbs = ir.cfg.blocks
    loop_header_bb = ir.cfg.blocks[loop_header]
    loop_header_preds = loop_header_bb.preds
    @assert length(loop_header_preds) == 2

    # Map the loop
    if dominates(domtree, loop_header, loop_header_preds[1])
        loop_latch = loop_header_preds[1]
        entry_block = loop_header_preds[2]
    else
        loop_latch = loop_header_preds[2]
        entry_block = loop_header_preds[1]
    end
    @assert dominates(domtree, loop_header, loop_latch)
    loop_body_blocks = Int[loop_latch]
    worklist = Int[loop_latch]
    while !isempty(worklist)
        item = pop!(worklist)
        for pred in bbs[item].preds
            if pred != loop_header
                push!(loop_body_blocks, pred)
                push!(worklist, pred)
            end
        end
    end

    loop_cond = ir.stmts[loop_header_bb.stmts[end]]
    @assert isa(loop_cond, GotoIfNot)
    (bb1, bb2) = (loop_header + 1, loop_cond.dest)
    if dominates(domtree, bb2, loop_latch)
        exit_block = bb1
    else
        exit_block = bb2
    end

    # First identify everything that needs to be in our loop carry
    used_from_outside = Union{Compiler.Argument, SSAValue}[]
    used_outside = SSAValue[]
    phi_nodes = Pair{SSAValue,PhiNode}[]
    for bb in Iterators.flatten((loop_header, loop_body_blocks))
        stmt_range = bbs[bb].stmts
        for stmt_idx in stmt_range
            stmt = ir.stmts[stmt_idx]
            # We handle PhiNodes separately
            if isa(stmt, PhiNode)
                @assert bb == loop_header
                push!(phi_nodes, SSAValue(stmt_idx)=>stmt)
                continue
            end
            urs = userefs(stmt)
            for op in urs
                val = op[]
                if isa(val, SSAValue)
                    # If it's in our current basic block, we don't care
                    val.id in stmt_range && continue
                    # Else find out which bb it's in
                    val_bb = block_for_inst(ir.cfg, val.id)
                    # TODO: If something from the loop header (other than a loop
                    # carried phi) is used in the loop latch it needs to be duplicated.
                    val_bb == loop_header && continue
                    val_bb in loop_body_blocks && continue
                elseif !isa(val, Compiler.Argument)
                    continue
                end
                push!(used_from_outside, val)
            end
        end
    end
    unique!(used_from_outside)
    for bb in dominated(domtree, exit_block)
        stmt_range = ir.cfg.blocks[bb].stmts
        for stmt_idx in stmt_range
            stmt = ir.stmts[stmt_idx]
            foreachssa(stmt) do val
                val_bb = block_for_inst(ir.cfg, val.id)
                (val_bb == loop_header) || return
                push!(used_outside, val)
            end
        end
        # We also need to scan the phi nodes of any successor blocks
        for succ_bb in ir.cfg.blocks[bb].succs
            if !dominates(domtree, exit_block, succ_bb)
                stmt_range = ir.cfg.blocks[succ_bb].stmts
                for stmt_idx in stmt_range
                    stmt = ir.stmts[stmt_idx]
                    (isa(stmt, Nothing) || isa(stmt, PhiNode)) || break
                    isa(stmt, Nothing) && continue
                    idx = findfirst(==(bb), stmt.edges)
                    idx === nothing && continue
                    isassigned(stmt.values, idx) || continue
                    val = stmt.values[idx]
                    isa(val, SSAValue) || continue
                    val_bb = block_for_inst(ir.cfg, val.id)
                    (val_bb == loop_header) || continue
                    push!(used_outside, val)
                end
            end
        end
    end
    types = Any[]
    ssa_order = Any[]
    for (val, _) in phi_nodes
        push!(ssa_order, val)
        push!(types, argextype(val, ir, sv.sp))
    end
    for val in used_from_outside
        push!(ssa_order, val)
        push!(types, argextype(val, ir, sv.sp))
    end
    # We put these in a tuple in the order (phi_nodes..., used_from_outside...)
    # Replace Phi nodes by getfield calls
    tuple_T = Compiler.tuple_tfunc(types)
    cond_ir = outline_loop_header(ir, domtree, loop_header, types, phi_nodes, used_from_outside)
    append!(cond_ir.argtypes, [nothing, tuple_T])
    Compiler.verify_ir(cond_ir)
    body_ir = outline_loop_body(ir, domtree, loop_body_blocks, types, phi_nodes, used_from_outside)
    append!(body_ir.argtypes, [nothing, tuple_T])
    Compiler.verify_ir(body_ir)

    # Now erase the blocks we outlined
    erase_block!.((ir,), [loop_header; loop_body_blocks])
    # Fix the CFG
    entry_succs, exit_preds = ir.cfg.blocks[entry_block].succs, ir.cfg.blocks[exit_block].preds
    empty!(entry_succs); empty!(exit_preds)
    push!(entry_succs, exit_block); push!(exit_preds, entry_block)
    # Insert a GotoNode at the end of the entry block
    insert_node!(ir, bbs[entry_block].stmts[end], Any, GotoNode(exit_block), true)
    # Assemble the initial execution tuple
    initial_args = Any[]
    for (_, node) in phi_nodes
        push!(initial_args, node.values[findfirst(==(entry_block), node.edges)])
    end
    append!(initial_args, used_from_outside)
    first_exit_pos = bbs[exit_block].stmts[1]
    initial = insert_node!(ir, first_exit_pos, tuple_T,
        Expr(:call, Core.tuple, initial_args...), false)
    # Insert our while node at the top of the exit_block
    while_return = insert_node!(ir, first_exit_pos, tuple_T,
        Expr(:call, _HloWhile(cond_ir, body_ir), initial), false)
    # Insert destructuring calls
    used_outside_replacement = Any[]
    for u in used_outside
        idx = findfirst(==(u), ssa_order)
        push!(used_outside_replacement,
            insert_node!(ir, first_exit_pos, types[idx],
                Expr(:call, Core.getfield, while_return, idx), false))
    end
    # Replace any outside uses of loop values
    for bb in dominated(domtree, exit_block)
        stmt_range = ir.cfg.blocks[bb].stmts
        for stmt_idx in stmt_range
            stmt = ir.stmts[stmt_idx]
            urs = userefs(stmt)
            for op in urs
                val = op[]
                if isa(val, SSAValue) && val in used_outside
                    op[] = used_outside_replacement[findfirst(==(val), used_outside)]
                end
            end
            ir.stmts[stmt_idx] = urs[]
        end
        for succ_bb in ir.cfg.blocks[bb].succs
            if !dominates(domtree, exit_block, succ_bb)
                stmt_range = ir.cfg.blocks[succ_bb].stmts
                for stmt_idx in stmt_range
                    stmt = ir.stmts[stmt_idx]
                    (isa(stmt, Nothing) || isa(stmt, PhiNode)) || break
                    isa(stmt, Nothing) && continue
                    idx = findfirst(==(bb), stmt.edges)
                    idx === nothing && continue
                    isassigned(stmt.values, idx) || continue
                    val = stmt.values[idx]
                    if isa(val, SSAValue) && val in used_outside
                        stmt.values[idx] = used_outside_replacement[findfirst(==(val), used_outside)]
                    end
                end
            end
        end
    end
    ir = Compiler.compact!(ir)
    Compiler.verify_ir(ir)
    ir
end

function outline_if!(ir, sv, domtree, split_block, join)
    # Simple version for experiments: Assume each branch of the if has
    # exactly one basic block
    @assert isa(ir.stmts[ir.cfg.blocks[split_block].stmts[end]], GotoIfNot)
    divergence_blocks = copy(ir.cfg.blocks[split_block].succs)
    @assert length(divergence_blocks) == 2
    # For our simple experiments only allow a single block
    join_block = ir.cfg.blocks[divergence_blocks[1]].succs[1]
    @assert all(divergence_blocks) do block
        preds = ir.cfg.blocks[block].preds
        succs = ir.cfg.blocks[block].succs
        length(preds) == 1 && length(succs) == 1 && succs[1] == join_block
    end
    phi_node = ir.stmts[ir.cfg.blocks[join_block].stmts[1]]

    all_used_from_outside = Vector{Any}[]
    for bb in divergence_blocks
        # TODO: Allow each branch to have multiple BBS
        used_from_outside = Vector{Any}()
        stmt_range = ir.cfg.blocks[bb].stmts
        function process_val(val)
            if isa(val, SSAValue)
                # If it's in our current basic block, we don't care
                val.id in stmt_range && return
            elseif !isa(val, Compiler.Argument)
                return
            end
            push!(used_from_outside, val)
        end
        for stmt_idx in stmt_range
            stmt = ir.stmts[stmt_idx]
            urs = userefs(stmt)
            for op in urs
                val = op[]
                process_val(val)
            end
        end
        process_val(phi_node.values[findfirst(==(bb), phi_node.edges)])
        unique!(used_from_outside)
        push!(all_used_from_outside, used_from_outside)
    end
    if_types = Any[argextype(arg, ir, sv.sp) for arg in all_used_from_outside[1]]
    if_tuple_T = Compiler.tuple_tfunc(if_types)
    if_func = outline_if_bb(ir, divergence_blocks[1], if_types, all_used_from_outside[1], phi_node)
    else_types = Any[argextype(arg, ir, sv.sp) for arg in all_used_from_outside[2]]
    else_tuple_T = Compiler.tuple_tfunc(else_types)
    else_func = outline_if_bb(ir, divergence_blocks[2], else_types, all_used_from_outside[2], phi_node)
    append!(if_func.argtypes, [nothing, if_tuple_T])
    append!(else_func.argtypes, [nothing, else_tuple_T])
    Compiler.verify_ir(if_func)
    Compiler.verify_ir(else_func)
    inst = _HloIf(if_func, else_func)
    condition = ir.stmts[ir.cfg.blocks[split_block].stmts[end]]
    @assert isa(condition, GotoIfNot)
    # Replace condition by direct branch to join block
    ir.stmts[ir.cfg.blocks[split_block].stmts[end]] = GotoNode(join_block)
    # Replace phi node by HloIf
    ## Look through the conversion to bool
    @assert isa(condition.cond, SSAValue)
    cond_stmt = ir.stmts[condition.cond.id]
    @assert isexpr(cond_stmt, :invoke)
    cond = cond_stmt.args[4]
    #
    insert_pos = ir.cfg.blocks[join_block].stmts[1]
    if_tuple = insert_node!(ir, insert_pos, if_tuple_T, Expr(:call, Core.tuple, all_used_from_outside[1]...))
    else_tuple = insert_node!(ir, insert_pos, else_tuple_T, Expr(:call, Core.tuple, all_used_from_outside[2]...))
    ir.stmts[insert_pos] = Expr(:call, inst, cond, if_tuple, else_tuple)
    # Delete divergence blocks
    erase_block!.((ir,), divergence_blocks)
    # Fix CFG
    let cb_succs = ir.cfg.blocks[split_block].succs, jb_preds = ir.cfg.blocks[join_block].preds
        empty!(cb_succs); empty!(jb_preds);
        push!(cb_succs, join_block); push!(jb_preds, split_block)
    end
    for bbidx in divergence_blocks
        block = ir.cfg.blocks[bbidx]
        empty!(block.preds); empty!(block.succs)
    end
    ir = Compiler.compact!(ir)
    Compiler.verify_ir(ir)
    ir
end

function outline_control_flow!(ir, sv)
    bbs = ir.cfg.blocks
    length(bbs) == 1 && return ir

    domtree = Compiler.construct_domtree(ir.cfg);

    todo = analyze_regions(domtree, ir.cfg)
    while !isempty(todo)
        item = popfirst!(todo)
        if isa(item, OutlineLoop)
            ir = outline_loop!(ir, sv, domtree, item.header)
        else
            ir = outline_if!(ir, sv, domtree, item.head, item.join)
        end
    end
    Compiler.verify_ir(ir)
    ir = Compiler.cfg_simplify!(ir)
    Compiler.verify_ir(ir)
    ir = Compiler.compact!(ir)
    ir
end


text:
"""
    abstract type NonlinearSolver <: GridapType end

- [`solve!(x::AbstractVector,nls::NonlinearSolver,op::NonlinearOperator)`](@ref)
- [`solve!(x::AbstractVector,nls::NonlinearSolver,op::NonlinearOperator, cache)`](@ref)
"""
abstract type NonlinearSolver <: GridapType end

"""
    solve!(x::AbstractVector,nls::NonlinearSolver,op::NonlinearOperator)

Usage:

   cache = solve!(x,nls,op)

The returned `cache` object can be used in subsequent solves:

   cache = solve!(x,nls,op,cache)
"""
function solve!(x::AbstractVector,nls::NonlinearSolver,op::NonlinearOperator)
  solve!(x,nls,op,nothing)
end

"""
    solve!(x::AbstractVector,nls::NonlinearSolver,op::NonlinearOperator,cache)

Solve using the cache object from a previous solve.
"""
function solve!(x::AbstractVector,nls::NonlinearSolver,op::NonlinearOperator,cache)
  @abstractmethod
end

"""
    solve(nls::NonlinearSolver,op::NonlinearOperator)

Creates and uses a zero initial guess.
"""
function solve(nls::NonlinearSolver,op::NonlinearOperator)
  x = zero_initial_guess(op)
  solve!(x,nls,op)
  x
end

"""
    test_nonlinear_solver(
      nls::NonlinearSolver,
      op::NonlinearOperator,
      x0::AbstractVector,
      x::AbstractVector,
      pred::Function=isapprox)
"""
function test_nonlinear_solver(
  nls::NonlinearSolver,
  op::NonlinearOperator,
  x0::AbstractVector,
  x::AbstractVector,
  pred::Function=isapprox)

  x1 = copy(x0)
  cache = solve!(x1,nls,op)
  @test pred(x1,x)

  x1 = copy(x0)
  cache = solve!(x1,nls,op,cache)
  @test pred(x1,x)

  x1 = copy(x0)
  cache = solve!(x1,nls,op,cache)
  @test pred(x1,x)

end


text:
"""
    handle_packets(link::Link, T)

Handle incoming packets on `link`, returning a type `T`.
"""
function handle_packets(link::Link, T)
    while true
        pkt = nextpacket(link)
        if pkt == PKT_RETURN
            return get(link, T)
        elseif pkt == PKT_TEXT
            print(get(link, String))
        elseif pkt == PKT_MESSAGE
            newpacket(link)
            @assert nextpacket(link) == PKT_TEXT
            @warn get(link, String)
        else
            error("Unsupported packet type $pkt")
        end
    end
end

function weval(link::Link, T, expr)
    put(link, W"EvaluatePacket"(expr))
    endpacket(link)
    handle_packets(link, T)
end

"""
    weval([T,] expr)

Evaluate expression `expr`, returning a value of type `T` (default = `Any`).
"""
weval(T, expr) = weval(_defaultlink(), T, expr)
weval(expr) = weval(Any, expr)

function wevalstr(link, T, str::AbstractString)
    weval(link, W"ToExpression"(str), T)
end
wevalstr(T, expr) = wevalstr(_defaultlink(), T, expr)
wevalstr(expr) = wevalstr(Any, expr)

function parseexpr(str::AbstractString)
    r = weval(W"ToExpression"(str, W"InputForm", W"Hold"))
    r.args[1]
end


text:
import Distances: chebyshev

##############################################################################
##
## Types and constructors
##
##############################################################################

##################################
##
## Household PROBLEM
##
##################################
"""
Consumption-saving problem with two assets and non-convex adjustment costs

## Fields

#### Parameters
- `γ::Float64`  : parameter on CRRA utility
- `ρ::Float64`  : discount factor
- `ξ::Float64`  : automatic deposit
- `σ::Float64`  : frisch elasticity of labor supply
- `ψ::Float64`  : disutility of labor

#### Prices
- `rᴬ::Float64`     : interest rate on illiquid asset
- `rᴮ::Float64`     : interest rate for savings in liquid asset
- `wedge::Float64`  : interest rate for borrowing in liquid asset
- `w::Float64`      : wage rate

#### Tax
- `τ::Float64`      : tax on income
- `T::Float64`      : Lump-sum transfer

#### Cost Function
- `χ₀::Float64` :
- `χ₁::Float64` :
"""
type TwoAssetsProb
    γ::Float64      # CRRA utility with parameter γ
    ρ::Float64      # discount rate
    ξ::Float64      # automatic deposit on illiquid asset
    σ::Float64
    ψ::Float64

    rᴬ ::Float64    # ret on illiquid asset
    rᴮ::Float64    # ret on liq asset
    wedge::Float64
    w  ::Float64    # wage rate

    τ::Float64
    T::Float64

    χ₀::Float64     # parameters on adjustment cost
    χ₁::Float64
end

function Base.show(io::IO, twoap::TwoAssetsProb)
    @printf io "    HOUSEHOLD PROBLEM\n"
    @printf io "\n"
    @printf io "    Parameters  \n"
    @printf io "   ------------ \n"
    @printf io "     γ:      %.3f\n" twoap.γ
    @printf io "     ρ:      %.3f\n" twoap.ρ
    @printf io "     ξ:      %.3f\n" twoap.ξ
    @printf io "\n"
    @printf io "    Prices  \n"
    @printf io "   -------- \n"
    @printf io "     rᴬ   :    %.3f \n" twoap.rᴬ
    @printf io "     rᴮ   :    %.3f \n" twoap.rᴮ
    @printf io "     wage :    %.3f \n" twoap.w
    @printf io "     τ    :    %.3f \n" twoap.τ
    @printf io "     T    :    %.3f \n" twoap.T
    @printf io "\n"
    @printf io "    Cost Function  \n"
    @printf io "   --------------- \n"

    m ="""
    χ₀ + χ₁ * (d/a)² * a
    """
    @printf io "     %s\n" m
end

function TwoAssetsProb2(pc::Prices; γ::Float64 = 2.0, ρ::Float64 = 0.06, ξ::Float64 = 0.1,
                            σ = 0.5, ψ = 27.0,
                            τ::Float64 = 0.0, T::Float64 = 0.,
                            χ₀::Float64 = 0.08, χ₁::Float64 = 3.0)

    # REVIEW: define some global functions that will now change on the go
    #== Substitute the functions to be used afterward ==#
    # global χ(d,a) = χ(d,a, [χ₀,χ₁])
    # global dχinv(pVb::Float64, pVa::Float64, a::Float64) = dχinv(pVb, pVa, a, χ₁)
    # global

    TwoAssetsProb(γ, ρ, ξ, σ, ψ, pc.rᴬ, pc.rᴮ, pc.wedge, pc.w, τ, T, χ₀, χ₁)
end

_unpackparams(twoap::TwoAssetsProb) =
twoap.γ, twoap.ρ, twoap.ξ, twoap.σ, twoap.ψ
_unpackprices(twoap::TwoAssetsProb) =
twoap.rᴬ, twoap.rᴮ, twoap.rᴮ+twoap.wedge, twoap.w
_unpacktax(twoap::TwoAssetsProb) =
twoap.τ, twoap.T

function utilfn(twoap::TwoAssetsProb, c::Float64, ℓutil::Float64)
    return (c-ℓutil)^(1.- twoap.γ)/(1.-twoap.γ)
end

function χ(twoap::TwoAssetsProb, d::Float64, a::Float64, abar::Float64 = 2.0)
    χ₀, χ₁ = twoap.χ₀, twoap.χ₁
    return χ₀*(abs(d)>0) + 0.5*χ₁*d^2.0/max(a,abar)
end
# χ(d::Float64,a::Float64, χval::Array{Float64,1}) = χval[1]*(abs(d)>0) + 0.5*χval[2]*d^2.0/max(a,1e-5)

function dχinv(twoap::TwoAssetsProb, ∂ratio::Float64, a::Float64, abar::Float64 = 2.0)
    χ₁ = twoap.χ₁
    return 1.0/χ₁ * ( ∂ratio-1.0 ) * max(a,abar)
end

function change_prices!(twoap::TwoAssetsProb, pc::Prices)

    twoap.rᴬ, twoap.rᴮ, twoap.w = pc.rᴬ, pc.rᴮ, pc.w
    return Void
end

function change_transfer!(twoap::TwoAssetsProb, τ::Float64, T::Float64)

    twoap.τ, twoap.T = τ, T
    return Void
end

# OLD function
# function foc_dep(∂Vb::Float64, ∂Va::Float64, a::Float64, χ₁::Float64)
#
#     d = ( ∂Va/∂Vb-1 ) * a/χ₁              # optimal d in case of adjustment
#
#     # NOTE:340 compute the change in utility and use that to decide whether to act or
#     #       not...
#
#     ∂V = ∂Va*d - ∂Vb*( d + χ(d,a)  )      # change in utility
#     d = d * (∂V > 0)
# end

##############################################################################
##
## Helping functions
##
##############################################################################

function powerspacegrid(init::Real, eend::Real, n::Int64, k::Real = 1)

    if n<=2
        println("n has to be larger than 2")
        return
    end

    x = collect(linspace(0,1,n))

    z = x.^(1.0/float(k))

    return init +(eend - init)*z
end

##################################
##
## FDSpec
##
##################################

abstract FDSol
abstract FDSpec

type SolutionExp <: FDSol

    ## Solution
    V::Vector{Vector{Float64}}      # Value function
    g::Vector{Vector{Float64}}      # Density over state space
    c::Array{Float64,3}             # Optimal consumption
    sc::Array{Float64,3}            # Savings without deposit
    d::Array{Float64,3}             # Optimal deposit flow

    ## Matrices to be filled
    A ::Vector{Base.SparseMatrix.SparseMatrixCSC{Float64, Int}}  # used on HJB to compute vⁿ⁺¹
    Au::Vector{Base.SparseMatrix.SparseMatrixCSC{Float64, Int}}  # used for KFE to compute transition dynamics

end

unpackcopy_fdsol(fdsol::SolutionExp) = deepcopy(fdsol.V), deepcopy(fdsol.g), deepcopy(fdsol.c), deepcopy(fdsol.sc), deepcopy(fdsol.d), deepcopy(fdsol.A), deepcopy(fdsol.Au)

"""
Specification for the Implicit-Explicit Finite Difference method
"""
type FDExp <: FDSpec

    ## GRID INFO
    b      ::Vector{Float64}
    Δbgrid ::Vector{Float64}
    ΔTbgrid::Vector{Float64}
    rbdrift::Vector{Float64}
    netbinc::Matrix{Float64}

    a      ::Vector{Float64}
    Δagrid ::Vector{Float64}
    ΔTagrid::Vector{Float64}
    radrift::Vector{Float64}
    netainc::Matrix{Float64}

    ΔTab::Vector{Float64}

    ## LaborSupply decisions
    ℓsupply::Vector{Float64}
    ℓutilgrid::Vector{Float64}

    ## Δ in finite difference scheme
    invΔ::Float64
    invΔᴷ::Float64

    ## Stochastic Information
    z::Vector{Float64}
    λ::Matrix{Float64}
    λdiag::Vector{Float64}
    λoff::Matrix{Float64}

    ## Solution
    sol::SolutionExp

    ## Storage
    b̃ ::Vector{Vector{Float64}}                                  # RHS of hjb
    B ::Vector{Base.SparseMatrix.SparseMatrixCSC{Float64, Int}}  # storage matrix
end

function Base.show(io::IO, fd::FDExp)
    fd.Δbgrid[2]-fd.Δbgrid[1] == 0 ? (binfo = "uniform ") : (binfo = "non-uniform ")
    fd.Δagrid[2]-fd.Δagrid[1] == 0 ? (ainfo = "uniform ") : (ainfo = "non-uniform ")
    @printf io "\n"
    @printf io "    Explicit-Implicit Finite Difference Method\n"
    @printf io "\n"
    @printf io "    Grids  \n"
    @printf io "   ------- \n"
    @printf io "     %12sb: %3d points in [% .0f, %.0f]\n" binfo length(fd.b) fd.b[1] fd.b[end]
    @printf io "     %12sa: %3d poitns in [% .0f, %.0f]\n" ainfo length(fd.a) fd.a[1] fd.a[end]
end

function FDExp(twoap::TwoAssetsProb;
z::Vector{Float64} = [.5,1.5], λ::Matrix{Float64} = [-1/3 1/3; 1/3 -1/3],
fixedΔb = true, bn::Int =100, bmin::Float64 =-1.0, bmax::Float64 = 40.0, bparam::Vector{Float64} = [0.35, 0.7],
fixedΔa = true, an::Int = 50, amin::Float64 = 1e-8, amax::Float64 = 70.0, aparam::Float64 = 0.7,
invΔ::Float64 = .025, invΔᴷ::Float64 = 0.025)
    # z = [.8,1.3];
    # λ = [-1/3 1/3;1/3 -1/3];
    #
    # fixedΔb = true;
    # bn =100;
    # bmin =-2.0;
    # bmax = 40.0;
    # bparam = [0.35;0.7];
    #
    # fixedΔa = true;
    # an = 50;
    # amin = 0.0;
    # amax = 70.0;
    # aparam = 0.7;
    #
    # invΔ = .025;
    # invΔᴷ = 0.075

    #== Household parameters and prices ==#
    γ, ρ, ξ, σ, ψ  = _unpackparams(twoap)
    rᴬ, rᴮ, rᴮ⁻, w = _unpackprices(twoap)
    τ, T           = _unpacktax(twoap)

    #== Construct the grids ==#
    if fixedΔb
        b   = powerspacegrid(bmin,bmax,bn)
    else
        ngpbPOS = Int(0.9*bn)
        ngpbNEG = bn - ngpbPOS

        bpos = powerspacegrid(1e-8, bmax, ngpbPOS, bparam[2])
        bneg = zeros(ngpbNEG)
        n = Int(ngpbNEG/2+1)
        bneg[1:n] = powerspacegrid(bmin, (bpos[1]+bmin)/2, n , bparam[1])
        b = [bneg;bpos]
        for i in n+1:ngpbNEG
            b[i] = bpos[1] - (b[ngpbNEG+2-i]-b[1])
        end
    end

    fixedΔa ?  (a = powerspacegrid(amin,amax,an)) : (a = powerspacegrid(amin,amax,an,aparam[1]))

    zn  = length(z); abn = an*bn

    #== Grid counting functions ==#
    global afromab(abi) = div(abi-1,bn)+1
    global bfromab(abi) = rem(abi-1,bn)+1
    global abfromab(ai,bi) = (ai-1)*bn + bi
    global abzfromabz(ai,bi,zi) = (zi-1)*an*bn + ( ai-1 )*bn +bi

    #== Grid info ==#
    Δbgrid = b[2:end]-b[1:end-1]
    Δagrid = a[2:end]-a[1:end-1]

    ΔTbgrid = zeros(bn); ΔTagrid = zeros(an)
    ΔTbgrid[1]       = 0.5*Δbgrid[1]
    ΔTbgrid[2:end-1] = 0.5*( Δbgrid[2:end] + Δbgrid[1:end-1] )
    ΔTbgrid[end]     = 0.5*Δbgrid[end]

    ΔTagrid[1]       = 0.5*Δagrid[1]
    ΔTagrid[2:end-1] = 0.5*( Δagrid[2:end] + Δagrid[1:end-1] )
    ΔTagrid[end]     = 0.5*Δagrid[end]

    #== Construct a vector nb*na with the terms Δb̃[bi] * Δã[ai] =#
    ΔTab = zeros(abn)
    for abi = 1:abn

        ΔTab[abi] = ΔTagrid[afromab(abi)] * ΔTbgrid[bfromab(abi)]
    end

    #== Labor grid info ==#
    ℓsupply    = zeros(zn)
    ℓsupply[:] = (1/ψ * (1-τ) * w)^σ
    ℓutilgrid   = ψ * (z .* ℓsupply.^( 1 + 1.0/σ ) ) / ( 1+1./σ )

    #== return drifts ==#
    rbdrift = (b.>= 0.0)*rᴮ + (b.<0.0)*rᴮ⁻
    rbdrift = b .* rbdrift

    # NOTE: trick on the upper part of grid
    #== impose tax on the very top part of grid ==#
    τa = 15
    τc = rᴬ * ( a[end] * 0.999 )^(1 - τa)
    radrift = rᴬ * a - τc * a .^ τa

    #== net?drifts ==#
    netbinc = zeros(bn,zn)
    netainc = zeros(an,zn)

    for zi in 1:zn
        netℓinc = (1-τ)* w * z[zi] * ℓsupply[zi] + T
        netbinc[:,zi] =  (1-ξ) * netℓinc + rbdrift
        netainc[:,zi] = radrift + ξ * netℓinc
    end

    #== Idiosyncratic shock ==#
    λdiag = diag(λ)
    λoff  = λ - diagm(λdiag)

    ### STORAGE ###
    A = SparseMatrixCSC{Float64, Int}[spdiagm(
        (ones(abn), ones(abn-1), ones(abn-1), ones( abn -bn ) , ones( abn -bn ) ),
        (0, 1, -1, bn, -bn) ) for zi =1:zn]

    # B = SparseMatrixCSC{Float64, Int}[spdiagm(
    #     (ones(abn), ones(abn-1), ones(abn-1), ones( abn -bn ) , ones( abn -bn ) ),
    #     (0, 1, -1, bn, -bn) ) for zi =1:zn]

    B  = Array(SparseMatrixCSC{Float64, Int},zn)
    Au = Array(SparseMatrixCSC{Float64, Int},zn)
    for zi=1:zn
        fill!(nonzeros(A[zi]), zero(Float64))
        B[zi]  = deepcopy(A[zi])
        Au[zi] = deepcopy(A[zi])
    end

    #== Initial Distribution ==#
    g      = zeros(abn,zn)
    bpos_ind = findfirst(b.>0)

    #REVIEW:70 change ydist assumption
    ydist = ones(zn)./zn
    for zi=1:zn
        for abi =1:abn
            ai = afromab(abi)
            bi = bfromab(abi)
            ai==1 && bi==bpos_ind && ( g[abi,zi] = ydist[zi]/(ΔTagrid[ai]*ΔTbgrid[bi]) )
        end
    end

    g  = Vector{Float64}[g[:,zi] for zi=1:zn]
    #----------------------------------------------------------------

    ### SOLUTION ###
    V = Vector{Float64}[Array(Float64, abn) for zi=1:zn]
    ij = zero(Int)
    for zi in 1:zn
        ij = Int(0)
        for ai in 1:an, bi in 1:bn
            ij += 1
            V[zi][ij] = ( (1- ξ)*w*z[zi] + rᴬ*a[ai] + rᴮ⁻*b[bi] ).^( 1- γ )/( 1- γ )/ρ
        end
    end

    # == storage arrays == #
    b̃  = Vector{Float64}[Array(Float64, abn) for zi=1:zn]         # RHS of the Bellman equation

    #== policies ==#
    c  = Array(Float64, bn, an, zn)     # consumption policy
    sc = copy(c); d  = copy(c)

    #== Create solution type ==#
    sol = SolutionExp(V, g, c, sc, d, A, Au)

FDExp(
b, Δbgrid, ΔTbgrid, rbdrift, netbinc,
a, Δagrid, ΔTagrid, radrift, netainc, ΔTab,
ℓsupply, ℓutilgrid,
invΔ, invΔᴷ,
z, λ, λdiag, λoff,
sol,
b̃, B)
end

function initialize_solexp(twoap::TwoAssetsProb, fd::FDSpec)

    #= precompute the lengths =#
    b = fd.b
    a = fd.a
    z = fd.z
    bn = length(b); an = length(a); zn = length(z); abn = bn*an

    #= extract some other from fd =#
    ΔTagrid = fd.ΔTagrid
    ΔTbgrid = fd.ΔTbgrid
    netbinc = fd.netbinc

    #== Create the A matrix ==#
    A = SparseMatrixCSC{Float64, Int}[spdiagm(
        (ones(abn), ones(abn-1), ones(abn-1), ones( abn -bn ) , ones( abn -bn ) ),
        (0, 1, -1, bn, -bn) ) for zi =1:zn]

    Au = Array(SparseMatrixCSC{Float64, Int},zn)
    for zi=1:zn
        fill!(nonzeros(A[zi]), zero(Float64))
        Au[zi] = deepcopy(A[zi])
    end

    #== Initial Distribution ==#
    g      = zeros(abn,zn)
    bpos_ind = findfirst(b.>0)

    #REVIEW:70 change ydist assumption
    ydist = ones(zn)./zn
    for zi=1:zn
        for abi =1:abn
            ai = afromab(abi)
            bi = bfromab(abi)
            ai==1 && bi==bpos_ind && ( g[abi,zi] = ydist[zi]/(ΔTagrid[ai]*ΔTbgrid[bi]) )
        end
    end

    g  = Vector{Float64}[g[:,zi] for zi=1:zn]
    #----------------------------------------------------------------

    ### SOLUTION ###
    V = Vector{Float64}[Array(Float64, abn) for zi=1:zn]
    ij = zero(Int)
    for zi in 1:zn
        ij = Int(0)
        for ai in 1:an, bi in 1:bn
            ij += 1
            V[zi][ij] = ( netbinc[bi,zi] ).^( 1- γ )/( 1- γ )/ρ
        end
    end

    #== policies ==#
    c  = Array(Float64, bn, an, zn)     # consumption policy
    sc = copy(c); d  = copy(c)

    #== Create solution type ==#
    SolutionExp(V, g, c, sc, d, A, Au)
end

"""
Update the grid for a different specification of HOUSEHOLD PROBLEM
"""
function change_grid!(fd::FDSpec, twoap::TwoAssetsProb)

    #== Parameters ==#
    _, _, ξ, σ, ψ  = _unpackparams(twoap)
    rᴬ, rᴮ, rᴮ⁻, w = _unpackprices(twoap)
    τ, T           = _unpacktax(twoap)

    b, a, z = fd.b, fd.a, fd.z
    bn = length(b)
    an = length(a)
    abn = bn * an
    zn = length(z)

    #== Labor grid info ==#
    ℓsupply    = zeros(zn)
    ℓsupply[:] = (1/ψ * (1-τ) * w)^σ
    fd.ℓsupply = ℓsupply
    fd.ℓutilgrid   = ψ * (z .* ℓsupply.^( 1 + 1.0/σ ) ) / ( 1+1./σ )

    #== Update the drifts ==#
    rbdrift    = ( b.>= 0.0)*rᴮ + (b.<0.0)*rᴮ⁻
    fd.rbdrift = b .* rbdrift

    τa = 15
    τc = rᴬ * ( a[end] * 0.999 )^(1 - τa)
    fd.radrift = rᴬ * a - τc * a .^ τa

    for zi in 1:zn
        netℓinc = (1-τ)* w *z[zi] * ℓsupply[zi] + T
        fd.netbinc[:,zi] = (1-ξ) * netℓinc + fd.rbdrift
        fd.netainc[:,zi] = fd.radrift + ξ * netℓinc
    end

end

##############################################################################
##
## Implicit scheme
##
##############################################################################

type SolutionImp <: FDSol

    ## Household Solution
    V::Vector{Float64}          # Value function
    g::Vector{Float64}          # Density over state space
    c::Array{Float64,3}         # Optimal consumption
    sc::Array{Float64,3}        # Savings without deposit
    d::Array{Float64,3}         # Optimal deposit flow

    ## matrices to be filled
    A  ::Base.SparseMatrix.SparseMatrixCSC{Float64, Int}  # intensity matrix A on vⁿ⁺¹
    Au ::Base.SparseMatrix.SparseMatrixCSC{Float64, Int}  # intensity matrix A for kfe

end
"""
Specification for the Implicit Finite Difference method
"""
type FDImp <: FDSpec

    ## GRID INFO
    b      ::Vector{Float64}
    Δbgrid ::Vector{Float64}
    ΔTbgrid::Vector{Float64}
    rbdrift::Vector{Float64}
    netbinc::Matrix{Float64}

    a      ::Vector{Float64}
    Δagrid ::Vector{Float64}
    ΔTagrid::Vector{Float64}
    radrift::Vector{Float64}
    netainc::Matrix{Float64}

    ## LaborSupply decisions
    ℓsupply  ::Vector{Float64}
    ℓutilgrid::Vector{Float64}

    ## Δ in finite difference scheme
    invΔ ::Float64
    invΔᴷ::Float64

    ## Stochastic Information
    z::Vector{Float64}
    λ::Matrix{Float64}

    ## OUTPUT
    sol::SolutionImp

    ## Storage
    b̃  ::Vector{Float64}                                  # RHS of hjb
    B  ::Base.SparseMatrix.SparseMatrixCSC{Float64, Int}  # B storage matrix
    Λ  ::Base.SparseMatrix.SparseMatrixCSC{Float64, Int}  # stochastic process terms (useful bc can fill before)
end

function Base.show(io::IO, fd::FDImp)
    fd.Δbgrid[2]-fd.Δbgrid[1] == 0 ? (binfo = "uniform ") : (binfo = "non-uniform ")
    fd.Δagrid[2]-fd.Δagrid[1] == 0 ? (ainfo = "uniform ") : (ainfo = "non-uniform ")
    @printf io "\n"
    @printf io "    Implicit Finite Difference Method\n"
    @printf io "\n"
    @printf io "    Grids  \n"
    @printf io "   ------- \n"
    @printf io "     %12sb: %3d points in [% .0f, %.0f]\n" binfo length(fd.b) fd.b[1] fd.b[end]
    @printf io "     %12sa: %3d poitns in [% .0f, %.0f]\n" ainfo length(fd.a) fd.a[1] fd.a[end]
    # @printf io "     z: %.2f\n" twoap.ξ
end

function FDImp(twoap::TwoAssetsProb;
z::Vector{Float64} = [.8,1.3], λ::Matrix{Float64} = [-1/3 1/3; 1/3 -1/3],
fixedΔb = true, bn::Int =100, bmin::Float64 =-2.0, bmax::Float64 = 40.0, bparam::Vector{Float64} = [0.35, 0.7],
fixedΔa = true, an::Int = 50, amin::Float64 = 0.0, amax::Float64 = 70.0, aparam::Float64 = 0.7,
invΔ::Float64 = .01, invΔᴷ::Float64 = .05)

    #== Parameters ==#
    γ, ρ, ξ, σ, ψ  = _unpackparams(twoap)
    rᴬ, rᴮ, rᴮ⁻, w = _unpackprices(twoap)
    τ, T           = _unpacktax(twoap)

    #== construct the grids ==#
    if fixedΔb
        b   = powerspacegrid(bmin,bmax,bn)
    else
        ngpbPOS = Int(0.8*bn)
        ngpbNEG = bn - ngpbPOS

        bpos = powerspacegrid(1e-8, bmax, ngpbPOS, bparam[2])
        bneg = zeros(ngpbNEG)
        n = Int(ngpbNEG/2+1)
        bneg[1:n] = powerspacegrid(bmin, (bpos[1]+bmin)/2, n , bparam[1])
        b = [bneg;bpos]
        for i in n+1:ngpbNEG
            b[i] = bpos[1] - (b[ngpbNEG+2-i]-b[1])
        end
    end
    fixedΔa ? (a = powerspacegrid(amin,amax,an)) : (a = powerspacegrid(amin,amax,an,aparam[1]))

    zn  = length(z)
    abn = an*bn

    #== Grid counting functions ==#
    global afromab(abi) = div(abi-1,bn)+1
    global bfromab(abi) = rem(abi-1,bn)+1
    global abfromab(ai,bi) = (ai-1)*bn + bi
    global abzfromabz(ai,bi,zi) = (zi-1)*an*bn + ( ai-1 )*bn +bi

    #== Grid info ==#
    Δbgrid = b[2:end]-b[1:end-1]
    Δagrid = a[2:end]-a[1:end-1]

    ΔTbgrid = zeros(bn); ΔTagrid = zeros(an)
    ΔTbgrid[1]       = 0.5*Δbgrid[1]
    ΔTbgrid[2:end-1] = 0.5*( Δbgrid[2:end] + Δbgrid[1:end-1] )
    ΔTbgrid[end]     = 0.5*Δbgrid[end]

    ΔTagrid[1]       = 0.5*Δagrid[1]
    ΔTagrid[2:end-1] = 0.5*( Δagrid[2:end] + Δagrid[1:end-1] )
    ΔTagrid[end]     = 0.5*Δagrid[end]

    #== Labor grid info ==#
    ℓsupply    = zeros(zn)
    ℓsupply[:] = (1/ψ * (1-τ) * w)^σ
    ℓutilgrid  = ψ * (z .* ℓsupply.^( 1 + 1.0/σ ) ) / ( 1+1./σ )

    rbdrift = (b.>= 0.0)*rᴮ + (b.<0.0)*rᴮ⁻
    rbdrift = b .* rbdrift

    #== impose tax on the very top part of grid ==#
    τa = 15
    τc = rᴬ * ( a[end] * 0.999 )^(1 - τa)
    radrift = rᴬ * a - τc * a .^ τa

    #== net_?drifts ==#
    netbinc = zeros(bn,zn)
    netainc = zeros(an,zn)

    for zi in 1:zn
        netℓinc = (1-τ)* w * z[zi] * ℓsupply[zi] + T
        netbinc[:,zi] =  (1-ξ) * netℓinc + rbdrift
        netainc[:,zi] = radrift + ξ * netℓinc
    end

    ## STORAGE

    # NOTE:250 create a sparse matrix of the RIGHT format that later will be filled with zeros
    #       here we have to be careful bc prob has 3 state variables (bi, aj, zk)
    onesλ = Array{Float64,1}[ones( (zn-i)*abn ) for i=1:zn-1]
    colλ  = Int64[(i)*abn for i=1:zn-1]

    A = spdiagm(
        (ones(abn*zn), ones(abn*zn-1), ones(abn*zn-1), ones( abn*zn -bn ) , ones( abn*zn -bn ),
        onesλ..., onesλ...),
        (0, 1, -1, bn, -bn, colλ..., -colλ...)
    )

    fill!(nonzeros(A), zero(Float64))
    Au = deepcopy(A)
    B = deepcopy(A)
    Λ = deepcopy(A)

    # == fill up matrix Λ == #
    ij = zero(Int)
    λt = λ'
    @inbounds for zi in 1:zn, abi in 1:abn
        ij += 1
        Λ[ abi:abn:(zn-1)*abn+abi , ij] = λt[:,zi]
    end

    # == storage arrays == #
    b̃  = Array(Float64, abn*zn)         # RHS of the Bellman equation
    c  = Array(Float64, bn, an, zn)     # consumption policy
    sc = copy(c); d  = copy(c)

    #== Initial value function guess ==#
    V = Array(Float64, abn*zn)
    ij = zero(Int)
    for zi in 1:zn, ai in 1:an, bi in 1:bn
        ij += 1
        V[ij] = ( (1- ξ)*w*z[zi] + rᴬ*a[ai] + rᴮ⁻*b[bi] ).^( 1- γ )/( 1- γ )/ρ
    end

    #== Initial Distribution ==#
    g = zeros(abn,zn)
    bpos_ind = findfirst(b.>0)

    #REVIEW:60 change ydist assumption
    ydist = ones(zn)./zn
    for zi=1:zn
        for abi =1:abn
            ai = afromab(abi)
            bi = bfromab(abi)
            ai==1 && bi==bpos_ind && ( g[abi,zi] = ydist[zi]/(ΔTagrid[ai]*ΔTbgrid[bi]) )
        end
    end
    g = g[:]

    sol = SolutionImp(V, g, c, sc, d, A, Au)

FDImp(
b, Δbgrid, ΔTbgrid, rbdrift, netbinc,
a, Δagrid, ΔTagrid, radrift, netainc,
ℓsupply, ℓutilgrid,
invΔ, invΔᴷ,
z, λ,
sol,
b̃, B, Λ)
end

##################################
##
## TwoAssetsFD
##
##################################
#
# type TwoAssetsFD
#     ## Parameters of the problem
#     twoap::TwoAssetsProb
#
#     ## Discretization
#     b::Vector{Float64}
#     Δb::Float64
#     rbdrift::Vector{Float64}
#     a::Vector{Float64}
#     Δa::Float64
#     radrift::Vector{Float64}
#     invΔ::Float64
#
#     ## Stochastic Information
#     λ::Matrix{Float64}
#     z::Vector{Float64}
#
#     ## Solution
#     V::Vector{Float64}          # Value function
#     gg::Vector{Float64}         # distribution
#
#     ## Storage
#     newV::Vector{Float64}                               # New Value function
#     u::Vector{Float64}                                  # utility term
#     g⁰::Vector{Float64}                                 # for KolmogorovForward
#     A::Base.SparseMatrix.SparseMatrixCSC{Float64, Int}  # rhs matrix on vⁿ⁺¹
#     B::Base.SparseMatrix.SparseMatrixCSC{Float64, Int}  # B = (indΔ + ρ)I - A
#     Λ::Base.SparseMatrix.SparseMatrixCSC{Float64, Int}  # stochastic process terms (useful bc can fill before)
# end
#
#
# function TwoAssetsFD(twoap::TwoAssetsProb,
#                      z::Vector{Float64} = [.8,1.3], λ::Matrix{Float64} = [-1/3 1/3; 1/3 -1/3],
#                      bn::Int =100, bmin::Float64 =-2.0, bmax::Float64 = 40.0,
#                      an::Int = 50, amin::Float64 = 0.0, amax::Float64 = 70.0,
#                      invΔ::Float64 = .01)
#
#     #== Parameters ==#
#     γ = twoap.γ; ρ = twoap.ρ; ξ = twoap.ξ; rᴬ = twoap.rᴬ;
#     rᴮ = twoap.rᴮ; ; w = twoap.w
#     rᴮ⁻ = twoap.rᴮ + twoap.wedge
#     b   = collect(linspace(bmin,bmax,bn))
#     a   = collect(linspace(amin,amax,an))
#     Δb = b[2]-b[1]; Δa = a[2]-a[1];
#
#     rbdrift = ( b.>= 0.0)*rᴮ + (b.<0.0)*rᴮ⁻
#     rbdrift = b .* rbdrift
#     radrift = rᴬ * a
#
#     zn  = length(z)
#     abn = an*bn
#
#     # NOTE:380 create a sparse matrix of the RIGHT format that later will be filled with zeros
#     #       here we have to be careful bc prob has 3 state variables (bi, aj, zk)
#     onesλ = Array{Float64,1}[ones( (zn-i)*abn ) for i=1:zn-1]
#     colλ  = Int64[(i)*abn for i=1:zn-1]
#
#     A = spdiagm(
#         (ones(abn*zn), ones(abn*zn-1), ones(abn*zn-1), ones( abn*zn -bn ) , ones( abn*zn -bn ),
#         onesλ..., onesλ...),
#         (0, 1, -1, bn, -bn, colλ..., -colλ...)
#     )
#
#     fill!(nonzeros(A), zero(Float64))
#     B = deepcopy(A)
#     Λ = deepcopy(A)
#
#     # == fill up matrix Λ == #
#     ij = zero(Int)
#     λt = λ'
#     @inbounds for zi in 1:zn, abi in 1:abn
#         ij += 1
#         Λ[ abi:abn:(zn-1)*abn+abi , ij] = λt[:,zi]
#     end
#
#     V = Array(Float64, abn*zn)
#     ij = zero(Int)
#     for zi in 1:zn, ai in 1:an, bi in 1:bn
#         ij += 1
#         V[ij] = ( (1- ξ)*w*z[zi] + rᴬ*a[ai] + rᴮ⁻*b[bi] ).^( 1- γ )/( 1- γ )/ρ
#     end
#
#     # == b such that Ag = b in plank
#     g⁰ = fill(zero(Float64), abn*zn)
#     i_fix = 1
#     g⁰[i_fix] = .1
#
#     # == storage arrays == #
#     gg = Array(Float64, abn*zn)
#     u  = Array(Float64, abn*zn)
#     newV = deepcopy(V)
#
# TwoAssetsFD(twoap,
#             b, Δb, rbdrift, a, Δa, radrift, invΔ,
#             λ, z,
#             V, gg,
#             newV, u, g⁰, A, B, Λ )
# end