text:
using RecipesBase
using Parameters

"""
Plot a univariate parametric function, marking the collocation points.

`N` is the number of points for the dense grid on which the function
is calculated. When and endpoint is infinite, `padding` is used to add
some space to the nearest collocation point, specified as a fraction
of the range covered by all collocation points. `label` labels the plot.
"""
@recipe function f(pf::ParametricFunction, label = "function", N = 100, padding = 0.1)
    @unpack family, θ = pf
    dom = domain(family)
    @assert isa(dom, AbstractInterval) "Plotting not implemented for domain $(dom)."
    left, right = extrema(dom)
    zs = points(family)
    if !isfinite(dom)
        pleft, pright = extrema(zs)
        pad = (pright - pleft)*padding
        if !isfinite(left)
            left = pleft - pad
        end
        if !isfinite(right)
            right = pright + pad
        end
    end
    xs = linspace(left, right, N)
    @series begin
        seriestype := :path
        label := label
        xs, pf.(xs)
    end
    @series begin
        primary := false
        seriestype := :scatter
        zs, pf.(zs)
    end
end


text:
using Scallops
using RayTraceEllipsoids
using DataFrames
using DataFramesMeta
using CSV

retina = :dis_ret
morphz = 1.1
photoreceptor_radius = 5
absorption_coefficient = (dis_ret = 0.00667, pro_ret = 0.00667)
distance = 1e5
layers, rc, ous, nodalz, focallengths = createobj(morphz, photoreceptor_radius, absorption_coefficient)
n = 5
apertures = range(200, 400, length=n)
df = DataFrame(method = String[], aperture = Float64[], Δθ = Float64[], fraction = Float64[], geometry = Float64[], sensitivity = Float64[])
for aperture in apertures
    Δθ = Scallops.getΔΘ(distance, aperture, rc, nodalz, ous, retina)
    cosa = cos(Δθ/2)
    fraction = Scallops.raytrace_fraction(distance, aperture, rc, ous, nodalz, cosa, retina)
    geometry = Scallops.unmodified_photons(distance, cosa, aperture, rc, nodalz)
    sensitivity = fraction*geometry
    push!(df, ("ray-traced", aperture, Δθ, fraction, geometry, sensitivity))
    Δθ = Scallops.getΔΘ(focallengths[retina], photoreceptor_radius)
    geometry = Scallops.modified_photons(aperture, Δθ)
    l = layers[retina].tissue.thick
    fraction = RayTraceEllipsoids.absorption(l, absorption_coefficient[retina])
    sensitivity = fraction*geometry
    push!(df, ("modified-Land", aperture, Δθ, fraction, geometry, sensitivity))
end

CSV.write("compare.csv", df)
for method in ("ray-traced", "modified-Land")
    d = @linq df |> where(:method .== method)
    CSV.write("$method.csv", d)
end

ext: jl


text:
module TestCore
include("preamble.jl")
using InitialValues: INIT
using Transducers: has, reform

@testset "Reduced" begin
    @testset "constructor" begin
        @test Reduced(1) isa Reduced{Int}
        @test Reduced(Reduced(1)) isa Reduced{Reduced{Int}}
    end
    @testset "=" begin
        x = [1]
        y = [2]
        @test Reduced(x) !== Reduced(copy(x))
        @test Reduced(x) == Reduced(copy(x))
        @test Reduced(x) != Reduced(y)
        @test Reduced(Reduced(x)) == Reduced(Reduced(copy(x)))
        @test Reduced(Reduced(x)) != Reduced(Reduced(y))
        @test Reduced(Reduced(x)) != Reduced(x)
        @test Reduced(x) != Reduced(Reduced(x))
    end
end

@testset "reduced" begin
    @test reduced(1) === Reduced(1)
    @test reduced(reduced(1)) === Reduced(1)
end

@testset "has" begin
    @test has(Count(), Count)
    @test has(Map(identity) |> Count(), Count)
    @test has(Count() |> Map(identity), Count)
end

@testset "reform" begin
    f = Completing(Transducers.push!!)  # bottom reducing function for eduction
    @testset for xf in [
            Count(),
            Map(identity) |> Count(),
            ZipSource(Map(identity) |> Count()),
            ]
        rf = eduction(xf, 1:1).rf
        @test rf === reform(rf, f)
        @test rf !== reform(rf, +)
    end
end

@testset "reducingfunction" begin
    @test reducingfunction(Map(identity), +) isa Function
    @testset "INIT" begin
        @test reducingfunction(Map(identity), +)(INIT, :anything) === :anything
        @test reducingfunction(Map(identity), +)(:anything, INIT) === :anything

        # Test that `INIT` interacts with the bottom reducing
        # function.  The second argument should still be processed
        # with the transducer.  Note that, unlike `Map(identity)`, the
        # arguments in reversed order does not work.
        @test reducingfunction(Map(string), +)(INIT, :anything) === "anything"
    end
end

@testset "OnInit" begin
    @testset "mapfoldl" begin
        @testset for xs in iterator_variants(1:3)
            @test mapfoldl(
                Zip(Map(identity), Map(string)),
                push!,
                xs,
                init = OnInit(() -> [])
            ) == [(1, "1"), (2, "2"), (3, "3")]
        end
    end
end

@testset "broadcastable(::Transducer)" begin
    @test collect.(Enumerate(), ["abc", "def"]) ==
        [[(1, 'a'), (2, 'b'), (3, 'c')], [(1, 'd'), (2, 'e'), (3, 'f')]]
end

end  # module


text:
#################################################
#
#  Graph concepts
#
#################################################

### the root type of all graphs
abstract AbstractGraph{V, E}

vertex_type{V,E}(g::AbstractGraph{V,E}) = V
edge_type{V,E}(g::AbstractGraph{V,E}) = E

### concepts

const graph_concept_symbols = Set([
    :vertex_list,
    :edge_list,
    :vertex_map,
    :edge_map,
    :adjacency_list,
    :incidence_list,
    :bidirectional_adjacency_list,
    :bidirectional_incidence_list,
    :adjacency_matrix])

### concepts checking functions

implements_vertex_list(g::AbstractGraph) = false
implements_edge_list(g::AbstractGraph) = false

implements_vertex_map(g::AbstractGraph) = false
implements_edge_map(g::AbstractGraph) = false

implements_bidirectional_adjacency_list(g::AbstractGraph) = false
implements_bidirectional_incidence_list(g::AbstractGraph) = false

# This ensures that
# When implements_bidirectional_adjlist is set to true,
# implements_adjlist is automatically true without being overrided

implements_adjacency_list(g::AbstractGraph) = implements_bidirectional_adjacency_list(g)
implements_incidence_list(g::AbstractGraph) = implements_bidirectional_incidence_list(g)

implements_adjacency_matrix(g::AbstractGraph) = false

### macro to simplify concept declaration

function _graph_implements_code(G::Symbol, concepts::Symbol...)
    stmts = Expr[]
    for c in concepts
        (c in graph_concept_symbols) || error("Invalid concept name: $c")
        fun = symbol(string("implements_", c))
        push!(stmts, :( $(fun)(::$(G)) = true ))
    end
    Expr(:block, stmts...)
end

macro graph_implements(G, concepts...)
    esc(_graph_implements_code(G, concepts...))
end

### macro to check interface requirements

function _graph_requires_code(g::Symbol, concepts::Symbol...)
    stmts = Expr[]
    for c in concepts
        (c in graph_concept_symbols) || error("Invalid concept name: $c")
        fun = symbol(string("implements_", c))
        msg = "The graph $(g) does not implement a required concept: $(c)."
        push!(stmts, :($(fun)($g) ? nothing : error($msg)) )
    end
    Expr(:block, stmts...)
end

macro graph_requires(g, concepts...)
    esc(_graph_requires_code(g, concepts...))
end


text:
# Build a workload, run it to get kernel times, and create a datastructure encoding the
# graph for later plotting runs.
function generate_timeline(backend, f, opt, cache; use_2lm_scratchpad = false, file = "temp.jls", kw...)
    fex = AutoTM.Optimizer.factory(
        backend,
        f,
        opt;
        cache = cache,
        kw...
    ) |> maybeunwrap

    # Run twice - we'll get timing data from the second iteration.
    @time fex()

    # Reset the performance counters so they aren't "contaminated" from the original run.
    #
    # The performance counters accumulate across runs, so if we don't do this, we will
    # essentially be averaging the runtime from the second run with that of the first run.
    nGraph.reset_counters(fex.ex)
    @time fex()

    # Create a `FunctionData` object from the nGraph function.
    # Couple it with the performance metrics.
    function_data = AutoTM.Profiler.FunctionData(fex.ex.ngraph_function, backend)
    times = nGraph.get_performance(fex.ex)

    println("Sum of Kernel Run Times: $(sum(values(times)))")

    # Record metadata about each tensor that will be used for plot generation.
    tensor_records = map(collect(AutoTM.Profiler.tensors(function_data))) do tensor
        users = AutoTM.Profiler.users(tensor)
        return Dict(
             "name" => nGraph.name(tensor),
             "users" => nGraph.name.(users),
             "user_indices" => getproperty.(users, :index),
             "sizeof" => sizeof(tensor),
             "offset" => Int(AutoTM.Profiler.getoffset(tensor)),
        )
    end

    # Collect metadata on each node.
    node_records = map(AutoTM.Profiler.nodes(function_data)) do node
        outputs = nGraph.name.(AutoTM.Profiler.outputs(node))
        inputs = nGraph.name.(AutoTM.Profiler.inputs(node))
        time = get(times, nGraph.name(node), 0)
        delete!(times, nGraph.name(node))
        return Dict(
            "name" => nGraph.name(node),
            "inputs" => inputs,
            "outputs" => outputs,
            "time" => time,
        )
    end

    @assert isempty(times)

    # Combine together to the final record that will be saved.
    record = Dict(
        "tensors" => tensor_records,
        "nodes" => node_records,
    )

    # file = make_filename(
    #     workload,
    #     parsed_args["mode"],
    #     "",
    #     parsed_args["use_2lm_scratchpad"];
    #     prefix = joinpath(@__DIR__, "serialized")
    # )
    save(file, record)
    return nothing
end

# Tool for inspecting and generating the heap allocation of an executable.
#
# The modified nGraph allows us to pass two Julia callbacks to the compilation process.
# The first callback will be called before the memory allocation pass, and the second will
# be called AFTER memory allocatio pass.
#
# Since our graphs may exceed the physical memory of a system - we will just insert the
# extraction of the intermediate tensors from the second callback.
struct TensorRecord
    start::Float64
    stop::Float64
    offset::Float64
    height::Float64

    # Indices of reading and writing
    read_starts::Vector{Float64}
    read_stops::Vector{Float64}
    write_stop::Float64
end

function Base.isless(a::TensorRecord, b::TensorRecord)
    # First - order by starting index.
    if a.start < b.start
        return true

    # Then, order by offset
    elseif a.start == b.start
        if a.offset < b.offset
            return true
        end
    end
    return false
end

# PGF Draw Rectangle
_draw_rect(x0, y0, x1, y1; fill = "black") = """
    \\draw [fill = $fill, $fill, thick]
        (axis cs:$x0,$y0) rectangle (axis cs:$x1,$y1);
    """
function _draw_rect(r::TensorRecord; read_color = :red, write_color = :blue)
    # First, we draw a large black rectangle for the full lifetime of the tensor.
    first = _draw_rect(r.start, r.offset, r.stop, r.offset + r.height)

    # Then, we emit red rectangles for each user. By emitting these afterwards, they should
    # be drawn on top of the the black rectangles.
    reads = map(zip(r.read_starts, r.read_stops)) do x
        start = x[1]
        stop = x[2]
        return _draw_rect(
            start,
            r.offset,
            stop,
            r.offset + r.height;
            fill = read_color
        )
    end
    write = _draw_rect(
        r.start,
        r.offset,
        r.write_stop,
        r.offset + r.height;
        fill = write_color
    )

    return vcat(first, reads, write)
end

function heap_plot(path)
    # Create a record for each tensor in the compiled graph.
    tensor_records = get_records(path)
    sort!(tensor_records)

    # Plot a rectangle for each tensor.
    rectangle_strings = Iterators.flatten(_draw_rect.(tensor_records))

    # Get the axis min and maxes
    ymin = 0
    ymax = maximum(x -> x.offset + x.height, tensor_records)
    xmin = 0
    xmax = maximum(x -> x.start, tensor_records)

    axs = @pgf Axis(
        {
            width = "20cm",
            height = "14cm",
            ymin = ymin,
            ymax = ymax,
            xmin = xmin,
            xmax = xmax,
            xlabel = "Time (s)",
            ylabel = "Memory Position (GB)",
        },
        rectangle_strings...,
    )

    name, _ = splitext(basename(path))

    pgfsave(joinpath(FIGDIR, "$name.pdf"), axs)
    return axs
end

function get_records(path)
    d = deserialize(path)

    tensors = d["tensors"]
    nodes = d["nodes"]

    # Generate aggregate runtimes.
    times = [node["time"] for node in nodes]
    cumsum!(times, times)

    # Convert time from microseconds to seconds
    times = times ./ 1E6
    @show times[end]

    # Add one extra item to the end of `times` so we can index one past the original lenght
    # when calculating stop times of reading nodes.
    #
    # This is not 100% correct but close enough.
    push!(times, times[end])

    # Extract all the things!! :D
    tensor_records = map(tensors) do tensor
        users = tensor["user_indices"]
        user_times = times[users]

        start = first(user_times)
        stop = last(user_times)

        read_starts = user_times
        read_stops = times[users .+ 1]

        write_stop = times[first(users) + 1]

        return TensorRecord(
            start,
            stop,
            tensor["offset"] / 1E9,
            tensor["sizeof"] / 1E9,
            read_starts,
            read_stops,
            write_stop
        )
    end

    return tensor_records
end

#####
##### Plot summary statistics
#####

function barplot(records::OrderedDict{String, <:Dict}, ks;
        modifier = identity,
        reducer = ssum,
        ylabel = "",
        title = "",
    )
    # For each entry in outer dict, take the total sum of the values of interest for each
    # inner dict.
    records = map(collect(keys(records))) do name
        v = records[name]
        return name => Dict(k => reducer(v, k) for k in ks)
    end |> OrderedDict

    # Bar plots to add.
    plots = []

    for (name, data) in records
        # Generate the `x` and `y` coordinates for this sample
        x = replace_.(ks)
        y = [modifier(data[k]) for k in ks]

        # Emit the plot
        push!(plots,
            @pgf(PlotInc(
                Coordinates(x, y),
            )),
            @pgf(LegendEntry(name)),
        )
    end

    axs = @pgf Axis(
        {
            ybar,
            symbolic_x_coords = replace_.(ks),
            enlarge_x_limits = 0.2,
            xtick="data",
            xticklabel_style = {
                rotate = 10,
            },
            ymajorgrids,
            title = titlecase(replace_(title)),
            ylabel = ylabel,
        },
        plots...
    )
    return axs
end

max_stars_repo_name: darchr/AutoTM

max_issues_repo_name: darchr/AutoTM

max_forks_repo_name: darchr/AutoTM


text:
#####
##### Upwind-biased 3rd-order advection scheme
#####

"""
    struct UpwindBiasedFifthOrder <: AbstractUpwindBiasedAdvectionScheme{2}

Upwind-biased fifth-order advection scheme.
"""
struct UpwindBiasedFifthOrder <: AbstractUpwindBiasedAdvectionScheme{2} end

const U5 = UpwindBiasedFifthOrder

@inline boundary_buffer(::U5) = 2

@inline symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, ::U5, c) = symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, centered_fourth_order, c)
@inline symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, ::U5, c) = symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, centered_fourth_order, c)
@inline symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, ::U5, c) = symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, centered_fourth_order, c)

@inline symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, ::U5, u) = symmetric_interpolate_xᶜᵃᵃ(i, j, k, grid, centered_fourth_order, u)
@inline symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, ::U5, v) = symmetric_interpolate_yᵃᶜᵃ(i, j, k, grid, centered_fourth_order, v)
@inline symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, ::U5, w) = symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, centered_fourth_order, w)

@inline left_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::U5, c) = @inbounds (- 3 * c[i+1, j, k] + 27 * c[i, j, k] + 47 * c[i-1, j, k] - 13 * c[i-2, j, k] + 2 * c[i-3, j, k]) / 60
@inline left_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::U5, c) = @inbounds (- 3 * c[i, j+1, k] + 27 * c[i, j, k] + 47 * c[i, j-1, k] - 13 * c[i, j-2, k] + 2 * c[i, j-3, k]) / 60
@inline left_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::U5, c) = @inbounds (- 3 * c[i, j, k+1] + 27 * c[i, j, k] + 47 * c[i, j, k-1] - 13 * c[i, j, k-2] + 2 * c[i, j, k-3]) / 60

@inline left_biased_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme::U5, u) = left_biased_interpolate_xᶠᵃᵃ(i+1, j, k, grid, scheme, u)
@inline left_biased_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme::U5, v) = left_biased_interpolate_yᵃᶠᵃ(i, j+1, k, grid, scheme, v)
@inline left_biased_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme::U5, w) = left_biased_interpolate_zᵃᵃᶠ(i, j, k+1, grid, scheme, w)

@inline right_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::U5, c) = @inbounds (2 * c[i+2, j, k] - 13 * c[i+1, j, k] + 47 * c[i, j, k] + 27 * c[i-1, j, k] - 3 * c[i-2, j, k] ) / 60
@inline right_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::U5, c) = @inbounds (2 * c[i, j+2, k] - 13 * c[i, j+1, k] + 47 * c[i, j, k] + 27 * c[i, j-1, k] - 3 * c[i, j-2, k] ) / 60
@inline right_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::U5, c) = @inbounds (2 * c[i, j, k+2] - 13 * c[i, j, k+1] + 47 * c[i, j, k] + 27 * c[i, j, k-1] - 3 * c[i, j, k-2] ) / 60

@inline right_biased_interpolate_xᶜᵃᵃ(i, j, k, grid, scheme::U5, u) = right_biased_interpolate_xᶠᵃᵃ(i+1, j, k, grid, scheme, u)
@inline right_biased_interpolate_yᵃᶜᵃ(i, j, k, grid, scheme::U5, v) = right_biased_interpolate_yᵃᶠᵃ(i, j+1, k, grid, scheme, v)
@inline right_biased_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme::U5, w) = right_biased_interpolate_zᵃᵃᶠ(i, j, k+1, grid, scheme, w)


text:
@testset "Test read_time_series_metadata" begin
    file = joinpath(FORECASTS_DIR, "ComponentsAsColumnsNoTime.json")
    forecasts = IS.read_time_series_metadata(file)
    @test length(forecasts) == 1

    for forecast in forecasts
        @test isfile(forecast.data_file)
    end
end

@testset "Test add_forecast from file" begin
    data = IS.SystemData()

    name = "Component1"
    component = IS.TestComponent(name, 5)
    IS.add_component!(data, component)
    @test !IS.has_forecasts(component)

    file = joinpath(FORECASTS_DIR, "ComponentsAsColumnsNoTime.json")
    IS.add_forecasts!(IS.InfrastructureSystemsType, data, file)
    @test IS.has_forecasts(component)

    forecasts = get_all_forecasts(data)
    @test length(forecasts) == 1
    forecast = forecasts[1]
    @test forecast isa IS.Deterministic

    forecast2 = IS.get_forecast(
        typeof(forecast),
        component,
        IS.get_initial_time(forecast),
        IS.get_label(forecast),
    )
    @test IS.get_horizon(forecast) == IS.get_horizon(forecast2)
    @test IS.get_initial_time(forecast) == IS.get_initial_time(forecast2)

    it = IS.get_initial_time(forecast)

    forecasts = get_all_forecasts(data)
    @test length(collect(forecasts)) == 1

    @test IS.get_forecast_initial_times(data) == [it]
    unique_its = Set{Dates.DateTime}()
    IS.get_forecast_initial_times!(unique_its, component) == [it]
    @test collect(unique_its) == [it]
    @test IS.get_forecasts_initial_time(data) == it
    @test IS.get_forecasts_interval(data) == IS.UNINITIALIZED_PERIOD
    @test IS.get_forecasts_horizon(data) == IS.get_horizon(forecast)
    @test IS.get_forecasts_resolution(data) == IS.get_resolution(forecast)
end

@testset "Test add_forecast" begin
    sys = IS.SystemData()
    name = "Component1"
    component_val = 5
    component = IS.TestComponent(name, component_val)
    IS.add_component!(sys, component)

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)
    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    forecast = IS.Deterministic("get_val", ta)
    IS.add_forecast!(sys, component, forecast)
    forecast = IS.get_forecast(IS.Deterministic, component, dates[1], "get_val")
    @test forecast isa IS.Deterministic

    name = "Component2"
    component2 = IS.TestComponent(name, component_val)
    @test_throws ArgumentError IS.add_forecast!(sys, component2, forecast)

    # The component name will exist but not the component.
    component3 = IS.TestComponent(name, component_val)
    @test_throws ArgumentError IS.add_forecast!(sys, component3, forecast)
end

# TODO: this is disabled because PowerSystems currently does not set labels correctly.
#@testset "Test add_forecast bad label" begin
#    sys = IS.SystemData()
#    name = "Component1"
#    component_val = 5
#    component = IS.TestComponent(name, component_val)
#    IS.add_component!(sys, component)
#
#    dates = collect(Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1) :
#                    Dates.DateTime("2020-01-01T23:00:00"))
#    data = collect(1:24)
#    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
#    forecast = IS.Deterministic("bad-label", ta)
#    @test_throws ArgumentError IS.add_forecast!(sys, component, forecast)
#end

@testset "Test add_forecast from TimeArray" begin
    sys = IS.SystemData()
    name = "Component1"
    component_val = 5
    component = IS.TestComponent(name, component_val)
    IS.add_component!(sys, component)

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)
    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    IS.add_forecast!(sys, ta, component, "get_val")
    forecast = IS.get_forecast(IS.Deterministic, component, dates[1], "get_val")
    @test forecast isa IS.Deterministic
end

@testset "Test forecast initial times" begin
    sys = IS.SystemData()

    @test_throws ArgumentError IS.get_forecasts_initial_time(sys)
    @test_throws ArgumentError IS.get_forecasts_last_initial_time(sys)

    dates1 = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    dates2 = collect(
        Dates.DateTime("2020-01-02T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-02T23:00:00"),
    )
    data = collect(1:24)
    components = []

    for i in 1:2
        name = "Component" * string(i)
        component = IS.TestComponent(name, i)
        IS.add_component!(sys, component)
        push!(components, component)
        if i == 1
            dates1_ = dates1
            dates2_ = dates2
        else
            dates1_ = dates1 .+ Dates.Hour(1)
            dates2_ = dates2 .+ Dates.Hour(1)
        end
        ta1 = TimeSeries.TimeArray(dates1_, data, [IS.get_name(component)])
        ta2 = TimeSeries.TimeArray(dates2_, data, [IS.get_name(component)])
        IS.add_forecast!(sys, ta1, component, "get_val")
        IS.add_forecast!(sys, ta2, component, "get_val")
    end

    initial_times = IS.get_forecast_initial_times(sys)
    @test length(initial_times) == 4

    first_initial_time = dates1[1]
    last_initial_time = dates2[1] + Dates.Hour(1)
    @test IS.get_forecasts_initial_time(sys) == first_initial_time
    @test IS.get_forecasts_last_initial_time(sys) == last_initial_time

    @test_logs(
        (:error, r"initial times don't match"),
        @test !IS.validate_forecast_consistency(sys)
    )
    @test_logs(
        (:error, r"initial times don't match"),
        @test_throws IS.DataFormatError !IS.check_forecast_consistency(sys)
    )

    @test IS.get_forecast_counts(sys) == (2, 4)

    IS.clear_forecasts!(sys)
    for component in components
        ta1 = TimeSeries.TimeArray(dates1, data, [IS.get_name(component)])
        ta2 = TimeSeries.TimeArray(dates2, data, [IS.get_name(component)])
        IS.add_forecast!(sys, ta1, component, "get_val")
        IS.add_forecast!(sys, ta2, component, "get_val")
    end

    expected = [dates1[1], dates2[1]]
    for component in components
        @test IS.get_forecast_initial_times(IS.Deterministic, component) == expected
    end

    @test IS.validate_forecast_consistency(sys)
    IS.get_forecasts_interval(sys) == dates2[1] - dates1[1]
end

@testset "Test remove_forecasts" begin
    data = create_system_data(; with_forecasts = true)
    components = collect(IS.iterate_components(data))
    @test length(components) == 1
    component = components[1]
    forecasts = get_all_forecasts(data)
    @test length(get_all_forecasts(data)) == 1

    forecast = forecasts[1]
    IS.remove_forecast!(
        typeof(forecast),
        data,
        component,
        IS.get_initial_time(forecast),
        IS.get_label(forecast),
    )

    @test length(get_all_forecasts(data)) == 0
    @test IS.get_num_time_series(data.time_series_storage) == 0
end

@testset "Test clear_forecasts" begin
    data = create_system_data(; with_forecasts = true)
    IS.clear_forecasts!(data)
    @test length(get_all_forecasts(data)) == 0
end

@testset "Test that remove_component removes forecasts" begin
    data = create_system_data(; with_forecasts = true)

    components = collect(IS.get_components(IS.InfrastructureSystemsType, data))
    @test length(components) == 1
    component = components[1]

    forecasts = collect(IS.iterate_forecasts(data))
    @test length(forecasts) == 1
    forecast = forecasts[1]
    @test IS.get_num_time_series(data.time_series_storage) == 1

    IS.remove_component!(data, component)
    @test length(collect(IS.iterate_forecasts(component))) == 0
    @test length(collect(IS.get_components(IS.InfrastructureSystemsType, data))) == 0
    @test length(get_all_forecasts(data)) == 0
    @test IS.get_num_time_series(data.time_series_storage) == 0
end

@testset "Test get_forecast_values" begin
    sys = IS.SystemData()
    name = "Component1"
    component_val = 5
    component = IS.TestComponent(name, component_val)
    IS.add_component!(sys, component)

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)
    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    IS.add_forecast!(sys, ta, component, "get_val")
    forecast = IS.get_forecast(IS.Deterministic, component, dates[1], "get_val")

    # Test both versions of the function.
    vals = IS.get_forecast_values(IS, component, forecast)
    @test TimeSeries.timestamp(vals) == dates
    @test TimeSeries.values(vals) == data .* component_val

    vals2 = IS.get_forecast_values(IS.Deterministic, IS, component, dates[1], "get_val")
    @test TimeSeries.timestamp(vals2) == dates
    @test TimeSeries.values(vals2) == data .* component_val
end

@testset "Test get subset of forecast" begin
    sys = create_system_data()
    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)

    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    IS.add_forecast!(sys, ta, component, "get_val")

    forecast = IS.get_forecast(IS.Deterministic, component, dates[1], "get_val")
    @test TimeSeries.timestamp(IS.get_data(forecast))[1] == dates[1]

    forecast = IS.get_forecast(IS.Deterministic, component, dates[3], "get_val", 3)
    @test TimeSeries.timestamp(IS.get_data(forecast))[1] == dates[3]
    @test length(forecast) == 3
end

@testset "Test copy forecasts no label mapping" begin
    sys = create_system_data()
    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    initial_time = Dates.DateTime("2020-01-01T00:00:00")
    dates = collect(initial_time:Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"))
    data = collect(1:24)

    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    label = "get_val"
    IS.add_forecast!(sys, ta, component, label)

    component2 = IS.TestComponent("component2", 6)
    IS.add_component!(sys, component2)
    IS.copy_forecasts!(component, component2)
    forecast = IS.get_forecast(IS.Deterministic, component2, initial_time, label)
    @test forecast isa IS.Deterministic
    @test IS.get_initial_time(forecast) == initial_time
    @test IS.get_label(forecast) == label
end

@testset "Test copy forecasts label mapping" begin
    sys = create_system_data()
    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    initial_time = Dates.DateTime("2020-01-01T00:00:00")
    dates = collect(initial_time:Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"))
    data = collect(1:24)

    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    label1 = "get_val1"
    IS.add_forecast!(sys, ta, component, label1)

    component2 = IS.TestComponent("component2", 6)
    IS.add_component!(sys, component2)
    label2 = "get_val2"
    label_mapping = Dict(label1 => label2)
    IS.copy_forecasts!(component, component2, label_mapping)
    forecast = IS.get_forecast(IS.Deterministic, component2, initial_time, label2)
    @test forecast isa IS.Deterministic
    @test IS.get_initial_time(forecast) == initial_time
    @test IS.get_label(forecast) == label2
end

@testset "Test copy forecasts label mapping, missing label" begin
    sys = create_system_data()
    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    initial_time1 = Dates.DateTime("2020-01-01T00:00:00")
    end_time1 = Dates.DateTime("2020-01-01T23:00:00")
    dates1 = collect(initial_time1:Dates.Hour(1):end_time1)
    initial_time2 = Dates.DateTime("2020-01-02T00:00:00")
    end_time2 = Dates.DateTime("2020-01-02T23:00:00")
    dates2 = collect(initial_time2:Dates.Hour(1):end_time2)
    data = collect(1:24)

    ta1 = TimeSeries.TimeArray(dates1, data, [IS.get_name(component)])
    ta2 = TimeSeries.TimeArray(dates2, data, [IS.get_name(component)])
    label1 = "get_val1"
    label2a = "get_val2a"
    IS.add_forecast!(sys, ta1, component, label1)
    IS.add_forecast!(sys, ta2, component, label2a)

    component2 = IS.TestComponent("component2", 6)
    IS.add_component!(sys, component2)
    label2b = "get_val2b"
    label_mapping = Dict(label2a => label2b)
    IS.copy_forecasts!(component, component2, label_mapping)
    forecast = IS.get_forecast(IS.Deterministic, component2, initial_time2, label2b)
    @test forecast isa IS.Deterministic
    @test IS.get_initial_time(forecast) == initial_time2
    @test IS.get_label(forecast) == label2b
    @test_throws ArgumentError IS.get_forecast(
        IS.Deterministic,
        component2,
        initial_time2,
        label2a,
    )
end

function validate_generated_initial_times(
    forecast_type::Type{<:IS.Forecast},
    component::IS.InfrastructureSystemsType,
    label::AbstractString,
    horizon::Int,
    initial_times::Vector{Dates.DateTime},
    initial_time::Dates.DateTime,
    interval::Dates.Period,
    exp_length::Int,
)
    @test length(initial_times) == exp_length
    for it in initial_times
        @test it == initial_time
        # Verify all possible forecast ranges.
        for i in 2:horizon
            forecast = IS.get_forecast(forecast_type, component, it, label, i)
            @test IS.get_horizon(forecast) == i
            @test IS.get_initial_time(forecast) == it
            # This will throw if the resolution isn't consistent throughout.
            IS.get_resolution(IS.get_data(forecast))
        end
        initial_time += interval
    end
end

@testset "Test subset from contiguous forecasts" begin
    sys = create_system_data()

    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    dates1 = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    dates2 = collect(
        Dates.DateTime("2020-01-02T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-02T23:00:00"),
    )
    dates3 = collect(
        Dates.DateTime("2020-01-03T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-03T23:00:00"),
    )
    data = collect(1:24)

    label = "get_val"
    ta1 = TimeSeries.TimeArray(dates1, data, [IS.get_name(component)])
    ta2 = TimeSeries.TimeArray(dates2, data, [IS.get_name(component)])
    ta3 = TimeSeries.TimeArray(dates3, data, [IS.get_name(component)])
    IS.add_forecast!(sys, ta1, component, label)
    IS.add_forecast!(sys, ta2, component, label)
    IS.add_forecast!(sys, ta3, component, label)
    initial_times = IS.get_forecast_initial_times(component)
    @test length(initial_times) == 3
    @test IS.are_forecasts_contiguous(component)
    @test IS.are_forecasts_contiguous(sys)

    interval = Dates.Hour(1)
    horizon = 55
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates1[1],
        interval,
        18,
    )

    invalid_it = Dates.DateTime("2020-01-20T00:00:00")
    @test_throws ArgumentError IS.get_forecast(
        IS.Deterministic,
        component,
        invalid_it,
        label,
        horizon,
    )
end

@testset "Test generate_initial_times" begin
    sys = create_system_data()
    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)

    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    label = "get_val"
    IS.add_forecast!(sys, ta, component, label)
    initial_times = IS.get_forecast_initial_times(component)
    @test length(initial_times) == 1

    horizon = 24
    interval = Dates.Hour(1)
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates[1],
        interval,
        1,
    )

    horizon = 12
    interval = Dates.Hour(1)
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates[1],
        interval,
        13,
    )

    horizon = 6
    interval = Dates.Hour(3)
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates[1],
        interval,
        7,
    )

    horizon = 6
    interval = Dates.Hour(4)
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates[1],
        interval,
        5,
    )

    # Test through the system.
    horizon = 6
    initial_times = IS.generate_initial_times(sys, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates[1],
        interval,
        5,
    )
    IS.clear_forecasts!(sys)
    @test_throws ArgumentError IS.generate_initial_times(sys, interval, 6)
end

@testset "Test generate_initial_times contiguous" begin
    sys = create_system_data()

    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    @test_throws ArgumentError IS.are_forecasts_contiguous(component)
    @test_throws ArgumentError IS.are_forecasts_contiguous(sys)

    dates1 = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    dates2 = collect(
        Dates.DateTime("2020-01-02T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-02T23:00:00"),
    )
    data = collect(1:24)

    label = "get_val"
    ta1 = TimeSeries.TimeArray(dates1, data, [IS.get_name(component)])
    ta2 = TimeSeries.TimeArray(dates2, data, [IS.get_name(component)])
    IS.add_forecast!(sys, ta1, component, label)
    IS.add_forecast!(sys, ta2, component, label)
    initial_times = IS.get_forecast_initial_times(component)
    @test length(initial_times) == 2
    @test IS.are_forecasts_contiguous(component)
    @test IS.are_forecasts_contiguous(sys)

    interval = Dates.Hour(1)
    horizon = 48
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates1[1],
        interval,
        1,
    )

    horizon = 24
    interval = Dates.Hour(1)
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates1[1],
        interval,
        25,
    )

    horizon = 12
    interval = Dates.Hour(1)
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates1[1],
        interval,
        37,
    )

    horizon = 6
    interval = Dates.Hour(3)
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates1[1],
        interval,
        15,
    )

    horizon = 6
    interval = Dates.Hour(4)
    initial_times = IS.generate_initial_times(component, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates1[1],
        interval,
        11,
    )

    # Run once on the system.
    horizon = 6
    initial_times = IS.generate_initial_times(sys, interval, horizon)
    validate_generated_initial_times(
        IS.Deterministic,
        component,
        label,
        horizon,
        initial_times,
        dates1[1],
        interval,
        11,
    )
end

@testset "Test generate_initial_times overlapping" begin
    sys = create_system_data()

    @test_throws ArgumentError IS.generate_initial_times(sys, Dates.Hour(3), 6)

    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    dates1 = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    dates2 = collect(
        Dates.DateTime("2020-01-01T01:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-02T00:00:00"),
    )
    data = collect(1:24)

    ta1 = TimeSeries.TimeArray(dates1, data, [IS.get_name(component)])
    ta2 = TimeSeries.TimeArray(dates2, data, [IS.get_name(component)])
    IS.add_forecast!(sys, ta1, component, "get_val")
    @test_throws IS.ConflictingInputsError IS.generate_initial_times(
        sys,
        Dates.Minute(30),
        6,
    )

    IS.add_forecast!(sys, ta2, component, "get_val")
    @test_throws ArgumentError IS.generate_initial_times(sys, Dates.Hour(3), 6)

    @test !IS.are_forecasts_contiguous(component)
    @test !IS.are_forecasts_contiguous(sys)
end

@testset "Test generate_initial_times non-contiguous" begin
    sys = create_system_data()

    @test_throws ArgumentError IS.generate_initial_times(sys, Dates.Hour(3), 6)

    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    dates1 = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    # Skip 1 hour.
    dates2 = collect(
        Dates.DateTime("2020-01-02T01:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-03T00:00:00"),
    )
    data = collect(1:24)

    ta1 = TimeSeries.TimeArray(dates1, data, [IS.get_name(component)])
    ta2 = TimeSeries.TimeArray(dates2, data, [IS.get_name(component)])
    IS.add_forecast!(sys, ta1, component, "get_val")
    @test_throws IS.ConflictingInputsError IS.generate_initial_times(
        sys,
        Dates.Minute(30),
        6,
    )

    IS.add_forecast!(sys, ta2, component, "get_val")
    @test_throws ArgumentError IS.generate_initial_times(sys, Dates.Hour(3), 6)

    @test !IS.are_forecasts_contiguous(component)
    @test !IS.are_forecasts_contiguous(sys)
end

@testset "Test generate_initial_times offset from first initial_time" begin
    sys = create_system_data()

    components = collect(IS.get_components(IS.InfrastructureSystemsType, sys))
    @test length(components) == 1
    component = components[1]

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)

    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    IS.add_forecast!(sys, ta, component, "get_val")
    resolution = IS.get_forecasts_resolution(sys)
    initial_times = IS.get_forecast_initial_times(component)
    @test length(initial_times) == 1

    horizon = 6
    offset = 1
    interval = resolution * 2
    initial_time = initial_times[1] + offset * resolution

    expected = collect(
        Dates.DateTime("2020-01-01T01:00:00"):interval:Dates.DateTime("2020-01-01T17:00:00"),
    )

    actual =
        IS.generate_initial_times(component, interval, horizon; initial_time = initial_time)
    @test actual == expected

    # Repeat on the system.
    actual = IS.generate_initial_times(sys, interval, horizon; initial_time = initial_time)
    @test actual == expected
end

@testset "Test component-forecast being added to multiple systems" begin
    sys1 = IS.SystemData()
    sys2 = IS.SystemData()
    name = "Component1"
    component = IS.TestComponent(name, 5)
    IS.add_component!(sys1, component)

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)
    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    IS.add_forecast!(sys1, ta, component, "get_val")

    @test_throws ArgumentError IS.add_component!(sys1, component)
end

@testset "Summarize forecasts" begin
    data = create_system_data(; with_forecasts = true)
    summary(devnull, data.forecast_metadata)
end

@testset "Test forecast forwarding methods" begin
    data = create_system_data(; with_forecasts = true)
    forecast = get_all_forecasts(data)[1]

    # Iteration
    size = 24
    @test length(forecast) == size
    i = 0
    for x in forecast
        i += 1
    end
    @test i == size

    # Indexing
    @test length(forecast[1:16]) == 16

    # when
    fcast = IS.when(forecast, TimeSeries.hour, 3)
    @test length(fcast) == 1
end

@testset "Test forecast head" begin
    data = create_system_data(; with_forecasts = true)
    forecast = get_all_forecasts(data)[1]
    fcast = IS.head(forecast)
    # head returns a length of 6 by default, but don't hard-code that.
    @test length(fcast) < length(forecast)

    fcast = IS.head(forecast, 10)
    @test length(fcast) == 10
end

@testset "Test forecast tail" begin
    data = create_system_data(; with_forecasts = true)
    forecast = get_all_forecasts(data)[1]
    fcast = IS.tail(forecast)
    # tail returns a length of 6 by default, but don't hard-code that.
    @test length(fcast) < length(forecast)

    fcast = IS.head(forecast, 10)
    @test length(fcast) == 10
end

@testset "Test forecast from" begin
    data = create_system_data(; with_forecasts = true)
    forecast = get_all_forecasts(data)[1]
    start_time = Dates.DateTime(Dates.today()) + Dates.Hour(3)
    fcast = IS.from(forecast, start_time)
    @test length(fcast) == 21
    @test TimeSeries.timestamp(IS.get_data(fcast))[1] == start_time
end

@testset "Test forecast from" begin
    data = create_system_data(; with_forecasts = true)
    forecast = get_all_forecasts(data)[1]
    for end_time in (
        Dates.DateTime(Dates.today()) + Dates.Hour(15),
        Dates.DateTime(Dates.today()) + Dates.Hour(15) + Dates.Minute(5),
    )
        fcast = IS.to(forecast, end_time)
        @test length(fcast) == 16
        @test TimeSeries.timestamp(IS.get_data(fcast))[end] <= end_time
    end
end

@testset "Test ScenarioBased forecasts" begin
    sys = IS.SystemData()
    name = "Component1"
    label = "get_val"
    component = IS.TestComponent(name, 5)
    IS.add_component!(sys, component)

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = ones(24, 2)
    ta = TimeSeries.TimeArray(dates, data)
    forecast = IS.ScenarioBased(label, ta)
    fdata = IS.get_data(forecast)
    @test length(TimeSeries.colnames(fdata)) == 2
    @test TimeSeries.timestamp(ta) == TimeSeries.timestamp(fdata)
    @test TimeSeries.values(ta) == TimeSeries.values(fdata)

    IS.add_forecast!(sys, component, forecast)
    forecast2 = IS.get_forecast(IS.ScenarioBased, component, dates[1], label)
    @test forecast2 isa IS.ScenarioBased
    fdata2 = IS.get_data(forecast2)
    @test length(TimeSeries.colnames(fdata2)) == 2
    @test TimeSeries.timestamp(ta) == TimeSeries.timestamp(fdata2)
    @test TimeSeries.values(ta) == TimeSeries.values(fdata2)

    no_forecast = 3
    forecast3 = IS.get_forecast(IS.ScenarioBased, component, dates[1], label, no_forecast)
    @test forecast3 isa IS.ScenarioBased
    fdata3 = IS.get_data(forecast3)
    @test length(TimeSeries.colnames(fdata3)) == 2
    @test TimeSeries.timestamp(ta)[1:no_forecast] == TimeSeries.timestamp(fdata3)
    @test TimeSeries.values(ta)[1:no_forecast, :] == TimeSeries.values(fdata3)
end

@testset "Test PiecewiseFunction forecasts" begin
    sys = IS.SystemData()
    name = "Component1"
    label = "get_val"
    component = IS.TestComponent(name, 5)
    IS.add_component!(sys, component)

    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = ones(24, 4)
    name = collect(Iterators.flatten([
        (Symbol("cost_bp$(ix)"), Symbol("load_bp$ix")) for ix in 1:2
    ]))
    ta = TimeSeries.TimeArray(dates, data, name)
    forecast = IS.PiecewiseFunction(label, ta)
    fdata = IS.get_data(forecast)
    @test length(TimeSeries.colnames(fdata)) == 4
    @test TimeSeries.timestamp(ta) == TimeSeries.timestamp(fdata)
    @test TimeSeries.values(ta) == TimeSeries.values(fdata)

    IS.add_forecast!(sys, component, forecast)
    forecast2 = IS.get_forecast(IS.PiecewiseFunction, component, dates[1], label)
    @test forecast2 isa IS.PiecewiseFunction
    fdata2 = IS.get_data(forecast2)
    @test length(TimeSeries.colnames(fdata2)) == 4
    @test TimeSeries.timestamp(ta) == TimeSeries.timestamp(fdata2)
    @test TimeSeries.values(ta) == TimeSeries.values(fdata2)

    no_forecast = 4
    forecast3 =
        IS.get_forecast(IS.PiecewiseFunction, component, dates[1], label, no_forecast)
    @test forecast3 isa IS.PiecewiseFunction
    fdata3 = IS.get_data(forecast3)
    @test length(TimeSeries.colnames(fdata3)) == 4
    @test TimeSeries.timestamp(ta)[1:no_forecast] == TimeSeries.timestamp(fdata3)
    @test TimeSeries.values(ta)[1:no_forecast, :] == TimeSeries.values(fdata3)
end

@testset "Add forecast to unsupported struct" begin
    struct TestComponentNoForecasts <: IS.InfrastructureSystemsType
        name::AbstractString
        internal::IS.InfrastructureSystemsInternal
    end

    function TestComponentNoForecasts(name)
        return TestComponentNoForecasts(name, IS.InfrastructureSystemsInternal())
    end

    sys = IS.SystemData()
    name = "component"
    component = TestComponentNoForecasts(name)
    IS.add_component!(sys, component)
    dates = collect(
        Dates.DateTime("2020-01-01T00:00:00"):Dates.Hour(1):Dates.DateTime("2020-01-01T23:00:00"),
    )
    data = collect(1:24)
    ta = TimeSeries.TimeArray(dates, data, [IS.get_name(component)])
    forecast = IS.Deterministic("get_val", ta)
    @test_throws ArgumentError IS.add_forecast!(sys, component, forecast)
end


text:
# using Revise
using ShaleDrillingLikelihood
using ShaleDrillingLikelihood.SDLParameters

# detect if using SLURM
if "SLURM_JOBID" in keys(ENV)
    SLURM_JOBID = ENV["SLURM_JOBID"]
    using ClusterManagers
else
    SLURM_JOBID = ""
end

using CountPlus, Distributed, JLD2
using ShaleDrillingLikelihood: value_function, EVobj, cost
using SparseArrays: nonzeros
using Formatting: generate_formatter
using Dates: today
using RCall

num_nonzeros(x) = length(nonzeros(x))
intstr = generate_formatter("%'d")

using ShaleDrillingLikelihood: learn, royalty, constr, tech, tax,
    PerfectInfo, MaxLearning, drill,
    simulationPrimitives_information,
    doSimulations, average_cost_df,
    Theta_NoTech, theta_revenue,
    time_idx, DateQuarter

# ------------------- number of simulations ----------------------

pargs = parse_commandline_counterfactuals()
print_parsed_args_counterfactuals(pargs)

DO_PAR         = !pargs["noPar"]
JLD2FILE       = pargs["jld2"]
DATE_STOP      = pargs["dateStop"]
TECH_YEAR_ZERO = pargs["techYearZero"]
RFILEDIR       = pargs["rFileDir"]

# --------------- create data ---------------

if "SLURM_JOBID" in keys(ENV)
    DATADIR = "/home/magerton/haynesville/intermediate_data"
else
    DATADIR = "E:/projects/haynesville/intermediate_data"
end

# JLD2FILE = "E:/projects/ShaleDrillingResults/2020-01/17493674/estimation-results-17493674-no-rigs-data_all_leases.jld2"
# JLD2FILE = "E:/projects/ShaleDrillingResults/2020-01/17493743/estimation-results-17493743-WITH-rigs-data_all_leases.jld2"
# JLD2FILE = "E:/projects/ShaleDrillingResults/2020-02/18273309/estimation-results-18273309-no-rigs-dbl-timefe-data_all_leases.jld2"

println_time_flush("Loading results from $JLD2FILE")
file = jldopen(JLD2FILE, "r")
    DATAPATH = file["DATAPATH"]
    M        = file["M"]
    ddmnovf  = file["ddm_novf"]
    theta    = file["theta1"]
close(file)
println_time_flush("Done!")

REWARD = reward(ddmnovf)
ddm = DynamicDrillModel(ddmnovf, REWARD)

# --------------- create data ---------------

# load in data from disk
rdatapath = joinpath(DATADIR, DATAPATH)
println_time_flush("loading $rdatapath")

data_royalty    = DataRoyalty(REWARD, rdatapath)
data_produce    = DataProduce(REWARD, rdatapath)
data_drill_prim = DataDrillPrimitive(REWARD, rdatapath)
data_drill      = DataDrill(data_drill_prim, ddm)

evdims = size(EVobj(value_function(ddm)))
println_time_flush("EV dimension is $evdims, implying $(intstr(prod(evdims))) states")

# full dataset
dataset_full = DataSetofSets(data_drill, data_royalty, data_produce, CoefLinks(data_drill))

println_time_flush("Data created")

let thts = split_thetas(dataset_full, theta)
    println("theta_drill = $(round.(thts[1]; digits=3))")
    println("theta_roy   = $(round.(thts[2]; digits=3))")
    println("theta_pdxn  = $(round.(thts[3]; digits=3))")
end

# ------------------- simulations -----------------------

@show wp = statespace(ddm)
PP = PerpetualProblem(wp)

# ------------------- average costs -----------------------

cost_df = average_cost_df(dataset_full, theta)
# using Plots, StatsPlots
# @df cost_df plot(:date, -3.843351*[:cost1, :cost2], ylims=(0,Inf))

# ------------------- simulation info -----------------------
R = ShaleDrillingLikelihood.revenue(REWARD)
C = ShaleDrillingLikelihood.cost(REWARD)
E = ShaleDrillingLikelihood.extend(REWARD)
S = ShaleDrillingLikelihood.scrap(REWARD)

DrillRev(lrn, roy) = DrillingRevenue(constr(R), tech(R), tax(R), lrn, roy)
DrillRwrd(lrn, roy) = DrillReward(DrillRev(lrn,roy), C, E, S)

theta_notech = Theta_NoTech(dataset_full, theta, TECH_YEAR_ZERO)
TSTOP = time_idx(zchars(drill(dataset_full)), DateQuarter(DATE_STOP))

simlist = [
    (DrillRwrd( learn(R),      royalty(R) ),   wp, theta, ),           # "Baseline"
    (DrillRwrd( learn(R),      royalty(R) ),   wp, theta_notech, ),    # "Where to drill"
    (DrillRwrd( NoLearn(),     royalty(R) ),   wp, theta, ),           # "No update ($\\psi_{it} = \\psi_i^0$)"
    (DrillRwrd( PerfectInfo(), royalty(R) ),   wp, theta, ),           # "Perfect information ($\\psi_{it} = \\psi_i^1$)"
    (DrillRwrd( MaxLearning(), royalty(R) ),   wp, theta, ),           # "Uninformative signals ($\\rho = 0$)"
    (DrillRwrd( learn(R),      NoRoyalty()),   wp, theta, ),           # "No royalty"
    (DrillRwrd( learn(R),      royalty(R) ),   PP, theta, ),           # "No expiration"
    (DrillRwrd( learn(R),      NoRoyalty()),   PP, theta, ),           # "Ownership"
    (DrillRwrd(NoLearn(),      NoRoyalty()  ), PP, theta_notech, ),    # "Price only"
    (DrillRwrd(NoLearn(),      NoRoyalty()  ), PP, theta, ),           # "Price + tech (no distortions)"
    (DrillRwrd(NoLearn(),      WithRoyalty()), wp, theta_notech, ),    # "Price + leases"
    (DrillRwrd(Learn(),        NoRoyalty()  ), PP, theta_notech, ),    # "Price + learning"
    (DrillRwrd(NoLearn(),      WithRoyalty()), PP, theta, ),           # "How to drill"
]

simlist_meta = [(a,b,copy(theta_revenue(dataset_full, c))) for (a,b,c) in simlist]

simulations_meta_info = simulationPrimitives_information(simlist_meta)

# ------------------- simulations -----------------------

if DO_PAR
    pids = start_up_workers(ENV)
    @everywhere using ShaleDrillingLikelihood
    println_time_flush("Library loaded on workers")
end

println_time_flush("Starting simulations")
df_d, df_D = doSimulations(dataset_full, simlist, TSTOP, M)
println_time_flush("Simulations done")

# ------------------- save -----------------------

dpath_no_rdata = replace(DATAPATH, ".RData" => "")
filenm = "simulations-$(SLURM_JOBID)-$(dpath_no_rdata)-$(today()).RData"
filepath = joinpath(RFILEDIR, filenm)

println_time_flush("Saving simulations to $filepath")

THETA = theta
FILEPATH = filepath

@rput df_d df_D cost_df simulations_meta_info THETA FILEPATH TSTOP

R"""
library(data.table)
setDT(df_D)
setDT(df_d)
setDT(cost_df)
setDT(simulations_meta_info)

save(simulations_meta_info, df_d, df_D, cost_df, THETA, TSTOP, file=FILEPATH)
rm(  simulations_meta_info, df_d, df_D, cost_df, THETA, TSTOP,      FILEPATH)
gc()
"""

println_time_flush("Mischief managed. :)")


text:
using Documenter
using Bootstrap

makedocs(modules = [Bootstrap],
    sitename = "Bootstrap.jl",
    format = Documenter.HTML(disable_git = true,  # disable source and edit links to github
        canonical = "https://juliangehring.github.io/Bootstrap.jl/stable/"),
    doctest = true,
    checkdocs = :exports,
    linkcheck = true,
    pages = [
        "Home" => "index.md",
        "Library" => "library.md",
        "Changelog" => "NEWS.md",
        "License" => "LICENSE.md"
    ])

deploydocs(repo = "github.com/juliangehring/Bootstrap.jl.git")


text:
@userplot missing_plot

@recipe function f(h::missing_plot)
    if length(h.args) != 1 || !(typeof(h.args[1]) <: PopObj)
        error("Missingness Plot needs to be given a single PopObj")
    end
    x, y = missing(h.args[1])

    # set up the subplots
    legend := false
    framestyle := :zerolines
    xforeground_color_axis := :transparent
    yforeground_color_axis := :transparent
    xgrid := false
    ytickfontsize := 10

    group := x.population,
    color_palette := [HSL(i, .50, .50) for i in range(0, stop=300, length=unique(x.population) |> length)]
    hover := x.name .* ", " .* ["$i" for i in x.missing]
    title := "Number of missing loci per population"
    markerstrokealpha := 0.2
    markersize = 5.5

    # missing by individual
    @series begin
        seriestype := :dotplot
        group := x.population
        subplot := 1
        x.missing
    end

    # missing by locus
    @series begin
        seriestype := :histogram
        bins := length(unique(:missing))
        title := "Distribution of locus missingness"
        xguide := "# times missing"
        yguide := "# loci"
        linealpha := 0.1
        fill := "rgb(217, 217, 217)"
        xguidefontsize := 15
        xtickfontsize := 10
        ytickfontsize := 10
        yguidefontsize := 15
        subplot := 2
        y.missing
    end
end

colors = [HSL(i, .50, .50) for i in range(0, stop=300, length=unique(sh_miss[1].population) |> length)]

p1 = @df sh_miss[1] dotplot(
                :population,
                :missing,
                group = :population,
                color_palette = [HSL(i, .50, .50) for i in range(0, stop=300, length=unique(:population) |> length)],
                hover = :name .* ", " .* ["$i" for i in :missing],
                title = "Number of missing loci per population",
                xgrid = false,
                framestyle = :zerolines,
                legend = false,
                #xtickfontsize = 10,
                xforeground_color_axis = :transparent,
                ytickfontsize = 10,
                yforeground_color_axis = :transparent,
                markerstrokealpha = 0.2,
                markersize = 5.5,
)

p2 = @df sh_miss[2] histogram(
        :missing,
        bins = length(unique(:missing)),
        legend = false,
        title = "Distribution of locus missingness",
        xguide = "# times missing",
        yguide = "# loci",
        framestyle = :zerolines,
        linealpha = 0.1,
        fill = "rgb(217, 217, 217)",
        #hover = [count(i -> i ==j, :missing) for j in sort(unique(:missing))],
        xgrid = false,
        xforeground_color_axis = :transparent,
        xguidefontsize = 15,
        xtickfontsize = 10,
        yforeground_color_axis = :transparent,
        ytickfontsize = 10,
        yguidefontsize = 15,
)

plot(p1, p2, layout = (1,2))


text:
export Corona,
    compute_corona_covering_factor,
    compute_corona_seed_luminosity,
    compute_corona_dissipated_luminosity,
    compute_corona_photon_index,
    compute_corona_radius,
    compute_corona_luminosity,
    compute_reprocessed_flux,
    compute_corona_photon_flux,
    compute_corona_spectral_luminosity

using QuadGK, Roots

struct Corona{T<:Float64,U<:Bool}
    bh::BlackHole
    hard_xray_fraction::T
    electron_energy::T
    radius::T
    reprocessing::U
end
function Corona(bh::BlackHole; hard_xray_fraction, electron_energy, reprocessing)
    radius = compute_corona_radius(bh, hard_xray_fraction)
    return Corona(
        bh,
        Float64(hard_xray_fraction),
        Float64(electron_energy),
        Float64(radius),
        reprocessing,
    )
end
function Corona(bh::BlackHole, parameters::Parameters)
    return Corona(
        bh,
        hard_xray_fraction = parameters.hard_xray_fraction,
        electron_energy = parameters.corona_electron_energy,
        reprocessing = parameters.reprocessing,
    )
end

"""
    compute_covering_factor(corona)

Corona covering factor as seen from the disk at a radius r > corna radius.
"""
function compute_corona_covering_factor(corona, radius)
    if radius < corona.radius
        error("Radius can't be smaller than the corona.")
        return 0.0
    end
    theta0 = asin(corona.radius / radius)
    ret = theta0 - 0.5 * sin(2 * theta0)
    return ret
end

"""
    compute_corona_seed_luminosity(corona)

Seed photon luminosity intercepted from the warm region and the outer disk. Calculated assuming a truncated disk and spherical hot flow geometry.
"""
function compute_corona_seed_luminosity(corona)
    function f(radius)
        return radius *
               compute_disk_temperature(corona.bh, radius)^4 *
               compute_corona_covering_factor(corona, radius)
    end
    integ, _ = quadgk(f, corona.radius, gravity_radius(corona.bh), atol = 0, rtol = 1e-3)
    seed_lumin = 4 * corona.bh.Rg^2 * SIGMA_SB * integ
    return seed_lumin
end

"""
    compute_corona_dissipated_luminosity(corona)

Intrinsic luminosity from the Corona. This is assumed to be a constant fraction of the Eddington luminosity, regardless of actual accretion rate.
"""
function compute_corona_dissipated_luminosity(bh, hard_xray_fraction)
    return hard_xray_fraction * compute_eddington_luminosity(bh)
end
compute_corona_dissipated_luminosity(corona) =
    compute_corona_dissipated_luminosity(corona.bh, corona.hard_xray_fraction)

"""
    compute_corona_photon_index
Photon index (Gamma) for the corona SED. The functional form is assumed to be L_nu = k nu ^(-alpha) = k nu^( 1 - gamma ), where alpha = gamma - 1.
Computed using equation 14 of Beloborodov (1999).
"""
function compute_corona_photon_index(corona)
    gamma_cor =
        7 / 3 *
        (
            compute_corona_dissipated_luminosity(corona) /
            compute_corona_seed_luminosity(corona)
        )^(-0.1)
    return gamma_cor
end

"""
    compute_corona_radius
Finds the radius of the corona
"""
function compute_corona_radius(bh, hard_xray_fraction)
    function corona_radius_kernel(bh, radius)
        truncated_disk_lumin = compute_disk_luminosity(bh, bh.isco, radius)
        diff =
            truncated_disk_lumin -
            compute_corona_dissipated_luminosity(bh, hard_xray_fraction)
        return diff
    end
    r_cor = find_zero(
        r -> corona_radius_kernel(bh, r),
        (bh.isco, gravity_radius(bh)),
        Bisection(),
        atol = 0,
        rtol = 1e-3,
    )
    return r_cor
end
compute_corona_radius(corona) = corona.radius

"""
    compute_corona_luminosity(corona)
Total luminosity of the corona.
"""
function compute_corona_luminosity(corona)
    return compute_corona_seed_luminosity(corona) +
           compute_corona_dissipated_luminosity(corona)
end

"""
    compute_reprocessed_flux(corona, radius)
Computes the reprocessed flux from the corona at radius `radius`.
"""
function compute_reprocessed_flux(corona, radius; albedo = 0.3)
    Lhot = compute_corona_luminosity(corona)
    height = corona.radius * corona.bh.Rg
    radius = radius * corona.bh.Rg
    isco = corona.bh.isco * corona.bh.Rg
    Mdot = compute_mass_accretion_rate(corona.bh)
    aux1 = (3 * G * corona.bh.M * Mdot) / (8 * π * radius^3)
    aux2 = 2 * Lhot / (Mdot * C^2)
    aux3 = height / isco * (1 - albedo)
    aux4 = (1 + (height / radius)^2)^(-1.5)
    return aux1 * aux2 * aux3 * aux4
end

function compute_corona_spectral_luminosity(corona::Corona, warm, energy_range)
    gamma = compute_corona_photon_index(corona)
    kt_e = corona.electron_energy
    kt_c = compute_disk_temperature(corona.bh, corona.radius) * K_B
    kt_c_kev = kt_c * ERG_TO_KEV
    ywarm = compute_ywarm(warm)
    params = [gamma, kt_e, kt_c_kev * exp(ywarm), 0, 0]
    photons_per_bin = compute_compton_photons_per_bin(energy_range, params)
    mid_energies = (energy_range[1:end-1] + energy_range[2:end]) / 2
    lumin_per_bin = vcat([0], photons_per_bin[2:end] .* mid_energies)
    total_lumin = sum(lumin_per_bin)
    target_luminosity = compute_corona_luminosity(corona)
    ret = lumin_per_bin * (target_luminosity / total_lumin)
    ret = ret[2:end] ./ diff(energy_range)
    return vcat([0], ret)
end


text:
using RecipesBase

@userplot StatePlot

@recipe function f(p::StatePlot)

    evs = p.args[1]

    xguide --> "Remaining workload"
    yguide --> "Remaining deadline"
    markershape --> :circle

    w = [ev.currentWorkload for ev in evs];
    d = [ev.currentDeadline for ev in evs];
    u = [ev.currentPower>0 for ev in evs];

    @series begin

        seriestype := :scatter
        label := "Charging"
        seriescolor := :blue

        w[u.>0], d[u.>0]
    end

    @series begin

        seriestype := :scatter
        label := "Waiting"
        seriescolor := :red

        w[u.==0], d[u.==0]
    end

end

@userplot ReversedStatePlot

@recipe function f(p::ReversedStatePlot)

    evs = p.args[1]

    xguide --> "Attained workload"
    yguide --> "Sojourn time"
    markershape --> :circle

    w = [ev.requestedEnergy - ev.currentWorkload for ev in evs];
    d = [ev.departureTime - ev.arrivalTime - ev.currentDeadline for ev in evs];
    u = [ev.currentPower>0 for ev in evs];

    @series begin

        seriestype := :scatter
        label := "Charging"
        seriescolor --> :blue

        w[u.>0], d[u.>0]
    end

    @series begin

        seriestype := :scatter
        label := "Waiting"
        seriescolor --> :red

        w[u.==0], d[u.==0]
    end

end

@userplot struct ServicePlot{T<:Tuple{AbstractVector}}
    args::T
end

@recipe function f(p::ServicePlot{Tuple{Array{EVQueues.EVinstance,1}}})

    evs = p.args[1]

    xguide --> "Requested Energy"
    yguide --> "Attained Energy"
    markershape --> :circle

    S = [ev.requestedEnergy for ev in evs];
    Sa = [ev.requestedEnergy-ev.departureWorkload for ev in evs];

    @series begin

        seriestype := :scatter
        label --> "Charged EVs"

        S,Sa
    end

end

@userplot ServiceCDF

@recipe function f(p::ServiceCDF)

    evs = p.args[1]

    xguide --> "Attained service"
    yguide --> "CDF"
    legend --> :bottomright

    S = sort([ev.requestedEnergy for ev in evs])
    Sa = sort([ev.requestedEnergy-ev.departureWorkload for ev in evs])
    n=length(Sa)

    @series begin

        seriestype := :steppost
        label := "Attained service"
        Sa,(1:n)/n
    end

    @series begin

        seriestype := :steppost
        label := "Requested service"
        S,(1:n)/n
    end

end

@recipe function f(trace::EVQueues.TimeTrace)

    T=trace.T
    X=trace.X
    Y=trace.Y
    P=trace.P

    size --> (600,800)
    xguide --> "Time"
    legend := :none
    layout := (3,1)

    @series begin

        subplot := 1
        seriestype := :steppost
        title --> "#EV in charge"
        T,X
    end

    @series begin

        subplot := 2
        seriestype := :steppost
        title --> "#EV already charged"
        T,Y
    end

    @series begin

        subplot := 3
        seriestype := :steppost
        title --> "Charging Power"
        T,P
    end

end


text:
mutable struct ILU0Preconditioner{Tv, Ti} <: AbstractPreconditioner{Tv,Ti}
    A::ExtendableSparseMatrix{Tv,Ti}
    xdiag::Array{Tv,1}
    idiag::Array{Ti,1}
    phash::UInt64
    function ILU0Preconditioner{Tv,Ti}() where {Tv,Ti}
        p=new()
        p.phash=0
        p
    end
end

"""
```
ILU0Preconditioner(;valuetype=Float64,indextype=Int64)
ILU0Preconditioner(matrix)
```

ILU preconditioner with zero fill-in.
"""
ILU0Preconditioner(;valuetype::Type=Float64, indextype::Type=Int64)=ILU0Preconditioner{valuetype,indextype}()

function update!(precon::ILU0Preconditioner{Tv,Ti}) where {Tv,Ti}
    flush!(precon.A)
    cscmatrix=precon.A.cscmatrix
    colptr=cscmatrix.colptr
    rowval=cscmatrix.rowval
    nzval=cscmatrix.nzval
    n=cscmatrix.n

    if precon.phash==0
        n=size(precon.A,1)
        precon.xdiag=Array{Tv,1}(undef,n)
        precon.idiag=Array{Ti,1}(undef,n)
    end

    xdiag=precon.xdiag
    idiag=precon.idiag

    # Find main diagonal index and
    # copy main diagonal values
    if precon.phash != precon.A.phash
        @inbounds for j=1:n
            @inbounds for k=colptr[j]:colptr[j+1]-1
                i=rowval[k]
                if i==j
                    idiag[j]=k
                    break
                end
            end
        end
        precon.phash=precon.A.phash
    end

    @inbounds for j=1:n
        xdiag[j]=one(Tv)/nzval[idiag[j]]
        @inbounds for k=idiag[j]+1:colptr[j+1]-1
            i=rowval[k]
            for l=colptr[i]:colptr[i+1]-1
                if rowval[l]==j
                    xdiag[i]-=nzval[l]*xdiag[j]*nzval[k]
                    break
                end
            end
        end
    end
    precon
end

function  LinearAlgebra.ldiv!(u::AbstractArray{T,1}, precon::ILU0Preconditioner, v::AbstractArray{T,1}) where T
    cscmatrix=precon.A.cscmatrix
    colptr=cscmatrix.colptr
    rowval=cscmatrix.rowval
    n=cscmatrix.n
    nzval=cscmatrix.nzval
    xdiag=precon.xdiag
    idiag=precon.idiag

    @inbounds for j=1:n
        x=zero(T)
        @inbounds for k=colptr[j]:idiag[j]-1
            x+=nzval[k]*u[rowval[k]]
        end
        u[j]=xdiag[j]*(v[j]-x)
    end

    @inbounds for j=n:-1:1
        x=zero(T)
        @inbounds for k=idiag[j]+1:colptr[j+1]-1
            x+=u[rowval[k]]*nzval[k]
        end
        u[j]-=x*xdiag[j]
    end
end

function LinearAlgebra.ldiv!(precon::ILU0Preconditioner, v::AbstractArray{T,1} where T)
    ldiv!(v, precon, v)
end


text:
"""
struct MCMCResults
Struct for holding the MCMC chain (a vector of NamedTuples which represent the accepted parameter sets), and the MCMC tree statistics
"""
struct MCMCResults
    chain
    logL::Vector{Float64}
    treestatistics::Vector{DynamicHMC.TreeStatisticsNUTS}
end

"""
    mutable struct CoVArea

Type for containing the information required to calculate the log-likelihood contribution of one sub-area.
"""

Base.@kwdef mutable struct CoVAreaModel
    areaname::String #Name of area
    PCR_cases::Matrix{Int64} = zeros(Int64,365,2) #Daily PCR cases: dim 1: day, dim 2: num pos and total
    sero_cases::Array{Int64,3} = zeros(Int64,365,17,2) #Daily serology tests: dim 1: day, dim2 : age group, dim 3: num pos and neg
    average_sero_init::Vector{Float64} # seroprevalence to replace roond three seroprevalence used for initial starting conditions for counties without any sero data in round 3
    init_sero ::Vector{Float64} # initial (round 3) seroprevelance by age
    deaths::Matrix{Int64} = zeros(Int64,365,17) #Daily deaths
    pop_data::NamedMatrix{Float64} # Population by age and county
    N  #Area's population size
    α::Float64 = 1/3.1  # 1/mean latent period
    ϵ::Float64 = 0.5    # Relative infectiousness of asymptomatic cases
    αP::Float64 = 1/1.9  # rate of movement from Pre symptomatic state
    γA::Float64 = 1/2.4 #Recovery rate from Asymptomatic state
    γM::Float64 = 1/2.4 #Recovery rate from Mild disease

    relative_testing_rate::Vector{Float64} = relative_testing_rate
    PCR_array::Vector{Float64} = PCR_array #Relative Sensitivity of PCR test on days post-infection
    PCR_sensitivity::Float64 = 1. #Base sensitivity for RT-PCR test
    PCR_specificity::Float64 = 0.995 #Base specificity for RT_PCR test
    baseline_sero_array::Vector{Float64} = rel_sero_array_26days #Relative sensitivity of serology test on days post-infection
    sero_sensitivity::Float64 = 0.927 #Base sensitivity for serology test
    sero_specificity::Float64 = 0.992 #Base specificity for serology test
    M_BB::Float64 = 194.0 #Shrinkage factor for BetaBinomial distribution that incorporates the uncertainty in the sensitivity (fitted to serology group)
    log_likelihood::Function = θ -> 0.
    log_priors::Function = θ -> 0.

    M_county_ho::Matrix{Float64}     # "M_county_ho: Normalised age mixing matrix, home setting"
    M_county_other::Matrix{Float64}  # "M_county_other: Normalised age mixing matrix, social/other settings"
    M_county_school::Matrix{Float64} # "M_county_school: Normalised age mixing matrix, schools setting"
    M_county_work::Matrix{Float64}   # "M_county_work: Normalised age mixing matrix, work setting"

    rapid::Bool = false  # introduce vaccination in a rapid manner: finish target doses in 180 days
    scenario::Int64      # Scenario to run: Choices from 1-4
    time_to_full_efficacy::Float64  = 14. # Time to full vaccine vacc_efficacy
    hosp_rate_by_age:: Vector{Float64} = zeros(17) # Hospitalisation rate by age
    p_symp:: Vector{Float64}  # risk of symptoms by age
    p_sus:: Vector{Float64} # susceptibility by age
    p_sus_w:: Vector{Float64} # relative susceptibility following primary infection

    ω::Float64 = 1/180 # Waning of natural protection, rate per day

    # Include vaccination parameters in the vector of parameters
    # Sources: vacce effiecay against infections and severe VE_severe_disease_risk
        # Voysey M, Costa Clemens SA, Madhi SA, Weckx LY, Folegatti PM, et al. Single-dose administration
        #and the influence of the timing of the booster dose on immunogenicity and efficacy of
        #ChAdOx1 nCoV-19 (AZD1222) vaccine: a pooled analysis of four randomised trials. Lancet
        #(2021). doi:10.1016/S0140-6736(21)00432-3.

        # https://www.medrxiv.org/content/10.1101/2021.09.28.21264260v1

    VE_acquisition = [0,0.639,0.599] #Vaccine efficacy against acquisition by dose
    VE_infectiousness = [0,0.18,0.63] #Vaccine efficacy against transmission by dose
    VE_severe_disease_risk = [0,0.76,0.813] #Vaccine efficacy against disease by dose

    model_start_time # start in 1st Dec 2020
    model_end_time  # end of model run for fitting
    janstarttime  # time t start including observations in likelihood
    marchendpoint  # 10th March when roudn 3 serosurvey data ends
    alpha_variant_time  # earliest time of alpha variant introduction
    delta_variant_time  # earliest time of delta variant introduction

    startdate # date to start model run for projections/fitting
    enddate  # date to end model run for projections/fitting

    prob::ODEProblem

    MCMC_results::Union{MCMCResults,Nothing} = nothing #This field gets added after the MCMC run
end

"""
    CoVAreaModel(name::String,loglk_func,log_prior_func;
                        case_data,sero_data,death_data,pop_data::NamedArray)
Constructor for a `CoVAreaModel` struct using country data objects.
"""
function CoVAreaModel(name::String, log_likelihood, log_priors;
    linelist_data_with_pos_neg,
    serology_data,
    average_sero_init,
    deaths,
    pop_data::NamedMatrix,
    M_county_ho::Matrix{Float64},
    M_county_other::Matrix{Float64},
    M_county_school::Matrix{Float64},
    M_county_work::Matrix{Float64},
    rapid::Bool,
    scenario::Int64,
    startdate,
    enddate) # relative susceptibility by age following first episeode

    # Data input area
    #--------------------------------------------------------------------------------------------------------------------------
    time_to_full_efficacy = 14.0  # assumed 14 days post vaccination. This is the duration of stay in the partial vaccine protection states D1p and D2p

    #Set times, all times after model_start_time are in reference to the model_start_time
    model_start_time = (Date(2020, 12, 1) - Date(2020, 2, 20)).value # start in 1st Dec 2020
    model_end_time = (Date(2021, 9, 24) - Date(2020, 12, 1)).value  # end of model run for fitting
    janstarttime = (Date(2021, 1, 1) - Date(2020, 12, 1)).value # time t start including observations in likelihood
    marchendpoint = (Date(2021, 3, 10) - Date(2020, 12, 1)).value # 10th March when roudn 3 serosurvey data ends
    alpha_variant_time = (Date(2021, 2, 1) - Date(2020, 12, 1)).value  # earliest time of alpha variant introduction
    delta_variant_time = (Date(2021, 4, 1) - Date(2020, 12, 1)).value  # earliest time of delta variant introduction

    model_days = (Date(2022, 6, 30) - Date(2020, 12, 1)).value   # Spread vaccination till 2022 June
    pre_vacc_period = (Date(2021, 3, 6) - Date(2020, 12, 1)).value # vaccination starts on 6th March 2021

    # relative susceptibility by age
    prob_sus_w_17 = [0.636022, 0.37745, 0.424016, 0.456355, 0.633717, 0.725298, 0.739908, 0.714375, 0.7324, 0.790052, 0.917996, 1.02082, 1.06435, 1.31636, 1.39356, 1.31638, 1.531594]
    # probability of developing symptoms
    prop_symp_17 = [0.0638201, 0.0139376, 0.0195662, 0.0242427, 0.0631479, 0.093603, 0.0992078, 0.0895518, 0.0963006, 0.120111, 0.18606, 0.253572, 0.286403, 0.532203, 0.628416, 0.532232, 0.8299886]  # From Matt Keeling
    # relative susceptibility forllowinf first episode
    p_sus_w = ones(6) .* 0.16

    #reduce susceptibility and symptom risk vectors to six age groups
    p_sus = vcat(mean(prob_sus_w_17[1:4]), mean(prob_sus_w_17[5:10]), mean(prob_sus_w_17[11:12]), mean(prob_sus_w_17[13:14]), mean(prob_sus_w_17[15:16]), prob_sus_w_17[17])
    p_symp = vcat(mean(prop_symp_17[1:4]), mean(prop_symp_17[5:10]), mean(prop_symp_17[11:12]), mean(prop_symp_17[13:14]), mean(prop_symp_17[15:16]), prop_symp_17[17])

    # Age specific risk of severe disease
    #Estimate from CDC COVID-19 Response Team. Severe Outcomes Among Patients with Coronavirus Disease 2019 (COVID-19) - United States, February 12-March 16, 2020. MMWR Morb. Mortal. Wkly. Rep. 69, 343–346 (2020).
    #Model age groups: [0-19], [20-49], [50-59], [60-69], [70-79], 80+
    hosp_rate_CDC = [mean([1.6, 2.5]),#0-19 year olds
        mean([14.3, 20.8]),#20-44 yos
        mean([21.2, 28.3]),#45-54
        mean([20.5, 30.1]),#55-64
        mean([28.6, 43.5]),#65-74
        mean([30.5, 58.7]),#75-84
        mean([31.3, 70.3])] ./ 100 #85+

    # number of doses planned to cover about 30% of pop>18y
    if scenario == 0 || scenario == 1
        doses_Kenya = 0.0
    else
        doses_Kenya = (1.25 + 9.76 + 4.9) .* 1000000  #Number of doses planned for Kenya over the ynext one year
    end

    # END Data input area
    #--------------------------------------------------------------------------------------------------------------------------

    #Hsopitalisation rate by age. These rates dont actually affect the hospitalisations incidence reported from this work. We fit an IHR and ICU rates based on modelled incidence and observed hospital and ICU admissions
    H_rate_by_age = [hosp_rate_CDC[1], hosp_rate_CDC[2], hosp_rate_CDC[3], hosp_rate_CDC[4], hosp_rate_CDC[5], hosp_rate_CDC[6]]
    hosp_rate_by_age = H_rate_by_age .* 0.1 # Taking 10% of Verity rates. The IFR estimate from out model and death data was 10% of the expected from Verity data

    uprname = uppercase(name)
    PCR_cases = sum(linelist_data_with_pos_neg.cases[:, linelist_data_with_pos_neg.areas.==uprname, :, :], dims = [2, 3])[:, 1, 1, :]
    sero_cases = sum(serology_data.sero[:, serology_data.areas.==uprname, :, :], dims = [2])[:, 1, :, :]

    # reduce sero cases to 6 age groups
    p_vec = [[1, 2, 3, 4], [5, 6, 7, 8, 9, 10], [11, 12], [13, 14], [15, 16], [17]]
    sero_cases_collaped = zeros(Int64, size(sero_cases, 1), 6, 2)
    for a in 1:6
        sero_cases_collaped[:, a:a, :] .= sum(sero_cases[:, p_vec[a][1]:p_vec[a][end], :], dims = 2)
    end
    sero_cases = sero_cases_collaped

    # Put the initial serology as the default value --- this is based on Ngere et al and therefore more recent that rd 2 serology

    init_sero = average_sero_init

    # #Change from default if more than 10 tests
    # if sum(sero_cases[(model_start_time:(model_start_time+marchendpoint)), :, :]) == 0  # no round 3 data for some counties in the group of 36 counties without significant serodata
    #     init_sero = average_sero_init  # replace with the average
    # else
    #     u_sero = sero_cases[(model_start_time:(model_start_time+marchendpoint)), :, :]
    #     u_sero = sum(u_sero, dims = 1)[1, :, :]
    #     init_sero = u_sero[:, 1] ./ sum(u_sero, dims = 2)[:]
    # end

    # if isnan(init_sero[1])  # action if first value is isnan: replace with next succeding available value
    #     a = 1
    #     while isnan(init_sero[1])
    #         init_sero[1] = init_sero[a+1]
    #         a += 1
    #     end
    # end

    # for a in 2:6
    #     if isnan(init_sero[a])
    #         init_sero[a] = init_sero[a-1] # use the seroprevalence in the preceding age groups
    #     end
    # end

    #

    deaths = deaths.deaths[:, deaths.areas.==uprname, :][:, 1, :][:, :]

    #dates = [Date(2020,12,1) + Day(k) for k = 1:size(PCR_cases[(model_start_time:(model_start_time + model_end_time)),:],1)]
    N = pop_data[:, name]
    N = vcat(sum(N[1:4]), sum(N[5:10]), sum(N[11:12]), sum(N[13:14]), sum(N[15:16]), N[17]) # reduce to six age categories

    age_dist_vacc, doses_daily = KenyaCoVaccines.age_dist([name], doses_Kenya, pop_data, model_days, pre_vacc_period, rapid)
    vacc_rate_1, vacc_rate_2 = KenyaCoVaccines.vacc_rate_data(doses_daily, model_days, pre_vacc_period, age_dist_vacc, pop_data, name, scenario, rapid)

    # Reduce contact matrix from 17 to 6 age groups
    M_county_ho = KenyaCoVaccines.reduce_age_categories(M_county_ho, pop_data, name)
    M_county_other = KenyaCoVaccines.reduce_age_categories(M_county_other, pop_data, name)
    M_county_school = KenyaCoVaccines.reduce_age_categories(M_county_school, pop_data, name)
    M_county_work = KenyaCoVaccines.reduce_age_categories(M_county_work, pop_data, name)

    prob = build_prob_vaccine_modelling(N, vacc_rate_1, vacc_rate_2,
        M_county_ho, M_county_other, M_county_school, M_county_work;
        time_to_full_efficacy = time_to_full_efficacy,
        hosp_rate_by_age = hosp_rate_by_age,
        p_symp = p_symp,
        p_sus = p_sus,
        p_sus_w = p_sus_w,
        startdate = startdate,
        enddate = enddate)

    return KenyaCoVaccines.CoVAreaModel(areaname = name,
        log_likelihood = log_likelihood,
        log_priors = log_priors;
        PCR_cases = PCR_cases,
        sero_cases = sero_cases,
        init_sero = init_sero,
        average_sero_init = average_sero_init,
        deaths = deaths,
        startdate = startdate,
        enddate = enddate,
        model_start_time = model_start_time,
        model_end_time = model_end_time,
        janstarttime = janstarttime,
        alpha_variant_time = alpha_variant_time,
        delta_variant_time = delta_variant_time,
        marchendpoint = marchendpoint,
        prob = prob,
        pop_data = pop_data,
        N = N,
        M_county_ho = M_county_ho,
        M_county_other = M_county_other,
        M_county_school = M_county_school,
        M_county_work = M_county_work,
        rapid = rapid,
        scenario = scenario,
        time_to_full_efficacy = time_to_full_efficacy,
        hosp_rate_by_age = hosp_rate_by_age,
        p_symp = p_symp,
        p_sus = p_sus,
        p_sus_w = p_sus_w)

end

function (model::KenyaCoVaccines.CoVAreaModel)(θ)
    model.log_likelihood(θ,model,1/180) + model.log_priors(θ)
end

function change_prob!(model::CoVAreaModel, fully_vacc_rate::AbstractVector;
    startdate::Date,
    enddate::Date,
    maximum_uptake = 0.8)
    _prob = build_prob_vaccine_modelling_simplified(model.N, fully_vacc_rate, model.M_county_ho, model.M_county_other, model.M_county_school, model.M_county_work;
        maximum_uptake = maximum_uptake,
        startdate = startdate,
        enddate = enddate)
    model.prob = _prob
    return nothing
end

function change_prob_immune_escape!(model::CoVAreaModel, vaccine_prop;
    startdate::Date,
    enddate::Date,
    maximum_uptake = 0.8)
    _prob = build_prob_vaccine_modelling_vaccine_escape(model.N, vaccine_prop, model.M_county_ho, model.M_county_other, model.M_county_school, model.M_county_work;
        maximum_uptake = maximum_uptake,
        startdate = startdate,
        enddate = enddate)
    model.prob = _prob
    return nothing
end


text:
##################################################################################
# This file is part of StateSpaceEcon.jl
# BSD 3-Clause License
# Copyright (c) 2020, Bank of Canada
# All rights reserved.
##################################################################################

using LinearAlgebra

@testset "E1.simple" begin
    m = E1.model
    m.α = 0.5
    m.β = 1.0 - m.α
    for T = 6:10
        p = Plan(m, 2:T - 1)
        data = zerodata(m, p)
        data[1, :] = [1 0]   # initial condition
        data[end, :] = [5 0] # final condition
        sim00 = simulate(m, data, p)
        exp00 = hcat(1.0:(5.0 - 1.0) / (T - 1):5.0, zeros(T))
        # @info "T=$T" sim00 exp00 data
        @test sim00 ≈ exp00
        # solution piecewise linear y - one line before and another line after the shock
        data1 = copy(data)
        shk = .1
        y2val = ((T - 2) * (1 + 2 * shk) + 5.0) / (T - 1)
        data1[2U,:y_shk] = shk  # shock at time 2
        sim01 = simulate(m, data1, p)
        exp01 = vcat([1.0 y2val:(5.0 - y2val) / (T - 2):5.0...], [0 shk zeros(T - 2)...])'
        @test sim01 ≈ exp01
        # exogenous-endogenous swap
        # - replicate the solution above backing out the shock that produces it
        data2 = copy(data)
        data2[2U, :y] = sim01[2U, :y]
        p2 = deepcopy(p)
        exogenize!(p2, :y, 2)
        endogenize!(p2, :y_shk, 2)
        sim02 = simulate(m, data2, p2)
        exp02 = sim01
        @test sim02 ≈ exp02
    end
    # test_simulation(E1.model, "data/M1_TestMatrix.csv")
end

using DelimitedFiles

function test_simulation(m, path)
    nvars = ModelBaseEcon.nvariables(m)
    nshks = ModelBaseEcon.nshocks(m)
    #
    data_chk = readdlm(path, ',', Float64)
    IDX = size(data_chk, 1)
    plan = Plan(m, 1 + m.maxlag:IDX - m.maxlead)
    #
    p01 = deepcopy(plan)
    autoexogenize!(p01, m, m.maxlag + 1:IDX - m.maxlead)
    data01 = zeroarray(m, p01)
    data01[1:m.maxlag,:] = data_chk[1:m.maxlag,:]
    data01[end - m.maxlead + 1:end,:] = data_chk[end - m.maxlead + 1:end,:]
    data01[m.maxlag + 1:end - m.maxlead,1:nvars] = data_chk[m.maxlag + 1:end - m.maxlead,1:nvars]
    res01 = simulate(m, p01, data01)
    @test res01 ≈ data_chk
    p02 = deepcopy(plan)
    data02 = zeroarray(m, p02)
    data02[1:m.maxlag,:] = data_chk[1:m.maxlag,:]
    data02[end - m.maxlead + 1:end,:] = data_chk[end - m.maxlead + 1:end,:]
    data02[m.maxlag + 1:end - m.maxlead,nvars .+ (1:nshks)] = data_chk[m.maxlag + 1:end - m.maxlead,nvars .+ (1:nshks)]
    res02 = simulate(m, p02, data02)
    @test res02 ≈ data_chk
end

@testset "E1.sim" begin
    test_simulation(E1.model, "data/M1_TestData.csv")
end

@testset "E2.sim" begin
    test_simulation(E2.model, "data/M2_TestData.csv")
end

@testset "E3.sim" begin
    test_simulation(E3.model, "data/M3_TestData.csv")
end

@testset "E6.sim" begin
    test_simulation(E6.model, "data/M6_TestData.csv")
end

#############################################################
# linearization tests

@testset "linearize" begin
    m3 = deepcopy(E3.model)
    clear_sstate!(m3)
    @test_throws ModelBaseEcon.LinearizationError linearize!(m3)

    sssolve!(m3)
    @test issssolved(m3)

    test_simulation(m3, "data/M3_TestData.csv")

    m3nl = deepcopy(E3nl.model)
    m3nl.sstate.values .= m3.sstate.values
    m3nl.sstate.mask .= m3.sstate.mask
    @test issssolved(m3nl)

    # test_simulation(m3nl, "data/M3_TestMatrix.csv")

    with_linearized(m3nl) do foo
        test_simulation(foo, "data/M3_TestData.csv")
    end

    @test isa(m3nl.evaldata, ModelBaseEcon.ModelEvaluationData)

    linearize!(m3)
    @test isa(m3.evaldata, ModelBaseEcon.LinearizedModelEvaluationData)
    test_simulation(m3, "data/M3_TestData.csv")

    m7 = deepcopy(E7.model)
    if !issssolved(m7)
        empty!(m7.sstate.constraints)
        @steadystate m7 linv = lc - 7;
        @steadystate m7 lc = 14;
        clear_sstate!(m7)
        sssolve!(m7);
    end
    # the aux variables are present
    @test_throws ModelBaseEcon.LinearizationError linearize!(m7)

    clear_sstate!(E7A.model)
    sssolve!(E7A.model)
    # non-zeros linear growth in the steady state
    @test_throws ModelBaseEcon.LinearizationError linearize!(E7A.model)

end

#############################################################
# Tests of unanticipated shocks

@testset "E1.unant" begin
    m = E1.model
    m.α = m.β = 0.5
    p = Plan(m, 1:3);
    data = zeroarray(m, p);
    data[1,1] = 1.0;
    data[end,1] = 5.0;
    data[3,2] = 0.1;
    res_u = simulate(m, p, data; anticipate=false)
    true_u = Float64[1 0; 2 0; 47 / 15 0.1; 61 / 15 0; 5 0]
    @test res_u ≈ true_u atol = 1e-12
end

@testset "E2.unant" begin
    m = E2.model
    nvars = length(m.variables)
    nshks = length(m.shocks)
    var_inds = 1:nvars
    shk_inds = nvars .+ (1:nshks)
    # Steady state
    m.sstate.values .= 0;
    m.sstate.mask   .= true;
    let atrue = readdlm("./data/M2_Ant.csv", ',', Float64),
        utrue = readdlm("./data/M2_Unant.csv", ',', Float64)
        # Set simulation ranges and plan
        IDX = size(atrue, 1)
        init = 1:m.maxlag
        sim = 1 + m.maxlag:IDX - m.maxlead
        term = IDX .+ (1 - m.maxlead:0)
        p = Plan(m, sim)
        # Make up the exogenous data for the experiment
        let adata = zeroarray(m, p),
            udata = zeroarray(m, p)
            adata[sim,shk_inds] .= atrue[sim,shk_inds]
            udata[sim,shk_inds] .= utrue[sim,shk_inds]
            adata[term,:] .= atrue[term,:]
            udata[term,:] .= utrue[term,:]
            # Run the simulations and test
            ares = simulate(m, p, adata)
            @test ares ≈ atrue atol = m.options.tol
            ures = simulate(m, p, udata; anticipate=false)
            @test ures ≈ utrue atol = m.options.tol
        end
        let adata = zeroarray(m, p),
            udata = zeroarray(m, p)
            autoexogenize!(p, m, sim)
            adata[sim,var_inds] .= atrue[sim,var_inds]
            udata[sim,var_inds] .= utrue[sim,var_inds]
            adata[term,:] .= atrue[term,:]
            udata[term,:] .= utrue[term,:]
            # Run the simulations and test
            ares = simulate(m, p, adata)
            @test ares ≈ atrue atol = m.options.tol
            ures = simulate(m, p, udata; anticipate=false)
            @test ures ≈ utrue atol = m.options.tol
        end
    end
    for s in (1, 4)
        atrue = readdlm("./data/M2_t$(s)_Ant.csv", ',', Float64)
        utrue = readdlm("./data/M2_t$(s)_Unant.csv", ',', Float64)
        # Set simulation ranges and plan
        IDX = size(atrue, 1)
        init = 1:m.maxlag
        sim = 1 + m.maxlag:IDX - m.maxlead
        term = IDX .+ (1 - m.maxlead:0)
        p = Plan(m, sim)
        # Make up the exogenous data for the experiment
        adata = zeroarray(m, p);
        udata = zeroarray(m, p);
        ygap_ind = indexin([:ygap], m.variables)[1]
        adata[sim[1:s], ygap_ind] .= -0.10
        udata[sim[1:s], ygap_ind] .= -0.10
        adata[term,:] .= atrue[term,:]
        udata[term,:] .= utrue[term,:]
        exogenize!(p, :ygap, sim[1:s])
        endogenize!(p, :ygap_shk, sim[1:s])
        # Run the simulations and test
        ares = simulate(m, p, adata)
        @test ares ≈ atrue atol = m.options.tol
        ures = simulate(m, p, udata; anticipate=false)
        @test ures ≈ utrue atol = m.options.tol
    end
end

@testset "E3.unant" begin
    m = E3.model
    nvars = length(m.variables)
    nshks = length(m.shocks)
    var_inds = 1:nvars
    shk_inds = nvars .+ (1:nshks)
    # Steady state
    m.sstate.values .= 0;
    m.sstate.mask   .= true;
    let atrue = readdlm("./data/M3_Ant.csv", ',', Float64),
        utrue = readdlm("./data/M3_Unant.csv", ',', Float64)
        # Set simulation ranges and plan
        IDX = size(atrue, 1)
        init = 1:m.maxlag
        sim = 1 + m.maxlag:IDX - m.maxlead
        term = IDX .+ (1 - m.maxlead:0)
        p = Plan(m, sim)
        # Make up the exogenous data for the experiment
        let adata = zeroarray(m, p),
            udata = zeroarray(m, p)
            adata[sim,shk_inds] .= atrue[sim,shk_inds]
            udata[sim,shk_inds] .= utrue[sim,shk_inds]
            adata[term,:] .= atrue[term,:]
            udata[term,:] .= utrue[term,:]
            # Run the simulations and test
            ares = simulate(m, p, adata)
            @test ares ≈ atrue atol = m.options.tol
            ures = simulate(m, p, udata; anticipate=false)
            @test ures ≈ utrue atol = m.options.tol
        end
        let adata = zeroarray(m, p),
            udata = zeroarray(m, p)
            autoexogenize!(p, m, sim)
            adata[sim,var_inds] .= atrue[sim,var_inds]
            udata[sim,var_inds] .= utrue[sim,var_inds]
            adata[term,:] .= atrue[term,:]
            udata[term,:] .= utrue[term,:]
            # Run the simulations and test
            ares = simulate(m, p, adata)
            @test ares ≈ atrue atol = m.options.tol
            ures = simulate(m, p, udata; anticipate=false)
            @test ures ≈ utrue atol = m.options.tol
        end
        # tests with different final conditions and expectation_horizon values.
        for (fc, eh) = Iterators.product((fcgiven, fclevel, fcslope, fcnatural), ( nothing, 30, 50, 100))
            # Make up the exogenous data for the experiment
            p = Plan(m, sim)
            udata = zeroarray(m, p)
            udata[sim,shk_inds] .= utrue[sim,shk_inds]
            udata[term,:] .= utrue[term,:]
            # Run the simulations and test
            ures = simulate(m, p, udata; anticipate=false, expectation_horizon=eh, fctype=fc)
            @test ures ≈ utrue atol = m.options.tol  rtol = maximum(abs, utrue[term,:])
            # Repeat with autoexogenize
            p = Plan(m, sim)
            udata = zeroarray(m, p)
            autoexogenize!(p, m, sim)
            udata[sim,var_inds] .= utrue[sim,var_inds]
            udata[term,:] .= utrue[term,:]
            # Run the simulations and test
            ures = simulate(m, p, udata; anticipate=false, expectation_horizon=eh, fctype=fc)
            @test ures ≈ utrue atol = m.options.tol rtol = maximum(abs, utrue[term,:])
        end
    end
    for s in (1, 4)
        atrue = readdlm("./data/M3_t$(s)_Ant.csv", ',', Float64)
        utrue = readdlm("./data/M3_t$(s)_Unant.csv", ',', Float64)
        # Set simulation ranges and plan
        IDX = size(atrue, 1)
        init = 1:m.maxlag
        sim = 1 + m.maxlag:IDX - m.maxlead
        term = IDX .+ (1 - m.maxlead:0)
        p = Plan(m, sim)
        # Make up the exogenous data for the experiment
        adata = zeroarray(m, p);
        udata = zeroarray(m, p);
        ygap_ind = indexin([:ygap], m.variables)[1]
        adata[sim[1:s], ygap_ind] .= -0.10
        udata[sim[1:s], ygap_ind] .= -0.10
        adata[term,:] .= atrue[term,:]
        udata[term,:] .= utrue[term,:]
        exogenize!(p, :ygap, sim[1:s])
        endogenize!(p, :ygap_shk, sim[1:s])
        # Run the simulations and test
        ares = simulate(m, p, adata)
        @test ares ≈ atrue atol = m.options.tol
        ures = simulate(m, p, udata; anticipate=false)
        @test ures ≈ utrue atol = m.options.tol
    end
end

@testset "FCTypes" begin
    let m = Model()
        @variables m x
        @equations m begin x[t] = x[t + 3] end
        @initialize m
        p = Plan(m, 1:1)
        ed = zerodata(m, p)
        @test_throws ArgumentError simulate(m, ed, p, fctype=fcnatural)
    end
    let m = E1.model
        p = Plan(m, 3:17)
        ed = zerodata(m, p)
        ed .= rand(Float64, size(ed))
        @test_logs (:error, r"The system is underdetermined.*") (@test_throws SingularException simulate(m, ed, p, fctype=fcnatural))
        sd = StateSpaceEcon.StackedTimeSolver.StackedTimeSolverData(m, p, fcgiven)
        @test_throws ErrorException StateSpaceEcon.StackedTimeSolver.update_plan!(sd, m, Plan(m, 3:8))
        x = rand(Float64, size(ed))
        R1, _ = StateSpaceEcon.StackedTimeSolver.global_RJ(x, ed, sd)
        R2 = StateSpaceEcon.StackedTimeSolver.global_R!(similar(R1), x, ed, sd)
        @test R1 ≈ R2
    end
    let m = E3.model
        p = Plan(m, 3:177)
        ed = zerodata(m, p)
        ed[3, 4:6] .= 0.2 * rand(3)
        ed[3:end, 1:3] .= rand(178, 3)
        empty!(m.sstate.constraints)
        clear_sstate!(m)
        @test_throws ArgumentError simulate(m, ed, p, fctype=fclevel)
        @test_throws ArgumentError simulate(m, ed, p, fctype=fcslope)
        vec1 = sssolve!(m)
        initial_sstate!(m, 0.4 * rand(Float64, size(m.sstate.values)))
        vec2 = sssolve!(m)
        @test vec1 ≈ vec2
        s1 = simulate(m, ed, p, fctype=fclevel)
        s2 = simulate(m, ed, p, fctype=fcslope)
        s3 = simulate(m, ed, p, fctype=fcnatural)
        @test maximum(abs, s1 - s2) < 1e-8
        @test maximum(abs, s1 - s3) < 1e-8
        @test maximum(abs, s3 - s2) < 1e-8

        sd = StateSpaceEcon.StackedTimeSolver.StackedTimeSolverData(m, p, fcgiven)
        x = rand(Float64, size(ed))
        R1, _ = StateSpaceEcon.StackedTimeSolver.global_RJ(x, ed, sd)
        R2 = StateSpaceEcon.StackedTimeSolver.global_R!(similar(R1), x, ed, sd)
        @test R1 ≈ R2
    end
    let m = E6.model
        empty!(m.sstate.constraints)
        @steadystate m lp  = 1.0
        @steadystate m lp + lyn = 1.5
        @steadystate m lp + ly = 1.2
        clear_sstate!(m)
        sssolve!(m)
        p = Plan(m, 3:177)
        ed = zerodata(m, p)
        ed.dlp_shk[3] = rand()
        ed.dly_shk[3] = rand()
        s2 = simulate(m, ed, p, fctype=fcslope)
        s3 = simulate(m, ed, p, fctype=fcnatural)
        @test maximum(abs, s3 - s2) < 1e-12

        sd = StateSpaceEcon.StackedTimeSolver.StackedTimeSolverData(m, p, fcgiven)
        x = rand(Float64, size(ed))
        R1, _ = StateSpaceEcon.StackedTimeSolver.global_RJ(x, ed, sd)
        R2 = StateSpaceEcon.StackedTimeSolver.global_R!(similar(R1), x, ed, sd)
    end
end


text:
using Solvent, Test, LinearAlgebra, Random, CUDA, KernelAbstractions

@kernel function laplace_kernel!(Ax, x, param)
    i = @index(Global, Linear)
    @inbounds begin
        Ax[i] =  x[i]
        if i > 1
            Ax[i] += param*x[i-1]
        end
        if i < length(x)
            Ax[i] += param * x[i+1]
        end
    end
end

"""
    laplace!(AX, X, param=-0.01)

Applies the 1d Laplace operator `A*x`  where
```
A[i,i] = 1
A[i,i-1] = A[i,i+1] = param
```
"""
function laplace!(AX, X, param=eltype(X)(-0.01))
    event = laplace_kernel!(Solvent.array_device(X), 256)(
        AX, X, eltype(X)(param); ndrange=length(X))
    wait(event)
end

for T in (Float32, Float64)
    Random.seed!(44)

    n = 200
    M = 20
    K = 10
    tol = sqrt(eps(T))

    x_ref = CuArray(randn(T,n))
    x = similar(x_ref)
    b = similar(x_ref)
    laplace!(b,x_ref)

    linearsolver = LinearSolver(
        laplace!,
        GeneralizedMinimalResidualMethod(M = M, K = K),
        x;
        rtol = tol,
        atol = tol,
    )

    linearsolve!(linearsolver, x, b)
    @test x ≈ x_ref
end

max_stars_repo_name: CliMA/SystemSolvers

max_issues_repo_name: CliMA/SystemSolvers


text:
function Base.rand(rng::Random.AbstractRNG, sampler::Random.SamplerTrivial{<:Consistent})
    prediction = rand(rng, sampler[].predictions)
    target = sample_target(rng, prediction)
    return (prediction, target)
end

# in-place sampling
function Random.rand!(
    rng::Random.AbstractRNG,
    out::AbstractArray{<:Tuple{<:Any,<:Any}},
    sampler::Random.SamplerTrivial{<:Consistent},
)
    s = Random.Sampler(rng, sampler[].predictions)
    for i in eachindex(out)
        prediction = rand(rng, s)
        target = sample_target(rng, prediction)
        @inbounds out[i] = (prediction, target)
    end
    return out
end

# sample targets
sample_target(rng::Random.AbstractRNG, sampler) = rand(rng, sampler)
function sample_target(rng::Random.AbstractRNG, prediction::Real)
    return rand(rng, typeof(float(prediction))) < prediction
end
function sample_target(rng::Random.AbstractRNG, prediction::AbstractVector{<:Real})
    return StatsBase.sample(rng, StatsBase.Weights(prediction, one(eltype(prediction))))
end