text:
variable_name(op::Operation, ::Nothing) = mangledvar(op)
variable_name(op::Operation, suffix) = Symbol(mangledvar(op), suffix, :_)

function append_inds!(ret, indices, deps)
    start = (first(indices) === Symbol("##DISCONTIGUOUSSUBARRAY##")) + 1# && return append_inds!(ret, @view(indices[2:end]), deps)
    for ind ∈ @view(indices[start:end])
        if ind isa Int
            push!(ret.args, ind - 1)
        elseif ind ∈ deps
            push!(ret.args, ind)
        else
            push!(ret.args, Expr(:call, :-, ind, 1))
        end
    end
    ret
end

function mem_offset(op::Operation)
    @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
    append_inds!(Expr(:tuple), op.ref.ref, op.dependencies)
end
function mem_offset(op::Operation, incr::Int, unrolled::Symbol)
    @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
    ret = Expr(:tuple)
    indices = op.ref.ref
    deps = op.dependencies
    if incr == 0
        append_inds!(ret, indices, deps)
    else
        for n ∈ 1:length(indices)
            ind = indices[n]
            n == 1 && ind === Symbol("##DISCONTIGUOUSSUBARRAY##") && continue
            if ind isa Int
                push!(ret.args, ind - 1)
            elseif ind === unrolled
                push!(ret.args, Expr(:call, :+, ind, incr))
            elseif ind ∈ deps
                push!(ret.args, ind)
            else
                push!(ret.args, Expr(:call, :-, ind, 1))
            end
        end
    end
    ret
end
function mem_offset(op::Operation, mul::Symbol, incr::Int, unrolled::Symbol)
    @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
    ret = Expr(:tuple)
    indices = op.ref.ref
    deps = op.dependencies
    if incr == 0
        append_inds!(ret, indices, deps)
    else
        for n ∈ 1:length(indices)
            ind = indices[n]
            n == 1 && ind === Symbol("##DISCONTIGUOUSSUBARRAY##") && continue
            if ind isa Int
                push!(ret.args, ind - 1)
            elseif ind === unrolled
                push!(ret.args, Expr(:call, :+, ind, Expr(:call, lv(:valmul), mul, incr)))
            elseif ind ∈ deps
                push!(ret.args, ind)
            else
                push!(ret.args, Expr(:call, :-, ind, 1))
            end
        end
    end
    ret
end

# function add_expr(q, incr)
#     if q.head === :call && q.args[2] === :+
#         qc = copy(q)
#         push!(qc.args, incr)
#         qc
#     else
#         Expr(:call, :+, q, incr)
#     end
# end
function varassignname(var::Symbol, u::Int, isunrolled::Bool)
    isunrolled ? Symbol(var, u) : var
end
# name_mo only gets called when vectorized
function name_mo(var::Symbol, op::Operation, u::Int, W::Symbol, vecnotunrolled::Bool, unrolled::Symbol)
    if u < 0 # sentinel value meaning not unrolled
        name = var
        mo = mem_offset(op)
    else
        name = Symbol(var, u)
        mo = vecnotunrolled ? mem_offset(op, u, unrolled) : mem_offset(op, W, u, unrolled)
    end
    name, mo
end
function pushvectorload!(q::Expr, op::Operation, var::Symbol, u::Int, U::Int, W::Symbol, mask, unrolled::Symbol, vecnotunrolled::Bool)
    ptr = refname(op)
    name, mo = name_mo(var, op, u, W, vecnotunrolled, unrolled)
    instrcall = Expr(:call, lv(:vload), W, ptr, mo)
    if mask !== nothing && (vecnotunrolled || u == U - 1)
        push!(instrcall.args, mask)
    end
    push!(q.args, Expr(:(=), name, instrcall))
end
function pushvectorgather!(
    q::Expr, op::Operation, var::Symbol, u::Int, U::Int, W::Symbol,
    mask, unrolled::Symbol, ustride::Symbol, vecnotunrolled::Bool
)
    ptr = refname(op)
    name, mo = name_mo(var, op, u, W, vecnotunrolled, unrolled)#, vecnotunrolled)
    instrcall = Expr(:call, lv(:gather), ptr, mo, ustride)
    if mask !== nothing && (vecnotunrolled || u == U - 1)
        push!(instrcall.args, mask)
    end
    push!(q.args, Expr(:(=), name, instrcall))
end
function lower_load_scalar!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    loopdeps = loopdependencies(op)
    @assert vectorized ∉ loopdeps
    var = variable_name(op, suffix)
    ptr = refname(op)
    isunrolled = unrolled ∈ loopdeps
    U = isunrolled ? U : 1
    for u ∈ 0:U-1
        varname = varassignname(var, u, isunrolled)
        push!(q.args, Expr(:(=), varname, Expr(:call, lv(:load),  ptr, mem_offset(op, u, unrolled))))
    end
    nothing
end
function lower_load_vectorized!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    loopdeps = loopdependencies(op)
    @assert vectorized ∈ loopdeps
    if unrolled ∈ loopdeps
        umin = 0
        U = U
    else
        umin = -1
        U = 0
    end
    # Urange = unrolled ∈ loopdeps ? 0:U-1 : 0
    var = variable_name(op, suffix)
    vecnotunrolled = vectorized !== unrolled
    if first(op.ref.ref) === vectorized # vload
        for u ∈ umin:U-1
            pushvectorload!(q, op, var, u, U, W, mask, unrolled, vecnotunrolled)
        end
    else
        sn = findfirst(x -> x === vectorized, op.ref.ref)::Int
        ustrides = Expr(:call, lv(:vmul), Expr(:call, :stride, refname(op), sn), Expr(:call, lv(:vrange), W))
        ustride = gensym(:ustride)
        push!(q.args, Expr(:(=), ustride, ustrides))
        for u ∈ umin:U-1
            pushvectorgather!(q, op, var, u, U, W, mask, unrolled, ustride, vecnotunrolled)
        end
    end
    nothing
end

# TODO: this code should be rewritten to be more "orthogonal", so that we're just combining separate pieces.
# Using sentinel values (eg, T = -1 for non tiling) in part to avoid recompilation.
function lower_load!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    # @show op.instruction
    # @show unrolled, loopdependencies(op)
    if vectorized ∈ loopdependencies(op)
        lower_load_vectorized!(q, op, vectorized, W, unrolled, U, suffix, mask)
    else
        lower_load_scalar!(q, op, vectorized, W, unrolled, U, suffix, mask)
    end
end
function reduce_range!(q::Expr, toreduct::Symbol, instr::Instruction, Uh::Int, Uh2::Int)
    for u ∈ 0:Uh-1
        tru = Symbol(toreduct, u)
        push!(q.args, Expr(:(=), tru, Expr(instr, tru, Symbol(toreduct, u + Uh))))
    end
    for u ∈ 2Uh:Uh2-1
        tru = Symbol(toreduct, u + 1 - 2Uh)
        push!(q.args, Expr(:(=), tru, Expr(instr, tru, Symbol(toreduct, u))))
    end
end
function reduce_range!(q::Expr, ls::LoopSet, Ulow::Int, Uhigh::Int)
    for or ∈ ls.outer_reductions
        op = ls.operations[or]
        var = mangledvar(op)
        temp = gensym(var)
        instr = op.instruction
        instr = get(REDUCTION_TRANSLATION, instr, instr)
        reduce_range!(q, var, instr, Ulow, Uhigh)
    end
end

function reduce_expr!(q::Expr, toreduct::Symbol, instr::Instruction, U::Int)
    U == 1 && return nothing
    instr = get(REDUCTION_TRANSLATION, instr, instr)
    Uh2 = U
    iter = 0
    while true # combine vectors
        Uh = Uh2 >> 1
        reduce_range!(q, toreduct, instr, Uh, Uh2)
        Uh == 1 && break
        # @show Uh
        Uh2 = Uh
        iter += 1; iter > 4 && throw("Oops! This seems to be excessive unrolling.")
    end
    # reduce last vector
    # push!(q.args, Expr(:(=), assignto, Expr(:call, reductfunc, Symbol(toreduct,:_0))))
    nothing
end

pvariable_name(op::Operation, ::Nothing) = mangledvar(first(parents(op)))
pvariable_name(op::Operation, suffix) = Symbol(pvariable_name(op, nothing), suffix, :_)

function reduce_unroll!(q, op, U, unrolled)
    loopdeps = loopdependencies(op)
    isunrolled = unrolled ∈ loopdeps
    if unrolled ∉ reduceddependencies(op)
        U = isunrolled ? U : 1
        return U, isunrolled
    end
    unrolled ∈ reduceddependencies(op) || return U
    var = mangledvar(op)
    instr = first(parents(op)).instruction
    reduce_expr!(q, var, instr, U) # assigns reduction to storevar
    1, isunrolled
end
function lower_store_reduction!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned}, isunrolled::Bool
)
    var = pvariable_name(op, suffix)
    ptr = refname(op)
    # need to find out reduction type
    instr = first(parents(op)).instruction
    reduct_instruct = CORRESPONDING_REDUCTION[instr]
    for u ∈ 0:U-1
        reducedname = varassignname(var, u, isunrolled)
        storevar = Expr(reduct_instruct, reducedname)
        push!(q.args, Expr(:call, lv(:store!), ptr, storevar, mem_offset(op, u, unrolled))) # store storevar
    end
    nothing
end
function lower_store_scalar!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned}, isunrolled::Bool
)
    var = pvariable_name(op, suffix)
    ptr = refname(op)
    for u ∈ 0:U-1
        varname = varassignname(var, u, isunrolled)
        push!(q.args, Expr(:call, lv(:store!), ptr, varname, mem_offset(op, u, unrolled)))
    end
    nothing
end
function lower_store_vectorized!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned}, isunrolled::Bool
)
    loopdeps = loopdependencies(op)
    @assert unrolled ∈ loopdeps
    var = pvariable_name(op, suffix)
    if isunrolled
        umin = 0
        U = U
    else
        umin = -1
        U = 0
    end
    ptr = refname(op)
    vecnotunrolled = vectorized !== unrolled
    if first(loopdependencies(op)) === vectorized # vstore!
        for u ∈ 0:U-1
            name, mo = name_mo(var, op, u, W, vecnotunrolled, unrolled)
            instrcall = Expr(:call,lv(:vstore!), ptr, name, mo)
            # @show mask, vecnotunrolled, u, U
            if mask !== nothing && (vecnotunrolled || u == U - 1)
                push!(instrcall.args, mask)
            end
            push!(q.args, instrcall)
        end
    else
        sn = findfirst(x -> x === unrolled, loopdependencies(op))::Int
        ustrides = Expr(:call, lv(:vmul), Expr(:call, :stride, ptr, sn), Expr(:call, lv(:vrange), W))
        for u ∈ 0:U-1
            name, mo = name_mo(var, op, u, W, vecnotunrolled, unrolled)
            instrcall = Expr(:call, lv(:scatter!), ptr, mo, ustrides, name)
            if mask !== nothing && (vecnotunrolled || u == U - 1)
                push!(instrcall.args, mask)
            end
            push!(q.args, instrcall)
        end
    end
    nothing
end
function lower_store!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    U, isunrolled = reduce_unroll!(q, op, U, unrolled)
    if vectorized ∈ reduceddependencies(op)
        lower_store_reduction!(q, op, vectorized, W, unrolled, U, suffix, mask, isunrolled)
    elseif vectorized ∈ loopdependencies(op)
        lower_store_vectorized!(q, op, vectorized, W, unrolled, U, suffix, mask, isunrolled)
    else
        lower_store_scalar!(q, op, vectorized, W, unrolled, U, suffix, mask, isunrolled)
    end
end
# A compute op needs to know the unrolling and tiling status of each of its parents.
#
function lower_compute_scalar!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    lower_compute!(q, op, vectorized, W, unrolled, tiled, U, suffix, mask, false)
end
function lower_compute_unrolled!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    lower_compute!(q, op, vectorized, W, unrolled, tiled, U, suffix, mask, true)
end
struct FalseCollection end
Base.getindex(::FalseCollection, i...) = false
function lower_compute!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing,
    opunrolled = unrolled ∈ loopdependencies(op)
)

    var = op.variable
    mvar = mangledvar(op)
    parents_op = parents(op)
    nparents = length(parents_op)
    if opunrolled
        parentsunrolled = Vector{Bool}(undef, nparents)
        for (p,opp) ∈ enumerate(parents_op)
            # if op is an inner reduction, one of its parents will be the initialization of op
            # They will share the same `variable` field. The initialization may not have
            # unrolled in its loop dependencies, but (if opunrolled) op itself is, so we return true
            parentsunrolled[p] = var === opp.variable ? true : (unrolled ∈ loopdependencies(opp))
        end
    else # maybe skip allocating this?
        parentsunrolled = fill(false, nparents)
    end
    parentstiled = if suffix === nothing
        optiled = false
        tiledouterreduction = -1
        FalseCollection()
    else
        tiledouterreduction = isouterreduction(op)
        suffix_ = Symbol(suffix, :_)
        if tiledouterreduction == -1
            mvar = Symbol(mvar, suffix_)
        end
        optiled = true
        [tiled ∈ loopdependencies(opp) for opp ∈ parents_op]
    end
    instr = op.instruction
    # cache unroll and tiling check of parents
    # not broadcasted, because we use frequent checks of individual bools
    # making BitArrays inefficient.
    # @show instr parentsunrolled
    # parentsyms = [opp.variable for opp ∈ parents(op)]
    Uiter = opunrolled ? U - 1 : 0
    maskreduct = mask !== nothing && isreduction(op) && any(opp -> opp.variable === var, parents_op)
    # if a parent is not unrolled, the compiler should handle broadcasting CSE.
    # because unrolled/tiled parents result in an unrolled/tiled dependendency,
    # we handle both the tiled and untiled case here.
    # bajillion branches that go the same way on each iteration
    # but smaller function is probably worthwhile. Compiler could theoreically split anyway
    # but I suspect that the branches are so cheap compared to the cost of everything else going on
    # that smaller size is more advantageous.
    modsuffix = 0
    for u ∈ 0:Uiter
        instrcall = Expr(instr) # Expr(:call, instr)
        varsym = if tiledouterreduction > 0 # then suffix !== nothing
            modsuffix = ((u + suffix*U) & 3)
            Symbol(mvar, modsuffix)
        elseif opunrolled
            Symbol(mvar, u)
        else
            mvar
        end
        for n ∈ 1:nparents
            parent = mangledvar(parents_op[n])
            if n == tiledouterreduction
                parent = Symbol(parent, modsuffix)
            else
                if parentstiled[n]
                    parent = Symbol(parent, suffix_)
                end
                if parentsunrolled[n]
                    parent = Symbol(parent, u)
                end
            end
            push!(instrcall.args, parent)
        end
        if maskreduct && u == Uiter # only mask last
            push!(q.args, Expr(:(=), varsym, Expr(:call, lv(:vifelse), mask, instrcall, varsym)))
        else
            push!(q.args, Expr(:(=), varsym, instrcall))
        end
    end
end
function lower!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    if isload(op)
        lower_load!(q, op, vectorized, W, unrolled, U, suffix, mask)
    elseif isstore(op)
        lower_store!(q, op, vectorized, W, unrolled, tiled, U, suffix, mask)
    elseif iscompute(op)
        lower_compute!(q, op, vectorized, W, unrolled, U, suffix, mask)
    else
        lower_constant!(q, op, vectorized, W, unrolled, U, suffix, mask)
    end
end
function lower_constant!(
    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Any = nothing
)
    instruction = op.instruction
    mvar = variable_name(op, suffix)

    constsym = instruction.instr
    if vectorized ∈ loopdependencies(op) || vectorized ∈ reduceddependencies(op)
        call = Expr(:call, lv(:vbroadcast), W, constsym)
        for u ∈ 0:U-1
            push!(q.args, Expr(:(=), Symbol(mvar, u), call))
        end
    else
        for u ∈ 0:U-1
            push!(q.args, Expr(:(=), Symbol(mvar, u), constsym))
        end
    end
    nothing
end
function lower!(
    q::Expr, ops::AbstractVector{Operation}, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    foreach(op -> lower!(q, op, vectorized, W, unrolled, tiled, U, suffix, mask), ops)
end
function lower_load!(
    q::Expr, ops::AbstractVector{Operation}, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    foreach(op -> lower_load!(q, op, vectorized, W, unrolled, U, suffix, mask), ops)
end
function lower_compute!(
    q::Expr, ops::AbstractVector{Operation}, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    foreach(op -> lower_compute!(q, op, vectorized, W, unrolled, tiled::Symbol, U, suffix, mask), ops)
end
function lower_store!(
    q::Expr, ops::AbstractVector{Operation}, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    foreach(op -> lower_store!(q, op, vectorized, W, unrolled, U, suffix, mask), ops)
end
function lower_constant!(
    q::Expr, ops::AbstractVector{Operation}, vectorized::Symbol, W::Symbol, unrolled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    foreach(op -> lower_constan!(q, op, vectorized, W, unrolled, U, suffix, mask), ops)
end

function lower!(
    q::Expr, ops::AbstractVector{<:AbstractVector{Operation}}, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
)
    @assert length(ops) == 4
    @inbounds begin
        foreach(op -> lower_constant!(q, op, vectorized, W, unrolled, U, suffix, mask), ops[1])
        foreach(op -> lower_load!(q, op, vectorized, W, unrolled, U, suffix, mask), ops[2])
        foreach(op -> lower_compute!(q, op, vectorized, W, unrolled, tiled, U, suffix, mask), ops[3])
        foreach(op -> lower_store!(q, op, vectorized, W, unrolled, U, suffix, mask), ops[4])
    end
end

tiledsym(s::Symbol) = Symbol("##outer##", s, "##outer##")
function lower_nest(
    ls::LoopSet, n::Int, vectorized::Symbol, U::Int, T::Int, loopq_old::Union{Expr,Nothing},
    loopstart::Union{Int,Symbol}, W::Symbol,
    mask::Union{Nothing,Symbol,Unsigned} = nothing, exprtype::Symbol = :while
)
    ops = oporder(ls)
    order = names(ls)
    istiled = T != -1
    loopsym = order[n]
    nloops = num_loops(ls)
    outer_reduce = length(ls.outer_reductions) > 0
    nisvectorized = loopsym === vectorized
    nisunrolled = false
    nistiled = false
    # @show n, mask
    if istiled
        if n == nloops
            loopsym = tiledsym(loopsym)
            nistiled = true
        elseif n == nloops - 1
            nisunrolled = true
        end
        unrolled = order[end-1]
    else
        unrolled = last(order)
        nisunrolled = n == nloops
    end
    # @show unrolled, order
    blockq = Expr(:block)
    n == 1 || push!(blockq.args, Expr(:(=), order[n-1], loopstart))
    loopq = if exprtype === :block
        blockq
    else
        @assert exprtype === :while || exprtype === :if "Expression type $exprtype not recognized."
        # if we have a mask, we use 1 as the incremement to pass to looprange, so that we get a iter < maxiter condition.
        lr = if mask !== nothing # on mask, we do evaluate the rest of the loop, so decrement only 1
            looprange(ls, loopsym, 1)
        elseif nisvectorized
            vec_looprange(ls, loopsym, nisunrolled, W, U) # may not be tiled
        else
            looprange(ls, loopsym, nisunrolled ? U : (nistiled ? T : 1))
        end
        Expr(exprtype, lr, blockq)
    end
    for prepost ∈ 1:2
        # !U && !T
        lower!(blockq, @view(ops[:,1,1,prepost,n]), vectorized, W, unrolled, last(order), U, nothing, mask)
        # for u ∈ 0:U-1     #  U && !T
        lower!(blockq, @view(ops[:,2,1,prepost,n]), vectorized, W, unrolled, last(order), U, nothing, mask)
        # end
        if sum(length, @view(ops[:,:,2,prepost,n])) > 0
            for t ∈ 0:T-1
                if t == 0
                    push!(blockq.args, Expr(:(=), last(order), tiledsym(last(order))))
                else
                    push!(blockq.args, Expr(:+=, last(order), 1))
                end
                # !U &&  T
                lower!(blockq, @view(ops[:,1,2,prepost,n]), vectorized, W, unrolled, last(order), U, t, mask)
                # for u ∈ 0:U-1 #  U &&  T
                lower!(blockq, @view(ops[:,2,2,prepost,n]), vectorized, W, unrolled, last(order), U, t, mask)
                # end
            end
        end
        if loopq_old !== nothing && n > 1 && prepost == 1
            push!(blockq.args, loopq_old)
        end
    end
    if nisvectorized
        if nisunrolled
            push!(blockq.args, Expr(:+=, loopsym, Expr(:call, lv(:valmul), W, U)))
        else
            push!(blockq.args, Expr(:(=), loopsym, Expr(:call, lv(:valadd), W, loopsym)))
        end
    else
        push!(blockq.args, Expr(:+=, loopsym, nisunrolled ? U : (nistiled ? T : 1)))
    end
    loopq
end
# Calculates nested loop set,
# if tiled, it will not lower the tiled iteration.
function add_vec_rem_iter(
    ls::LoopSet, n::Int, vectorized::Symbol, U::Int, T::Int, loopqold, loopstart, W, Uexprtype, order
)
    loopq = lower_nest(ls, n, vectorized, U, T, loopqold, loopstart, W, nothing, Uexprtype)
    if order[n] === vectorized
        loopq = Expr(
            :block, loopq,
            lower_nest(ls, n, vectorized, U, T, loopqold, loopstart, W, Symbol("##mask##"), :if)
        )
    end
    loopq
end
function lower_set(ls::LoopSet, vectorized::Symbol, U::Int, T::Int, W::Symbol, ::Nothing, Uexprtype::Symbol)
    # @show U, T, W
    loopstart = 0
    istiled = T != -1
    order = names(ls)
    unrolled = order[end - istiled]
    unrolled === vectorized && return lower_set_unrolled_is_vectorized(ls, vectorized, U, T, W, nothing, Uexprtype)
    ns = 1
    nl = num_loops(ls) - istiled
    exprtype = nl == ns ? Uexprtype : :while
    nomaskq = lower_nest(ls, 1, vectorized, U, T, nothing, loopstart, W, nothing, exprtype)
    maskq = lower_nest(ls, 1, vectorized, U, T, nothing, loopstart, W, Symbol("##mask##"), order[ns] === vectorized ? :if : exprtype)
    while order[ns] !== vectorized
        ns += 1
        exprtype = nl == ns ? Uexprtype : :while
        nomaskq = lower_nest(ls, ns, vectorized, U, T, nomaskq, loopstart, W, nothing, exprtype)
        maskq = lower_nest(ls, ns, vectorized, U, T, maskq, loopstart, W, Symbol("##mask##"), order[ns] === vectorized ? :if : exprtype)
    end
    ns += 1
    loopq = Expr(:block, nomaskq, maskq)
    for n ∈ ns:nl
        exprtype = n == nl ? Uexprtype : :while
        loopq = lower_nest(ls, n, vectorized, U, T, loopq, loopstart, W, nothing, exprtype)
        # loopq = add_vec_rem_iter( ls, n, vectorized, U, T, loopq, 0, W, exprtype, order )
    end
    loopq
end
function lower_set(ls::LoopSet, vectorized::Symbol, U::Int, T::Int, W::Symbol, mask::Symbol, Uexprtype::Symbol)
    lower_set_unrolled_is_vectorized(ls, vectorized, U, T, W, mask, Uexprtype)
end
function lower_set_unrolled_is_vectorized(ls::LoopSet, vectorized::Symbol, U::Int, T::Int, W::Symbol, mask, Uexprtype::Symbol)
    nl = num_loops(ls) - (T != -1)
    loopq = lower_nest(ls, 1, vectorized, U, T, nothing, 0, W, mask, nl == 1 ? Uexprtype : :while)
    for n ∈ 2:nl
        exprtype = n == nl ? Uexprtype : :while
        loopq = lower_nest(ls, n, vectorized, U, T, loopq, 0, W, mask, exprtype)
    end
    loopq
end
function initialize_outer_reductions!(
    q::Expr, op::Operation, Umin::Int, Umax::Int, W::Symbol, typeT::Symbol, unrolled::Symbol, suffix::Union{Symbol,Nothing} = nothing
)
    # T = op.elementbytes == 8 ? :Float64 : :Float32
    z = Expr(:call, REDUCTION_ZERO[op.instruction], typeT)
    if unrolled ∈ reduceddependencies(op)
        z = Expr(:call, lv(:vbroadcast), W, z)
    end
    mvar = variable_name(op, suffix)
    for u ∈ Umin:Umax-1
        push!(q.args, Expr(:(=), Symbol(mvar, u), z))
    end
    nothing
end
function initialize_outer_reductions!(q::Expr, ls::LoopSet, Umin::Int, Umax::Int, W::Symbol, typeT::Symbol, unrolled::Symbol, suffix::Union{Symbol,Nothing} = nothing)
    foreach(or -> initialize_outer_reductions!(q, ls.operations[or], Umin, Umax, W, typeT, unrolled, suffix), ls.outer_reductions)
end
function initialize_outer_reductions!(ls::LoopSet, Umin::Int, Umax::Int, W::Symbol, typeT::Symbol, unrolled::Symbol, suffix::Union{Symbol,Nothing} = nothing)
    initialize_outer_reductions!(ls.preamble, ls, Umin, Umax, W, typeT, unrolled, suffix)
end
function add_upper_outer_reductions(ls::LoopSet, loopq::Expr, Ulow::Int, Uhigh::Int, W::Symbol, typeT::Symbol, unrolledloop::Loop)
    unrolled = unrolledloop.itersymbol
    ifq = Expr(:block)
    initialize_outer_reductions!(ifq, ls, Ulow, Uhigh, W, typeT, unrolled)
    push!(ifq.args, loopq)
    reduce_range!(ifq, ls, Ulow, Uhigh)
    comparison = if unrolledloop.hintexact
        Expr(:call, :<, unrolledloop.rangehint, Expr(:call, lv(:valmul), W, Uhigh))
    else
        Expr(:call, :<, unrolledloop.rangesym, Expr(:call, lv(:valmul), W, Uhigh))
    end
    ncomparison = Expr(:call, :!, comparison)
    Expr(:if, ncomparison, ifq)
end
function reduce_expr!(q::Expr, ls::LoopSet, U::Int)
    for or ∈ ls.outer_reductions
        op = ls.operations[or]
        var = op.variable
        mvar = mangledvar(op)
        instr = op.instruction
        reduce_expr!(q, mvar, instr, U)
        push!(q.args, Expr(:(=), var, Expr(:call, REDUCTION_SCALAR_COMBINE[instr], var, Symbol(mvar, 0))))
    end
end
function gc_preserve(ls::LoopSet, q::Expr)
    length(ls.includedarrays) == 0 && return q # is this even possible?
    gcp = Expr(:macrocall, Expr(:(.), :GC, QuoteNode(Symbol("@preserve"))), LineNumberNode(@__LINE__, @__FILE__))
    for (array,_) ∈ ls.includedarrays
        push!(gcp.args, array)
    end
    push!(q.args, nothing)
    push!(gcp.args, q)
    Expr(:block, gcp)
end
function determine_eltype(ls::LoopSet)
    # length(ls.includedarrays) == 0 && return REGISTER_SIZE >>> 3
    if length(ls.includedarrays) == 1
        return Expr(:call, :eltype, first(first(ls.includedarrays)))
    end
    promote_q = Expr(:call, :promote_type)
    for (array,_) ∈ ls.includedarrays
        push!(promote_q.args, Expr(:call, :eltype, array))
    end
    promote_q
end
function determine_width(ls::LoopSet, typeT::Symbol, unrolled::Symbol)
    unrolledloop = ls.loops[unrolled]
    if unrolledloop.hintexact
        Expr(:call, lv(:pick_vector_width_val), Expr(:call, Expr(:curly, :Val, unrolledloop.rangehint)), typeT)
    else
        Expr(:call, lv(:pick_vector_width_val), typeT)
    end
end
function lower_unrolled!(
    q::Expr, ls::LoopSet, vectorized::Symbol, U::Int, T::Int, W::Symbol, typeT::Symbol, unrolledloop::Loop
)
    # q = if unrolledloop.hintexact
        # Ureduct = U
        # lower_unrolled_static!( q, ls, U, T, W, typeT, unrolledloop )
    # else
    Ureduct = min(U,4)
    q = lower_unrolled_dynamic!( q, ls, vectorized, U, T, W, typeT, unrolledloop )
    # end
    if T == -1
        q = gc_preserve(ls, q)
    end
    manageouterreductions = T == -1 && length(ls.outer_reductions) > 0
    manageouterreductions && reduce_expr!(q, ls, Ureduct)
    q
end
#=
function lower_unrolled_static!(
    q::Expr, ls::LoopSet, U::Int, T::Int, W::Int, unrolledloop::Loop
)
    unrolled = unrolledloop.itersymbol
    unrolled_numiter = unrolledloop.rangehint
    Urem = unrolled_numiter
    # if static, we use Urem to indicate remainder.
    if unrolled_numiter ≥ 2U*W # we need at least 2 iterations
        Uexprtype = :while
    elseif unrolled_numiter ≥ U*W # complete unroll
        Uexprtype = :block
    else# we have only a single block
        Uexprtype = :skip
    end
    manageouterreductions = T == -1 && length(ls.outer_reductions) > 0
    if manageouterreductions
        # Umax = (!static_unroll && U > 2) ? U >> 1 : U
        Ureduct = U
        initialize_outer_reductions!(q, ls, 0, Ureduct, W, last(names(ls)))
    else
        Ureduct = -1
    end
    Wt = W
    Ut = U
    while true
        if Uexprtype !== :skip
            loopq = if Urem == unrolled_numiter || Urem == -1 # static, no mask
                lower_set(ls, Ut, T, Wt, nothing, Uexprtype)
            else # static, need mask
                lower_set(ls, Ut, T, Wt, VectorizationBase.unstable_mask(Wt, Urem), Uexprtype)
            end
            push!(q.args, loopq)
        end
        if Urem == unrolled_numiter
            remUiter = unrolled_numiter % (U*W)
            if remUiter == 0 # no remainder, we're done with the unroll
                break
            else # remainder, requires another iteration; what size?
                Ut, Urem = divrem(remUiter, W)
                if Urem == 0 # Ut iters of W
                    Urem = -1
                else
                    if Ut == 0 # if Urem == unrolled_numiter, we may already be done, othererwise, we may be able to shrink Wt
                        if Urem == unrolled_numiter && Uexprtype !== :skip
                            break
                        else
                            Wt = VectorizationBase.nextpow2(Urem)
                            if Wt == Urem # no mask needed
                                Urem = -1
                            end
                        end
                    end
                    # because initial Urem > 0 (it either still is, or we shrunk Wt and made it a complete iter)
                    # we must increment Ut (to perform masked or shrunk complete iter)
                    Ut += 1
                end
                Uexprtype = :block
            end
        else
            break
        end
    end
    q
end
=#
function lower_unrolled_dynamic!(
    q::Expr, ls::LoopSet, vectorized::Symbol, U::Int, T::Int, W::Symbol, typeT::Symbol, unrolledloop::Loop
)
    unrolled = unrolledloop.itersymbol
    unrolled_numitersym = unrolledloop.rangesym
    Urem = 0
    Uexprtype = :while
    manageouterreductions = T == -1 && length(ls.outer_reductions) > 0
    if manageouterreductions
        # Umax = (!static_unroll && U > 2) ? U >> 1 : U
        Ureduct = U > 6 ? 4 : U
        initialize_outer_reductions!(q, ls, 0, Ureduct, W, typeT, last(names(ls)))
    else
        Ureduct = -1
    end
    Ut = U
    vecisunrolled = unrolled === vectorized
    local remblock::Expr
    firstiter = true
    while true
        if firstiter # first iter
            loopq = lower_set(ls, vectorized, Ut, T, W, nothing, Uexprtype)
            if T == -1 && manageouterreductions && U > 4
                loopq = add_upper_outer_reductions(ls, loopq, Ureduct, U, W, typeT, unrolledloop)
            end
            push!(q.args, loopq)
        elseif U == 1 #
            if vecisunrolled
                push!(remblock.args, lower_set(ls, vectorized, Ut, T, W, Symbol("##mask##"), :block))
            else
                push!(remblock.args, lower_set(ls, vectorized, Ut, T, W, nothing, :block))
            end
        else
            remblocknew = if vecisunrolled
                itercount = if unrolledloop.hintexact
                    Expr(:call, :-, unrolledloop.rangehint, Expr(:call, lv(:valmuladd), W, Ut, 1))
                else
                    Expr(:call, :-, unrolled_numitersym, Expr(:call, lv(:valmuladd), W, Ut, 1))
                end
                comparison = Expr(:call, :>, unrolled, itercount)
                Expr(Ut == 1 ? :if : :elseif, comparison, lower_set(ls, vectorized, Ut, T, W, Symbol("##mask##"), :block))
            else
                comparison = if unrolledloop.hintexact
                    Expr(:call, :>, unrolled, unrolledloop.rangehint - (Ut + 1))
                else
                    Expr(:call, :>, unrolled, Expr(:call, :-, unrolled_numitersym, Ut + 1))
                end
                Expr(Ut == 1 ? :if : :elseif, comparison, lower_set(ls, vectorized, Ut, T, W, nothing, :block))
                # Expr(Ut == 1 ? :if : :elseif, comparison, lower_set(ls, vectorized, Ut, T, W, Symbol("##mask##"), :block))
            end
            push!(remblock.args, remblocknew)
            remblock = remblocknew
        end
        if firstiter
            firstiter = false
            if manageouterreductions && Ureduct < U
                Udiff = U - Ureduct
                loopq = lower_set(ls, vectorized, Udiff, T, W, nothing, :if)
                push!(q.args, loopq)
            end
            Ut = 1
            # setup for branchy remainder calculation
            comparison = if unrolledloop.hintexact
                Expr(:call, :(!=), unrolledloop.rangehint, unrolled)
            else
                Expr(:call, :(!=), unrolled_numitersym, unrolled)
            end
            remblock = Expr(:block)
            push!(q.args, Expr(:if, comparison, remblock))
        elseif !(Ut < U - 1 + vecisunrolled) || Ut == Ureduct
            break
        else
            Ut += 1
        end
    end
    q
end
function maskexpr(W::Symbol, looplimit, allon::Bool)
    rem = Expr(:call, lv(:valrem), W, looplimit)
    Expr(:(=), Symbol("##mask##"), Expr(:call, lv(allon ? :masktable : :mask), W, rem))
end
function definemask(loop::Loop, W::Symbol, allon::Bool)
    if loop.hintexact
        maskexpr(W, loop.rangehint, allon)
    else
        maskexpr(W, loop.rangesym, allon)
    end
end
function setup_Wmask!(ls::LoopSet, W::Symbol, typeT::Symbol, vectorized::Symbol, unrolled::Symbol, U::Int)
    pushpreamble!(ls, Expr(:(=), typeT, determine_eltype(ls)))
    pushpreamble!(ls, Expr(:(=), W, determine_width(ls, typeT, unrolled)))
    pushpreamble!(ls, definemask(ls.loops[vectorized], W, U > 1 && unrolled === vectorized))
end
function lower_tiled(ls::LoopSet, vectorized::Symbol, U::Int, T::Int)
    order = ls.loop_order.loopnames
    tiled    = order[end]
    unrolled = order[end-1]
    mangledtiled = tiledsym(tiled)
    W = gensym(:W)
    typeT = gensym(:T)
    setup_Wmask!(ls, W, typeT, vectorized, unrolled, U)
    tiledloop = ls.loops[tiled]
    static_tile = tiledloop.hintexact
    unrolledloop = ls.loops[unrolled]
    initialize_outer_reductions!(ls, 0, 4, W, typeT, unrolled)
    q = Expr(:block, Expr(:(=), mangledtiled, 0))
    # we build up the loop expression.
    Trem = Tt = T
    nloops = num_loops(ls);
    firstiter = true
    mangledtiled = tiledsym(tiled)
    local qifelse::Expr
    while Tt > 0
        tiledloopbody = Expr(:block)
        lower_unrolled!(tiledloopbody, ls, vectorized, U, Tt, W, typeT, unrolledloop)
        tiledloopbody = lower_nest(ls, nloops, vectorized, U, Tt, tiledloopbody, 0, W, nothing, :block)
        if firstiter
            push!(q.args, (static_tile && tiledloop.rangehint < 2T) ? tiledloopbody : Expr(:while, looprange(ls, tiled, Tt, mangledtiled, tiledloop), tiledloopbody))
        elseif static_tile
            push!(q.args, tiledloopbody)
        else # not static, not firstiter
            comparison = Expr(:call, :(==), mangledtiled, Expr(:call, :-, tiledloop.rangesym, Tt))
            qifelsenew = Expr(:elseif, comparison, tiledloopbody)
            push!(qifelse.args, qifelsenew)
            qifelse = qifelsenew
        end
        if static_tile
            if Tt == T
                Texprtype = :block
                Tt = looprangehint(ls, tiled) % T
                # Recalculate U
                U = solve_tilesize_constT(ls, Tt)
            else
                Tt = 0 # terminate
            end
            nothing
        else
            if firstiter
                comparison = Expr(:call, :(==), mangledtiled, tiledloop.rangesym)
                qifelse = Expr(:if, comparison, Expr(:block)) # do nothing
                push!(q.args, qifelse)
                Tt = 0
            end
            Tt += 1 # we start counting up by 1
            if Tt == T # terminate on Tt = T
                Tt = 0
            end
            nothing
        end
        firstiter = false
    end
    q = gc_preserve(ls, q)
    reduce_expr!(q, ls, 4)
    Expr(:block, ls.preamble, q)
end
function lower_unrolled(ls::LoopSet, vectorized::Symbol, U::Int)
    order = ls.loop_order.loopnames
    # @show order
    unrolled = last(order)
    # W = VectorizationBase.pick_vector_width(ls, unrolled)
    W = gensym(:W)
    typeT = gensym(:T)
    setup_Wmask!(ls, W, typeT, vectorized, unrolled, U)
    q = lower_unrolled!(Expr(:block, Expr(:(=), unrolled, 0)), ls, vectorized, U, -1, W, typeT, ls.loops[unrolled])
    Expr(:block, ls.preamble, q)
end

# Here, we have to figure out how to convert the loopset into a vectorized expression.
# This must traverse in a parent -> child pattern
# but order is also dependent on which loop inds they depend on.
# Requires sorting
function lower(ls::LoopSet)
    order, vectorized, U, T = choose_order(ls)
    # @show order, U, T
    # @show ls.loop_order.loopnames
    istiled = T != -1
    fillorder!(ls, order, istiled)
    # @show order, ls.loop_order.loopnames
    istiled ? lower_tiled(ls, vectorized, U, T) : lower_unrolled(ls, vectorized, U)
end

Base.convert(::Type{Expr}, ls::LoopSet) = lower(ls)
Base.show(io::IO, ls::LoopSet) = println(io, lower(ls))


text:
struct RootNode <: Node
    children::Tuple{Any}
    function RootNode(node)
        new((node,))
    end
end

function _traversal(node, visitfn::Function, visitfnargs::Union{Any, Nothing}, visited)
    if haskey(visited, node)
        return visited[node]
    else
        new_children = [_traversal(child, visitfn, visitfnargs, visited) for child in node.children]
        result = visitfnargs !== nothing ?
            visitfn(visited, node, visitfnargs, new_children...) :
                visitfn(visited, node, new_children...)
        visited[node] = result
        return result
    end
end

"""
visitfn requires the interface node::Node, visitfnargs, _posttraversal
or node::Node, _posttraversal if visitfnargs is nothing.
"""
function _posttraversal(node, visitfn::Function, visitfnargs::Union{Any, Nothing}, visited)
    if haskey(visited, node)
        return visited[node]
    else
        result = visitfnargs !== nothing ?
            visitfn(node, visitfnargs, (x -> _posttraversal(x, visitfn, visitfnargs, visited))) :
                visitfn(node, (x -> _posttraversal(x, visitfn, visitfnargs, visited)))
        visited[node] = result
        return result
    end
end

function traversal(node, pretraversalfn::Function, visitfn::Function,
        pretraversalfnargs::Union{Any, Nothing}, visitfnargs::Union{Any, Nothing},
        visited::Union{Dict, Nothing}, posttraversal::Bool)
    if visited === nothing
        visited = Dict{Any, Any}()
    end
    root = RootNode(node)
    root = pretraversalfnargs !== nothing ? pretraversalfn(root, pretraversalfnargs...) : pretraversalfn(root)
    root = posttraversal ? _posttraversal(
        root, visitfn, visitfnargs, visited) : _traversal(
                root, visitfn, visitfnargs, visited)
    return root.children[1]
end

function traversal(node, pretraversalfn::Function, visitfn::Function,
        pretraversalfnargs::Union{Any, Nothing}, visitfnargs::Union{Any, Nothing},
        visited::Union{Dict, Nothing})
    return traversal(node, pretraversalfn, visitfn, pretraversalfnargs,
            visitfnargs, visited, false)
end

ext: jl

max_stars_repo_name: Sohva/TensorDSL

max_issues_repo_name: Sohva/TensorDSL

max_forks_repo_name: Sohva/TensorDSL


text: ROOT_FILEPATH = "/home/alex/stanford_dev/gail-driver/"

max_stars_repo_name: sisl/gail-driver

max_issues_repo_name: fcdtc/gail-driver

max_forks_repo_name: fcdtc/gail-driver


text: Dict{String, Any}("longpi" => Dict{String, Any}("value" => "3.141592653589793", "type" => "float"), "neglongpi" => Dict{String, Any}("value" => "-3.141592653589793", "type" => "float"))


text:
mutable struct BadAgent <: AbstractAgent
    useless::Int
    pos::Int
end
mutable struct BadAgentId <: AbstractAgent
    id::Float64
end
struct ImmutableAgent <: AbstractAgent
    id::Int
end
mutable struct DiscreteVelocity <: AbstractAgent
    id::Int
    pos::NTuple{2,Float64}
    vel::NTuple{2,Int}
    diameter::Float64
end
mutable struct ParametricAgent{T<:Integer} <: AbstractAgent
    id::T
    pos::NTuple{2,T}
    weight::T
    info::String
end

@testset "Model construction" begin
    # Shouldn't use ImmutableAgent since it cannot be edited
    agent = ImmutableAgent(1)
    @test_logs (
        :warn,
        "AgentType should be mutable. Try adding the `mutable` keyword infront of `struct` in your agent definition.",
    ) ABM(agent)
    # Warning is suppressed if flag is set
    @test Agents.agenttype(ABM(agent; warn = false)) <: AbstractAgent
    # Cannot use BadAgent since it has no `id` field
    @test_throws ArgumentError ABM(BadAgent)
    agent = BadAgent(1, 1)
    @test_throws ArgumentError ABM(agent)
    # Cannot use BadAgent in a grid space context since `pos` has an invalid type
    @test_throws ArgumentError ABM(BadAgent, GridSpace((1, 1)))
    @test_throws ArgumentError ABM(agent, GridSpace((1, 1)))
    # Cannot use BadAgentId since `id` has an invalid type
    @test_throws ArgumentError ABM(BadAgentId)
    agent = BadAgentId(1.0)
    @test_throws ArgumentError ABM(agent)
    # Cannot use Agent0 in a grid space context since it has no `pos` field
    @test_throws ArgumentError ABM(Agent0, GridSpace((1, 1)))
    agent = Agent0(1)
    @test_throws ArgumentError ABM(agent, GridSpace((1, 1)))
    # Cannot use Agent3 in a graph space context since `pos` has an invalid type
    @test_throws ArgumentError ABM(Agent3, GraphSpace(Agents.Graph(1)))
    agent = Agent3(1, (1, 1), 5.3)
    @test_throws ArgumentError ABM(agent, GraphSpace(Agents.Graph(1)))
    # Cannot use Agent3 in a continuous space context since `pos` has an invalid type
    @test_throws ArgumentError ABM(Agent3, ContinuousSpace((1, 1), 0.1))
    @test_throws ArgumentError ABM(agent, ContinuousSpace((1, 1), 0.1))
    # Cannot use Agent4 in a continuous space context since it has no `vel` field
    @test_throws ArgumentError ABM(Agent4, ContinuousSpace((1, 1), 0.1))
    agent = Agent4(1, (1, 1), 5)
    @test_throws ArgumentError ABM(agent, ContinuousSpace((1, 1), 0.1))
    # Shouldn't use DiscreteVelocity in a continuous space context since `vel` has an invalid type
    @test_logs (
        :warn,
        "`vel` field in Agent struct should be of type `NTuple{<:AbstractFloat}` when using ContinuousSpace.",
    ) ABM(DiscreteVelocity, ContinuousSpace((1, 1), 0.1))
    agent = DiscreteVelocity(1, (1, 1), (2, 3), 2.4)
    @test_logs (
        :warn,
        "`vel` field in Agent struct should be of type `NTuple{<:AbstractFloat}` when using ContinuousSpace.",
    ) ABM(agent, ContinuousSpace((1, 1), 0.1))
    # Warning is suppressed if flag is set
    @test Agents.agenttype(ABM(agent, ContinuousSpace((1, 1), 0.1); warn = false)) <:
          AbstractAgent
    # Shouldn't use ParametricAgent since it is not a concrete type
    @test_logs (
        :warn,
        "AgentType is not concrete. If your agent is parametrically typed, you're probably seeing this warning because you gave `Agent` instead of `Agent{Float64}` (for example) to this function. You can also create an instance of your agent and pass it to this function. If you want to use `Union` types for mixed agent models, you can silence this warning.",
    ) ABM(ParametricAgent, GridSpace((1, 1)))
    # Warning is suppressed if flag is set
    @test Agents.agenttype(ABM(ParametricAgent, GridSpace((1, 1)); warn = false)) <:
          AbstractAgent
    # ParametricAgent{Int} is the correct way to use such an agent
    @test Agents.agenttype(ABM(ParametricAgent{Int}, GridSpace((1, 1)))) <: AbstractAgent
    #Type inferance using an instance can help users here
    agent = ParametricAgent(1, (1, 1), 5, "Info")
    @test Agents.agenttype(ABM(agent, GridSpace((1, 1)))) <: AbstractAgent
    #Mixed agents
    @test Agents.agenttype(ABM(Union{Agent0,Agent1}; warn = false)) <: AbstractAgent
    @test_logs (
        :warn,
        "AgentType is not concrete. If your agent is parametrically typed, you're probably seeing this warning because you gave `Agent` instead of `Agent{Float64}` (for example) to this function. You can also create an instance of your agent and pass it to this function. If you want to use `Union` types for mixed agent models, you can silence this warning.",
    ) ABM(Union{Agent0,Agent1})
    @test_throws ArgumentError ABM(Union{Agent0,BadAgent}; warn = false)
    @test_throws ArgumentError ABM(Agent6, GridSpace((50, 50)))
    @test_throws ErrorException Agents.notimplemented(ABM(Agent0))
    # Test @agent macro
    @agent A3 GridAgent{2} begin
        weight::Float64
    end
    @test A3 <: AbstractAgent
    @test fieldnames(A3) == fieldnames(Agent3)
    @test A3.types == Agent3.types
end

@testset "sample!" begin
    Random.seed!(6459)
    model = ABM(Agent2)
    for i in 1:20
        add_agent!(model, rand() / rand())
    end
    allweights = [i.weight for i in values(model.agents)]
    mean_weights = sum(allweights) / length(allweights)
    sample!(model, 12, :weight)
    @test nagents(model) == 12
    allweights = [i.weight for i in values(model.agents)]
    mean_weights_new = sum(allweights) / length(allweights)
    @test mean_weights_new > mean_weights
    sample!(model, 40, :weight)
    @test nagents(model) == 40
    allweights = [i.weight for i in values(model.agents)]
    mean_weights_new = sum(allweights) / length(allweights)
    @test mean_weights_new > mean_weights

    Random.seed!(6459)
    model2 = ABM(Agent3, GridSpace((10, 10)))
    for i in 1:20
        add_agent_single!(Agent3(i, (1, 1), rand() / rand()), model2)
    end
    @test sample!(model2, 10) == nothing
    @test sample!(model2, 10, :weight) == nothing
    allweights = [i.weight for i in values(model2.agents)]
    mean_weights = sum(allweights) / length(allweights)
    sample!(model2, 12, :weight)
    @test nagents(model2) == 12
    allweights = [i.weight for i in values(model2.agents)]
    mean_weights_new = sum(allweights) / length(allweights)
    @test mean_weights_new > mean_weights

    sample!(model2, 40, :weight)
    @test nagents(model2) == 40

    Random.seed!(6459)
    #Guarantee all starting weights are unique
    model3 = ABM(Agent2)
    while true
        for i in 1:20
            add_agent!(model3, rand() / rand())
        end
        allweights = [i.weight for i in values(model3.agents)]
        allunique(allweights) && break
    end
    # Cannot draw 50 samples out of a pool of 20 without replacement
    @test_throws ErrorException sample!(model3, 50, :weight; replace = false)
    sample!(model3, 15, :weight; replace = false)
    allweights = [i.weight for i in values(model3.agents)]
    @test allunique(allweights)
    model3 = ABM(Agent2)
    while true
        for i in 1:20
            add_agent!(model3, rand() / rand())
        end
        allweights = [i.weight for i in values(model3.agents)]
        allunique(allweights) && break
    end
    sample!(model3, 100, :weight; replace = true)
    allweights = [i.weight for i in values(model3.agents)]
    @test !allunique(allweights)
end

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

end

@testset "random agent" begin
    space = GridSpace((10, 10))
    model = ABM(Union{Daisy,Land}, space; warn = false)
    fill_space!(Daisy, model, "black")
    add_agent!(Land(999, (1, 1), 999), model)

    a = random_agent(model)
    @test typeof(a) <: Union{Daisy,Land}

    c1(a) = a isa Land
    a = random_agent(model, c1)
    @test a.id == 999

    c2(a) = a isa Float64
    a = random_agent(model, c2)
    @test isnothing(a)
end

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

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


text:
using BenchmarkTools
using FinancialMonteCarlo, CUDA
CUDA.allowscalar(false)

S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01;
D = 90.0;

Nsim = 10000;
Nstep = 30;
sigma = 0.2;
lam = 5.0;
mu1 = 0.03;
sigma1 = 0.02;
mc = MonteCarloConfiguration(Nsim, Nstep);

mc_2 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.CudaMode());
toll = 0.8;

rfCurve = ZeroRate(r);

FwdData = Forward(T)
EUData = EuropeanOption(T, K)
AMData = AmericanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model = MertonProcess(sigma, lam, mu1, sigma1, Underlying(S0, d));

@show "CUDA_1 fwd"

@show "STD fwd"
@btime FwdPrice = pricer(Model, rfCurve, mc, FwdData);
@show "CUDA_1 fwd"

@show "CUDA_2 fwd"
@btime FwdPrice = pricer(Model, rfCurve, mc_2, FwdData);
@show "std eu"
@btime EuPrice = pricer(Model, rfCurve, mc, EUData);
@show "CUDA_1 eu"

@show "CUDA_2 eu"
@btime EuPrice = pricer(Model, rfCurve, mc_2, EUData);
@show "std am"
@btime AmPrice = pricer(Model, rfCurve, mc, AMData);
@show "CUDA_1 am"

@show "CUDA_2 am"
@btime AmPrice = pricer(Model, rfCurve, mc_2, AMData);


text:
using DoubleFloats
using LinearAlgebra

using SemiLagrangian:
    Advection,
    magicsplit,
    nosplit,
    sizeall,
    AdvectionData,
    getdata,
    advection!,
    UniformMesh,
    getrotationvar,
    AbstractInterpolation,
    Lagrange,
    Hermite,
    B_SplineLU,
    B_SplineFFT,
    interpolate!,
    NoTimeAlg,
    ABTimeAlg_ip
"""

   exact(tf, mesh1, mesh2)

   Julia function to compute exact solution "

```math
\frac{d f}{dt} +  (y \frac{df}{delta1} - x \frac{df}{delta2}) = 0
```

"""
function exact!(f, mesh1::UniformMesh{T}, mesh2::UniformMesh{T}, tf::T) where {T}
    for (i, x) in enumerate(mesh1.points), (j, y) in enumerate(mesh2.points)
        s, c = sincos(tf)
        xn, yn = c * x - s * y, +s * x + c * y
        f[i, j] = exp(-2 * ((xn)^2 + (yn + T(6 // 5))^2))
    end
end

function test_rotation(
    sz::NTuple{2,Int},
    interp_sp::AbstractInterpolation{T},
    interp_v::AbstractInterpolation{T},
    nbdt::Int,
) where {T}
@time begin
    spmin, spmax, nsp = T(-5), T(5), sz[1]
    vmin, vmax, nv = -T(12 // 2), T(9 // 2), sz[2]

    mesh_sp = UniformMesh(spmin, spmax, nsp)
    mesh_v = UniformMesh(vmin, vmax, nv)
end
@time begin
    dt = T(2big(pi) / nbdt)
    adv = Advection(
        (mesh_sp,mesh_v,),
        [interp_sp, interp_v,],
        dt,
        [([1,2], 1, 1, true),([2,1], 1, 2, true)];
        tab_coef=magicsplit(dt)
    )
    #    adv = Advection1d((mesh_sp,), (mesh_v,), (interp_sp,), (interp_v,), dt; tab_coef=[1], tab_fct=[identity])
end

@time begin

    sz = sizeall(adv)
    tabref = zeros(T, sz)
    exact!(tabref, mesh_sp, mesh_v, T(0))

    pvar = getrotationvar(adv)

    advdata = AdvectionData(adv, tabref, pvar)

    diffmax = 0
    data = getdata(advdata)
end

    for ind = 1:nbdt
        while advection!(advdata)
        end
        exact!(tabref, mesh_sp, mesh_v, dt * ind)
        diff = norm(data .- tabref, Inf)
        diffmax = max(diffmax, diff)
        @show ind, diff
    end
    println("test_rotation sz=$sz interp=$interp_sp, $interp_v nbdt=$nbdt diffmax=$diffmax")
    return diffmax

end
function maxind(tab::Array{T,N}) where {T,N}
    v=-Inf
    res = CartesianIndex(0,0)
    for ind in CartesianIndices(tab)
        if tab[ind] > v
            res = ind
            v = tab[ind]
        end
    end
    return res
end

function test_rotation2d(
    sz::NTuple{2,Int},
    interp::Vector{I},
    nbdt::Int,
    timealg,
    ordalg
) where {T, I <: AbstractInterpolation{T} }
    spmin, spmax, nsp = T(-5), T(5), sz[1]
    vmin, vmax, nv = -T(5), T(5), sz[2]

    mesh_sp = UniformMesh(spmin, spmax, nsp)
    mesh_v = UniformMesh(vmin, vmax, nv)
    dt = T(2big(pi) / nbdt)
    adv = Advection(
        (mesh_sp,mesh_v,),
        interp,
        dt,
        [([1,2], 2, 1, false),];
        tab_coef=nosplit(dt),
        timealg=timealg,
        ordalg=ordalg
    )
    #    adv = Advection1d((mesh_sp,), (mesh_v,), (interp_sp,), (interp_v,), dt; tab_coef=[1], tab_fct=[identity])

    sz = sizeall(adv)
    tabref = zeros(T, sz)
    exact!(tabref, mesh_sp, mesh_v, T(0))

    pvar = getrotationvar(adv)

    advdata = AdvectionData(adv, tabref, pvar)

    diffmax = 0
    data = getdata(advdata)

    for ind = 1:nbdt
        while advection!(advdata)
        end
        exact!(tabref, mesh_sp, mesh_v, dt * ind)

        indmaxdata = maxind(advdata.data)
        indmaxexact = maxind(tabref)

        diff = norm(data .- tabref, Inf)
        diffmax = max(diffmax, diff)
        @show ind, diff, indmaxdata, indmaxexact
    end
    println("test_rotation sz=$sz diffmax=$diffmax")
    return diffmax

end
function test_rotation(
    sz::NTuple{2,Int},
    interp::Vector{I},
    nbdt::Int,
) where {T,edge,order, I <: AbstractInterpolation{T,edge,order}}
    spmin, spmax, nsp = T(-4.5), T(5), sz[1]
    vmin, vmax, nv = -T(5.2), T(5), sz[2]

    mesh_sp = UniformMesh(spmin, spmax, nsp)
    mesh_v = UniformMesh(vmin, vmax, nv)

    dt = T(2big(pi) / nbdt)

    # dec1 = -2tan(dt/2) / step(mesh_sp) * mesh_v.points
    # dec2 = 2tan(dt/2) / step(mesh_v) * mesh_sp.points
    dec1 = zeros(T, sz)
    dec2 = zeros(T, sz)
    tgdt = 2tan(dt / 2)
    coef = 1 / (1 + tgdt^2 / 4)
    for i = 1:sz[1], j = 1:sz[2]
        x = mesh_sp.points[i]
        y = mesh_v.points[j]
        v = [
            -tgdt * coef * (y + tgdt * x / 2) / step(mesh_sp),
            (tgdt * coef * (x - tgdt * y / 2)) / step(mesh_v),
        ]
        # nv = norm(v)
        # nr = norm([x,y])
        # if nv != 0
        #     v *= cor*nr/nv
        # end
        dec1[i, j], dec2[i, j] = v
    end

    @show minimum(dec1), maximum(dec1)
    @show minimum(dec2), maximum(dec2)

    #    @show dec1, dec2

    tabref = zeros(T, sz)
    data = zeros(T, sz)

    exact!(tabref, mesh_sp, mesh_v, T(0))
    copyto!(data, tabref)

    buf = zeros(T, sz)
    diffmax = 0
    for ind = 1:nbdt
        interpolate!(buf, data, (i,ind) -> (i==1) ? dec1[ind] : dec2[ind], interp)
        copyto!(data, buf)
        exact!(tabref, mesh_sp, mesh_v, dt * ind)
        diff = norm(data .- tabref, Inf)
        diffmax = max(diffmax, diff)
        @show ind, diff

    end
    println("test_rotation sz=$sz interp=$interp nbdt=$nbdt diffmax=$diffmax")
    return diffmax
end

@testset "test rotation" begin
    T = Float64
#    @time @test test_rotation((100, 122), [Lagrange(5, T), Lagrange(5, T)], 11) < 1e-3
# @time ret1 = test_rotation2d((1000, 1000), [Lagrange(9, T), Lagrange(9, T)], 20, ABTimeAlg, 1)
# @time ret2 = test_rotation2d((1000, 1000), [Lagrange(9, T), Lagrange(9, T)], 40, ABTimeAlg, 1)
# @show ret1, ret2, ret1/ret2

# @test ret2 < (ret1*1.1)/2

# @time ret1 = test_rotation2d((200, 102), [Lagrange(9, T), Lagrange(9, T)], 20, ABTimeAlg, 2)
# @time ret2 = test_rotation2d((200, 102), [Lagrange(9, T), Lagrange(9, T)], 40, ABTimeAlg, 2)
# @show ret1, ret2, ret1/ret2

# @test ret2 < (ret1*1.1)/4

# @time ret1 = test_rotation2d((200, 102), [Lagrange(9, T), Lagrange(9, T)], 20, ABTimeAlg, 3)
# @time ret2 = test_rotation2d((200, 102), [Lagrange(9, T), Lagrange(9, T)], 40, ABTimeAlg, 3)
# @show ret1, ret2, ret1/ret2

# @test ret2 < (ret1*1.1)/8

# @time ret1 = test_rotation2d((200, 102), [Lagrange(9, T), Lagrange(9, T)], 20, ABTimeAlg, 4)
# @time ret2 = test_rotation2d((200, 102), [Lagrange(9, T), Lagrange(9, T)], 40, ABTimeAlg, 4)
# @show ret1, ret2, ret1/ret2

# @test ret2 < (ret1*1.1)/16
@time @test test_rotation((2000, 1022), Lagrange(5, T), Lagrange(5, T), 11) < 1e-3
@time @test test_rotation((2000, 1022), Lagrange(5, T), Lagrange(5, T), 11) < 1e-3
@time @test test_rotation((2000, 1022), Hermite(5, T), Hermite(5, T), 11) < 1e-3
@time @test test_rotation((2000, 1022), Hermite(5, T), Hermite(5, T), 11) < 1e-3
    @time @test test_rotation(
        (128, 256),
        B_SplineLU(5, 128, T),
        B_SplineLU(5, 256, T),
        11,
    ) < 1e-3
    @time @test test_rotation(
        (128, 256),
        B_SplineFFT(5, 128, T),
        B_SplineFFT(5, 256, T),
        11,
    ) < 1e-3
    T = Double64
#    @time @test test_rotation((100, 122), [Lagrange(15, T),Lagrange(15, T)], 11) < 1e-6
    @time @test test_rotation((200, 220), Lagrange(15, T), Lagrange(15, T), 11) < 1e-7
    @time @test test_rotation(
        (128, 256),
        B_SplineLU(15, 128, T),
        B_SplineLU(15, 256, T),
        11,
    ) < 1e-8
    @time @test test_rotation(
        (128, 256),
        B_SplineFFT(15, 128, T),
        B_SplineFFT(15, 256, T),
        11,
    ) < 1e-8
    T = BigFloat
#    @time @test test_rotation((100, 122), [Lagrange(21, T),Lagrange(21, T)], 11) < 1e-7
    @time @test test_rotation((200, 220), Lagrange(25, T), Lagrange(25, T), 11) < 1e-9
    @time @test test_rotation(
        (128, 256),
        B_SplineLU(25, 128, T),
        B_SplineLU(25, 256, T),
        11,
    ) < 1e-9
    @time @test test_rotation(
        (128, 256),
        B_SplineFFT(25, 128, T),
        B_SplineFFT(25, 256, T),
        11,
    ) < 1e-9
end


text:
"""

abstract type FlowChannel{Sensors<:Tuple,Actuators<:Tuple} <: Component end

"""

# channels always sense the reversal of the currents they actuate
FlowChannel(T) = FlowChannel{Tuple{Reversal{T}},Tuple{T}}
FlowChannel(S, A) = FlowChannel{Tuple{S,Reversal{A}},Tuple{A}}
FlowChannel() = FlowChannel{Tuple{Nothing},Tuple{Nothing}}

sensed(::FlowChannel{S,A}) where {S,A} = typeflatten(S)
actuated(::FlowChannel{S,A}) where {S,A} = typeflatten(A)
sensed(::FlowChannel{Tuple{Nothing},A}) where {A} = Vector{DataType}()
actuated(::FlowChannel{S,Tuple{Nothing}}) where {S} = Vector{DataType}()

function typeflatten(s::DataType)
    map(fieldtypes(s)) do el
        el <: Tuple && return typeflatten(el)
        return (el,)
    end |> Iterators.flatten |> collect |> unique!
end

voltage(el) = (Voltage,)
ionic(x) = x <: Ion

currents(i::FlowChannel) =
    map(filter(ionic, actuated(i))) do thing
        Current{thing}
    end

sensedvars(i::FlowChannel) =
    map(sensed(i)) do thing
        Symbol(thing |> shorthand_name)
    end

mrna(i::FlowChannel) =
    map(filter(ionic, actuated(i))) do thing
        mRNA{thing}
    end

reversals(i::FlowChannel) =
    map(filter(ionic, actuated(i))) do thing
        Reversal{thing}
    end

conductances(i::FlowChannel) =
    map(filter(ionic, actuated(i))) do thing
        Conductance{thing}
    end

sensed_ions(i::FlowChannel) = filter(ionic, sensed(i))

instantiate_hooks(pre::Compartment, to::Component, args::Function...) =
    map(args) do fun
        map(fun(to)) do thing
            _name = shorthand_name(thing)
            if has_dynamics(pre, thing)
                return @variables $_name(t)
            else
                return @parameters $_name
            end
        end |> Iterators.flatten |> collect
    end |> Iterators.flatten |> collect

instantiate_variables(c::Component, args::Function...) =
    map(args) do fun
        map(fun(c)) do thing
            _name = shorthand_name(thing)
            return @variables $_name(t)
        end |> Iterators.flatten |> collect
    end |> Iterators.flatten |> collect

instantiate_parameters(c::Component, args::Function...) =
    map(args) do fun
        map(fun(c)) do thing
            _name = shorthand_name(thing)
            return @parameters $_name
        end |> Iterators.flatten |> collect
    end |> Iterators.flatten |> collect

instantiate_variables(v::Vector{Symbol}) =
    map(v) do f
        @variables $f(t)
    end |> Iterators.flatten |> collect

instantiate_variables(s::Symbol) = @variables $s(t)

instantiate_parameters(v::Vector{Symbol}) =
    map(v) do f
        @variables $f
    end |> Iterators.flatten |> collect

instantiate_parameters(s::Symbol) = @parameters $f

get_actuator(c::FlowChannel, sys::ODESystem, v::Type{Voltage}) =
    sum(currents(c)) do I
        ModelingToolkit.getvar(sys, shorthand_name(I); namespace=true)
    end

function get_actuator(c::FlowChannel, sys::ODESystem, f::DataType)
    indx = findfirst(x -> x == f, actuated(c))
    isnothing(indx) && return Num(0.0)
    return ModelingToolkit.getvar(sys, shorthand_name.(currents(c))[indx]; namespace=true)
end

function get_sensor(c::FlowChannel, sys::ODESystem, f::DataType)
    if f ∈ sensed(c)
        return ModelingToolkit.getvar(sys, shorthand_name(f); namespace=true)
    end
end
get_sensor(c::FlowChannel, sys::ODESystem, f::Type{Voltage}) =
    ModelingToolkit.getvar(sys, shorthand_name(f); namespace=true)


text:
@doc raw"""
    StateSpaceSystem

Abstract type that unifies the definition of state space models matrices such as ``y, Z, d, T, c, R, H, Q`` for
linear models.
"""
abstract type StateSpaceSystem end

@doc raw"""
    LinearUnivariateTimeInvariant{Fl}(
        y::Vector{Fl},
        Z::Vector{Fl},
        T::Matrix{Fl},
        R::Matrix{Fl},
        d::Fl,
        c::Vector{Fl},
        H::Fl,
        Q::Matrix{Fl},
    ) where Fl <: AbstractFloat

Definition of the system matrices ``y, Z, d, T, c, R, H, Q`` for linear univariate time invariant state space models.

```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  Z\alpha_{t} + d + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, H)\\
        \alpha_{t+1} &= T\alpha_{t} + c + R\eta_{t} \quad &\eta_{t} \sim \mathcal{N}(0, Q)\\
    \end{aligned}
\end{gather*}
```

where:

* ``y_{t}`` is a scalar
* ``Z`` is a ``m \times 1`` vector
* ``d`` is a scalar
* ``T`` is a ``m \times m`` matrix
* ``c`` is a ``m \times 1`` vector
* ``R`` is a ``m \times r`` matrix
* ``H`` is a scalar
* ``Q`` is a ``r \times r`` matrix
"""
mutable struct LinearUnivariateTimeInvariant{Fl<:AbstractFloat} <: StateSpaceSystem
    y::Vector{Fl}
    Z::Vector{Fl}
    T::Matrix{Fl}
    R::Matrix{Fl}
    d::Fl
    c::Vector{Fl}
    H::Fl
    Q::Matrix{Fl}

    function LinearUnivariateTimeInvariant{Fl}(
        y::Vector{Fl},
        Z::Vector{Fl},
        T::Matrix{Fl},
        R::Matrix{Fl},
        d::Fl,
        c::Vector{Fl},
        H::Fl,
        Q::Matrix{Fl},
    ) where Fl <: AbstractFloat
        mz = length(Z)
        mt1, mt2 = size(T)
        mr, rr = size(R)
        mc = length(c)
        rq1, rq2 = size(Q)

        dim_str =
            "Z is $(mz)x1, T is $(mt1)x$(mt2), R is $(mr)x$(rr), " *
            "d is a scalar, c is $(mc)x1, H is a scalar."

        !(mz == mt1 == mt2 == mr == mc) && throw(DimensionMismatch(dim_str))
        !(rr == rq1 == rq2) && throw(DimensionMismatch(dim_str))

        return new{Fl}(y, Z, T, R, d, c, H, Q)
    end
end

num_states(system::LinearUnivariateTimeInvariant) = size(system.T, 1)

@doc raw"""
    LinearUnivariateTimeVariant{Fl}(
        y::Vector{Fl},
        Z::Vector{Vector{Fl}},
        T::Vector{Matrix{Fl}},
        R::Vector{Matrix{Fl}},
        d::Vector{Fl},
        c::Vector{Vector{Fl}},
        H::Vector{Fl},
        Q::Vector{Matrix{Fl}},
    ) where Fl <: AbstractFloat

Definition of the system matrices ``y, Z, d, T, c, R, H, Q`` for linear univariate time variant state space models.

```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  Z_{t}\alpha_{t} + d_{t} + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, H_{t})\\
        \alpha_{t+1} &= T_{t}\alpha_{t} + c_{t} + R_{t}\eta_{t} \quad &\eta_{t} \sim \mathcal{N}(0, Q_{t})\\
    \end{aligned}
\end{gather*}
```

where:

* ``y_{t}`` is a scalar
* ``Z_{t}`` is a ``m \times 1`` vector
* ``d_{t}`` is a scalar
* ``T_{t}`` is a ``m \times m`` matrix
* ``c_{t}`` is a ``m \times 1`` vector
* ``R_{t}`` is a ``m \times r`` matrix
* ``H_{t}`` is a scalar
* ``Q_{t}`` is a ``r \times r`` matrix
"""
mutable struct LinearUnivariateTimeVariant{Fl<:AbstractFloat} <: StateSpaceSystem
    y::Vector{Fl}
    Z::Vector{Vector{Fl}}
    T::Vector{Matrix{Fl}}
    R::Vector{Matrix{Fl}}
    d::Vector{Fl}
    c::Vector{Vector{Fl}}
    H::Vector{Fl}
    Q::Vector{Matrix{Fl}}

    function LinearUnivariateTimeVariant{Fl}(
        y::Vector{Fl},
        Z::Vector{Vector{Fl}},
        T::Vector{Matrix{Fl}},
        R::Vector{Matrix{Fl}},
        d::Vector{Fl},
        c::Vector{Vector{Fl}},
        H::Vector{Fl},
        Q::Vector{Matrix{Fl}},
    ) where Fl <: AbstractFloat

        # TODO assert dimensions

        return new{Fl}(y, Z, T, R, d, c, H, Q)
    end
end

num_states(system::LinearUnivariateTimeVariant) = size(system.T[1], 1)

"""
TODO
"""
mutable struct LinearMultivariateTimeInvariant{Fl<:Real} <: StateSpaceSystem
    y::Matrix{Fl}
    Z::Matrix{Fl}
    T::Matrix{Fl}
    R::Matrix{Fl}
    d::Vector{Fl}
    c::Vector{Fl}
    H::Matrix{Fl}
    Q::Matrix{Fl}

    function LinearMultivariateTimeInvariant{Fl}(
        y::Matrix{Fl},
        Z::Matrix{Fl},
        T::Matrix{Fl},
        R::Matrix{Fl},
        d::Vector{Fl},
        c::Vector{Fl},
        H::Matrix{Fl},
        Q::Matrix{Fl},
    ) where Fl

        # TODO assert dimensions

        return new{Fl}(y, Z, T, R, d, c, H, Q)
    end
end

hasnan(y::Vector{Fl}) where Fl = findfirst(isnan, y) === nothing ? false : true
num_states(system::LinearMultivariateTimeInvariant) = size(system.T, 1)

"""
TODO
"""
mutable struct LinearMultivariateTimeVariant{Fl<:Real} <: StateSpaceSystem
    y::Matrix{Fl}
    Z::Vector{Matrix{Fl}}
    T::Vector{Matrix{Fl}}
    R::Vector{Matrix{Fl}}
    d::Vector{Vector{Fl}}
    c::Vector{Vector{Fl}}
    H::Vector{Matrix{Fl}}
    Q::Vector{Matrix{Fl}}

    function LinearMultivariateTimeVariant{Fl}(
        y::Matrix{Fl},
        Z::Vector{Matrix{Fl}},
        T::Vector{Matrix{Fl}},
        R::Vector{Matrix{Fl}},
        d::Vector{Vector{Fl}},
        c::Vector{Vector{Fl}},
        H::Vector{Matrix{Fl}},
        Q::Vector{Matrix{Fl}},
    ) where Fl

        # TODO assert dimensions

        return new{Fl}(y, Z, T, R, d, c, H, Q)
    end
end

# Some useful functions for dealing with these systems
function fill_system_matrice_with_value_in_time(vector::Vector{Fl}, value::Fl) where Fl
    for t in axes(vector, 1)
        vector[t] = value
    end
end

function repeat_system_matrice_n_times(value::T, n::Int) where T
    vector_of_matrice = Vector{T}(undef, n)
    for i in 1:n
        vector_of_matrice[i] = value
    end
    return vector_of_matrice
end

function to_multivariate_Z(Z::Vector)
    return permutedims(Z)
end
function to_multivariate_Z(Z::Vector{Vector{Fl}}) where Fl
    multivariate_Z = Vector{Matrix{Fl}}(undef, length(Z))
    for i in 1:length(Z)
        multivariate_Z[i] = permutedims(Z[i])
    end
    return multivariate_Z
end
function to_multivariate_H(H::Fl) where Fl
    return [H][:, :]
end
function to_multivariate_H(H::Vector{Fl}) where Fl
    multivariate_H = Vector{Matrix{Fl}}(undef, length(H))
    for i in 1:length(H)
        multivariate_H[i] = permutedims([H[i]])
    end
    return multivariate_H
end
function to_multivariate_d(d::Fl) where Fl
    return [d]
end
function to_multivariate_d(d::Vector{Fl}) where Fl
    multivariate_d = Vector{Vector{Fl}}(undef, length(d))
    for i in 1:length(d)
        multivariate_d[i] = [d[i]]
    end
    return multivariate_d
end

# Bridges between Linear systems

# LinearUnivariateTimeInvariant to (LinearUnivariateTimeVariant,
#                                   LinearMultivariateTimeInvariant, TODO
#                                   LinearMultivariateTimeVariant)
function to_univariate_time_variant(system::LinearUnivariateTimeInvariant{Fl}) where Fl
    n = length(system.y)
    y = system.y
    Z = repeat_system_matrice_n_times(system.Z, n)
    T = repeat_system_matrice_n_times(system.T, n)
    R = repeat_system_matrice_n_times(system.R, n)
    d = repeat_system_matrice_n_times(system.d, n)
    c = repeat_system_matrice_n_times(system.c, n)
    H = repeat_system_matrice_n_times(system.H, n)
    Q = repeat_system_matrice_n_times(system.Q, n)
    return LinearUnivariateTimeVariant{Fl}(y, Z, T, R, d, c, H, Q)
end
function to_multivariate_time_variant(system::LinearUnivariateTimeInvariant)
    univariate_time_variant = to_univariate_time_variant(system)
    return to_multivariate_time_variant(univariate_time_variant)
end

# LinearUnivariateTimeVariant to (LinearMultivariateTimeVariant)
function to_multivariate_time_variant(system::LinearUnivariateTimeVariant{Fl}) where Fl
    y = system.y[:, :]
    Z = to_multivariate_Z(system.Z)
    T = system.T
    R = system.R
    d = to_multivariate_d(system.d)
    c = system.c
    H = to_multivariate_H(system.H)
    Q = system.Q
    return LinearMultivariateTimeVariant{Fl}(y, Z, T, R, d, c, H, Q)
end

to_multivariate_time_variant(system::LinearMultivariateTimeVariant) = system

# Functions for simulations
function simulate(
    sys::LinearUnivariateTimeInvariant{Fl},
    initial_state::Vector{Fl},
    n::Int;
    return_simulated_states::Bool=false,
) where Fl
    m = size(sys.T, 1)
    y = Vector{Fl}(undef, n)
    alpha = Matrix{Fl}(undef, n + 1, m)
    # Sampling errors
    chol_H = sqrt(sys.H)
    chol_Q = cholesky(sys.Q)
    standard_ε = randn(n)
    standard_η = randn(n + 1, size(sys.Q, 1))

    # The first state of the simulation is the update of a_0
    alpha[1, :] = sys.T * initial_state + sys.c + sys.R * chol_Q.L * standard_η[1, :]
    # Simulate scenarios
    for t in 1:n
        y[t] = dot(sys.Z, alpha[t, :]) + sys.d + chol_H * standard_ε[t]
        alpha[t + 1, :] = sys.T * alpha[t, :] + sys.c + sys.R * chol_Q.L * standard_η[t, :]
    end

    if return_simulated_states
        return y, alpha[1:n, :]
    end
    return y
end


text:
# Inspired from https://github.com/jump-dev/JuMP.jl/blob/master/docs/make.jl

using SumOfSquares
using Documenter, Literate

const _TUTORIAL_DIR = joinpath(@__DIR__, "src", "tutorials")
const _OUTPUT_DIR   = joinpath(@__DIR__, "src", "generated")
const _TUTORIAL_SUBDIR = [
    "Getting started",
    "Polynomial Optimization",
    "Systems and Control",
    "Other Applications",
    "Noncommutative and Hermitian",
    "Sparsity",
    "Symmetry",
    "Extension",
]

function literate_directory(dir)
    output_dir = joinpath(_OUTPUT_DIR, dir)
    for filename in readdir(joinpath(_TUTORIAL_DIR, dir))
        if filename[1] == '_'
            continue
        end
        path = joinpath(_TUTORIAL_DIR, dir, filename)
        Literate.markdown(path, output_dir; documenter = true)
        Literate.notebook(path, output_dir; documenter = true)
        Literate.script(path, output_dir; documenter = true)
    end
end

literate_directory.(_TUTORIAL_SUBDIR)

makedocs(
    sitename = "SumOfSquares",
    # See https://github.com/JuliaDocs/Documenter.jl/issues/868
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
    # See https://github.com/jump-dev/JuMP.jl/issues/1576
    strict = true,
    pages = [
        "Index" => "index.md",
        "Sum-of-Squares Programming" => "sumofsquares.md",
        "Variables" => "variables.md",
        "Constraints" => "constraints.md",
        "Tutorials" => map(
            subdir ->
                subdir => map(
                    file -> joinpath("generated", subdir, file),
                    filter(
                        file -> endswith(file, ".md"),
                        sort(readdir(joinpath(_OUTPUT_DIR, subdir))),
                    ),
                ),
            _TUTORIAL_SUBDIR,
        ),
    ],
    # The following ensures that we only include the docstrings from
    # this module for functions define in Base that we overwrite.
    modules = [SumOfSquares, PolyJuMP]
)

deploydocs(
    repo   = "github.com/jump-dev/SumOfSquares.jl.git",
    push_preview = true,
)

max_stars_repo_name: blegat/SumOfSquares

max_issues_repo_name: blegat/SumOfSquares

max_forks_repo_name: blegat/SumOfSquares


text:
"""
```
initialize_state_draws(s0::Vector{Float64}, F_ϵ::Distribution, Φ::Function,
                       n_parts::Int; burn::Int = 10000, thin::Int = 5)
```

### Inputs

- `s0::Vector`: An initial guess/starting point to begin iterating the
    states forward.
- `F_ϵ::Distribution`: The shock distribution: ϵ ~ F_ϵ
- `Φ::Function`: The state transition function: s_t = Φ(s_t-1, ϵ_t)
- `n_parts::Int`: The number of particles (draws) to generate

### Keyword Arguments

- `initialize::Bool`: Flag indicating whether one is solving for incremental
    weights during the initialization of weights; default is `false`.
- `burn::Int`: The number of draws to burn in before the draws are actually
    collected. This is under the assumption that the s_t reaches its stationary
    distribution post burn-in.
- `thin::Int`: The number of draws to thin by to minimize serial correlation

### Output

- `s_init`: A matrix (# of states x # of particles) containing the initial
    draws of states to start
the tpf algorithm from.
"""
function initialize_state_draws(s0::Vector{Float64}, F_ϵ::Distribution,
                                Φ::Function, n_parts::Int;
                                burn::Int = 10000, thin::Int = 5)
    s_init = zeros(length(s0), n_parts)
    s_old  = s0
    for i in 1:(burn + thin*n_parts)
        ϵ     = rand(F_ϵ)
        s_new = Φ(s_old, ϵ)

        if i > burn && i % thin == 0
            draw_index = convert(Int, (i - burn)/thin)
            s_init[:, draw_index] = s_new
        end
        s_old = s_new
    end
    return s_init
end


text:
@testset "Generating set search" begin

    ss = RectSearchSpace(3, (0.0, 1.0))
    @test BlackBoxOptim.calc_initial_step_size(ss) == (0.5 * (1.0 - 0.0))

    ss = RectSearchSpace(3, (-1.2, 42.0))
    @test BlackBoxOptim.calc_initial_step_size(ss, 0.80) == (0.80 * (42.0 + 1.2))

end


text:
module PSIS

using LogExpFunctions: LogExpFunctions
using Printf: @sprintf
using RecipesBase: RecipesBase
using Requires: Requires
using Statistics: Statistics

export PSISResult
export psis, psis!, ess_is
export ParetoShapePlot, paretoshapeplot, paretoshapeplot!

const PLOTTING_BACKEND = Ref(:Plots)

include("utils.jl")
include("generalized_pareto.jl")
include("core.jl")
include("ess.jl")
include("recipes/definitions.jl")
include("recipes/plots.jl")

function __init__()
    Requires.@require Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" begin
        using .Makie
        include("recipes/makie.jl")
    end
    Requires.@require Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" begin
        using .Distributions: Distributions
        include("distributions.jl")
    end
end

end


text:
using Gen, Genify, GenParticleFilters, Distributions, StatsBase
using DataFrames, CSV
using ProgressMeter: @showprogress

include("model.jl")
include("utils.jl")
include("inference/proposals.jl")
include("inference/resimulation_mh.jl")
include("inference/block_mh.jl")
include("inference/single_site_mh.jl")
include("inference/basic_smc.jl")
include("inference/drift_smc.jl")
include("inference/data_driven_smc.jl")

show_plots = true
save_plots = false

run_resimulation_mh = (trace, T, N) -> begin
    observations = merge(case_count_obs(trace), location_obs(trace, frac=0.75));
    t_start = time()
    trs, scores, data = resimulation_mh(T, observations, N)
    t_stop = time()

    if show_plots || save_plots
        plot_obs(trace, [1]); plot_obs!(trs, [1], lw=1, color=:grey)
        plot!(legend=:topleft, title="Resimulation MH, N=$(N)")
        if save_plots savefig("images/resimulation_mh.png") end
        if show_plots display(plot!()) end
    end

    mean_score = logsumexp(scores) - log(N)
    return (t_stop - t_start), trs, scores, ones(N), data.β, mean_score, plot!()
end

run_block_mh = (trace, T, N) -> begin
    observations = merge(case_count_obs(trace), location_obs(trace, frac=0.75));
    t_start = time()
    trs, scores, data = block_mh(T, observations, N)
    t_stop = time()

    if show_plots || save_plots
        plot_obs(trace, [1]); plot_obs!(trs, [1], lw=1, color=:grey)
        plot!(legend=:topleft, title="Block resimulation MH, N=$(N)")
        if save_plots savefig("images/block_mh.png") end
        if show_plots display(plot!()) end
    end

    mean_score = logsumexp(scores) - log(N)
    return (t_stop - t_start), trs, scores, ones(N), data.β, mean_score, plot!()
end

run_single_site_mh = (trace, T, N) -> begin
    observations = merge(case_count_obs(trace), location_obs(trace, frac=0.75));
    t_start = time()
    trs, scores, data = single_site_mh(T, observations, N)
    t_stop = time()

    if show_plots || save_plots
        plot_obs(trace, [1]); plot_obs!(trs, [1], lw=1, color=:grey)
        plot!(legend=:topleft, title="Single site MH, N=$(N)")
        if save_plots savefig("images/single_site_mh.png") end
        if show_plots display(plot!()) end
    end

    mean_score = logsumexp(scores) - log(N)
    return (t_stop - t_start), trs, scores, ones(N), data.β, mean_score, plot!()
end

run_basic_smc = (trace, T, N) -> begin
    observations = merge(case_count_obs(trace), location_obs(trace, frac=0.75));
    t_start = time()
    pf_state = basic_smc(T, observations, N);
    t_stop = time()
    lml_est = log_ml_estimate(pf_state)
    βs = [tr[:β] for tr in pf_state.traces]
    trs, ws = get_traces(pf_state), get_norm_weights(pf_state);

    if show_plots || save_plots
        plot_obs(trace, [1]); plot_obs!(trs, ws, [1], lw=1, color=:grey)
        plot!(legend=:topleft, title="Vanilla SMC, N=$(N)")
        if save_plots savefig("images/basic_smc.png") end
        if show_plots display(plot!()) end
    end

    scores = collect(get_score.(trs))
    mean_score = logsumexp(sample(scores, StatsBase.weights(ws), N)) - log(N)
    return (t_stop - t_start), trs, scores, ws, βs, mean_score, plot!()
end

run_drift_smc = (trace, T, N) -> begin
    observations = merge(case_count_obs(trace), location_obs(trace, frac=0.75));
    t_start = time()
    pf_state = drift_smc(T, observations, N);
    t_stop = time()
    lml_est = log_ml_estimate(pf_state)
    β_hat, β_var = mean(pf_state, :β), var(pf_state, :β)
    βs = [tr[:β] for tr in pf_state.traces]
    trs, ws = get_traces(pf_state), get_norm_weights(pf_state);

    if show_plots || save_plots
        plot_obs(trace, [1]); plot_obs!(trs, ws, [1], lw=1, color=:grey)
        plot!(legend=:topleft, title="SMC with Gaussian drift rejuvenation, N=$(N)")
        if save_plots savefig("images/drift_smc.png") end
        if show_plots display(plot!()) end
    end

    scores = collect(get_score.(trs))
    mean_score = logsumexp(sample(scores, StatsBase.weights(ws), N)) - log(N)
    return (t_stop - t_start), trs, scores, ws, βs, mean_score, plot!()
end

run_data_driven_smc = (trace, T, N) -> begin
    observations = merge(case_count_obs(trace), location_obs(trace, frac=0.75));
    t_start = time()
    pf_state = data_driven_smc(T, observations, N);
    t_stop = time()
    lml_est = log_ml_estimate(pf_state)
    β_hat, β_var = mean(pf_state, :β), var(pf_state, :β)
    βs = [tr[:β] for tr in pf_state.traces]
    trs, ws = get_traces(pf_state), get_norm_weights(pf_state);

    if show_plots || save_plots
        plot_obs(trace, [1]); plot_obs!(trs, ws, [1], lw=1, color=:grey)
        plot!(legend=:topleft, title="SMC with data-driven migration proposals, N=$(N)")
        if save_plots savefig("images/data_driven_smc.png") end
        if show_plots display(plot!()) end
    end

    scores = collect(get_score.(trs))
    mean_score = logsumexp(sample(scores, StatsBase.weights(ws), N)) - log(N)
    return (t_stop - t_start), trs, scores, ws, βs, mean_score, plot!()
end

run_experiments = (trace, T, repeats) -> begin
    algs = Dict(
        :basic_smc => (run_basic_smc, 100),
        :data_driven_smc => (run_data_driven_smc, 100),
        :resimulation_mh => (run_resimulation_mh, 100),
        :block_mh => (run_block_mh, 100)
    )
    df = DataFrame(alg = Symbol[], dur = Float64[], dur_std = Float64[],
                   score = Float64[], score_std = Float64[],
                   β_hat = Float64[], β_std = Float64[], β_rmse = Float64[])
    for (name, (alg, N)) in sort(collect(algs), by=first)
        durs = Vector{Float64}(undef, repeats)
        scores = Vector{Float64}(undef, repeats)
        β_hats = Vector{Float64}(undef, repeats)
        for i in 1:repeats
            dur, _, _, ws, data, score, _ = alg(trace, T, N)
            durs[i], scores[i] = dur, score
            β_hats[i] = mean(data, StatsBase.weights(ws))
            GC.gc() # Reclaim unused memory
        end
        β_rmse = mean((β_hats .- trace[:β]).^2).^0.5
        push!(df, [name, mean(durs), std(durs), mean(scores), std(scores),
                   mean(β_hats), std(β_hats), β_rmse])
        println(df[end, :])
        GC.gc() # Reclaim unused memory
    end
    return df
end

# Generate ground truth trace.
trace, _ = generate(bayesian_sir, (50, 0.01), choicemap(:β => 0.5));
plot_obs(trace)

show_plots = false
save_plots = false

df = run_experiments(trace, 50, 10)
CSV.write("results.csv", df)

# dur, trs, scores, ws, data, plt = fill(nothing, 6); GC.gc()
# dur, trs, scores, ws, data, plt = run_resimulation_mh(trace, 50, 100);
# dur, trs, scores, ws, data, plt = run_block_mh(trace, 50, 10);
# dur, trs, scores, ws, data, plt = run_single_site_mh(trace, 50, 100);
# dur, trs, scores, ws, data, plt = run_basic_smc(trace, 50, 100);
# dur, trs, scores, ws, data, plt = run_drift_smc(trace, 50, 100);
# dur, trs, scores, ws, data, plt = run_data_driven_smc(trace, 50, 100);


text:
using Base.Test
using Graphtypes

#simple directed graph
g = simple_graph(10)
connect!(g,1,5)
@test connected(g,1,5)
@test !connected(g,5,1)
reverse!(g,1,5)
@test !connected(g,1,5)
@test connected(g,5,1)

connect!(g,9,10)
connect!(g,9,10)
connect!(g,9,7)
connect!(g,9,4)
connect!(g,5,9)
connect!(g,9,9)

@test list_outgoing(g,9) == [4, 7, 9, 10]

connect!(g,4,5)
connect!(g,4,5)
connect!(g,5,4)

@test length(list_edges(g)) == 2 * 10

@test count_connections(g,4,5) == 2
@test total_vertices(g) == 10

#undirected structure graph
r = structure_graph(4)
connect!(r, 1, 3)
connect!(r, 2, 1)
connect!(r, 4, 1)
@test connected(r, 1, 3)
@test connected(r, 3, 1)
@test total_vertices(r) == 4
@test list_outgoing(r, 1) == [2,3,4]

@test length(list_edges(r)) == 6
connect!(r, 2, 2)
lv = list_edges(r)
@test length(lv) == 2*4

max_stars_repo_name: AaltoScienceInst/pebblegame

max_issues_repo_name: AaltoScienceInst/pebblegame

max_forks_repo_name: AaltoScienceInst/pebblegame


text:
@testset "basic session rendering" begin
    session = Session()
    obs1 = Observable(1)
    TestLib = JSServe.Dependency(:Test, ["bla.js"])
    dom = DOM.div(test=obs1, onload=js"$(TestLib).lol()")
    jsrender(session, dom)
    @test haskey(session.observables, obs1.id)
    # Obs registered correctly
    @test session.observables[obs1.id][1] # is registered
    @test obs1.listeners[1] isa JSServe.JSUpdateObservable
    @test session.message_queue[1][:msg_type] == JSServe.RegisterObservable
    @test session.message_queue[1][:id] == obs1.id

    # Will deregisters correctly on session close
    @test session.deregister_callbacks[1].observable === obs1
    @test length(session.dependencies) == 1
    @test first(session.dependencies).local_path == joinpath(pwd(), "bla.js")

    # onjs callback to update attributes
    @test session.message_queue[2][:msg_type] == JSServe.OnjsCallback
    @test occursin("JSServe.update_node_attribute", string(session.message_queue[2][:payload]))
end

@testset "Page rendering & cleanup" begin
    io = IOBuffer()
    open_session = JSServe.Session(Base.RefValue{Union{Nothing, JSServe.WebSocket, IOBuffer}}(io))
    put!(open_session.js_fully_loaded, true)
    page = JSServe.Page(session=open_session)
    @test isopen(page.session)
    rslider = JSServe.RangeSlider(1:100; value=[10, 80])

    for obs_field in (:range, :value, :connect, :orientation, :tooltips, :ticks)
        obs = getfield(rslider, obs_field)
        @test isempty(obs.listeners)
    end

    sliderapp = JSServe.App(rslider)
    html = JSServe.show_in_page(page, sliderapp)
    @test length(page.child_sessions) == 1
    c_session = first(page.child_sessions)[2]
    # ID of div correctly set to child session
    @test occursin("<div id=\"$(c_session.id)\"", html)

    @testset "assets" begin
        dependencies = getfield.(c_session.dependencies, :local_path)
        @test length(dependencies) == 2
        @test JSServe.noUiSlider.assets[1].local_path in dependencies
        @test JSServe.noUiSlider.assets[2].local_path in dependencies
    end

    for obs_field in (:range, :value, :connect, :orientation, :tooltips, :ticks)
        obs = getfield(rslider, obs_field)
        @test length(obs.listeners) == 1
    end

    @testset "closing" begin
        c_session.connection[] = nothing
        close(c_session)
        delete!(page.child_sessions, c_session.id)
        # All listeners should get disconnected!
        for obs_field in (:range, :value, :connect, :orientation, :tooltips, :ticks)
            obs = getfield(rslider, obs_field)
            @test isempty(obs.listeners)
        end
        # all obs need to get removed
        @test isempty(open_session.observables)
    end

    @testset "dependency second include" begin
        html2 = JSServe.show_in_page(page, sliderapp)
        # Test that dependencies only get loaded one time!
        @test !occursin("-nouislider.min.js", html2)
        @test !occursin("nouislider.min.css", html2)

        id, session = first(page.child_sessions)
        session.connection[] = nothing
        close(session)
        delete!(page.child_sessions, id)
    end

    @test isempty(page.session.observables)
    @test isempty(page.session.unique_object_cache)
    # Wrap in func, to make sure we dont capture any reference globally
    function gc_avoid_func()
        dublicate = rand(1000, 1000)
        dub_ref = JSServe.pointer_identity(dublicate)
        dom = JSServe.App() do session
            JSServe.evaljs(session, js"$(dublicate); $(dublicate)")
            JSServe.evaljs(session, js"$(dublicate); $(dublicate)")
            return JSServe.DOM.div()
        end
        html = JSServe.show_in_page(page, dom)
        id, session = first(page.child_sessions)
        # There should only be one obs registered, which is the init obs
        (id, (reg, init_obs)) = first(open_session.observables)

        @test init_obs[] isa Bool
        @test !init_obs[] # shouldn't be initialized yet
        @test isopen(session)
        session.connection[] = nothing
        @test !isopen(session)
        # there should be 3 messages in here
        # 2 for evaljs, one for registering the int_obs
        @test length(session.message_queue) == 3
        @test isempty(open_session.unique_object_cache)
        init_obs[] = true
        @test !isempty(session.message_queue)
        # There should only be two messages now, 1 fused message that sends all messages in one go
        # and one to update obs
        @test length(session.message_queue) == 2
        @test session.message_queue[1][:msg_type] == JSServe.FusedMessage

        session.connection[] = open_session.connection[]
        messages = copy(session.message_queue)
        empty!(session.message_queue)
        for msg in messages
            JSServe.send(session, msg)
        end
        @test !isempty(session.unique_object_cache)
        @test haskey(session.unique_object_cache, dub_ref)
        obj = session.unique_object_cache[dub_ref]
        @test obj.value === dublicate
        session.connection[] = nothing
        close(session)
        dublicate = nothing
        obj = nothing
        @test isempty(session.message_queue)
        @test isempty(session.deregister_callbacks)

        return dub_ref
    end

    dub_ref = gc_avoid_func()
    # Ok, somehow GC is bitchy - but after some time it works
    function test_ref()
        tstart = time()
        while time() - tstart < 30
            id, obj = first(open_session.unique_object_cache)
            isnothing(obj.value) && return true
            yield()
            GC.gc(true)
        end
        return false
    end

    # @test test_ref()
    # # cache get cleaned up on any serialization event.
    # # the object, doesn't really matter
    # obj = Dict("helo" => 22)
    # binary = JSServe.serialize_binary(open_session, obj)
    # data_unpacked = JSServe.MsgPack.unpack(transcode(JSServe.GzipDecompressor, binary))
    # @show data_unpacked
    # @test isempty(data_unpacked["update_cache"]["to_register"])
    # @test data_unpacked["update_cache"]["to_remove"][1] == dub_ref
end