text:
using HDF5, JLD

include("CenterOfMass.jl")
include("setup_parameters.jl")

plt[:close]("all")

actuator_limit = 4.0
max_robots = 4
csqmi_ratios = [1.0; 2.0; 5.0; 10.0]
#csqmi_ratios = [1.0; 10.0]
#n_trial = 10
n_trial = 20
#n_trial = 5
n_measurement = 25

prior = initialize_prior(circle_ps, resolution, interior_q, masses)

belief = deepcopy(prior)

type TrialData
  belief::Histogram
  estimate
  error
  applied_f
  f_hat
  csqmi
  robots_measurement
  robots
end

data_array = Array{TrialData}(length(csqmi_ratios), n_trial, n_measurement)

rand_com() = r_attach * rand_in_circle()[1:2] # hack
rand_mass() = min_mass + rand() * (max_mass - min_mass)

rand_parameters() = [rand_com(); rand_mass()]

thetas = hcat(map(x->rand_parameters(), 1:n_trial)...)

for (ii, csqmi_ratio) = enumerate(csqmi_ratios)
  for jj = 1:n_trial
    theta = thetas[:,jj]
    gravity_w = theta[3]*g*point_to_wrench(theta[1:2])
    belief = deepcopy(prior)
    robots = Array{Int64,1}()
    for kk = 1:n_measurement
      println("\n$ii,$jj: Measurement: $(kk), robots: $(robots)\n")

      @time @show csqmi_current, robots_current = maximize_csqmi_available_robots(
        robots, belief, circle_ws, attachment_ws, sigma, interior_q, actuator_limit)

      csqmi = csqmi_current
      robots_measurement = robots_current

      if length(robots) < max_robots
        @time @show csqmi_additional, robots_additional, additional_robot =
          maximize_csqmi_additional_robot(robots, belief,
          circle_ws, attachment_ws, sigma, interior_q, actuator_limit)

        if length(robots) == 0 || csqmi_additional / csqmi_current > csqmi_ratio
          println("\nAdding robot $(additional_robot), csqmi: $(csqmi_additional) vs: $(csqmi_current)\n")

          csqmi = csqmi_additional
          robots_measurement = robots_additional
          push!(robots, additional_robot)
        end
      end

      applied_w = mean(attachment_ws[robots_measurement])
      total_actuator_limit = length(robots_measurement) * actuator_limit

      applied_f = critical_force_iterative(applied_w, gravity_w,
      circle_ws, total_actuator_limit)

      f_hat = applied_f + sigma * randn()

      println("\nRobots measurement: $(robots_measurement)")
      println("Measurement: $(applied_f)/$(f_hat), total limit: $(total_actuator_limit)\n")

      critical_forces = get_critical_values(circle_ws, applied_w, belief,
        interior_q, total_actuator_limit)

      update_prior!(belief, critical_forces, f_hat, sigma)

      estimate = weighted_average(belief)

      error = normalized_error(estimate - theta)
      println("Estimate: $(estimate), actual: $(theta), normalized error: $(error)")

      data = TrialData(deepcopy(belief), deepcopy(estimate), error, applied_f, f_hat, csqmi,
        deepcopy(robots_measurement), deepcopy(robots))
      data_array[ii,jj,kk] = data
    end
  end
end

@save "multi_robot/data" data_array csqmi_ratios thetas actuator_limit
Void

max_stars_repo_name: mcorah/CenterOfMass

max_issues_repo_name: mcorah/CenterOfMass

max_forks_repo_name: mcorah/CenterOfMass


text:
export Subsampling_1D
###### define the subsampling function
using Random
# regular subsample
function Regular_subsample(ns::Int64,p::Int64)
    # ns: the number of receivers
    #   p=4 --> 75 % subsampling
    #   p=2 --> 50 % subsampling
    selected_locs = collect(1:p:ns);
    return selected_locs
end

# uniform subsampling % white noise subsample
function Uniform_subsample(ns::Int64,p::Int64)
    # ns: the number of receivers
    # p
    ratio = 1/p;
    Num = Int(round(ns*ratio));
	rand_locs = randperm(ns);
    selected_locs = rand_locs[1:Num];
    # return the subsampling locations for these data
    return sort(selected_locs)
end

#Optimal jittered subsampling %kind of blue noise subsample
function Optimal_Jittered_subsample(ns::Int64, p::Int64)
    # ns: the number of receivers
    # p
    b = collect(1:p)
    ind = []
    for i = 1:Int(round(ns/p))
        push!(ind,rand(b)+(i-1)*p)
    end
    return ind
end

#Blue noise subsampling w/ Farthest Point sampling
#Eldar, Yuval, et al. "The farthest point strategy for progressive image sampling." IEEE Transactions on Image Processing 6.9 (1997): 1305-1315.
function Blue_Noise_subsample(ns::Int64, p::Int64)
    # ns: the number of receivers
    # p
    ind = []
    #step 1: randomly choose one point
    push!(ind, randperm(ns)[1])
    ratio = 1/p;
    Num = Int(round(ns*ratio));
    for i = 1:Num-1
        #step 2: randomly choose a set of points
        cand_ind = rand(filter(x->!(x in ind), collect(1:ns)), p);
        Min_ind = []
        #step 3: calculate the smallest distance between the random set and the existed subpoints
        for j = 1:length(cand_ind)
            push!(Min_ind,minimum(abs.(cand_ind[j] .- ind)));
        end
        #step 4: find the largest distance among the random set
        Idx = findall(x-> x == maximum(Min_ind), Min_ind);
        #step 5: add one point
        push!(ind, cand_ind[Idx][1])
    end
    return sort(unique(ind))
end

function Subsampling_1D(ns::Int64, p::Int64; SubMethod = "Optimal_Jittered")
    #Regular: Period subsampling
    #Uniform: Uniform random subsampling (White noise subsampling)
    #Optimal_Jittered: Optimal Jittered subsampling
    #Blue_Noise: Blue noise subsampling w/ Farthest Point sampling
    SubMethod == "Regular" && return Regular_subsample(ns,p)
    SubMethod == "Uniform" && return Uniform_subsample(ns,p)
    SubMethod == "Blue_Noise" && return Blue_Noise_subsample(ns,p)
    return Optimal_Jittered_subsample(ns, p)
end

ext: jl


text:
module AccretionGeometry

import GeodesicTracer: tracegeodesics, DiscreteCallback, terminate!
import GeodesicRendering: rendergeodesics, prerendergeodesics
import Base: in
import RecursiveArrayTools: ArrayPartition
import GeometryBasics
import LinearAlgebra: ×, ⋅

using GeodesicBase
using StaticArrays

include("geometry.jl")
include("meshes.jl")
include("intersections.jl")
include("discs.jl")

function tracegeodesics(
    m::AbstractMetricParams{T},
    init_positions,
    init_velocities,
    time_domain::Tuple{T,T},
    accretion_geometry;
    callback = nothing,
    kwargs...,
) where {T}
    cbs = add_collision_callback(callback, accretion_geometry)
    tracegeodesics(
        m,
        init_positions,
        init_velocities,
        time_domain;
        callback = cbs,
        kwargs...,
    )
end

function rendergeodesics(
    m::AbstractMetricParams{T},
    init_pos,
    max_time::T,
    accretion_geometry;
    callback = nothing,
    kwargs...,
) where {T}
    cbs = add_collision_callback(callback, accretion_geometry)
    rendergeodesics(m, init_pos, max_time; callback = cbs, kwargs...)
end

function prerendergeodesics(
    m::AbstractMetricParams{T},
    init_pos,
    max_time::T,
    accretion_geometry;
    callback = nothing,
    kwargs...,
) where {T}
    cbs = add_collision_callback(callback, accretion_geometry)
    prerendergeodesics(m, init_pos, max_time; callback = cbs, kwargs...)
end

add_collision_callback(::Nothing, accretion_geometry) =
    build_collision_callback(accretion_geometry)
add_collision_callback(callback::Base.AbstractVecOrTuple, accretion_geometry) =
    (callback..., build_collision_callback(accretion_geometry))
add_collision_callback(callback, accretion_geometry) =
    (callback, build_collision_callback(accretion_geometry))

function build_collision_callback(geometry::AbstractAccretionGeometry{T}) where {T}
    DiscreteCallback(collision_callback(geometry), i -> terminate!(i, :Intersected))
end

collision_callback(m::AbstractAccretionGeometry{T}) where {T} =
    (u, λ, integrator) -> intersects_geometry(m, line_element(u, integrator))

export tracegeodesics, rendergeodesics, prerendergeodesics

end # module


text:
using Test, Batched
import Base.Broadcast: BroadcastStyle, combine_styles, broadcasted

A = BatchedArray(rand(2, 3, 10));
B = BatchedArray(rand(3, 2, 10));
C = rand(2, 2)

@test combine_styles(A, B) == Batched.BatchedArrayStyle{2, 1}()
@test combine_styles(A, B, C) == Batched.BatchedArrayStyle{2, 1}()

bc = broadcasted(*, A, B)

import Batched: combine_batch_axes

@test combine_batch_axes(A, B) == (Base.OneTo(10), ) # A * B
@test combine_batch_axes(A, B, C) == (Base.OneTo(10), ) # A * B * C

import Batched: combine_element_axes

combine_element_axes(*, A, B) == (Base.OneTo(2), Base.OneTo(2))

copy(bc)

A .* B

# using LinearAlgebra
# tr.(A)


text:
include("gen_code_mem.jl")
include("gen_code_snippets.jl")

include("multilincomb.jl");

include("gen_c_code.jl")
include("gen_julia_code.jl")
include("gen_matlab_code.jl")

export gen_code

# Every language needs:
# comment(::Lang,s)
# slotname(::Lang,i) #
# assign_coeff(::Lang,v,i)
# function_definition(::Lang,graph,T,funname)
# function_init(lang::Lang,T,mem,graph)
# init_mem(lang::Lang,max_nof_nodes)
# function_end(lang::Lang,graph,mem)
# execute_operation!(lang::Lang,T,graph,node,
#                    dealloc_list,   mem)

# Fallback. Default to no main function
function gen_main(lang, T, fname, funname)
    return init_code(lang)
end

function preprocess_codegen(graph, lang)
    return graph  # Fallback to no preprocessing
end

"""
    gen_code(
        fname,
        graph;
        priohelp = Dict{Symbol,Float64}(),
        lang = LangJulia(),
        funname = "dummy",
        precomputed_nodes = [:A],
    )

Generates the code for the `graph` in the language specified in `lang` and
writes it into the file `fname`. The string `funname` is the function name.
Topological order of the nodes is comptued using [`get_topo_order`](@ref) and
`priohelp` can be used to influence the order. The nodes listed in
`precomputed_nodes` are viewed as inputs, and code to compute these nodes are
not computed.

Currently supported languages: [`LangC_MKL`](@ref), [`LangC_OpenBLAS`](@ref),
[`LangJulia`](@ref), [`LangMatlab`](@ref).
"""
function gen_code(
    fname,
    graph;
    priohelp = Dict{Symbol,Float64}(),
    lang = LangJulia(),
    funname = "dummy",
    precomputed_nodes = [:A],
)
    # Make dispatch possible for lang
    return _gen_code(fname, graph, lang, priohelp, funname, precomputed_nodes)
end

# Most gen code calls will call this. Can be overloaded with lang
function _gen_code(fname, graph, lang, priohelp, funname, precomputed_nodes)

    # Error if graph is trivial (no operations) or has trivial nodes.
    if isempty(graph.operations)
        error("Unable to generate code for graphs without operations.")
    end
    if has_trivial_nodes(graph)
        error("Please run compress_graph!() on the graph first.")
    end
    T = eltype(eltype(typeof(graph.coeffs.vals)))

    if (fname isa String)
        file = open(abspath(fname), "w+")
    else
        # Lazy: Print out to stdout if no filename
        file = Base.stdout
    end

    graph = preprocess_codegen(graph, lang)

    (order, can_be_deallocated, max_nof_slots) =
        get_topo_order(graph; priohelp = priohelp)

    # max_nof_slots is the path width which gives
    # a bound on the number memory slots needed.

    println(
        file,
        to_string(
            function_definition(lang, graph, T, funname, precomputed_nodes),
        ),
    )

    # We do a double sweep in order to determine exactly how many memory slots
    # are needed. The first sweep we carry out all operations but store only the
    # maximum number of memory slots needed. The second sweep generates the
    # code.

    mem = init_mem(lang, max_nof_slots + 3, precomputed_nodes)
    function_init(lang, T, mem, graph, precomputed_nodes)

    # Sweep 1: Determine exactly the number of slots needed
    nof_slots = 0
    for (i, node) in enumerate(order)
        if (node in precomputed_nodes) # Nothing to do for precomputed nodes
            continue
        end

        (exec_code, result_variable) =
            execute_operation!(lang, T, graph, node, can_be_deallocated[i], mem)

        # How many slots needed to reach this point
        if (!isnothing(findlast(mem.slots .!= :Free)))
            nof_slots = max(nof_slots, findlast(mem.slots .!= :Free))
        end
    end

    # Sweep 2:
    mem = init_mem(lang, nof_slots, precomputed_nodes)
    function_init_code = function_init(lang, T, mem, graph, precomputed_nodes)
    push_comment!(
        function_init_code,
        "Computation order: " * join(string.(order), " "),
    )
    println(file, to_string(function_init_code))

    for (i, node) in enumerate(order)
        if (node in precomputed_nodes)  # Nothing to do for precomputed nodes
            continue
        end
        (exec_code, result_variable) =
            execute_operation!(lang, T, graph, node, can_be_deallocated[i], mem)
        println(file, to_string(exec_code))
    end
    println(file, to_string(function_end(lang, graph, mem)))

    # Generate main function, if necessary.
    exec_code = gen_main(lang, T, fname, funname)
    println(file, to_string(exec_code))

    if (fname isa String)
        close(file)
    end
end


text:
# """
# k: an index from subspace
# indices: array of indices from comspace
# """
# @inline function subspace_mul_kernel_diagonal!(S::AbstractVector{T}, indices, U, k::Int, offset::Int) where T
#     @inbounds for i in axes(U, 1)
#         idx_i = k + indices[i] + offset
#         S[idx_i] = U[i, i] * S[idx_i]
#     end
#     return
# end

# @inline function subspace_mul_kernel_diagonal_4x4!(S::AbstractVector{T}, indices, U, k::Int, offset::Int) where T
#     @inbounds begin
#         Base.Cartesian.@nexprs 4 i -> begin
#             idx_i = k + indices[i] + offset
#             S[idx_i] = U[i, i] * S[idx_i]
#         end
#     end
#     return
# end


text:
export GRDstack, GRDarray

const GRD_INDEX_ORDER = ForwardIndex()
const GRD_X_ARRAY = ForwardArray()
const GRD_Y_ARRAY = ReverseArray()
const GRD_BAND_ARRAY = ForwardArray()
const GRD_X_RELATION = ForwardRelation()
const GRD_Y_RELATION = ReverseRelation()
const GRD_BAND_RELATION= ForwardRelation()

const GRD_DATATYPE_TRANSLATION = Dict{String, DataType}(
    "LOG1S" => Bool,
    "INT1S" => Int8,
    "INT2S" => Int16,
    "INT4S" => Int32,
    "INT8S" => Int64,
    "INT1U" => UInt8,
    "INT2U" => UInt16,
    "INT4U" => UInt32,
    "INT8U" => UInt64,
    "FLT4S" => Float32,
    "FLT8S" => Float64
)
const REVGRDfile_DATATYPE_TRANSLATION =
    Dict{DataType, String}(v => k for (k,v) in GRD_DATATYPE_TRANSLATION)

# GRD attributes wrapper. Only used during file load, for dispatch.
struct GRDattrib{T,F}
    filename::F
    attrib::Dict{String,String}
    write::Bool
end
function GRDattrib(filename::AbstractString; write=false)
    filename = first(splitext(filename))
    lines = readlines(filename * ".grd")
    entries = filter!(x -> !isempty(x) && !(x[1] == '['), lines)
    matches = (match(r"([^=]+)=(.*)", st) for st in entries)
    captures = (string.(strip.(m.captures[1:2])) for m in matches)
    pairs = map(c -> Pair(c[1], c[2]), captures)
    attrib = Dict(pairs...)
    T = GRD_DATATYPE_TRANSLATION[attrib["datatype"]]
    GRDattrib{T,typeof(filename)}(filename, attrib, write)
end

attrib(grd::GRDattrib) = grd.attrib
filename(grd::GRDattrib) = grd.filename
filekey(grd::GRDattrib, key::Nothing) = get(attrib(grd), "layername", Symbol(""))

function DD.dims(grd::GRDattrib, crs=nothing, mappedcrs=nothing)
    attrib = grd.attrib
    crs = crs isa Nothing ? ProjString(attrib["projection"]) : crs

    ncols, nrows, nbands = size(grd)

    xbounds = parse.(Float64, (attrib["xmin"], attrib["xmax"]))
    ybounds = parse.(Float64, (attrib["ymin"], attrib["ymax"]))

    # Always intervals
    xspan = (xbounds[2] - xbounds[1]) / ncols
    yspan = (ybounds[2] - ybounds[1]) / nrows

    # Not fully implemented yet
    xy_metadata = Metadata{GRDfile}(Dict())

    xmode = Projected(
        order=Ordered(GRD_INDEX_ORDER, GRD_X_ARRAY, GRD_X_RELATION),
        span=Regular(xspan),
        sampling=Intervals(Start()),
        crs=crs,
        mappedcrs=mappedcrs,
    )
    ymode = Projected(
        order=Ordered(GRD_INDEX_ORDER, GRD_Y_ARRAY, GRD_Y_RELATION),
        span=Regular(yspan),
        sampling=Intervals(Start()),
        crs=crs,
        mappedcrs=mappedcrs,
    )
    x = X(LinRange(xbounds[1], xbounds[2] - xspan, ncols), xmode, xy_metadata)
    y = Y(LinRange(ybounds[1], ybounds[2] - yspan, nrows), ymode, xy_metadata)
    band = Band(1:nbands; mode=Categorical(Ordered()))
    x, y, band
end

DD.name(grd::GRDattrib) = Symbol(get(grd.attrib, "layername", ""))

function DD.metadata(grd::GRDattrib, args...)
    metadata = Metadata{GRDfile}()
    for key in ("creator", "created", "history")
        val = get(grd.attrib, key, "")
        if val != ""
            metadata[key] = val
        end
    end
    metadata
end

function missingval(grd::GRDattrib{T}) where T
    if haskey(grd.attrib, "nodatavalue")
        ndv = grd.attrib["nodatavalue"]
        ndv === "nothing" && return nothing
        try
            return parse(T, ndv)
        catch
            @warn "nodatavalue $(ndv) is not convertible to data type $T. `missingval` set to `nothing`."
            return nothing
        end
    else
        return nothing
    end
end

Base.eltype(::GRDattrib{T}) where T = T

function Base.size(grd::GRDattrib)
    ncols = parse(Int, grd.attrib["ncols"])
    nrows = parse(Int, grd.attrib["nrows"])
    nbands = parse(Int, grd.attrib["nbands"])
    # The order is backwards to jula array order
    ncols, nrows, nbands
end

Base.Array(grd::GRDattrib) = _mmapgrd(Array, grd)

# Array ########################################################################

@deprecate GRDarray(args...; kw...) GeoArray(args...; source=GRDfile, kw...)

function FileArray(grd::GRDattrib, filename=filename(grd); kw...)
    filename = first(splitext(filename))
    size_ = size(grd)
    eachchunk = DiskArrays.GridChunks(size_, size_)
    haschunks = DiskArrays.Unchunked()
    T = eltype(grd)
    N = length(size_)
    FileArray{GRDfile,T,N}(filename, size_; eachchunk, haschunks, kw...)
end

# Base methods

"""
    Base.write(filename::AbstractString, ::Type{GRDfile}, s::AbstractGeoArray)

Write a `GeoArray` to a .grd file with a .gri header file.
The extension of `filename` will be ignored.

Returns `filename`.
"""
function Base.write(filename::String, ::Type{GRDfile}, A::AbstractGeoArray)
    A = maybe_typemin_as_missingval(filename, A)
    if hasdim(A, Band)
        correctedA = permutedims(A, (X, Y, Band)) |>
            a -> reorder(a, GRD_INDEX_ORDER) |>
            a -> reorder(a, (X(GRD_X_RELATION), Y(GRD_Y_RELATION), Band(GRD_BAND_RELATION)))
        checkarrayorder(correctedA, (GRD_X_ARRAY, GRD_Y_ARRAY, GRD_BAND_ARRAY))
        nbands = length(val(dims(correctedA, Band)))
    else
        correctedA = permutedims(A, (X, Y)) |>
            a -> reorder(a, GRD_INDEX_ORDER) |>
            a -> reorder(a, (X(GRD_X_RELATION), Y(GRD_Y_RELATION)))
            checkarrayorder(correctedA, (GRD_X_ARRAY, GRD_Y_ARRAY))
        nbands = 1
    end
    # Remove extension
    filename = splitext(filename)[1]
    minvalue = minimum(filter(x -> x !== missingval(A), parent(A)))
    maxvalue = maximum(filter(x -> x !== missingval(A), parent(A)))
    _write_grd(filename, eltype(A), dims(A), missingval(A), minvalue, maxvalue, name(A))

    # Data: gri file
    open(filename * ".gri", write=true) do IO
        write(IO, parent(correctedA))
    end

    return filename * ".grd"
end

function _write_grd(filename, T, dims, missingval, minvalue, maxvalue, name)
    filename = splitext(filename)[1]

    x, y = map(DD.dims(dims, (X(), Y()))) do d
        convertmode(Projected, d)
    end

    nbands = hasdim(dims, Band) ? length(DD.dims(dims, Band)) : 1
    ncols, nrows = length(x), length(y)
    xmin, xmax = bounds(x)
    ymin, ymax = bounds(y)
    proj = convert(String, convert(ProjString, crs(x)))
    datatype = REVGRDfile_DATATYPE_TRANSLATION[T]
    nodatavalue = missingval

    # Metadata: grd file
    open(filename * ".grd"; write=true) do IO
        write(IO,
            """
            [general]
            creator=GeoData.jl
            created= $(string(now()))
            [georeference]
            nrows= $nrows
            ncols= $ncols
            xmin= $xmin
            ymin= $ymin
            xmax= $xmax
            ymax= $ymax
            projection= $proj
            [data]
            datatype= $datatype
            nodatavalue= $nodatavalue
            byteorder= little
            nbands= $nbands
            minvalue= $minvalue
            maxvalue= $maxvalue
            [description]
            layername= $name
            """
        )
    end
end

function create(filename, ::Type{GRDfile}, T::Type, dims::DD.DimTuple;
    name="layer", metadata=nothing, missingval=nothing, keys=(name,),
)
    # Remove extension
    basename = splitext(filename)[1]
    minvalue = maxvalue = zero(T)
    sze = map(length, DD.dims(dims, (XDim, YDim, Band)))

    # Metadata: grd file
    _write_grd(basename, T, dims, missingval, minvalue, maxvalue, name)

    # Data: gri file
    open(basename * ".gri", write=true) do IO
        write(IO, FillArrays.Zeros(sze))
    end
    return filename
end

# AbstractGeoStack methods

@deprecate GRDstack(args...; kw...) GeoStack(args...; source=GRDfile, kw...)

# Custom `open` because the data and metadata objects are separate
# Here we _mmapgrd instead of `_open`
function Base.open(f::Function, A::FileArray{GRDfile}, key...; write=A.write)
    _mmapgrd(mm -> f(GeoDiskArray{GRDfile}(mm, A.eachchunk, A.haschunks)), A; write)
end

_open(f, ::Type{GRDfile}, filename; key=nothing, write=false) = f(GRDattrib(filename; write))

haslayers(::Type{GRDfile}) = false

# Utils ########################################################################

function _mmapgrd(f, x::Union{FileArray,GRDattrib}; kw...)
    _mmapgrd(f, filename(x), eltype(x), size(x); kw...)
end
function _mmapgrd(f, filename::AbstractString, T::Type, size::Tuple; write=false)
    arg = write ? "r+" : "r"
    basename = splitext(filename)[1]
    open(basename * ".gri", arg) do io
        mmap = Mmap.mmap(io, Array{T,length(size)}, size)
        output = f(mmap)
        close(io)
        output
    end
end

# precompilation

T = UInt16
for T in (Any, UInt8, UInt16, Int16, UInt32, Int32, Int64, Float32, Float64)
    precompile(GRDattrib, (String,))
    DS = GeoData.GRDattrib{T,String}
    precompile(crs, (DS,))
    precompile(GeoData.FileArray, (DS, String))
    precompile(dims, (DS,))
    precompile(dims, (DS,WellKnownText{GeoFormatTypes.CRS,String},Nothing))
    precompile(dims, (DS,WellKnownText{GeoFormatTypes.CRS,String},EPSG))
    precompile(dims, (DS,WellKnownText{GeoFormatTypes.CRS,String},ProjString))
    precompile(dims, (DS,WellKnownText{GeoFormatTypes.CRS,String},WellKnownText{GeoFormatTypes.CRS,String}))
    precompile(metadata, (DS, ))
    precompile(metadata, (DS, Symbol))
    precompile(missingval, (DS,))
    precompile(GeoArray, (DS, String, Nothing))
    precompile(GeoArray, (DS, String, Symbol))
end


text:
export bmitest

"""
    Test if discrete sets are independent.

    test(X, Y) -> p

    Computes the Independence Test by Margaritis and Thrun for X and Y.
    The test returns the probability of X and Y being independent.
"""
function bmitest(X::Vector{Int}, Y::Vector{Int}; α = 1.0)

    uniqueX = unique(X)
    uniqueY = unique(Y)

    if any(uniqueX .≠ 1:length(uniqueX))
        t = zeros(Int, size(X));
        for k = 1:length(uniqueX)
            t[X .== uniqueX[k]] .= Int(k)
        end
        X = copy(t)
    end

    if any(uniqueY .≠ 1:length(uniqueY))
        t = zeros(Int, size(Y));
        for k = 1:length(uniqueY)
            t[Y .== uniqueY[k]] .= Int(k)
        end
        Y = copy(t)
    end

    L  = length(X)
    NX = length(uniqueX)
    NY = length(uniqueY)

    hX = counts(X, 0:NX)
    hY = counts(Y, 0:NY)
    h  = counts((X.-1)*NY .+ Y, 0:(NX*NY))

    PRD = lgamma(α*NX*NY) - lgamma(α*NX*NY + L) + sum( lgamma.(h .+ α) ) - NX*NY*lgamma(α)
    PRI = lgamma(α*NX) - lgamma(α*NX + L) + sum( lgamma.(hX .+ α) ) - NX * lgamma(α) +
            lgamma(α*NY) - lgamma(α*NY + L) + sum( lgamma.(hY .+ α) ) - NY * lgamma(α)

    m = max(PRD, PRI)

    logP = PRI - m - log(exp(PRD-m) + exp(PRI-m))
    return exp(logP)
end


text:
using CovarianceMatrices
using Statistics
using CSV
using BenchmarkTools

X = randn(10000,20);
k = BartlettKernel(prewhiten=true)
cfg = CovarianceMatrices.HACConfig(X, k)
CovarianceMatrices.variance(X, k, cfg)

@benchmark CovarianceMatrices.variance(X, k, cfg, calculatechol=false)


text:
module Normalization
export
norm_libsize
# norm_deseq

function norm_libsize(X)
    lbsz=sum(X,dims=1)
    X=(X./lbsz)*1e4;
    return X
end

# Not sure how to include another function in the same package -- pls help!
#
# include("../differentialexpression/differentialexpression.jl")
# using .DifferentialExpression

# function norm_deseq(X)
#    X = DifferentialExpression.median_of_ratios_normalization(X);
#    return X
# end

end


text:
using CLIMAParameters
using CLIMAParameters.Planet: planet_radius
using ..Diagnostics
using ..Mesh.Interpolation

"""
    DiagnosticsConfiguration

Container for all the `DiagnosticsGroup`s to be used for a simulation.
"""
mutable struct DiagnosticsConfiguration
    groups::Array{DiagnosticsGroup, 1}

    DiagnosticsConfiguration(
        groups::Array{DG, 1},
    ) where {DG <: DiagnosticsGroup} = new(groups)
end

function InterpolationConfiguration(
    ::StackedBrickTopology,
    driver_config,
    boundaries,
    axes,
)
    grid = driver_config.grid
    return InterpolationBrick(grid, boundaries, axes[1], axes[2], axes[3])
end

function InterpolationConfiguration(
    ::StackedCubedSphereTopology,
    driver_config,
    boundaries,
    axes;
    nr_toler = nothing,
)
    grid = driver_config.grid
    info = driver_config.config_info
    return InterpolationCubedSphere(
        grid,
        info.vert_range,
        info.nelem_horz,
        axes[1],
        axes[2],
        axes[3];
        nr_toler = nr_toler,
    )
end

"""
    InterpolationConfiguration(
        driver_config::DriverConfiguration,
        boundaries::Array,
        resolution = nothing;
        axes = nothing;
    )

Creates an `InterpolationTopology` (either an `InterpolationBrick` or an
`InterpolationCubedSphere`) to be used with a `DiagnosticsGroup`. Either
`resolution` is specified, in which case the axes are set up with
equi-distant points, or the `axes` may be specified directly (in
lat/lon/lvl or x/y/z order).
"""
function InterpolationConfiguration(
    driver_config::DriverConfiguration,
    boundaries::Array,
    resolution = nothing;
    axes = nothing,
)
    @assert isnothing(resolution) || isnothing(axes)
    if isnothing(axes)
        axes = (
            collect(range(
                boundaries[1, 1],
                boundaries[2, 1],
                step = resolution[1],
            )),
            collect(range(
                boundaries[1, 2],
                boundaries[2, 2],
                step = resolution[2],
            )),
            collect(range(
                boundaries[1, 3],
                boundaries[2, 3],
                step = resolution[3],
            )),
        )
    end

    return InterpolationConfiguration(
        driver_config.grid.topology,
        driver_config,
        boundaries,
        axes,
    )
end


text:
using MathProgBase
using Cbc
using Test

function knapsack()
    f = Float64[3, 2, 3, 1, 2]
    A = [7. 2. 9. 3. 1.]
    capacity = 10.
    solution = mixintprog(-f, A, '<', capacity, :Int, 0, 1, CbcSolver(logLevel=0))
    println("Solution status: ", solution.status)
    println("Optimal value: ", solution.objval)
    print("Solution vector: ")
    show(solution.sol)
    println()
    @test solution.objval == -7
end

knapsack()


text:
"""
    CountingExperiment()

Building a Poisson counting model of n>1 components, each
added individual, including uncertainties.

# Example
```jldoctest
julia> m = CountingExperiment()

julia> add_component!(m, "Signal", 20.0)
julia> add_component!(m, "Background", 30.0; σ=5.0)

julia> set_counts!(m, 55)
julia> results = minimize!(m)
```
"""
mutable struct CountingExperiment <: Model
  params
  pdflist
  nll::NLogLikelihood
  counts::Constant
  function CountingExperiment()
    new( [], [], NLogLikelihood(),
        Constant("counts", 0.0) )
  end
end

function add_component!( m::CountingExperiment, name::String, μ; σ=Inf )
  p = Parameter( name; init=μ, bounds=(0.0, Inf) )
  push!(m.params, p)
  if σ != Inf
    a = Constant( name*"_mu", μ )
    b = Constant( name*"_sig", σ )
    nlp = NLogPDF( "lognormal", p, a, b )
    push!(m.pdflist, nlp)
  end
end

function set_counts!( m::CountingExperiment, counts::Number )
  m.counts.init = Float64(counts)
end

function minimize!( m::CountingExperiment )
  poisson_pdf = NLogPDF("logpoisson", m.counts, (m.params)...)
  likelihood = NLogLikelihood([m.pdflist..., poisson_pdf])
  add_likelihood!( m, likelihood )
  optimize_model!( m, likelihood )
  #optimize_model!( m, likelihood;
  #                 lower_bounds=[0.0 for a in 1:length(m.params)] )
end


text:
#define some function stubs
#users will need to define these to run their stuff, but
#we need the names in the namespace

nnAct(z) = error("You must define the activation function nnAct().");
nnDAct(z) = error("You must define the activation function derivative nnDAct().");
nnCost(datay,y) = error("You must define the cost function nnCost(dataOutput,output).");
nnDCost(datay,y) = error("You must define the cost function derivative nnDCost(dataOutput,output).");

max_stars_repo_name: krometis/SimpleNeuralNet

max_issues_repo_name: krometis/SimpleNeuralNet

max_forks_repo_name: krometis/SimpleNeuralNet


text:
module TestExecutors

using FLoops
using Test

function f_copy(executor)
    xs = 1:10
    ys = zeros(10)
    @floop executor for i in eachindex(xs, ys)
        ys[i] = xs[i]
    end
    return ys
end

function f_sum(executor)
    @floop executor for x in 1:10
        @reduce(s += x)
    end
    return s
end

function f_filter_sum(executor)
    @floop executor for x in 1:10
        if isodd(x)
            @reduce(s += x)
        end
    end
    return s
end

function f_sum_nested_loop(executor)
    @floop executor for x in 1:10
        for y in 1:x
            @reduce(s += y)
        end
    end
    return s
end

function f_sum_update(executor)
    @floop executor for x in 1:10
        if isodd(x)
            @reduce(s += 2x)
        end
    end
    return s
end

function f_sum_op_init(executor)
    @floop executor for x in 1:10
        if isodd(x)
            @reduce(s = 0 + 2x)
        end
    end
    return s
end

function f_count_update(executor)
    @floop executor for x in 1:10
        if isodd(x)
            @reduce(s += 1)
        end
    end
    return s
end

function f_count_op_init(executor)
    @floop executor for x in 1:10
        if isodd(x)
            @reduce(s = 0 + 1)
        end
    end
    return s
end

function f_sum_continue(executor)
    @floop executor for x in 1:10
        x > 4 && continue
        @reduce(s += x)
    end
    return s
end

function f_sum_break(executor)
    @floop executor for x in 1:10
        @reduce(s += x)
        x == 3 && break
    end
    return s
end

function f_find_return(executor)
    @floop executor for x in 1:10
        @reduce(s += x)
        x == 3 && return (:found, x)
    end
    return s
end

function f_find_goto(executor)
    @floop executor for x in 1:10
        @reduce() do (s; x)
            s = x
        end
        x == 3 && @goto FOUND
    end
    return s
    @label FOUND
    return (:found, s)
end

TESTDATA = [
    (f_copy, ==, 1:10, false),
    (f_sum, ===, 55, true),
    (f_filter_sum, ===, 25, true),
    (f_sum_nested_loop, ===, 220, true),
    (f_sum_update, ===, 50, true),
    (f_sum_op_init, ===, 50, true),
    (f_count_update, ===, 5, true),
    (f_count_op_init, ===, 5, true),
    (f_sum_continue, ===, 10, true),
    (f_sum_break, ===, 6, true),
    (f_find_return, ===, (:found, 3), true),
    (f_find_goto, ===, (:found, 3), true),
]

function test()
    @testset "$(args[1])" for args in TESTDATA
        test(args...)
    end
end

function test(f, ≛, desired, distributed)
    # Make sure that `executor` is used
    err = try
        f("string is not executor")
        nothing
    catch err
        @debug("Expected exception from `$f(nothing)`", exception = (err, catch_backtrace()))
        err
    end
    @test err isa MethodError
    # `maybe_set_simd` is happened to be the first function that
    # touches executor:
    @test err.f === FLoops._fold

    @test f(SequentialEx()) ≛ desired
    @test f(SequentialEx(simd = true)) ≛ desired
    @test f(ThreadedEx()) ≛ desired
    @testset for basesize in 2:10
        @test f(ThreadedEx(basesize = 2)) ≛ desired
    end
    distributed || return
    @test f(DistributedEx()) ≛ desired
    @testset for threads_basesize in 2:10
        @test f(DistributedEx(threads_basesize = threads_basesize)) ≛ desired
    end
end

end  # module


text:
using RigidBodyDynamics.OdeIntegrators

"""
$(SIGNATURES)

A trivial controller that simply sets the torques to zero.
"""
function zero_torque!(torques::AbstractVector, t, state::MechanismState)
    torques .= 0
end

"""
$(SIGNATURES)

Basic `Mechanism` simulation: integrate the state from time ``0`` to `final_time`
starting from the initial state `state0`. Return a `Vector` of times, as well as
`Vector`s of configuration vectors and velocity vectors at these times.

Optionally, a function (or other callable) can be passed in as the third argument (`control!`).
`control!` will be called at each time step of the simulation and allows you to specify joint torques
given the time and the state of the `Mechanism`. It should look like this:

```julia
function control!(torques::AbstractVector, t, state::MechanismState)
    rand!(torques) # for example
end
```

The integration time step can be specified using the `Δt` keyword argument (defaults to `1e-4`).

$stabilization_gains_doc

Uses `MuntheKaasIntegrator`. See [`RigidBodyDynamics.OdeIntegrators.MuntheKaasIntegrator`](@ref) for a lower
level interface with more options.
"""
function simulate(state0::MechanismState{X}, final_time, control! = zero_torque!;
        Δt = 1e-4, stabilization_gains=default_constraint_stabilization_gains(X)) where X
    T = cache_eltype(state0)
    result = DynamicsResult{T}(state0.mechanism)
    control_torques = similar(velocity(state0))
    closed_loop_dynamics! = let result=result, control_torques=control_torques, stabilization_gains=stabilization_gains # https://github.com/JuliaLang/julia/issues/15276
        function (v̇::AbstractArray, ṡ::AbstractArray, t, state)
            control!(control_torques, t, state)
            dynamics!(result, state, control_torques; stabilization_gains=stabilization_gains)
            copyto!(v̇, result.v̇)
            copyto!(ṡ, result.ṡ)
            nothing
        end
    end
    tableau = runge_kutta_4(T)
    storage = ExpandingStorage{T}(state0, ceil(Int64, final_time / Δt * 1.001)) # very rough overestimate of number of time steps
    integrator = MuntheKaasIntegrator(state0, closed_loop_dynamics!, tableau, storage)
    integrate(integrator, final_time, Δt)
    storage.ts, storage.qs, storage.vs
end


text:
# Toy linear ODE with 2 DOFs
# u_1_t - a * u_1 = 0
# u_2_t - b * u_1 - c * u_2 = 0

# Toy 2nd order ODE with 2 DOFs
# u_1_tt + b * u_1_t - a * u_1 = 0
# u_2_tt + a * u_1_t - b * u_1 - c * u_2 = 0

import Gridap.ODEs.ODETools: ODEOperator
import Gridap.ODEs.ODETools: AffineODEOperator
import Gridap.ODEs.ODETools: ConstantODEOperator
import Gridap.ODEs.ODETools: allocate_cache
import Gridap.ODEs.ODETools: update_cache!
import Gridap.ODEs.ODETools: allocate_residual
import Gridap.ODEs.ODETools: jacobian!
import Gridap.ODEs.ODETools: jacobians!
import Gridap.ODEs.ODETools: allocate_jacobian
import Gridap.ODEs.ODETools: residual!
using SparseArrays: spzeros

struct ODEOperatorMock{T<:Real,C} <: ODEOperator{C}
  a::T
  b::T
  c::T
  order::Integer
end

get_order(op::ODEOperatorMock) = op.order

function residual!(r::AbstractVector,op::ODEOperatorMock,t::Real,x::NTuple{2,AbstractVector},ode_cache)
  u,u_t = x
  r .= 0
  r[1] = u_t[1] - op.a * u[1]
  r[2] = u_t[2] - op.b * u[1] - op.c * u[2]
  r
end

function residual!(r::AbstractVector,op::ODEOperatorMock,t::Real,x::NTuple{3,AbstractVector},ode_cache)
  u,u_t,u_tt = x
  r .= 0
  r[1] = u_tt[1] + op.b * u_t[1] - op.a * u[1]
  r[2] = u_tt[2] + op.a * u_t[1]- op.b * u[1] - op.c * u[2]
  r
end

function allocate_residual(op::ODEOperatorMock,u::AbstractVector,cache)
  zeros(2)
end

function jacobian!(J::AbstractMatrix,
  op::ODEOperatorMock,
  t::Real,
  x::NTuple{2,AbstractVector},
  i::Int,
  γᵢ::Real,
  ode_cache)
  @assert get_order(op) == 1
  @assert 0 < i <= get_order(op)+1
  if i==1
    J[1,1] += -op.a
    J[2,1] += -op.b
    J[2,2] += -op.c
  elseif i==2
    J[1,1] += 1.0*γᵢ
    J[2,2] += 1.0*γᵢ
  end
  J
end

function jacobian!(J::AbstractMatrix,
  op::ODEOperatorMock,
  t::Real,
  x::NTuple{3,AbstractVector},
  i::Int,
  γᵢ::Real,
  ode_cache)
  @assert get_order(op) == 2
  @assert 0 < i <= get_order(op)+1
  if i==1
    J[1,1] += -op.a
    J[2,1] += -op.b
    J[2,2] += -op.c
  elseif i==2
    J[1,1] += op.b*γᵢ
    J[2,2] += op.a*γᵢ
  elseif i==3
    J[1,1] += 1.0*γᵢ
    J[2,2] += 1.0*γᵢ
  end
  J
end

function jacobians!(
  J::AbstractMatrix,
  op::ODEOperatorMock,
  t::Real,
  x::Tuple{Vararg{AbstractVector}},
  γ::Tuple{Vararg{Real}},
  ode_cache)
  @assert length(γ) == get_order(op) + 1
  for order in 1:get_order(op)+1
    jacobian!(J,op,t,x,order,γ[order],ode_cache)
  end
  J
end

function allocate_jacobian(op::ODEOperatorMock,u::AbstractVector,cache)
  spzeros(2,2)
end

allocate_cache(op::ODEOperatorMock) = nothing
allocate_cache(op::ODEOperatorMock,v::AbstractVector,a::AbstractVector) = (v,a,nothing)
update_cache!(cache,op::ODEOperatorMock,t::Real) = cache


text:
# The following variables need to be defined before `include`-ing this file:
# `testset_name_suffix`
# `n_values`
# `k_values`
# `m_values`

@time @testset "Matrix Multiply Float32 $(testset_name_suffix)" begin
    T = Float32
    for n ∈ n_values
        for k ∈ k_values
            for m ∈ m_values
                A = rand(T, m, k)
                B = rand(T, k, n)
                A′ = permutedims(A)'
                B′ = permutedims(B)'
                AB = A * B;
                @info "" T n k m
                @test @time(Octavian.matmul(A, B)) ≈ AB
                @test @time(Octavian.matmul(A′, B)) ≈ AB
                @test @time(Octavian.matmul(A, B′)) ≈ AB
                @test @time(Octavian.matmul(A′, B′)) ≈ AB
                @test @time(Octavian.matmul_serial(A, B)) ≈ AB
                @test @time(Octavian.matmul_serial(A′, B)) ≈ AB
                @test @time(Octavian.matmul_serial(A, B′)) ≈ AB
                @test @time(Octavian.matmul_serial(A′, B′)) ≈ AB
            end
        end
    end
    m = k = n = max(8Octavian.OCTAVIAN_NUM_TASKS[], 400)
    A = rand(T, m, k);
    B = rand(T, k, n);
    A′ = permutedims(A)';
    B′ = permutedims(B)';
    AB = A * B;
    @test matmul_pack_ab!(similar(AB), A, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A, B′) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B′) ≈ AB
end

@time @testset "Matrix Multiply Float64 $(testset_name_suffix)" begin
    T = Float64
    for n ∈ n_values
        for k ∈ k_values
            for m ∈ m_values
                A = rand(T, m, k)
                B = rand(T, k, n)
                A′ = permutedims(A)'
                B′ = permutedims(B)'
                AB = A * B;
                @info "" T n k m
                @test @time(Octavian.matmul(A, B)) ≈ AB
                @test @time(Octavian.matmul(A′, B)) ≈ AB
                @test @time(Octavian.matmul(A, B′)) ≈ AB
                @test @time(Octavian.matmul(A′, B′)) ≈ AB
                @test @time(Octavian.matmul_serial(A, B)) ≈ AB
                @test @time(Octavian.matmul_serial(A′, B)) ≈ AB
                @test @time(Octavian.matmul_serial(A, B′)) ≈ AB
                @test @time(Octavian.matmul_serial(A′, B′)) ≈ AB
            end
        end
    end
    m = k = n = max(8Octavian.OCTAVIAN_NUM_TASKS[], 400)
    A = rand(T, m, k);
    B = rand(T, k, n);
    A′ = permutedims(A)';
    B′ = permutedims(B)';
    AB = A * B;
    @test matmul_pack_ab!(similar(AB), A, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A, B′) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B′) ≈ AB
end

@time @testset "Matrix Multiply Int32 $(testset_name_suffix)" begin
    T = Int32
    for n ∈ n_values
        for k ∈ k_values
            for m ∈ m_values
                A = rand(T, m, k)
                B = rand(T, k, n)
                A′ = permutedims(A)'
                B′ = permutedims(B)'
                AB = A * B;
                @info "" T n k m
                @test @time(Octavian.matmul(A, B)) ≈ AB
                @test @time(Octavian.matmul(A′, B)) ≈ AB
                @test @time(Octavian.matmul(A, B′)) ≈ AB
                @test @time(Octavian.matmul(A′, B′)) ≈ AB
                @test @time(Octavian.matmul_serial(A, B)) ≈ AB
                @test @time(Octavian.matmul_serial(A′, B)) ≈ AB
                @test @time(Octavian.matmul_serial(A, B′)) ≈ AB
                @test @time(Octavian.matmul_serial(A′, B′)) ≈ AB
            end
        end
    end
    m = k = n = max(8Octavian.OCTAVIAN_NUM_TASKS[], 400)
    A = rand(T, m, k);
    B = rand(T, k, n);
    A′ = permutedims(A)';
    B′ = permutedims(B)';
    AB = A * B;
    @test matmul_pack_ab!(similar(AB), A, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A, B′) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B′) ≈ AB
end

@time @testset "Matrix Multiply Int64 $(testset_name_suffix)" begin
    T = Int64
    for n ∈ n_values
        for k ∈ k_values
            for m ∈ m_values
                A = rand(T, m, k)
                B = rand(T, k, n)
                A′ = permutedims(A)'
                B′ = permutedims(B)'
                AB = A * B;
                @info "" T n k m
                @test @time(Octavian.matmul(A, B)) ≈ AB
                @test @time(Octavian.matmul(A′, B)) ≈ AB
                @test @time(Octavian.matmul(A, B′)) ≈ AB
                @test @time(Octavian.matmul(A′, B′)) ≈ AB
                @test @time(Octavian.matmul_serial(A, B)) ≈ AB
                @test @time(Octavian.matmul_serial(A′, B)) ≈ AB
                @test @time(Octavian.matmul_serial(A, B′)) ≈ AB
                @test @time(Octavian.matmul_serial(A′, B′)) ≈ AB
            end
        end
    end
    m = k = n = max(8Octavian.OCTAVIAN_NUM_TASKS[], 400)
    A = rand(T, m, k);
    B = rand(T, k, n);
    A′ = permutedims(A)';
    B′ = permutedims(B)';
    AB = A * B;
    @test matmul_pack_ab!(similar(AB), A, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A, B′) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B′) ≈ AB
end

@time @testset "zero-sized-matrices" begin
    @test Octavian.matmul_serial(randn(0,0), randn(0,0)) == zeros(0, 0)
    @test Octavian.matmul_serial(randn(2,3), randn(3,0)) == zeros(2, 0)
    @test Octavian.matmul_serial(randn(2,0), randn(0,2)) == zeros(2, 2)
    @test Octavian.matmul_serial!(ones(2,2),randn(2,0), randn(0,2), 1.0, 2.0) == ones(2, 2) .* 2
    @test Octavian.matmul(randn(0,0), randn(0,0)) == zeros(0, 0)
    @test Octavian.matmul(randn(2,3), randn(3,0)) == zeros(2, 0)
    @test Octavian.matmul(randn(2,0), randn(0,2)) == zeros(2, 2)
    @test Octavian.matmul!(ones(2,2),randn(2,0), randn(0,2), 1.0, 2.0) == ones(2, 2) .* 2
end


text:
module TestExamples2DEulerMulticomponent

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "tree_2d_dgsem")

@testset "MHD Multicomponent" begin

  @trixi_testset "elixir_mhdmulti_ec.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhdmulti_ec.jl"),
      l2   = [0.04300299195675897, 0.042987505670835945, 0.025747180552589767, 0.1621856170457937,
              0.017453693413025828, 0.0174545523206645, 0.026873190440613162, 1.364647699274761e-15,
              0.012124340829605002, 0.024248681659210004],
      linf = [0.31371522041799105, 0.3037839783173047, 0.21500228807094351, 0.904249573054642,
              0.0939809809658183, 0.09470282020962761, 0.1527725397829759, 8.245701827530042e-15,
              0.0787460541210726, 0.1574921082421452])
  end

  @trixi_testset "elixir_mhdmulti_ec.jl with flux_derigs_etal" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhdmulti_ec.jl"),
      l2   = [0.04301155595653799, 0.04299735787276207, 0.025745530869947714,
              0.16206102676791553, 0.017454384272339165, 0.01745523378100091,
              0.026879482381500154, 0.0002038008756963954, 0.012094208262809778,
              0.024188416525619556],
      linf = [0.3156206778985397, 0.30941696929809526, 0.21167563519254176,
              0.9688251298546122, 0.09076254289155083, 0.09160589769498295,
              0.15698032974768705, 0.006131914796912965, 0.07839287555951036,
              0.1567857511190207],
      volume_flux = (flux_derigs_etal, flux_nonconservative_powell),
      surface_flux = (flux_derigs_etal, flux_nonconservative_powell))
  end

  @trixi_testset "elixir_mhdmulti_es.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhdmulti_es.jl"),
      l2   = [0.042511527162267, 0.04250603277530184, 0.02385422747993974, 0.11555081362726903,
              0.016366641053738043, 0.01636681584592762, 0.02581748418797907, 0.00023394429554818215,
              0.010834603551662698, 0.021669207103325396],
      linf = [0.23454607703107877, 0.23464789247380322, 0.11898832084115452, 0.5331209602648022,
              0.061744814466827336, 0.061767127585091286, 0.09595041452184983, 0.004421037168524759,
              0.06186597801911198, 0.12373195603822396])
  end

  @trixi_testset "elixir_mhdmulti_convergence.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhdmulti_convergence.jl"),
      l2   = [0.0003808877028249613, 0.0003808877028249593, 0.0005155994511260122, 0.000570394227652563,
              0.000439568811048544, 0.0004395688110485541, 0.0005074093477702055, 0.0003859005258180428,
              7.4611207452221e-5, 0.000149222414904442, 0.000298444829808884],
      linf = [0.0013324014301672943, 0.0013324014301669181, 0.002684449324758791, 0.0016236816790307085,
              0.0019172373117153363, 0.0019172373117148922, 0.002664932274107224, 0.0011872396664042962,
              0.0002855492944235094, 0.0005710985888470188, 0.0011421971776940376])
  end

  @trixi_testset "elixir_mhdmulti_rotor.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhdmulti_rotor.jl"),
      l2   = [0.6574605535168556, 0.6623234319361953, 0.0, 0.689806698245354,
              0.04883686128677976, 0.08382459729494686, 0.0, 0.0021114516459281177,
              0.15909290019096098, 0.07954645009548049],
      linf = [9.362339085941425, 9.169838118652539, 0.0, 10.600957847359556,
              0.6628317732399827, 1.4185626901435056, 0.0, 0.06914316292003836,
              3.328770801731456, 1.664385400865728],
      tspan = (0.0, 0.01))
end

end

end # module


text:
# to get one realization

function SDE(traceT::Array{Float64,1},drift,D::Union{Float64},Y0::Union{Float64,Array{Float64,1}})
	b=sqrt(2*D);
    dts = diff(traceT);
    steps = length(dts);
    sigma=sqrt.(dts);  # variance=sigma^2
    dim = length(Y0);
    dW = sigma.*randn((steps,dim)) .+ 0;
    if ~isa(Y0, Array)
        Y0 = [Y0];
    end
    if ~isa(drift, Array)
        drift = [drift];
    end

    if length(drift)!= dim
        error("Number of `drift` function should be the same as the number of the initial condition `Y0`");
    end

    traceY = fill(NaN,(length(traceT),dim));
    traceY[1,:] = Y0;

    for j = 1:dim
        y = Y0[j];
        for i = 2:length(traceT)
            y = y + drift[j](y)*dts[i-1] + b*dW[i-1, j]
            traceY[i, j] = y;
        end
    end
    return traceY
end

function SDE(dt,drift::Function,energyFunc::Function,D,Y0)
# CHECKPOINT:
# 1. Create Union for ::Function and Array{Function,1}
# 2. energyFunc is for the way calculating energy of an event.
# 3. Try to make it compatible with 2D process

	b=sqrt(2*D);
    dts = diff(traceT);
    steps = length(dts);
    sigma=sqrt.(dts);  # variance=sigma^2
    traceY2 = fill(0,size(traceT))
    for k = 1:numrealization
        dW = sigma.*randn((steps,)) .+ 0
        traceY = fill(NaN,size(traceT))
        traceY[1] = Y0;
        y = Y0;
        for i = 2:length(traceT)
            y = y + drift(y)*dts[i-1] + b*dW[i-1]
            traceY[i] = y;
        end
        traceY2 = traceY .+ traceY2;
    end
    return traceY2
end

# ensemble prediction
function SDE(dt,predictat,numrealizations::Int,drift::Function,D,Y0)
	b=sqrt(2*D);
    steps = floor(predictat/dt) |> Int64;
    dt = predictat/steps; # adjust dt a little bit to make dt*steps exactly the time we want to predict.
    sigma = sqrt(dt);  # variance=sigma^2
    lenY = steps +1;
    ensembleY = [];
    for k = 1:numrealizations
        dW = sigma.*randn((steps,)) .+ 0
        y = Y0;
        for i = 2:lenY
            y = y + drift(y)*dt + b*dW[i-1]
        end
        push!(ensembleY, y);
    end
    return ensembleY
end


text:
#=
Plots of Lucas tree price for different gamma.

=#

using Plots, LaTeXStrings
pyplot()

include("asset_pricing.jl")

gammas = [1.2, 1.4, 1.6, 1.8, 2.0]
ap = AssetPriceModel()
states = ap.mc.state_values

lines = []
labels = []

for gamma in gammas
    ap.gamma = gamma
    v = tree_price(ap)
    label="gamma = $gamma"
    push!(labels, label)
    push!(lines, v)
end

plot(lines,
    labels=labels',
    title="Price-divdend ratio as a function of the state",
    ylabel="price-dividend ratio",
    xlabel="state")

max_stars_repo_name: mwhchen/quantecon

max_issues_repo_name: mwhchen/quantecon

max_forks_repo_name: mwhchen/quantecon


text:
export misfit, reducedMisfit
"""
function misfit(param,M,w;doDerivative,pwr)

Compute the following quantities

misfit:             1/pwr * || R*(M*w - d) ||_pwr^pwr
residual:           M*w - d
weighted residual:  R*(M*w - d)

and optionally the gradient and Hessian of the misfit with respect to w.

Here, R is the inverse square root of the covariance matrix Gamma.

If M is a 3D array, all values are computed separately for each binary matrix
stored in M.

Inputs:

	1. param  - StrainReconParam
	2. M      - binary matrix or a 3D binary array
	3. w      - vector of weights

Keyword arguments:

    4. doDerivative   - flag for computing derivatives (default = false)
    5. pwr            - exponent in the likelihood function (default = param.pwr)

Output:
	1. Jc             - misfit
	2. res            - residual
	3. wres           - weighted residual
	4. dJ             - gradient of misfit w.r.t. w (if doDerivative = true)
	5. H              - Hessian of misfit w.r.t. w (if doDerivative = true)
"""
function misfit(param::StrainReconParam,M::BitArray{2},w::Vector{Float64}; doDerivative::Bool=false,pwr::Float64=param.pwr)
	d    = param.d
	R    = param.R

	res  = M*w-d
	wres = R*res
	Jc   = 1/pwr * sum(abs.(wres).^pwr)

	if doDerivative
		if pwr != 2.0
		    error("derivative for pwr!=2 not implemented");
		end
		dJ = M'*(R'*wres)
		H  = M'*((R'*R)*M)
		return Jc, res, wres, dJ, H
	else
		return Jc, res, wres
	end
end

function misfit(param::StrainReconParam,M::BitArray{3},w::Vector{Float64}; doDerivative::Bool=false,pwr::Float64=param.pwr)
	if doDerivative
		error("derivative for multiple binary matrices not implemented")
	end

	d    = param.d
	n    = param.n
	nM   = size(M,3)
	m    = length(d)

	# allocate space
	Jc   = zeros(nM)
	res  = zeros(m,nM)
	wres = zeros(m,nM)

	# compute misfit and residuals for each M
	for k=1:nM
		Jc[k],res[:,k],wres[:,k] = misfit(param,M[:,:,k],w,pwr=pwr)
	end
	return Jc, res, wres
end

"""
function reducedMisfit(param,w;nsol,doDerivative,pwr)

Compute reduced misfit

min_M 1/pwr * || R*(M*w - d) ||_pwr^pwr

where the minimum is taken over all binary matrices M in Omega_M.

If nsol>1, then nsol matrices, corresponding to nsol smallest misfit values, are
considered and the output is similar to misfit with 3D input array.

Input:

	1. param  - StrainReconParam
	2. w      - vector of weights

Keyword arguments:

	3. nsol           - consider nsol binary matrices (default = 1)
    4. doDerivative   - flag for computing derivatives (default = false)
	5. p              - exponent in the likelihood function (default = param.pwr)

Output:
	1. Jc             - misfit
	2. res            - residual
	3. wres           - weighted residual
	4. M              - the minimizing matrix (or a 3D array if nsol>1)
	5. dJ             - gradient of misfit w.r.t. w (if doDerivative = true)
	6. H              - Hessian of misfit w.r.t. w (if doDerivative = true)

"""
function reducedMisfit(param::StrainReconParam,w::Vector{Float64};nsol::Int=1,doDerivative::Bool=false,pwr::Float64=param.pwr)
	M = findM(param,w,nsol=nsol,pwr=pwr)
	if !doDerivative
		Jc, res, wres = misfit(param,M,w,pwr=pwr)
		return Jc, res, wres, M
	else
		Jc, res, wres, dJ, H = misfit(param,M,w,doDerivative=doDerivative,pwr=pwr)
		return Jc, res, wres, M, dJ, H
	end
end


text:
"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	curtailable_variable_renewable(EP::Model, inputs::Dict, Reserves::Int, CapacityReserveMargin::Int)
This function defines the constraints for operation of variable renewable energy (VRE) resources whose output can be curtailed ($y \in \mathcal{VRE}$), such as utility-scale solar PV or wind power resources or run-of-river hydro resources that can spill water.
The operational constraints for VRE resources are a function of each technology's time-dependent hourly capacity factor (or availability factor, $\rho^{max}_{y,z,t}$), in per unit terms, and the total available capacity ($\Delta^{total}_{y,z}$).
**Power output in each time step**
For each VRE technology type $y$ and model zone $z$, the model allows for incorporating multiple bins with different parameters for resource quality ($\rho^{max}_{y,z,t}$), maximum availability ($\overline{\Omega_{y,z}}$) and investment cost ($\Pi^{INVEST}_{y,z}$, for example, due to interconnection cost differences). We define variables related to installed capacity ($\Delta_{y,z}$) and retired capacity ($\Delta_{y,z}$) for all resource bins for a particular VRE resource type $y$ and zone $z$ ($\overline{\mathcal{VRE}}^{y,z}$). However, the variable corresponding to power output in each timestep is only defined for the first bin. Parameter $VREIndex_{y,z}$, is used to keep track of the first bin, where $VREIndex_{y,z}=1$ for the first bin and $VREIndex_{y,z}=0$ for the remaining bins. This approach allows for modeling many different bins per VRE technology type and zone while significantly reducing the number of operational variable (related to power output for each time step from each bin) added to the model with every additional bin. Thus, the maximum power output for each VRE resource type in each zone is given by the following equation:
```math
\begin{aligned}
	\Theta_{y,z,t} \leq \sum_{(x,z)\in \overline{\mathcal{VRE}}^{x,z}}{\rho^{max}_{x,z,t} \times \Delta^{total}_{x,z}}  \hspace{2 cm}  \forall y,z \in \{(y,z)|VREIndex_{y,z}=1, z \in \mathcal{Z}\},t \in \mathcal{T}
\end{aligned}
```
The above constraint is defined as an inequality instead of an equality to allow for VRE power output to be curtailed if desired. This adds the possibility of introducing VRE curtailment as an extra degree of freedom to guarantee that generation exactly meets demand in each time step.
Note that if ```Reserves=1``` indicating that frequency regulation and operating reserves are modeled, then this function calls ```curtailable_variable_renewable_reserves()```, which replaces the above constraints with a formulation inclusive of reserve provision.
"""
function curtailable_variable_renewable(EP::Model, inputs::Dict, Reserves::Int, CapacityReserveMargin::Int)
	## Controllable variable renewable generators
	### Option of modeling VRE generators with multiple availability profiles and capacity limits -  Num_VRE_Bins in Generators_data.csv  >1
	## Default value of Num_VRE_Bins ==1
	println("Dispatchable Resources Module")

	dfGen = inputs["dfGen"]

	T = inputs["T"]     # Number of time steps (hours)
	Z = inputs["Z"]     # Number of zones
	G = inputs["G"] 	# Number of generators

	VRE = inputs["VRE"]

	VRE_POWER_OUT = intersect(dfGen[dfGen.Num_VRE_Bins.>=1,:R_ID], VRE)
	VRE_NO_POWER_OUT = setdiff(VRE, VRE_POWER_OUT)

	### Expressions ###

	## Power Balance Expressions ##

	@expression(EP, ePowerBalanceDisp[t=1:T, z=1:Z],
	sum(EP[:vP][y,t] for y in intersect(VRE, dfGen[dfGen[!,:Zone].==z,:][!,:R_ID])))

	EP[:ePowerBalance] += ePowerBalanceDisp

	# Capacity Reserves Margin policy
	if CapacityReserveMargin > 0
		@expression(EP, eCapResMarBalanceVRE[res=1:inputs["NCapacityReserveMargin"], t=1:T], sum(dfGen[y,Symbol("CapRes_$res")] * EP[:eTotalCap][y] * inputs["pP_Max"][y,t]  for y in VRE))
		EP[:eCapResMarBalance] += eCapResMarBalanceVRE
	end

	### Constratints ###
	# For resource for which we are modeling hourly power output
	for y in VRE_POWER_OUT
		# Define the set of generator indices corresponding to the different sites (or bins) of a particular VRE technology (E.g. wind or solar) in a particular zone.
		# For example the wind resource in a particular region could be include three types of bins corresponding to different sites with unique interconnection, hourly capacity factor and maximim available capacity limits.
		VRE_BINS = intersect(dfGen[dfGen[!,:R_ID].>=y,:R_ID], dfGen[dfGen[!,:R_ID].<=y+dfGen[!,:Num_VRE_Bins][y]-1,:R_ID])

		# Constraints on contribution to regulation and reserves
		if Reserves == 1
			EP = curtailable_variable_renewable_reserves(EP, inputs)
		else
			# Maximum power generated per hour by renewable generators must be less than
			# sum of product of hourly capacity factor for each bin times its the bin installed capacity
			# Note: inequality constraint allows curtailment of output below maximum level.
			@constraint(EP, [t=1:T], EP[:vP][y,t] <= sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS))
		end

	end

	# Set power variables for all bins that are not being modeled for hourly output to be zero
	for y in VRE_NO_POWER_OUT
		fix.(EP[:vP][y,:], 0.0, force=true)
	end
	##CO2 Polcy Module VRE Generation by zone
	@expression(EP, eGenerationByVRE[z=1:Z, t=1:T], # the unit is GW
		sum(EP[:vP][y,t] for y in intersect(inputs["VRE"], dfGen[dfGen[!,:Zone].==z,:R_ID]))
	)
	EP[:eGenerationByZone] += eGenerationByVRE
	return EP
end

@doc raw"""
	curtailable_variable_renewable_reserves(EP::Model, inputs::Dict)
When modeling operating reserves, this function is called by ```curtailable_variable_renewable()```, which modifies the constraint for maximum power output in each time step from VRE resources to account for procuring some of the available capacity for frequency regulation ($f_{y,z,t}$) and upward operating (spinning) reserves ($r_{y,z,t}$).
```math
\begin{aligned}
	\Theta_{y,z,t} + f_{y,z,t} + r_{y,z,t} \leq \sum_{(x,z)\in \overline{\mathcal{VRE}}^{x,z}}{\rho^{max}_{x,z,t}\times \Delta^{total}_{x,z}}  \hspace{0.1 cm}   \forall y,z \in \{(y,z)|VREIndex_{y,z}=1, z \in \mathcal{Z}\},t \in \mathcal{T}
\end{aligned}
```
The amount of downward frequency regulation reserves also cannot exceed the current power output.
```math
\begin{aligned}
	f_{y,z,t} \leq \Theta_{y,z,t}
	\forall y,z \in \{(y,z)|VREIndex_{y,z}=1, z \in \mathcal{Z}\},t \in \mathcal{T}
\end{aligned}
```
The amount of frequency regulation and operating reserves procured in each time step is bounded by the user-specified fraction ($\upsilon^{reg}_{y,z}$,$\upsilon^{rsv}_{y,z}$) of available capacity in each period for each reserve type, reflecting the maximum ramp rate for the VRE resource in whatever time interval defines the requisite response time for the regulation or reserve products (e.g., 5 mins or 15 mins or 30 mins). These response times differ by system operator and reserve product, and so the user should define these parameters in a self-consistent way for whatever system context they are modeling.
```math
\begin{aligned}
	r_{y,z,t} \leq \upsilon^{rsv}_{y,z} \sum_{(x,z)\in \overline{\mathcal{VRE}}^{x,z}}{\rho^{max}_{x,z,t}\times \Delta^{total}_{x,z}}  \hspace{1 cm}   \forall y,z \in \{(y,z)|VREIndex_{y,z}=1, z \in \mathcal{Z}\},t \in \mathcal{T} \\
	f_{y,z,t} \leq \upsilon^{reg}_{y,z} \sum_{(x,z)\in \overline{\mathcal{VRE}}^{x,z}}{\rho^{max}_{x,z,t}\times \Delta^{total}_{x,z}}  \hspace{1 cm}   \forall y,z \in \{(y,z)|VREIndex_{y,z}=1, z \in \mathcal{Z}\},t \in \mathcal{T}
\end{aligned}
```
"""
function curtailable_variable_renewable_reserves(EP::Model, inputs::Dict)

	dfGen = inputs["dfGen"]
	T = inputs["T"]

	VRE_POWER_OUT = intersect(dfGen[dfGen.Num_VRE_Bins.>=1,:R_ID], inputs["VRE"])

	for y in VRE_POWER_OUT
		# Define the set of generator indices corresponding to the different sites (or bins) of a particular VRE technology (E.g. wind or solar) in a particular zone.
		# For example the wind resource in a particular region could be include three types of bins corresponding to different sites with unique interconnection, hourly capacity factor and maximim available capacity limits.
		VRE_BINS = intersect(dfGen[dfGen[!,:R_ID].>=y,:R_ID], dfGen[dfGen[!,:R_ID].<=y+dfGen[!,:Num_VRE_Bins][y]-1,:R_ID])

		if y in inputs["REG"] && y in inputs["RSV"] # Resource eligible for regulation and spinning reserves
			@constraints(EP, begin
				# For VRE, reserve contributions must be less than the specified percentage of the capacity factor for the hour
				[t=1:T], EP[:vREG][y,t] <= dfGen[!,:Reg_Max][y]*sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS)
				[t=1:T], EP[:vRSV][y,t] <= dfGen[!,:Rsv_Max][y]*sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS)

				# Power generated and regulation reserve contributions down per hour must be greater than zero
				[t=1:T], EP[:vP][y,t]-EP[:vREG][y,t] >= 0

				# Power generated and reserve contributions up per hour by renewable generators must be less than
				# hourly capacity factor. Note: inequality constraint allows curtailment of output below maximum level.
				[t=1:T], EP[:vP][y,t]+EP[:vREG][y,t]+EP[:vRSV][y,t] <= sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS)
			end)
		elseif y in inputs["REG"] # Resource only eligible for regulation reserves
			@constraints(EP, begin
				# For VRE, reserve contributions must be less than the specified percentage of the capacity factor for the hour
				[t=1:T], EP[:vREG][y,t] <= dfGen[!,:Reg_Max][y]*sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS)

				# Power generated and regulation reserve contributions down per hour must be greater than zero
				[t=1:T], EP[:vP][y,t]-EP[:vREG][y,t] >= 0

				# Power generated and reserve contributions up per hour by renewable generators must be less than
				# hourly capacity factor. Note: inequality constraint allows curtailment of output below maximum level.
				[t=1:T], EP[:vP][y,t]+EP[:vREG][y,t] <= sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS)
			end)

		elseif y in inputs["RSV"] # Resource only eligible for spinning reserves - only available in up, no down spinning reserves

			@constraints(EP, begin
				# For VRE, reserve contributions must be less than the specified percentage of the capacity factor for the hour
				[t=1:T], EP[:vRSV][y,t] <= dfGen[!,:Rsv_Max][y]*sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS)

				# Power generated and reserve contributions up per hour by renewable generators must be less than
				# hourly capacity factor. Note: inequality constraint allows curtailment of output below maximum level.
				[t=1:T], EP[:vP][y,t]+EP[:vRSV][y,t] <= sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS)
			end)
		else # Resource not eligible for reserves
			# Maximum power generated per hour by renewable generators must be less than
			# sum of product of hourly capacity factor for each bin times its the bin installed capacity
			# Note: inequality constraint allows curtailment of output below maximum level.
			@constraint(EP, [t=1:T], EP[:vP][y,t] <= sum(inputs["pP_Max"][yy,t]*EP[:eTotalCap][yy] for yy in VRE_BINS))
		end
	end

	return EP
end

max_stars_repo_name: Betristor/LTESOM

max_issues_repo_name: Betristor/LTESOM

max_forks_repo_name: Betristor/LTESOM


text:
# Called at end of each time-step in the DAE solver (mainly to calculate L_tilde for derivatives)
function store_static_callback(u, t, integrator, p)
    @unpack results, parameters, settings, L_1, L_2, ω, Ω, Ξ₁ = p
    P = length(u) - 3
    v_1_t = u[1]
    g_t = u[P+1]
    z_hat_t = max(u[P+2], 1.)
    E_t = u[P+3]
    Ω_t = Ω(t)
    S_t = S(g_t, parameters)
    L_tilde_t = L_tilde(g_t, z_hat_t, Ω_t, E_t, S_t, parameters)
    entry_residual_t = entry_residual(v_1_t, Ξ₁, parameters)
    push!(results, (t = t, g = g_t, z_hat = z_hat_t, Ω = Ω_t, E = E_t, v_1 = v_1_t, L_tilde = L_tilde_t, entry_residual = entry_residual_t))
end

# Dynamic kernel
function f!(residual,du,u,p,t)
    @unpack results, parameters, settings, L_1, L_2, ω, Ω, Ξ₁, z = p
    @unpack z_ex, T = settings
    @unpack ρ, δ, σ, μ, υ, κ, d, ζ = parameters
    residual .= 0
    P = length(residual) - 3 # note u[1:P]  = v(t) in the solution iterators
    v_1 = u[1]
    g = u[P+1]
    z_hat = max(u[P+2], 1.)
    E = u[P+3]

    # Get static equilibrium values and calculate the L_tilde growth rate
    @unpack S, L_tilde, z_bar, π_min, π = static_equilibrium(Ξ₁, v_1, g, z_hat, E, Ω(t), z, parameters)
    L_tilde_log_derivative = 0.0 # Default to Float literal.
    if (t < T)
        t_forward = results.t[end]
        L_tilde_forward = results.L_tilde[end]
        L_tilde_log_derivative = (log(1 - L_tilde_forward) - log(1 - L_tilde))/(t_forward - t) # See note under (40)
    end

    # Combine (40) with (52) to get A_t = (ρ + δ + L_tilde_log_derivative - (σ - 1) * (μ - g + (σ - 1) * υ^2 / 2))*I - (μ - g + (σ-1)*υ^2)*L_1 - (υ^2/2)*L_2
    # Then building the A(t) * v(t) - v'(t) residual directly for the non-algebraic equations
    residual[1:P] = (ρ + δ + L_tilde_log_derivative - (σ - 1) * (μ - g + (σ - 1) * υ^2 / 2))*u[1:P] # (40) into (52) then (53)
    residual[1:P] .-= (μ - g + (σ-1)*υ^2)*L_1*u[1:P] # (52) 2nd term in (53)
    residual[1:P] .-= (υ^2/2)*L_2*u[1:P] # (52) third term in (53)
    residual[1:P] .-= π # (53) final term
    residual[1:P] .-= du[1:P] # (53) subtracting the v'(t) to form the residual
    residual[P+1] = Ξ₁*v_1 + ζ - dot(ω, u[1:P])  # (54)
    residual[P+2] = z_hat^(σ-1) - κ * d^(σ-1) / π_min  # (55)
    residual[P+3] = entry_residual(v_1, Ξ₁, parameters)  # (56)
end

# Calculate the transition dynamics given a fixed Ω(t) function
function solve_dynamics(parameters, stationary_sol_T, settings, Ω)
    @unpack ρ, σ, N, θ, γ, d, κ, ζ, η, Theta, χ, υ, μ, δ = parameters
    @unpack z_ex, T, tstops = settings
    z = z_ex[2:end-1]
    P = length(z)
    @assert γ ≈ 1 # These are the only supported parameters for the transition dynamics at this point
    @assert η == 0

    # unpack the stationary solution
    v_T = stationary_sol_T.v_tilde
    g_T = stationary_sol_T.g
    z_hat_T = stationary_sol_T.z_hat
    L_tilde_T = stationary_sol_T.L_tilde
    Ω_T = stationary_sol_T.Ω
    Ξ₁ = 1/(1 - (σ-1)*(z[1] - z_ex[1]))  # (24), with ξ = (σ-1)
    entry_residual_T = entry_residual(v_T[1], Ξ₁, parameters)

    # Define the results data frame we'll be using and push the stationary onto it.
    results = DataFrame(t = T, g = g_T, z_hat = z_hat_T, Ω = Ω_T, E = δ, v_1 = v_T[1], L_tilde = L_tilde_T, entry_residual = entry_residual_T)

    # Define intermediate quantitities.
    ω = ω_weights(z_ex, θ, σ-1) # Quadrature weights.
    bc = (Mixed(ξ = σ-1), Mixed(ξ = σ-1)) # boundary conditions for differential operators
    L_1 = L₁₋bc(z_ex, bc) # use backward difference as the drift is negative
    L_2 = L₂bc(z_ex, bc)

    # initial conditions and parameters for the solver (under backwards time)
    u0 = [v_T; g_T; z_hat_T; δ] # i.e. E(T) = δ
    du0 = zeros(P+3)

    p = (results = results, parameters = parameters, settings = settings,
        L_1 = L_1, L_2 = L_2, ω = ω, Ξ₁ = Ξ₁, Ω = Ω, z = z)

    dae_prob = DAEProblem(f!, du0, u0, (T, 0.0), p, differential_vars = [trues(P); false; false; false])

    # Solve that DAE passing in the callbacks to store "future" L_tilde
    cb = FunctionCallingCallback((u, t, integrator) -> store_static_callback(u, t, integrator, p), tdir = -1, func_start = false) # Callback object.
    sol = DifferentialEquations.solve(dae_prob, callback = cb, tstops = tstops)

    # postprocess and return
    sort!(results) # sort results to be forward in time

    return (results = results, # dataframe of results
            sol = sol, # DAE solver output
            dae_parameters = p) # parameters bundle we passed to the DAE
end

# takes an E_hat guess and spits out a new E_hat guess
function dynamics_fixedpoint(E_nodes_interior, params_T, stationary_sol_0, stationary_sol_T, settings)
    # calculate the Ω(t) from the E_hat. Throws away the E (which the DAE calculates)
    Ω = Ω_from_E_hat(E_nodes_interior, stationary_sol_T.Ω, stationary_sol_0.Ω, params_T, settings)
    @unpack δ = params_T
    solution = solve_dynamics(params_T, stationary_sol_T, settings, Ω)

    # turn the DAE solution E into an E_hat (i.e., scale it to [-1, 0])
    filter!(row -> row.t ∈ settings.tstops, solution.results) # cut down to tstops
    E = solution.results[:E]
    E_hat = (E .- δ)/(δ - E[1]) # (43)

    return E_hat[2:end-1]
end

function solve_transition(parameters, settings)
    @unpack T, tstops, fixedpoint_beta, fixedpoint_ftol, fixedpoint_iterations, fixedpoint_show_trace, fixedpoint_m, fixedpoint_x0  = settings
    x0 = fixedpoint_x0(parameters, settings)
    @unpack d_0, d_T = parameters
    @assert d_0 !== d_T "Will lead to a divide by 0 error."

    # calculate the steady states at 0 and T
    params_T = merge(parameters, (d = d_T,))
    params_0 = merge(parameters, (d = d_0,))
    stationary_sol_T = stationary_numerical(params_T, settings)
    stationary_sol_0 = stationary_numerical(params_0, settings)

    @assert tstops[1] ≈ 0.0
    @assert tstops[end] ≈ T

    fp_sol = fixedpoint(E_nodes_interior -> dynamics_fixedpoint(E_nodes_interior, params_T, stationary_sol_0, stationary_sol_T, settings),
                        x0,
                        iterations = fixedpoint_iterations,
                        ftol = fixedpoint_ftol,
                        m = fixedpoint_m,
                        beta = fixedpoint_beta,
                        show_trace = fixedpoint_show_trace)
    converged(fp_sol) || @warn "fixed point didn't converge."
    E_nodes_interior = fp_sol.zero

    # regenerate full solution from the E_nodes_interior
    Ω = Ω_from_E_hat(E_nodes_interior, stationary_sol_T.Ω, stationary_sol_0.Ω, params_T, settings)
    @unpack δ = params_T
    solve_dynamics_output = solve_dynamics(params_T, stationary_sol_T, settings, Ω)
    return (solve_dynamics_output = solve_dynamics_output,
            nodes = E_nodes_interior,
            results = prepare_results(solve_dynamics_output, stationary_sol_T, stationary_sol_0))
end

function prepare_results(solve_dynamics_output, stationary_T, stationary_0)
    @unpack results, sol, dae_parameters = solve_dynamics_output
    @unpack Ω, parameters, settings, Ξ₁, ω = dae_parameters
    @unpack ρ, σ, N, θ, γ, d, κ, ζ, η, Theta, χ, υ, μ, δ = parameters
    @unpack z_ex, T, tstops = settings
    z = z_ex[2:end-1]
    P = length(z)

    # time-T values for welfare integrals
    L_tilde_T = stationary_T.L_tilde
    g_T = stationary_T.g
    z_hat_T = stationary_T.z_hat
    Ω_T = stationary_T.Ω
    log_c_T = log(c(L_tilde_T, Ω_T, z_bar(z_hat_T, Ω_T, parameters)))

    # interpolates for welfare integrals
    interpolated_L_tilde(t) = L_tilde(sol(t)[P+1], sol(t)[P+2], Ω(t), sol(t)[P+3], S(sol(t)[P+1], parameters), parameters) # recall that the P+1-th element of sol is g, P+2-nd is z_hat, P+3-rd is E
    interpolated_z_bar(t) = z_bar(sol(t)[P+2], Ω(t), parameters)
    interpolated_c(t) = c(interpolated_L_tilde(t), Ω(t), interpolated_z_bar(t))
    log_c(t) = log(interpolated_c(t))

    # welfare rows for dataframe
    log_M(t) = quadgk(t -> sol(t)[P+1], 0, t)[1]
    U(t) = quadgk(τ -> exp(-ρ*τ)*(log_M(t+τ) + log_c(t+τ)), 0, (T-t))[1] + exp(-ρ*(T-t))*(g_T + ρ*(log_c_T + g_T * T))/(ρ^2) # (C.83)

    # build out dataframe
    results = @transform(results, Ω = Ω.(:t))
    results = @transform(results, λ_ii = λ_ii.(:z_hat, Ref(parameters)))
    results = @transform(results, S = S.(:g, Ref(parameters)))
    results = @transform(results, z_bar = z_bar.(:z_hat, :Ω, Ref(parameters)))
    results = @transform(results, c = c.(:L_tilde, :Ω, :z_bar))
    results = @transform(results, π_min = π_min.(:L_tilde, :z_bar, Ref(parameters)))
    results = @transform(results, log_M = log_M.(:t))
    results = @transform(results, U = U.(:t))
    results = @transform(results, π_rat = π_rat.(:z_hat, Ref(parameters)))
    results = @transform(results, L_tilde_a = L_tilde_a.(:Ω, :S, Ref(parameters)))
    results = @transform(results, L_tilde_x = L_tilde_x.(:z_hat, :Ω, Ref(parameters)))
    results = @transform(results, L_tilde_E = L_tilde_E.(:Ω, :E, Ref(parameters)))
    results = @transform(results, w = w.(:z_bar, Ref(parameters)))

    results.r = ones(Float64, nrow(results)) # filler, to be overwritten
    for i in 1:nrow(results)
        t = results.t[i]
        c = results.c[i]
        g = results.g[i]
        log_c_forward = (i < nrow(results)) ? (log(results.c[i+1]) - log(c))/(results.t[i+1] - t) : 0.0
        if (i != nrow(results))
            @assert results.t[i+1] > t
        end
        results.r[i] = ρ +  δ + γ*(g + log_c_forward) # (C.56)
    end

    # prepend pre-shock steadystate to dataframe
    reduced_stationary = delete(stationary_0, :F, :a, :b, :ν, :x, :y, :U_bar, :v_tilde)
    stationary_data = merge(reduced_stationary, (E = δ, U = stationary_0.U_bar, v_1 = stationary_0.v_tilde[1], entry_residual = Ξ₁*stationary_0.v_tilde[1] - dot(stationary_0.v_tilde, ω) + ζ))
    ts_0 = settings.pre_shock_times
    for t in ts_0
        push!(results, merge(stationary_data, (t = t, log_M = t * stationary_data.g)))
    end

    return sort!(results) # sort results again to keep time in the right order
end

function Ω_from_E_hat(E_nodes_interior, Ω_T, Ω_0, parameters, settings)
    @unpack δ = parameters
    @unpack T, tstops, interp = settings
    E_nodes = [-1.0; E_nodes_interior; 0.0] # See footnote 19.  Without pinning a corner there is an extra degree of freedom in the scale
    ts = range(0.0, T, length=length(E_nodes))

    E_hat_interpolation = interp(ts, E_nodes) # might worth trying cubic spline
    E_hat(t) = E_hat_interpolation(t)
    E_hat_integral = quadgk(E_hat, 0, T)[1] # (42)
    Q = log(Ω_T/Ω_0) / E_hat_integral # (42) when Ω_T = Ω_0, then Q = 0 so that E(t) is constant with δ as expected
    E(t) = Q*E_hat(t) + δ # (43)

    Ω_derivative(Ω,p,t) = Q*E_hat(t)*Ω # (44)
    Ω_solution = DifferentialEquations.solve(ODEProblem(Ω_derivative, Ω_0, (0.0, T)), reltol = 1e-15, tstops = tstops) # if this fails, error will be thrown

    function Ω(t)
        if t > T
            return Ω_T
        elseif t > 0 && t <= T
            return Ω_solution(t)
        else # useful for fixed point approach
            return Ω_0
        end
    end

    return Ω
end


text:
using ErgodicControl
using ErgodicControlPlots

# Set up different domains with different discretizations
d = Domain([1,1], 100)
num_agents = 10

# Set up distribution and ergodic manager
K = 5
means = [[.3,.7], [.7,.3]]
Sigmas = [.025*eye(2), .025*eye(2)]
 phi = gaussian(d, means, Sigmas)
#phi = circle(d, [0.5, 0.5], 0.25)
em = ErgodicManagerR2(d, phi, K)

# Set up first trajectory manager
x0 = [0.49,0.01]
N = 80
h = 0.2
ci = ConstantInitializer([0.0, 0.0])
tm1 = TrajectoryManager(em, x0, h, N, ci)
dynamics!(tm1, SingleIntegrator(2,h))

tm2 = deepcopy(tm1)
tm2.x0= [.79,.99]
tm2.xy= [.79,.99]

tm3 = deepcopy(tm1)
tm3.x0 = [.10,.24]
tm3.xy = [.10,.24]

tm4 = deepcopy(tm1)
tm4.x0 = [.5,.02]
tm4.xy = [.5,.02]

tm5 = deepcopy(tm1)
tm5.x0 = [.23,.056]
tm5.xy = [.23,.056]

tm6 = deepcopy(tm1)
tm6.x0 = [.06,.1]
tm6.xy = [.06,.1]

tm7 = deepcopy(tm1)
tm7.x0 = [.23,.65]
tm7.xy = [.23,.65]

tm8 = deepcopy(tm1)
tm8.x0 = [.7,.11]
tm8.xy = [.7,.11]

tm9 = deepcopy(tm1)
tm9.x0 = [.8,.8]
tm9.xy = [.8,.8]

tm10 = deepcopy(tm1)
tm10.x0 = [.1,.41]
tm10.xy = [.1,.41]

tm11 = deepcopy(tm1)
tm11.x0 = [rand(), rand()]
tm11.xy = tm11.x0

tm12 = deepcopy(tm1)
tm12.x0 = [rand(), rand()]
tm12.xy = tm12.x0

tm13 = deepcopy(tm1)
tm13.x0 = [rand(), rand()]
tm13.xy = tm13.x0

tm14 = deepcopy(tm1)
tm14.x0 = [rand(), rand()]
tm14.xy = tm14.x0

tm15 = deepcopy(tm1)
tm15.x0 = [rand(), rand()]
tm15.xy = tm15.x0

tm16 = deepcopy(tm1)
tm16.x0 = [rand(), rand()]
tm16.xy = tm16.x0

tm17 = deepcopy(tm1)
tm17.x0 = [rand(), rand()]
tm17.xy = tm17.x0

tm18 = deepcopy(tm1)
tm18.x0 = [rand(), rand()]
tm18.xy = tm18.x0

tm19 = deepcopy(tm1)
tm19.x0 = [rand(), rand()]
tm19.xy = tm19.x0

tm20 = deepcopy(tm1)
tm20.x0 = [rand(), rand()]
tm20.xy = tm20.x0

tm21 = deepcopy(tm1)
tm21.x0 = [rand(), rand()]
tm21.xy = tm21.x0

tm22 = deepcopy(tm1)
tm22.x0 = [rand(), rand()]
tm22.xy = tm22.x0

tm23 = deepcopy(tm1)
tm23.x0 = [rand(), rand()]
tm23.xy = tm23.x0

tm24 = deepcopy(tm1)
tm24.x0 = [rand(), rand()]
tm24.xy = tm24.x0

tm25 = deepcopy(tm1)
tm25.x0 = [rand(), rand()]
tm25.xy = tm25.x0

tm26 = deepcopy(tm1)
tm26.x0 = [rand(), rand()]
tm26.xy = tm26.x0

tm27 = deepcopy(tm1)
tm27.x0 = [rand(), rand()]
tm27.xy = tm27.x0

tm28 = deepcopy(tm1)
tm28.x0 = [rand(), rand()]
tm28.xy = tm28.x0

tm29 = deepcopy(tm1)
tm29.x0 = [rand(), rand()]
tm29.xy = tm29.x0

tm30 = deepcopy(tm1)
tm30.x0 = [rand(), rand()]
tm30.xy = tm30.x0

tm31 = deepcopy(tm1)
tm31.x0 = [rand(), rand()]
tm31.xy = tm31.x0

tm32 = deepcopy(tm1)
tm32.x0 = [rand(), rand()]
tm32.xy = tm32.x0

tm33 = deepcopy(tm1)
tm33.x0 = [rand(), rand()]
tm33.xy = tm33.x0

tm34 = deepcopy(tm1)
tm34.x0 = [rand(), rand()]
tm34.xy = tm34.x0

tm35 = deepcopy(tm1)
tm35.x0 = [rand(), rand()]
tm35.xy = tm35.x0

tm36 = deepcopy(tm1)
tm36.x0 = [rand(), rand()]
tm36.xy = tm36.x0

tm37 = deepcopy(tm1)
tm37.x0 = [rand(), rand()]
tm37.xy = tm37.x0

tm38 = deepcopy(tm1)
tm38.x0 = [rand(), rand()]
tm38.xy = tm38.x0

tm39 = deepcopy(tm1)
tm39.x0 = [rand(), rand()]
tm39.xy = tm39.x0

tm40 = deepcopy(tm1)
tm40.x0 = [rand(), rand()]
tm40.xy = tm40.x0

tm41 = deepcopy(tm1)
tm41.x0 = [rand(), rand()]
tm41.xy = tm41.x0

tm42 = deepcopy(tm1)
tm42.x0 = [rand(), rand()]
tm42.xy = tm42.x0

tm44 = deepcopy(tm1)
tm44.x0 = [rand(), rand()]
tm44.xy = tm44.x0

tm43 = deepcopy(tm1)
tm43.x0 = [rand(), rand()]
tm43.xy = tm43.x0

tm45 = deepcopy(tm1)
tm45.x0 = [rand(), rand()]
tm45.xy = tm45.x0

tm46 = deepcopy(tm1)
tm46.x0 = [rand(), rand()]
tm46.xy = tm46.x0

tm47 = deepcopy(tm1)
tm47.x0 = [rand(), rand()]
tm47.xy = tm47.x0

tm48 = deepcopy(tm1)
tm48.x0 = [rand(), rand()]
tm48.xy = tm48.x0

tm49 = deepcopy(tm1)
tm49.x0 = [rand(), rand()]
tm49.xy = tm49.x0

tm50 = deepcopy(tm1)
tm50.x0 = [rand(), rand()]
tm50.xy = tm50.x0

tm51 = deepcopy(tm1)
tm51.x0 = [rand(), rand()]
tm51.xy = tm51.x0

tm52 = deepcopy(tm1)
tm52.x0 = [rand(), rand()]
tm52.xy = tm52.x0

tm53 = deepcopy(tm1)
tm53.x0 = [rand(), rand()]
tm53.xy = tm53.x0

tm54 = deepcopy(tm1)
tm54.x0 = [rand(), rand()]
tm54.xy = tm54.x0

tm55 = deepcopy(tm1)
tm55.x0 = [rand(), rand()]
tm55.xy = tm55.x0

tm56 = deepcopy(tm1)
tm56.x0 = [rand(), rand()]
tm56.xy = tm56.x0

tm57 = deepcopy(tm1)
tm57.x0 = [rand(), rand()]
tm57.xy = tm57.x0

tm58 = deepcopy(tm1)
tm58.x0 = [rand(), rand()]
tm58.xy = tm58.x0

tm59 = deepcopy(tm1)
tm59.x0 = [rand(), rand()]
tm59.xy = tm59.x0

tm60 = deepcopy(tm1)
tm60.x0 = [rand(), rand()]
tm60.xy = tm60.x0

tm61 = deepcopy(tm1)
tm61.x0 = [rand(), rand()]
tm61.xy = tm61.x0

tm62 = deepcopy(tm1)
tm62.x0 = [rand(), rand()]
tm62.xy = tm62.x0

tm63 = deepcopy(tm1)
tm63.x0 = [rand(), rand()]
tm63.xy = tm63.x0

tm64 = deepcopy(tm1)
tm64.x0 = [rand(), rand()]
tm64.xy = tm64.x0

tm65 = deepcopy(tm1)
tm65.x0 = [rand(), rand()]
tm65.xy = tm65.x0

tm66 = deepcopy(tm1)
tm66.x0 = [rand(), rand()]
tm66.xy = tm66.x0

tm67 = deepcopy(tm1)
tm67.x0 = [rand(), rand()]
tm67.xy = tm67.x0

tm68 = deepcopy(tm1)
tm68.x0 = [rand(), rand()]
tm68.xy = tm68.x0

tm69 = deepcopy(tm1)
tm69.x0 = [rand(), rand()]
tm69.xy = tm69.x0

tm70 = deepcopy(tm1)
tm70.x0 = [rand(), rand()]
tm70.xy = tm70.x0

# array of trajectory managers
# vtm = [tm1, tm2, tm3, tm4, tm5, tm6, tm7, tm8, tm9, tm10, tm11, tm12, tm13, tm14, tm15, tm16, tm17,tm18, tm19, tm20,tm21,tm22, tm23,tm24, tm25, tm26, tm27, tm28, tm29, tm30,tm31, tm32, tm33, tm34, tm35, tm36, tm37, tm38, tm39, tm40, tm41,tm42, tm43, tm44, tm45, tm46,tm47,tm48, tm49, tm50,tm51,tm52,tm53, tm54,tm55,tm56,tm57,tm58,tm59,tm60,tm61,tm62,tm63, tm64, tm65,tm66,tm67,tm68,tm69,tm70]
vtm = [tm1, tm2, tm3, tm4]

# Generate the trajectories
ddc = 1e-4
tm1.barrier_cost = 1
tm2.barrier_cost = 1
 # xd, ud = smc_trajectory_s(em, vtm, 0.2)
xd, ud = smc_trajectory(em, vtm)

# plotting
plot(em, xd, vtm)


text:
using Mimi

# components
@defcomp Comp1 begin
    par_1_1 = Parameter(index=[time])      # external input
    var_1_1 = Variable(index=[time])       # computed
    foo = Parameter()

    function run_timestep(p, v, d, t)
        v.var_1_1[t] = p.par_1_1[t]
    end
end

@defcomp Comp2 begin
    par_2_1 = Parameter(index=[time])      # connected to Comp1.var_1_1
    par_2_2 = Parameter(index=[time])      # external input
    var_2_1 = Variable(index=[time])       # computed
    foo = Parameter()

    function run_timestep(p, v, d, t)
        v.var_2_1[t] = p.par_2_1[t] + p.foo * p.par_2_2[t]
    end
end

@defcomp Comp3 begin
    par_3_1 = Parameter(index=[time])      # connected to Comp2.var_2_1
    var_3_1 = Variable(index=[time])       # external output
    foo = Parameter(default=30)

    function run_timestep(p, v, d, t)
        # @info "Comp3 run_timestep"
        v.var_3_1[t] = p.par_3_1[t] * 2
    end
end

@defcomp Comp4 begin
    par_4_1 = Parameter(index=[time])      # connected to Comp2.var_2_1
    var_4_1 = Variable(index=[time])        # external output
    foo = Parameter(default=300)

    function run_timestep(p, v, d, t)
        # @info "Comp4 run_timestep"
        v.var_4_1[t] = p.par_4_1[t] * 2
    end
end

@defcomposite A begin
    Component(Comp1)
    Component(Comp2)

    foo1 = Parameter(Comp1.foo)
    foo2 = Parameter(Comp2.foo)

    var_2_1 = Variable(Comp2.var_2_1)

    connect(Comp2.par_2_1, Comp1.var_1_1)
    connect(Comp2.par_2_2, Comp1.var_1_1)
end

@defcomposite B begin
    Component(Comp3)
    Component(Comp4)

    foo3 = Parameter(Comp3.foo)
    foo4 = Parameter(Comp4.foo)

    var_3_1 = Variable(Comp3.var_3_1)
end

@defcomposite top begin
    Component(A)

    fooA1 = Parameter(A.foo1)
    fooA2 = Parameter(A.foo2)

    # TBD: component B isn't getting added to mi
    Component(B)
    foo3 = Parameter(B.foo3)
    foo4 = Parameter(B.foo4)

    var_3_1 = Variable(B.var_3_1)

    connect(B.par_3_1, A.var_2_1)
    connect(B.par_4_1, B.var_3_1)
end

# model
m = Model()
md = m.md
set_dimension!(m, :time, 2005:2020)
add_comp!(m, top, nameof(top))

set_param!(m, :fooA1, 1)
set_param!(m, :fooA2, 2)
set_param!(m, :foo3, 10)
set_param!(m, :foo4, 20)
set_param!(m, :par_1_1, collect(1:length(time_labels(md))))
run(m)


text:
const qasm_compatible_intrinsics = Dict{Any, String}(
    Core.Intrinsics.add_float => "+",
    Core.Intrinsics.add_float_fast => "+",
    Core.Intrinsics.add_int => "+",
    Core.Intrinsics.sub_float => "-",
    Core.Intrinsics.sub_float_fast => "-",
    Core.Intrinsics.sub_int => "-",
    Core.Intrinsics.div_float => "/",
    Core.Intrinsics.div_float_fast => "/",
    Core.Intrinsics.mul_float => "*",
    Core.Intrinsics.mul_float_fast => "*",
    Core.Intrinsics.mul_int => "*",
    Core.Intrinsics.neg_float => "-",
    Core.Intrinsics.neg_float_fast => "-",
    Core.Intrinsics.neg_int => "-",
)

function validate(target::OpenQASMTarget, ci::CodeInfo)
    for (v, stmt) in enumerate(ci.code)
        @switch stmt begin
            # valid operations
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :apply), _, _, QuoteNode(locs::Locations))
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :apply), _, _, QuoteNode(locs::Locations), QuoteNode(ctrl::CtrlLocations))
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :measure), _, QuoteNode(locs::Locations))
                target.toplevel || error("cannot have measure inside QASM gate statement")
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :barrier), _, QuoteNode(locs::Locations))
            @case Expr(:invoke, _, GlobalRef(YaoCompiler, :measure_cmp), ::SSAValue, ::Int)
                target.toplevel || error("cannot have measure_cmp inside QASM gate statement")
            @case Expr(:invoke, _, GlobalRef(YaoCompiler, :measure_cmp), ::Int, ::SSAValue)
                target.toplevel || error("cannot have measure_cmp inside QASM gate statement")
            @case Expr(:new, ::Type{<:NamedTuple}, _...)
            @case Expr(:new, ::Type{<:IntrinsicRoutine}, _...)
            @case Expr(:invoke, mi, &(QASM.sin), _)
            @case Expr(:invoke, mi, &(QASM.cos), _)
            @case Expr(:invoke, mi, &(QASM.tan), _)
            @case Expr(:invoke, mi, &(QASM.exp), _)
            @case Expr(:invoke, mi, &(QASM.log), _)
            @case Expr(:invoke, mi, &(QASM.sqrt), _)

            # if we are lucky they are not inlined, we can still convert them without overlay
            # TODO: remove this in 1.7+ using external method table
            @case Expr(:invoke, mi, GlobalRef(_, :sin), _)
                mi.def.module === Base.Math || error("invalid function: $(mi.def.module).sin")
            @case Expr(:invoke, mi, GlobalRef(_, :cos), _)
                mi.def.module === Base.Math || error("invalid function: $(mi.def.module).cos")
            @case Expr(:invoke, mi, GlobalRef(_, :tan), _)
                mi.def.module === Base.Math || error("invalid function: $(mi.def.module).tan")
            @case Expr(:invoke, mi, GlobalRef(_, :exp), _)
                mi.def.module === Base.Math || error("invalid function: $(mi.def.module).exp")
            @case Expr(:invoke, mi, GlobalRef(_, :log), _)
                mi.def.module === Base.Math || error("invalid function: $(mi.def.module).log")
            @case Expr(:invoke, mi, GlobalRef(_, :sqrt), _)
                mi.def.module === Base.Math || error("invalid function: $(mi.def.module).sqrt")
            @case ::ReturnNode
            @case ::GotoIfNot
            # packed cargs
            @case Expr(:call, GlobalRef(&Core, :tuple), _...)

            # gate statements
            @case Expr(:call, GlobalRef(&Base, :getfield), Argument(3), QuoteNode(:args))
                target.toplevel && error("cannot have parameters in toplevel QASM program")
            @case Expr(:call, GlobalRef(&Base, :getfield), ::SSAValue, ::Int, true)
                target.toplevel && error("cannot have parameters in toplevel QASM program")

            # intrinsic math function
            @case Expr(:call, GlobalRef(mod, name), _...)
                f = Core.Compiler.abstract_eval_global(mod, name)
                f === Any && error("cannot determine function: $mod.$name")
                f = f.val
                f isa Core.IntrinsicFunction && f in keys(qasm_compatible_intrinsics) ||
                    error("function $f is not compatible with QASM $(target.version)")

            # errors
            @case Expr(:invoke, _, GlobalRef(YaoCompiler, :measure_cmp), ::SSAValue, ::SSAValue)
                error("can only compare measurement result with a constant integer in QASM $(target.version)")
            @case Expr(:invoke, _, GlobalRef(YaoCompiler, :measure_cmp), _, _)
                error("if statement condition must be comparision between measurement results in QASM $(target.version)")
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :expect), _, QuoteNode(locs::Locations))
                error("intrinsic expect is not supported for QASM")
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :apply), _, locs)
                error("location is not constant, got apply($locs)")
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :apply), _, locs, ctrl)
                error("location is not constant, got apply($locs, $ctrl)")
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :measure), _, locs)
                error("location is not constant, got measure($locs)")
            @case Expr(:invoke, _, GlobalRef(Intrinsics, :barrier), _, locs)
                error("location is not constant, got barrier($locs)")

            @case _
                error("unsupported operation: $stmt")
        end
    end
end


text:
using NESSie.Format

@testset "writeskel (surface model)" begin
    for T in testtypes
        nodes = Vector{T}[[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]]
        model = Model(
            nodes,
            Triangle{T}[
                Triangle(nodes[1], nodes[2], nodes[3]),
                Triangle(nodes[1], nodes[2], nodes[4])
            ]
        )
        (fname, fh) = mktemp()
        try
            writeskel(fh, model)
            seek(fh, 0)
            @test readline(fh) == "SKEL"
            @test split(readline(fh)) == ["4", "2"]
            for node in nodes
                @test [parse(T, e) for e in split(readline(fh))] == node
            end
            @test split(readline(fh)) == ["4", "0", "1", "2", "0"]
            @test split(readline(fh)) == ["4", "0", "1", "3", "0"]
            @test strip(readline(fh)) == ""
            @test eof(fh)
        finally
            close(fh)
            rm(fname)
        end
    end
end

@testset "writeskel (volume model)" begin
    for T in testtypes
        nodes = Vector{T}[[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [-2, 0, 0]]
        model = Model(
            nodes,
            Tetrahedron{T}[
                Tetrahedron(nodes[1], nodes[2], nodes[3], nodes[4]),
                Tetrahedron(nodes[1], nodes[2], nodes[3], nodes[5])
            ]
        )
        (fname, fh) = mktemp()
        try
            writeskel(fh, model)
            seek(fh, 0)
            @test readline(fh) == "SKEL"
            @test split(readline(fh)) == ["5", "2"]
            for node in nodes
                @test [parse(T, e) for e in split(readline(fh))] == node
            end
            @test split(readline(fh)) == ["7", "0", "1", "2", "3", "1", "3", "0"]
            @test split(readline(fh)) == ["7", "0", "1", "2", "4", "1", "4", "0"]
            @test strip(readline(fh)) == ""
            @test eof(fh)
        finally
            close(fh)
            rm(fname)
        end
    end
end