text:
import LinearAlgebra.triu

#### function ####

"""
    symmetrize(arr::Array)

Symmetrize the matrix. ``A = (A + A') / 2``
"""
function symmetrize!(arr::Array)
    axpby!(0.5, arr', 0.5, arr)
    copy!(arr, Symmetric(arr, :L))
end

"""
    heaviside(X)

Heaviside's step function.
"""
function heaviside(X)
    return (X .> 0.0) - (X .< 0.0)
end

#### macro ####

macro get(dict, key)
    return esc(quote
        if haskey($dict, $key)
            $dict[$key]
        else
            $nothing
        end
    end)
end

macro getarray(dict, key, shape, default_arr)
    return esc(quote
        if (!haskey($dict, $key) || size($dict[$key]) != $shape)
            $dict[$key] = convert(Array, $default_arr)
        end
        $dict[$key]
    end)
end

#### inspection ####

const OPTIONAL_ITEMS = Dict(
    :alpha => 0.01,
    :beta => 1.0,
    :lr => 1.0,
    :epochs => 20,
    :n_iter => 100 # Option used within `fit` and `annal`
    )

const SAMPLING_ITEMS = Dict(
    :n_sweep => 1000,
    :n_read => 1,
    :beta_min => 5.0, # initial beta
    :beta_max => 15.0 # final beta to stop Simulated Annealing.
    )

"""
    option_warner(opts::Dict)

Display warning if unknown option finded.
"""
function option_warner(opts::Dict, items::Dict)

    OPTIONAL_KEYWARDS = keys(items)

    for opt in keys(opts)
        if !in(opt, OPTIONAL_KEYWARDS)
            @warn("Keyward '$opt' is an invalid optional entry, ignoring")
        end
    end
end

"""
    option_setter!(opts::Dict)

Sets options when a specific key is found.
"""
function option_setter!(opts::Dict, items::Dict)

    OPTIONAL_KEYWARDS = keys(items)

    for key in OPTIONAL_KEYWARDS
        get!(opts, key, items[key])
    end
end

function opt_check(opts::Dict, items::Dict)
    option_warner(opts, items)
    option_setter!(opts, items)
end

option_checker!(opts::Dict{Any, Any}) = opt_check(opts, OPTIONAL_ITEMS)
option_checker!(opts::Dict{Symbol, Real}) = opt_check(opts, SAMPLING_ITEMS)

"""
    variable_check(X::Array)

Check if an element is a spin value: {+1, -1}
"""
function variable_check(X::Array)
    bit_matrix = (X .== 1.0) + (X .== -1.0)
    if !all(Bool.(bit_matrix))
        @warn("Elements of the input array must be spin values: {1, -1}")
    end
end

"""
    elements_match(key_list::Vector)

Check that the types of all elements match.
"""
function elements_match(key_list::Vector)
    flag = all(typeof.(key_list) .== typeof(key_list[1]))
    @assert flag "invalid types of linear and quadratic"
end

#### converter ####

"""
    trans_type(dType::DataType, A::AbstarctArray)

Convert type of matrix elements to `dType`.
"""
function trans_type(dType::DataType, A::AbstractArray)
    (eltype(A) != dType) ? map(dType, A) : A
end

"""
    decode(arr::AbstractArray)

Convert from upper triaugular matrix to dictionary type.

# Example
```jldoctest
julia> a = reshape(1:9, 3, 3)

julia> a
3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64:
 1  4  7
 2  5  8
 3  6  9

julia> decode(a)
OrderedCollections.OrderedDict{Tuple{Int64,Int64},Int64} with 3 entries:
  (1, 2) => 4
  (1, 3) => 7
  (2, 3) => 8
```
"""
function decode(arr::AbstractArray)
    triu_mask = triu(trues(size(arr)), 1)
    idx = CartesianIndices(arr)[triu_mask]
    decoded_data = Dict(zip(Tuple.(idx), arr[idx]))
    return sort(decoded_data, by = x -> x[1])
end

"""
    trans(dict::AbstractDict, len::Int64)

Convert from dictionary `dict` to adjacency matrix with size `len` × `len`.

# Examples
```jldoctest
julia> d = Dict((1, 2) => 5, (2, 3) => -1)

julia> trans(d, 4)
4×4 Array{Float64,2}:
 0.0   5.0   0.0  0.0
 5.0   0.0  -1.0  0.0
 0.0  -1.0   0.0  0.0
 0.0   0.0   0.0  0.0
```
"""
function trans(dict::AbstractDict, len::Int64)

    arr = zeros((len, len))
    for (key, val) in zip(keys(dict), values(dict))
        arr[key...]  = val
        arr[reverse(key)...] = val # symmetric
    end
    return arr
end

"""
    trans(dict::AbstarctDict)

Convert from dictionary `dict` to adjacency matrix.

# Example
```jldoctest
julia> d = Dict((1, 2) => -10, (2, 3) => -5, (4, 2) => 1)

julia> trans(d)
4×4 Array{Float64,2}:
   0.0  -10.0   0.0  0.0
 -10.0    0.0  -5.0  1.0
   0.0   -5.0   0.0  0.0
   0.0    1.0   0.0  0.0
```
"""
function trans(dict::AbstractDict)
    len = maximum(maximum.(keys(dict)))
    trans(dict, len)
end

max_stars_repo_name: oilneck/InverseIsing

max_issues_repo_name: oilneck/InverseIsing

max_forks_repo_name: oilneck/InverseIsing


text:
using YaoBase, SparseArrays, BitBasis, YaoArrayRegister, SymEngine
export @ket_str, @bra_str
export SymReg, AdjointSymReg, SymRegOrAdjointSymReg, expand

YaoArrayRegister._warn_type(raw::AbstractArray{Basic}) = nothing

const SymReg{B, MT} = ArrayReg{B, Basic, MT}
const AdjointSymReg{B, MT} = AdjointArrayReg{B, Basic, MT}
const SymRegOrAdjointSymReg{B, MT} = Union{SymReg{B, MT}, AdjointSymReg{B, MT}}

function ket_m(s)
    v, N = parse_str(s)
    st = spzeros(Basic, 1 << N, 1)
    st[v+1] = 1
    return ArrayReg{1}(st)
end

const MAX_SYM_QUBITS = 10

function Base.show(io::IO, r::SymRegOrAdjointSymReg{1})
    if nqubits(r) < MAX_SYM_QUBITS
        print_sym_state(io, r)
    else
        summary(io, r)
        print(io, "\n    active qubits: ", nactive(r), "/", nqubits(r))
    end
end

Base.:(*)(x::SymReg{B, MT}, y::SymReg{B, MT}) where {B, MT} = SymReg{B, MT}(kron(state(x), state(y)))
Base.:(^)(x::SymReg{B, MT}, n::Int) where {B, MT} = SymReg{B, MT}(kron(state(x) for _ in 1:n))

Base.:(*)(x::AdjointSymReg{B, MT}, y::AdjointSymReg{B, MT}) where {B, MT} = adjoint(parent(x) * parent(y))
Base.:(^)(x::AdjointSymReg{B, MT}, n::Int) where {B, MT} = adjoint(parent(x)^n)

SymEngine.expand(x::SymReg{B}) where B = ArrayReg{B}(expand.(state(x)))

ext: jl


text:
module JuMPExtension
# Simple example of JuMP extension used in the tests to check that JuMP works well with extensions
# The main difference between `JuMP.Model` and `JuMPExtension.MyModel` is the fact that in `add_variable` (resp. `add_constraint`),
# `JuMP.Model` applies the modification to its `moi_backend` field while
# `JuMPExtension.MyModel` stores the `AbstractVariable` (resp. `AbstractConstraint`) in a list.

using MathOptInterface
const MOI = MathOptInterface
import JuMP

struct ConstraintIndex
    value::Int # Index in `model.constraints`
end

mutable struct MyModel <: JuMP.AbstractModel
    nextvaridx::Int                                 # Next variable index is nextvaridx+1
    variables::Dict{Int, JuMP.ScalarVariable}       # Map varidx -> variable
    var_to_name::Dict{Int, String}                  # Map varidx -> name
    name_to_var::Union{Dict{String, Int}, Nothing}  # Map varidx -> name
    nextconidx::Int                                 # Next constraint index is nextconidx+1
    constraints::Dict{ConstraintIndex,
                      JuMP.AbstractConstraint}      # Map conidx -> variable
    con_to_name::Dict{ConstraintIndex, String}      # Map conidx -> name
    name_to_con::Union{Dict{String, ConstraintIndex},
                       Nothing}                     # Map name -> conidx
    objectivesense::MOI.OptimizationSense
    objective_function::JuMP.AbstractJuMPScalar
    obj_dict::Dict{Symbol, Any}                     # Same that JuMP.Model's field `obj_dict`
    function MyModel()
        new(0, Dict{Int, JuMP.AbstractVariable}(),
            Dict{Int, String}(), nothing,                        # Variables
            0, Dict{ConstraintIndex, JuMP.AbstractConstraint}(),
            Dict{ConstraintIndex, String}(), nothing,            # Constraints
            MOI.FEASIBILITY_SENSE,
            zero(JuMP.GenericAffExpr{Float64, MyVariableRef}),
            Dict{Symbol, Any}())
    end
end
Base.broadcastable(model::MyModel) = Ref(model)

JuMP.object_dictionary(model::MyModel) = model.obj_dict

# Variables
struct MyVariableRef <: JuMP.AbstractVariableRef
    model::MyModel # `model` owning the variable
    idx::Int       # Index in `model.variables`
end
Base.copy(v::MyVariableRef) = v
Base.copy(v::MyVariableRef, new_model::MyModel) = MyVariableRef(new_model, v.idx)

Base.:(==)(v::MyVariableRef, w::MyVariableRef) = v.model === w.model && v.idx == w.idx
Base.broadcastable(v::MyVariableRef) = Ref(v)
JuMP.isequal_canonical(v::MyVariableRef, w::MyVariableRef) = v == w
JuMP.variable_type(::MyModel) = MyVariableRef
function JuMP.add_variable(m::MyModel, v::JuMP.AbstractVariable, name::String="")
    m.nextvaridx += 1
    vref = MyVariableRef(m, m.nextvaridx)
    m.variables[vref.idx] = v
    JuMP.set_name(vref, name)
    vref
end
function JuMP.delete(model::MyModel, vref::MyVariableRef)
    @assert JuMP.is_valid(model, vref)
    delete!(model.variables, vref.idx)
    delete!(model.var_to_name, vref.idx)
end
function JuMP.is_valid(model::MyModel, vref::MyVariableRef)
    return (model === vref.model &&
            vref.idx in keys(model.variables))
end
JuMP.num_variables(m::MyModel) = length(m.variables)

# Internal functions
variable_info(vref::MyVariableRef) = vref.model.variables[vref.idx].info
function update_variable_info(vref::MyVariableRef, info::JuMP.VariableInfo)
    vref.model.variables[vref.idx] = JuMP.ScalarVariable(info)
end

JuMP.has_lower_bound(vref::MyVariableRef) = variable_info(vref).has_lb
function JuMP.lower_bound(vref::MyVariableRef)::Float64
    @assert !JuMP.is_fixed(vref)
    return variable_info(vref).lower_bound
end
function JuMP.set_lower_bound(vref::MyVariableRef, lower)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(true, lower,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.delete_lower_bound(vref::MyVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(false, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
JuMP.has_upper_bound(vref::MyVariableRef) = variable_info(vref).has_ub
function JuMP.upper_bound(vref::MyVariableRef)::Float64
    @assert !JuMP.is_fixed(vref)
    return variable_info(vref).upper_bound
end
function JuMP.set_upper_bound(vref::MyVariableRef, upper)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           true, upper,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.delete_upper_bound(vref::MyVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           false, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
JuMP.is_fixed(vref::MyVariableRef) = variable_info(vref).has_fix
function JuMP.fix_value(vref::MyVariableRef)::Float64
    return variable_info(vref).fixed_value
end
function JuMP.fix(vref::MyVariableRef, value; force::Bool = false)
    info = variable_info(vref)
    if !force && (info.has_lb || info.has_ub)
        error("Unable to fix $(vref) to $(value) because it has existing bounds.")
    end
    update_variable_info(vref, JuMP.VariableInfo(
        false, info.lower_bound, false, info.upper_bound, true, value,
        info.has_start, info.start, info.binary, info.integer)
    )
    return
end
function JuMP.unfix(vref::MyVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           false, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.start_value(vref::MyVariableRef)::Union{Nothing, Float64}
    return variable_info(vref).start
end
function JuMP.set_start_value(vref::MyVariableRef, start)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           true, start,
                                           info.binary, info.integer))
end
JuMP.is_binary(vref::MyVariableRef) = variable_info(vref).binary
function JuMP.set_binary(vref::MyVariableRef)
    @assert !JuMP.is_integer(vref)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           true, info.integer))
end
function JuMP.unset_binary(vref::MyVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           false, info.integer))
end
JuMP.is_integer(vref::MyVariableRef) = variable_info(vref).integer
function JuMP.set_integer(vref::MyVariableRef)
    @assert !JuMP.is_binary(vref)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, true))
end
function JuMP.unset_integer(vref::MyVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, false))
end

# Constraints
const MyConstraintRef = JuMP.ConstraintRef{MyModel, ConstraintIndex}
JuMP.constraint_type(::MyModel) = MyConstraintRef
function JuMP.add_constraint(model::MyModel, c::JuMP.AbstractConstraint,
                             name::String="")
    model.nextconidx += 1
    index = ConstraintIndex(model.nextconidx)
    cref = JuMP.ConstraintRef(model, index, JuMP.shape(c))
    model.constraints[index] = c
    JuMP.set_name(cref, name)
    return cref
end
function JuMP.delete(model::MyModel, constraint_ref::MyConstraintRef)
    @assert JuMP.is_valid(model, constraint_ref)
    delete!(model.constraints, constraint_ref.index)
    delete!(model.con_to_name, constraint_ref.index)
end
function JuMP.is_valid(model::MyModel, constraint_ref::MyConstraintRef)
    return (model === constraint_ref.model &&
            constraint_ref.index in keys(model.constraints))
end
function JuMP.constraint_object(cref::MyConstraintRef)
    return cref.model.constraints[cref.index]
end

# Objective
function JuMP.set_objective(m::MyModel, sense::MOI.OptimizationSense,
                            f::JuMP.AbstractJuMPScalar)
    m.objectivesense = sense
    m.objective_function = f
end
function JuMP.set_objective(m::MyModel, sense::MOI.OptimizationSense, f::Real)
    m.objectivesense = sense
    m.objective_function = JuMP.GenericAffExpr{Float64, MyVariableRef}(f)
end
JuMP.objective_sense(model::MyModel) = model.objectivesense
function JuMP.set_objective_sense(model::MyModel, sense)
    model.objectivesense = sense
end
JuMP.objective_function_type(model::MyModel) = typeof(model.objective_function)
JuMP.objective_function(model::MyModel) = model.objective_function
function JuMP.objective_function(model::MyModel, FT::Type)
    # InexactError should be thrown, this is needed in `objective.jl`
    if !(model.objective_function isa FT)
        throw(InexactError(:objective_function, FT,
                           typeof(model.objective_function)))
    end
    return model.objective_function::FT
end

# Names
JuMP.name(vref::MyVariableRef) = vref.model.var_to_name[vref.idx]
function JuMP.set_name(vref::MyVariableRef, name::String)
    vref.model.var_to_name[vref.idx] = name
    vref.model.name_to_var = nothing
end
function JuMP.variable_by_name(model::MyModel, name::String)
    if model.name_to_var === nothing
        # Inspired from MOI/src/Utilities/model.jl
        model.name_to_var = Dict{String, Int}()
        for (var, var_name) in model.var_to_name
            if haskey(model.name_to_var, var_name)
                # -1 is a special value that means this string does not map to
                # a unique variable name.
                model.name_to_var[var_name] = -1
            else
                model.name_to_var[var_name] = var
            end
        end
    end
    index = get(model.name_to_var, name, nothing)
    if index isa Nothing
        return nothing
    elseif index == -1
        error("Multiple variables have the name $name.")
    else
        return MyVariableRef(model, index)
    end
end
JuMP.name(cref::MyConstraintRef) = cref.model.con_to_name[cref.index]
function JuMP.set_name(cref::MyConstraintRef, name::String)
    cref.model.con_to_name[cref.index] = name
    cref.model.name_to_con = nothing
end
function JuMP.constraint_by_name(model::MyModel, name::String)
    if model.name_to_con === nothing
        # Inspired from MOI/src/Utilities/model.jl
        model.name_to_con = Dict{String, ConstraintIndex}()
        for (con, con_name) in model.con_to_name
            if haskey(model.name_to_con, con_name)
                # -1 is a special value that means this string does not map to
                # a unique constraint name.
                model.name_to_con[con_name] = ConstraintIndex(-1)
            else
                model.name_to_con[con_name] = con
            end
        end
    end
    index = get(model.name_to_con, name, nothing)
    if index isa Nothing
        return nothing
    elseif index.value == -1
        error("Multiple constraints have the name $name.")
    else
        # We have no information on whether this is a vector constraint
        # or a scalar constraint
        return JuMP.ConstraintRef(model, index, JuMP.ScalarShape())
    end
end

# Show
function JuMP.show_backend_summary(io::IO, model::MyModel) end
function JuMP.show_objective_function_summary(io::IO, model::MyModel)
    println(io, "Objective function type: ",
            JuMP.objective_function_type(model))
end
function JuMP.objective_function_string(print_mode, model::MyModel)
    return JuMP.function_string(print_mode, JuMP.objective_function(model))
end
_plural(n) = (isone(n) ? "" : "s")
function JuMP.show_constraints_summary(io::IO, model::MyModel)
    n = length(model.constraints)
    print(io, "Constraint", _plural(n), ": ", n)
end
function JuMP.constraints_string(print_mode, model::MyModel)
    strings = String[]
    # Sort by creation order, i.e. ConstraintIndex value
    constraints = sort(collect(model.constraints), by = c -> c.first.value)
    for (index, constraint) in constraints
        push!(strings, JuMP.constraint_string(print_mode, constraint))
    end
    return strings
end

end


text:
# using Parameters, Interpolations, Plots, Test
# using LinearAlgebra, ToeplitzMatrices, FileIO, JLD2
#
# abstract type Field end
# abstract type Vortex end
# abstract type CoreShape end
#
# @with_kw mutable struct Torus <: Field
#     ψ::Array{Complex{Float64},2}
#     x::Vector{Float64}
#     y::Vector{Float64}
# end
#
# @with_kw mutable struct Sphere <: Field
#     ψ::Array{Complex{Float64},2}
#     x::Vector{Float64}
#     y::Vector{Float64}
# end
#
# @with_kw mutable struct PointVortex <: Vortex
#     xv::Float64
#     yv::Float64
#     qv::Int64
# end
#
# PointVortex(v::Array{Float64,2}) = PointVortex.(v[:,1],v[:,2],v[:,3])
# rawData(v::PointVortex) = [v.xv v.yv v.qv]
# rawData(v::Array{PointVortex,1}) = reduce(vcat,rawData.(v))
#
# function uniform(a,b)
#     @assert a<b
#     return a + (b-a)*rand()
# end
#
# function uniform(a,b,n)
#     return uniform.(a*ones(n),b)
# end
#
# uniform() = uniform(-.5,.5)
# uniform(n) = uniform(-.5,.5,n)
#
# randcharge() = rand([-1 1],)
# randcharge(n) = rand([-1 1],n)
# randPointVortex(n) = PointVortex.(uniform(n), uniform(n), randcharge(n))
# randPointVortex() = PointVortex(uniform(),uniform(),randcharge())
#
# function randPointVortex(n,psi::Field)
#     @unpack ψ,x,y = psi
#     dx = x[2]-x[1]; dy = y[2]-y[1]
#     xi,xf = x[1]+2dx,x[end]-2dx
#     yi,yf = y[1]+2dy,y[end]-2dy
#     return PointVortex.(uniform(xi,xf,n),uniform(yi,yf,n),randcharge(n))
# end
#
# randPointVortex(psi::Field) = randPointVortex((),psi::Field)
#
# # Detection
# include("../src/detection.jl")
#
# function findwhere(A)
#     I = findall(!iszero,A)
#     v = A[I]
#     ix = [I[i][1] for i in eachindex(I)]
#     iy = [I[i][2] for i in eachindex(I)]
#     return ix,iy,v
# end
#
# function findvortices_jumps(psi::Field;shift=true)
#     @unpack x,y,ψ = psi
#     phase = angle.(ψ)
#
#     # Compute all nearest neighbour jumps
#     diffx = phasejumps(phase,1); diffy = phasejumps(phase,2)
#
#     # Compute all plaquette loop integrals with in-place memory recycling
#     circshift!(phase,diffx,(0,1))
#     diffx .-= phase; diffx .-= diffy
#     circshift!(phase,diffy,(1,0))
#     diffx .+= phase
#
#     # Find all windings
#     ixp,iyp,vp = findwhere(diffx .> 0.0)
#     xp = x[ixp]; yp = y[iyp]; np = length(vp)
#     ixn,iyn,vn = findwhere(diffx .< 0.0)
#     xn = x[ixn]; yn = y[iyn]; nn = length(vn)
#
#     if shift
#         dx = x[2]-x[1]; dy = y[2] - y[1]
#         xp .-= dx/2; yp .-= dy/2; xn .-= dx/2; yn .-= dy/2
#     end
#
#     vortices = [xn yn -vn; xp yp vp]
#
#     return PointVortex(vortices)
# end
#
# function findvortices_grid(psi::Torus;shift=true)
#     vort = findvortices_jumps(psi,shift=shift)
#     rawvort = rawData(vort)
#     vort = sortslices(rawvort,dims=1)
#     return PointVortex(vort)
# end
#
# function findvortices_grid(psi::Sphere;shift=true)
#     @unpack ψ = psi
#     windvals = phasejumps(angle.(ψ),2)
#     vort = findvortices_jumps(psi,shift=shift)
#     rawvort = rawData(vort)
#
#     # North pole winding. - sign means polar vortex co-rotating with w > 0
#     w1 = sum(windvals[1,:])
#     (sign(w1) != 0) && (rawvort = [rawvort; 0.0 0.0 -w1])
#     # South pole winding. + sign means co-rotating with w > 0
#     w2 = sum(windvals[end,:])
#     (sign(w2) != 0) && (rawvort = [rawvort; pi 0.0 w2])
#
#     vort = sortslices(vort,dims=1)
#     return PointVortex(vort)
# end
#
# function findvortices_interp(psi::Field)
#     vort = findvortices_grid(psi,shift=true)
#     vort = removeedgevortices(vort,psi)
#     #TODO: allow for interp with periodic data (here edges are stripped)
#
#     for (j,vortex) in enumerate(vort)
#         try
#         vortz,psiz = corezoom(vortex,psi)
#         vortz,psiz = corezoom(vortz,psiz)
#         vortz,psiz = corezoom(vortz,psiz)
#         vort[j] = vortz
#         catch nothing
#         end
#         return vort
#     end
# end
#
# """
#     vortices = findvortices(psi::T,interp=true) where T<:Field
#
# Locates vortices as 2π phase windings around plaquettes on a cartesian spatial grid. Uses an optimized plaquette method followed by recursive interpolation.
#
# Requires a 2D wavefunction ψ(x,y) on a cartesian grid specified by vectors x, y.
#
# `vortices` - array of vortex coordinates `xv,yv` and charges `qv`. Each row is of the form `[xv, yv, cv]`, and the array is sorted into lexical order according to the `xv` coordinates
# """
# function findvortices(psi::Field,interp=true)
#     if interp
#         return findvortices_interp(psi)
#     else
#         return findvortices_grid(psi)
#     end
# end
#
#
# """
# `vortices = removeedgevortices(vort::Array{PointVortex,1},x,y,edge=1)`
#
# Strips edgevortices due to periodic phase differencing."""
# function removeedgevortices(vort::Array{PointVortex,1},psi::Field,edge=1)
#     @unpack x,y = psi; dx=x[2]-x[1]; dy=y[2]-y[1]
#     keep = []
#     for j = 1:length(vort)
#         xi,yi,qi = rawData(vort[j])
#         xedge = isapprox(xi,x[1],atol=edge*dx) || isapprox(xi,x[end],atol=edge*dx)
#         yedge = isapprox(yi,y[1],atol=edge*dy) || isapprox(yi,y[end],atol=edge*dy)
#         not_edge = !(xedge || yedge)
#         not_edge && push!(keep,j)
#     end
#     return vort[keep]
# end
#
# """
#     vortz,psiz = corezoom(vortex,psi::T,winhalf=2,Nz=30) where T <: Field
#
# Uses local interpolation to resolve core location.
# """
# function corezoom(vortex::PointVortex,psi::T,winhalf=2,Nz=30) where T<:Field
#     @unpack ψ,x,y = psi
#     xv,yv,qv = rawData(vortex)
#     dx=x[2]-x[1]; dy=y[2]-y[1]
#     ixv = isapprox.(x,xv,atol=dx) |> findlast
#     iyv = isapprox.(y,yv,atol=dy) |> findlast
#     ixwin = (ixv-winhalf):(ixv+winhalf-1)
#     iywin = (iyv-winhalf):(iyv+winhalf-1)
#     xw = x[ixwin]; yw = y[iywin]; psiw = ψ[ixwin,iywin]
#     xz = LinRange(xw[1],xw[end],Nz)
#     yz = LinRange(yw[1],yw[end],Nz)
#     knots = (xw,yw)
#     itp = interpolate(knots, psiw, Gridded(Linear()))
#     psiz = itp(xz,yz)
#     ψv = T(psiz,xz |> Vector,yz |> Vector)
#     vortz = findvortices_grid(ψv,shift=false)
#     vortz = removeedgevortices(vortz,ψv)[1]
#     return vortz,ψv
# end
#
# corezoom(vortex::Array{PointVortex,1},psi::Field,winhalf=2,Nz=30) = corezoom(vortex[1],psi,2,30)

# Masking
# function findvortexmask(psi::T,R) where T<: Field
#     @unpack x, y, ψ = psi
#     ψ = circmask(psi, 1.1 * R)
#     ϕ = T(x, y, ψ)
#     vortices = findvortices(ϕ)
#
#     # TODO more tests
#     # remove vortices found outside mask boundary
#         for (i, xv) in enumerate(vortices[:,1]), yv in vortices[i,2]
#             (norm([xv,yv]) > R) && (vortices[i,:] = [0. 0. 0.])
#         end
#
#         vortfound = Array{Float64,2}[]
#         for (i, sv) in enumerate(vortices[:,3])
#             (sv != 0) && push!(vortfound, vortices[i,:]')
#         end
#     # currently will fail for more that one vortex (?)
#     return vortfound[1]
# end

# Helpers
# linspace(a,b,n) = LinRange(a,b,n) |> collect
#
# function edgemask!(psi::T) where T<:Field
#     @unpack ψ,x,y = psi
#     ψ[:,1] = zero(x) |> complex
#     ψ[:,end] = zero(x) |> complex
#     ψ[1,:] = zero(y') |> complex
#     ψ[end,:] = zero(y') |> complex
#     return T(ψ,x,y)
# end
#
# function circmask(psi::T,R) where T<:Field
#     @unpack ψ,x,y = psi
#     for j in eachindex(x), k in eachindex(y)
#         (x[j]^2+y[k]^2 > R^2) && (ψ[j,k] = complex(0.))
#     end
#     return T(x,y,ψ)
# end

# function circmask!(phi::T,psi::T,R) where T<: Field
#     @unpack ψ,x,y = psi
#     phi.ψ .= psi.ψ
#     for j in eachindex(x), k in eachindex(y)
#             (x[j]^2+y[k]^2 > R^2) && (phi.ψ[j,k] = complex(0.))
#     end
#     return nothing
# end
#
# function isinterior(a,b,x,y)
#     Lx = x[end]-x[1]; Ly = y[end] - y[1]
#     dx = x[2] - x[1]; dy = y[2] - y[1]
#     return (-Lx/2 + 2dx < a < Lx/2 - 2dx && -Ly/2 + 2dy < b < Ly/2 - 2dy)
# end

# #TODO update this function
# function checkvortexlocations(testvort,vortices,x,y,Nv)
# #check detection to 2 x grid resolution
#     dx = x[2] - x[1]; dy = y[2] - y[1]
#     vortfound = 0
#     for j=1:Nv
#         if (isapprox(testvort[j,1],vortices[j,1],atol=2dx) && isapprox(testvort[j,2],vortices[j,2],atol=2dy))
#         vortfound+=1
#         end
#     end
#     return vortfound
# end

# Creation
# include("../src/creation.jl")
# scalaransatz(x) = sqrt(x^2/(1+x^2))

using Parameters, Interpolations, Plots, Test
using LinearAlgebra, ToeplitzMatrices, FileIO, JLD2

include("../src/types.jl")
include("../src/pointvortex.jl")
include("../src/detection.jl")
include("../src/creation.jl")

# y,ψ,res = gpecore_exact(1)
# ψa = interpolate(tuple(y[1:99]), scalaransatz.(y[1:99]), Gridded(Linear()))
#
# ψi = interpolate(tuple(y[1:99]), ψ[1:99], Gridded(Linear()))
# @save "./src/ansatzcore.jld2" ψa
# @save "./src/exactcore.jld2" ψi

# @load "./src/exactcore.jld2" ψi
# @load "./src/ansatzcore.jld2" ψa

# r(x,y) = sqrt(x^2+y^2)
#
# struct Ansatz <: CoreShape
#     f::Interpolations.GriddedInterpolation{Float64,1,Float64,Gridded{Linear},Tuple{Array{Float64,1}}}
#     ξ::Float64
#     Λ::Float64
# end
#
# Ansatz() = Ansatz(ψa,1.0,0.8249)
#
# function (core::Ansatz)(x)
#     @unpack f,ξ,Λ = core
#     return f(Λ*x/ξ)
# end
# (core::Ansatz)(x,y) = core(r(x,y))
#
# struct Exact <: CoreShape
#     f::Interpolations.GriddedInterpolation{Float64,1,Float64,Gridded{Linear},Tuple{Array{Float64,1}}}
#     ξ::Float64
# end
# Exact() = Exact(ψi,1.0)
#
# function (core::Exact)(x)
#     @unpack ξ,f = core
#     return f(x/ξ)
# end
# (core::Exact)(x,y) = core(r(x,y))
#
# struct ScalarVortex{T<:CoreShape} <: Vortex
#     core::T
#     vort::PointVortex
# end
# ScalarVortex(vort::PointVortex) = ScalarVortex(Exact(),vort)
#
# function (s::ScalarVortex{T})(x,y) where T <: CoreShape
#     @unpack xv,yv,qv = s.vort
#     return s.core(x - xv, y - yv)*exp(im*qv*atan(y - yv,x - xv))
# end
#
# randPointVortex() = PointVortex(uniform(),uniform(),randcharge())
# randPointVortex(n) = PointVortex.(uniform(n), uniform(n), randcharge(n))
# function randPointVortex(psi::Field)
#     @unpack x,y = psi; a,b = first(x),last(x); c,d = first(y), last(y)
#     return PointVortex(uniform(a,b,),uniform(c,d,),randcharge())
# end
# function randPointVortex(n,psi::Field)
#     @unpack x,y = psi;
#     dx = x[2]-x[1]; dy = y[2] - y[1]
#     a,b = first(x)+2dx,last(x)-2dx; c,d = first(y)+2dy, last(y)-2dy
#     return PointVortex.(uniform(a,b,n),uniform(c,d,n),randcharge(n))
# end
#
# randScalarVortex() = ScalarVortex(randPointVortex())
# randScalarVortex(n) = ScalarVortex.(randPointVortex(n))
# randScalarVortex(psi::Field) = ScalarVortex(Exact(),randPointVortex(psi))
# randScalarVortex(n,psi::Field) = ScalarVortex.([Exact()],randPointVortex(n,psi))
#
# randVortex() = randScalarVortex()
# randVortex(n) = randScalarVortex(n)
# randVortex(n,psi::Field) = randScalarVortex(n,psi)
#
# function vortex!(psi::F,vort::ScalarVortex{T}) where {T <: CoreShape, F<:Field}
#     @unpack ψ,x,y = psi
#     @. ψ *= vort.(x,y')
#     @pack! psi = ψ
# end
#
# function vortex!(psi::F,vort::Array{S}) where {F <: Field, S <: Vortex}
#     for j in eachindex(vort)
#         vortex!(psi,vort[j])
#     end
# end

#=
# tests
x = LinRange(0,1,500)
ansatz = Ansatz(ψa,1.,0.8249)
ansatz(.1)
@time ansatz.(x)
@time ansatz.(x,x')
@time Ansatz(ψa,1.2,0.8249).(x,x')

exact = Exact()

exact(.1)
@time exact.(x)
@time exact.(x,x')
@time Exact(ψi,1.2).(x,x')
@time ψi.(x)
@time scalaransatz.(x)

vort1 = randPointVortex()
scalar_exact = ScalarVortex(exact,vort1)
scalar_ansatz = ScalarVortex(ansatz,vort1)
scalar_test = ScalarVortex(vort1)

scalar_ansatz.(x,x')
scalar_exact.(x,x')
scalar_test.(x,x')
@test ScalarVortex(vort1).(x,x') == scalar_test.(x,x')

using Plots
gr(transpose=true,aspectratio=1)
N = 200
L = 100.0
x = LinRange(-L/2,L/2,N)
y = x

psi0 = one.(x*y') |> complex; psi = Torus(psi0,x,y)
vtest = randPointVortex(psi)
vort1 = ScalarVortex(exact,vtest)
vort2 = ScalarVortex(Exact(),vtest)

vort3 = ScalarVortex(vtest)
psi0 = one.(x*y') |> complex; psi = Torus(psi0,x,y)
vortex!(psi,vort3)
@unpack ψ = psi
heatmap(x,y,angle.(ψ))
heatmap(x,y,abs2.(ψ))

vfound = findvortices(psi)

vort4 = ScalarVortex(Ansatz(),vtest)
psi0 = one.(x*y') |> complex
psi = Torus(psi0,x,y)
vortex!(psi,vort4)
@unpack ψ = psi
heatmap(x,y,angle.(ψ))
heatmap(x,y,abs2.(ψ))

# timing
x = LinRange(0, 6, 300)

vort4.(x,x')
@time vort4.(x,x')
@time vort3.(x,x')
@time vort2.(x,x')

# ==== old tests

# exact core
# y,ψ,res = gpecore_exact(1)
# @load "./src/exactcore.jld2" ψi

# Testing
gr(transpose=true,aspectratio=1)
N = 200
L = 100.0
x = LinRange(-L/2,L/2,N)
y = x

psi0 = one.(x*y') |> complex
psi = Torus(psi0,x,y)
vtest = randVortex(psi)

@time vortex!(psi,vtest)
@unpack ψ = psi
heatmap(x,y,angle.(ψ))
heatmap(x,y,abs2.(ψ))

# Testing
gr(transpose=true,aspectratio=1)
N = 200
L = 100.0
x = LinRange(-L/2,L/2,N)
y = x

psi0 = one.(x*y') |> complex
psi = Torus(psi0,x,y)
vtest = randVortex(psi)
vortex!(psi,vtest)
@unpack ψ = psi
heatmap(x,y,angle.(ψ))
heatmap(x,y,abs2.(ψ))

# test detection
vgridraw = findvortices_grid(psi,shift=false)
vgridraw = removeedgevortices(vgridraw,psi)
scatter!([rawData(vgridraw)[1]],[rawData(vgridraw)[2]])

vgrid = findvortices_grid(psi)
vgrid = removeedgevortices(vgrid,psi)
scatter!([rawData(vgrid)[1]],[rawData(vgrid)[2]])
=#


text:
export TimeStamp, PreciseTime

struct TimeStamp end
struct PreciseTime end

#=
timeStamp:
- Julia version:
- Author: Dan
- Date: 2021-06-02
=#


text:
# replicate the first experiment from Optimal Best Arm Identification
# with Fixed Confidence (Garivier and Kaufmann 2016).
#
# running on my laptop this takes
# 931.314152 seconds (9.09 M allocations: 1.541 GiB, 0.03% gc time)

using JLD2;
using Distributed;
using Printf;

@everywhere include("runit.jl");
@everywhere include("thresholds.jl");
include("experiment_helpers.jl");

# setup

dist = Gaussian();
μ = [.9, .7, .5];
pep = BestArm(dist);

srs = [TrackAndStop(CTracking)];

δs = (0.1,);
βs = GK16.(δs); # Recommended in section 6 of paper

N = 10;
seed = 1234;

# compute

@time data = pmap(
    ((sr,i),) -> runit(seed+i, sr, μ, pep, βs),
    Iterators.product(srs, 1:N)
);

dump_stats(pep, μ, δs, βs, srs, data);

# save

@save isempty(ARGS) ? "experiment_bai1.dat" : ARGS[1]  dist μ pep srs data δs βs N seed

# visualise by loading viz_bai1.jl


text:
module InlineTest

@static if isdefined(Base, :Experimental) &&
             isdefined(Base.Experimental, Symbol("@optlevel"))
    Base.Experimental.@optlevel 0
end

export @testset, @testset_macro

# from Test:
export Test,
    @test, @test_throws, @test_broken, @test_skip,
    @test_warn, @test_nowarn,
    @test_logs, @test_deprecated,
    @inferred,
    detect_ambiguities, detect_unbound_args

using Test: Test,
    @test, @test_throws, @test_broken, @test_skip,
    @test_warn, @test_nowarn,
    @test_logs, @test_deprecated,
    @inferred,
    detect_ambiguities, detect_unbound_args

# it seems difficult to use gensym here, as other gensym calls (from other
# modules) can lead to collisions or mismatches, whether we statically
# initialize INLINE_TEST or at __init__ time; so let's use a long enough
# static random hexadecimal string
const INLINE_TEST = Symbol("##InlineTest-01b48f5c342f65df7fcd07f28f0d2cacbb09f0a0")

const TESTED_MODULES = Union{Module,Nothing}[]
const TESTSET_MACROS = Symbol[]

get_tests(m::Module) = getfield(m, INLINE_TEST).tests

function register(m::Module, macros::Vector{Symbol})
    push!(TESTED_MODULES, m)
    foreach(register_macro, macros)
end

register_macro(mac::Symbol) = mac ∉ TESTSET_MACROS && push!(TESTSET_MACROS, mac)
# uniquify just in case, to keep TESTSET_MACROS as short as possible

function retest end

"""
    @testset args...

Similar to `Test.@testset args...`, but the contained tests are not run
immediately, and are instead stored for later execution, triggered by
[`retest()`](@ref) or `runtests()`.

Besides the `@testset` body (last argument) and a description string,
arguments of `@testset` can be:
* the `verbose` option, with a *literal* boolean value (e.g.
  `verbose=true`)
* a literal symbol, serving as a label which can be used for
  testset filtering (see [`retest`](@ref)'s docstring for details).
  All nested testsets inherit such labels.

Invocations of `@testset` can be nested, but qualified invocations of
`ReTest.@testset` can't.

A `@testset` can contain a nested `Test.@testset`, or call a function
which defines a `Test.@testset`: in this case, the `Test.@testset`
will be run whenever the parent testset is run, but `retest` won't
know about it: it won't be taken into account
during the filtering phase, and won't be printed in dry mode.

Internally, `@testset` expressions are converted to an equivalent of
`Test.@testset` at execution time.
"""
macro testset(args...)
    let mod = __module__
        while !isdefined(mod, :runtests)
            @eval mod begin
                """
                    $($mod).runtests(pattern...; kwargs...)

                Equivalent to `ReTest.retest($($mod), pattern...; kwargs...)`.
                This function is defined automatically in any module containing
                a `@testset`, possibly nested within submodules.
                """
                function runtests(specs...; kwargs...)
                    $retest($mod, specs...; kwargs...)
                end
            end
            mod = parentmodule(mod)
        end
    end
    get_inline_mod!(__module__)
    # this must take effect at compile/run time rather than parse time, e.g.
    # if the @testset if in a `if false` branch
    # TODO: test that
    quote
        push!(get_tests($__module__).news, (ts=$args, source=$(QuoteNode(__source__))))
        nothing
    end
end

function get_inline_mod!(mod)::Module
    if isdefined(mod, INLINE_TEST)
        getfield(mod, INLINE_TEST)
    else
        # Credit to Takafumi Arakaki for the idea of creating a submodule
        # within modules using InlineTest in order to be able to define
        # an `__init__` function which does the registration
        # (as we can't directly create/modify `__init__` for the given module).
        @eval mod module $INLINE_TEST
            const tests = (tests=[], news=[], map=Dict{Union{String,Expr},Int}())
            const TESTSET_MACROS = Symbol[]
            __init__() = $register($mod, TESTSET_MACROS)
        end
    end
end

"""
    @testset_macro @mac

Declare `@mac` as a macro which must be expanded statically by `retest` so that
contained `@testset`s can be discovered.

Consider this pattern with `Test` which factors out testsets in a function:

```julia
using Test

function test_iseven(x)
    @testset "iseven \$x" begin
        @test iseven(x)
    end
end

@testset "test \$x" for x=2:2:4
    test_iseven(x)
end
```

This doesn't translate directly with `ReTest`, as the call to `test_iseven` will be
performed at run-time, and will end up declaring a new `@testset "iseven \$x"` at
toplevel (this is a problem similar to having `include` inside testsets).
So on the first run of `retest()`, no `@test` is run, and on the second one,
it fails because `x` is not defined at global scope.

The alternative is to turn `test_iseven` into a macro and declare it with `@testset_macro`:

```julia
using ReTest

macro test_iseven(x)
    quote
        @testset "iseven \$(\$x)" begin
            @test iseven(\$x)
        end
    end
end

@testset_macro @test_iseven

@testset "test \$x" for x=2:2:4
    @test_iseven(x)
end
```
Then, running `retest("iseven", verbose=2)` gives:
```
                    Pass
test 2          |      1
  iseven 2      |      1
test 4          |      1
  iseven 4      |      1
Main            |      2
```
"""
macro testset_macro(mac)
    Meta.isexpr(mac, :macrocall, 2) ||
        error("usage: @testset_macro @macro_to_be_registered")
    inline_mod = get_inline_mod!(__module__)
    mac = QuoteNode(mac.args[1]::Symbol)
    quote
        push!($(inline_mod.TESTSET_MACROS), $mac)
        # need also to register right away, e.g. in case __init__ was already run
        InlineTest.register_macro($mac)
        nothing
    end
end

end # InlineTest

max_issues_repo_name: rfourquet/InlineTest

max_forks_repo_name: rfourquet/InlineTest


text:
module TOML

    # -------------------------------------------------------------------------
    # =========================================================================
    # -------------------------------------------------------------------------

    # All TOML errors inherit from the abstract Error
    abstract Error <: Exception

    # Errors that arise in parsing input
    type ParserError <: Error
        message::String
    end
    Base.showerror(io::IO, e::ParserError) = print(io,
        join(["TOML.ParserError: ", e.message]))

    # Errors that arise in Table structure
    type TableError <: Error
        message::String
    end
    Base.showerror(io::IO, e::TableError) = print(io,
        join(["TOML.TableError: ", e.message]))

    # Errors that arise in handling types within TOML
    type TypeConversionError <: Error
        message::String
    end
    Base.showerror(io::IO, e::TypeConversionError) = print(io,
        join(["TOML.TypeConversionError: ", e.message]))

    # Errors that arise in dealing with Values
    type ValueError <: Error
        message::String
    end
    Base.showerror(io::IO, e::ValueError) = print(io,
        join(["TOML.ValueError: ", e.message]))

    # -------------------------------------------------------------------------
    # =========================================================================
    # -------------------------------------------------------------------------

    # Check if the character is a valid letter (a-zA-Z).
    function is_letter(c::Char)
        if c == 'a' || c == 'A'
            return(true)
        elseif c == 'b' || c == 'B'
            return(true)
        elseif c == 'c' || c == 'C'
            return(true)
        elseif c == 'd' || c == 'D'
            return(true)
        elseif c == 'e' || c == 'E'
            return(true)
        elseif c == 'f' || c == 'F'
            return(true)
        elseif c == 'g' || c == 'G'
            return(true)
        elseif c == 'h' || c == 'H'
            return(true)
        elseif c == 'i' || c == 'I'
            return(true)
        elseif c == 'j' || c == 'J'
            return(true)
        elseif c == 'k' || c == 'K'
            return(true)
        elseif c == 'l' || c == 'L'
            return(true)
        elseif c == 'm' || c == 'M'
            return(true)
        elseif c == 'n' || c == 'N'
            return(true)
        elseif c == 'o' || c == 'O'
            return(true)
        elseif c == 'p' || c == 'P'
            return(true)
        elseif c == 'q' || c == 'Q'
            return(true)
        elseif c == 'r' || c == 'R'
            return(true)
        elseif c == 's' || c == 'S'
            return(true)
        elseif c == 't' || c == 'T'
            return(true)
        elseif c == 'u' || c == 'U'
            return(true)
        elseif c == 'v' || c == 'V'
            return(true)
        elseif c == 'w' || c == 'W'
            return(true)
        elseif c == 'x' || c == 'X'
            return(true)
        elseif c == 'y' || c == 'Y'
            return(true)
        elseif c == 'z' || c == 'Z'
            return(true)
        else
            return(false)
        end
    end # function is_letter

    # =========================================================================

    # Check if the character is a valid digit (0-9).
    function is_digit(c::Char)
        if c == '0'
            return(true)
        elseif c == '1'
            return(true)
        elseif c == '2'
            return(true)
        elseif c == '3'
            return(true)
        elseif c == '4'
            return(true)
        elseif c == '5'
            return(true)
        elseif c == '6'
            return(true)
        elseif c == '7'
            return(true)
        elseif c == '8'
            return(true)
        elseif c == '9'
            return(true)
        else
            return(false)
        end
    end # function is_digit

    # =========================================================================

    # Convert valid digit characters to the correct number
    function to_digit(c::Char)
        if c == '0'
            return(0::Int64)
        elseif c == '1'
            return(1::Int64)
        elseif c == '2'
            return(2::Int64)
        elseif c == '3'
            return(3::Int64)
        elseif c == '4'
            return(4::Int64)
        elseif c == '5'
            return(5::Int64)
        elseif c == '6'
            return(6::Int64)
        elseif c == '7'
            return(7::Int64)
        elseif c == '8'
            return(8::Int64)
        elseif c == '9'
            return(9::Int64)
        else
            throw(ParserError(join(["Character \"", string(c),
                "\" is not a digit."])))
        end
    end # function to_digit

    # -------------------------------------------------------------------------
    # =========================================================================
    # -------------------------------------------------------------------------

    # eat whitespace at start of string and return what's left
    function consume_whitespace(s::String)
        idx = 1
        while idx <= length(s) && (s[idx] == ' ' || s[idx] == '\t')
            idx += 1
        end
        return(s[idx:end])
    end # function consume_whitespace

    # =========================================================================

    # eat the specified character or throw an error
    function consume_character(c::Char, s::String)
        if s[1] == c
            return(s[2:end])
        else
            throw(ParserError(join(["Expected \"", c, "\" at start of \"", s,
                "\"."])))
        end
    end

    # =========================================================================

    function consume_to_eol(s::String)
        remain = consume_whitespace(s)
        if length(remain) != 0 && remain[1] != '#'
            throw(ParserError(join(
                ["Non-comment characters remaining in line \"", s, "\"."])))
        end
    end

    # -------------------------------------------------------------------------
    # =========================================================================
    # -------------------------------------------------------------------------

    # Attempt to parse a string into a string Value (which follows a particular
    # format).  Return the string and the trailing characters.  Throw a
    # ParserError if parsing fails.
    function parse_string_value(s::String)
        idx = 1
        if s[idx] != '"'
            throw(ParserError("Unable to parse as a string."))
        end
        idx += 1
        temp_str::UTF8String = ""
        while idx <= length(s)
            if s[idx] == '\\'
                idx += 1
                if s[idx] == '"'
                    temp_str = join([temp_str, '"'])
                    idx += 1
                elseif s[idx] == '\\'
                    temp_str = join([temp_str, '\\'])
                    idx += 1
                elseif s[idx] == 'b'
                    temp_str = join([temp_str, "\b"])
                    idx += 1
                elseif s[idx] == 't'
                    temp_str = join([temp_str, "\t"])
                    idx += 1
                elseif s[idx] == 'n'
                    temp_str = join([temp_str, "\n"])
                    idx += 1
                elseif s[idx] == 'f'
                    temp_str = join([temp_str, "\f"])
                    idx += 1
                elseif s[idx] == 'r'
                    temp_str = join([temp_str, "\r"])
                    idx += 1
                else
                    throw(ParserError(join(["Unknown escale character \"\\",
                        s[idx], "\"."])))
                end
            elseif s[idx] == '"'
                break
            else
                temp_str = join([temp_str, s[idx]])
                idx += 1
            end
        end
        if s[idx] != '"'
            throw(ParserError("Unable to parse as a string."))
        end
        idx += 1
        return(temp_str, s[idx:end])
    end # function parse_string_value

    # =========================================================================

    # Attempt to parse a string as a number.  Return the NumberValue struct or
    # raise a ParserError.
    function parse_number_value(s::String)
        num_str = split(s)[1]
        remain  = s[length(num_str)+1:end]
        if (contains(num_str, ".") || contains(num_str, "e") ||
            contains(num_str, "E"))
            try
                return(parsefloat(Float64, num_str), remain)
            catch
                throw(ParserError("Unable to parse as a float."))
            end
        else
            try
                return(parseint(Int64, num_str), remain)
            catch
                throw(ParserError("Unable to parse as an integer."))
            end
        end
    end # function parse_number_value

    # =========================================================================

    # Attempt to parse a string as a boolean.  Return the boolean or raise a
    # ParserError.
    function parse_boolean_value(s::String)
        temp_bool::Bool = false
        if length(s) >= 4 && s[1:4] == "true"
            idx = 5
            temp_bool = true
        elseif length(s) >= 5 && s[1:5] == "false"
            idx = 6
            temp_bool = false
        else
            throw(ParserError("Unable to parse as a boolean."))
        end
        return(temp_bool, s[idx:end])
    end # function parse_boolean_value

    # -------------------------------------------------------------------------
    # =========================================================================
    # -------------------------------------------------------------------------

    # Return a value based on parsing the String
    function analyze_value(s::String)
        # Clear the current internal values and flags
        clear(v)

        # Trim leading whitespace
        s = consume_whitespace(s)

        # Ensure there is something (non-comment) left in the string
        if length(s) == 0 || s[1] == '#'
            throw(ParserError("Empty value."))
        end

        # Choose which type to parse
        if s[1] == '"'
            # This is either a String or nothing
            value, s = parse_string_value(s)
        elseif s[1] == 't' || s[1] == 'f'
            # This is either a Boolean or nothing
            value, s = parse_boolean_value(s)
        elseif s[1] == '-' || s[1] == '+' || s[1] == '.' || is_digit(s[1])
            # This is either an Integer, a Float, both, or nothing
            value, s = parse_number_value(s)
        else
            # This is nothing
            throw(ParserError(join(["Unable to parse \"", s,
                "\" to a value."])))
        end

        # Return the remainder of the string
        return value, s
    end # function analyze_value

    # =========================================================================

    function serialize(v::UTF8String)
        output = ["\""]
        idx = 1
        while idx <= length(v)
            if v[idx] == '"'
                push!(output, "\\\"")
            elseif v[idx] == '\\'
                push!(output, "\\\\")
            elseif v[idx] == '\b'
                push!(output, "\\b")
            elseif v[idx] == '\t'
                push!(output, "\\t")
            elseif v[idx] == '\n'
                push!(output, "\\\n")
            elseif v[idx] == '\f'
                push!(output, "\\f")
            elseif v[idx] == '\r'
                push!(output, "\\r")
            else
                push!(output, v[idx])
            end
            idx += 1
        end
        pushd!(output, "\"")
        return(convert(UTF8String, join(output)))
    end # function serialize(UTF8String)

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    function serialize(v::Bool)
        return(convert(UTF8String, string(v)))
    end # function serialize(Bool)

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    function serialize(v::Int64)
        return(convert(UTF8String, string(v)))
    end # function serialize(Int64)

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    function serialize(v::Float64)
        return(convert(UTF8String, string(v)))
    end # function serialize(Float64)

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    function serialize(v)
        throw(ValueError("Value cannot be serialized."))
    end # function serialize(Any)

    # -------------------------------------------------------------------------
    # =========================================================================
    # -------------------------------------------------------------------------

    type Table
        data_map::Dict{UTF8String,Any}
    end

    function Table()
        return Table(Dict{UTF8String,Any}())
    end function # Table

    # =========================================================================

    function empty!(t::Table)
        empty!(t.data_map)
    end

    # =========================================================================

    function get_table(t::Table, key::String)
        if has(t, key)
            v = t.data_map[key]
            if isa(v, Table)
                return(v)
            else
                throw(TableError("Requested key is not a table."))
            end
        else
            throw(TableError("Requested key does not exist in table."))
        end
    end

    # =========================================================================

    function has(t::Table, key::String)
        return(key in keys(t.data_map))
    end

    # =========================================================================

    function has(t::Table, path::Array{String,1})
        current_table = t
        for p in path[1:length(path)-1]
            try
                current_table = get_table(current_table, p)
            catch e
                isa(e, TableError) || rethrow(e)
                return(false)
            end
        end
        return(has(current_table, path[length(path)]))
    end

    # =========================================================================

    function analyze_quoted_key(s::String)
        # Quoted keys follow the same rules as string values, so we can just
        # analyze it as a value and return the string if a value string results
        try
            str, s = parse_string_value(s)
        catch e
            isa(e, ParserError) || rethrow(e)
            throw(ParserError("Could not parse quoted key."))
        end
        if length(str) == 0
            throw(ParserError("Cannot have an empty quoted key."))
        end
        return(str, s)
    end

    # =========================================================================

    function analyze_bare_key(s::String)
        key = String[]
        idx = 1
        while (length(s) != 0 && (is_digit(s[idx]) || is_letter(s[idx]) ||
            s[idx] == '_' || s[idx] == '-'))
            push!(key, s[idx])
            idx += 1
        end
        if length(key) == 0
            throw(ParserError("Empty bare key."))
        end
        return(join(key), s[idx:end])
    end

    # =========================================================================

    function analyze_key(s::String)
        if s[1] == '"'
            return(analyze_quoted_key(s))
        else
            return(analyze_bare_key(s))
        end
    end

    # =========================================================================

    function analyze_table_name(s::String)
        tname = consume_whitespace(s)
        p, tname = analyze_key(tname)
        path = [p]
        tname = consume_whitespace(tname)
        while tname[1] == '.'
            tname = consume_character('.', tname)
            tname = consume_whitespace(tname)
            p, tname = analyze_key(tname)
            push!(path, p)
            tname = consume_whitespace(tname)
        end
        return(path)
    end

    # =========================================================================

    function parse_string_table(root::Table, s::String)
        empty!(root)
        current_Table = root
        lines = split(s, "\n")
        try
            for line in lines
                # Strip leading whitespace
                line = consume_whitespace(line)
                # What kind of line is it?
                if length(line) == 0 || line[1] == '#'
                    # If the line is empty or comment-only, skip it
                    continue
                elseif line[1] == '['
                    # This is the start of a new Table
                    # Note: All paths in a file will be specified from the root
                    line = consume_character('[', line)
                    path, line = analyze_table_name(line)
                    line = consume_character(']', line)
                    consume_to_eol(line)
                    # The TOML standard does not allow re-entering a Table
                    # after you've already created it and then moved to another
                    # Table.  Thus we generate an error if the Table already
                    # exists.
                    # TODO -- The TOML standard actually allows a slightly more
                    #         complex behavior: If you define Table [a.b], you
                    #         can then go back and fill in Table [a] so long as
                    #         [a] only exists because you built it as an
                    #         intermediary to build [a.b].  Thus I will need a
                    #         more complex bookkeeping mechanism to specify
                    #         whether a Table exists because it was directly
                    #         defined or because it was built as an
                    #         intermediary.
                    if has(root, path)
                        message = ["Key \"", path[1]]
                        for p in path[2:end]
                            push!(message, ".")
                            push!(message, p)
                        end
                        push!(message, "\" is not unique.")
                        throw(ParserError(join(message)))
                    end
                    # Create the Table (and all intermediaries)
                    current_table = get_table(root, path, true)
                else
                    # This is a key pair
                    key, line = analyze_key(line)
                    if has(current_table, key)
                        throw(ParserError(
                            join(["Key \"", key, "\" is not unique."])))
                    end
                    line = consume_whitespace(line)
                    line = consume_character("=", line)
                    line = consume_whitespace(line)
                    if line[1] == '['
                        # This is an array
                        arr = []
                        line = consume_character('[', line)
                        line = consume_whitespace(line)
                        while line[1] != ']'
                            value, line = parse_string_value(line)
                            if length(arr) == 0
                                arr = [value]
                            else
                                if typeof(arr[1]) == typeof(value)
                                    push!(arr, value)
                                else
                                    throw(ParserError(
                                        "Mixed-type arrays not allowed."))
                                end
                            end
                            line = consume_whitespace(line)
                            if line[1] == ','
                                line = consume_character(",", line)
                                line = consume_whitespace(line)
                            elseif line[1] != ']'
                                throw(ParserError(
                                    "Malformed array of values."))
                            end
                            end
                            line = consume_character("]", line)
                            consume_to_eol(line)
                            add(current_table, key, arr)
                        else
                            # This is a value
                            value, line = parse_string_value(line)
                            consume_to_eol(line)
                            add(current_table, key, value)
                        end
                    end
                end # "what kind of line" if-else ladder
            end # loop over lines
        catch e
            if isa(e, ParserError)
                empty!(root)
            end
            rethrow(e)
        end
    end # function parse_string_table
end # module TOML

max_stars_repo_name: bkkrueger/config

max_issues_repo_name: bkkrueger/TOML

max_forks_repo_name: bkkrueger/TOML


text:
using Test
using Catlab.Graphs
using Catlab.CategoricalAlgebra
using Catlab.Present
using Catlab.Theories

using CSetAutomorphisms
using Catlab.Graphs.BasicGraphs: TheoryGraph

using Random

# Auxillary function tests
##########################

@test common([],[]) == 0
@test common([1],Int[]) == 0
@test common([1],Int[1]) == 1
@test common([1],Int[1,2]) == 1
@test common([1,2,3],Int[1,2]) == 2

# Helper functions for writing automorphism tests
#################################################

# Tests
#######

G,H = Graph(4), Graph(4);
add_edges!(G,[1,2,4,4,3],[2,4,3,3,2]);
add_edges!(H,[2,3,1,4,4],[1,1,4,3,3]);
test_iso(G,H, TheoryGraph) # 196 automorphisms

Triangle = Graph(3) # f;g = h
add_edges!(Triangle, [1,1,2], [2,3,3]) # f,h,g

Tri_, (G,H) = init_graphs(:Tri, Triangle,[2,2,2])
for i in 1:3 set_subpart!(G, Symbol("e$i"), [1,1]) end
for i in 1:3 set_subpart!(H, Symbol("e$i"), [2,2]) end
test_iso(G, H, Tri_)

Loop = Graph(1)
add_edge!(Loop, 1, 1)
Loo_, (G, H) = init_graphs(:Loo, Loop, [3])
set_subpart!(G, Symbol("e1"), [3,2,1])
set_subpart!(H, Symbol("e1"), [1,3,2])
test_iso(G, H, Loo_)

cyclel, cycler = Graph(3), Graph(3)
add_edges!(cyclel,[1,2,3],[2,3,1])
add_edges!(cycler,[3,2,1],[2,1,3])
test_iso(cyclel, cycler, TheoryGraph)

Loop2 = Graph(1)
add_edges!(Loop2, [1,1],[1,1])

Loo2_, (G,H) = init_graphs(:Loo2, Loop2, [2])
set_subpart!(G, :e1, [2,1])
set_subpart!(G, :e2, [2,1])
set_subpart!(H, :e1, [1,1])
set_subpart!(H, :e2, [2,2])
test_iso(G, H, Loo2_)

# Example from Hartke and Radcliffe exposition of Nauty.
# G is their optimal ordering. H is the original.
G, H = Graph(9), Graph(9)
add_edges!(G,[1,1,2,2,3,3,4,4,5,6,7,8],
             [7,8,5,6,6,8,5,7,9,9,9,9])
add_edges!(H,[1,1,3,3,7,7,9,9,2,4,6,8],
             [2,4,2,6,4,8,6,8,5,5,5,5])
res, tree = autos(H);
# When branching is restricted to :V as is the case in Nauty
# length should be 13 without auto pruning
# length is 10 with auto pruning tactic #1
# length is 6 with auto pruning tactic #2 too
# However, we can branch on :E too. This leads to just length 4 soln.

test_iso(G,H) # nauty doesn't seem to terminate

"""Graph corresponding to schema for finite limit sketch for categories"""
catschema = @acset Graph begin
  V = 7
  E = 17
  src = [2,2,1, 3,3,3, 4,4,5,5,4, 6,6,6, 7,7,7]
  tgt = [1,1,2, 2,2,2, 2,2,2,2,5, 1,2,3, 1,2,3]
end
random_perm = Dict([:V=>randperm(7), :E=>randperm(17)])
catschema2 = apply_automorphism(catschema, random_perm)
test_iso(catschema,catschema2, TheoryGraph)

# ACSet Tests
#############

G = @acset Labeled{String} begin
  V = 4
  E = 4
  src = [1,2,3,4]
  tgt = [2,3,4,1]
  dec = ["a","b","c","d"]
end;

H = @acset Labeled{String} begin
  V = 4
  E = 4
  src = [1,3,2,4]
  tgt = [3,2,4,1]
  dec = ["a","b","c","d"]
end;

test_iso(G, H, TheoryDecGraph) # vertices permuted

I = @acset Labeled{String} begin
V = 4
E = 4
src = [1,2,3,4]
tgt = [2,3,4,1]
dec = ["b","c","d","a"]
end;

test_iso(G, I, TheoryDecGraph) # labels permuted

N = @acset Labeled{String} begin
  V = 4
  E = 4
  src = [1,2,3,4]
  tgt = [2,3,4,1]
  dec = ["a","a","b","c"]
end;

test_iso(G,N) # label mismatch

K = @acset Labeled{String} begin
  V = 4
  E = 4
  src = [1,3,2,4]
  tgt = [2,3,4,1]
  dec = ["a","d","b","c"]
end;

test_iso(G, K) # vertex mismatch

G1 = @acset Labeled{String} begin
  V = 1
  E = 1
  src = [1]
  tgt = [1]
  dec = ["a"]
end;

H1 = @acset Labeled{String} begin
  V = 1
  E = 1
  src = [1]
  tgt = [1]
  dec = ["b"]
end;

test_iso(G1, H1) # label values different


text:
(::Type{SA})(x::Tuple{Tuple{Tuple{<:Tuple}}}) where {SA <: TupleVector} =
    throw(DimensionMismatch("No precise constructor for $SA found. Val of input was $(length(x[1][1][1]))."))

@inline (::Type{SA})(x...) where {SA <: TupleVector} = SA(x)
@inline (::Type{SA})(a::TupleVector) where {SA<:TupleVector} = SA(Tuple(a))
@propagate_inbounds (::Type{SA})(a::AbstractArray) where {SA <: TupleVector} = convert(SA, a)

# this covers most conversions and "statically-sized reshapes"
#@inline Base.convert(::Type{SA}, sa::TupleVector) where {SA<:TupleVector} = SA(Tuple(sa))
#@inline Base.convert(::Type{SA}, sa::SA) where {SA<:TupleVector} = sa
@inline Base.convert(::Type{SA}, x::Tuple) where {SA<:TupleVector} = SA(x) # convert -> constructor. Hopefully no loops...

# support conversion to AbstractArray
Base.AbstractArray{T}(sa::TupleVector{N,T}) where {N,T} = sa
Base.AbstractArray{T,1}(sa::TupleVector{N,T}) where {N,T} = sa
Base.AbstractArray{T}(sa::TupleVector{N,U}) where {N,T,U} = similar_type(typeof(sa),T,Val(N))(sa)
Base.AbstractArray{T,1}(sa::TupleVector{N,U}) where {N,T,U} = similar_type(typeof(sa),T,Val(N))(sa)

# Constructing a Tuple from a TupleVector
@inline Base.Tuple(a::TupleVector{N}) where N = unroll_tuple(a, Val(N))

@noinline function dimension_mismatch_fail(SA::Type, a::AbstractArray)
    throw(DimensionMismatch("expected input array of length $(length(SA)), got length $(length(a))"))
end

@propagate_inbounds function Base.convert(::Type{SA}, a::AbstractArray) where {SA<:TupleVector{N}} where N
    @boundscheck if length(a) != length(SA)
        dimension_mismatch_fail(SA, a)
    end

    return _convert(SA, a, Val(N))
end

@inline _convert(SA, a, l::Val) = SA(unroll_tuple(a, l))
@inline _convert(SA::Type{<:TupleVector{N,T}}, a, ::Val{0}) where {N,T} = similar_type(SA, T)(())
@inline _convert(SA, a, ::Val{0}) = similar_type(SA, eltype(a))(())

@pure length_val(n::Val{S}) where S = n
@pure length_val(n::TupleVector{S}) where S = Val(S)
@pure length_val(n::Type{<:TupleVector{S}}) where S = Val(S)
@pure length_val(a::T) where {T <: TupleMatrixLike{S}} where S = Val(S)
@pure length_val(a::Type{T}) where {T<:TupleMatrixLike{S}} where S = Val(S)

@generated function unroll_tuple(a::AbstractArray, ::Val{N}) where N
    exprs = [:(a[$j]) for j = 1:N]
    quote
        @_inline_meta
        @inbounds return $(Expr(:tuple, exprs...))
    end
end


text:
# Define a dummy function for automatic differentiation:
function G3(param::Array{T,1}) where {T <: Real}
return G3(param[1],param[2])
end

function G3(gamma::T,beta::T;gc=convert(T,0.5)) where {T <: Real}
sqb = sqrt(abs(beta))
#x = sqb*s
if gamma < gc
  return G3_series(gamma,beta)
else
  if beta >= 0
    return (gamma-sin(gamma))/(sqb*beta)
  else
    return (gamma-sinh(gamma))/(sqb*beta)
  end
end
end

y = zeros(2)
dG3 = y -> ForwardDiff.gradient(G3,y);

function H2(gamma::T,beta::T;gc=convert(T,0.5)) where {T <: Real}
sqb = sqrt(abs(beta))
#x=sqb*s
if gamma < gc
  return H2_series(gamma,beta)
else
  if beta >= 0
    return (sin(gamma)-gamma*cos(gamma))/(sqb*beta)
  else
    return (sinh(gamma)-gamma*cosh(gamma))/(sqb*beta)
  end
end
end

function H1(gamma::T,beta::T;gc=convert(T,0.5)) where {T <: Real}
#x=sqrt(abs(beta))*s
if gamma < gc
  return H1_series(gamma,beta)
else
  if beta >= 0
    return (4sin(0.5*gamma)^2 -gamma*sin(gamma))/beta^2
  else
    return (-4sinh(0.5*gamma)^2+gamma*sinh(gamma))/beta^2
  end
end
end

function G3_series(gamma::T,beta::T) where {T <: Real}
epsilon = eps(gamma)
# Computes G_3(\beta,s) using a series tailored to the precision of s.
#x2 = -beta*s^2
x2 = -sign(beta)*gamma^2
term = one(T)
g3 = one(T)
g31 = 2g3
g32 = 2g3
n=0
iter = 0
ITMAX = 100
# Terminate series when required precision reached:
#while abs(term) > epsilon*abs(g3)
while true
  g32 = g31
  g31 = g3
  n += 1
  term *= x2/((2n+3)*(2n+2))
  g3 += term
  iter +=1
  if iter >= ITMAX || g3 == g32 || g3 == g31
    break
  end
end
g3 *= gamma^3/(6*sqrt(abs(beta^3)))
return g3::T
end

dG3_series = y -> ForwardDiff.gradient(G3_series,y);

# Define a dummy function for automatic differentiation:
function G3_series(param::Array{T,1}) where {T <: Real}
return G3_series(param[1],param[2])
end

function H2_series(gamma::T,beta::T) where {T <: Real}
# Computes H_2(\beta,s) using a series tailored to the precision of s.
epsilon = eps(gamma)
#x2 = -beta*s^2
x2 = -sign(beta)*gamma^2
term = one(T)
h2 = one(T)
h21 = 2h2
h22 = 2h2
n=0
iter = 0
ITMAX = 100
# Terminate series when required precision reached:
#while abs(term) > epsilon*abs(h2)
while true
  h22 = h21
  h21 = h2
  n += 1
  term *= x2
  term /= (4n+6)*n
  h2 += term
  iter += 1
  if iter >= ITMAX || h2 == h22 || h2 == h21
    break
  end
end
h2 *= gamma^3/(3*sqrt(abs(beta^3)))
return h2::T
end

function H1_series(gamma::T,beta::T) where {T <: Real}
# Computes H_1(\beta,s) using a series tailored to the precision of s.
epsilon = eps(gamma)
#x2 = -beta*s^2
x2 = -sign(beta)*gamma^2
term = one(T)
h1 = one(T)
h11 = 2h1
h12 = 2h1
n=0
iter = 0
ITMAX = 100
# Terminate series when required precision reached:
#while abs(term) > epsilon*abs(h1)
while true
  h12 = h11
  h11 = h1
  n += 1
  term *= x2*(n+1)
  term /= (2n+4)*(2n+3)*n
  h1 += term
  iter +=1
  if iter >= ITMAX || h1 == h12 || h1 == h11
    break
  end
end
h1 *= gamma^4/(12*beta^2)
return h1::T
end

function H5(gamma::T,beta::T;gc=convert(T,0.5)) where {T <: Real}
# This is G_1 G_2 -(2+G_0) G_3 = H_2 - 2 G_3:
if gamma < gc
  return H5_series(gamma,beta)
else
  if beta >= 0
    return (3sin(gamma)-2gamma-gamma*cos(gamma))/(beta*sqrt(beta))
  else
    return (3sinh(gamma)-2gamma-gamma*cosh(gamma))/(beta*sqrt(-beta))
  end
end
end

function H5_series(gamma::T,beta::T) where {T <: Real}
# Computes H_5(\beta,s) using a series tailored to the precision of gamma:
epsilon = eps(gamma)
#x2 = -beta*s^2
x2 = -sign(beta)*gamma^2
term = one(T)/60
h5 = copy(term)
h51 = 2h5
h52 = 2h5
n=0
iter = 0
ITMAX = 100
# Terminate series when required precision reached:
#while abs(term) > epsilon*abs(h5)
while true
  h52 = h51
  h51 = h5
  n += 1
  term *= x2*(n+1)
  term /= (2n+5)*(2n+4)*n
  h5 += term
  iter += 1
  if iter >= ITMAX || h5 == h52 || h5 == h51
    break
  end
end
h5 *= -gamma^5/(beta*sqrt(abs(beta)))
return h5::T
end

function H6(gamma::T,beta::T;gc=convert(T,0.5)) where {T <: Real}
# This is 2 G_2^2 -3 G_1 G_3:
if gamma < gc
  return H6_series(gamma,beta)
else
  if beta >= 0
    return (9-8cos(gamma) -cos(2gamma) -6gamma*sin(gamma))/(2beta^2)
  else
    return (9-8cosh(gamma)-cosh(2gamma)+6gamma*sinh(gamma))/(2beta^2)
  end
end
end

function H6_series(gamma::T,beta::T) where {T <: Real}
# Computes H_6(\beta,s) using a series tailored to the precision of gamma:
epsilon = eps(gamma)
#x2 = -beta*s^2
x2 = -sign(beta)*gamma^2
term = convert(T,1//360)
h6 = convert(T,1//40)
h61 = 2h6
h62 = 2h6
n=0
iter = 0
ITMAX = 100
# Terminate series when required precision reached:
#while abs(term) > epsilon*abs(h6)
while true
  h62 = h61
  h61 = h6
  n += 1
  term *= x2
  term /= (2n+5)*(2n+6)
  h6 += term*(4^(n+2)-3*n-7)
  iter +=1
  if iter >= ITMAX || h6 == h62 || h6 == h61
    break
  end
end
h6 *= gamma^6/(beta*abs(beta))
return h6::T
end

function H3(gamma::T,beta::T;gc=convert(T,0.5)) where {T <: Real}
# This is H_3 = G_1 G_2 - 3 G_3:
if gamma < gc
  return H3_series(gamma,beta)
else
  if beta >= 0
    return (4*sin(gamma)-sin(gamma)*cos(gamma)-3*gamma)/(beta*sqrt(beta))
  else
    return (4*sinh(gamma)-sinh(gamma)*cosh(gamma)-3*gamma)/(beta*sqrt(-beta))
  end
end
end

function H3_series(gamma::T,beta::T) where {T <: Real}
# Computes H_3(\beta,s) using a series tailored to the precision of gamma:
epsilon = eps(gamma)
#x2 = -beta*s^2
x2 = -sign(beta)*gamma^2
term = convert(T,1//30)
h3 = convert(T,1//10)
h31 = 2h3
h32 = 2h3
n=0
iter = 0
ITMAX = 100
# Terminate series when required precision reached:
#while abs(term) > epsilon*abs(h3)
while true
  h32 = h31
  h31 = h3
  n += 1
  term *= x2
  term /= (2n+4)*(2n+5)
  h3 += term*(4^(n+1)-1)
  iter += 1
  if iter >= ITMAX || h3 == h32 || h3 == h31
    break
  end
end
h3 *= -gamma^5/(beta*sqrt(abs(beta)))
return h3::T
end

function H7(gamma::T,beta::T;gc=convert(T,0.5)) where {T <: Real}
# This is H_7 = G_1 G_2 (1-2 G_0) + 3 G_0^2 G_3:
if gamma < gc
  return H7_series(gamma,beta)
else
  if beta >= 0
    return (3*cos(gamma)*(gamma*cos(gamma)-sin(gamma))+sin(gamma)^3)/(beta*sqrt(beta))
  else
    return (3*cosh(gamma)*(gamma*cosh(gamma)-sinh(gamma))-sinh(gamma)^3)/(beta*sqrt(-beta))
  end
end
end

function H7_series(gamma::T,beta::T) where {T <: Real}
# Computes H_7(\beta,s) using a series tailored to the precision of gamma:
epsilon = eps(gamma)
#x2 = -beta*s^2
x2 = -sign(beta)*gamma^2
term = -convert(T,3//20160)*x2
h7 = convert(T,1//10-11//840*x2)
h71 = 2h7
h72 = 2h7
n=0
iter = 0
ITMAX = 100
# Terminate series when required precision reached:
#while abs(term) > epsilon*abs(h7)
while true
  h72 = h71
  h71 = h7
  n += 1
  term *= x2
  term /= (2n+6)*(2n+7)
  h7 += term*(9^(3+n)-1-(5+2*n)*2^(7+2*n))
  iter += 1
  if iter >= ITMAX || h7 == h72 || h7 == h71
    break
  end
end
h7 *= gamma^5/(beta*sqrt(abs(beta)))
return h7::T
end

function H8(gamma::T,beta::T;gc=convert(T,0.5)) where {T <: Real}
# This is H_8 = G_1 G_2 - 3 G_0 G_3:
if gamma < gc
  return H8_series(gamma,beta)
else
  if beta >= 0
    return (-3gamma*cos(gamma) +sin(gamma) +sin(2gamma))/(beta*sqrt(beta))
  else
    return (-3gamma*cosh(gamma)+sinh(gamma)+sinh(2gamma))/(beta*sqrt(-beta))
  end
end
end

function H8_series(gamma::T,beta::T) where {T <: Real}
# Computes H_8(\beta,s) using a series tailored to the precision of gamma:
epsilon = eps(gamma)
#x2 = -beta*s^2
x2 = -sign(beta)*gamma^2
term = convert(T,1//120)
h8 = convert(T,3//20)
h81 = 2h8
h82 = 2h8
n=0
iter = 0
ITMAX = 100
# Terminate series when required precision reached:
#while abs(term) > epsilon*abs(h8)
while true
  h82 = h81
  h81 = h8
  n += 1
  term *= x2
  term /= (2n+4)*(2n+5)
  h8 += term*(2^(5+2n)-14-6n)
  iter +=1
  if iter >= ITMAX || h8 == h82 || h8 == h81
    break
  end
end
h8 *= gamma^5/(beta*sqrt(abs(beta)))
return h8::T
end


text:
module MLMetrics

using LearnBase
using MLLabelUtils
using MLLabelUtils: LabelEncoding, BinaryLabelEncoding
import LossFunctions
using LossFunctions: @_dimcheck
using UnicodePlots
using Statistics
using Base.Broadcast

export

    AvgMode,
    LabelEnc,

    # counts
    true_positives,
    true_negatives,
    false_positives, type_1_errors,
    false_negatives, type_2_errors,
    condition_positive,
    condition_negative,
    predicted_positive,
    predicted_negative,
    correctly_classified,
    incorrectly_classified, misclassified,

    # class metrics
    prevalence,
    positive_predictive_value, precision_score,
    negative_predictive_value,
    true_positive_rate, sensitivity, recall,
    false_positive_rate,
    false_negative_rate,
    true_negative_rate, specificity,
    false_discovery_rate,
    false_omission_rate,
    accuracy,
    f_score, f1_score,
    positive_likelihood_ratio,
    negative_likelihood_ratio,
    diagnostic_odds_ratio,
    matthews_corrcoef,

    # roc
    confusions,
    confusions_at_sensitivity,
    roc,
    specificity_at_sensitivity,
    precision_at_sensitivity,
    accuracy_at_sensitivity,
    auc, auc_from_rates,

    # regression
    absolute_error,
    percent_error,
    log_error,
    squared_error,
    squared_log_error,
    absolute_percent_error,
    mean_error,
    mean_absolute_error,
    median_absolute_error,
    mean_percent_error,
    median_percent_error,
    mean_squared_error,
    median_squared_error,
    sum_squared_error,
    mean_squared_log_error,
    mean_absolute_percent_error,
    median_absolute_percent_error,
    symmetric_mean_absolute_percent_error,
    symmetric_median_absolute_percent_error,
    mean_absolute_scaled_error,
    total_variance_score,
    explained_variance_score,
    unexplained_variance_score,
    r2_score

include("classification/averagemode.jl")
include("classification/upstream.jl")
include("classification/docstring_const.jl")
include("classification/counts.jl")
include("classification/fraction.jl")
include("classification/metrics.jl")
include("classification/roc.jl")
include("regression.jl")

end # module MLMetrics


text:
using DataStructures: PriorityQueue
using LinearAlgebra: dot
using NNlib: σ
using Statistics: mean
using Flux: param

struct HeirarchicalSoftmax{A<:AbstractArray{<:AbstractFloat, 2}, B<:AbstractArray{<:AbstractArray{<:Integer, 1}}}
    tree_vectors::A
    nodepaths::B
    branchpaths::B
end

function HeirarchicalSoftmax(wordvector_dims::Int, pq::PriorityQueue, vocab_hash::Dict, vocab::Array)
    return HeirarchicalSoftmax(param(randn(length(keys(vocab_hash)), wordvector_dims)), allpaths(HuffmanTree(pq, vocab_hash), vocab_hash, vocab)...)
end

function (layer::HeirarchicalSoftmax)(hiddenvector::AbstractArray{<:AbstractFloat, 1}, outputword::Integer)
    reduce(*, map((node, branch)->(_nodeprob(hiddenvector, layer.tree_vectors, node, branch)), layer.nodepaths[outputword], layer.branchpaths[outputword]))
end

function (layer::HeirarchicalSoftmax)(hiddenvectors::AbstractArray{<:AbstractFloat, 2}, outputwords::AbstractArray{<:Integer, 1})
    p = []
    size(hiddenvectors)[2] ≠ length(outputwords) ? error("X and y batch sizes do not match") : nothing
    for i in 1:length(outputwords)
        hiddenvector = hiddenvectors[:, i]
        outputword = outputwords[i]
        push!(p, reduce(*, map((node, branch)->(_nodeprob(hiddenvector, layer.tree_vectors, node, branch)), layer.nodepaths[outputword], layer.branchpaths[outputword])))
    end
    return mean(p)
end

(layer::HeirarchicalSoftmax)(incoming::Tuple) = layer(incoming...)

function _nodeprob(hiddenvector, tree_vectors, node, branch)
    branch == 1 ? ◫ = 1 : ◫ = -1
    return σ(◫*dot(tree_vectors[node,:],hiddenvector))
end


text:
# ---
# title: 989. Add to Array-Form of Integer
# id: problem989
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/add-to-array-form-of-integer/description/>
# hidden: true
# ---
#
# For a non-negative integer `X`, the  _array-form of`X`_ is an array of its
# digits in left to right order.  For example, if `X = 1231`, then the array
# form is `[1,2,3,1]`.
#
# Given the array-form `A` of a non-negative integer `X`, return the array-form
# of the integer `X+K`.
#
#
#
# **Example 1:**
#
#
#
#     Input: A = [1,2,0,0], K = 34
#     Output: [1,2,3,4]
#     Explanation: 1200 + 34 = 1234
#
#
# **Example 2:**
#
#
#
#     Input: A = [2,7,4], K = 181
#     Output: [4,5,5]
#     Explanation: 274 + 181 = 455
#
#
# **Example 3:**
#
#
#
#     Input: A = [2,1,5], K = 806
#     Output: [1,0,2,1]
#     Explanation: 215 + 806 = 1021
#
#
# **Example 4:**
#
#
#
#     Input: A = [9,9,9,9,9,9,9,9,9,9], K = 1
#     Output: [1,0,0,0,0,0,0,0,0,0,0]
#     Explanation: 9999999999 + 1 = 10000000000
#
#
#
#
# **Note：**
#
#   1. `1 <= A.length <= 10000`
#   2. `0 <= A[i] <= 9`
#   3. `0 <= K <= 10000`
#   4. If `A.length > 1`, then `A[0] != 0`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
const EA_AS_PKG = Symbol(@__MODULE__) !== :Base # develop EA as an external package

module EAUtils

import ..EA_AS_PKG
if EA_AS_PKG
    import ..EscapeAnalysis
else
    import Core.Compiler.EscapeAnalysis: EscapeAnalysis
    Base.getindex(estate::EscapeAnalysis.EscapeState, @nospecialize(x)) =
        Core.Compiler.getindex(estate, x)
end
const EA = EscapeAnalysis
const CC = Core.Compiler

# entries
# -------

@static if EA_AS_PKG
import InteractiveUtils: gen_call_with_extracted_types_and_kwargs

@doc """
    @code_escapes [options...] f(args...)

Evaluates the arguments to the function call, determines its types, and then calls
[`code_escapes`](@ref) on the resulting expression.
As with `@code_typed` and its family, any of `code_escapes` keyword arguments can be given
as the optional arguments like `@code_escpase interp=myinterp myfunc(myargs...)`.
"""
macro code_escapes(ex0...)
    return gen_call_with_extracted_types_and_kwargs(__module__, :code_escapes, ex0)
end
end # @static if EA_AS_PKG

"""
    code_escapes(f, argtypes=Tuple{}; [world], [interp]) -> result::EscapeResult
    code_escapes(tt::Type{<:Tuple}; [world], [interp]) -> result::EscapeResult

Runs the escape analysis on optimized IR of a genefic function call with the given type signature.
Note that the escape analysis runs after inlining, but before any other optimizations.

```julia
julia> mutable struct SafeRef{T}
           x::T
       end

julia> Base.getindex(x::SafeRef) = x.x;

julia> Base.isassigned(x::SafeRef) = true;

julia> get′(x) = isassigned(x) ? x[] : throw(x);

julia> result = code_escapes((String,String,String)) do s1, s2, s3
           r1 = Ref(s1)
           r2 = Ref(s2)
           r3 = SafeRef(s3)
           try
               s1 = get′(r1)
               ret = sizeof(s1)
           catch err
               global g = err # will definitely escape `r1`
           end
           s2 = get′(r2)      # still `r2` doesn't escape fully
           s3 = get′(r3)      # still `r2` doesn't escape fully
           return s2, s3
       end
#3(X _2::String, ↑ _3::String, ↑ _4::String) in Main at REPL[7]:2
2  X  1 ── %1  = %new(Base.RefValue{String}, _2)::Base.RefValue{String}   │╻╷╷ Ref
3  *′ │    %2  = %new(Base.RefValue{String}, _3)::Base.RefValue{String}   │╻╷╷ Ref
4  ✓′ └─── %3  = %new(SafeRef{String}, _4)::SafeRef{String}               │╻╷  SafeRef
5  ◌  2 ── %4  = \$(Expr(:enter, #8))                                      │
   ✓′ │    %5  = ϒ (%3)::SafeRef{String}                                  │
   *′ └─── %6  = ϒ (%2)::Base.RefValue{String}                            │
6  ◌  3 ── %7  = Base.isdefined(%1, :x)::Bool                             │╻╷  get′
   ◌  └───       goto #5 if not %7                                        ││
   X  4 ──       Base.getfield(%1, :x)::String                            ││╻   getindex
   ◌  └───       goto #6                                                  ││
   ◌  5 ──       Main.throw(%1)::Union{}                                  ││
   ◌  └───       unreachable                                              ││
7  ◌  6 ──       nothing::typeof(Core.sizeof)                             │╻   sizeof
   ◌  │          nothing::Int64                                           ││
   ◌  └───       \$(Expr(:leave, 1))                                       │
   ◌  7 ──       goto #10                                                 │
   ✓′ 8 ── %17 = φᶜ (%5)::SafeRef{String}                                 │
   *′ │    %18 = φᶜ (%6)::Base.RefValue{String}                           │
   ◌  └───       \$(Expr(:leave, 1))                                       │
   X  9 ── %20 = \$(Expr(:the_exception))::Any                             │
9  ◌  │          (Main.g = %20)::Any                                      │
   ◌  └───       \$(Expr(:pop_exception, :(%4)))::Any                      │
11 ✓′ 10 ┄ %23 = φ (#7 => %3, #9 => %17)::SafeRef{String}                 │
   *′ │    %24 = φ (#7 => %2, #9 => %18)::Base.RefValue{String}           │
   ◌  │    %25 = Base.isdefined(%24, :x)::Bool                            ││╻   isassigned
   ◌  └───       goto #12 if not %25                                      ││
   ↑  11 ─ %27 = Base.getfield(%24, :x)::String                           │││╻   getproperty
   ◌  └───       goto #13                                                 ││
   ◌  12 ─       Main.throw(%24)::Union{}                                 ││
   ◌  └───       unreachable                                              ││
12 ↑  13 ─ %31 = Base.getfield(%23, :x)::String                           │╻╷╷ get′
13 ↑  │    %32 = Core.tuple(%27, %31)::Tuple{String, String}              │
   ◌  └───       return %32                                               │
```

The symbols in the side of each call argument and SSA statements represents the following meaning:
- `◌`: this value is not analyzed because escape information of it won't be used anyway (when the object is `isbitstype` for example)
- `✓`: this value never escapes (`has_no_escape(result.state[x])` holds)
- `↑`: this value can escape to the caller via return (`has_return_escape(result.state[x])` holds)
- `X`: this value can escape to somewhere the escape analysis can't reason about like escapes to a global memory (`has_all_escape(result.state[x])` holds)
- `*`: this value's escape state is between the `ReturnEscape` and `AllEscape` in the `EscapeLattice`, e.g. it has unhandled `ThrownEscape`
and additional `′` indicates that field analysis has been done successfully on that value.

For testing, escape information of each call argument and SSA value can be inspected programmatically as like:
```julia
julia> result.state[Core.Argument(3)]
ReturnEscape

julia> result.state[Core.SSAValue(3)]
NoEscape′
```
"""
function code_escapes(@nospecialize(args...);
                      world = get_world_counter(),
                      interp = Core.Compiler.NativeInterpreter(world))
    interp = EscapeAnalyzer(interp)
    results = code_typed(args...; optimize=true, world, interp)
    isone(length(results)) || throw(ArgumentError("`code_escapes` only supports single analysis result"))
    return EscapeResult(interp.ir, interp.state, interp.linfo)
end

# AbstractInterpreter
# -------------------

# imports
import .CC:
    AbstractInterpreter,
    NativeInterpreter,
    WorldView,
    WorldRange,
    InferenceParams,
    OptimizationParams,
    get_world_counter,
    get_inference_cache,
    lock_mi_inference,
    unlock_mi_inference,
    add_remark!,
    may_optimize,
    may_compress,
    may_discard_trees,
    verbose_stmt_info,
    code_cache,
    @timeit,
    get_inference_cache,
    convert_to_ircode,
    slot2reg,
    compact!,
    ssa_inlining_pass!,
    sroa_pass!,
    adce_pass!,
    type_lift_pass!,
    JLOptions,
    verify_ir,
    verify_linetable
# usings
import Core:
    CodeInstance, MethodInstance, CodeInfo
import .CC:
    OptimizationState, IRCode
import .EA:
    analyze_escapes, cache_escapes!

mutable struct EscapeAnalyzer{State} <: AbstractInterpreter
    native::NativeInterpreter
    ir::IRCode
    state::State
    linfo::MethodInstance
    EscapeAnalyzer(native::NativeInterpreter) = new{EscapeState}(native)
end

CC.InferenceParams(interp::EscapeAnalyzer)    = InferenceParams(interp.native)
CC.OptimizationParams(interp::EscapeAnalyzer) = OptimizationParams(interp.native)
CC.get_world_counter(interp::EscapeAnalyzer)  = get_world_counter(interp.native)

CC.lock_mi_inference(::EscapeAnalyzer,   ::MethodInstance) = nothing
CC.unlock_mi_inference(::EscapeAnalyzer, ::MethodInstance) = nothing

CC.add_remark!(interp::EscapeAnalyzer, sv, s) = add_remark!(interp.native, sv, s)

CC.may_optimize(interp::EscapeAnalyzer)      = may_optimize(interp.native)
CC.may_compress(interp::EscapeAnalyzer)      = may_compress(interp.native)
CC.may_discard_trees(interp::EscapeAnalyzer) = may_discard_trees(interp.native)
CC.verbose_stmt_info(interp::EscapeAnalyzer) = verbose_stmt_info(interp.native)

CC.get_inference_cache(interp::EscapeAnalyzer) = get_inference_cache(interp.native)

const GLOBAL_CODE_CACHE = IdDict{MethodInstance,CodeInstance}()
__clear_code_cache!() = empty!(GLOBAL_CODE_CACHE)

function CC.code_cache(interp::EscapeAnalyzer)
    worlds = WorldRange(get_world_counter(interp))
    return WorldView(GlobalCache(), worlds)
end

struct GlobalCache end

CC.haskey(wvc::WorldView{GlobalCache}, mi::MethodInstance) = haskey(GLOBAL_CODE_CACHE, mi)

CC.get(wvc::WorldView{GlobalCache}, mi::MethodInstance, default) = get(GLOBAL_CODE_CACHE, mi, default)

CC.getindex(wvc::WorldView{GlobalCache}, mi::MethodInstance) = getindex(GLOBAL_CODE_CACHE, mi)

function CC.setindex!(wvc::WorldView{GlobalCache}, ci::CodeInstance, mi::MethodInstance)
    GLOBAL_CODE_CACHE[mi] = ci
    add_callback!(mi) # register the callback on invalidation
    return nothing
end

function add_callback!(linfo)
    if !isdefined(linfo, :callbacks)
        linfo.callbacks = Any[invalidate_cache!]
    else
        if !any(@nospecialize(cb)->cb===invalidate_cache!, linfo.callbacks)
            push!(linfo.callbacks, invalidate_cache!)
        end
    end
    return nothing
end

function invalidate_cache!(replaced, max_world, depth = 0)
    delete!(GLOBAL_CODE_CACHE, replaced)

    if isdefined(replaced, :backedges)
        for mi in replaced.backedges
            mi = mi::MethodInstance
            if !haskey(GLOBAL_CODE_CACHE, mi)
                continue # otherwise fall into infinite loop
            end
            invalidate_cache!(mi, max_world, depth+1)
        end
    end
    return nothing
end

function CC.optimize(interp::EscapeAnalyzer, opt::OptimizationState, params::OptimizationParams, @nospecialize(result))
    ir = run_passes_with_ea(interp, opt.src, opt)
    return CC.finish(interp, opt, params, ir, result)
end

function run_passes_with_ea(interp::EscapeAnalyzer, ci::CodeInfo, sv::OptimizationState)
    @timeit "convert"   ir = convert_to_ircode(ci, sv)
    @timeit "slot2reg"  ir = slot2reg(ir, ci, sv)
    # TODO: Domsorting can produce an updated domtree - no need to recompute here
    @timeit "compact 1" ir = compact!(ir)
    @timeit "Inlining"  ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
    # @timeit "verify 2" verify_ir(ir)
    @timeit "compact 2" ir = compact!(ir)
    nargs = let def = sv.linfo.def
        isa(def, Method) ? Int(def.nargs) : 0
    end
    local state
    try
        @timeit "collect escape information" state = analyze_escapes(ir, nargs)
    catch err
        @info "error happened within `analyze_escapes`, insepct `Main.ir` and `Main.nargs`"
        @eval Main (ir = $ir; nargs = $nargs)
        rethrow(err)
    end
    cacheir = Core.Compiler.copy(ir)
    # cache this result
    cache_escapes!(sv.linfo, state, cacheir)
    # return back the result
    interp.ir = cacheir
    interp.state = state
    interp.linfo = sv.linfo
    @timeit "SROA"      ir = sroa_pass!(ir)
    @timeit "ADCE"      ir = adce_pass!(ir)
    @timeit "type lift" ir = type_lift_pass!(ir)
    @timeit "compact 3" ir = compact!(ir)
    if JLOptions().debug_level == 2
        @timeit "verify 3" (verify_ir(ir); verify_linetable(ir.linetable))
    end
    return ir
end

# printing
# --------

import Core: Argument, SSAValue
import .CC: widenconst, singleton_type
import .EA: EscapeLattice, EscapeState, ⊑, ⊏

# in order to run a whole analysis from ground zero (e.g. for benchmarking, etc.)
__clear_caches!() = (__clear_code_cache!(); EA.__clear_escape_cache!())

function get_name_color(x::EscapeLattice, symbol::Bool = false)
    getname(x) = string(nameof(x))
    if x === EA.⊥
        name, color = (getname(EA.NotAnalyzed), "◌"), :plain
    elseif EA.has_no_escape(x)
        name, color = (getname(EA.NoEscape), "✓"), :green
    elseif EA.has_all_escape(x)
        name, color = (getname(EA.AllEscape), "X"), :red
    elseif EA.NoEscape() ⊏ (EA.ignore_thrownescapes ∘ EA.ignore_aliasescapes)(x) ⊑ EA.AllReturnEscape()
        name = (getname(EA.ReturnEscape), "↑")
        color = EA.has_thrown_escape(x) ? :yellow : :cyan
    else
        name = (nothing, "*")
        color = EA.has_thrown_escape(x) ? :yellow : :bold
    end
    name = symbol ? last(name) : first(name)
    if name !== nothing && !isa(x.AliasEscapes, Bool)
        name = string(name, "′")
    end
    return name, color
end

# pcs = sprint(show, collect(x.EscapeSites); context=:limit=>true)
function Base.show(io::IO, x::EscapeLattice)
    name, color = get_name_color(x)
    if isnothing(name)
        Base.@invoke show(io::IO, x::Any)
    else
        printstyled(io, name; color)
    end
end
function Base.show(io::IO, ::MIME"application/prs.juno.inline", x::EscapeLattice)
    name, color = get_name_color(x)
    if isnothing(name)
        return x # use fancy tree-view
    else
        printstyled(io, name; color)
    end
end

struct EscapeResult
    ir::IRCode
    state::EscapeState
    linfo::Union{Nothing,MethodInstance}
    EscapeResult(ir::IRCode, state::EscapeState, linfo::Union{Nothing,MethodInstance} = nothing) =
        new(ir, state, linfo)
end
Base.show(io::IO, result::EscapeResult) = print_with_info(io, result.ir, result.state, result.linfo)
@eval Base.iterate(res::EscapeResult, state=1) =
    return state > $(fieldcount(EscapeResult)) ? nothing : (getfield(res, state), state+1)

# adapted from https://github.com/JuliaDebug/LoweredCodeUtils.jl/blob/4612349432447e868cf9285f647108f43bd0a11c/src/codeedges.jl#L881-L897
function print_with_info(io::IO,
    ir::IRCode, state::EscapeState, linfo::Union{Nothing,MethodInstance})
    # print escape information on SSA values
    function preprint(io::IO)
        ft = ir.argtypes[1]
        f = singleton_type(ft)
        if f === nothing
            f = widenconst(ft)
        end
        print(io, f, '(')
        for i in 1:state.nargs
            arg = state[Argument(i)]
            i == 1 && continue
            c, color = get_name_color(arg, true)
            printstyled(io, c, ' ', '_', i, "::", ir.argtypes[i]; color)
            i ≠ state.nargs && print(io, ", ")
        end
        print(io, ')')
        if !isnothing(linfo)
            def = linfo.def
            printstyled(io, " in ", (isa(def, Module) ? (def,) : (def.module, " at ", def.file, ':', def.line))...; color=:bold)
        end
        println(io)
    end

    # print escape information on SSA values
    # nd = ndigits(length(ssavalues))
    function preprint(io::IO, idx::Int)
        c, color = get_name_color(state[SSAValue(idx)], true)
        # printstyled(io, lpad(idx, nd), ' ', c, ' '; color)
        printstyled(io, rpad(c, 2), ' '; color)
    end

    print_with_info(preprint, (args...)->nothing, io, ir)
end

function print_with_info(preprint, postprint, io::IO, ir::IRCode)
    io = IOContext(io, :displaysize=>displaysize(io))
    used = Base.IRShow.stmts_used(io, ir)
    # line_info_preprinter = Base.IRShow.lineinfo_disabled
    line_info_preprinter = function (io::IO, indent::String, idx::Int)
        r = Base.IRShow.inline_linfo_printer(ir)(io, indent, idx)
        idx ≠ 0 && preprint(io, idx)
        return r
    end
    line_info_postprinter = Base.IRShow.default_expr_type_printer
    preprint(io)
    bb_idx_prev = bb_idx = 1
    for idx = 1:length(ir.stmts)
        preprint(io, idx)
        bb_idx = Base.IRShow.show_ir_stmt(io, ir, idx, line_info_preprinter, line_info_postprinter, used, ir.cfg, bb_idx)
        postprint(io, idx, bb_idx != bb_idx_prev)
        bb_idx_prev = bb_idx
    end
    max_bb_idx_size = ndigits(length(ir.cfg.blocks))
    line_info_preprinter(io, " "^(max_bb_idx_size + 2), 0)
    postprint(io)
    return nothing
end

end # module EAUtils


text:
@testset "Prim" begin
    g4 = complete_graph(4)

    distmx = [
      0  1  5  6
      1  0  4  10
      5  4  0  3
      6  10  3  0
    ]

    vec_mst = Vector{Edge}([Edge(1, 2), Edge(2, 3), Edge(3, 4)])
    for g in testgraphs(g4)
        # Testing Prim's algorithm
        mst = @inferred(prim_mst(g, distmx))
        @test mst == vec_mst
    end

    #second test
    distmx_sec = [
        0       0       0.26    0       0.38    0       0.58    0.16
        0       0       0.36    0.29    0       0.32    0       0.19
        0.26    0.36    0       0.17    0       0       0.4     0.34
        0       0.29    0.17    0       0       0       0.52    0
        0.38    0       0       0       0       0.35    0.93    0.37
        0       0.32    0       0       0.35    0       0       0.28
        0.58    0       0.4     0.52    0.93    0       0       0
        0.16    0.19    0.34    0       0.37    0.28    0       0
    ]

    vec2 = Vector{Edge}([Edge(8, 2), Edge(1, 3), Edge(3, 4), Edge(6, 5), Edge(8, 6), Edge(3, 7), Edge(1, 8)])
    gx = SimpleGraph(distmx_sec)
    for g in testgraphs(gx)
        mst2 = @inferred(prim_mst(g, distmx_sec))
        @test mst2 == vec2
    end
end