text:
function ChronologicalComparator()
    return ChronologicalComparator(())
end

function compare(obj::ChronologicalComparator, arg0::TimeStamped, arg1::TimeStamped)
    return jcall(obj, "compare", jint, (TimeStamped, TimeStamped), arg0, arg1)
end

function reversed(obj::Comparator)
    return jcall(obj, "reversed", Comparator, ())
end

function then_comparing(obj::Comparator, arg0::Comparator)
    return jcall(obj, "thenComparing", Comparator, (Comparator,), arg0)
end

function then_comparing(obj::Comparator, arg0::Function)
    return jcall(obj, "thenComparing", Comparator, (Function,), arg0)
end

function then_comparing(obj::Comparator, arg0::Function, arg1::Comparator)
    return jcall(obj, "thenComparing", Comparator, (Function, Comparator), arg0, arg1)
end

function then_comparing_double(obj::Comparator, arg0::ToDoubleFunction)
    return jcall(obj, "thenComparingDouble", Comparator, (ToDoubleFunction,), arg0)
end

function then_comparing_int(obj::Comparator, arg0::ToIntFunction)
    return jcall(obj, "thenComparingInt", Comparator, (ToIntFunction,), arg0)
end

function then_comparing_long(obj::Comparator, arg0::ToLongFunction)
    return jcall(obj, "thenComparingLong", Comparator, (ToLongFunction,), arg0)
end

ext: jl


text:
@testset "fit moments" begin
    D = Exponential(2.1)
    M = Moments(mean(D))
    Dfit = fit(Exponential, M)
    @test D ≈ Dfit
    # handle not given any moment
    @test_throws Exception fit(Exponential, Moments())
end;

@testset "fit to quantilepoint and mode - ignore mode" begin
    d = Exponential(0.8)
    qp = @qp(quantile(d,0.95),0.95)
    dfit = fit_mode_quantile(Exponential, mode(d), qp)
    @test dfit ≈ d
    dfit = fit(Exponential, mode(d), qp, Val(:mode))
    @test dfit ≈ d
    # with lower quantile
    qp = @qp(quantile(d,0.025),0.025)
    dfit = fit_mode_quantile(Exponential, mode(d), qp)
    @test dfit ≈ d
end;

@testset "fit two quantiles same" begin
    qpl = @qp_m(3)
    d = fit(Exponential, qpl, qpl);
    @test quantile.(d, [qpl.p]) ≈ [qpl.q]
end;

@testset "fit two quantiles" begin
    qpl = @qp_m(3)
    qpu = @qp_u(5)
    d = fit(Exponential, qpl, qpu);
    #plot(d);
    #plot!(fit_mode_quantile(Exponential, NaN, qpl))
    #plot!(fit_mode_quantile(Exponential, NaN, qpu))
    #vline!([3,5])
    @test quantile(d, qpl.p) < qpl.q
    @test quantile(d, qpu.p) > qpu.q
end;


text:
module Aerodynamics

# using LinearAlgebra
using StaticArrays, ComponentArrays
using UnPack

using Flight.Modeling
import Flight.Modeling: init_x0, init_y0, init_u0, init_d0, f_cont!, f_disc!

using Flight.Plotting
import Flight.Plotting: plots

using Flight.Attitude
using Flight.Kinematics
using Flight.Dynamics
using Flight.Airdata

using Flight.Components
import Flight.Components: WrenchTrait, AngularMomentumTrait, get_wr_b

export SimpleDrag, TestAerodynamics

abstract type AbstractAerodynamics <: SystemDescriptor end

WrenchTrait(::System{<:AbstractAerodynamics}) = HasWrench()
AngularMomentumTrait(::System{<:AbstractAerodynamics}) = HasNoAngularMomentum()

#################### SimpleDrag ######################

Base.@kwdef struct SimpleDrag <: AbstractAerodynamics
    c_d::Float64 = 0.5
    A::Float64 = 1.0
end

Base.@kwdef struct SimpleDragY
    D::Float64 = 0.0
    wr_b::Wrench = Wrench()
end

init_y0(::SimpleDrag) = SimpleDragY()

#in case of type instability use this one instead (or maybe a function barrier)
# function f_cont!(sys::System{SimpleDrag}, air::AirData, ::KinData, ::Any, ::Any)
function f_cont!(sys::System{SimpleDrag}, air::AirData, args...)

    @unpack α_b, β_b = air
    q_bw = get_wind_axes(α_b, β_b)

    D = 0.5 * air.q * sys.params.c_d * sys.params.A
    F_Ob_w = SVector{3,Float64}(-D, 0.0, 0.0)
    F_Ob_b = q_bw(F_Ob_w)

    #now the wrench is trivially computed
    wr_b = Wrench(F = F_Ob_b)

    sys.y = SimpleDragY(D, wr_b)

end

f_disc!(::System{SimpleDrag}) = false

Components.get_wr_b(sys::System{SimpleDrag}) = sys.y.wr_b

function plots(t, data::AbstractVector{SimpleDragY}; mode, save_path, kwargs...)

    @unpack D, wr_b = StructArray(data)

    pd = Dict{String, Plots.Plot}()

    pd["01_D"] = thplot(t, D;
        plot_title = "Aerodynamic Drag",
        label = "",
        ylabel = L"$D (N)$",
        kwargs...)

    pd["02_wr_Ob_b"] = thplot(t, wr_b;
        plot_title = "Aerodynamic Wrench [Airframe]",
        wr_source = "aero", wr_frame = "b",
        kwargs...)

    save_plots(pd; save_path)

end

#################### TestAerodynamics ################

#this subtype showcases the solutions to two potential scenarios:
#1) what to do if the aerodynamic data is expressed in a different frame (a)
#   than that chosen for the airframe (b), and we actually care about the
#   velocity lever arm between them

# answer: having v_wOb_b in AirData, any SystemDescriptor to which AirData is
# passed is free to transform v_wOb_b into its own frame. this may be
# particularly useful for an Aerodynamics component. if its frame is f(Oa, εa):

# v_wOb_b = v_eOb_b - v_ew_b
# v_eOa_b = v_eOb_b + ω_eb_b × r_ObOa_b
# v_wOa_b = v_eOa_b - v_ew_b = v_eOb_b + ω_eb_b × r_ObOa_b - v_ew_b
# v_wOa_b = v_wOb_b + ω_eb_b × r_ObOa_b
# v_wOa_a = q_ba'(v_wOa_b)

#generally, v_wOa ≈ v_wOb, so the whole conversion won't be necessary; at most,
#we will need to reproject v_wOb_b into v_wOb_a to comply with the axes used by
#the aerodynamics database, and from v_wOb_a we can compute airflow angles for
#those axes

#2) what to do if there is dependence of α_dot and β_dot in the aerodynamic
#   coefficients

# answer: we make the AerodynamicsSystem stateful, with filtered versions of α
# and β as states.

# we implement a first order filter as:
# ẋ.α = 1/τ * (α_in - x.α)
# ẋ.β = 1/τ * (β_in - x.β)

# now we have filtered versions of α_dot and β_dot available to use in the
# aerodynamics data.

# note that having strictly α_dot and β_dot as inputs (not their filtered
# counterparts) would make the complete world model implicit: α_dot and β_dot
# depend on v̇_eOb_b and v̇_ew_n, which in turn depend on aerodynamic forces and
# moments, which in turn depend on α_dot and β_dot. not only that, but this
# formulation would mean that any discontinuous changes in wind velocity would
# cause these derivatives to go to infinity. so... having filtered versions of
# these variables is the way to go.

#the remaining question is: what do we input as airflow angles to the
#aerodynamic database? the plain α_c, β_c inputs or the filtered states α, β?
#the first option seems better, because they are direct feedthrough, so unlike
#the filtered states, they do not introduce delays. all aerodynamic
#contributions not involving α_dot and β_dot are then delay-free. however, we
#should assume that in any case τ will have been chosen so that the bandwidth of
#the airflow angle filter extends well beyond that of aircraft dynamics

Base.@kwdef struct TestAerodynamics <: AbstractAerodynamics
    frame::FrameTransform = FrameTransform()
    τ::Float64 = 0.1 #time constant for first-order airflow angle filters
end

Base.@kwdef struct TestAerodynamicsY
    α_dot::Float64 = 0
    α_in::Float64 = 0
    α::Float64 = 0
    β_dot::Float64 = 0
    β_in::Float64 = 0
    β::Float64 = 0
    wr_c::Wrench = Wrench()
    wr_b::Wrench = Wrench()
end

init_x0(::TestAerodynamics) = ComponentVector(α = 0.0, β = 0.0)
init_y0(::TestAerodynamics) = TestAerodynamicsY()

function f_cont!(sys::System{TestAerodynamics}, air::AirData, kin::KinData, ::Any, ::Any)

    @unpack ẋ, x, y, params = sys
    @unpack α, β = x

    t_bc = params.frame

    # v_wOc_b = v_wOb_b #simply use the airframe origin velocity
    v_wOc_b = air.v_wOb_b + kin.ω_eb_b × t_bc.r
    v_wOc_c = t_bc.q'(v_wOc_b)
    α_in, β_in = get_airflow_angles(v_wOc_c) #airflow angles in component frame

    α_dot = 1/τ * (α_in - x.α)
    β_dot = 1/τ * (β_in - x.β)

    #here we would actually compute aerodynamic forces and moments using either
    #α_c, β_c or α_c_in and β_c_in
    wr_c = Wrench()
    wr_b = t_bc(wr_c)

    ẋ.α = α_dot
    ẋ.β = β_dot

    sys.y = TestAerodynamicsY(α_dot, α_in, α, β_dot, β_in, β, wr_c, wr_b)

end

end #module


text:
include("base.jl")
include("optimizers.jl")
include("tiling.jl")


text:
module CoreTests

using Test
using UnstructuredGrids
using UnstructuredGrids.Core
using UnstructuredGrids.Kernels
using UnstructuredGrids.Factories
using UnstructuredGrids.VTK

c = Connections([[1,2,6,3,],[1,4,4],[1]])
s = """
    1 -> [1, 2, 6, 3]
    2 -> [1, 4, 4]
    3 -> [1]
    """
@test s == string(c)

grid = UGrid(domain=(0,1,-1,0),partition=(2,2))

cell_to_ctype = celltypes(grid)

ctype_to_refcell = refcells(grid)

cell_to_vertices = connections(grid)

@test ndims(grid) == 2

c = cell_to_vertices
l = [1, 2, 4, 5, 2, 3, 5, 6, 4, 5, 7, 8, 5, 6, 8, 9]
p = [1, 5, 9, 13, 17]
@test list(c) == l
@test ptrs(c) == p

vertex_to_cells = generate_dual_connections(cell_to_vertices)

c = vertex_to_cells
l = [1, 1, 2, 2, 1, 3, 1, 2, 3, 4, 2, 4, 3, 3, 4, 4]
p = [1, 2, 4, 5, 7, 11, 13, 14, 16, 17]
@test list(c) == l
@test ptrs(c) == p

cell_to_faces = generate_cell_to_faces(
  cell_to_vertices,vertex_to_cells,cell_to_ctype,ctype_to_refcell,1)

c = cell_to_faces
l = [1, 2, 3, 4, 5, 6, 4, 7, 2, 8, 9, 10, 6, 11, 10, 12]
p = [1, 5, 9, 13, 17]
@test list(c) == l
@test ptrs(c) == p

fgrid = UGrid(grid,1)

@test ndims(fgrid) == 1

face_to_vertices = connections(fgrid)

c = face_to_vertices
l = [1, 2, 4, 5, 1, 4, 2, 5, 2, 3, 5, 6, 3, 6, 7, 8, 4, 7, 5, 8, 8, 9, 6, 9]
p = [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25]
@test list(c) == l
@test ptrs(c) == p

fgrid = UGrid(grid,1,vertex_to_cells)

face_to_vertices = connections(fgrid)

c = face_to_vertices
@test list(c) == l
@test ptrs(c) == p

fgrid = UGrid(grid,1,vertex_to_cells,cell_to_faces)

face_to_vertices = connections(fgrid)

c = face_to_vertices
@test list(c) == l
@test ptrs(c) == p

vertex_to_faces = generate_dual_connections(face_to_vertices)

cell_to_faces = find_cell_to_faces(
  cell_to_vertices, vertex_to_faces, cell_to_ctype, ctype_to_refcell, 1)

c = cell_to_faces
l = [1, 2, 3, 4, 5, 6, 4, 7, 2, 8, 9, 10, 6, 11, 10, 12]
p = [1, 5, 9, 13, 17]
@test list(c) == l
@test ptrs(c) == p

cell_to_faces = find_cell_to_faces(grid,fgrid)
@test c.list == cell_to_faces.list
@test c.ptrs == cell_to_faces.ptrs

face_to_cells = generate_dual_connections(cell_to_faces)

_face_to_isboundary = generate_facet_to_isboundary(face_to_cells)

face_to_isboundary = Bool[
  true, false, true, false, true, false, true, true, true, false, true, true]

@test face_to_isboundary == _face_to_isboundary

_vertex_to_isboundary = generate_face_to_isboundary(
  face_to_isboundary, vertex_to_faces)

vertex_to_isboundary = Bool[
  true, true, true, true, false, true, true, true, true]

@test vertex_to_isboundary == _vertex_to_isboundary

primal, dual = generate_grid_graph(grid)

@test primal[1+1] == cell_to_faces
@test primal[0+1] == cell_to_vertices

@test dual[1+1] == face_to_cells
@test dual[0+1] == vertex_to_cells

fullgraph = generate_full_grid_graph(grid)

@test fullgraph[2+1,1+1] == cell_to_faces
@test fullgraph[1+1,2+1] == face_to_cells

@test fullgraph[1+1,0+1] == face_to_vertices
@test fullgraph[0+1,1+1] == vertex_to_faces

@test fullgraph[2+1,0+1] == cell_to_vertices
@test fullgraph[0+1,2+1] == vertex_to_cells

@test isa(fullgraph[0+1,0+1],Connections)
@test isa(fullgraph[1+1,1+1],Connections)
@test isa(fullgraph[2+1,2+1],Connections)

@test length(fullgraph[0+1,0+1].ptrs) == length(vertex_to_cells.ptrs)
@test length(fullgraph[1+1,1+1].ptrs) == length(face_to_vertices.ptrs)
@test length(fullgraph[2+1,2+1].ptrs) == length(cell_to_vertices.ptrs)

grid = UGrid(domain=(0,1,-1,0,0,1),partition=(2,2,2))

fullgraph = generate_full_grid_graph(grid)

end # module CoreTests


text:
"Simulator that returns end state of simulation."
@kwdef struct EndStateSimulator <: Simulator
    max_steps::Union{Int,Nothing} = nothing
end

function simulate(sim::EndStateSimulator, domain::Domain, state::State, actions)
    state = copy(state)
    for (t, act) in enumerate(actions)
        state = transition!(domain, state, act)
        sim.max_steps !== nothing && t >= sim.max_steps && break
    end
    return state
end


text:
function erf(::Type{Erf}, arg0::jdouble)
    return jcall(Erf, "erf", jdouble, (jdouble,), arg0)
end

function erf(::Type{Erf}, arg0::jdouble, arg1::jdouble)
    return jcall(Erf, "erf", jdouble, (jdouble, jdouble), arg0, arg1)
end

function erf_inv(::Type{Erf}, arg0::jdouble)
    return jcall(Erf, "erfInv", jdouble, (jdouble,), arg0)
end

function erfc(::Type{Erf}, arg0::jdouble)
    return jcall(Erf, "erfc", jdouble, (jdouble,), arg0)
end

function erfc_inv(::Type{Erf}, arg0::jdouble)
    return jcall(Erf, "erfcInv", jdouble, (jdouble,), arg0)
end


text:
module SpiralExample

using Flux
using DiffEqFlux
using OrdinaryDiffEq
using Random
using Distributions

const noise_std = 0.1 # standard deviation of random noise added to samples
const noise_logvar = 2*log(noise_std)

# The elements of `SpiralSample`s are vectors with two entries for the two dimensions.
const SpiralSample = Vector{Vector{Float32}}

logp_x_z(x, z) = sum(log_normal_pdf(x, z, noise_logvar))

function spiral_clockwise(; start, stop, ntotal, a, b)
    zs = stop .+ 1. .- range(start, length = ntotal, stop=stop)
    rs = a .+ (b .* zs)
    xs = rs .* cos.(zs) .- 5.
    ys = rs .* sin.(zs)
    hcat(xs, ys)
end

function spiral_counterclockwise(; start, stop, ntotal, a, b)
    zs = range(start, length = ntotal, stop = stop)
    rw = a .+ (b .* 20 ./ zs)
    xs = rw .* cos.(zs) .+ 5.
    ys = rw .* sin.(zs)
    hcat(xs, ys)
end

function spiral_samples(;
        nspiral::Int = 1000, # no of spirals generated
        nsample::Int = 100, # no of trajectory points drawn from each spiral
        start::Float64 = 0.0, # spiral starting phi value
        stop::Float64 = 6*pi, # spiral ending phi value
        a::Float64 = 0.0, # parameters defining shape of Archimedean spiral
        b::Float64 = 1.0)
    # Parametric formula for 2d spiral is `r = a + b * phi`.

    # returns named tuple:
    # - orig_trajs: vector of `SpiralSample`s over a full spiral
    # - samp_trajs: noisy observations, vector of `SpiralSample`s
    # - orig_ts: vector of length `ntotal` containing the timepoints for `orig_trajs`
    # - orig_ts: vector of length `nsample` containing the timepoints for `samp_trajs`

    ntotal = 5*nsample # no of trajectory points for each spiral sample

    orig_traj_cc = spiral_counterclockwise(start = start, stop = stop, ntotal = ntotal, a = a, b = b)
    orig_traj_cw = spiral_clockwise(start = start, stop = stop, ntotal = ntotal, a = a, b = b)

    # sample starting timestamps
    orig_ts = range(start, length = ntotal, stop = stop)
    samp_ts = orig_ts[1:nsample]

    orig_trajs = []
    samp_trajs = []

    for _ in 1:nspiral
        # don't sample t0 very near the start or the end
        pvec = [1. ./ (ntotal.- 4. .* nsample) for _ in 1:(ntotal .- (4 .* nsample))]
        t0_idx = rand(Multinomial(1, pvec))
        t0_idx = argmax(t0_idx) + nsample

        cc = rand() > .5 # uniformly select rotation
        orig_traj = cc ? orig_traj_cc : orig_traj_cw
        push!(orig_trajs, Float32.(orig_traj))

        samp_traj = deepcopy(orig_traj[t0_idx:2:(t0_idx + 2*(nsample-1)),:])
        samp_traj += randn(size(samp_traj)) .* noise_std
        push!(samp_trajs, Float32.(samp_traj))
    end

    datatrafo(x) = [[x[i][j,:] for j in 1:size(x[1],1)] for i in 1:length(x)]

    return (orig_trajs = datatrafo(orig_trajs),
            samp_trajs = datatrafo(samp_trajs),
            orig_ts = orig_ts, samp_ts = samp_ts)
end

struct LatentTimeSeriesVAE
    rnn
    latentODEfunc
    latentODEparams
    decoder
end

function LatentTimeSeriesVAE(; latent_dim, obs_dim, rnn_nhidden, f_nhidden, dec_nhidden)
    rnn = Chain(RNN(obs_dim, rnn_nhidden), Dense(rnn_nhidden, latent_dim*2))

    latentODEfunc = Chain(Dense(latent_dim, f_nhidden, Flux.elu),
                          Dense(f_nhidden, f_nhidden, Flux.elu),
                          Dense(f_nhidden, latent_dim))
    latentODEparams, re = Flux.destructure(latentODEfunc)

    decoder = Chain(Dense(latent_dim, dec_nhidden, Flux.relu),
                    Dense(dec_nhidden, obs_dim))
    LatentTimeSeriesVAE(rnn, re, latentODEparams, decoder)
end

latentz0(μ, logσ) = μ .+ exp.(logσ) .* randn(Float32)

function n_ode(model, z0, t)
    tspan = (t[1], t[end])
    latent_dynamics(u, latentODEparams, t) = model.latentODEfunc(latentODEparams)(u)
    latentODEprob = ODEProblem(latent_dynamics, z0, tspan)
    Array(solve(latentODEprob, Tsit5(), u0=z0, p=model.latentODEparams, saveat=t))
end

function latent_mu_logsd(model, x::SpiralSample)
    latent_dim = nhiddennodes(model.latentODEfunc.m)
    [model.rnn(xi) for xi in reverse(x, dims=1)[1:(end-1)]]
    rnn_encoded = model.rnn(reverse(x, dims=1)[end])
    μ = rnn_encoded[1:latent_dim]
    logσ = rnn_encoded[(latent_dim+1):end]
    μ, logσ
end

# p(x,z)
function log_normal_pdf(x,mean,logvar)
    constant = log(2*pi)
    -0.5*(constant .+ logvar .+ ((x.-mean).^2 ./ exp.(logvar)))
end

# Kullback-Leibler-divergence
kl_q_p(μ, logσ) = 0.5 * sum(exp.(2 .* logσ) + μ.^2 .- 1 .- (2 .* logσ))

nhiddennodes(d::Dense) = size(d.W, 1)
nhiddennodes(c::Chain) = nhiddennodes(c.layers[end])

# loss function - ELBO
function elbo(model, x::SpiralSample, t)
    empmu, emplogsd = latent_mu_logsd(model, x)
    Flux.reset!(model.rnn)
    z0 = latentz0(empmu, emplogsd)
    pred_z = n_ode(model, z0, t)
    sumlogp_x_z = sum([logp_x_z(x[i], model.decoder(pred_z[:,i])) for i in 1:size(pred_z,2)])
    sumlogp_x_z - kl_q_p(empmu, emplogsd)
end

function train!(model, xs, t;
        epochs = 20, learningrate = 0.01, monitoring = (args...) -> nothing)

    opt = ADAM(learningrate)
    ps = Flux.params(model.rnn, model.latentODEparams, model.decoder)
    cumloss = 0.0f0

    function loss(x)
        -elbo(model, x, t) + 0.01 * sum(x->sum(x.^2), Flux.params(model.rnn))
    end

    for epoch in 1:epochs
        local curloss
        for x in xs
            grads = Flux.gradient(ps) do
              curloss = loss(x)
              return curloss
            end
            cumloss += curloss
            Flux.Optimise.update!(opt, ps, grads)
        end
        monitoring(epoch, cumloss)
        cumloss = 0.0f0
    end
end

function predictspiral(model, x::SpiralSample, t)
    predμ, predlogσ = latent_mu_logsd(model, x)
    predz0 = latentz0(predμ, predlogσ)
    predz = n_ode(model, predz0, t)
    predx = [model.decoder(predz[:,i]) for i in 1:size(predz,2)]
end

end # module SpiralExample


text:
include("common.jl")
using Unitful

@testset "Units" begin
    prob_notinplace =
        DiffEqProblemLibrary.DDEProblemLibrary.build_prob_dde_1delay_long_scalar_notinplace(1.0u"N", 1.0u"s")
    prob_inplace = DiffEqProblemLibrary.DDEProblemLibrary.build_prob_dde_1delay_long(1.0u"N", 1.0u"s")

    @testset for prob in (prob_notinplace, prob_inplace)
        @testset "correct" begin
            # default tolerances
            alg1 = MethodOfSteps(Tsit5(), constrained=false, max_fixedpoint_iters=100)
            sol1 = solve(prob, alg1)

            # with correct units
            alg2 = MethodOfSteps(Tsit5(), constrained=false, max_fixedpoint_iters=100,
                                 fixedpoint_abstol=1e-6u"N", fixedpoint_reltol=1e-3u"N")
            sol2 = solve(prob, alg2)

            @test sol1.t == sol2.t && sol1.u == sol2.u

            # with correct units as vectors
            if typeof(prob.u0) <: AbstractArray
                alg3 = MethodOfSteps(Tsit5(), constrained=false, max_fixedpoint_iters=100,
                                     fixedpoint_abstol=[1e-6u"N"], fixedpoint_reltol=[1e-3u"N"])
                sol3 = solve(prob, alg3)

                @test sol1.t == sol3.t && sol1.u == sol3.u
            end
        end

        @testset "incorrect" begin
            # without units
            alg1 = MethodOfSteps(Tsit5(), constrained=false, max_fixedpoint_iters=100,
                                 fixedpoint_abstol=1e-6, fixedpoint_reltol=1e-3)
            @test_throws Unitful.DimensionError solve(prob, alg1)

            # with incorrect units for absolute tolerance
            alg2 = MethodOfSteps(Tsit5(), constrained=false, max_fixedpoint_iters=100,
                                fixedpoint_abstol=1e-6u"s", fixedpoint_reltol=1e-3u"N")
            @test_throws Unitful.DimensionError solve(prob, alg2)

            # with incorrect units for relative tolerance
            alg3 = MethodOfSteps(Tsit5(), constrained=false, max_fixedpoint_iters=100,
                                fixedpoint_abstol=1e-6u"N", fixedpoint_reltol=1e-3u"s")
            @test_throws Unitful.DimensionError solve(prob, alg3)
        end
    end
end


text:
#! /usr/bin/env julia
# Parsing
str = read("./3.in", String)
input = split(str, "\n")[1:end-1]
squares = [
    begin
        # Thanks to u/mserrano!
        id, x, y, w, h = Tuple(map(x -> parse(Int64, x.match), eachmatch(r"\d+", i)));
        (
            id = id,
            # Plus one 'cuz of Julia's indexing starting at one not zero
            idx = (x + 1 : x + w, y + 1 : y + h)
        )
    end
    for i in input
]

# Calculating the fabric with claims per square inch
fabric = zeros(Int64, 1000, 1000)
for s in squares
    v = view(fabric, s.idx...)
    v.+= 1
end

println("Solution for part 1:")
println(length(fabric[fabric .> 1]))

println("Solution for part 2:")
for s in squares
    if sum(fabric[s.idx[1], s.idx[2]] .- 1) == 0
        println(s.id)
    end
end

# Solution 1: 107663
# Solution 2: 1166

max_stars_repo_name: vypxl/aoc

max_issues_repo_name: vypxl/aoc

max_forks_repo_name: vypxl/aoc


text:
module TestGradient

using Printf

export test_gradient

struct GradientComponent
  i::Int
  g::Float64
  gbest::Float64
  error::Float64
  stepbest::Float64
end
Base.show(io::IO, c::GradientComponent) =
  @printf("%9i %12.5e %12.5e %12.5e %12.5e", c.i, c.g, c.gbest, c.error, c.stepbest)

title() = println("Component   Analytical     Discrete        Error    Best_Step")
function Base.show(io::IO, ::MIME"text/plain", c::AbstractVector{GradientComponent} )
  title()
  for i in 1:min(length(c),3)
    println(c[i])
  end
  if length(c) > 7
    @printf("%30s\n","⋮ ")
  end
  for i in max(4,length(c)-2):length(c)
    println(c[i])
  end
  max_error = findmax(map(x->x.error,c))
  print("""

  Maximum error = $(max_error[1]) at component: $(max_error[2])
  """)
  title()
  println(c[max_error[2]])
end

function discret(icomp,x,step,f)
  save = x[icomp]
  x[icomp] = save + step
  fplus = f(x)
  x[icomp] = save - step
  fminus = f(x)
  x[icomp] = save
  return (fplus - fminus) / (2 * step)
end

"""

```
test_gradient(f::Function,g!::Function,x::AbstractVector{T}) where T<:Real
```

Function that performs finite difference and analytical gradient
comparision. Used only for test purpouses. The output of `test_gradient`
is a vector of type `GradientComponent`, containing the index of the
component, the analytical and discrete gradientes, and the error associated
to the best step found for that component.

### Example

```julia-repl
julia> f(x) = sum(x.^2)
f (generic function with 1 method)

julia> g!(g,x) = 2*x
g! (generic function with 1 method)

julia> x = rand(20);

julia> test_gradient(f,g!,x)

Comparing analytical and finite-difference gradients... may take a while.

Function value on test point: 8.671516638852811

Component   Analytical     Discrete        Error    Best_Step
        1  1.38356e+00  1.38356e+00  8.66637e-15  1.00000e-02
        2  1.97347e+00  1.97347e+00  4.20806e-14  1.00000e-02
        3  1.53930e-01  1.53930e-01  3.17351e-14  1.00000e-02
                            ⋮
       18  1.89896e-01  1.89896e-01  1.12252e-13  1.00000e-02
       19  1.37622e+00  1.37622e+00  5.84064e-14  1.00000e-02
       20  1.62172e+00  1.62172e+00  2.19071e-15  1.00000e-02

Maximum error = 1.1225244035991744e-13 at component: 18
Component   Analytical     Discrete        Error    Best_Step
       18  1.89896e-01  1.89896e-01  1.12252e-13  1.00000e-02

```

"""
function test_gradient(f::Function,g!::Function,x::AbstractVector{T}) where T<:Real
  n = length(x)

  println("""

  Comparing analytical and finite-difference gradients... may take a while.
  """)

  fx = f(x)
  println("Function value on test point: ", fx)
  println()

  g = similar(x)
  g = g!(g,x)
  output = Vector{GradientComponent}(undef,0)
  t0 = time()
  eworst = 0.
  iworst = 1
  for i in 1:n
    step = 1.e-2
    stepbest = step
    error = +Inf
    gbest = 0.
    if mod(time() - t0,10) == 0
      println("Computing the $(i)th of $n components. Worst error: $eworst")
    end
    while error > 1.e-6 && step >= 1.e-20
      gcomp = discret(i,x,max(abs(g[i]*step),1e-20),f)
      if g[i] ≈ 0.
        steperror = abs(gcomp - g[i])
      else
        steperror = abs( ( gcomp - g[i] ) / g[i] )
      end
      if steperror < error
        error = steperror
        gbest = gcomp
        stepbest = step
      end
      step = step / 10
    end
    push!(output,GradientComponent(i, g[i], gbest, error, stepbest))
    if error > eworst
      iworst = i
      eworst = error
    end
  end
  return output
end

"""

```
test_gradient(f::Function,g!::Function,x::AbstractVector{T}) where T
```

Try to reinterpret the type of input into an array of `Float64`, and then
call `test_gradient(f,g,x::AbstractVector{<:Real})`

### Example

```julia-repl
julia> x = rand(SVector{2,Float64},10);

julia> f(x) = sum(v[1]^2 + v[2]^2 for v in x)
f (generic function with 1 method)

julia> g!(g,x) = @. g = (SVector{2,Float64}(2*v[1],2*v[2]) for v in x)
g! (generic function with 1 method)

julia> test_gradient(f,g!,x)

Comparing analytical and finite-difference gradients... may take a while.

Function value on test point: 6.5919481328739655

Computing the 1th of 20 components. Worst error: 0.0
Component   Analytical     Discrete        Error    Best_Step
        1  1.57869e+00  1.57869e+00  4.78215e-15  1.00000e-02
        2  1.04203e+00  1.04203e+00  1.91779e-14  1.00000e-02
        3  1.72248e+00  1.72248e+00  1.03128e-15  1.00000e-02
                            ⋮
       18  3.32460e-01  3.32460e-01  2.67153e-15  1.00000e-02
       19  1.03341e+00  1.03341e+00  8.16492e-15  1.00000e-02
       20  9.99455e-01  9.99455e-01  8.88663e-16  1.00000e-02

Maximum error = 2.2306785702197458e-12 at component: 7
Component   Analytical     Discrete        Error    Best_Step
        7  1.21440e-02  1.21440e-02  2.23068e-12  1.00000e-02

```

"""
function test_gradient(f::Function, g!::Function, x::AbstractVector{T}) where T
  test_gradient(
    x -> f(reinterpret(T,x)),
    (g,x) -> reinterpret(
      Float64,
      g!(
        reinterpret(T,g),
        reinterpret(T,x)
      ),
    ),
    reinterpret(Float64,x)
  )
end

end # module


text:
abstract type Tile end
start_location(::Tile) = [4,1]

struct L <: Tile end
data(tile::L) = [0 0 1; 1 1 1] * 2  # orange

struct J <: Tile end
data(tile::J) = [1 0 0; 1 1 1] * 6  # blue

struct S <: Tile end
data(tile::S) = [0 1 1; 1 1 0] * 4  # green

struct Z <: Tile end
data(tile::Z) = [1 1 0; 0 1 1] * 1  # red

struct T <: Tile end
data(tile::T) = [1 1 1; 0 1 0] * 7  # magenta

struct I <: Tile end
data(tile::I) = [1 1 1 1] * 5       # cyan
start_location(::I) = [4,2]

struct O <: Tile end
data(tile::O) = [1 1; 1 1] * 3      # yellow
start_location(::O) = [5,1]

const Tiles = [T, L, J, S, Z, I, O]

max_stars_repo_name: ChristianKurz/REPLTetris

max_issues_repo_name: ChristianKurz/REPLTetris

max_forks_repo_name: ChristianKurz/REPLTetris


text:
# This file is auto-generated by AWSMetadata.jl
using AWS
using AWS.AWSServices: sagemaker_edge
using AWS.Compat
using AWS.UUIDs

"""
    get_device_registration(device_fleet_name, device_name)
    get_device_registration(device_fleet_name, device_name, params::Dict{String,<:Any})

Use to check if a device is registered with SageMaker Edge Manager.

# Arguments
- `device_fleet_name`: The name of the fleet that the device belongs to.
- `device_name`: The unique name of the device you want to get the registration status from.

"""
get_device_registration(DeviceFleetName, DeviceName; aws_config::AbstractAWSConfig=global_aws_config()) = sagemaker_edge("POST", "/GetDeviceRegistration", Dict{String, Any}("DeviceFleetName"=>DeviceFleetName, "DeviceName"=>DeviceName); aws_config=aws_config)
get_device_registration(DeviceFleetName, DeviceName, params::AbstractDict{String, <:Any}; aws_config::AbstractAWSConfig=global_aws_config()) = sagemaker_edge("POST", "/GetDeviceRegistration", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("DeviceFleetName"=>DeviceFleetName, "DeviceName"=>DeviceName), params)); aws_config=aws_config)

"""
    send_heartbeat(agent_version, device_fleet_name, device_name)
    send_heartbeat(agent_version, device_fleet_name, device_name, params::Dict{String,<:Any})

Use to get the current status of devices registered on SageMaker Edge Manager.

# Arguments
- `agent_version`: Returns the version of the agent.
- `device_fleet_name`: The name of the fleet that the device belongs to.
- `device_name`: The unique name of the device.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"AgentMetrics"`: For internal use. Returns a list of SageMaker Edge Manager agent
  operating metrics.
- `"Models"`: Returns a list of models deployed on the the device.
"""
send_heartbeat(AgentVersion, DeviceFleetName, DeviceName; aws_config::AbstractAWSConfig=global_aws_config()) = sagemaker_edge("POST", "/SendHeartbeat", Dict{String, Any}("AgentVersion"=>AgentVersion, "DeviceFleetName"=>DeviceFleetName, "DeviceName"=>DeviceName); aws_config=aws_config)
send_heartbeat(AgentVersion, DeviceFleetName, DeviceName, params::AbstractDict{String, <:Any}; aws_config::AbstractAWSConfig=global_aws_config()) = sagemaker_edge("POST", "/SendHeartbeat", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("AgentVersion"=>AgentVersion, "DeviceFleetName"=>DeviceFleetName, "DeviceName"=>DeviceName), params)); aws_config=aws_config)


text:
function addcover!(G::NGrid{D, B}, L::Vector{Int}) where {D,B}
    any(all(G.covers .== L', 2)) && return
    bf = B == Linear ? cc_bf_l : cc_bf_q
    x   = kron([SparseInterp.cc_dg(i) for i in L])
    ind = repmat(map(Int16, L'), size(x, 1))
    level_M = map(i -> Int16(SparseInterp.M(level(i))), G.grid)

    for id = 1:size(x, 1)
        push!(G.Bs, [])
        push!(G.IDs, Int[])
        for i = 1:length(G)
            b = 1.0
            for d = 1:length(G.L)
                b *= bf(x[id, d], G.grid[i, d], level_M[i, d])
            end
            b > 0 && push!(G.Bs[end], -b)
            b > 0 && push!(G.IDs[end], i)
        end
        push!(G.Bs[end], 1.0)
        push!(G.IDs[end], length(G) + id)
    end
    G.grid = [G.grid; x]
    ind = map(level, G.grid)

    G.covers = [G.covers;L']
    G.covers_loc = Int32[findfirst(all(G.covers[i:i, :] .== ind, 2)) for i = 1:size(G.covers, 1)]
    G.covers_dM = map(x -> SparseInterp.dM(Int(x)), G.covers)
    G.adapt.active = [G.adapt.active;zeros(Bool, size(x, 1))]
    G.L = vec(maximum(G.covers, 1)) - 1
    return
end

function Base.sort!(G::NGrid)
    lev = level(G)
    try
        for i = 2:length(G)
            @assert lev[i - 1] ≤ lev[i]
        end
    catch
        id = zeros(Int, 0)
        for l = 0:maximum(lev)
            id = [id;find(lev .== l)]
        end
        G.grid = G.grid[id, :]
        ind = map(level, G.grid)
        G.covers = map(UInt16, unique(ind, 1))
        G.covers_dM = map(x -> dM(Int(x)), G.covers)
        G.covers_loc = Int32[findfirst(all(G.covers[i:i, :] .== ind, 2)) for i = 1:size(G.covers, 1)]
        SparseInterp.buildW(G)
    end
end

function iscoverfull(G::NGrid, L::Vector{Int})
    @assert length(L) == length(G.L)
    n = 0
    D = length(L)
    for i = 1:length(G)
        b = true
        for j = 1:D
            b = b && level(G.grid[i, j]) == L[j]
        end
        b && (n += 1)
    end
    return n == prod(map(dM, L))
end

looseid(G::NGrid) = G.covers_loc[end] + prod(G.covers_dM[end, :]):size(G.grid, 1)

"""
    grow!(G,id,bounds)

Adapt grid G at node number 'id'. 2*D nodes are added (except when G[id] is on
an edge in which case 2*D-e nodes are added where e is the number of edges G[id]
sits on). 'bounds' limits the maximum level the grid can grow to.
"""
function grow!(G::NGrid{D, BF}, id::Int, bounds::Vector{Int} = 12 * ones(Int, length(G.L))) where {D,BF}
    !G.adapt.active[id] && (return nothing)
    G.adapt.active[id] = false

    x       = G.grid[id, :]
    n       = D * 2 - sum((x[d] == 0.0 || x[d] == 1.0) for d = 1:D)
    Lid     = sum(map(level, G.grid[id, :])) - D
    X, ind   = SmolyakGrid(min(map(level, x) + 1, bounds), Lid + (1:1))

    id1     = (hX = mapslices(hash, G.grid, 2);Bool[!in(hash(X[i, :]), hX) for i = 1:size(X, 1)])
    all(!id1) && (return nothing)
    X = X[id1, :]
    id1 = sortperm(Float64[norm(X[i, :] - vec(x)) for i = 1:size(X, 1)])[1:min(n, size(X, 1))]
    X = X[id1, :]
    ind = map(level, X)

    hind = mapslices(hash, ind, 2)
    firstloc = length(G)
    for i = 1:n
        Hind = mapslices(hash, map(level, G.grid), 2) # need not be completely recalculated, can insert new hash
        if in(hind[i], Hind)
            insertloc = findlast(Hind, hind[i])
        else
            insertloc = max(findlast(sum(map(level, G.grid), 2), sum(ind[i, :])), length(G))
        end
        firstloc = min(firstloc, insertloc)

        G.grid          = [G.grid[1:insertloc, :];X[i, :]'G.grid[insertloc + 1:end, :]]
        G.Bs            = [G.Bs[1:insertloc];[Float64[]];G.Bs[insertloc + 1:end]]
        G.IDs           = [G.IDs[1:insertloc];[Float64[]];G.IDs[insertloc + 1:end]]
        G.adapt.active  = [G.adapt.active[1:insertloc];true;G.adapt.active[insertloc + 1:end]]
    end

    G.adapt.overlap = getparents(G)
    rebuildW(G, firstloc:length(G))

    return
end

function grow!(G::NGrid, id::Vector{Int}, bounds::Vector{Int} = 12 * ones(Int, length(G.L)))
    X = [G.grid[i, :] for i in id]
    for x in X
        grow!(G, findrow(G.grid, x), bounds)
    end
end

function rebuildW(G::NGrid{D, BF}, idr) where {D,BF}
    bf = (BF == Linear ? cc_bf_l : cc_bf_q)
    level_M = map(i -> Int16(SparseInterp.M(level(i))), G.grid)
    for j = idr
        G.Bs[j] = Float64[]
        G.IDs[j] = Int[]
        for i = 1:j - 1
            b = 1.0
            for d = 1:D
                b *= bf(G.grid[j, d], G.grid[i, d], level_M[i, d])
            end
            if b > 0
                push!(G.Bs[j], -b)
                push!(G.IDs[j], i)
            end
        end
        push!(G.Bs[j], 1.0)
        push!(G.IDs[j], length(G) + j)
    end
end

mutable struct box{D}
    lower::Vector{Float64}
    upper::Vector{Float64}
end

function box(G::NGrid{D, BF}, i::Int) where {D,BF}
    lower = Array{Float64}(D)
    upper = Array{Float64}(D)
    for d = 1:D
        p = level(G.grid[i, d])
        lower[d] = clamp(G.grid[i, d] - 1 / 2^(p - 1), 0, 1) + eps(Float64)
        upper[d] = clamp(G.grid[i, d] + 1 / 2^(p - 1), 0, 1) - eps(Float64)
    end
    return box{D}(lower, upper)
end

function Base.intersect(a::box{D}, b::box{D}) where {D}
    nool = a.lower[1] > b.upper[1] ||
           b.lower[1] > a.upper[1]
    for d = 2:D
        nool = nool ||
           a.lower[d] > b.upper[d] ||
           b.lower[d] > a.upper[d]
    end
    return !nool
end

"""
   getparents(G)

For each node of grid G this generates a Boolean Vector indicating whether the
basis function overlaps with that of the other nodes.
"""
function getparents(G::NGrid)
    T = [zeros(Bool, length(G)) for i = 1:length(G)]
    bx = [box(G, i) for i = 1:length(G)]
    @threads for i = 1:length(G)
        for j = 1:length(G)
            T[i][j] = intersect(bx[i], bx[j])
        end
    end
    return T
end

# Generate interpolation functions for single function approximation
for b in [(Linear, Lbj, cc_bf_l), (Quadratic, Qbj, cc_bf_q)]
    for D = 2:15
        coverloop = :(@inbounds for ii in nc
            b   = B[G.covers[ii, $D], $D] * B[G.covers[ii, 1], 1]
            id1 = J[G.covers[ii, $D], $D]
        end)
        for d in D - 1:-1:2
            push!(coverloop.args[2].args[2].args, :(b  *= B[G.covers[ii, $d], $d]))
            push!(coverloop.args[2].args[2].args, :(id1 = id1 * G.covers_dM[ii, $d] + J[G.covers[ii, $d], $d]))
        end

        push!(coverloop.args[2].args[2].args, :(id1 = id1 * G.covers_dM[ii, 1] + G.covers_loc[ii] + J[G.covers[ii, 1], 1];yi += b * w[id1]))

        f = :(function jl_ainterp(G::NGrid{$D, $(b[1])}, A::Vector{Float64}, xi::Array{Float64, 2}, y = zeros(Float64, size(xi, 1)))
            w       = getW(G, A)
            x       = nXtoU(xi, G.bounds)
            nx      = size(x, 1)
            nc      = 1:size(G.covers, 1)
            nG      = length(G)
            ns      = (G.covers_loc[end] + prod(G.covers_dM[end, :]))
            nc2     = ns:UInt32(nG)
            dr1 = 1:$D
            mL      = maximum(G.L) + 1
            J         = zeros(Int, mL, $D)
            B         = ones(mL, $D)
            # dM = map(x -> M(level(x)), G.grid)

            @threadsfixed [J, B] for i = 1:nx
                $(b[2])
                yi = 0.0
                coverloop

                i1, i2 = UInt32(1), UInt32(1)
                iswitch = false

                @inbounds for ii = nc2
                    if G.adapt.overlap[i1][ii] && G.adapt.overlap[i2][ii]
                        b = 1.0
                        for d = 1:$D
                            # b *= $(b[3])(x[i, d], G.grid[ii, d], dM[ii, d])
                            b *= $(b[3])(x[i, d], G.grid[ii, d], M(level(G.grid[ii, d])))
                            b == 0 && break
                        end
                        if b > 0
                            (yi += w[ii] * b)
                            iswitch = !iswitch
                            iswitch ? (i1 = ii::UInt32) : (i2 = ii::UInt32)
                        end
                    end
                end
                y[i] = yi
            end
            y
        end)

        subs!(f, :coverloop => coverloop)
        eval(current_module(), f)
    end
end

# Generate interpolation functions for multiple function approximation
for b in [(Linear, Lbj, cc_bf_l), (Quadratic, Qbj, cc_bf_q)]
    for D = 2:15
        coverloop = quote
            fill!(yi, 0.0)
            @inbounds for ii in nc
                b   = B[G.covers[ii, $D], $D] * B[G.covers[ii, 1], 1]
                id1 = J[G.covers[ii, $D], $D]
            end
        end
        for d in D - 1:-1:2
            push!(coverloop.args[4].args[2].args[2].args, :(b  *= B[G.covers[ii, $d], $d]))
            push!(coverloop.args[4].args[2].args[2].args, :(id1 = id1 * G.covers_dM[ii, $d] + J[G.covers[ii, $d], $d]))
        end
        push!(coverloop.args[4].args[2].args[2].args, :(id1 = id1 * G.covers_dM[ii, 1] + G.covers_loc[ii] + J[G.covers[ii, 1], 1]))

        push!(coverloop.args[4].args[2].args[2].args, :(for d = 1:nA yi[d] += b * w[id1, d] end))
        push!(coverloop.args, :(@inbounds for d = 1:nA y[i, d] = yi[d] end))

        f = :(function jl_ainterp(G::NGrid{$D, $(b[1])}, A::Array{Float64, 2}, xi::Array{Float64, 2}, y = zeros(Float64, size(xi, 1), size(A, 2)))
            w       = getW(G, A)
            x       = nXtoU(xi, G.bounds)
            nx      = size(x, 1)
            nA      = size(A, 2)
            nc      = 1:size(G.covers, 1)
            nG      = length(G)
            ns      = (G.covers_loc[end] + prod(G.covers_dM[end, :]))
            nc2     = ns:UInt32(nG)
            dr1 = 1:$D
            mL      = maximum(G.L) + 1
            J         = zeros(Int, mL, $D)
            B         = ones(mL, $D)
            dM = map(xij -> M(level(xij)), G.grid)
            yi = zeros(nA)

            @threadsfixed [J, B, yi] for i = 1:nx
                $(b[2])
                coverloop

                i1, i2 = UInt32(1), UInt32(1)
                iswitch = false

                @inbounds for ii = nc2
                    if G.adapt.overlap[i1][ii] && G.adapt.overlap[i2][ii]
                        b = 1.0
                        for d = 1:$D
                            b *= $(b[3])(x[i, d], G.grid[ii, d], dM[ii, d])
                            b == 0 && break
                        end
                        if b > 0
                            (yi += w[ii, :] * b)
                            iswitch = !iswitch
                            iswitch ? (i1 = ii::UInt32) : (i2 = ii::UInt32)
                        end
                    end
                end
                y[i, :] = yi
            end
            y
        end)

        subs!(f, :coverloop => coverloop)
        eval(current_module(), f)
    end
end


text:
function optdata(ex::Expr)
    allindices = getallindices(ex)
    cost = Power{:χ}(1,1)
    return Dict{Any, typeof(cost)}(i=>cost for i in allindices)
end

function optdata(optex::Expr, ex::Expr)
    if optex.head == :tuple
        isempty(optex.args) && return nothing
        args = optex.args
        if all(x -> isa(x, Expr) && x.head == :call && x.args[1] == :(=>), args)
            indices, costs = _optdata(map(x -> x.args[2], args), map(x -> x.args[3], args))
            costtype = promote_type(typeof.(costs)...)
            costs = convert(Vector{costtype}, costs)
        elseif all(x -> isa(x, Expr) && x.head == :(=), args)
            indices, costs = _optdata(map(x -> x.args[1], args), map(x -> x.args[2], args))
            costtype = promote_type(typeof.(costs)...)
            costs = convert(Vector{costtype}, costs)
        else
            indices = map(normalizeindex, args)
            costtype = Power{:χ,Int}
            costs = fill(Power{:χ,Int}(1,1), length(args))
        end
        return Dict{Any, costtype}(k=>v for (k,v) in zip(indices, costs))
    elseif optex.head == :call && optex.args[1] == :!
        allindices = unique(getallindices(ex))
        excludeind = map(normalizeindex, optex.args[2:end])
        cost = Power{:χ}(1,1)
        d = Dict{Any, typeof(cost)}(i=>cost for i in allindices)
        for i in excludeind
            d[i] = 1
        end
        return d
    else
        error("invalid index cost specification")
    end
end

function _optdata(indexvec, costvec)
    indices = Vector{Any}()
    costs = Vector{Any}()
    for (index, cost) in zip(indexvec, costvec)
        if typeof(index) != Symbol && index.head == :tuple
            for index_ in index.args
                push!(indices, normalizeindex(index_))
                push!(costs, parsecost(cost))
            end
        else
            push!(indices, normalizeindex(index))
            push!(costs, parsecost(cost))
        end
    end
    indices, costs
end

# Process index cost specification for @tensoropt and friends
function parsecost(ex::Expr)
    if ex.head == :call && ex.args[1] == :*
        return *(map(parsecost, ex.args[2:end])...)
    elseif ex.head == :call && ex.args[1] == :+
        return +(map(parsecost, ex.args[2:end])...)
    elseif ex.head == :call && ex.args[1] == :-
        return -(map(parsecost, ex.args[2:end])...)
    elseif ex.head == :call && ex.args[1] == :^
        return ^(map(parsecost, ex.args[2:end])...)
    elseif ex.head == :call && ex.args[1] == :/
        return /(map(parsecost, ex.args[2:end])...)
    elseif ex.head == :call && ex.args[1] == :big
        return big(map(parsecost, ex.args[2:end])...)
    elseif ex.head == :call && ex.args[1] == :float
        return float(map(parsecost, ex.args[2:end])...)
    elseif ex.head == :call && ex.args[1] == :Int128
        return Int128(map(parsecost, ex.args[2:end])...)
    else
        error("invalid index cost specification: $ex")
    end
end
parsecost(ex::Number) = ex
parsecost(ex::Symbol) = Power{ex}(1,1)


text:
using DelayDiffEq, DiffEqProblemLibrary.DDEProblemLibrary
using Test

DDEProblemLibrary.importddeproblems()

const prob_ip = DDEProblemLibrary.prob_dde_constant_1delay_ip
const prob_scalar = DDEProblemLibrary.prob_dde_constant_1delay_scalar
const ts = 0:0.1:10

# ODE algorithms
const algs = [Rosenbrock23(), Rosenbrock32(), ROS3P(), Rodas3(),
              RosShamp4(), Veldd4(), Velds4(), GRK4T(), GRK4A(),
              Ros4LStab(), Rodas4(), Rodas42(), Rodas4P(), Rodas5()]

@testset "Algorithm $(nameof(typeof(alg)))" for alg in algs
  println(nameof(typeof(alg)))

  stepsalg = MethodOfSteps(alg)
  sol_ip = solve(prob_ip, stepsalg)
  sol_scalar = solve(prob_scalar, stepsalg)

  @test sol_ip(ts, idxs=1) ≈ sol_scalar(ts)
  @test sol_ip.t ≈ sol_scalar.t
  @test sol_ip[1, :] ≈ sol_scalar.u
end

max_stars_repo_name: JuliaDiffEq/DelayDiffEq

max_issues_repo_name: JuliaDiffEq/DelayDiffEq

max_forks_repo_name: JuliaDiffEq/DelayDiffEq


text:
# this file is just a bunch of ugly code to make sure that the Julia Array{UInt8,1} becomes a JS Uint8Array() instead of a normal array. This improves performance of the client.
# ignore it if you are not interested in that kind of stuff

import Dates
import UUIDs: UUID
import MsgPack
import .Configuration
import Pkg

# MsgPack.jl doesn't define a serialization method for MIME and UUID objects, so we write these ourselves:
MsgPack.msgpack_type(::Type{<:MIME}) = MsgPack.StringType()
MsgPack.msgpack_type(::Type{UUID}) = MsgPack.StringType()
MsgPack.msgpack_type(::Type{VersionNumber}) = MsgPack.StringType()
MsgPack.msgpack_type(::Type{Pkg.Types.VersionRange}) = MsgPack.StringType()
MsgPack.to_msgpack(::MsgPack.StringType, m::MIME) = string(m)
MsgPack.to_msgpack(::MsgPack.StringType, u::UUID) = string(u)
MsgPack.to_msgpack(::MsgPack.StringType, v::VersionNumber) = string(v)
MsgPack.to_msgpack(::MsgPack.StringType, v::Pkg.Types.VersionRange) = string(v)

# Support for sending Dates
MsgPack.msgpack_type(::Type{Dates.DateTime}) = MsgPack.ExtensionType()
MsgPack.to_msgpack(::MsgPack.ExtensionType, d::Dates.DateTime) = let
    millisecs_since_1970_because_thats_how_computers_work = Dates.value(d - Dates.DateTime(1970))
    MsgPack.Extension(0x0d, reinterpret(UInt8, [millisecs_since_1970_because_thats_how_computers_work]))
end

# Our Configuration types:
MsgPack.msgpack_type(::Type{Configuration.Options}) = MsgPack.StructType()
MsgPack.msgpack_type(::Type{Configuration.EvaluationOptions}) = MsgPack.StructType()
MsgPack.msgpack_type(::Type{Configuration.CompilerOptions}) = MsgPack.StructType()
MsgPack.msgpack_type(::Type{Configuration.ServerOptions}) = MsgPack.StructType()
MsgPack.msgpack_type(::Type{Configuration.SecurityOptions}) = MsgPack.StructType()

# Don't try to send callback functions which can't be serialized (see ServerOptions.event_listener)
MsgPack.msgpack_type(::Type{Function}) = MsgPack.NilType()
MsgPack.to_msgpack(::MsgPack.NilType, ::Function) = nothing

# We want typed integer arrays to arrive as JS typed integer arrays:
const JSTypedIntSupport = [Int8, UInt8, Int16, UInt16, Int32, UInt32, Float32, Float64]
JSTypedInt = Union{Int8,UInt8,Int16,UInt16,Int32,UInt32,Float32,Float64}

MsgPack.msgpack_type(::Type{Vector{T}}) where T <: JSTypedInt = MsgPack.ExtensionType()
MsgPack.to_msgpack(::MsgPack.ExtensionType, x::Vector{T}) where T <: JSTypedInt = let
    type = findfirst(isequal(T), JSTypedIntSupport) + 0x10
    MsgPack.Extension(type, reinterpret(UInt8, x))
end
MsgPack.msgpack_type(::Type{Vector{Union{}}}) = MsgPack.ArrayType()

# The other side does the same (/frontend/common/MsgPack.js), and we decode it here:
function decode_extension_and_addbits(x::MsgPack.Extension)
    if x.type == 0x0d
        # the datetime type
        millisecs_since_1970_because_thats_how_computers_work = reinterpret(Int64, x.data)[1]
        Dates.DateTime(1970) + Dates.Millisecond(millisecs_since_1970_because_thats_how_computers_work)
        # TODO? Dates.unix2datetime does exactly this ?? - DRAL
    else
        # the array types
        julia_type = JSTypedIntSupport[x.type - 0x10]
        if eltype(x.data) == julia_type
            x.data
        else
            reinterpret(julia_type, x.data)
        end
    end
end

function decode_extension_and_addbits(x::Dict)
    # we mutate the dictionary, that's fine in our use case and saves memory?
    for (k, v) in x
        x[k] = decode_extension_and_addbits(v)
    end
    x
end
decode_extension_and_addbits(x::Array) = map(decode_extension_and_addbits, x)

# We also convert everything (except the JS typed arrays) to 64 bit numbers, just to make it easier to work with.
decode_extension_and_addbits(x::T) where T <: Union{Signed,Unsigned} = Int64(x)
decode_extension_and_addbits(x::T) where T <: AbstractFloat = Float64(x)
decode_extension_and_addbits(x::Any) = x

function pack(args...)
    MsgPack.pack(args...)
end

function unpack(args...)
    MsgPack.unpack(args...) |> decode_extension_and_addbits
end


text:
"""
```
HetDSGEGovDebt{T} <: AbstractHeterogeneousModel{T}
```

### Fields

#### Parameters and Steady-States
* `parameters::Vector{AbstractParameter}`: Vector of all time-invariant model
  parameters.

* `steady_state::Vector{AbstractParameter}`: Model steady-state values, computed
  as a function of elements of `parameters`.

* `keys::OrderedDict{Symbol,Int}`: Maps human-readable names for all model
  parameters and steady-states to their indices in `parameters` and
  `steady_state`.

#### Inputs to Measurement and Equilibrium Condition Equations

The following fields are dictionaries that map human-readable names to row and
column indices in the matrix representations of of the measurement equation and
equilibrium conditions.

* `endogenous_states::OrderedDict{Symbol,Int}`: Maps each state to a column in
  the measurement and equilibrium condition matrices.

* `exogenous_shocks::OrderedDict{Symbol,Int}`: Maps each shock to a column in
  the measurement and equilibrium condition matrices.

* `expected_shocks::OrderedDict{Symbol,Int}`: Maps each expected shock to a
  column in the measurement and equilibrium condition matrices.

* `equilibrium_conditions::OrderedDict{Symbol,Int}`: Maps each equlibrium
  condition to a row in the model's equilibrium condition matrices.

* `observables::OrderedDict{Symbol,Int}`: Maps each observable to a row in the
  model's measurement equation matrices.

#### Model Specifications and Settings

* `spec::String`: The model specification identifier, \"an_schorfheide\", cached
  here for filepath computation.

* `subspec::String`: The model subspecification number, indicating that some
  parameters from the original model spec (\"ss0\") are initialized
  differently. Cached here for filepath computation.

* `settings::Dict{Symbol,Setting}`: Settings/flags that affect computation
  without changing the economic or mathematical setup of the model.

* `test_settings::Dict{Symbol,Setting}`: Settings/flags for testing mode

#### Other Fields

* `rng::MersenneTwister`: Random number generator. Can be is seeded to ensure
  reproducibility in algorithms that involve randomness (such as
  Metropolis-Hastings).

* `testing::Bool`: Indicates whether the model is in testing mode. If `true`,
  settings from `m.test_settings` are used in place of those in `m.settings`.

* `observable_mappings::OrderedDict{Symbol,Observable}`: A dictionary that
  stores data sources, series mnemonics, and transformations to/from model units.
  DSGE.jl will fetch data from the Federal Reserve Bank of
  St. Louis's FRED database; all other data must be downloaded by the
  user. See `load_data` and `Observable` for further details.

"""
mutable struct HetDSGEGovDebt{T} <: AbstractHetModel{T}
    parameters::ParameterVector{T}               # vector of all time-invariant model parameters
    steady_state::ParameterVector{T}             # model steady-state values

    # Temporary to get it to work. Need to
    # figure out a more flexible way to define
    # "grids" that are not necessarily quadrature
    # grids within the model
    grids::OrderedDict{Symbol,Union{Grid, Array}}
    keys::OrderedDict{Symbol,Int}                    # Human-readable names for all the model
                                              # parameters and steady-states

    state_variables::Vector{Symbol}            # Vector of symbols of the state variables
    jump_variables::Vector{Symbol}                   # Vector of symbols of the jump variables
    normalized_model_states::Vector{Symbol}          # All of the distributional model
                                                     # state variables that need to be normalized
    # Vector of unnormalized ranges of indices
    # endogenous_states_unnormalized::OrderedDict{Symbol,UnitRange}

    # Vector of ranges corresponding to normalized (post Klein solution) indices
    endogenous_states::OrderedDict{Symbol,UnitRange}
    endogenous_states_original::OrderedDict{Symbol,UnitRange}

    exogenous_shocks::OrderedDict{Symbol,Int}
    expected_shocks::OrderedDict{Symbol,Int}
    equilibrium_conditions::OrderedDict{Symbol,UnitRange}
    endogenous_states_augmented::OrderedDict{Symbol, Int}
    observables::OrderedDict{Symbol,Int}

    spec::String                                     # Model specification number (eg "m990")
    subspec::String                                  # Model subspecification (eg "ss0")
    settings::Dict{Symbol,Setting}                   # Settings/flags for computation
    test_settings::Dict{Symbol,Setting}              # Settings/flags for testing mode
    rng::MersenneTwister                             # Random number generator
    testing::Bool                                    # Whether we are in testing mode or not
    observable_mappings::OrderedDict{Symbol, Observable}
end

description(m::HetDSGEGovDebt) = "HetDSGEGovDebt, $(m.subspec)"

"""
`init_model_indices!(m::HetDSGEGovDebt)`

Arguments:
`m:: HetDSGEGovDebt`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and equilibrium conditions.
"""
function init_model_indices!(m::HetDSGEGovDebt, states::Vector{Symbol}, jumps::Vector{Symbol})

    endogenous_states = collect(vcat(states, jumps))

    # Exogenous shocks
    exogenous_shocks = collect([:b_sh,:g_sh,:z_sh,:μ_sh,:λ_w_sh, :λ_f_sh,:rm_sh])

    # Equilibrium conditions
    equilibrium_conditions = collect([:eq_euler,:eq_kolmogorov_fwd,
                                      :eq_market_clearing,:eq_lambda, :eq_transfers,
                                      :eq_investment,:eq_tobin_q,:eq_capital_accumulation,
                                      :eq_wage_phillips,:eq_price_phillips,:eq_marginal_cost,
                                      :eq_gdp,:eq_optimal_kl,:eq_taylor, :eq_fisher,
                                      :eq_nominal_wage_inflation, :eq_fiscal_rule,
                                      :eq_g_budget_constraint, :LI,:LY,
                                      :LW,:LX,:eq_b,:eq_g,:eq_z,:eq_μ,:eq_λ_w,
                                      :eq_λ_f, :eq_rm])

    # Additional states added after solving model
    # Lagged states and observables measurement error
    endogenous_states_augmented = [:C_t1]

    # Observables
    observables = keys(m.observable_mappings)

    ########################################################################################
    # Setting indices of endogenous_states and equilibrium conditions manually for now

    setup_indices!(m)
    endo = m.endogenous_states #_unnormalized
    m.endogenous_states_original = deepcopy(endo)
    m <= Setting(:n_model_states_original, m.endogenous_states_original[jumps[end]][end])#first(collect(values(m.endogenous_states_original))[end]))
    eqcond = equilibrium_conditions
    ########################################################################################

    m.normalized_model_states = [:kf′_t]

    for (i,k) in enumerate(exogenous_shocks); m.exogenous_shocks[k] = i end
    m.observables[:obs_gdp]         = 1
    m.observables[:obs_hours]       = 2
    m.observables[:obs_wages]       = 3
    m.observables[:obs_gdpdeflator] = 4
    m.observables[:obs_nominalrate] = 5
    m.observables[:obs_consumption] = 6
    m.observables[:obs_investment]  = 7
    #for (i,k) in enumerate(observables);      m.observables[k]      = i end
end

function HetDSGEGovDebt(subspec::String="ss0";
                        custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                        testing = false, testing_gamma::Bool = false,
                        ref_dir = "")

    # Model-specific specifications
    spec               = "het_dsge"
    subspec            = subspec
    settings           = Dict{Symbol,Setting}()
    test_settings      = Dict{Symbol,Setting}()
    rng                = MersenneTwister(0)

    # initialize empty model
    m = HetDSGEGovDebt{Float64}(
            # model parameters and steady state values
            Vector{AbstractParameter{Float64}}(), Vector{Float64}(),
            # grids and keys
            OrderedDict{Symbol,Union{Grid, Array}}(), OrderedDict{Symbol,Int}(),

            # normalized_model_states, state_inds, jump_inds
            Vector{Symbol}(), Vector{Symbol}(), Vector{Symbol}(),

            # model indices
            # endogenous states unnormalized, endogenous states normalized
            OrderedDict{Symbol,UnitRange}(), OrderedDict{Symbol,UnitRange}(),
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,UnitRange}(), # OrderedOrderedDict{Symbol,UnitRange}(),
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),

            spec,
            subspec,
            settings,
            test_settings,
            rng,
            testing,
            OrderedDict{Symbol,Observable}())

    default_settings!(m)

    m <= Setting(:ref_dir, ref_dir, "Absolute filepath to reference directory")

    # # Set observable transformations
    init_observable_mappings!(m)

    # Set settings
    model_settings!(m)

    # default_test_settings!(m)
    for custom_setting in values(custom_settings)
        m <= custom_setting
    end

    # Endogenous states
    states = collect([:kf′_t,:k′_t,:i′_t1, :y′_t1,:w′_t1,:I′_t1, :bg′_t,
                      :b′_t,:g′_t,:z′_t,:μ′_t,:λ_w′_t, :λ_f′_t,:rm′_t])

    jumps = collect([:l′_t,:C′_t,:R′_t,:i′_t,:t′_t,:w′_t, :L′_t,:π′_t,:π_w′_t,
                     :margutil′_t,:y′_t, :I′_t,:mc′_t,:Q′_t,:capreturn′_t, :tg′_t])

    # Initialize model indices
    init_model_indices!(m, states, jumps)

    # Init what will keep track of # of states, jumps, and states/jump indices
    # (need model indices first)
    init_states_and_jumps!(m, states, jumps)

    # Initialize parameters
    init_parameters!(m, testing_gamma = testing_gamma)

    # Initialize aggregate steady state parameters (necessary for grid construction)
    aggregate_steadystate!(m)

    # Initialize grids
    init_grids!(m)

    # Solve for the steady state
    #steadystate!(m)

    # So that the indices of m.endogenous_states reflect the normalization
    normalize_model_state_indices!(m)

    endogenous_states_augmented = [:C_t1]

    for (i,k) in enumerate(endogenous_states_augmented)
        m.endogenous_states_augmented[k] = i + first(m.endogenous_states[get_setting(m, :jumps)[end]])#first(collect(values(m.endogenous_states))[end])
    end
    m <= Setting(:n_model_states_augmented, get_setting(m, :n_model_states) +
                 length(m.endogenous_states_augmented))

    init_subspec!(m)

    return m
end

"""
```
init_parameters!(m::HetDSGEGovDebt)
```

Initializes the model's parameters, as well as empty values for the steady-state
parameters (in preparation for `steadystate!(m)` being called to initialize
those).
"""
function init_parameters!(m::HetDSGEGovDebt; testing_gamma::Bool = false)
    ######################################
    # Parameters that affect steady-state
    ######################################
    m <= parameter(:α, 0.3, fixed = false, (1e-5, 0.999), (1e-5, 0.999),
                   SquareRoot(), Normal(0.30, 0.05),
                   description = "α: Capital elasticity in the intermediate goods" *
                   "sector's production function (also known as the capital share).",
                   tex_label = "\\alpha")
    # Check this: Previously the description was "Aggregate hours worked"
    m <= parameter(:H, 1.0, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed = true,
                   description = "Lmean: Mean level of hours", tex_label = "H")
    m <= parameter(:δ, 0.03, fixed = true,
                   description = "δ: The capital depreciation rate", tex_label = "\\delta")
    m <= parameter(:μ_sp, 0.0, fixed = true,
                   description = "μ_sp: The trend in the skill process",
                   tex_label = "\\mu_{sp}")

    # Exogenous processes - level
    # Uncomment scaling once adjusted properly in the code
    if testing_gamma == true
        m <= parameter(:γ, 0.0, (-5.0, 5.0), (-5., 5.), Untransformed(),
                       Normal(0.4, 0.1), fixed = false, scaling = x -> x/100,
                       description = "γ: The log of the steady-state growth rate of technology",
                       tex_label="100\\gamma")
    else
        m <= parameter(:γ, 0.5, (-5.0, 5.0), (-5., 5.), Untransformed(),
                       Normal(0.4, 0.1), fixed = false, scaling = x -> x/100,
                       description = "γ: The log of the steady-state growth rate of technology",
                       tex_label="100\\gamma")
    end

    m <= parameter(:r, 0.5, (1e-5, 10.0), (1e-5, 10.0), ModelConstructors.Exponential(),
                   GammaAlt(0.5, 0.5), fixed = false, scaling = x -> x/100,
                   description= "r: Quarterly steady-state real interest rate.",
                   tex_label= "100*r^{HetDSGE}")

    m <= parameter(:g, 1/(1-0.01), fixed = true,
                   description = "g_star: 1 - (c_star + i_star)/y_star",
                   tex_label = "g_*")

    m <= parameter(:β_save, 0.0, fixed = true,
                   description = "saving the betas per particle",
                   tex_label = "\\beta_save")
    m <= parameter(:sH_over_sL, 6.33333, (3.0, 9.0), (3.0, 9.0), Untransformed(),
                   Uniform(3.0, 9.0), fixed = true,
                   description = "Ratio of high to low earners", tex_label = "s_H / s_L")

    m <= parameter(:pLH, 0.005, (0.0025, 0.095), (0.0025, 0.095), Untransformed(),
                   Uniform(0.005, 0.095), fixed = false,
                   description = "Prob of going from low to high persistent skill",
                   tex_label = "p(s_L \\mid s_H)")
    m <= parameter(:pHL, 0.03, (0.0025, 0.095), (0.0025, 0.095), Untransformed(),
                   Uniform(0.0025, 0.095), fixed = false,
                   description = "Prob of going from high to low persistent skill",
                   tex_label = "p(s_H \\mid s_L)")

    m <= parameter(:BoverY, 0.26, fixed = true, description = "B / Y", tex_label = "B / Y")

    m <= parameter(:zlo, 0.0323232, (1e-18, 0.8-eps()), (1e-18, 0.8-eps()), Untransformed(),
                   Uniform(1e-18, 0.8-eps()), fixed = true,
                   description = "Lower bound on second income shock to mollify actual income",
                   tex_label = "\\underbar{z}")

    m <= parameter(:zhi, 2-m[:zlo].value, fixed = true,
                   description = "Upper bound on second income shock to mollify actual income",
                   tex_label = "\\bar{z}")

    m <= parameter(:mpc, 0.23395,  fixed = true, tex_label = "MPC")
    m <= parameter(:pc0, 0.071893, fixed = true, description = "Number of people at 0 income",
                   tex_label = "pc0")

    # Give model new parameters
    m <= parameter(:varlinc, 0.0, fixed = true, tex_label = "varlinc",
                   description = "var(log(annual income))")
    m <= parameter(:vardlinc, 0.0, fixed = true, tex_label = "vardlinc",
                   description = "var(log(deviations in annual income))")

    # Not in m1002
    m <= parameter(:η, 0.0, description = "η: Borrowing constraint (normalized by TFP)",
                   tex_label = "\\eta")

    ####################################################
    # Parameters that affect dynamics (not steady-state)
    ####################################################
    m <= parameter(:spp, 4., (-15., 15.), (-15., 15.), Untransformed(),
                   Normal(4., 1.5), fixed = false,
                   description = "S'': The second derivative of households' cost " *
                   "of adjusting investment.", tex_label = "S''")

    m <= parameter(:ϕh, 2., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   Normal(2, 0.75), fixed = false, description = "inverse frisch elasticity",
                   tex_label = "\\phi_h")

    m <= parameter(:κ_p, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = false, description = "κ_p : The slope of the Price Phillips curve",
                   tex_label = "\\kappa_p")
    m <= parameter(:κ_w, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = false, description = "κ_w: The slope of the Wage Phillips curve",
                   tex_label = "\\kappa_w")

    m <= parameter(:ρR , 0.75, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.75, 0.10), fixed = false,
                   description = "ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho_R")
    m <= parameter(:ψπ , 1.5, (1e-5, 10.), (1e-5, 10.0), ModelConstructors.Exponential(),
                   Normal(1.5, 0.25), fixed = false,
                   description = "ψ1: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1")
    m <= parameter(:ψy , 0.5, (-0.5, 0.5), (-0.5, 0.5), Untransformed(),
                   Normal(0.12, 0.05), fixed = false,
                   description = "ψy: Weight on output growth in monetary policy rule",
                   tex_label = "\\psi_y")

    m <= parameter(:δb, 1., (0.0, 1.0), (0.0, 1.0), Untransformed(),
                   Uniform(0.0, 1.0), fixed = false, description = "=1 means balanced budget",
                   tex_label = "\\delta_b")

    # Exogenous processes - autocorrelation
    m <= parameter(:ρ_g, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = false,
                   description = "ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g")
    m <= parameter(:ρ_b, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = false,
                   description = "ρ_b: AR(1) coefficient in intertemporal preference " *
                   "shift process.", tex_label = "\\rho_B")
    m <= parameter(:ρ_μ, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = false,
                   description = "ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label = "\\rho_{\\mu}")
    m <= parameter(:ρ_z, 0.5,(1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = false,
                   description = "ρ_z: AR(1) coefficient in the technology process.",
                   tex_label = "\\rho_z")
    m <= parameter(:ρ_λ_f, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = false,
                   description = "ρ_λ_f: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_f}")
    m <= parameter(:ρ_λ_w, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = false,
                   description = "ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}")
    m <= parameter(:ρ_rm, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = false,
                   description = "ρ_rm: AR(1) coefficient in the monetary policy shock process.",
                   tex_label = "\\rho_{r^m}")

    m <= parameter(:σ_g, 0.15, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = false,
                   description = "σ_g: standard dev. of the government spending process.",
                   tex_label = "\\sigma_{g}")
    m <= parameter(:σ_b, 0.15, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = false,
                   description = "σ_b: standard dev. of the intertemporal preference " *
                   "shifter process.", tex_label = "\\sigma_{b}")
    m <= parameter(:σ_μ, 0.15, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = false,
                   description = "σ_μ: standard dev. of the exogenous marginal efficiency" *
                   " of investment shock process.", tex_label = "\\sigma_{\\mu}")
    m <= parameter(:σ_z, 0.15,(1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = false,
                   description = "σ_z: standard dev. of the process describing the " *
                   "stationary component of productivity.", tex_label = "\\sigma_z")
    m <= parameter(:σ_λ_f, 0.15, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = false,
                   description = "σ_λ_f: mean of the process that generates the price " *
                   "elasticity of the composite good. Specifically, the elasticity is " *
                   "(1+λ_{f,t})/(λ_{f,t}).", tex_label = "\\sigma_{\\lambda_f}")
    m <= parameter(:σ_λ_w, 0.15, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = false,
                   description = "σ_λ_w", tex_label = "\\sigma_{\\lambda_w}")
    m <= parameter(:σ_rm, 0.15, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = false,
                   description = "σ_r_m: standard dev. of the monetary policy shock.",
                   tex_label = "\\sigma_{r^m}")

    m <= parameter(:π_star, 0.7000, (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   GammaAlt(0.62, 0.1), fixed=false, scaling = x -> 1 + x/100,
                   description="π_star: steady-state rate of inflation.",
                   tex_label="\\pi_*")

    m <= parameter(:Lmean, -45.9364, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed=false,
                   description="Lmean: Mean level of hours.", tex_label="\\bar{L}")

    m <= parameter(:e_y, 0.0, fixed = true,
                   description = "e_y: Measurement error on GDP", tex_label = "e_y")
    m <= parameter(:e_L, 0.0, fixed = true,
                   description = "e_L: Measurement error on hours worked", tex_label = "e_L")
    m <= parameter(:e_w, 0.0, fixed = true, description = "e_w: Measurement error on wages",
                   tex_label ="e_w")
    m <= parameter(:e_π, 0.0, fixed = true,
                   description = "e_π: Measurement error on GDP deflator", tex_label = "e_π")
    m <= parameter(:e_R, 0.0, fixed = true,
                   description = "e_R: Measurement error on nominal rate of interest",
                   tex_label = "e_R")
    m <= parameter(:e_c, 0.0, fixed = true,
                   description = "e_c: Measurement error on consumption", tex_label = "e_c")
    m <= parameter(:e_i, 0.0, fixed = true,
                   description = "e_i: Measurement error on investment", tex_label = "e_i")

    # Setting steady-state parameters
    nx = get_setting(m, :nx1_state) + get_setting(m, :nx2_state)
    ns = get_setting(m, :ns)

    # Steady state grids for functional/distributional variables and market-clearing discount rate
    m <= SteadyStateParameterGrid(:lstar, fill(NaN, nx*ns),
                                  description = "Steady-state expected discounted
                                  marginal utility of consumption", tex_label = "l_*")
    m <= SteadyStateParameterGrid(:cstar, fill(NaN, nx*ns),
                                  description = "Steady-state consumption", tex_label = "c_*")
    m <= SteadyStateParameterGrid(:μstar, fill(NaN, nx*ns), description = "Steady-state" *
                                  " cross-sectional density of cash on hand",
                                  tex_label = "\\mu_*")
    m <= SteadyStateParameter(:βstar, NaN, description = "Steady-state discount factor",
                                  tex_label = "\\beta_*")
end

"""
```
aggregate_steadystate!(m::HetDSGEGovDebt)
```
Steady state of aggregate scalar variables (Break out the "analytic" steady-state solution here instead of in steadystate!() since some of these parameters are required to construct the grids that are used in the functional/distributional steady state variable calculation (where init_grids! is called before steadystate! in model construction).
"""
function aggregate_steadystate!(m::HetDSGEGovDebt)
    # FOR POSTERITY
    #=m <= SteadyStateParameter(:rprior, 100*(m[:r] - m[:γ]),
                              description= "100*(r* - γ)",
                              tex_label= "100*(r^{HetDSGE}-\\gamma^{FRBNY})")
    =#
    m <= SteadyStateParameter(:Rkstar, m[:r] + m[:δ],
                              description = "Rental rate on capital", tex_label = "Rk_*")
    m <= SteadyStateParameter(:ωstar,
                              (m[:α]^(m[:α]/(1-m[:α])))*(1-m[:α])*m[:Rkstar]^(-m[:α]/(1-m[:α])),
                              description = "Real wage", tex_label = "\\omega_*")
    m <= SteadyStateParameter(:klstar, (m[:α]/(1-m[:α]))*(m[:ωstar]/m[:Rkstar])*exp(m[:γ]),
                              description = "Capital/Labor ratio", tex_label = "kl_*")
    m <= SteadyStateParameter(:kstar, m[:klstar]*m[:H], description = "Capital",
                              tex_label = "k_*")
    m <= SteadyStateParameter(:xstar, (1-(1-m[:δ])*exp(-m[:γ]))*m[:kstar],
                              description = "Investment", tex_label = "x_*")
    m <= SteadyStateParameter(:ystar, (exp(-m[:α]*m[:γ])*(m[:kstar]^m[:α])*m[:H]^(1-m[:α])),
                              description = "GDP", tex_label = "y_*")
    m <= SteadyStateParameter(:bg, m[:BoverY]*m[:ystar]*exp(m[:γ]), description = "Govt Debt",
                              tex_label = "bg")
    m <= SteadyStateParameter(:Tg, (exp(-m[:γ]) - 1/(1+m[:r]))*m[:bg] + (1-(1/m[:g]))*m[:ystar],
                              description = "Net lump sump taxes", tex_label = "Tg")
    m <= SteadyStateParameter(:Tstar, m[:Rkstar]*m[:kstar]*exp(-m[:γ]) - m[:xstar] - m[:Tg],
                              description = "Net transfer to households", tex_label = "T_*")
    return m
end

"""
```
init_grids!(m::HetDSGEGovDebt)
```
This function initializes the grids used for distribution/functional states. the xgrid/agrid is cash on hand and the sgrid is skill (right now we only have 2 sgrid points: low and high skill)
They are constructed with uniform weights on each gridpoint and the weights sum up to the kwarg passed into `scale' (which defaults to the difference between the max and min of the grid. e.g. if your max and min are 12.0 amnd 0.0, the weights wqill sum to 12.0).

f is the transition matrix to and from different skills levels
"""
function init_grids!(m::HetDSGEGovDebt)

    nx      = get_setting(m, :nx)
    ns      = get_setting(m, :ns)
    λ       = get_setting(m, :λ)

    grids = OrderedDict()

    # Skill grid
    f, sgrid, swts, sscale = persistent_skill_process(m[:sH_over_sL].value, m[:pLH].value,
                                                      m[:pHL].value, get_setting(m, :ns))
    # Markov transition matrix for skill
    grids[:sgrid] = Grid(sgrid, swts, sscale)
    grids[:fgrid] = f

    xgrid, xwts, xlo, xhi, xscale = cash_grid(sgrid, m[:ωstar].value, m[:H].value,
                                              m[:r].scaledvalue, m[:η].value, m[:γ].scaledvalue,
                                              m[:Tstar].value, m[:zlo].value, nx)

    grids[:xgrid] = Grid(uniform_quadrature(xscale), xlo, xhi, nx, scale = xscale)

    m <= Setting(:xlo, xlo)
    m <= Setting(:xhi, xhi)
    m <= Setting(:xscale, xscale)

    # Total grid vectorized across both dimensions
    grids[:sgrid_total] = kron(sgrid, ones(nx))
    grids[:xgrid_total] = kron(ones(ns), grids[:xgrid].points)
    grids[:weights_total] = kron(swts, grids[:xgrid].weights)

    m.grids = grids
end

function model_settings!(m::HetDSGEGovDebt)
    default_settings!(m)

    # Defaults
    # Data settings for released and conditional data. Default behavior is to set vintage
    # of data to today's date.
    vint = Dates.format(now(), DSGE_DATE_FORMAT)
    m <= Setting(:data_vintage, vint, true, "vint", "Data vintage")

    saveroot = normpath(joinpath(dirname(@__FILE__), "../../../../","save"))
    datapath = normpath(joinpath(dirname(@__FILE__), "../../../../","save","input_data"))

    m <= Setting(:saveroot, saveroot, "Root of data directory structure")
    m <= Setting(:dataroot, datapath, "Input data directory path")

    # Solution method
    m <= Setting(:solution_method, :klein, false, "", "This should always be klein")

    # Likelihood method
    m <= Setting(:use_chand_recursion, true)

    # Anticipated shocks
    m <= Setting(:n_anticipated_shocks, 0,
                 "Number of anticipated policy shocks")
    # m <= Setting(:n_anticipated_shocks_padding, 20,
                 # "Padding for anticipated policy shocks")

    # Whether one wishes to re-compute βstar
    m <= Setting(:use_last_βstar, false, "Flag to avoid recomputing steady-state.")
    m <= Setting(:estimate_only_non_steady_state_parameters, false, "Whether one wants
                 to only estimate parameters that do not affect the steady state. Note,
                 this setting is set automatically in subspec 11, and should not be set to
                 true otherwise.")

    # Number of states and jumps
    m <= Setting(:normalize_distr_variables, true, "Whether or not to perform the
                 normalization of the distributional states in the Klein solution step")
    m <= Setting(:n_predetermined_variables, 0, "Number of predetermined variables after
                 removing the densities. Calculated with the Jacobian normalization.
                 This set to 0 is just the default setting, since it will always be
                 overwritten once the Jacobian is calculated.")

    m <= Setting(:reduce_ell, true, false, "", "Whether to reduce the ell function (the inverse of consumption)")

    m <= Setting(:policy_damp, 0.5, "Dampening parameter for policy function iteration")
    m <= Setting(:policy_maxit, 500, true,
                 "maxit", "Maximum number of iterations in policy function")

    # Mollifier setting parameters
    m <= Setting(:In, 0.443993816237631, "Normalizing constant for the mollifier")

    # s: Skill Distribution/ "Units of effective labor" Grid Setup
    m <= Setting(:ns, 2, "Skill distribution grid points")
    m <= Setting(:λ, 2.0, "λ parameter in the Tauchen distribution calculation")

    # x: Cash on Hand Grid Setup
    m <= Setting(:nx1_state, 300, "Cash on hand distribution grid points (Lo)")
    m <= Setting(:nx2_state, 300, "Cash on hand distribution grid points (hi)")
    m <= Setting(:nx1_jump,  300, "Cash on hand distribution grid points (Lo)")
    m <= Setting(:nx2_jump,  300, "Cash on hand distribution grid points (hi)")
    m <= Setting(:nx,        300, "Cash on hand distribution grid points")

    m <= Setting(:binsize, 4, false, "", "Number of people to put in each bin when doing avg proxy reduction") # Setting binsize=1 gives us what we had before doing the binning reduction
    m <= Setting(:poor_man_reduc, true) #note that we're actually doing more than the "poor man reduction" now however this turns ont both poorman truncation and binning reduction

    # Set targets
    m <= Setting(:calibration_targets, [0.7, 0.23],
                 "Targets for: [var(log(annual income)), var(one year changes in " *
                 "log(annual income))]")
    m <= Setting(:calibration_targets_lb, [3.0, 1e-18], "Lower bounds on [sH/sL, zlo]")
    m <= Setting(:calibration_targets_ub, [9.0, 0.8-eps()], "Upper bounds on [sH/sL, zlo]")

    # Important settings for likelihood penalty
    m <= Setting(:use_likelihood_penalty, true)
    m <= Setting(:ψ_likelihood, 1.0,
                 "Multiplier on time series likelihood in penalty function")
    m <= Setting(:ψ_penalty, 1.0,
                 "Multiplier on moment penalties in penalty function")

    m <= Setting(:targets, [0.16, 0.10],
                 "Targets for: [MPC, proportion of individuals with 0 income]")
    m <= Setting(:target_vars, [:mpc, :pc0], "Symbols of variables we're targeting")
    m <= Setting(:target_σt, [0.2, 0.1], "Target \\sigma_t for MPC and pc0")

    m <= Setting(:calibrate_income_targets, true, "Calibrate for varlinc and vardlinc")
    m <= Setting(:steady_state_only, false, "Testing setting")
    m <= Setting(:auto_reject, false, "This flag is set when policy function doesn't converge")

    # Steady state constants
    m <= Setting(:ni, 10000) # This and the next line have to do with the random draws we draw for steadystate but are presaved in a file (the u's and z's)
    m <= Setting(:nz, 1000)
    m <= Setting(:fix_random_matrices, true, "Determines if use fixed matrices")
    m <= Setting(:us, load(get_setting(m, :ref_dir) * "/us_zs.jld2","us"))
    m <= Setting(:zs, load(get_setting(m, :ref_dir) * "/us_zs.jld2","zs"))

    # Misc
    m <= Setting(:trunc_distr, false) #old setting, we dont do this any more
    m <= Setting(:rescale_weights, true)
    m <= Setting(:mindens, 1e-8)

    # Function-valued variables include distributional variables
    m <= Setting(:n_function_valued_backward_looking_states, 1, "Number of function-valued" *
                 " backward looking state variables")
    m <= Setting(:n_backward_looking_distributional_vars, 1,
                 "Number of state variables that are distributional variables.")
    m <= Setting(:n_function_valued_jumps, 1, "Number of function-valued jump variables")
    m <= Setting(:n_jump_distributional_vars, 1,
                 "Number of jump variables that are distributional variables.")

    # Note, these settings assume normalization.
    # The n degrees of freedom removed depends on the distributions/dimensions
    # of heterogeneity that we have discretized over, in this case,
    # cash on hand and the skill distribution. In general the rule of
    # thumb is, remove one degree of freedom for the first endogenous distribution (cash on
    # hand), then one additional degree of freedom for each exogenous distribution (skill
    # distribution). Multiple endogenous distributions only permit removing a single degree
    # of freedom since it is then non-trivial to obtain the marginal distributions.
    m <= Setting(:n_degrees_of_freedom_removed_state, 2,
                 "Number of degrees of freedom from the distributional variables to remove.")
    m <= Setting(:n_degrees_of_freedom_removed_jump, 0,
                 "Number of degrees of freedom from the distributional variables to remove.")
end

function setup_indices!(m::HetDSGEGovDebt)
    nx = get_setting(m, :nx)
    ns = get_setting(m, :ns)
    endo = m.endogenous_states
    eqconds = m.equilibrium_conditions
    nxns_state = (get_setting(m, :nx1_state) + get_setting(m, :nx2_state)) #*ns
    nxns_jump = (get_setting(m, :nx1_jump) + get_setting(m, :nx2_jump)) #*ns

    # Endogenous function-valued states
    endo[:kf′_t]  = 1:nxns_state    #  combination of lagged ell and m functions, predicts m

    # Endogenous scalar-valued states
    endo[:k′_t]   = nxns_state+1:nxns_state+1   # capital (NOT steadystate object K)
    endo[:i′_t1]  = nxns_state+2:nxns_state+2   # lagged nominal interest rate
    endo[:y′_t1]  = nxns_state+3:nxns_state+3   # lagged gdp
    endo[:w′_t1]  = nxns_state+4:nxns_state+4   # lagged real wages
    endo[:I′_t1]  = nxns_state+5:nxns_state+5   # lagged investment (NOT nominal interest rate, i)
    endo[:bg′_t]  = nxns_state+6:nxns_state+6   # govt debt

    # Exogenous scalar-valued states:
    endo[:b′_t]   = nxns_state+7:nxns_state+7   # discount factor shock
    endo[:g′_t]   = nxns_state+8:nxns_state+8   # govt spending
    endo[:z′_t]   = nxns_state+9:nxns_state+9   # tfp growth
    endo[:μ′_t]   = nxns_state+10:nxns_state+10 # investment shock
    endo[:λ_w′_t] = nxns_state+11:nxns_state+11 # wage markup
    endo[:λ_f′_t] = nxns_state+12:nxns_state+12 # price markup
    endo[:rm′_t]  = nxns_state+13:nxns_state+13 # monetary policy shock

    # Function-valued jumps
    endo[:l′_t]   = nxns_state+14:nxns_state+nxns_jump+13 # ell function

    # Scalar-valued jumps
    endo[:C′_t]         = nxns_state+nxns_jump+14:nxns_state+nxns_jump+14 # real interest rate
    endo[:R′_t]         = nxns_state+nxns_jump+15:nxns_state+nxns_jump+15 # real interest rate
    endo[:i′_t]         = nxns_state+nxns_jump+16:nxns_state+nxns_jump+16 # nominal interest rate
    endo[:t′_t]         = nxns_state+nxns_jump+17:nxns_state+nxns_jump+17 # transfers + dividends
    endo[:w′_t]         = nxns_state+nxns_jump+18:nxns_state+nxns_jump+18 # real wage
    endo[:L′_t]         = nxns_state+nxns_jump+19:nxns_state+nxns_jump+19 # hours worked
    endo[:π′_t]         = nxns_state+nxns_jump+20:nxns_state+nxns_jump+20 # inflation
    endo[:π_w′_t]       = nxns_state+nxns_jump+21:nxns_state+nxns_jump+21 # nominal wage inflation
    endo[:margutil′_t]  = nxns_state+nxns_jump+22:nxns_state+nxns_jump+22 # avg marginal utility
    endo[:y′_t]         = nxns_state+nxns_jump+23:nxns_state+nxns_jump+23 # gdp
    endo[:I′_t]         = nxns_state+nxns_jump+24:nxns_state+nxns_jump+24 # investment
    endo[:mc′_t]        = nxns_state+nxns_jump+25:nxns_state+nxns_jump+25 # marginal cost - this is ζ in HetDSGEGovDebtₖd.pdf
    endo[:Q′_t]         = nxns_state+nxns_jump+26:nxns_state+nxns_jump+26 # Tobin's qfunction
    endo[:capreturn′_t] = nxns_state+nxns_jump+27:nxns_state+nxns_jump+27 # return on capital
    endo[:tg′_t]        = nxns_state+nxns_jump+28:nxns_state+nxns_jump+28

    # Function blocks which output a function
    eqconds[:eq_euler]                  = 1:nxns_state
    eqconds[:eq_kolmogorov_fwd]         = nxns_state+1:2*nxns_state

    # Function blocks which map functions to scalars
    eqconds[:eq_agg_consumption]        = 2*nxns_state+1:2*nxns_state+1
    eqconds[:eq_lambda]                 = 2*nxns_state+2:2*nxns_state+2

    # Scalar blocks involving endogenous variables
    eqconds[:eq_transfers]              = 2*nxns_state+3:2*nxns_state+3   # transfers
    eqconds[:eq_investment]             = 2*nxns_state+4:2*nxns_state+4   # investment
    eqconds[:eq_tobin_q]                = 2*nxns_state+5:2*nxns_state+5   # tobin's q
    eqconds[:eq_capital_accumulation]   = 2*nxns_state+6:2*nxns_state+6   # capital accumulation
    eqconds[:eq_wage_phillips]          = 2*nxns_state+7:2*nxns_state+7   # wage phillips curve
    eqconds[:eq_price_phillips]         = 2*nxns_state+8:2*nxns_state+8   # price phillips curve
    eqconds[:eq_marginal_cost]          = 2*nxns_state+9:2*nxns_state+9   # marginal cost
    eqconds[:eq_gdp]                    = 2*nxns_state+10:2*nxns_state+10 # gdp
    eqconds[:eq_optimal_kl]             = 2*nxns_state+11:2*nxns_state+11 # optimal K/L ratio
    eqconds[:eq_taylor]                 = 2*nxns_state+12:2*nxns_state+12 # taylor rule
    eqconds[:eq_fisher]                 = 2*nxns_state+13:2*nxns_state+13 # fisher eqn
    eqconds[:eq_nominal_wage_inflation] = 2*nxns_state+14:2*nxns_state+14 # nominal wage inflation
    eqconds[:eq_fiscal_rule]            = 2*nxns_state+15:2*nxns_state+15
    eqconds[:eq_g_budget_constraint]    = 2*nxns_state+16:2*nxns_state+16
    eqconds[:eq_resource_constraint]    = 2*nxns_state+17:2*nxns_state+17

    # Lagged Variables
    eqconds[:LR] = 2*nxns_state+17:2*nxns_state+17 # LR
    eqconds[:LI] = 2*nxns_state+18:2*nxns_state+18 # LI
    eqconds[:LY] = 2*nxns_state+19:2*nxns_state+19 # LY
    eqconds[:LW] = 2*nxns_state+20:2*nxns_state+20 # LW
    eqconds[:LX] = 2*nxns_state+21:2*nxns_state+21 # LX

    # Shocks
    eqconds[:eq_b]   = 2*nxns_state+22:2*nxns_state+22 # discount factor B
    eqconds[:eq_g]   = 2*nxns_state+23:2*nxns_state+23 # govt spending G
    eqconds[:eq_z]   = 2*nxns_state+24:2*nxns_state+24 # tfp growth Z
    eqconds[:eq_μ]   = 2*nxns_state+25:2*nxns_state+25 # investment MU
    eqconds[:eq_λ_w] = 2*nxns_state+26:2*nxns_state+26 # wage mkup LAMW
    eqconds[:eq_λ_f] = 2*nxns_state+27:2*nxns_state+27 # price mkup LAMF
    eqconds[:eq_rm]  = 2*nxns_state+28:2*nxns_state+28 # monetary policy MON

    # Total grid x*s
    m <= Setting(:n_state, (get_setting(m, :nx1_state) +get_setting(m, :nx2_state)),
                 "Total grid size, multiplying across grid dimensions.")
    m <= Setting(:n_jump, (get_setting(m, :nx1_jump) +get_setting(m, :nx2_jump)),
             "Total grid size, multiplying across grid dimensions.")
    m <= Setting(:nvars,     2*get_setting(m, :n_state) + 28, "num variables")
    m <= Setting(:nscalars,  28, "num eqs which output scalars")
    m <= Setting(:nyscalars, 15, "num scalar jumps")
    m <= Setting(:nxscalars, get_setting(m, :nscalars) - get_setting(m, :nyscalars),
                 "num scalar states")
end

function init_states_and_jumps!(m::AbstractModel, states::Vector{Symbol},
                                jumps::Vector{Symbol}, states_only::Bool = false)

    endo = m.endogenous_states #_unnormalized

    m <= Setting(:states, states)
    m <= Setting(:jumps,  jumps)

    state_indices = stack_indices(endo, states)
    jump_indices  = stack_indices(endo, jumps)

    m <= Setting(:state_indices, state_indices, "Indices of m.endogenous_states " *
                 "corresponding to backward-looking state variables")
    m <= Setting(:jump_indices, jump_indices, "Indices of m.endogenous_states " *
                 "corresponding to jump
                 variables")

    n_dof_removed_state = get_setting(m, :n_degrees_of_freedom_removed_state)
    n_dof_removed_jump  = get_setting(m, :n_degrees_of_freedom_removed_jump)

    ####################################################
    # Calculating the number of backward-looking states
    ####################################################
    n_backward_looking_distr_vars = get_setting(m, :n_backward_looking_distributional_vars)
    m <= Setting(:backward_looking_states_normalization_factor,
                 n_dof_removed_state*n_backward_looking_distr_vars,
                 "Number of dimensions removed from the backward looking state variables
                  for the normalization.")

    nxns_state = (get_setting(m, :nx1_state) + get_setting(m, :nx2_state))* get_setting(m, :ns)
    nxns_jump  = (get_setting(m, :nx1_jump)  + get_setting(m, :nx2_jump)) * get_setting(m, :ns)
    nx1_state  = get_setting(m, :nx1_state)
    nx2_state  = get_setting(m, :nx2_state)
    nx1_jump   = get_setting(m, :nx1_jump)
    nx2_jump   = get_setting(m, :nx2_jump)

    n_backward_looking_vars = length(get_setting(m, :state_indices))
    n_backward_looking_function_valued_vars = get_setting(m,
                                               :n_function_valued_backward_looking_states)
    n_backward_looking_scalar_vars = get_setting(m, :nxscalars)

    m <= Setting(:n_backward_looking_states, (nx1_state + nx2_state) *
                 n_backward_looking_distr_vars +
                 n_backward_looking_scalar_vars -
                 get_setting(m, :backward_looking_states_normalization_factor),
                 "Number of state variables, in the true sense (fully
                  backward looking) accounting for the discretization across the grid
                  of function-valued variables and the normalization of
                  distributional variables.")

    ##################################
    # Calculating the number of jumps
    ##################################
    n_jump_distr_vars = get_setting(m, :n_jump_distributional_vars)
    m <= Setting(:jumps_normalization_factor,
                 n_dof_removed_jump*n_jump_distr_vars, "number of dimensions removed from" *
                 " the jump variables for the normalization.")

    n_jump_vars = length(get_setting(m, :jump_indices))
    n_jump_function_valued_vars = get_setting(m, :n_function_valued_jumps)
    n_jump_scalar_vars = get_setting(m, :nyscalars) #n_jump_vars - (nx1+nx2)*n_jump_function_valued_vars

    m <= Setting(:n_jumps, (nx1_jump+nx2_jump)*n_jump_function_valued_vars +
                 n_jump_scalar_vars - get_setting(m, :jumps_normalization_factor),
                 "Number of jump variables (forward looking) accounting for
                  the discretization across the grid of function-valued variables and the
                  normalization of distributional variables.")

    m <= Setting(:n_model_states, get_setting(m, :n_backward_looking_states) +
                 get_setting(m, :n_jumps),
                 "Number of 'states' in the state space model. Because backward and forward
                 looking variables need to be explicitly tracked for the Klein solution
                 method, we have n_states and n_jumps")
end

"""
```
init_grids!(m::HetDSGEGovDebt)
```
This is a very important function. ANy time you reduce the grid, you need to restore to the full grid before evaluating steadystate/jacobian/likelihood again. Otherwise, you'll keep redsucing and after just a few times, you'll be left with NOTHING. So PSA: reset your grid when you're done with whatever you need to do with the smaller grid!
"""
function reset_grids!(m)
    m <= Setting(:nx1_state, 300)
    m <= Setting(:nx2_state, 300)
    m <= Setting(:nx1_jump,  300)
    m <= Setting(:nx2_jump,  300)

    setup_indices!(m)
    init_states_and_jumps!(m, get_setting(m, :states), get_setting(m, :jumps))
    init_grids!(m)

    # So that the indices of m.endogenous_states reflect the normalization
    normalize_model_state_indices!(m)

    endogenous_states_augmented = [:C_t1]
    for (i,k) in enumerate(endogenous_states_augmented)
        m.endogenous_states_augmented[k] = i +
            first(m.endogenous_states[get_setting(m, :jumps)[end]])
    end
    m <= Setting(:n_model_states_augmented, get_setting(m, :n_model_states) +
                 length(m.endogenous_states_augmented))
end


text:
module LyceumAI

# stdlib
using Random
using Random: default_rng
using Statistics, LinearAlgebra

# 3rd party
import LearnBase
using DocStringExtensions
import Flux: params, params!
import Flux.Optimise: update!
using UnsafeArrays,
      StaticArrays,
      Distances,
      FastClosures,
      IterativeSolvers,
      Parameters,
      EllipsisNotation,
      Flux,
      Zygote,

      Shapes,
      LyceumBase

# Lyceum
using Zygote: Params, Grads
using Base: promote_eltype, @propagate_inbounds, require_one_based_indexing
using MLDataPattern: eachbatch, nobs
using MacroTools: @forward
using SpecialArrays
using LyceumBase: @mustimplement

zerofn(args...) = 0
noop(args...) = nothing

export
       # Algorithms
      MPPI,
      NaturalPolicyGradient,

       # Models
      DiagGaussianPolicy,
      grad_loglikelihood!,
      loglikelihood,

       # Flux tools
      FluxTrainer,
      FluxTrainerIterator,
      orthonormal,
      multilayer_perceptron,

       # Miscellaneous
      ControllerIterator

const AbsVec = AbstractVector
const AbsMat = AbstractMatrix

infinitecb(x...) = false

abstract type AbstractModel{T} end
Base.eltype(m::AbstractModel{T}) where {T} = T

@mustimplement params(m::AbstractModel)
@mustimplement params!(ps, m::AbstractModel)
@mustimplement params!(m::AbstractModel, ps)

@mustimplement update!(m::AbstractModel, gs)

include("util/misc.jl")
include("vectorproducts.jl")

include("flux.jl")

include("models/policy.jl")
export DiagGaussianPolicy, grad_loglikelihood!, loglikelihood

abstract type AbstractTrainer end
# (o::AbstractTrainer{M})(m::M) where M
#@mustimplement fit!(m::AbstractModel, data); export fit!
# --

include("algorithms/cg.jl")
include("algorithms/MPPI.jl")
include("algorithms/naturalpolicygradient.jl")

include("controller.jl")

end # module


text:
module LazyFym

using Transducers

import Base: size
using ProgressMeter

export Fym,  # types
       # internalAPI
       ∫,  # integration
       Sim,  # simulation
       # tools
       sequentialise  # postprocess

export InputAffineQuadraticCostEnv

include("types.jl")
include("internalAPI.jl")
include("integration.jl")
include("simulation.jl")
include("tools.jl")
include("postprocess.jl")

include("fymenvs.jl")

end  # module