text:
#========== cost "model" ==========#

BLOCK = Ref(2^19)
# matmul: crossover about 70x70 on my laptop, 70^3 = 343_000, log2(70^3) = 18.3, but only 30% effect at 100^3=10^6
# batchmul: crossover between 20 & 30, log2(20^4) == 17.3, log2(30^4) == 19.6
# contract01: 1500 * 100, length 15_000, doesn't want threading
# cosine01: block 65_536, not sure if it wants
# log: vector crossover about length 10_000

COSTS = Dict(:* => 0, :/ => 2, :log => 10, :exp => 10) # plus 1 initially

callcost(sy, store) =
    if haskey(COSTS, sy)
        store.cost += COSTS[sy]
    end

# Then block = BLOCK[] ÷ store.cost[] is the number of iterations at which threading is turned on.

#========== runtime functions ==========#

"""
    threader(f!,T, Z, (A,B), (1:5,1:6), (1:7); block=100, keep=nothing)

Calling `f!(T, Z,A,B, 1:5,1:6, 1:7, nothing)` should do the work.
But if there are enough elements (meaning `5*6*7 > 100`)
then this will call `f!` many times in different threads.
(`block=nothing` turns the whole thing off.)

The first tuple of ranges are supposed to be safe to thread over,
probably the axes of the output `Z`.
It will subdivide the longest until either there are too few elements,
or it has spent its spawning budget, `nthreads()`.

For a scalar reduction the first tuple will be empty, and `length(Z)==1`.
Then it divides up the other axes, each accumulating in its own copy of `Z`.

`keep=nothing` means that it overwrites the array, anything else (`keep=true`) adds on.
"""
function threader(fun!::Function, T::Type, Z::AbstractArray, As::Tuple, I0s::Tuple, J0s::Tuple; block, keep=nothing)
    if !all(r -> r isa AbstractUnitRange, I0s) || !all(r -> r isa AbstractUnitRange, J0s)
        fun!(T, Z, As..., I0s..., J0s..., keep) # don't thread ranges like 10:-1:1
        return nothing
    end
    Is = map(UnitRange, I0s)
    Js = map(UnitRange, J0s)
    if isnothing(block)
        fun!(T, Z, As..., Is..., Js..., keep)
    elseif length(Is) >= 1
        thread_halves(fun!, T, (Z, As...), Is, Js, block, Threads.nthreads(), keep)
    elseif length(Z) == 1 && eltype(Z) <: Number
        thread_scalar(fun!, T, Z, As, Js, block, Threads.nthreads(), keep)
    else
        fun!(T, Z, As..., Is..., Js..., keep)
    end
    return nothing
end

"""
    ∇threader(f!,T, (dA,dB,dZ,A,B), (1:5), (1:6,1:7); block=100)

Again, calling `f!(T, dA,dB,dZ,A,B, 1:5,1:6, 1:7)` should do the work.

The first tuple of ranges should be safe to thread over, e.g. those in common
to all output arrays. If there are none, then it takes a second strategy
of dividing up the other ranges into blocks disjoint in every index,
and giving those to different threads.
"""
function ∇threader(fun!::Function, T::Type, As::Tuple, I0s::Tuple, J0s::Tuple; block)
    Is = map(UnitRange, I0s)
    Js = map(UnitRange, J0s)
    if isnothing(block)
        fun!(T, As..., Is..., Js...)
    elseif length(Is) >= 1
        thread_halves(fun!, T, As, Is, Js, block, Threads.nthreads())
    else
        thread_quarters(fun!, T, As, Js, block, Threads.nthreads())
    end
    nothing
end

function thread_halves(fun!::Function, T::Type, As::Tuple, Is::Tuple, Js::Tuple, block::Int, spawns::Int, keep=nothing)
    if spawns >= 2 && productlength(Is,Js) > block && productlength(Is) > 2
        I1s, I2s = cleave(Is, maybe32divsize(T))
        Base.@sync begin
            Threads.@spawn thread_halves(fun!, T, As, I1s, Js, block, div(spawns,2), keep)
            Threads.@spawn thread_halves(fun!, T, As, I2s, Js, block, div(spawns,2), keep)
        end
    elseif length(Is) + length(Js) >= 2
        block_halves(fun!, T, As, Is, Js, keep)
    else
        fun!(T, As..., Is..., Js..., keep)
    end
    nothing
end

const MINIBLOCK = Ref(64^3) # 2x quicker matmul at size 1000

function block_halves(fun!::Function, T::Type, As::Tuple, Is::Tuple, Js::Tuple, keep=nothing)
    if productlength(Is,Js) <= MINIBLOCK[]
        return fun!(T, As..., Is..., Js..., keep)
    elseif maximumlength(Is) > maximumlength(Js)
        I1s, I2s = cleave(Is)
        block_halves(fun!, T, As, I1s, Js, keep)
        block_halves(fun!, T, As, I2s, Js, keep)
    else
        J1s, J2s = cleave(Js)
        block_halves(fun!, T, As, Is, J1s, keep)
        block_halves(fun!, T, As, Is, J2s, true)
    end
    nothing
end

#=

using Tullio
Tullio.MINIBLOCK[] = 4
Z = zeros(Int, 11,9);
cnt = 0
f!(::Type, Z, i, j, keep) = begin
    global cnt
    Z[i,j] .= (global cnt+=1)
end
Tullio.block_halves(f!, Array, (Z,), UnitRange.(axes(Z)), (), nothing)
Z

  1   1   3   3   5   5   7   8   8
  1   1   3   3   5   5   7   8   8
  2   2   4   4   6   6   9  10  10
  2   2   4   4   6   6   9  10  10
 11  11  13  13  19  19  21  21  21
 12  12  14  14  20  20  22  23  23
 12  12  14  14  20  20  22  23  23
 15  15  16  16  24  24  26  27  27
 15  15  16  16  24  24  26  27  27
 17  17  18  18  25  25  28  29  29
 17  17  18  18  25  25  28  29  29

using TiledIteration
function colour!(A, n=1)
    for (i,t) in enumerate(TileIterator(axes(A), ntuple(_->n, ndims(A))))
        A[t...] .= i
    end
    A
end;
colour!(zeros(Int, 11,9), 2)

 1  1   7   7  13  13  19  19  25
 1  1   7   7  13  13  19  19  25
 2  2   8   8  14  14  20  20  26
 2  2   8   8  14  14  20  20  26
 3  3   9   9  15  15  21  21  27
 3  3   9   9  15  15  21  21  27
 4  4  10  10  16  16  22  22  28
 4  4  10  10  16  16  22  22  28
 5  5  11  11  17  17  23  23  29
 5  5  11  11  17  17  23  23  29
 6  6  12  12  18  18  24  24  30

=#

function thread_scalar(fun!::Function, T::Type, Z::AbstractArray, As::Tuple, Js::Tuple, block::Int, spawns::Int, keep=nothing)
    if productlength(Js) <= block || spawns < 2
        # @info "thread_scalar on $(Threads.threadid())" Js
        return fun!(T, Z, As..., Js..., keep)
    else
        Z1, Z2 = similar(Z), similar(Z)
        J1s, J2s = cleave(Js)
        Base.@sync begin
            Threads.@spawn thread_scalar(fun!, T, Z1, As, J1s, block, div(spawns,2))
            Threads.@spawn thread_scalar(fun!, T, Z2, As, J2s, block, div(spawns,2))
        end
        if keep === nothing
            Z .= Z1 .+ Z2
        else
            Z .+= Z1 .+ Z2
        end
    end
    nothing
end

function thread_quarters(fun!::Function, T::Type, As::Tuple, Js::Tuple, block::Int, spawns::Int)
    if productlength(Js) <= block || count(r -> length(r)>=2, Js) < 2 || spawns < 4
        return fun!(T, As..., Js...)
    else
        Q11, Q12, Q21, Q22 = quarter(Js, maybe32divsize(T))
        Base.@sync begin
            Threads.@spawn thread_quarters(fun!, T, As, Q11, block, div(spawns,4))
            Threads.@spawn thread_quarters(fun!, T, As, Q22, block, div(spawns,4))
        end
        Base.@sync begin
            Threads.@spawn thread_quarters(fun!, T, As, Q12, block, div(spawns,4))
            Threads.@spawn thread_quarters(fun!, T, As, Q21, block, div(spawns,4))
        end
    end
    nothing
end

productlength(Is::Tuple) = prod(length.(Is))
productlength(Is::Tuple, Js::Tuple) = productlength(Is) * productlength(Js)

maximumlength(Is::Tuple) = max(length.(Is)...)
maximumlength(::Tuple{}) = 0

maybe32divsize(::Type{<:AbstractArray{T}}) where T<:Number = max(1, 32 ÷ sizeof(T))
maybe32divsize(::Type) = 4

"""
    cleave((1:10, 1:20, 5:15)) -> lo, hi
Picks the longest of a tuple of ranges, and divides that one in half.
"""
function cleave(ranges::Tuple{Vararg{<:UnitRange,N}}, step::Int=4) where {N}
    longest = foldl((l,r) -> length(l) >= length(r) ? l : r, ranges; init=1:0)
    c = foldl((l,i) -> ranges[i]==longest ? i : l, ntuple(identity, N); init=0)

    cleft = if length(longest) >= 2*step
        minimum(longest) - 1 + step * div(length(longest), step * 2)
    else
        minimum(longest) - 1 + div(length(longest), 2)
    end
    alpha = ntuple(Val(N)) do i
        ri = ranges[i]
        i == c ? (minimum(ri):cleft) : (minimum(ri):maximum(ri))
    end
    beta = ntuple(Val(N)) do i
        ri = ranges[i]
        i == c ? (cleft+1:maximum(ri)) : (minimum(ri):maximum(ri))
    end
    return alpha, beta
end
cleave(::Tuple{}, n::Int=4) = (), ()

#=
@btime Tullio.cleave((1:100, 1:50, 50:90)) # was OK
@code_warntype Tullio.cleave((1:100, 1:50, 50:90), 4)
=#

"""
    quarter((1:10, 1:20, 3:4)) -> Q11, Q12, Q21, Q22
Picks the longest two ranges, divides each in half, and returns the four quadrants.
"""
function quarter(ranges::Tuple{Vararg{<:UnitRange,N}}, step::Int=4) where {N}
    c::Int, long::Int = 0, 0
    ntuple(Val(N)) do i
        li = length(ranges[i])
        if li > long
            c = i
            long = li
        end
    end
    d::Int, second::Int = 0,0
    ntuple(Val(N)) do j
        j == c && return
        lj = length(ranges[j])
        if lj > second
            d = j
            second = lj
        end
    end

    cleft = if long >= 2*step
        minimum(ranges[c]) - 1 + step * div(long, step * 2)
    else
        minimum(ranges[c]) - 1 + div(long, 2)
    end
    delta = if second >= 2*step
        minimum(ranges[d]) - 1 + step * div(second, step * 2)
    else
        minimum(ranges[d]) - 1 + div(second, 2)
    end

    Q11 = ntuple(Val(N)) do i
        ri = ranges[i]
        (i == c) ? (minimum(ri):cleft) : (i==d) ? (minimum(ri):delta) : (minimum(ri):maximum(ri))
    end
    Q12 = ntuple(Val(N)) do i
        ri = ranges[i]
        (i == c) ? (minimum(ri):cleft) : (i==d) ? (delta+1:maximum(ri)) : (minimum(ri):maximum(ri))
    end
    Q21 = ntuple(Val(N)) do i
        ri = ranges[i]
        (i == c) ? (cleft+1:maximum(ri)) : (i==d) ? (minimum(ri):delta) : (minimum(ri):maximum(ri))
    end
    Q22 = ntuple(Val(N)) do i
        ri = ranges[i]
        (i == c) ? (cleft+1:maximum(ri)) : (i==d) ? (delta+1:maximum(ri)) : (minimum(ri):maximum(ri))
    end
    return Q11, Q12, Q21, Q22
end

#========== the end ==========#


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

struct FilesRemoteApi <: SwaggerApi
    client::Swagger.Client
end

function _swaggerinternal_filesRemoteAdd(_api::FilesRemoteApi; external_id=nothing, title=nothing, filetype=nothing, external_url=nothing, preview_image=nothing, indexable_file_contents=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "POST", Dict{String, Any}, "/files.remote.add", ["slackAuth"])
    Swagger.set_param(_ctx.form, "external_id", external_id)  # type String
    Swagger.set_param(_ctx.form, "title", title)  # type String
    Swagger.set_param(_ctx.form, "filetype", filetype)  # type String
    Swagger.set_param(_ctx.form, "external_url", external_url)  # type String
    Swagger.set_param(_ctx.form, "preview_image", preview_image)  # type String
    Swagger.set_param(_ctx.form, "indexable_file_contents", indexable_file_contents)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/x-www-form-urlencoded"] : [_mediaType])
    return _ctx
end

"""

Adds a file from a remote service
Param: external_id::String
Param: title::String
Param: filetype::String
Param: external_url::String
Param: preview_image::String
Param: indexable_file_contents::String
Return: Dict{String, Any}
"""
function filesRemoteAdd(_api::FilesRemoteApi; external_id=nothing, title=nothing, filetype=nothing, external_url=nothing, preview_image=nothing, indexable_file_contents=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteAdd(_api; external_id=external_id, title=title, filetype=filetype, external_url=external_url, preview_image=preview_image, indexable_file_contents=indexable_file_contents, _mediaType=_mediaType)
    Swagger.exec(_ctx)
end

function filesRemoteAdd(_api::FilesRemoteApi, response_stream::Channel; external_id=nothing, title=nothing, filetype=nothing, external_url=nothing, preview_image=nothing, indexable_file_contents=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteAdd(_api; external_id=external_id, title=title, filetype=filetype, external_url=external_url, preview_image=preview_image, indexable_file_contents=indexable_file_contents, _mediaType=_mediaType)
    Swagger.exec(_ctx, response_stream)
end

function _swaggerinternal_filesRemoteInfo(_api::FilesRemoteApi; file=nothing, external_id=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", Dict{String, Any}, "/files.remote.info", ["slackAuth"])
    Swagger.set_param(_ctx.query, "file", file)  # type String
    Swagger.set_param(_ctx.query, "external_id", external_id)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/x-www-form-urlencoded"] : [_mediaType])
    return _ctx
end

"""

Retrieve information about a remote file added to Slack
Param: file::String
Param: external_id::String
Return: Dict{String, Any}
"""
function filesRemoteInfo(_api::FilesRemoteApi; file=nothing, external_id=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteInfo(_api; file=file, external_id=external_id, _mediaType=_mediaType)
    Swagger.exec(_ctx)
end

function filesRemoteInfo(_api::FilesRemoteApi, response_stream::Channel; file=nothing, external_id=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteInfo(_api; file=file, external_id=external_id, _mediaType=_mediaType)
    Swagger.exec(_ctx, response_stream)
end

function _swaggerinternal_filesRemoteList(_api::FilesRemoteApi; channel=nothing, ts_from=nothing, ts_to=nothing, limit=nothing, cursor=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", Dict{String, Any}, "/files.remote.list", ["slackAuth"])
    Swagger.set_param(_ctx.query, "channel", channel)  # type String
    Swagger.set_param(_ctx.query, "ts_from", ts_from)  # type Float32
    Swagger.set_param(_ctx.query, "ts_to", ts_to)  # type Float32
    Swagger.set_param(_ctx.query, "limit", limit)  # type Int32
    Swagger.set_param(_ctx.query, "cursor", cursor)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/x-www-form-urlencoded"] : [_mediaType])
    return _ctx
end

"""

Retrieve information about a remote file added to Slack
Param: channel::String
Param: ts_from::Float32
Param: ts_to::Float32
Param: limit::Int32
Param: cursor::String
Return: Dict{String, Any}
"""
function filesRemoteList(_api::FilesRemoteApi; channel=nothing, ts_from=nothing, ts_to=nothing, limit=nothing, cursor=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteList(_api; channel=channel, ts_from=ts_from, ts_to=ts_to, limit=limit, cursor=cursor, _mediaType=_mediaType)
    Swagger.exec(_ctx)
end

function filesRemoteList(_api::FilesRemoteApi, response_stream::Channel; channel=nothing, ts_from=nothing, ts_to=nothing, limit=nothing, cursor=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteList(_api; channel=channel, ts_from=ts_from, ts_to=ts_to, limit=limit, cursor=cursor, _mediaType=_mediaType)
    Swagger.exec(_ctx, response_stream)
end

function _swaggerinternal_filesRemoteRemove(_api::FilesRemoteApi; file=nothing, external_id=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "POST", Dict{String, Any}, "/files.remote.remove", ["slackAuth"])
    Swagger.set_param(_ctx.form, "file", file)  # type String
    Swagger.set_param(_ctx.form, "external_id", external_id)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/x-www-form-urlencoded"] : [_mediaType])
    return _ctx
end

"""

Remove a remote file.
Param: file::String
Param: external_id::String
Return: Dict{String, Any}
"""
function filesRemoteRemove(_api::FilesRemoteApi; file=nothing, external_id=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteRemove(_api; file=file, external_id=external_id, _mediaType=_mediaType)
    Swagger.exec(_ctx)
end

function filesRemoteRemove(_api::FilesRemoteApi, response_stream::Channel; file=nothing, external_id=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteRemove(_api; file=file, external_id=external_id, _mediaType=_mediaType)
    Swagger.exec(_ctx, response_stream)
end

function _swaggerinternal_filesRemoteShare(_api::FilesRemoteApi; file=nothing, external_id=nothing, channels=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", Dict{String, Any}, "/files.remote.share", ["slackAuth"])
    Swagger.set_param(_ctx.query, "file", file)  # type String
    Swagger.set_param(_ctx.query, "external_id", external_id)  # type String
    Swagger.set_param(_ctx.query, "channels", channels)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/x-www-form-urlencoded"] : [_mediaType])
    return _ctx
end

"""

Share a remote file into a channel.
Param: file::String
Param: external_id::String
Param: channels::String
Return: Dict{String, Any}
"""
function filesRemoteShare(_api::FilesRemoteApi; file=nothing, external_id=nothing, channels=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteShare(_api; file=file, external_id=external_id, channels=channels, _mediaType=_mediaType)
    Swagger.exec(_ctx)
end

function filesRemoteShare(_api::FilesRemoteApi, response_stream::Channel; file=nothing, external_id=nothing, channels=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteShare(_api; file=file, external_id=external_id, channels=channels, _mediaType=_mediaType)
    Swagger.exec(_ctx, response_stream)
end

function _swaggerinternal_filesRemoteUpdate(_api::FilesRemoteApi; file=nothing, external_id=nothing, title=nothing, filetype=nothing, external_url=nothing, preview_image=nothing, indexable_file_contents=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "POST", Dict{String, Any}, "/files.remote.update", ["slackAuth"])
    Swagger.set_param(_ctx.form, "file", file)  # type String
    Swagger.set_param(_ctx.form, "external_id", external_id)  # type String
    Swagger.set_param(_ctx.form, "title", title)  # type String
    Swagger.set_param(_ctx.form, "filetype", filetype)  # type String
    Swagger.set_param(_ctx.form, "external_url", external_url)  # type String
    Swagger.set_param(_ctx.form, "preview_image", preview_image)  # type String
    Swagger.set_param(_ctx.form, "indexable_file_contents", indexable_file_contents)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/x-www-form-urlencoded"] : [_mediaType])
    return _ctx
end

"""

Updates an existing remote file.
Param: file::String
Param: external_id::String
Param: title::String
Param: filetype::String
Param: external_url::String
Param: preview_image::String
Param: indexable_file_contents::String
Return: Dict{String, Any}
"""
function filesRemoteUpdate(_api::FilesRemoteApi; file=nothing, external_id=nothing, title=nothing, filetype=nothing, external_url=nothing, preview_image=nothing, indexable_file_contents=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteUpdate(_api; file=file, external_id=external_id, title=title, filetype=filetype, external_url=external_url, preview_image=preview_image, indexable_file_contents=indexable_file_contents, _mediaType=_mediaType)
    Swagger.exec(_ctx)
end

function filesRemoteUpdate(_api::FilesRemoteApi, response_stream::Channel; file=nothing, external_id=nothing, title=nothing, filetype=nothing, external_url=nothing, preview_image=nothing, indexable_file_contents=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_filesRemoteUpdate(_api; file=file, external_id=external_id, title=title, filetype=filetype, external_url=external_url, preview_image=preview_image, indexable_file_contents=indexable_file_contents, _mediaType=_mediaType)
    Swagger.exec(_ctx, response_stream)
end

export filesRemoteAdd, filesRemoteInfo, filesRemoteList, filesRemoteRemove, filesRemoteShare, filesRemoteUpdate

ext: jl


text:
# This file includees the callbacks tests

@testset  "CallbackTestSet" begin
    @info "Running CallbackTestSet ..."
    condition(obj) = obj.x >= 5
    action(obj) = println("Callaback activated . obj.x = ", obj.x)
    clb = Callback(condition=condition, action=action)
    @test isenabled(clb)
    disable!(clb)
    @test !isenabled(clb)

    mutable struct Object{CB}
        x::Int
        clb::CB
    end
    obj = Object(1, clb)
    for val in 1 : 10
        obj.x  = val
        obj.clb(obj)
    end

    mutable struct Object2{CB}
        x::Int
        callbacks::CB
    end

    obj2 = Object2(4, Callback(condition=condition, action=action))
    for val in 1 : 10
        obj2.x = val
        applycallbacks(obj2)
    end

    @info "Done CallbackTestSet."
end # testset


text:
module sensor_msgs

using Cxx
using CxxStd: StdString, StdVector

cxx"""
    #include <sensor_msgs/Imu.h>
    #include <sensor_msgs/LaserScan.h>
    """

for dtype in [ :Imu, :LaserScan ]
    cppname = string("sensor_msgs::", dtype)
    cxxtval = Expr(:macrocall, Symbol("@cxxt_str"), cppname)
    symtval = Symbol(dtype, :Msg)

    @eval global const $(symtval) = $(cxxtval)

    cxxcstr = Expr(:macrocall, Symbol("@icxx_str"), string(cppname, "();"))
    @eval (::Type{$(symtval)})() = $(cxxcstr)
end

end #module

eval(sensor_msgs, :(Imu	      = ImuMsg))
eval(sensor_msgs, :(LaserScan = LaserScanMsg))

export sensor_msgs

max_stars_repo_name: moyhig/CxxROS

max_issues_repo_name: moyhig/CxxROS

max_forks_repo_name: moyhig/CxxROS


text:
using OceanParameterizations
using Makie, MakieLayout

files =  ["free_convection", "strong_wind", "strong_wind_no_coriolis", "weak_wind_strong_cooling",
          "strong_wind_weak_cooling", "strong_wind_weak_heating"]

file_labels = Dict(
    "free_convection" => "Free convection",
    "strong_wind" => "Strong wind",
    "strong_wind_no_coriolis" => "Strong wind, no rotation",
    "weak_wind_strong_cooling" => "Weak wind, strong cooling",
    "strong_wind_weak_cooling" => "Strong wind, weak cooling",
    "strong_wind_weak_heating" => "Strong wind, weak heating"
)

x_lims = Dict(
    "uw" => (-0.0007,0.0001),
    "vw" => (-0.0001,0.0005),
    "wT" => (-0.00009,0.00006)
)

f = Dict(
    "uw" => 𝒟 -> 𝒟.uw.coarse, "vw" => 𝒟 -> 𝒟.vw.coarse, "wT" => 𝒟 -> 𝒟.wT.coarse, "T"  => 𝒟 -> 𝒟.T.coarse
)

scaling_factor = Dict(
    "uw" => 1e4,
    "vw" => 1e4,
    "wT" => 1e4,
    "T" => 1
)

x_labels = Dict(
    "uw" => "U'W' x 10⁴ (m²/s²)",
    "vw" => "V'W' x 10⁴ (m²/s²)",
    "wT" => "W'T' x 10⁴ (C⋅m/s)",
    "T" => "T (C)"
)

titles = Dict(
    "uw" => "Zonal momentum flux, U'W'",
    "vw" => "Meridional momentum flux, V'W'",
    "wT" => "Temperature flux, W'T'",
    "T" => "Temperature, T",
)

function generate_scene(name)
    scene = Scene()
    for (file, T) in Ts
        z = T.uw.z
        if name=="T"; z = T.T.z end
        lines!(scene, f[name](T)[:,288].*scaling_factor[name], z, linewidth=3)
    end
    scene
end

function z(name, T)
    if name=="T"; return T.T.z; end
    return T.uw.z
end
lins = [lines!(f[name](T)[:,288].*scaling_factor[name], z(name, T), linewidth=3) for (file,T) in Ts]

# leg = legend(lins, files)
vbox(generate_scene("uw"), generate_scene("vw"), generate_scene("wT"), generate_scene("T"))


text:
# Support Python iterators in Julia and vice versa

#########################################################################
# Iterating over Python objects in Julia

Base.IteratorSize(::Type{PyObject}) = Base.SizeUnknown()
function _start(po::PyObject)
    disable_sigint() do
        o = PyObject(@pycheckn ccall((@pysym :PyObject_GetIter), PyPtr, (PyPtr,), po))
        nxt = PyObject(@pycheck ccall((@pysym :PyIter_Next), PyPtr, (PyPtr,), o))

        return (nxt,o)
    end
end

"""
    PyIterator{T}(pyobject)

Wrap `pyobject::PyObject` into an iterator, that produces items of type `T`. To be more precise `convert(T, item)` is applied in each iteration. This can be useful to avoid automatic conversion of items into corresponding julia types.
```jldoctest
julia> using PyCall

julia> l = PyObject([PyObject(1), PyObject(2)])
PyObject [1, 2]

julia> piter = PyCall.PyIterator{PyAny}(l)
PyCall.PyIterator{PyAny,Base.HasLength()}(PyObject [1, 2])

julia> collect(piter)
2-element Array{Any,1}:
    1
    2

julia> piter = PyCall.PyIterator(l)
PyCall.PyIterator{PyObject,Base.HasLength()}(PyObject [1, 2])

julia> collect(piter)
2-element Array{PyObject,1}:
    PyObject 1
    PyObject 2
```
"""
struct PyIterator{T,S}
    o::PyObject
end

function _compute_IteratorSize(o::PyObject)
    S = try
        length(o)
        Base.HasLength
    catch err
        if !(err isa PyError && pyisinstance(err.val, @pyglobalobjptr :PyExc_TypeError))
            rethrow()
        end
        Base.SizeUnknown
    end
end
function PyIterator(o::PyObject)
    PyIterator{PyObject}(o)
end
function (::Type{PyIterator{T}})(o::PyObject) where {T}
    S = _compute_IteratorSize(o)
    PyIterator{T,S}(o)
end

Base.eltype(::Type{<:PyIterator{T}}) where T = T
Base.eltype(::Type{<:PyIterator{PyAny}}) = Any
Base.length(piter::PyIterator) = length(piter.o)

Base.IteratorSize(::Type{<: PyIterator{T,S}}) where {T,S} = S()

_start(piter::PyIterator) = _start(piter.o)

function Base.iterate(piter::PyIterator{T}, s=_start(piter)) where {T}
    ispynull(s[1]) && return nothing
    disable_sigint() do
        nxt = PyObject(@pycheck ccall((@pysym :PyIter_Next), PyPtr, (PyPtr,), s[2]))
        return (convert(T,s[1]), (nxt, s[2]))
    end
end
function Base.iterate(po::PyObject, s=_start(po))
    # avoid the constructor that calls length
    # since that might be an expensive operation
    # even if length is cheap, this adds 10% performance
    piter = PyIterator{PyAny, Base.SizeUnknown}(po)
    iterate(piter, s)
end

# issue #216
function Base.collect(::Type{T}, o::PyObject) where T
    a = T[]
    for x in o
        push!(a, x)
    end
    return a
end
Base.collect(o::PyObject) = collect(Any, o)

#########################################################################
# Iterating over Julia objects in Python

const jlWrapIteratorType = PyTypeObject()

# tp_iternext object of a jlwrap_iterator object, similar to PyIter_Next
function pyjlwrap_iternext(self_::PyPtr)
    try
        iter, iter_result_ref = unsafe_pyjlwrap_to_objref(self_)
        iter_result = iter_result_ref[]
        if iter_result !== nothing
            item, state = iter_result
            iter_result_ref[] = iterate(iter, state)
            return pyreturn(item)
        end
    catch e
        @pyraise e
    end
    return PyPtr_NULL
end

# the tp_iter slot of jlwrap object: like PyObject_GetIter, it
# returns a reference to a new jlwrap_iterator object
function pyjlwrap_getiter(self_::PyPtr)
    try
        self = unsafe_pyjlwrap_to_objref(self_)
        return pystealref!(jlwrap_iterator(self))
    catch e
        @pyraise e
    end
    return PyPtr_NULL
end

# Given a Julia object o, return a jlwrap_iterator Python iterator object
# that wraps the Julia iteration protocol with the Python iteration protocol.
# Internally, the jlwrap_iterator object stores the tuple (o, Ref(iterate(o))),
# where the Ref is used to store the (element, state)-iterator result tuple
# (which updates during iteration) and also the nothing termination indicator.
function jlwrap_iterator(o::Any)
    if jlWrapIteratorType.tp_name == C_NULL # lazily initialize
        pyjlwrap_type!(jlWrapIteratorType, "PyCall.jlwrap_iterator") do t
            t.tp_iter = @cfunction(pyincref_, PyPtr, (PyPtr,)) # new reference to same object
            t.tp_iternext = @cfunction(pyjlwrap_iternext, PyPtr, (PyPtr,))
        end
    end
    iter_result = iterate(o)
    return pyjlwrap_new(jlWrapIteratorType, (o, Ref{Union{Nothing,typeof(iter_result)}}(iterate(o))))
end

#########################################################################
# Broadcasting: if the object is iterable, return collect(o), and otherwise
#               return o.
function Base.Broadcast.broadcastable(o::PyObject)
    iter = ccall((@pysym :PyObject_GetIter), PyPtr, (PyPtr,), o)
    if iter == C_NULL
        pyerr_clear()
        return Ref(o)
    else
        ccall(@pysym(:Py_DecRef), Cvoid, (PyPtr,), iter)
        return collect(o)
    end
end


text:
"""
createPage(port::Int) -> String

Creates the proxy page attaching to the given WebSocket port (Julia side).
This page allows communication between Julia and the Chrome DevTools interface.
"""
function createPage(port::Int)
  tmppath = tempname() * ".html"

  open(tmppath, "w") do io
    Mustache.render(io, String(read(joinpath(@__DIR__, "wsproxy.html"))), port=port)
  end

  tmppath
end

"""
findfreeport([port_hint::Int]) -> Int

Finds the first available port on localhost.
"""
function findfreeport(porthint::Int=5000)
    xport, sock = listenany(porthint)
    close(sock)
    Int(xport)
end

"""
launchServer(outchan::Channel, inchan::Channel, port::Int)

Starts the WebSocket server that will send and receive messages from the
proxy page.
"""
function launchServer(outchan::Channel, inchan::Channel, port::Int)
  wsh = WebSocketHandler() do req,client

    @async begin  # listening loop
      DEBUG && info("starting inbound loop")
      while isopen(client) && isopen(inchan)
        try
          msg = String(read(client))
          put!(inchan, msg)
        catch e
          warn("error in reception loop : $e")
        end
      end
      DEBUG && info("exiting inbound loop (port $port)")
    end

    # outgoing message loop
    DEBUG && info("starting outbound loop")
    for m in outchan
      DEBUG && info("sending $m")
      try
        write(client, m)
      catch e
      end
    end

    DEBUG && info("exiting outbound loop (port $port)")
  end

  handler = HttpHandler() do req::Request, res::Response
    rsp = Response(100)
    rsp.headers["Access-Control-Allow-Origin"] = "http://localhost:8080"
    rsp.headers["Access-Control-Allow-Credentials"] = "true"
    rsp
  end

  handler.events["error"]  =
    (client, err) -> DEBUG ? println(err) : nothing
  handler.events["listen"] =
    (port) -> DEBUG ? println("Listening on $port...") : nothing

  server = Server(handler, wsh)
  @async run(server, port)
  server
end


text:
# include("Wavefunction.jl")
# include("SwapWavefunction.jl")
# include("SwapWavefunctionNS.jl")

type Estimate{T}
    sum::T
    n::Int
    avg::T
    Estimate() = new(T(0.), 0, T(0.))
end

function tally!{T}(est::Estimate{T}, value::T)
    est.sum += value
    est.n += 1
end

function average!(est::Estimate)
    est.avg = est.sum / est.n
end

abstract MeasurementData{T}

# for a single estimate
type MeasurementDataSimple{T} <: MeasurementData{T}
    description::ASCIIString
    steps_per_measurement::Int # a parameter
    measurement_steps::Int # a counter
    estimate::Estimate{T}

    function MeasurementDataSimple(description::ASCIIString, spm::Int)
        @assert spm > 0
        return new(description, spm, 0, Estimate{T}())
    end
end

# for many estimates corresponding to, say, different lattice site combinations in a correlation function
type MeasurementDataLattice{T} <: MeasurementData{T}
    description::ASCIIString
    steps_per_measurement::Int # a parameter
    measurement_steps::Int # a counter
    # FIXME: estimates and site_labels should be more closely tied together, through an ordered dict, say?
    # e.g., estimates[site_labels] = our estimate
    site_labels::Vector{Vector{Int}} # the part specific to being on a lattice
    estimates::Vector{Estimate{T}}

    function MeasurementDataLattice(description::ASCIIString, spm::Int)
        @assert spm > 0
        return new(description, spm, 0, Vector{Int}[], Estimate{T}[])
    end
end

# for holding subsystem occupation number probability measurements, \delta_{nA, NA}, for a single species
type MeasurementDataSubsystemOccupationProbability <: MeasurementData{Float64}
    description::ASCIIString
    steps_per_measurement::Int # a parameter
    measurement_steps::Int # a counter
    n::Vector{Int} # the subsystem occupation numbers
    species::Int
    subsystem::Vector{Int}
    estimates::Vector{Estimate{Float64}}

    function MeasurementDataSubsystemOccupationProbability(description::ASCIIString, spm::Int, n::Vector{Int}, species::Int, subsystem::Vector{Int})
        @assert spm > 0
        estimates = Estimate{Float64}[]
        for nn in n
            push!(estimates, Estimate{Float64}())
        end
        return new(description, spm, 0, n, species, subsystem, estimates)
    end
end

type Measurement
    data::MeasurementData
    measure_function::Function
end

# lots of average! functions, nice!
average!(meas::MeasurementDataSimple) = average!(meas.estimate)

function average!(meas::Union{MeasurementDataLattice, MeasurementDataSubsystemOccupationProbability})
    for e in meas.estimates
        average!(e)
    end
end

average!(measurement::Measurement) = average!(measurement.data)

function average!(measurements::Vector{Measurement})
    for m in measurements
        average!(m)
    end
end

function average!(measurements::Dict{ASCIIString, Measurement})
    for k in keys(measurements)
        average!(measurements[k])
    end
end

# the following two functions get MeasurementDataLattice into a form suitable for HDF5 and plotting
function package_averaged_data(data::MeasurementDataLattice)
    @assert length(data.site_labels) == length(data.estimates)
    out = Any[]
    for p in 1:length(data.site_labels)
        push!(out, Any[data.site_labels[p], data.estimates[p].avg])
    end
    return out
end

# data::Vector{Any}) should be of the form outputted by package_average_data
function corr_func_matrix(data::Vector{Any})
    N = sqrt(length(data))
    @assert floor(N) == N
    out = zeros(typeof(data[1][2]), (Int(N), Int(N)))
    for p in data
        out[p[1][1], p[1][2]] = p[2]
    end
    return out
end

# the following function gets MeasurementDataSubsystemOccupationProbability into a form suitable for HDF5 and plotting
function subsystem_occupatation_probabililities_matrix(data::MeasurementDataSubsystemOccupationProbability)
    @assert length(data.n) == length(data.estimates)
    out = zeros(Float64, (length(data.n), 2))
    for i in 1:length(data.n)
        out[i, :] = [data.n[i], data.estimates[i].avg].'
    end
    return out
end

# the following few functions are essentially constructors for appropriate MeasurementData* types
function initialize_2ptcorr_lattice_measurement(wf::Wavefunction, description::ASCIIString, T::Type, spm::Int)
    corr = MeasurementDataLattice{T}(description, spm)
    for i in 1:get_Nsites(wf)
        for j in 1:get_Nsites(wf)
            push!(corr.site_labels, [i, j])
            push!(corr.estimates, Estimate{T}())
        end
    end
    return corr
end

function initialize_1ptcorr_lattice_measurement(wf::Wavefunction, description::ASCIIString, T::Type, spm::Int)
    corr = MeasurementDataLattice{T}(description, spm)
    for i in 1:get_Nsites(wf)
        push!(corr.site_labels, [i])
        push!(corr.estimates, Estimate{T}())
    end
    return corr
end

function initialize_subsystem_occupation_probability_measurement(wf::Wavefunction, description::ASCIIString,
                                                                 spm::Int, subsystem::Vector{Int}, species::Int = 1)
    @assert 0 < length(subsystem) <= get_Nsites(wf)
    @assert minimum(subsystem) >= 1 && maximum(subsystem) <= get_Nsites(wf)
    n = collect(0:get_Nparticles(wf)[species])
    meas = MeasurementDataSubsystemOccupationProbability(description, spm, n, species, subsystem)
end

# this is what gets called by iterate! in the Simulation class
# FIXME: take current Wavefunction --> RegularWavefunction, and make RegularWavefunction and SwapWavefunction subtypes of Wavefunction
function possibly_measure(wf::Union{Wavefunction, SwapWavefunction, SwapWavefunctionNS}, corr::MeasurementData, measure_func::Function)
    corr.measurement_steps += 1
    if mod(corr.measurement_steps, corr.steps_per_measurement) == 0
        measure_func(wf, corr)
    end
end

# FIXME: put bangs on the names of the measure functions:  measure_* --> measure_*!

# FIXME: move this stuff to the Wavefunction class,
# but first need to separate out specific wf implementations from Wavefunction.jl (e.g., SBM)
function measure_SzSz(wf::SBMWavefunction, corr::MeasurementDataLattice{Float64})
    @assert is_valid_Gutz_spinhalf_configuration(wf)
    for p in 1:length(corr.site_labels)
        i = corr.site_labels[p][1]
        j = corr.site_labels[p][2]
        Szi = (wf.config.lattice_status[1][i] != 0) ? 0.5 : -0.5
        Szj = (wf.config.lattice_status[1][j] != 0) ? 0.5 : -0.5
        tally!(corr.estimates[p], Szi * Szj)
    end
end

function measure_SpSm{T}(wf::SBMWavefunction{T}, corr::MeasurementDataLattice{T})
    @assert is_valid_Gutz_spinhalf_configuration(wf)
    for p in 1:length(corr.site_labels)
        i = corr.site_labels[p][1]
        j = corr.site_labels[p][2]
        if i == j
            tally!(corr.estimates[p], (wf.config.lattice_status[1][i] != 0) ? T(1.) : T(0.))
        elseif wf.config.lattice_status[1][i] != 0 && wf.config.lattice_status[2][j] != 0
            m = wf.config.lattice_status[1][i] # = index of up particle at site i
            n = wf.config.lattice_status[2][j] # = index of down particle at site j
            tally!(corr.estimates[p],
                -1. * detrat_from_columns_update(wf.cmat_up, m, columns_for_detmat(wf.orbitals_up, j))
                    * detrat_from_columns_update(wf.cmat_down, n, columns_for_detmat(wf.orbitals_down, i)))
        else
            tally!(corr.estimates[p], T(0.))
        end
    end
end

function measure_ninj(wf::HardcoreWavefunction, corr::MeasurementDataLattice{Float64})
    for p in 1:length(corr.site_labels)
        i = corr.site_labels[p][1]
        j = corr.site_labels[p][2]
        ni = is_occupied(wf.config, i, 1) ? 1. : 0.
        nj = is_occupied(wf.config, j, 1) ? 1. : 0.
        tally!(corr.estimates[p], ni * nj)
    end
end

function measure_ni(wf::HardcoreWavefunction, corr::MeasurementDataLattice{Float64})
    for p in 1:length(corr.site_labels)
        i = corr.site_labels[p][1]
        ni = is_occupied(wf.config, i, 1) ? 1. : 0.
        tally!(corr.estimates[p], ni)
    end
end

function measure_subsystem_occupation_probability(wf::Wavefunction, meas::MeasurementDataSubsystemOccupationProbability)
    for i in 1:length(meas.n)
        # FIXME: gross...this is why estimates should be a dict or something, but who cares
        tally!(meas.estimates[i], get_Nparticles_in_subsystem(wf.config, meas.subsystem)[meas.species] == meas.n[i] ? 1. : 0.)
    end
end

function measure_Greens_function{T}(wf::FreeFermionWavefunction{T}, corr::MeasurementDataLattice{T})
    for p in 1:length(corr.site_labels)
        i = corr.site_labels[p][1]
        j = corr.site_labels[p][2]
        if i == j
            tally!(corr.estimates[p], is_occupied(wf.config, i, 1) ? T(1.) : T(0.))
        elseif is_occupied(wf.config, i, 1) && !is_occupied(wf.config, j, 1)
            m = wf.config.lattice_status[1][i]
            tally!(corr.estimates[p], detrat_from_columns_update(wf.cmat, m, columns_for_detmat(wf.orbitals, j)))
        else
            tally!(corr.estimates[p], T(0.))
        end
    end
end

function measure_Greens_function{T}(wf::CFLWavefunction{T}, corr::MeasurementDataLattice{T})
    for p in 1:length(corr.site_labels)
        i = corr.site_labels[p][]
        j = corr.site_labels[p][2]
        if i == j
            tally!(corr.estimates[p], is_occupied(wf.config, i, 1) ? T(1.) : T(0.))
        elseif is_occupied(wf.config, i, 1) && !is_occupied(wf.config, j, 1)
            m = wf.config.lattice_status[1][i]
            tally!(corr.estimates[p],
                detrat_from_columns_update(wf.cmat_d1, m, columns_for_detmat(wf.orbitals_d1, j))
              * detrat_from_columns_update(wf.cmat_d2, m, columns_for_detmat(wf.orbitals_d2, j))
              * detrat_from_columns_update(wf.cmat_f, m, columns_for_detmat(wf.orbitals_f, j)))
        else
            tally!(corr.estimates[p], T(0.))
        end
    end
end

function measure_Greens_function{Td,Tf}(wf::BosonCFLWavefunction{Td,Tf}, corr::MeasurementDataLattice{Td})
    for p in 1:length(corr.site_labels)
        i = corr.site_labels[p][]
        j = corr.site_labels[p][2]
        if i == j
            tally!(corr.estimates[p], is_occupied(wf.config, i, 1) ? Td(1.) : Td(0.))
        elseif is_occupied(wf.config, i, 1) && !is_occupied(wf.config, j, 1)
            m = wf.config.lattice_status[1][i]
            tally!(corr.estimates[p],
                detrat_from_columns_update(wf.cmat_d, m, columns_for_detmat(wf.orbitals_d, j))
              * detrat_from_columns_update(wf.cmat_f, m, columns_for_detmat(wf.orbitals_f, j)))
        else
            tally!(corr.estimates[p], Td(0.))
        end
    end
end

# should be pretty robust
function measure_swap_sign(wf::Union{SwapWavefunction, SwapWavefunctionNS}, meas::MeasurementDataSimple{Complex128})
    theta = 0.
    theta -= angle(get_phialpha1(wf))
    theta -= angle(get_phialpha2(wf))
    theta += angle(get_phibeta1(wf))
    theta += angle(get_phibeta2(wf))
    tally!(meas.estimate, exp(im * theta))
    return nothing
end

# gets the phibeta's exactly from the alpha configs and orbitals
function measure_swap_mod_exact(wf::Union{SwapWavefunction, SwapWavefunctionNS}, meas::MeasurementDataSimple{Float64})
    tally!(meas.estimate,
        abs((get_phibeta1_exact_from_swap(wf) / get_phialpha1(wf))
          * (get_phibeta2_exact_from_swap(wf) / get_phialpha2(wf))))
    return nothing
end

# pretty much garbage
function measure_swap_mod_naive(wf::SwapWavefunction, meas::MeasurementDataSimple{Float64})
    tally!(meas.estimate,
        abs((get_phibeta1(wf) / get_phialpha1(wf))
          * (get_phibeta2(wf) / get_phialpha2(wf))))
    return nothing
end

# Could alternatively have functions in SwapWavefunction.jl which return the detrats,
# and have measure_swap_mod_from_SMW(wf::SwapWavefunction, ...) take that function as an argument.
# This applies equally well to, e.g., measure_Greens_function, above.
# Do we want to put the specificity of things with the measurements (here) or with the wavefunctions?
function measure_swap_mod_from_SMW(wf::Union{FreeFermionSwapWavefunction, FreeFermionSwapWavefunctionNS}, meas::MeasurementDataSimple{Float64})
    config_beta1, config_beta2 = swap(wf.config_alpha1, wf.config_alpha2, subsystemA)
    # FIXME: check for sites which are occupied in both the swapped and unswapped configs, so we can get away with moving less particles;
    # since we take abs(detrats), shouldn't need to worry at all about signs popping up in that process.
    tally!(meas.estimate,
        abs(detrat_from_columns_update(wf.cmat_phialpha1, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals, config_beta1.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_phialpha2, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals, config_beta2.particle_positions[1][1:get_nA(wf)]))))
    return nothing
end

function measure_swap_mod_from_SMW(wf::BosonCFLSwapWavefunction, meas::MeasurementDataSimple{Float64})
    config_beta1, config_beta2 = swap(wf.config_alpha1, wf.config_alpha2, subsystemA)
    tally!(meas.estimate,
        abs(detrat_from_columns_update(wf.cmat_d_phialpha1, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_d, config_beta1.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_f_phialpha1, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_f, config_beta1.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_d_phialpha2, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_d, config_beta2.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_f_phialpha2, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_f, config_beta2.particle_positions[1][1:get_nA(wf)]))))
    return nothing
end

function measure_swap_mod_from_SMW(wf::CFLSwapWavefunction, meas::MeasurementDataSimple{Float64})
    config_beta1, config_beta2 = swap(wf.config_alpha1, wf.config_alpha2, subsystemA)
    tally!(meas.estimate,
        abs(detrat_from_columns_update(wf.cmat_d1_phialpha1, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_d1, config_beta1.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_d2_phialpha1, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_d2, config_beta1.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_f_phialpha1, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_f, config_beta1.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_d1_phialpha2, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_d1, config_beta2.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_d2_phialpha2, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_d2, config_beta2.particle_positions[1][1:get_nA(wf)]))
          * detrat_from_columns_update(wf.cmat_f_phialpha2, collect(1:get_nA(wf)), columns_for_detmat(wf.orbitals_f, config_beta2.particle_positions[1][1:get_nA(wf)]))))
    return nothing
end


text:
function HLADataset(s::String)
    HLADataset(Val(Symbol(s)))
end

function HLADataset(::Val{:Rousseau})
    hla_annotation = rousseau_hla_types()

    function extract_hla_types(file::String)
        hla_types = Vector{HLAType}()

        for record in FASTA.Reader(open(file, "r"))
            identifier = FASTA.identifier(record)
            patient_id = split(identifier, '_')[1]
            patient_id = endswith(patient_id, '.') ? patient_id[1:end-1] : patient_id
            accession = length(split(identifier, '_')) == 1 ? "" : split(identifier, '_')[2]

            if patient_id in keys(hla_annotation)
                hla_type = hla_annotation[patient_id]
            elseif accession in keys(hla_annotation)
                hla_type = hla_annotation[accession]
            else
                hla_type = HLAType(missing)
            end

            push!(hla_types, hla_type)
        end

        return hla_types
    end

    gag_file = joinpath(dirname(@__DIR__), "data", "test_sequences", "gag_rousseau.fasta")
    gag_hla_types = extract_hla_types(gag_file)
    gag_hla_data = HLAData("gag", gag_file, gag_hla_types, missing, missing)

    nef_file = joinpath(dirname(@__DIR__), "data", "test_sequences", "nef_rousseau.fasta")
    nef_hla_types = extract_hla_types(nef_file)
    nef_hla_data = HLAData("nef", nef_file, nef_hla_types, missing, missing)

    HLADataset("Rousseau", [gag_hla_data, nef_hla_data])
end

function rousseau_hla_types()
    hla_file = joinpath(dirname(@__DIR__), "data", "test_sequences", "hla_rousseau.csv")
    hla_data = readdlm(hla_file, '\t')[2:end, :]

    hla_types = Dict{String, HLAType}()

    for i in 1:size(hla_data, 1)
        patient_id, accession = string.(hla_data[i, 1:2])

        hla_alleles = parse_allele.(split(hla_data[i, 3], ' '))
        while length(hla_alleles) < 6
            push!(hla_alleles, missing)
        end

        hla_type = HLAType(tuple(hla_alleles...))
        if patient_id != "" hla_types[patient_id] = hla_type end
        if accession != "" hla_types[accession] = hla_type end
    end

    return hla_types
end


text:
using QuantumLab
using Test
using Dates
using LinearAlgebra

INFO(str) = @info "$(Dates.format(now(),dateformat"Y/m/d HH:MM:SS,sss"))  $str"

if (!@isdefined(indent))
  indent = ""
end

INFO("$(indent)READING:   h2o.xyz -> h2o::Geometry ")
h2o = readGeometryXYZ("h2o.xyz")

INFO("$(indent)OBTAINING:   BasisSetExchange -> sto3g::BasisSet")
sto3g = BasisSet("sto-3g")

INFO("$(indent)COMPUTING:   sto3g,h2o -> bas")
bas = computeBasis(sto3g,h2o)

INFO("$(indent)COMPUTING:   bas,h2o -> matrixOverlap, matrixKinetic, matrixNuclearAttraction, ERIs")
matrixOverlap = computeMatrixOverlap(bas)
matrixKinetic = computeMatrixKinetic(bas)
matrixNuclearAttraction = computeMatrixNuclearAttraction(bas,h2o)
ERIs = computeTensorElectronRepulsionIntegrals(bas)

INFO("$(indent)COMPUTING:   h2o -> matrixSADguess")
matrixSADguess = computeDensityGuessSAD("HF","STO-3G",h2o)

INFO("$(indent)COMPUTING:   sto3g, h2o -> shells, shells_native")
shells = computeBasisShellsLibInt2(sto3g,h2o)
shells_native = computeBasisShells(sto3g,h2o)

INFO("$(indent)COMPUTING:   (HartreeFock)  shells, h2o, matrixSADguess -> density")
density = evaluateSCF(shells,h2o,sum(matrixSADguess)/2,info=false,detailedinfo=false)[3]

INFO("$(indent)COMPUTING:   density, matrixKinetic, matrixNuclearAttraction, ERIs -> matrixFock")
matrixFock = computeMatrixFock(density,matrixKinetic,matrixNuclearAttraction,ERIs)

densityvirt = inv(matrixOverlap) - density

INFO("$(indent)COMPUTING:   shells, density -> S, T, J, K, Fock")
S = computeMatrixOverlap(shells)
T = computeMatrixKinetic(shells)
J = computeMatrixCoulomb(shells,density)
K = computeMatrixExchange(shells,density)
V = computeMatrixNuclearAttraction(shells,h2o)
fock = computeMatrixFock(T,V,J,K)
moenergies = eigvals(Symmetric(fock),Symmetric(S))


text:
using Statistics
import OnlineStats
d1, w1 = fill(1, 40), fill(4, 40)
d2, w2 = fill(2, 30), fill(3, 30)
d3, w3 = fill(3, 20), fill(2, 20)
d4, w4 = fill(4, 10), fill(1, 10)

d, wh = vcat(d1, d2, d3, d4), vcat(w1, w2, w3, w4)

@testset "WeightedAdaptiveHist fit!" begin
    h = (unique(d),
            map(*,
                sort(collect(keys(countmap(wh))), rev = true),
                sort(collect(values(countmap(wh))), rev = true)
            )
        )
    ws = sum(wh)

    o = WeightedAdaptiveHist(4)
    for i in 1:length(d)
        fit!(o, d[i], wh[i])
    end

    hfor, vsfor = value(o), weightsum(o)
    hval, vsval = value(fit!(WeightedAdaptiveHist(4), d, wh)),
                    weightsum(fit!(WeightedAdaptiveHist(4), d, wh))
    hzip, vszip = value(fit!(WeightedAdaptiveHist(4), zip(d, wh))),
                    weightsum(fit!(WeightedAdaptiveHist(4), zip(d, wh)))

    @test (h, ws) == (hfor, vsfor)
    @test (h, ws) == (hval, vsval)
    @test (h, ws) == (hzip, vszip)
end

@testset "WeightedAdaptiveHist merge" begin
    h = (unique(d),
            map(*,
                sort(collect(keys(countmap(wh))), rev = true),
                sort(collect(values(countmap(wh))), rev = true)
            )
        )
    ws = sum(wh)

    oh = map(d, wh) do di, whi
        fit!(WeightedAdaptiveHist(4), di, whi)
    end;

    r = reduce(merge!, oh)
    r2 = merge!(
        fit!(WeightedAdaptiveHist(4), d[1:end ÷ 2], wh[1:end ÷ 2]),
        fit!(WeightedAdaptiveHist(4), d[end ÷ 2 + 1:end], wh[end ÷ 2 + 1:end])
    )

    rh, rws = value(r), weightsum(r)
    rh2, rws2 = value(r2), weightsum(r2)

    @test (h, ws) == (rh, rws)
    @test (h, ws) == (rh2, rws2)
end

@testset "WeightedAdaptiveHist copy" begin
    o1 = WeightedAdaptiveHist(20)
    o2 = copy(o1)

    @test Ref(o2.alg.value, 1) != Ref(o1.alg.value, 1)
    @test Ref(o2.alg.b)        != Ref(o1.alg.b)
    @test Ref(o2.alg.ex)       != Ref(o1.alg.ex)
end

@testset "WeightedAdaptiveHist constructor" begin
    @test WeightedAdaptiveHist{WeightedAdaptiveBins{Float64}}(
        WeightedAdaptiveBins{Float64}(
            Pair{Float64, Float64}[], 20, Extrema(Float64)
        )
    ) == WeightedAdaptiveHist(20)
    @test WeightedAdaptiveHist{WeightedAdaptiveBins{Float32}}(
        WeightedAdaptiveBins{Float32}(
            Pair{Float32, Float32}[], 20, Extrema(Float32)
        )
    ) == WeightedAdaptiveHist(Float32, 20)
    @test WeightedAdaptiveHist(20) == WeightedAdaptiveHist(Float64, 20)
end

@testset "WeightedHist" begin
    @testset "fit!" begin
        h = WeightedHist(-3:1:1)

        fit!(h, -2.5,1.3)
        @test h.counts == [1,0,0,0]
        @test h.meanw  == [1.3,0,0,0]

        fit!(h, (-2.1, 1.0))
        @test value(h).y == [2.3, 0,0,0]
        @test h.counts == [2,0,0,0]
        @test h.meanw == [1.15,0,0,0]

        fit!(h, (-20, 2.0))
        @test value(h).y == [2.3, 0,0,0]
        @test h.outcount == [1, 0, 0]
        @test h.meanwout == [2.0,0,0]

        fit!(h, 20, 1.7)
        @test value(h).y == [2.3, 0,0,0]
        @test h.meanwout == [2.0, 0.0, 1.7]
        @test h.outcount == [1, 0, 1]

        fit!(h, -0.1, 1.1)
        @test value(h).y == [2.3, 0,1.1,0]
        @test h.edges === -3:1:1
    end

    @testset "merge!" begin
        h1 = WeightedHist(-2:2:2)
        fit!(h1, (-1, 5))
        fit!(h1, (1, 10))
        h1_copy = deepcopy(h1)
        @test merge!(h1, WeightedHist([-2,0,2])) == h1_copy
        @test value(merge!(h1_copy, h1)).y == [10, 20.]
    end

    @testset "stats" begin
        h = WeightedHist(-2:2:2)
        ho = OnlineStats.Hist(-2.:2:2)
        for _ in 1:rand(1:20)
            x = randn()
            fit!(h, x, 1.0)
            fit!(ho, x)
        end
        @test mean(h) ≈ mean(ho)
        @test std(h) ≈ std(ho)
        @test median(h) ≈ median(ho)
        @test nobs(h) ≈ nobs(ho)
        @test var(h) ≈ var(ho)
        @test all(extrema(h) .≈ extrema(ho))
    end

    @testset "N-dimensional Hist" begin
      h = WeightedHist((-2:2:2,0:3:6))
      fit!(h,(-1.5,1.5),1.5)
      @test h.counts == [1 0; 0 0]
      @test h.meanw  == [1.5 0; 0 0]

      fit!(h,(-0.5,0.2),1.1)
      @test h.counts == [2 0; 0 0]
      @test h.meanw  == [1.3 0;0 0]

      fit!(h,(-3.0,0.0),1.5)
      @test h.counts == [2 0; 0 0]
      @test h.meanw  == [1.3 0;0 0]
      @test h.outcount == [0 1 0; 0 0 0; 0 0 0]

      fit!(h,(-10,-10),1.1)
      @test h.counts == [2 0; 0 0]
      @test h.meanw  == [1.3 0;0 0]
      @test h.outcount == [1 1 0; 0 0 0; 0 0 0]

      fit!(h,(1.5,4.5),2.6)
      @test h.counts == [2 0; 0 1]
      @test h.meanw  == [1.3 0; 0 2.6]
      @test value(h) == (x=(-2:2:2, 0:3:6),y=[2.6 0;0 2.6])

      @test mean(h) == (0.0,3.0)
      @test var(h)  == (1.2380952380952381, 2.785714285714286)
      @test std(h)  == (1.1126972805283737, 1.6690459207925605)
      @test median(h) == (-1.0, 1.5)
      @test nobs(h) == 5
      @test extrema(h) == ((-1,1),(1.5,4.5))
    end

end


text:
using Test
using SafeTestsets
using Plots

unicodeplots()
ENV["GRAPE_LINESEARCH_ANALYSIS_VERBOSE"] = "1"

include(joinpath(@__DIR__, "generate_example_tests.jl"))

include(joinpath(@__DIR__, "download_dumps.jl"))

# Note: comment outer @testset to stop after first @safetestset failure
@time @testset verbose = true "GRAPE.jl Package" begin

    print("\n* Example 1 (examples/simple_state_to_state.jl):")
    @time @safetestset "Example 1" begin
        include(joinpath("examples", "simple_state_to_state.jl"))
    end

    print("\n* Example 2 (examples/perfect_entanglers.jl):")
    @time @safetestset "Example 2" begin
        include(joinpath("examples", "perfect_entanglers.jl"))
    end

    println("")

end


text:
# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control

module test_SimulationExamples

import ModiaMath

# Desired:
#   using Test
#
# In order that Test needs not to be defined in the user environment, it is included via ModiaMath:
using ModiaMath.Test

include(joinpath(ModiaMath.path, "examples", "Simulate_Pendulum.jl"))
import .Simulate_Pendulum
pendulumModel = Simulate_Pendulum.Pendulum(L=0.5, m=1.0, d=0.1, phi0_deg=60.0)
pendulum1 = ModiaMath.SimulationModel(pendulumModel; structureOfDAE=ModiaMath.DAE_NoSpecialStructure)
pendulum2 = ModiaMath.SimulationModel(pendulumModel; structureOfDAE=ModiaMath.DAE_ExplicitDerivatives)
pendulum3 = ModiaMath.SimulationModel(pendulumModel; structureOfDAE=ModiaMath.DAE_LinearDerivativesAndConstraints)

@testset "ModiaMath: examples/Simulate_Pendulum.jl" begin
    result = ModiaMath.simulate!(pendulum1, stopTime=10.0, tolerance=1e-8, interval=0.1, log=true)
    phi = result.series["phi"]
    w   = result.series["w"]

    @test isapprox(phi[end], 0.120289; atol=1e-3 )
    @test isapprox(w[end], 0.273965; atol=1e-2 )

    result = ModiaMath.simulate!(pendulum2, stopTime=10.0, tolerance=1e-8, interval=0.1, log=true)
    phi = result.series["phi"]
    w   = result.series["w"]

    @test isapprox(phi[end], 0.120289; atol=1e-3 )
    @test isapprox(w[end], 0.273965; atol=1e-2 )

    result = ModiaMath.simulate!(pendulum3, stopTime=10.0, tolerance=1e-8, interval=0.1, log=true)
    phi = result.series["phi"]
    w   = result.series["w"]

    @test isapprox(phi[end], 0.120289; atol=1e-3 )
    @test isapprox(w[end], 0.273965; atol=1e-2 )
end

include(joinpath(ModiaMath.path, "examples", "Simulate_FreeBodyRotation.jl"))
import .Simulate_FreeBodyRotation.FreeBodyRotation
import .Simulate_FreeBodyRotation.result

@testset "ModiaMath: examples/Simulate_FreeBodyRotation.jl" begin
    simulationModel = ModiaMath.SimulationModel(FreeBodyRotation(), stopTime=5.0, tolerance=1e-8)
    result2         = ModiaMath.simulate!(simulationModel, log=true)

    q1 = result.series["q"]
    w1 = result.series["w"]
    q2 = result2.series["q"]
    w2 = result2.series["w"]

    @test isapprox(q1[end,:], q2[end,:]; atol=1e-3)
    @test isapprox(w1[end,:], w2[end,:]; atol=1e-3)
end

end


text:
using Optim

function count_histogram(counts)
    max_counts = 0
    for e in counts
        max_counts = max(max_counts, maximum(e))
    end
    hist = zeros(max_counts + 1)
    for e in counts
        for c in e
            hist[c + 1] += 1
        end
    end
    hist
end

function bipoisson(n, λ0, λ1, bright_ratio)
    if (bright_ratio < 0 || bright_ratio > 1 || λ0 < 0 || λ1 < 0)
        0.0
    else
        (exp(-λ0) * (1 - bright_ratio) * λ0.^n +
            exp(-λ1) * bright_ratio * λ1.^n) ./ gamma(n + 1)
    end
end

function fit_count_histogram(histo)
    ns = collect(0:(length(histo) - 1))
    function negloglike(params)
       -sum(histo .* log(bipoisson(ns, params...)))
    end

    mle = optimize(negloglike, [2.0, 40.0, 0.5])
    mle.minimum
end

function count_threshold(λ0, λ1)
    floor((λ1 - λ0) / log(λ1 / λ0))
end

function apply_threshold(counts, threshold)
    map(c -> count(n -> n > threshold, c) / length(c), counts)
end

function auto_threshold(shots)
    λ0, λ1, ratio = fit_count_histogram(count_histogram(shots))
    apply_threshold(shots, count_threshold(λ0, λ1))
end


text:
# Module containing functions and types relevant to the
# gutzwiller projected free fermion system.

# needed functions:
#  computing ratio
# doing a swap operation
# yellowbook which is "inverse" of R_up and R_down: it gives the spin
# and fermion label # on each site
# so for example, if we had 4 sites with R_up = [1,3] and R_down = [2,4]
# then the log book would look like [(1,1), (1,-1), (2,1), (2,-1)]

# another example: let R_up = [2 5 1] R_down = [4 3 6]
# book = [(3,1), (1,1), (2,-1), (1,-1), (2,1), (3,-1)]
# so the format is (electron #, spin)
# should it be the other way around? maybe thats better
# or maybe even better is to write it as
# book = [(3,"Up"), (1,"Up"), (2,"Down"), (1,"Down"), (2,"Up"), (3,"Down")]

module FFG

include("mcbase.jl")

using StatsBase: sample, sample!
using IterTools: nth
using LightGraphs:
    SimpleGraph, edges, src, dst, rem_edge!, add_edge!, neighbors,
    adjacency_matrix, ne, nv, vertices
using SimpleWeightedGraphs: SimpleWeightedGraph
using LinearAlgebra: dot, det, inv, eigen, cond

# abstract type definitions
abstract type FreeFermionGutzwillerPolicy <: AbstractPolicy end

# concrete type defs
mutable struct GutzwillerChain <: AbstractChain
    basechain::MarkovChain
    GutzwillerChain() = new()
end
get_MarkovChain(c::GutzwillerChain) = c.basechain

struct DoubleGutzwillerChain <: ChainCollection
    """ Struct for working in the replica space. The item replicas
    is a tuple of whole GutzwillerChain objects. All the regular chain
    functions can be applied to DoubleGutzwillerChain by making them elementwise
    i.e. init_Chain!.(DGC) """
    replicas::Tuple{GutzwillerChain,GutzwillerChain} # replicas is a tuple of (gutz1, gutz2)
    basechain::MarkovChain # store "global" info like the observable and data
    DoubleGutzwillerChain() = new((GutzwillerChain(),GutzwillerChain()),MarkovChain())
end
get_Replicas(c::DoubleGutzwillerChain) = c.replicas # return the individual markov chains

struct InvalidLatticeError <: Exception
    var::Symbol
end

struct Electron
    """ for example, c_{1,up} """
    label::Int
    spin::String
end

struct SpinConfiguration
    """ This is borderline pointless, it's just a typed array"""
    sc::Array
end

struct ExchangeMove <: AbstractMove
    sites::Tuple
end

struct SwapNeighborsPolicy <: FreeFermionGutzwillerPolicy
end

mutable struct GutzwillerState <: AbstractState
     r_up::Array{Int64} # position of spin up electrons
     r_down::Array{Int64} # position of spin down electrons
     business_directory::Array{Electron} # map from site to electron index.
     spin_config::SpinConfiguration # a map of the spin configuration onto the graph
     bonds::SimpleGraph{Int64} # which bonds have disaligned spins
     wavefunctions::Array # wavefunctions ϕ_j(r_i) where ϵ_j < ϵ_fermi
     det_A_up::Any # could be real or complex
     det_A_down::Any
     A_inv_up::Array
     A_inv_down::Array
end

# Load in helper functions
include("FFGmath.jl")
include("FFGhelpers.jl")
include("FFGobservables.jl")
include("swap_routines.jl")

""" High level functions """
function get_Move(chain::GutzwillerChain)
    """ Get a move from the policy of the chain """
    policy = get_Policy(chain)
    move = get_move_from_policy(chain,policy)
    return move
end

function update_chain!(
        accept::Bool, chain::GutzwillerChain, move ; extras=nothing)
    """ Annoyingly, the order in which things are updated is important.
    In particular, update_Rs! needs the current business directory and
    update_Bonds! needs the current spin config"""
    if accept
        state = get_State(chain)
        update_Rs!(state,move)
        update_Bonds!(chain,move)
        update_Determinants!(chain,move,extras)
        update_Inverses!(state,move)
        update_Spin_config!(state,move)
        update_Business_directory!(state,move) # important that BD is updated last
    end
    update_Diagnostics!(chain,accept)
end

function update_state!(chain::GutzwillerChain,move;extras=nothing)
    """ A subset of update_chain! used for debugging and in swap """
    state = get_State(chain)
    update_Rs!(state,move)
    update_Bonds!(chain,move)
    update_Determinants!(chain,move,extras)
    update_Inverses!(state,move)
    update_Spin_config!(state,move)
    update_Business_directory!(state,move) # important that BD is updated last
end

function get_Wavefunctions(chain::GutzwillerChain)::Array
    """ This is specific to the model, so it is here and not in MCbase """
    model = get_Model(chain)
    lattice = model["lattice"]
    fermi_energy = model["fermi_energy"]
    hamiltonian = model["hamiltonian"] # eigen object

    wavefunctions = hamiltonian.vectors # wavefunctions[i,j] =   ϕ_j(r_i)
    eigenvalues = hamiltonian.values
    filled_mask = eigenvalues.<fermi_energy # find index of filled wavefunctions
    filled_states = wavefunctions[:,filled_mask] # ϕ_{ϵ(j) < ϵ_f} (r_i)

    # if ef = 0 then we need to just manually select the first N/2 states
    # since there are usually degenerate states at energy 0
    if fermi_energy == 0
        half = Int64(length(eigenvalues)/2)
        filled_states = wavefunctions[:,1:half]
    end

    return filled_states
end

function get_init_state(chain::GutzwillerChain)::GutzwillerState
    # get all the model information from the chain object
    model = get_Model(chain)
    lattice = model["lattice"]
    filling = model["filling"]
    fermi_energy = model["fermi_energy"]
    hamiltonian = model["hamiltonian"]
    dims = model["dims"]

    filled_states = get_Wavefunctions(chain)

    # Distribute electrons over the sites of the lattice
    num_sites = prod(dims)
    sites = collect(1:num_sites)

    # for now just assume filling = 1 and paramagnetic state
    # later i can add functionality for magnetic states
    # or states with holes
    n_up = Int(num_sites*filling/2)
    n_down = Int(num_sites*filling/2)

    # in one dimension at half filling, n_σ must be odd
    # https://bit.ly/2FPiRsw eq 31. this probably resolves
    # the conditioning error I was encountering sometimes.

    if length(dims)==1
        @assert (n_up % 2) == 1 "Need N_σ odd in one dimension"
        @assert (n_down % 2) == 1 "Need N_σ odd in one dimension"
    end

    # find a well conditioned initial state
    conditioning_tol = 10^6

    R_up, R_down, filled_states = get_conditioned_state(
        chain,n_up,n_down, conditioning_tol)

    # check whether filling is compatible with # states less than fermi energy
    @assert isacceptablelattice(filled_states,filling) "Lattice not playing well with filling,
         # please try different dimensions"

    # make the business directory
    business_directory = make_business_directory(R_up,R_down)

    # make the spin configuration
    spin_configuration = make_spin_config(R_up,R_down)

    # create list of neighbors which can be swapped in current configuration
    bonds = get_init_bonds(lattice.graph,spin_configuration)

    # calculate determinants
    det_A_up = get_Determinant(R_up,filled_states)
    det_A_down = get_Determinant(R_down,filled_states)

    # calculate inverse matrices for A_up and A_down
    A_inv_up = get_Inverse_Matrix(R_up,filled_states)
    A_inv_down = get_Inverse_Matrix(R_down,filled_states)

    # make the state object and return
    state = GutzwillerState(
        R_up,R_down, business_directory, spin_configuration,bonds,
        filled_states, det_A_up, det_A_down, A_inv_up,A_inv_down
        )
    return state
end

function get_test_state(chain::GutzwillerChain,R_up, R_down)::GutzwillerState
    """ Function that creates states from input r_up and r_down """
    model = get_Model(chain)
    lattice = model["lattice"]
    filling = model["filling"]
    fermi_energy = model["fermi_energy"]
    hamiltonian = model["hamiltonian"]
    dims = model["dims"]

    filled_states = get_Wavefunctions(chain)

    # Distribute electrons over the sites of the lattice
    num_sites = prod(dims)
    sites = collect(1:num_sites)

    # for now just assume filling = 1 and paramagnetic state
    # later i can add functionality for magnetic states
    # or states with holes
    n_up = Int(num_sites*filling/2)
    n_down = Int(num_sites*filling/2)

    # find a well conditioned initial state
    conditioning_tol = 10^6

    filled_states = get_Wavefunctions(chain)

    # check whether filling is compatible with # states less than fermi energy
    @assert isacceptablelattice(filled_states,filling) "Lattice not playing well with filling,
         # please try different dimensions"

    # make the business directory
    business_directory = make_business_directory(R_up,R_down)

    # make the spin configuration
    spin_configuration = make_spin_config(R_up,R_down)

    # create list of neighbors which can be swapped in current configuration
    bonds = get_init_bonds(lattice.graph,spin_configuration)

    # calculate determinants
    det_A_up = get_Determinant(R_up,filled_states)
    det_A_down = get_Determinant(R_down,filled_states)

    # calculate inverse matrices for A_up and A_down
    A_inv_up = get_Inverse_Matrix(R_up,filled_states)
    A_inv_down = get_Inverse_Matrix(R_down,filled_states)

    # make the state object and return
    state = GutzwillerState(
        R_up,R_down, business_directory, spin_configuration,bonds,
        filled_states, det_A_up, det_A_down, A_inv_up,A_inv_down
        )
    return state
end

function get_init_diagnostics(chain::GutzwillerChain)::Dict
    d0 = Dict("mc_steps" => 0,
            "accepted_moves" => 0,
            "num_measurements" => 1,
            "acceptance_ratio" => 0.0)
    return d0
end

function get_init_diagnostics(chain::DoubleGutzwillerChain)::Dict
    d0 = Dict("num_measurements" => 1,
            )
    return d0
end

""" Lower level functions """

function isacceptablelattice(filled_states,filling)::Bool
    """ Checks whether filled states is a square matrix """
    return Int64(size(filled_states)[1]*filling/2) == size(filled_states)[2]
end

function update_Diagnostics!(chain::GutzwillerChain,accept::Bool)
    diag = get_Diagnostics(chain)
    diag["mc_steps"] += 1
    diag["accepted_moves"] += Int64(accept)
end

function get_move_from_policy(chain::GutzwillerChain,policy::SwapNeighborsPolicy)::ExchangeMove
    """ Specific function for gutzwiller chain with swap neighbor policy"""
    bonds = get_State(chain).bonds
    move = get_ExchangeMove(bonds)
    return move
end

function compute_ratio(chain::GutzwillerChain,move::ExchangeMove)
    """ Computes transition amplitude. Returns transition amplitude,
    det_ratio_up, det_ratio_down. First return argument must always be ratio"""
    state = get_State(chain)
    wavefunctions = state.wavefunctions

    # update vectors for up spin determinant (1) and down spin determinant (2)
    u1, v1, u2, v2 = get_update_vectors(state,move)

    det_ratio_up = det_ratio_factor(state.A_inv_up,u1,v1)
    det_ratio_down = det_ratio_factor(state.A_inv_down,u2,v2)

    # we also need the ratio of proposal factors
    pf_ratio = get_proposal_factor_ratio(chain,move)

    return abs(det_ratio_up)^2*abs(det_ratio_down)^2*pf_ratio, (det_ratio_up, det_ratio_down)
end

function get_proposal_factor_ratio(chain::GutzwillerChain,move::ExchangeMove)
    state = get_State(chain)
    bonds = state.bonds
    configuration_factor_forwards = 1.0/Float64(count_bonds(bonds))
    new_bonds = get_updated_bonds(chain,move)
    configuration_factor_backwards = 1.0/Float64(count_bonds(new_bonds))
    return configuration_factor_backwards/configuration_factor_forwards
end

function make_spin_config(R_up,R_down)::SpinConfiguration
    nsites = length([R_up R_down])
    sc = zeros(Int64,nsites)
    sc[R_up] .= +1
    sc[R_down] .= -1

    return SpinConfiguration(sc)
end

function get_init_bonds(graph::SimpleGraph{Int64},
        spin_config::SpinConfiguration)::SimpleGraph{Int64}

    # when the graph is created, it is fully connected.
    # now we loop over the edges and use the spin configuration to remove bonds
    # where necessary. we do not need to add any bonds in the initial step

    # we need to make a deep copy of the graph because as we operate
    # on the graph its list of edges will change. instead we operate on a
    # copy of graph, and return that
    bonds = deepcopy(graph)

    for edge in deepcopy(edges(bonds))
        u, v = src(edge), dst(edge)
        if spin_config.sc[u] == spin_config.sc[v]
            rem_edge!(bonds,edge)
        end
    end

    return bonds
end

function get_init_bonds(graph::SimpleWeightedGraph,
        spin_config::SpinConfiguration)::SimpleGraph{Int64}

    # when the graph is created, it is fully connected.
    # now we loop over the edges and use the spin configuration to remove bonds
    # where necessary. we do not need to add any bonds in the initial step

    # we need to make a deep copy of the graph because as we operate
    # on the graph its list of edges will change. instead we operate on a
    # copy of graph, and return that
    graph = SimpleGraph(adjacency_matrix(graph))
    bonds = deepcopy(graph)

    for edge in deepcopy(edges(bonds))
        u, v = src(edge), dst(edge)
        if spin_config.sc[u] == spin_config.sc[v]
            rem_edge!(bonds,edge)
        end
    end

    return bonds
end

end


text:
using PyPlot
using JLD
function plot_rmse_vs_Np()
    X = load("../data/mean_rmse_vs_Np_full.jld")
    Np = X["Np"]
	mean_rmse = X["mean_rmse"]
	fig, ax = subplots()
	ax.grid(true)
	ax.plot(Np, mean_rmse, "bo--",lw=1.0,ms=20,label="RMSE")
	ax.legend(fontsize=28)
	ax.set_xlabel("# of particles", fontsize=28)
	ax.xaxis.set_tick_params(labelsize=28)
	ax.yaxis.set_tick_params(labelsize=28)
end


text:
module SodShockTube

using NLsolve: nlsolve
using PartialFunctions
using Documenter

export solve, ShockTubeProblem

sound_speed(γ, p, ρ) = √(γ * p / ρ)

function shock_tube_fn!(p1, p5, ρ1, ρ5, γ, p4)
    z = (p4[1] / p5 - 1)
    c1 = sound_speed(γ, p1, ρ1)
    c5 = sound_speed(γ, p5, ρ5)

    gm1 = γ - 1
    gp1 = γ + 1
    g2 = 2 * γ

    fact = gm1 / g2 * (c5 / c1) * z / √(1 + gp1 / g2 * z)
    fact = (1 - fact)^(g2 / gm1)
    return p1 * fact - p4[1]
end

function calculate_regions(pℓ, uℓ, ρℓ, pr, ur, ρr, γ=1.4)
    if pℓ < pr
        ρ1, p1, u1 = ρr, pr, ur
        ρ5, p5, u5 = ρℓ, pℓ, uℓ
    else
        ρ1, p1, u1 = ρℓ, pℓ, uℓ
        ρ5, p5, u5 = ρr, pr, ur
    end

    p4 = nlsolve(shock_tube_fn!$(p1, p5, ρ1, ρ5, γ), [p1]).zero[1]

    z = (p4 / p5 - 1)
    c5 = sound_speed(γ, p5, ρ5)

    gm1 = γ - 1
    gp1 = γ + 1

    gmfac1 = 0.5 * gm1 / γ
    gmfac2 = 0.5 * gp1 / γ

    fact = sqrt(1 + gmfac2 * z)
    u4 = c5 * z / (γ * fact)
    ρ4 = ρ5 * (1 + gmfac2 * z) / (1 + gmfac1 * z)

    w = c5 * fact
    p3 = p4
    u3 = u4
    ρ3 = ρ1 * (p3 / p1)^(1 / γ)
    region1 = (p1, ρ1, u1)
    region3 = (p3, ρ3, u3)
    region4 = (p4, ρ4, u4)
    region5 = (p5, ρ5, u5)
    return region1, region3, region4, region5, w
end

function calc_positions(pℓ, pr, region1, region3, w, xi, t, γ)
    p1, ρ1, u1 = region1
    p3, ρ3, u3 = region3

    c1 = sound_speed(γ, p1, ρ1)
    c3 = sound_speed(γ, p3, ρ3)

    if pℓ > pr
        xsh = xi + w * t
        xcd = xi + u3 * t
        xft = xi + (u3 - c3) * t
        xhd = xi - c1 * t
    else
        xsh = xi - w * t
        xcd = xi - u3 * t
        xft = xi - (u3 - c3) * t
        xhd = xi + c1 * t
    end

    return xhd, xft, xcd, xsh
end

function region_states(pℓ, pr, region1, region3, region4, region5)
    if pℓ > pr
        return Dict(
            "Region 1" => region1,
            "Region 2" => "RAREFACTION",
            "Region 3" => region3,
            "Region 4" => region4,
            "Region 5" => region5,
        )
    else
        return Dict(
            "Region 1" => region5,
            "Region 2" => region4,
            "Region 3" => region2,
            "Region 4" => "RAREFACTION",
            "Region 5" => region1,
        )
    end
end

function create_arrays(
    pℓ, pr, xℓ, xr, positions, state1, state3, state4, state5, x_arr, γ, t, xi
)
    xhd, xft, xcd, xsh = positions
    p1, ρ1, u1 = state1
    p3, ρ3, u3 = state3
    p4, ρ4, u4 = state4
    p5, ρ5, u5 = state5

    npts = length(x_arr)
    ρ = zeros(npts)
    p = zeros(npts)
    u = zeros(npts)
    c1 = sound_speed(γ, p1, ρ1)
    gm1 = γ - 1
    gp1 = γ + 1

    if pℓ > pr
        for (i, x) in enumerate(x_arr)
            if x < xhd
                ρ[i], p[i], u[i] = ρ1, p1, u1
            elseif x < xft
                u2 = 2 / gp1 * (c1 + (x - xi) / t)
                fact = 1 - 0.5 * gm1 * u2 / c1
                ρ2 = ρ1 * fact^(2 / gm1)
                p2 = p1 * fact^(2 * γ / gm1)
                ρ[i], p[i], u[i] = ρ2, p2, u2
            elseif x < xcd
                ρ[i], p[i], u[i] = ρ3, p3, u3
            elseif x < xsh
                ρ[i], p[i], u[i] = ρ4, p4, u4
            else
                ρ[i], p[i], u[i] = ρ5, p5, u5
            end
        end
    else
        for (i, x) in enumerate(x_arr)
            if x < xhd
                ρ[i], p[i], u[i] = ρ5, p5, u5
            elseif x < xft
                ρ[i], p[i], u[i] = ρ4, p4, u4
            elseif x < xcd
                ρ[i], p[i], u[i] = ρ3, p3, u3
            elseif x < xsh
                u2 = -2 / gp1 * (c1 + (xi - x) / t)
                fact = 1 + 0.5 * (γ - 1) * u2 / c1
                ρ2 = ρ1 * fact^(2 / gm1)
                p2 = p1 * fact^(2 * γ / gm1)
                ρ[i], p[i], u[i] = ρ2, p2, u2
            else
                ρ[i], p[i], u[i] = ρ1, p1, u1
            end
        end
    end
    return x_arr, p, ρ, u
end

function solve(left_state, right_state, geometry, t, γ, x_arr)
    pℓ, ρℓ, uℓ = left_state.p, left_state.ρ, left_state.u
    pr, ρr, ur = right_state.p, right_state.ρ, right_state.u
    xℓ, xr, xi = geometry

    if xℓ ≥ xr
        error("xℓ has to be less than xr")
    end

    if xℓ ≥ xr || xi ≤ xℓ
        error("xi has to be in between xℓ and xr")
    end

    region1, region3, region4, region5, w = calculate_regions(pℓ, uℓ, ρℓ, pr, ur, ρr, γ)
    regions = region_states(pℓ, pr, region1, region3, region4, region5)
    x_positions = calc_positions(pℓ, pr, region1, region3, w, xi, t, γ)

    pos_descriptions = (
        "Head of rarefaction", "Foot of rarefaction", "Contact Discontinuity", "Shock"
    )
    positions = Dict(desc => pos for (desc, pos) in zip(pos_descriptions, x_positions))

    x, p, ρ, u = create_arrays(
        pℓ, pr, xℓ, xr, x_positions, region1, region3, region4, region5, x_arr, γ, t, xi
    )
    energy = @. p / ((γ - 1)) + 0.5 * u^2
    values = (x=x, ρ=ρ, p=p, u=u, e=energy)
    return positions, regions, values
end

"""
    ShockTubeProblem
Contains the parameters of a shock tube problem

# Fields
`geometry::Tuple{Float64, Float64, Float64}` Contains the locations of the (left edge, right edge, initial shock location)

`left_state` Completely specified thermodynamic state of the left side of the discontinuity (NamedTuple of p, ρ, u)

`right_state` Completely specified thermodynamic state of the right side of the discontinuity (NamedTuple of p, ρ, u)

`t::Float64` The time at which the shock tube problem will be solved

`γ::Float64` The heat capacity ratio of the gas in the shock tube
"""
Base.@kwdef struct ShockTubeProblem
    geometry::Tuple{Float64,Float64,Float64}
    left_state
    right_state
    t::Float64
    γ::Float64
end

"""
    solve(s::ShockTubeProblem, x_arr)
Solve the given shock tube problem at the provided x locations.

# Returns

positions: A `Dictionary` which maps descriptive names of the regions to x coordinates

regions: A `Dictionary` which maps regions ("Region 1", "Region 2", etc) to thermodynamic states (ρ, p, u) in the shock tube solution

values: A `NamedTuple` (;x, ρ, p, e) containing the x coordinates, the density, pressure, and stagnation energy, respectively

# Example

```jldoctest;setup = :(using SodShockTube)
julia> problem = ShockTubeProblem(
    geometry = (0.0, 1.0, 0.5),
    left_state = (ρ = 1.0, u = 0.0, p = 1.0),
    right_state = (ρ = 0.125, u = 0.0, p = 0.1),
    t = 0.1,
    γ = 1.4
);

julia> xs = LinRange(0.0, 1.0, 500);

julia> positions, regions, values = solve(problem, xs);

julia> positions
Dict{String, Float64} with 4 entries:
  "Shock"                 => 0.850431
  "Foot of rarefaction"   => 0.485945
  "Head of rarefaction"   => 0.263357
  "Contact Discontinuity" => 0.685491

julia> regions
Dict{String, Any} with 5 entries:
  "Region 5" => (0.1, 0.125, 0.0)
  "Region 1" => (1.0, 1.0, 0.0)
  "Region 4" => (0.30313, 0.265574, 0.927453)
  "Region 3" => (0.30313, 0.426319, 0.927453)
  "Region 2" => "RAREFACTION"
"""
function solve(s::ShockTubeProblem, x_arr)
    return solve(s.left_state, s.right_state, s.geometry, s.t, s.γ, x_arr)
end

end

max_stars_repo_name: archermarx/SodShockTube

max_issues_repo_name: archermarx/SodShockTube

max_forks_repo_name: archermarx/SodShockTube


text:
@coeff function Tau(reactor::Reactor)
  K_tau * K_T^1.5 * (log(8/epsilon) - 2) * sqrt(B_0^3 * R_0^7)
end


text:
using REPL
using REPL.Terminals
using REPL.TerminalMenus
using ObjectFile.ELF
using HTTP
import PkgLicenses
using JSON
using MbedTLS
using JLD2

# It's Magic!
export run_wizard

include("wizard/state.jl")
include("wizard/utils.jl")
include("wizard/obtain_source.jl")
include("wizard/interactive_build.jl")
include("wizard/deploy.jl")

# This is here so that if the wizard crashes, we may have a shot at resuming.
wizard_cache_path = joinpath(dirname(dirname(pathof(@__MODULE__))), "deps", "wizard.state")
function save_last_wizard_state(state::WizardState)
    global wizard_cache_path
    jldopen(wizard_cache_path, "w") do f
        serialize(f, state)
    end

    return state
end

function load_last_wizard_state()
    global wizard_cache_path
    try
        state = jldopen(wizard_cache_path, "r") do f
            return unserialize(f)
        end

        if !(state.step in (:done, :step1))
            # Looks like we had an incomplete build; ask the user if they want to continue it
            terminal = TTYTerminal("xterm", state.ins, state.outs, state.outs)
            choice = request(terminal,
                "Would you like to resume the previous incomplete wizard run?",
                RadioMenu([
                    "Resume previous run",
                    "Start from scratch",
                ]),
            )

            if choice == 1
                return state
            else
                return WizardState()
            end
        end
    catch
    end
    # Either something went wrong, or there was nothing interesting stored.
    # Either way, just return a blank slate.
    return WizardState()
end

function run_wizard(state::Union{Nothing,WizardState} = nothing)
    global last_wizard_state

    if state === nothing
        # If we weren't given a state, check to see if we'd like to resume a
        # previous run or start from scratch again.
        state = load_last_wizard_state()
    end

    if state.step == :step1
        print_wizard_logo(state.outs)

        println(state.outs,
            "Welcome to the BinaryBuilder wizard.\n"*
            "We'll get you set up in no time.\n")
    end

    try
        while state.step != :done
            if state.step == :step1
                step1(state)
                state.step = :step2
            elseif state.step == :step2
                step2(state)
                state.step = :step3
            elseif state.step == :step3
                step34(state)
            elseif state.step == :step3_retry
                step3_retry(state)
            elseif state.step == :step5a
                step5a(state)
            elseif state.step == :step5b
                step5b(state)
            elseif state.step == :step5c
                step5c(state)
                state.step = :step6
            elseif state.step == :step6
                step6(state)
            elseif state.step == :step7
                step7(state)
                state.step = :done
            end

            # Save it every step along the way
            save_last_wizard_state(state)
        end
    catch err
        # If anything goes wrong, immediately save the current wizard state
        save_last_wizard_state(state)
        if isa(err, InterruptException)
            msg = "\n\nWizard stopped, use run_wizard() to resume.\n\n"
            printstyled(state.outs, msg, bold=true, color=:red)
        else
            bt = catch_backtrace()
            Base.showerror(stderr, err, bt)
            println(state.outs, "\n")
        end
        return state
    end

    # We did it!
    save_last_wizard_state(state)

    println(state.outs, "\nWizard Complete. Press any key to exit...")
    read(state.ins, Char)

    state
end


text:
"""
    count_word_occurrence_in_string(text::String, word::String)

Count how often word appears in text.
Example: if `text` is "one two one two three four"
and `word` is "one", then this function returns 2
"""
function count_word_occurrence_in_string(text::String, word::String)

    return count(word, text)

end


text:
using Bijectors, LinearAlgebra, Optim, Plots, Random, StatsFuns, Stheno, WarpedRegressors

# Specify a Warped-GP.
f = GP(kernel(EQ(); l=2.0, s=0.5), GPC())
σ² = 0.0001
ϕ = Bijectors.Exp()
g = warp(f, ϕ)

# Generate some training data.
rng = MersenneTwister(123456)
Ntr = 35
xtr = 3 * randn(rng, Ntr)
g_xtr = g(xtr, σ²)
ytr = rand(rng, g_xtr)

# Specify parameter handler.
unpack(θ) = (
    σ² = log1pexp(θ[1]) + 1e-6,
    l  = log1pexp(θ[2]) + 1e-6,
    s  = log1pexp(θ[3]) + 1e-6,
)

# Infer σ², l, and s.
function nlml(θ::AbstractVector{<:Real})
    θ = unpack(θ)
    f = GP(kernel(EQ(); l=θ.l, s=θ.s), GPC())
    g = warp(f, ϕ)
    return -logpdf(g(xtr, θ.σ²), ytr)
end
θ0 = randn(3)
results = Optim.optimize(nlml, θ0, NelderMead())
θ = unpack(results.minimizer)

# Construct posterior Warped GP
f_learned = GP(kernel(EQ(); l=θ.l, s=θ.s), GPC())
g_learned = warp(f_learned, ϕ)
g_posterior = posterior(g_learned(xtr, θ.σ²), ytr)

# Visualise the posterior.
let
    xpr = range(-10.0, 10.0; length=250)
    Ypr = hcat([rand(rng, g_posterior(xpr, θ.σ²)) for _ in 1:1000]...)
    median_pr = map(ys->quantile(ys, 0.5), eachrow(Ypr));
    q95_pr = map(ys->quantile(ys, 0.95), eachrow(Ypr));
    q05_pr = map(ys->quantile(ys, 0.05), eachrow(Ypr));

    plotly()
    plt = plot()
    plot!(plt, xpr, Ypr[:, 1:10]; linecolor=:blue, linewidth=0.5, linealpha=0.2, label="")
    plot!(plt, xpr, median_pr; linecolor=:blue, linewidth=1.0, linealpha=1.0, label="")
    plot!(plt, xpr, [median_pr, median_pr];
        linewidth=0.0,
        fillrange=[q05_pr, q95_pr],
        fillalpha=0.3,
        fillcolor=:blue,
        label="",
    )
    scatter!(plt, xtr, ytr; label="", markersize=2)
    display(plt)
end


text:
abstract type AbstractMeanMethod end
abstract type AbstractExactMeanMethod <: AbstractMeanMethod end
abstract type AbstractTraceEstimationMeanMethod <: AbstractMeanMethod end

struct ExactMean{TF} <: AbstractExactMeanMethod
    F::TF
end
ExactMean() = ExactMean(nothing)

struct ExactSVDMean{TUS} <: AbstractExactMeanMethod
    US::TUS
    n::Int
end
ExactSVDMean() = ExactSVDMean(nothing, 0)
function ExactSVDMean(F::SparseMatrixCSC)
    rows = unique(F.rowval)
    Fc = Matrix(F[rows, :])
    svdfact = svd(Fc)
    threshold = 1e-6
    @show length(svdfact.S)
    inds = findall(x -> x > threshold, svdfact.S)
    @show length(inds)
    US_dense = svdfact.U[:, inds] * Diagonal(svdfact.S[inds])
    I = Int[]
    J = Int[]
    V = eltype(F)[]
    for j in 1:length(inds)
        for i in 1:length(rows)
            push!(I, rows[i])
            push!(J, inds[j])
            push!(V, US_dense[i, j])
        end
    end
    US = sparse(I, J, V, size(F, 1), length(inds))
    return ExactSVDMean(US, size(F, 2))
end

struct TraceEstimationMean{TF,TV,TM} <: AbstractTraceEstimationMeanMethod
    F::TF
    V::TV
    sample_once::Bool
    sample_method::TM
end
function TraceEstimationMean(
    F::SparseMatrixCSC, nv::Int, sample_once::Bool=true, sample_method=hutch_rand!
)
    V = zeros(eltype(F), size(F, 2), nv)
    sample_method(V)
    return TraceEstimationMean(F, V, sample_once, sample_method)
end
struct TraceEstimationSVDMean{TUS,TV,TM} <: AbstractTraceEstimationMeanMethod
    US::TUS
    n::Int
    V::TV
    sample_once::Bool
    sample_method::TM
end
function TraceEstimationSVDMean(
    F::SparseMatrixCSC, nv::Int, sample_once::Bool=true, sample_method=hutch_rand!
)
    US = ExactSVDMean(F).US
    V = zeros(eltype(F), size(US, 2), nv)
    sample_method(V)
    sample_once = sample_method === hadamard! || sample_once
    return TraceEstimationSVDMean(US, size(F, 2), V, sample_once, sample_method)
end
function TraceEstimationSVDMean(
    F::SparseMatrixCSC, V::AbstractMatrix, sample_once::Bool=true, sample_method=hutch_rand!
)
    US = ExactSVDMean(F).US
    sample_once = sample_method === hadamard! || sample_once
    return TraceEstimationSVDMean(US, size(F, 2), V, sample_once, sample_method)
end

abstract type AbstractDiagonalMethod end
abstract type AbstractExactDiagonalMethod <: AbstractDiagonalMethod end
abstract type AbstractDiagonalEstimationMethod <: AbstractDiagonalMethod end

struct ExactDiagonal{TF,TY,Ttemp} <: AbstractExactDiagonalMethod
    F::TF
    Y::TY # K^-1 F
    temp::Ttemp # f' K^-1 dK/d(filtered x_e) K^-1 f for all e and any one f
end
function ExactDiagonal(F::SparseMatrixCSC, nE::Int)
    return ExactDiagonal(F, zeros(eltype(F), size(F)...), zeros(eltype(F), nE))
end

struct ExactSVDDiagonal{TF,TUS,TV,TQ,TY,Ttemp} <: AbstractExactDiagonalMethod
    F::TF
    US::TUS
    V::TV
    Q::TQ # K^-1 US
    Y::TY # Q' dK/d(filtered x_e) Q without the filter
    temp::Ttemp
end
function ExactSVDDiagonal(F::SparseMatrixCSC, nE::Int)
    rows = unique(F.rowval)
    Fc = Matrix(F[rows, :])
    svdfact = svd(Fc)
    threshold = 1e-6
    @show length(svdfact.S)
    inds = findall(x -> x > threshold, svdfact.S)
    @show length(inds)
    US_dense = svdfact.U[:, inds] * Diagonal(svdfact.S[inds])
    V = svdfact.V[:, inds]
    I = Int[]
    J = Int[]
    vals = eltype(F)[]
    for j in 1:length(inds)
        for i in 1:length(rows)
            push!(I, rows[i])
            push!(J, inds[j])
            push!(vals, US_dense[i, j])
        end
    end
    US = sparse(I, J, vals, size(F, 1), length(inds))
    Q = zeros(eltype(F), size(US))
    Y = zeros(eltype(F), length(inds), length(inds))
    temp = zeros(eltype(F), nE)
    return ExactSVDDiagonal(F, US, V, Q, Y, temp)
end

struct DiagonalEstimation{TF,TY,TQ,TV,Ttemp,TM} <: AbstractDiagonalEstimationMethod
    F::TF
    V::TV # all v_i
    Y::TY # K^-1 F v_i for all i
    Q::TQ # K^-1 F * (w .* v_i) for all i
    temp::Ttemp # (w .* v_i)' F' K^-1 dK/d(filtered x_e) K^-1 F v_i for all e and any one v_i
    sample_once::Bool
    sample_method::TM
end
function DiagonalEstimation(
    F::SparseMatrixCSC, nv::Int, nE::Int, sample_once::Bool=true, sample_method=hadamard!
)
    V = zeros(eltype(F), size(F, 2), nv)
    sample_method(V)
    Y = zeros(eltype(F), size(F, 1), nv)
    Q = similar(Y)
    temp = zeros(eltype(F), nE)
    sample_once = sample_method === hadamard! || sample_once
    return DiagonalEstimation(F, V, Y, Q, temp, sample_once, sample_method)
end
function DiagonalEstimation(
    F::SparseMatrixCSC,
    V::AbstractMatrix,
    nE::Int,
    sample_once::Bool=true,
    sample_method=hadamard!,
)
    nv = size(V, 2)
    Y = zeros(eltype(F), size(F, 1), nv)
    Q = similar(Y)
    temp = zeros(eltype(F), nE)
    sample_once = sample_method === hadamard! || sample_once
    return DiagonalEstimation(F, V, Y, Q, temp, sample_once, sample_method)
end

for T in (
    :ExactMean,
    :ExactSVDMean,
    :TraceEstimationMean,
    :TraceEstimationSVDMean,
    :ExactDiagonal,
    :ExactSVDDiagonal,
    :DiagonalEstimation,
)
    @eval baretype(::$T) = $T
end


text:
# common facilities

### value type conversion

f64(x::Real) = Float64(x)

### Common type hierarchy

##
#  The base type for all likelihood model
#
abstract LikelihoodModel

### Auxiliary functions

# half of squared L2-norm
hsqrnorm(x::AbstractArray) = vecnorm(x)^2 / 2


text:
@testset "679.24-game.jl" begin
    @test judgePoint24([1, 2, 1, 2]) == false

    # julia> 8 / (2 - (5 / 3))
    # 24.000000000000007
    @test judgePoint24([8, 2, 5, 3]) == true
end


text:
u = Union{UInt8, Bool}          # A new type
@assert typeof(u) == Union
@assert isa(u, Type)            # A union type is a type

@assert isa(0x08, u)            # UInt8 values belong
@assert isa(false, u)           # Boolean values belong
@assert isa(true, u)
@assert !isa(256, u)            # Not a UInt8, does not belong

max_stars_repo_name: idrougge/ple

max_issues_repo_name: idrougge/ple

max_forks_repo_name: idrougge/ple


text:
using DataStructure101
const DS = DataStructure101

export prims!, kruskals!, mst!

"""
    prims(graph::DS.WeightedAdjacencyList{Float64})

Pirim's algorithm for MST, by using DataStructure101.MimHeap.
Support DataStructure101.WeightedAdjacencyList{Float64} only.

Note: Weights must be unique and Int.

# Examples
```julia-repl
julia> g = DataStructure101Graph(7, :w_list, T=Int64)
julia> DS.relate!(g, 2, 5, 1)
julia> DS.relate!(g, 1, 6, 2)
julia> DS.relate!(g, 1, 4, 3)
julia> DS.relate!(g, 2, 3, 4)
julia> DS.relate!(g, 1, 5, 5)
julia> DS.relate!(g, 1, 2, 6)
julia> DS.relate!(g, 5, 6, 7)
julia> DS.relate!(g, 5, 7, 8)
julia> DS.relate!(g, 2, 7, 9)
julia> DS.relate!(g, 6, 7, 10)
julia> g
WeightedAdjacencyList{Int64}(
    1[6(2), 4(3), 5(5), 2(6), ],
    2[5(1), 3(4), 1(6), 7(9), ],
    3[2(4), ],
    4[1(3), ],
    5[2(1), 1(5), 6(7), 7(8), ],
    6[1(2), 5(7), 7(10), ],
    7[5(8), 2(9), 6(10), ],
)

julia> prims!(g)
WeightedAdjacencyList{Int64}(
    1[6(2), 4(3), 5(5), ],
    2[5(1), 3(4), ],
    3[2(4), ],
    4[1(3), ],
    5[1(5), 2(1), 7(8), ],
    6[1(2), ],
    7[5(8), ],
)
```
"""
function prims!(g::DS.WeightedAdjacencyList{Int64})
    intree = [1, ]
    edges = Dict()
    sorted_weight = DS.MinHeap()

    # loop over all vertex
    for vertex in intree
        # loop over all neighbor
        neighbors = DS.neighbor(g, vertex)
        i=1
        while i <= length(neighbors)
            neighbor = neighbors[i]
            (neighbor in intree) && (i += 1; continue)
            # 1. push weight into MinHeap
            weight = g.weight[vertex][i]
            push!(edges, weight=>(vertex, neighbor))
            push!(sorted_weight, weight)
            # 2. unrelate edges
            DS.unrelate!(g, vertex, neighbor)
        end
        DS.build!(sorted_weight)

        # loop until estebish an edge
        while length(sorted_weight) > 0
            weight = pop!(sorted_weight)
            edge = edges[weight]
            (edge[2] in intree) && continue
            DS.relate!(g, edge[1], edge[2], weight)
            push!(intree, edge[2])
            break
        end
    end

    return g
end

struct Edge
    vertex::Tuple{Int, Int}
    weight::Number
end

function bubble_insert!(edges::Vector{Edge}, edge::Edge)
    len = length(edges)
    for i in 1:len
        (edges[i].weight > edge.weight) && (insert!(edges, i, edge); return)
    end
    push!(edges, edge)
end

function kruskals!(g::DS.WeightedAdjacencyList{Int64})
    forest = DS.DisjointSet(DS.nv(g))
    edges = Edge[]

    # loop over all vertices
    for vertex in 1:(DS.nv(g))
        # loop over all neighbor
        neighbors = DS.neighbor(g, vertex)
        while length(neighbors) > 0
            neighbor = neighbors[1]
            # 1. push weight into sorted array
            bubble_insert!(edges, Edge((vertex, neighbor), g.weight[vertex][1]))
            # 2. unrelate edges
            DS.unrelate!(g, vertex, neighbor)
        end
    end

    # loop over all sorted edges
    for edge in edges
        v1 = edge.vertex[1]
        v2 = edge.vertex[2]
        if DS.find(forest, v1) != DS.find(forest, v2)
            DS.union!(forest, v1, v2)
            DS.relate!(g, v1, v2, edge.weight)
        end
    end

    return g
end

function mst!(g::DS.WeightedAdjacencyList{Int64}; algo::Symbol)
    if algo == :prims
        return prims!(g)
    elseif algo == :kruskals
        return kruskals!(g)
    end
end


text:
function reset_aggregated_jumps!(integrator,uprev = nothing)
     reset_aggregated_jumps!(integrator,uprev,integrator.opts.callback)
     nothing
end

function reset_aggregated_jumps!(integrator,uprev,callback::Nothing)
     nothing
end

function reset_aggregated_jumps!(integrator,uprev,callback::CallbackSet)
    if !isempty(callback.discrete_callbacks)
        reset_aggregated_jumps!(integrator,uprev,callback.discrete_callbacks...)
    end
    nothing
end

function reset_aggregated_jumps!(integrator,uprev,cb::DiscreteCallback,cbs...)
    if typeof(cb.condition) <: AbstractSSAJumpAggregator
        cb.condition(cb,integrator.u,integrator.t,integrator)
    end
    reset_aggregated_jumps!(integrator,uprev,cbs...)
    nothing
end

function reset_aggregated_jumps!(integrator,uprev,cb::DiscreteCallback)
    if typeof(cb.condition) <: AbstractSSAJumpAggregator
        cb.condition(cb,integrator.u,integrator.t,integrator)
    end
    nothing
end


text:
# adverserial interations module
module adv
dir_path = @__DIR__
push!(LOAD_PATH, dir_path * "/../sim") # simulation
push!(LOAD_PATH, dir_path * "/../vis") # visualization
push!(LOAD_PATH, dir_path * "/../dis") # discretization
push!(LOAD_PATH, dir_path * "/../pomdp")
using sim
using vis
using dis
using pomdp

using Serialization
using Printf
using LinearAlgebra

# include(dir * "/adv_weak_v2.jl")
# include(dir * "/adv_strong.jl")
# include(dir * "/adv_medium.jl")

export predict_collision

function replan_adv(agent_self::Agent, world::World, t::Float64=0.0)
	c_closest = pomdp.Collision(Inf, Inf, pomdp.NO_COLLISION)
	cv_closest = [NaN, NaN]
  # Check for collisions with all other agents
  # If collisions are detected then get the closest one to the agent
  for agent in world.agents
    if agent.id == agent_self.id
      continue
    end

		(c, cv) = predict_collision(agent_self.x, agent.x, world, t)

    if c.ctype != pomdp.NO_COLLISION && c.t < c_closest.t
			c_closest = c
			cv_closest = cv
    end
  end

	return (pomdp.sample_adv_a(agent_self.custom[5], c_closest), cv_closest)

end

function empty_controller!(x::AbstractArray{Float64},
													 u::AbstractArray{Float64},
													 dx::AbstractArray{Float64},
													 agent_world::Pair{Agent, World},
													 t::Float64)
	return
end

function parametrize(x::AbstractArray{Float64},
                     world::World)

  (sx, sy) = sim.sp2sxsy(x[1], x[3], world.road.path)
	n = norm([sx, sy])
	sx = sx/n
	sy = sy/n

  ds = x[2]

	th = atan(sy, sx)
  s = mod(x[1], world.road.path.S[end])
  p = x[3]

	v_x = ds*sx
	v_y = ds*sy
  (x_pos, y_pos) = sim.sp2xy(s, p, world.road.path)

  P = [x_pos v_x;
       y_pos v_y]

	return (P, th)
end

function closest_approach(p1::AbstractArray{Float64}, p2::AbstractArray{Float64})
  a = p1[1,1]
  b = p1[1,2]
  d = p1[2,1]
  y = p1[2,2]

  g = p2[1,1]
  h = p2[1,2]
  z = p2[2,1]
  k = p2[2,2]

  t = (-a*b + b*g + a*h - g*h + d*k - d*y - k*z + y*z) / (b^2 - 2*b*h + h^2 + k^2 - 2*k*y + y^2)

  v1 = [a + b*t, d + y*t]
  v2 = [g + h*t, z + k*t]

  return (v1, v2, t)
end

function predict_collision(x_self::Array{Float64},
													 x_opp::Array{Float64},
                           world::World,
                           time::Float64=0.0)
	# Parametrize both agents
	# Compute closest approach
	# Compute Collision angle
	# Return r, t, c_th, collision_type
	# Predicted Collision Types:
	# 0 - No collision detected
	# 1 - agent_self crashes into agent_opp
	# 2 - agent_self is crashed into by agent opp

	# Parametrize both agents
	(P_self, th_self) = parametrize(x_self, world)
	(P_opp, th_opp) = parametrize(x_opp, world)

	# Compute Closest approach collision
	(vc_self, vc_opp, t_c) = closest_approach(P_self, P_opp)

	cv = vc_self
	rv = vc_opp - vc_self

	# Compute collision angle
	th_collision = th_self - atan(rv[2], rv[1])

	# Determine collision type
	collision_type = pomdp.NO_COLLISION

  # Predicted collision criteria
	if norm(rv) < pomdp.max_d && 0.0 < t_c < pomdp.max_t
		if -pi / 2 < th_collision < pi / 2
			collision_type = pomdp.HITTING
		else
			collision_type = pomdp.BEING_HIT
		end
	end

  if collision_type == pomdp.NO_COLLISION
    cv = fill(NaN, length(cv))
  end

	# Compute distance of self to collision
	d = norm(rv)

	return (pomdp.Collision(d, t_c, collision_type), cv)

end

end