text:
# fbp2.jl

export fbp2

using MIRT: ImageGeom # todo ImageGeoms.jl
using MIRT: SinoGeom, SinoPar, SinoFan, SinoMoj

abstract type FBPplan end

struct Moj{Th,Tg}
    H::Th
    G::Tg
end

Moj() = Moj{AbstractMatrix{<:Real},AbstractMatrix{<:Real}}(zeros(Real,1,1),zeros(Real,1,1))
#Moj(a,b) = Moj{AbstractMatrix{<:Real},AbstractMatrix{<:Real}}(zeros(Real,1,1),zeros(Real,1,1)) # todo method overwritten warning

struct NormalPlan{S} <: FBPplan
    sg::S
    ig::ImageGeom
    window::Union{Symbol, AbstractVector{<:Real}}
    parker_weight::AbstractMatrix{<:Real}
    moj::Moj
end

NormalPlan(
    sg::SinoGeom,
    ig::ImageGeom,
    window::Union{Symbol, AbstractVector{<:Real}},
    parker_weight::AbstractMatrix{<:Real}
) = NormalPlan{typeof(sg)}(sg,ig,window,parker_weight,Moj())

struct DfPlan <: FBPplan
    sg::SinoGeom
    ig::ImageGeom
    window::Union{Symbol,AbstractVector{<:Real}}
    parallel_beam_parker_weight::AbstractMatrix{<:Real}
end

struct MojPlan <: FBPplan
    sg::SinoGeom
    ig::ImageGeom
    window::Union{Symbol,AbstractVector{<:Real}}
    parallel_beam_parker_weight::AbstractMatrix{<:Real}
    moj::Moj
end

struct TabPlan <: FBPplan
    sg::SinoGeom
    ig::ImageGeom
    window::Union{Symbol,AbstractVector{<:Real}}
    parallel_beam_parker_weight::AbstractMatrix{<:Real}
end

reale = (x) -> (@assert x ≈ real(x); real(x))

"""
    plan = fbp2(sg, ig; how=:normal, window=:none)

FBP 2D tomographic image reconstruction for parallel-beam or fan-beam cases,
with either flat or arc detector for fan-beam case.

To use this, you first call it with the sinogram and image geometry.
The routine returns the initialized "plan". Thereafter, to
to perform FBP reconstruction, you call this routine with the plan
(perhaps numerous times for the same geometry).

in
- `sg::SinoGeom`
- `ig::ImageGeom`

options
- `how::Symbol`             how to reconstruct
    * `:normal`             default
    * `:mojette`            use mojette rebinning and Gtomo2_table
- `window::Symbol`          e.g. `:hann` (default: `:none`)
-`T::DataType`              type of sino elements (default: `Float32`)

out
- `plan::FBPplan`            initialized plan

"""
function fbp2(
    sg::SinoGeom,
    ig::ImageGeom;
    how::Symbol = :normal,
    window::Symbol = :none,
    T::DataType = Float32,
#   nthread::Int = Threads.nthreads()
)

    if how === :normal
        plan = fbp2_setup_normal(sg, ig,window, T)
    elseif how === :dsc
        #plan = fbp2_setup_dsc(sg,ig,how,window)
    elseif how === :df
        #plan = fbp2_setup_df(sg,ig,how,window)
    elseif how === :mojette
        #plan = fbp2_setup_moj(sg,ig,how,window)
    elseif how === :table
        #plan = fbp2_setup_tab(sg,ig,how,window)
    else
        throw("unknown type: $how")
    end

    return plan
end

function fbp2_par_parker_wt(sg::SinoGeom)
    orbit = abs(sg.orbit)
    na = sg.na
    nb = sg.nb
    ad = abs(sg.ad - sg.orbit_start)

    if orbit < 180
        @warn("orbit $orbit < 180")
        return sg.ones
    end

    orbit > 360 && throw("only 180 ≤ orbit ≤ 360 supported for Parker weighting")
    extra = orbit - 180 # extra beyond 180

    wt = ones(T,na)

    ii = ad .< extra
    wt[ii] = abs2.(sin(ad[ii] ./ extra .* pi ./ 2))
    ii = ad .>= 180
    wt[ii] = abs2.(sin((orbit .- ad[ii]) ./ extra .* pi ./ 2))
    wt = wt * orbit / 180 # trick because of the back-projector normalization
    return repeat(wt, nb, 1) # [nb na] sinogram sized
end

function fbp2_setup_normal(sg::SinoPar, ig::ImageGeom, window::Symbol, T::DataType)
    weight = ones(T,sg.nb,sg.na)
    if abs(sg.orbit) != 180 && abs(sg.orbit) != 360
        weight = fbp2_par_parker_wt(sg)
    end
    return NormalPlan(sg,ig,window,weight)
end

function fbp2_setup_normal(sg::SinoFan, ig::ImageGeom, window::Symbol, T::DataType)
    weight = ones(T, sg.nb, sg.na)
    sg.orbit != 360 && @warn("short-scan fan-beam Parker weighting not done")
    return NormalPlan(sg, ig, window, weight)
end

function fbp2_setup_normal(sg::SinoMoj, ig::ImageGeom, window::Symbol, T::DataType)
    weight = ones(T, sg.nb, sg.na)
    if sg.dx == abs(ig.dx)
        plan.moj.G = Gtomo2_table(sg, ig, ["mojette,back1"], nthread=nthread)
    else
        d = sg.dx
        dx = ig.dx
        @warn("mojette sinogram with d=$d vs image with dx=$dx")
    end

    plan.moj.H = fbp2_make_sino_filter_moj(sg.nb, sg.na, sg.dx, sg.orbit, sg.orbit_start, window)
    return NormalPlan{SinoMoj}(sg,ig,window,weight,moj)
end

"""
    image, sino_filt = fbp2(plan, sino)

Filtered back-projection (FBP) reconstruction,
returning image and filtered sinogram

in
- `plan::FBPplan`
- `sino::AbstractMatrix{<:Number}`

out
- `image::AbstractMatrix{<:Number}`       reconstructed image(s)
- `sino_filt::AbstractMatrix{<:Number}`   filtered sinogram(s)

"""
function fbp2(plan::NormalPlan, sino::AbstractMatrix{<:Number})

    plan.sg.dim != size(sino) && throw("bad sino size")

    if plan.sg isa SinoPar
        sino = sino .* plan.parker_weight

        sino,_,_,_ = fbp2_sino_filter(plan.sg, sino, ds = plan.sg.dr, window = plan.window)

        #image = fbp2_back(plan.sg, plan.ig, sino, do_r_mask=true)
        image = fbp2_back(plan.sg, plan.ig, sino)

    elseif plan.sg isa SinoFan

        dfs = plan.sg.dfs

        dfs != 0 && ~isinf(dfs) && throw("only arc or flat fan done")

        sino = fbp2_sino_weight(plan.sg, sino) #todo

        sino,_,_,_ = fbp2_sino_filter(
            plan.sg, sino,
            ds = plan.sg.ds, dsd = plan.sg.dsd,
            window = plan.window,
        )
        image = fbp2_back(plan.sg, plan.ig, sino)

    elseif plan.sg isa SinoMoj #TODO (incomplete)

        sino = fbp2_apply_sino_filter_moj(sino, plan.moj.H)

        if plan.sg.dx == abs(plan.ig.dx)
            image = plan.moj.G' * sino # backproject
            image = image * (pi / plan.sg.na) # account for "dphi" in integral
        else # revert to conventional pixel driven
            ig = plan.ig
            sg = plan.sg
            arg1 = [uint8(ig.mask), ig.dx, ig.dy, ig.offset_x, sign(ig.dy) * ig.offset_y] # trick: old backproject
            arg2 = [sg.d(1:sg.na), sg.offset, sg.orbit, sg.orbit_start]
# todo      image = jf_mex("back2", arg1[:], arg2[:], int32(arg.nthread), single(sino))
            image = image .* plan.ig.mask
        end

    else
        throw("not done")
    end

    return image, sino
end

# fbp2_recon_moj
function fpb2(plan::MojPlan, sino::AbstractMatrix{<:Number})
    moj = plan.moj
    # bin sinogram to mojette sampling, whether it is fan or parallel
    msino = moj.ob_rebin * sino
    msino = fbp2_apply_sino_filter_moj(msino, moj.H) # filter mojette sinogram
    image = moj.G' * msino # backproject
    return image * (π / size(sino,2)) # account for "dphi" in integral
end

# todo: needs broadcasts?
function fbp_make_sino_filter_moj(nb, na, dx, orbit, orbit_start, window)
    ang = deg2rad(orbit_start .+ (0:na-1)./na .* orbit)
    npad = 2^ceil(log2(2*nb-1)) # padded size

    dr = dx * max(abs(cos(ang)), abs(sin(ang)))
    if true
        junk, H = fbp2_sino_filter(:flat, ones(nb,1), ds=1, window=:window, decon1=0)

        # trick: ramp filter usually has 1/dr^2 in it, but convolution by sum
        # requires "dr" so we need dr / dr^2 = 1 / dr
        return H * (1 ./ dr) # [npad,na]

    else # trick: make sure cutoff frequencies match
        u0 = 1/2/dx # standard cutoff for coarsest sampling
        r = -(npad/2):(npad/2-1) * dr # [nb na]
        h = u0^2 .* (2 .* sinc.(2*u0*r) - sinc.(u0*r).^2)
        h = h .* repeat(dr, [npad 1]) # extra dr for discrete-space convolution
   #    clf, plot(r, h, '.'), keyboard
        return reale(fft(fftshift(h,1), 1))
   #    H = fbp_apodize(H, nb, window)
        !isempty(window) && throw("window not done yet due to dr")
    end
end

function fbp2_apply_sino_filter_moj(sino, H)
    nb = size(sino,1)
    npad = 2^ceil(log2(2*nb-1)) # padded size
    sinopad = [sino; zeros(npad-nb,size(sino,2))] # padded sinogram
    sino = ifft(reale(fft(sinopad, 1) .* H), 1)
    sino = sino[1:nb,:]
end


text:
using ScalarRadau
using DifferentialEquations
using BenchmarkTools
using Statistics
using PyPlot

##

#stiff, scalar ODE for testing
f_sr(x, y, _) = 50*(cos(x) - y)
f_de(u, _, t) = 50*(cos(t) - u)
#known solution
s(x) = (50*(-50*exp(-50*x) + sin(x) + 50*cos(x)))/2501

##

#number of tolerances to sample
N = 10
#tolerance values
tols = 10 .^ LinRange(-6, -2, N)
#times and errors for ScalarRadau evaluations
t_sr = zeros(N)
e_sr = zeros(N)
#times and errors for DifferentialEquations evaluations
t_de = zeros(N)
e_de = zeros(N)
#starting point
x₀ = 0.0
y₀ = 0.0
#finishing coordinate
xₙ = 3.0
#exact solution end values
sₙ = s(xₙ)
#the problem for DifferentialEquations
prob = ODEProblem(f_de, x₀, (x₀,xₙ))

##

for (i,tol) ∈ enumerate(tols)

    #ScalarRadau error
    yₙ = radau!(y, x, f_sr, y₀, x₀, xₙ, rtol=tol, atol=tol)
    e_sr[i] = abs(yₙ - sₙ)/abs(sₙ)
    #ScalarRadau timing
    b = @benchmark radau($f_sr, $y₀, $x₀, $xₙ, rtol=$tol, atol=$tol)
    t_sr[i] = mean(b.times)

    #DifferentialEquations error
    sol = solve(prob, RadauIIA5(), reltol=tol, abstol=tol, save_start=false, saveat=xₙ)
    e_de[i] = abs(sol.u[end] - sₙ)/abs(sₙ)
    #DifferentialEquations timing
    b = @benchmark solve($prob, $RadauIIA5(), reltol=$tol, abstol=$tol, save_start=false, saveat=$xₙ)
    t_de[i] = mean(b.times)

end

ext: jl


text:
export NoDiffusion

struct NoDiffusion <: DiffusionComponent end

max_stars_repo_name: lars-frogner/Whirl

max_issues_repo_name: lars-frogner/Whirl

max_forks_repo_name: lars-frogner/Whirl


text:
using Logging
using JuMP, CPLEX
using LightGraphs, SimpleWeightedGraphs
using PrettyTables, Formatting, ProgressMeter, Crayons.Box

function value_rounded(v)
    return round(Int, value(v))
end

function basic_kmst!(model,
        graph :: SimpleWeightedGraph,
        k :: Int;
        lowerbound :: Int = typemin(Int),
        upperbound :: Int = typemax(Int),
        isdebug=false)
    n :: Int = nv(graph)
    m :: Int = ne(graph)
    progress::Progress = Progress(1 + m + 2 * m + m + 3 + 1,
                                    dt=0.5,
                                    barglyphs=BarGlyphs('[', '#', ['-','~','+','*','=','>','#'], ' ', ']'),
                                    enabled=isdebug)
    es = collect(SimpleWeightedEdge, edges(graph))
    sort!(es, by=weight)

    lb::Int = sum(round(Int, weight(e)) for e in Iterators.take(Iterators.drop(es, n-1), k-1))
    if lb > lowerbound
        @debug "Objective has $(string(CYAN_FG(string(lb)))) as natural lowerbound."
    end
    lb = max(lb, lowerbound)
    ub::Int = sum(round(Int, weight(e)) for e in Iterators.drop(es, n-(k-1)))
    if ub < upperbound
        @debug "Objective has $(string(MAGENTA_FG(string(ub)))) as natural upperbound."
    end
    ub = min(ub, upperbound)
    sort!(es, by=dst, alg=MergeSort)
    sort!(es, by=src, alg=MergeSort)
    @variable(model, o, Int, lower_bound=lb, upper_bound=ub)
    next!(progress)
    @variable(model, x[e=1:m], Bin, lower_bound=0, upper_bound=1)
    ProgressMeter.update!(progress, m)
    y::Dict{Tuple{Int,Int}, VariableRef} = Dict{Tuple{Int,Int}, VariableRef}()
    for e in es
        i::Int = src(e)
        j::Int = dst(e)
        y[i,j] = @variable(model, base_name="y[$i,$j]", binary=true, lower_bound=0, upper_bound=1) # Set bounds, that solver doesn't have to set it implicitly
        next!(progress)
        y[j,i] = @variable(model, base_name="y[$j,$i]", binary=true, lower_bound=0, upper_bound=1)
        next!(progress)
    end

    for i in 1:m
        e = es[i]
        s = src(e)
        d = dst(e)
        @constraint(model, y[s, d] + y[d, s] == x[i]) # Constraint 2
        next!(progress)
    end

    @constraints(model, begin
        sum(x[i] for i in 1:m if src(es[i]) != 1 && dst(es[i]) != 1) == k-1 # Constraint 3
        sum(x[e] * Int(round(Int, weight(es[e]))) for e in 1:m) == o # Constraint 4
    end)
    ProgressMeter.update!(progress, 3)
    @objective(model, Min, o)
    next!(progress)
    finish!(progress)
    #MOI.set(model, MOI.HeuristicCallback(), heuristic!(model, graph, k))
end

function basic_kmst_warmstart!(model, graph :: SimpleWeightedGraph, k :: Int, solution :: Vector{Edge{Int}})
    n :: Int = nv(graph)
    m :: Int = ne(graph)
    es :: Vector{SimpleWeightedEdge} = collect(edges(graph))
    sort!(es, by=dst)
    sort!(es, by=src, alg=MergeSort)
    x :: Vector{Int} = zeros(Int, m)
    totalweight :: Int = 0
    for s in solution
        for r in 1:m
            e = es[r]
            i = src(e)
            j = dst(e)
            w = weight(e)
            if (src(s) == i && dst(s) == j) || (dst(s) == i && src(s) == j)
                x[r] = 1
                totalweight += Int(round(Int ,w))
                break
            end
        end
    end
    for r in 1:m
        e = es[r]
        s = src(e)
        d = dst(e)
        if src(e) != 1 && dst(e) != 1
            set_start_value(variable_by_name(model, "x[$r]"), x[r])
        end
    end
    set_start_value(variable_by_name(model, "o"), totalweight)
end

function check_kmst_warmstart!(model, graph :: SimpleWeightedGraph, k :: Int) :: Int
    n :: Int = nv(graph)
    m :: Int = ne(graph)
    es :: Vector{SimpleWeightedEdge} = collect(edges(graph))
    sort!(es, by=dst)
    sort!(es, by=src, alg=MergeSort)
    x :: Vector{Union{Nothing,Int}} = fill(nothing, m)
    y :: Matrix{Union{Nothing,Int}} = fill(nothing, n, n)
    filled :: Int = 0
    totalweight :: Int = 0
    for r in 1:m
        e = es[r]
        i = src(e)
        j = dst(e)
        x[r] = start_value(variable_by_name(model, "x[$r]"))
        if x[r] == 1
            totalweight += Int(round(Int, weight(e)))
        end
        y[i, j] = start_value(variable_by_name(model, "y[$i,$j]"))
        y[j, i] = start_value(variable_by_name(model, "y[$j,$i]"))
        if isnothing(y[i,j]) && !isnothing(y[j,i]) && !isnothing(x[r])
            y[i,j] = x[r] - y[j,i]
            set_start_value(variable_by_name(model, "y[$i,$j]"), y[i,j])
            filled += 1
        elseif !isnothing(y[i,j]) && isnothing(y[j,i]) && !isnothing(x[r])
            y[j,i] = x[r] - y[i,j]
            set_start_value(variable_by_name(model, "y[$j,$i]"), y[j,i])
            filled += 1
        elseif !isnothing(y[i,j]) && !isnothing(y[j,i]) && isnothing(x[r])
            if 0 <= y[i,j] + y[j,i] <= 1
                x[r] = y[i,j] + y[j,i]
                set_start_value(variable_by_name(model, "x[$r]"), x[r])
                filled += 1
            else
                @warn "0 > y[$i,$j] + y[$j,$i] > 1 | ($(y[i,j]) + $(y[j,i]))"
            end
        elseif !isnothing(y[i,j]) && !isnothing(y[j,i]) && !isnothing(x[r]) && x[r] != y[i,j] + y[j,i]
            @warn "x[$r] != y[$i,$j] + y[$j,$i] | ($(x[r]) != $(y[i,j]) + $(y[j,i])"
        end
    end

    return filled

end

function heuristic!(model,
                    graph :: SimpleWeightedGraph,
                    k :: Int)
    y::Dict{Tuple{Int, Int}, VariableRef} = Dict{Tuple{Int, Int}, VariableRef}()
    y_val::Dict{Tuple{Int, Int}, Float64} = Dict{Tuple{Int, Int}, Float64}()
    x::Dict{Int, VariableRef} = Dict{Int, VariableRef}()
    x_val::Dict{Int, Float64} = Dict{Int, Float64}()
    o::VariableRef = variable_by_name(model, "o")
    o_val::Float64 = 0
    es = collect(edges(graph))
    sort!(es, by=dst)
    sort!(es, by=src, alg=MergeSort)
    r::Int = 0
    for e in es
        r += 1
        x[r] = variable_by_name(model, "x[$r]")
        i::Int = src(e)
        j::Int = dst(e)
        y[i,j] = variable_by_name(model, "y[$i,$j]")
        y[j,i] = variable_by_name(model, "y[$j,$i]")
    end
    return function(cb_data)
        variables::Vector{VariableRef} = Vector{VariableRef}()
        values::Vector{Float64} = Vector{Float64}()
        status = callback_node_status(cb_data, model)
        o_val = callback_value(cb_data, o)
        for ((i,j), v) in y
            y_val[i,j] = callback_value(cb_data, v)
        end
        for (k, v) in x
            x_val[k] = callback_value(cb_data, v)
        end
        if status == MOI.CALLBACK_NODE_STATUS_FRACTIONAL
            @debug "Found fractional solution. Attempting to round."
            arcs::Vector{Tuple{Int,Int,Int}} = Vector{Tuple{Int,Int,Int}}()
            r = 0
            for e in es
                r += 1
                i::Int = src(e)
                j::Int = dst(e)
                push!(arcs, (r,i,j))
                push!(arcs, (r,j,i))
            end

            sort!(arcs, by=(a -> a[3]))
            sort!(arcs, by=(a -> a[2]), alg=MergeSort)
            sort!(arcs, by=(a -> weight(es[a[1]])), alg=MergeSort)
            sort!(arcs, by=(a -> y_val[a[2],a[3]]), alg=MergeSort, rev=true)
            sort!(arcs, by=(a -> a[2] == 1 || a[3] == 1), alg=MergeSort)

            visited::Vector{Bool} = zeros(Bool, nv(graph))
            parents::Vector{Int} = zeros(Int, nv(graph))
            lastnode::Int = 1
            treesize::Int = 1
            visited[1] = true
            parents[1] = -1

            includededges::Vector{Bool} = zeros(Bool, length(es))

            while treesize < k+1
                for (r,i,j) in arcs
                    (visited[i] && !visited[j]) || continue
                    visited[j] = true
                    includededges[r] = true
                    parents[i] = lastnode
                    parents[j] = i
                    lastnode = j
                    treesize += 1
                    @debug "Including $r ($i -> $j)"
                    break
                end
                filter!(v -> !visited[v[3]], arcs)
            end
            @assert count(visited) == k+1 "Tree should have size $k, but has size $(count(visited))."

            push!(variables, o)
            push!(values, 0.0)

            r = 0
            for e in es
                r += 1
                if includededges[r]
                    values[1] += round(weight(e))
                    push!(variables, x[r])
                    push!(values, 1.0)
                end
            end

            # for ((i,j),v) in y
            #    push!(variables, v)
            #    push!(values, Float64(visited[i] && visited[j] && (parents[j] == i || parents[i] == j)))
            # end

            status = MOI.submit(model, MOI.HeuristicSolution(cb_data), variables, values)

            debugstring = "Found solution with weight $(round(Int, values[1])), "
            if status == MOI.HEURISTIC_SOLUTION_ACCEPTED
                debugstring *= string(GREEN_FG("accepted"))
            elseif status == MOI.HEURISTIC_SOLUTION_REJECTED
                debugstring *= string(RED_FG("rejected"))
            else
                @assert status == MOI.HEURISTIC_SOLUTION_UNKNOWN
                debugstring *= string(YELLOW_FG("unknown status"))
            end
            @debug debugstring * "."

        elseif status == MOI.CALLBACK_NODE_STATUS_INTEGER
            @debug "Found integer solution. Attempting to permute."
        else
            @assert status == MOI.CALLBACK_NODE_STATUS_UNKNOWN
        end
    end
end

struct KMSTSolution
    termination_status :: MOI.TerminationStatusCode
    solution_graph :: Union{SimpleWeightedGraph, Nothing}
    objective_value :: Float64
    relative_gap :: Float64
    solve_time_sec :: Float64
    bnb_nodes :: Union{Int, Nothing}
    violated_ineq :: Union{Int, Nothing}
end

function print_variable_table(model, graph :: SimpleWeightedGraph)
    if !has_values(model)
        @debug "Cannot print values, model has no values."
        return
    end
    @assert has_values(model) "Model should have values."
    n::Int= nv(graph)
    m::Int= ne(graph)
    es::Vector{SimpleWeightedEdge}= collect(edges(graph))
    sort!(es, by=dst)
    sort!(es, by=src, alg=MergeSort)

    all_vars::Vector{VariableRef} = all_variables(model)

    different_variables::Int = 0

    header :: Vector{String} = ["", "1-Vars"]
    names :: Vector{String} = []
    multidimvarslines :: Dict{String, Int} = Dict()
    multidimvarsdims :: Dict{String, Int} = Dict()
    multilines::Int = 0
    onevarsline::Int = 1
    onevarslines::Int = 0

    for var in all_vars
        name::String = JuMP.name(var)

        delim::Union{Int, Nothing} = findfirst('[', name)

        if isnothing(delim)
            onevarslines += 1
        elseif delim > 1
            basename::String = name[1:(delim-1)]
            if basename in keys(multidimvarsdims)
                @assert basename in keys(multidimvarslines) "If basename $basename is in multidimvarsdims it should be in multidimvarslines."
                @assert basename in names "If basename $basename is in multidimvarsdims it should be in names."
                multidimvarslines[basename] += 1
            else
                @assert !(basename in keys(multidimvarslines)) "If basename $basename is not in multidimvarsdims it should not be in multidimvarslines."
                @assert !(basename in names) "If basename $basename is in multidimvarsdims it should not be in names."
                multidimvarslines[basename] = 1
                multidimvarsdims[basename] = count(v -> v == ',', name) + 1
                push!(header, "")
                push!(header, basename)
                push!(names, basename)
            end
            @assert basename in keys(multidimvarslines) "Basename $basename should be present in multidimvarslines"
            @assert basename in keys(multidimvarsdims) "Basename $basename should be present in multidimvarsdims"
        end
    end
    if !isempty(multidimvarslines)
        multilines = maximum(values(multidimvarslines))
    end

    width::Int = length(header)
    height::Int = max(onevarslines, multilines)

    data::Matrix{Union{String, Int, Float64}} = fill("", height, width)

    onevarslines = 1
    for k in keys(multidimvarslines)
        multidimvarslines[k] = 1
    end

    for var in all_vars
        name :: String = JuMP.name(var)
        delim::Union{Int, Nothing} = findfirst('[', name)
        if isnothing(delim)
            @assert onevarslines <= height "1-vars exceeds height."
            data[onevarslines, 1] = name
            data[onevarslines, 2] = value(var)
            onevarslines += 1
        elseif delim > 1
            if round(value(var)) != 0
                basename::String = name[1:(delim-1)]
                @assert basename in keys(multidimvarslines) "Basename $basename never seen before. Not found in multidimvarslines."
                dims::String = name[delim:end]
                @assert basename in names "Basename $basename never seen before. Not found in names."
                index::Int = filter(v -> v[2] == basename, collect(enumerate(names)))[1][1]
                @assert 1 <= index * 2 + 1 <= width "$basename index ($index) exceeds width."
                @assert 1 <= multidimvarslines[basename] <= height "$name exceeds height."
                data[multidimvarslines[basename], index * 2 + 1] = dims
                if basename == "x"
                    e = es[parse(Int, dims[2:(end-1)])]
                    data[multidimvarslines[basename], index * 2 + 1] *= " = [$(src(e)),$(dst(e))]"
                end
                data[multidimvarslines[basename], index * 2 + 2] = value(var)
                multidimvarslines[basename] += 1
            end
        end
    end

    if !isempty(multidimvarslines)
        multilines = maximum(values(multidimvarslines))
    end
    height = max(onevarslines, multilines)

    data_view = view(data, 1:(height-1), :)

    if height > 45
        pretty_table(data_view, header=header, backend=Val(:html))
    else
        pretty_table(data_view, header=header)
    end

end

function solve!(model,
                graph :: SimpleWeightedGraph,
                k :: Int,
                mode :: SolvingMode) :: KMSTSolution
    ts :: MOI.TerminationStatusCode = MOI.INTERRUPTED
    @debug "Starting to optimize."
    optimizationtime = @elapsed optimize!(model)
    @debug "Finished optimizing in $(format_seconds_readable(optimizationtime, 2))."
    ts = termination_status(model)
    st :: Float64 = solve_time(model)
    if ts != MOI.OPTIMAL
        warning = "Optimizer did not find an optimal solution."
        if ts == MOI.TIME_LIMIT
            tl::Float64 = time_limit_sec(model)
            warning *= " Time limit of $(format_seconds_readable(tl)) reached"
            if tl < st
                warning *= " after $(format_seconds_readable(st)), overdrawn by $(format_seconds_readable(st - tl))"
            end
            warning *= "."
        elseif ts == MOI.MEMORY_LIMIT
            warning *= " Memory limit reached."
        else
           if ts == MOI.INFEASIBLE && mode != cec && mode != dcc
              if lowercase(solver_name(model)) != "cplex"
                  @debug "Setting optimizer from $(solver_name(model)) to CPLEX."
                  set_optimizer(model, CPLEX.Optimizer)
              end
              compute_conflict!(model)

              constraint_types::Vector{Tuple{DataType, DataType}} = list_of_constraint_types(model)

              debugstring::String = "Following constraints are part of the conflict:"

              for (ref, typ) in constraint_types
                  !(ref == VariableRef && typ == MOI.Integer) || continue
                  !(ref == VariableRef && typ == MOI.ZeroOne) || continue
                  !(ref == GenericAffExpr{Float64, VariableRef} && typ == MOI.Interval{Float64}) || continue
                  all_cons::Vector{ConstraintRef} = all_constraints(model, ref, typ)
                  @debug "$ref: $typ"
                  for con in all_cons
                      if MOI.get(model.moi_backend, MOI.ConstraintConflictStatus(), con.index) != MOI.NOT_IN_CONFLICT
                          debugstring *= "\n$con"
                      end
                  end
              end
              warning *= " Problem is infeasible."
              @debug debugstring
          else
              warning *= " $ts"
          end
          print_variable_table(model, graph)
        end
        @warn warning
    end

    bnbn::Union{Nothing, Int} = nothing
    if lowercase(solver_name(model)) == "cplex"
        bnbn = node_count(model)
    end
    lc::Union{Nothing, Int} = nothing
    if mode == cec
        lc = lazy_cec_constraints
    end
    if mode == dcc
        lc = lazy_dcc_constraints
    end
    if !has_values(model)
        return KMSTSolution(ts, nothing, Inf64, Inf64, st, bnbn, lc)
    end

    rg::Float64 = relative_gap(model)
    n::Int = nv(graph)
    solution_graph::SimpleWeightedGraph = get_solution_graph(model, graph)
    components = connected_components(solution_graph)
    bug::Bool = false
    if ne(solution_graph) != k-1
        @error "Solution has wrong number of edges."
        @debug "Should have $(k - 1) edges but has $(ne(solution_graph))."
        bug = true
    end
    if n - k + 1 != length(components)
        @error "Solution has wrong number of components."
        @debug "Should have $(n - k + 1) components but has $(length(components))."
        bug = true
    end
    for component in components
        cl::Int = length(component)
        if cl != 1 && cl != k
            @error "Solution is not connected." maxlog=1
            @debug "Should only have components with size 1 or $k but found with size $cl: $(component)."
            bug = true
        end
    end
    ov::Float64 = objective_value(model)
    ovrounded::Int = Int(round(ov))
    w::Int = sum(Int(round(Int,weight(e))) for e in edges(solution_graph))
    if w != ovrounded
        @warn "MILP solver claims objective value is $ovrounded ($ov), however got $(w)."
        bug = true
    end
    if bug
        print_variable_table(model, graph)
    end
    return KMSTSolution(ts, solution_graph, ov, rg, st, bnbn, lc)
end

function get_solution_graph(model, graph :: SimpleWeightedGraph)
    n :: Int = nv(graph)
    solution_graph = SimpleWeightedGraph(n)
    es :: Vector{SimpleWeightedEdge} = collect(edges(graph))
    sort!(es, by=weight)
    sort!(es, by=dst, alg=MergeSort)
    sort!(es, by=src, alg=MergeSort)
    for e in es
        i = src(e)
        j = dst(e)
        w = weight(e)
        if i > 1 && j > 1
            x_ij = variable_by_name(model, "y[$i,$j]")
            xv = Int(round(value(x_ij)))
            if xv == 1
                add_edge!(solution_graph, i, j, w)
            else
                x_ji = variable_by_name(model, "y[$j,$i]")
                xv = Int(round(value(x_ji)))
                if xv == 1
                    add_edge!(solution_graph, i, j, w)
                end
            end
        end
    end

    return solution_graph
end

max_stars_repo_name: Entze/jkMST

max_issues_repo_name: Entze/jkMST

max_forks_repo_name: Entze/jkMST


text:
# Wilcoxon.jl
# Wilcoxon rank sum (Mann-Whitney U) tests
#
# Copyright (C) 2012   Simon Kornblith
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

export MannWhitneyUTest, ExactMannWhitneyUTest, ApproximateMannWhitneyUTest

## COMMON MANN-WHITNEY U

# Automatic exact/normal selection
"""
    MannWhitneyUTest(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Perform a Mann-Whitney U test of the null hypothesis that the probability that an
observation drawn from the same population as `x` is greater than an observation drawn
from the same population as `y` is equal to the probability that an observation drawn
from the same population as `y` is greater than an observation drawn from the same
population as `x` against the alternative hypothesis that these probabilities are not
equal.

The Mann-Whitney U test is sometimes known as the Wilcoxon rank-sum test.

When there are no tied ranks and ≤50 samples, or tied ranks and ≤10 samples,
`MannWhitneyUTest` performs an exact Mann-Whitney U test. In all other cases,
`MannWhitneyUTest` performs an approximate Mann-Whitney U test. Behavior may be further
controlled by using [`ExactMannWhitneyUTest`](@ref) or [`ApproximateMannWhitneyUTest`](@ref)
directly.

Implements: [`pvalue`](@ref)
"""
function MannWhitneyUTest(x::AbstractVector{S}, y::AbstractVector{T}) where {S<:Real,T<:Real}
    (U, ranks, tieadj, nx, ny, median) = mwustats(x, y)
    if nx + ny <= 10 || (nx + ny <= 50 && tieadj == 0)
        ExactMannWhitneyUTest(U, ranks, tieadj, nx, ny, median)
    else
        ApproximateMannWhitneyUTest(U, ranks, tieadj, nx, ny, median)
    end
end

# Get U, ranks, and tie adjustment for Mann-Whitney U test
function mwustats(x::AbstractVector{S}, y::AbstractVector{T}) where {S<:Real,T<:Real}
    nx = length(x)
    ny = length(y)
    if nx <= ny
        (ranks, tieadj) = tiedrank_adj([x; y])
        U = sum(ranks[1:nx]) - nx*(nx+1)/2
    else
        (ranks, tieadj) = tiedrank_adj([y; x])
        U = nx*ny - sum(ranks[1:ny]) + ny*(ny+1)/2
    end
    (U, ranks, tieadj, nx, ny, median(x)-median(y))
end

## EXACT MANN-WHITNEY U TEST
struct ExactMannWhitneyUTest{T<:Real} <: HypothesisTest
    U::Float64              # test statistic: Mann-Whitney-U statistic
    ranks::Vector{Float64}  # ranks
    tie_adjustment::Float64 # adjustment for ties
    nx::Int                 # number of observations
    ny::Int
    median::T               # sample median
end

"""
    ExactMannWhitneyUTest(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Perform an exact Mann-Whitney U test of the null hypothesis that the probability that an
observation drawn from the same population as `x` is greater than an observation drawn
from the same population as `y` is equal to the probability that an observation drawn
from the same population as `y` is greater than an observation drawn from the same
population as `x` against the alternative hypothesis that these probabilities are not
equal.

When there are no tied ranks, the exact p-value is computed using the `pwilcox` function
from the `Rmath` package. In the presence of tied ranks, a p-value is computed by exhaustive
enumeration of permutations, which can be very slow for even moderately sized data sets.

Implements: [`pvalue`](@ref)
"""
ExactMannWhitneyUTest(x::AbstractVector{S}, y::AbstractVector{T}) where {S<:Real,T<:Real} =
    ExactMannWhitneyUTest(mwustats(x, y)...)

testname(::ExactMannWhitneyUTest) = "Exact Mann-Whitney U test"
population_param_of_interest(x::ExactMannWhitneyUTest) = ("Location parameter (pseudomedian)", 0, x.median) # parameter of interest: name, value under h0, point estimate
default_tail(test::ExactMannWhitneyUTest) = :both

function show_params(io::IO, x::ExactMannWhitneyUTest, ident)
    println(io, ident, "number of observations in each group: ", [x.nx, x.ny])
    println(io, ident, "Mann-Whitney-U statistic:             ", x.U)
    println(io, ident, "rank sums:                            ", [sum(x.ranks[1:x.nx]), sum(x.ranks[x.nx+1:end])])
    println(io, ident, "adjustment for ties:                  ", x.tie_adjustment)
end

# Enumerate all possible Mann-Whitney U results for a given vector,
# determining left-and right-tailed p values
function mwuenumerate(x::ExactMannWhitneyUTest)
    # Get the other U if inverted by mwu_stats
    n = min(x.nx, x.ny)
    U = x.U
    if x.ny > x.nx
        U = x.nx*x.ny -x.U
    end
    le = 0
    gr = 0
    tot = 0
    k = n*(n+1)/2
    for comb in combinations(x.ranks, n)
        Up = sum(comb) - k
        tot += 1
        le += Up <= U
        gr += Up >= U
    end
    (le/tot, gr/tot)
end

function pvalue(x::ExactMannWhitneyUTest; tail=:both)
    if x.tie_adjustment == 0
        # Compute exact p-value using method from Rmath, which is fast but
        # cannot account for ties
        if tail == :both
            if x.U < x.nx * x.ny / 2
                2 * pwilcox(x.U, x.nx, x.ny, true)
            else
                2 * pwilcox(x.U - 1, x.nx, x.ny, false)
            end
        elseif tail == :left
            pwilcox(x.U, x.nx, x.ny, true)
        elseif tail == :right
            pwilcox(x.U - 1, x.nx, x.ny, false)
        else
            throw(ArgumentError("tail=$(tail) is invalid"))
        end
    else
        # Compute exact p-value by enumerating possible ranks in the tied data
        if tail == :both
            min(1, 2 * minimum(mwuenumerate(x)))
        elseif tail == :left
            mwuenumerate(x)[1]
        elseif tail == :right
            mwuenumerate(x)[2]
        else
            throw(ArgumentError("tail=$(tail) is invalid"))
        end
    end
end

struct ApproximateMannWhitneyUTest{T<:Real} <: HypothesisTest
    U::Float64              # test statistic: Mann-Whitney-U statistic
    ranks::Vector{T}        # ranks
    tie_adjustment::Float64 # adjustment for ties
    nx::Int                 # number of observations
    ny::Int
    median::Float64         # sample median
    mu::Float64             # normal approximation: mean
    sigma::Float64          # normal approximation: std
end

## APPROXIMATE MANN-WHITNEY U TEST
"""
    ApproximateMannWhitneyUTest(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Perform an approximate Mann-Whitney U test of the null hypothesis that the probability that
an observation drawn from the same population as `x` is greater than an observation drawn
from the same population as `y` is equal to the probability that an observation drawn
from the same population as `y` is greater than an observation drawn from the same
population as `x` against the alternative hypothesis that these probabilities are not
equal.

The p-value is computed using a normal approximation to the distribution of the
Mann-Whitney U statistic:
```math
    \\begin{align*}
        μ & = \\frac{n_x n_y}{2}\\\\
        σ & = \\frac{n_x n_y}{12}\\left(n_x + n_y + 1 - \\frac{a}{(n_x + n_y)(n_x +
            n_y - 1)}\\right)\\\\
        a & = \\sum_{t \\in \\mathcal{T}} t^3 - t
    \\end{align*}
```
where ``\\mathcal{T}`` is the set of the counts of tied values at each tied position.

Implements: [`pvalue`](@ref)
"""
function ApproximateMannWhitneyUTest(U::Real, ranks::AbstractVector{T},
    tie_adjustment::Real, nx::Int, ny::Int, median::Float64) where T<:Real
    mu = U - nx * ny / 2
    sigma = sqrt((nx * ny * (nx + ny + 1 - tie_adjustment /
        ((nx + ny) * (nx + ny - 1)))) / 12)
    ApproximateMannWhitneyUTest(U, ranks, tie_adjustment, nx, ny, median, mu, sigma)
end
ApproximateMannWhitneyUTest(x::AbstractVector{S}, y::AbstractVector{T}) where {S<:Real,T<:Real} =
    ApproximateMannWhitneyUTest(mwustats(x, y)...)

testname(::ApproximateMannWhitneyUTest) = "Approximate Mann-Whitney U test"
population_param_of_interest(x::ApproximateMannWhitneyUTest) = ("Location parameter (pseudomedian)", 0, x.median) # parameter of interest: name, value under h0, point estimate
default_tail(test::ApproximateMannWhitneyUTest) = :both

function show_params(io::IO, x::ApproximateMannWhitneyUTest, ident)
    println(io, ident, "number of observations in each group: ", [x.nx, x.ny])
    println(io, ident, "Mann-Whitney-U statistic:             ", x.U)
    println(io, ident, "rank sums:                            ", [sum(x.ranks[1:x.nx]), sum(x.ranks[x.nx+1:end])])
    println(io, ident, "adjustment for ties:                  ", x.tie_adjustment)
    println(io, ident, "normal approximation (μ, σ):          ", (x.mu, x.sigma))
end

function pvalue(x::ApproximateMannWhitneyUTest; tail=:both)
    if x.mu == x.sigma == 0
        1
    else
        if tail == :both
            2 * ccdf(Normal(), abs(x.mu - 0.5 * sign(x.mu))/x.sigma)
        elseif tail == :left
            cdf(Normal(), (x.mu + 0.5)/x.sigma)
        elseif tail == :right
            ccdf(Normal(), (x.mu - 0.5)/x.sigma)
        end
    end
end


text:
"""
module FymCore

The main core codes for `FymEnvs`.

"""
module FymCore

import FymEnvs: close!

using Reexport
include("FymLogging.jl")
@reexport using .FymLogging

using ProgressMeter
# using Debugger

export BaseEnv
export BaseSystem
export Clock

export update!, close!, reset!, render
export observe_array, observe_dict, observe_flat
export systems!, dyn!, step!, sys
export time_over, time

############ Clock ############
"""
    Clock

# Arguments
max_t: maximum (minimum) terminal time for forward (backward) integration,
respectively.
"""
mutable struct Clock
    t::Float64
    dt::Float64
    max_t::Float64
    max_len::Int64
    thist::Array
    ode_step_len::Int64
    Clock(args...; kwargs...) = init!(new(), args...; kwargs...)
end

function init!(clock::Clock, dt, ode_step_len; max_t=10.0)
    clock.dt = Float64(dt)
    clock.max_t = Float64(max_t)
    clock.max_len = floor(Int, (max_t / dt) + 1)
    clock.ode_step_len = Int64(ode_step_len)
    clock.thist = collect(range(0.0,
                                stop=dt, length=clock.ode_step_len + 1))
    return clock
end

function reset!(clock::Clock; t=0.0)
    clock.t = t
    return clock
end

function tick!(clock::Clock)
    clock.t += clock.dt
end

function time!(clock::Clock, t)
    clock.t = t  # array
end

"Get the current time from clock."
function Base.time(clock::Clock)
    return clock.t
end

"Check if the time is larger than max_t."
function time_over(clock::Clock; t=nothing)
    if t == nothing
        t = time(clock)
    end
    return sign(clock.dt) * (t - clock.max_t) > 0.5 * abs(clock.dt)
end

function thist(clock::Clock)
    thist = clock.thist .+ time(clock)
    if time_over(clock, t=thist[end])
        index = findfirst(sign(clock.dt)*(thist .- clock.max_t) .> 0.5 * abs(clock.dt))
        if index == nothing
            return thist
        else
            return thist[1:index[1]]
        end
    else
        return thist
    end
end

############ BaseSystem ############
mutable struct BaseSystem
    initial_state::Array
    state::Array
    state_size::Tuple
    state_length::Int64
    flat_index::AbstractRange
    dot::Array
    name::Union{String, Nothing}
    BaseSystem(; kwargs...) = init!(new(); kwargs...)
end

function _show(sys::BaseSystem; i=0)
    result = []
    _stack_property!(result, sys,
                     [:name, :state, :dot, :initial_state,
                      :state_size, :flat_index], i=i)
    return join(result, "\n")
end

function Base.show(io::IO, sys::BaseSystem)
    res = _show(sys)
    println(io, res)
end

function init!(sys::BaseSystem; initial_state=nothing,
               state_size=(1, 1), name=nothing)
    if initial_state == nothing
        initial_state = zeros(state_size)
    end
    sys.initial_state = convert(Array{Float64}, initial_state)
    sys.state_size = size(sys.initial_state)
    sys.name = name
    return sys
end

function reset!(sys::BaseSystem)
    state!(sys, sys.initial_state)
    return sys.state
end

function state(sys::BaseSystem)
    return sys.state
end

function state!(sys::BaseSystem, state)
    sys.state = state
end

function dot(sys::BaseSystem)
    return sys.dot
end

function dot!(sys::BaseSystem, dot)
    sys.dot = dot
end

############ BaseEnv ############
mutable struct BaseEnv
    name
    systems::Dict{Any, Union{BaseEnv, BaseSystem}}
    dyn
    step

    initial_time
    dt::Float64
    clock::Clock
    progressbar
    logger

    solver
    ode_func
    ode_option::Dict
    dot::Array
    flat_index::AbstractRange
    state_length::Int
    state_size::Tuple

    # params::Dict

    BaseEnv(args...; kwargs...) = init!(new(), args...; kwargs...)
end

function _add_space(string, i; space=" "^4)
    return space^i * string
end

function _stack_property!(result, sys, symbol_array; i=0, space="")
    for symbol in symbol_array
        if isdefined(sys, symbol)
            value = getproperty(sys, symbol)
        else
            value = "undef"
        end
        if symbol == :name
            space = "+---"
        else
            space = "|   "
        end
        push!(result, _add_space("$(String(symbol)): $value", i, space=space))
        # push!(result, _add_space("$(String(symbol)): $value", 0))
    end
end

function _show(env::BaseEnv; i=0)
    result = []
    _stack_property!(result, env, [:name])  # from env
    if i == 0
        _stack_property!(result, env.clock, [:max_t, :dt])  # from env.clock
    end
    for system in _systems(env)
        if typeof(system) == BaseSystem
            v_str = _show(system, i=i+1)
        elseif typeof(system) == BaseEnv
            v_str = _show(system, i=i+1)
        end
        push!(result, v_str)
    end
    res = join(result, "\n")
end

function Base.show(io::IO, env::BaseEnv)
    res = _show(env)
    println(res)
end

function init!(env::BaseEnv;
               systems=Dict(), dyn=nothing, step=nothing,
               # params=Dict(),
               initial_time=0.0, dt=0.01, max_t=1.0,
               ode_step_len=1,
               logger=nothing, ode_option=Dict(), solver="rk4",
               name=nothing,
              )
    env.name = name
    systems!(env, systems)

    env.clock = Clock(dt, ode_step_len, max_t=max_t)
    env.dt = env.clock.dt
    env.initial_time = initial_time

    env.logger = logger

    if solver == "rk4"
        env.solver = rk4
    else
        error("Unsupported solver")
    end
    dyn!(env, dyn)
    step!(env, step)
    env.ode_func = ode_wrapper(env)
    env.ode_option = ode_option
    env.progressbar = nothing
    # env.params = params
    return env
end

function ode_wrapper(env::BaseEnv)
    wrapper = function(y, kwargs, t)
        for system in _systems(env)
            state!(system, reshape(y[system.flat_index],
                                   system.state_size))
        end
        env.dyn(env, t; kwargs...)
        res = vcat([dot(system)[:] for system in _systems(env)]...)
        return res
    end
    return wrapper
end

function indexing!(env::BaseEnv)
    start = 0
    for system in _systems(env)
        system.state_length = prod(system.state_size)
        system.flat_index = (start+1):(start + system.state_length)
        start += system.state_length
    end
    tmp = [system.state_length for system in _systems(env)]
    if length(tmp) == 0
        env.state_size = (0,)
    else
        env.state_size = (sum(tmp),)
    end
end

"""
    reset!(env::BaseEnv)

Reset BaseEnv.
All initial_state will be assigned to state of BaseSystems.
It is recommended to extend this method when using custom FymEnv.
"""
function reset!(env::BaseEnv)
    for system in _systems(env)
        reset!(system)
    end
    reset!(env.clock, t=env.initial_time)
end

function state(env::BaseEnv)
    return observe_flat(env)
end

function sys(env::BaseEnv, name)
    return env.systems[String(name)]
end

function _systems(env::BaseEnv)
    return values(systems(env))
end

function systems(env::BaseEnv)
    return env.systems
end

"""
    systems!(env::BaseEnv, systems::Dict)

Set systems of `env`. Required before reset!.

# Examples
```julia
systems = Dict("sys" => BaseSystem(initial_state=zeros(3)))
env = BaseEnv()
systems!(env, systems)
```
"""
function systems!(env::BaseEnv, systems::Dict)
    env.systems = systems
    indexing!(env)
end

"""
    dyn!(env::BaseEnv, dyn)
Set dynamics of each systems in `env`. Required before reset!.
`dot` of all systems of `env` should be assigned in function `dyn`.

# Examples
```julia
function set_dyn(env, t)
    sys = env.systems["sys"]
    x = sys.state
    A = Matrix(I, 3, 3)
    sys.dot = -A * x
end
env = BaseEnv()
dyn!(env, set_dyn)
```
"""
function dyn!(env::BaseEnv, dyn)
    env.dyn = dyn
end

"""
    step!(env::BaseEnv, step)
Set transition behaviour of `env` for each step.
Required before reset!.
It must contain `update!`.

# Examples
```julia
function step(env)
    t = time(env.clock)
    sys = env.systems["sys"]
    x = sys.state
    update!(env)
    next_obs = sys.state
    reward = zeros(1)
    done = time_over(env.clock)
    info = Dict("time" => t, "state" => x)
    return next_obs, reward, done, info
end
env = BaseEnv()
step!(env, step)
```
"""
function step!(env::BaseEnv, step)
    _step(args...; kwargs...) = step(env, args...; kwargs...)
    env.step = _step
end

function dyn(env::BaseEnv)
    return env.dyn
end

function state!(env::BaseEnv, state)
    for system in _systems(env)
        state!(system, reshape(state[system.flat_index],
                               system.state_size))
    end
end

function dot(env::BaseEnv)
    return vcat([dot(system)[:]
                 for system in _systems(env)]...)  # flatten
end

function dot!(env::BaseEnv, dot)
    for system in _systems(env)
        system.dot = reshape(dot[system.flat_index], system.state_size)
    end
end

function observe_array(env::BaseEnv; y=nothing)
    res = []
    if y == nothing
        for system in _systems(env)
            if typeof(system) == BaseSystem
                push!(res, state(system))
            elseif typeof(system) == BaseEnv
                push!(res, observe_array(system))
            end
        end
    else
        raise_unsupported_error()
    end
    return res
end

function observe_dict(env::BaseEnv; y=nothing)
    res = Dict()
    if y == nothing
        for (name, system) in systems(env)
            if typeof(system) == BaseSystem
                res[name] = state(system)
            elseif typeof(system) == BaseEnv
                res[name] = observe_dict(system)
            end
        end
    else
        for (name, system) in systems(env)
            if typeof(system) == BaseSystem
                res[name] = reshape(y[system.flat_index],
                                    system.state_size)
            elseif typeof(system) == BaseEnv
                res[name] = observe_dict(system,
                                         y=y[system.flat_index])
            end
        end
    end
    return res
end

function observe_flat(env::BaseEnv)
    return vcat([state(system)[:]
                 for system in _systems(env)]...)  # flatten
end

function ProgressMeter.update!(env::BaseEnv; kwargs...)
    t_hist = thist(env.clock)
    ode_hist = env.solver(
                          env.ode_func,
                          observe_flat(env),
                          t_hist,
                          kwargs,
                          env.ode_option...,
                         )
    done = false
    # TODO: low priority; add eager stop; should we?
    tfinal, yfinal = t_hist[end], ode_hist[end]
    # Update the systems' state
    for system in _systems(env)
        state!(system, reshape(yfinal[system.flat_index],
                               system.state_size))
    end

    # TODO: low priority; Log the inner history of states
    if env.logger != nothing
        for (t, y) in zip(t_hist[1:end-1], ode_hist[1:end-1])
            state_dict = observe_dict(env, y=y)
            _info = Dict("time" => t, "state" => state_dict)
            if kwargs != nothing
                if haskey(Dict(kwargs), "time") || haskey(Dict(kwargs),
                                                          "state")
                    error("invalid kwargs")
                else
                    info = merge(_info, Dict(kwargs))
                end
            else
                info = _info
            end
            record(env.logger, info)
        end
        # TODO: low priority; add logger_callback
    end
    time!(env.clock, tfinal)

    # TODO: low priority; add delay; should we?
    return t_hist, ode_hist, done || time_over(env.clock)
end

"Close `env`.
To save data in env's logger, you must `close!` after simulation."
function close!(env::BaseEnv)
    if env.logger != nothing
        close!(env.logger)
    end
end

function render(env::BaseEnv;
                mode="ProgressMeter", desc="", dt=0.01, kwargs...)
    if mode == "ProgressMeter"
        if env.progressbar == nothing || time(env.clock) == 0
            env.progressbar = Progress(env.clock.max_len, dt,
                                       desc, kwargs...)
        end
        next!(env.progressbar)
    end
end

function set_delay(env::BaseEnv, systems::Array, T)
    raise_unsupported_error()
end

function raise_unsupported_error()
    error("Unsupported function yet")
    # TODO
end

function rk4(func, y0, t, kwargs)
    n = length(t)
    y = [zero(y0) for i in 1:n]
    y[1] = y0
    for i in 1:n-1
        h = t[i+1] - t[i]
        k1 = func(y[i], kwargs, t[i])
        k2 = func(y[i] + k1 * h / 2.0, kwargs, t[i] + h / 2.0)
        k3 = func(y[i] + k2 * h / 2.0, kwargs,  t[i] + h / 2.0)
        k4 = func(y[i] + k3 * h, kwargs, t[i] + h)
        y[i+1] = y[i] + (h / 6.0) * (k1 + 2*k2 + 2*k3 + k4)
    end
    return y
end

end


text:
_basename(v::AbstractArray) = basename.(v)
_basename(::Nothing) = nothing
_basename(v) = basename(v)

"""
`filepicker(label="Choose a file..."; multiple=false, accept="*")`

Create a widget to select files.
If `multiple=true` the observable will hold an array containing the paths of all
selected files. Use `accept` to only accept some formats, e.g. `accept=".csv"`
"""
function filepicker(theme::WidgetTheme, lbl="Choose a file..."; attributes=PropDict(),
    label=lbl, className="", multiple=false, value=multiple ? String[] : "",  kwargs...)

    (value isa AbstractObservable) || (value = Observable{Any}(value))
    filename = Observable{Any}(_basename(value[]))

    if multiple
        onFileUpload = js"""function (data, e){
            var files = e.target.files;
            var fileArray = Array.from(files);
            this.filename(fileArray.map(function (el) {return el.name;}));
            return this.path(fileArray.map(function (el) {return el.path;}));
        }
        """
    else
        onFileUpload = js"""function(data, e) {
            var files = e.target.files;
            this.filename(files[0].name);
            return this.path(files[0].path);
        }
        """
    end
    multiple && (attributes=merge(attributes, PropDict(:multiple => true)))
    attributes = merge(attributes, PropDict(:type => "file", :style => "display: none;",
        Symbol("data-bind") => "event: {change: onFileUpload}"))
    className = mergeclasses(getclass(theme, :input, "file"), className)
    template = dom"div[style=display:flex; align-items:center;]"(
        node(:label, className=getclass(theme, :input, "file", "label"))(
            node(:input; className=className, attributes=attributes, kwargs...),
            node(:span,
                node(:span, (node(:i, className = getclass(theme, :input, "file", "icon"))), className=getclass(theme, :input, "file", "span", "icon")),
                node(:span, label, className=getclass(theme, :input, "file", "span", "label")),
                className=getclass(theme, :input, "file", "span"))
        ),
        node(:span, attributes = Dict("data-bind" => " text: filename() == '' ? 'No file chosen' : filename()"),
            className = getclass(theme, :input, "file", "name"))
    )

    observs = ["path" => value, "filename" => filename]
    ui = knockout(template, observs, methods = ["onFileUpload" => onFileUpload])
    slap_design!(ui, theme)
    Widget{:filepicker}(observs, scope = ui, output = ui["path"], layout = node(:div, className = "field interact-widget")∘Widgets.scope)
end

"""
`opendialog(; value = String[], label = "Open", icon = "far fa-folder-open", options...)`

Creates an [Electron openDialog](https://electronjs.org/docs/api/dialog#dialogshowopendialogbrowserwindow-options-callback).
`value` is the list of selected files or folders. `options` (given as keyword arguments) correspond to
`options` of the Electron dialog. This widget will not work in the browser but only in an Electron window.

## Examples

```jldoctest
julia> ui = InteractBase.opendialog(; properties = ["showHiddenFiles", "multiSelections"], filters = [(; name = "Text", extensions = ["txt", "md"])]);

julia> ui[]
0-element Array{String,1}
```
"""
opendialog(::WidgetTheme; value = String[], label = "Open", icon = "far fa-folder-open", kwargs...) =
    dialog(js"showOpenDialog"; value = value, label = label, icon = icon, kwargs...)

"""
`savedialog(; value = String[], label = "Open", icon = "far fa-folder-open", options...)`

Create an [Electron saveDialog](https://electronjs.org/docs/api/dialog#dialogshowsavedialogbrowserwindow-options-callback).
`value` is the list of selected files or folders. `options` (given as keyword arguments) correspond to
`options` of the Electron dialog. This widget will not work in the browser but only in an Electron window.

## Examples

```jldoctest
julia> ui = InteractBase.savedialog(; properties = ["showHiddenFiles"], filters = [(; name = "Text", extensions = ["txt", "md"])]);

julia> ui[]
""
```
"""
savedialog(::WidgetTheme; value = "", label = "Save", icon = "far fa-save", kwargs...) =
    dialog(js"showSaveDialog"; value = value, label = label, icon = icon, kwargs...)

function dialog(dialogtype; value, className = "", label = "dialog", icon = nothing, options...)
    (value isa AbstractObservable) || (value = Observable(value))
    scp = Scope()
    setobservable!(scp, "output", value)
    clicks = Observable(scp, "clicks", 0)
    callback = @js function (val)
        $value[] = val
    end
    onimport(scp, js"""
    function () {
        const { dialog } = require('electron').remote;
        this.dialog = dialog;
    }
    """)
    onjs(clicks, js"""
    function (val) {
        console.log(this.dialog.$dialogtype($options, $callback));
    }
    """)
    theme = gettheme()
    className = mergeclasses(getclass(theme, :button), className)
    content = if icon === nothing
        (label,)
    else
        iconNode = node(:span, node(:i, className = icon), className = "icon")
        (iconNode, node(:span, label))
    end
    btn = node(:button, content...,
        events=Dict("click" => @js event -> ($clicks[] = $clicks[] + 1)),
        className = className)
    scp.dom = btn
    slap_design!(scp, theme)
    Widget{:dialog}([]; output = value, scope = scp, layout = Widgets.scope)
end

_parse(::Type{S}, x) where{S} = parse(S, x)
function _parse(::Type{Dates.Time}, x)
    segments = split(x, ':')
    length(segments) >= 2 && all(!isempty, segments) || return nothing
    h, m = parse.(Int, segments)
    Dates.Time(h, m)
end

function _string(x::Dates.Time)
    h = Dates.hour(x)
    m = Dates.minute(x)
    string(lpad(h, 2, "0"), ":", lpad(m, 2, "0"))
end
_string(x::Dates.Date) = string(x)

"""
`datepicker(value::Union{Dates.Date, Observable, Nothing}=nothing)`

Create a widget to select dates.
"""
function datepicker end

"""
`timepicker(value::Union{Dates.Time, Observable, Nothing}=nothing)`

Create a widget to select times.
"""
function timepicker end

for (func, typ, str, unit) in [(:timepicker, :(Dates.Time), "time", Dates.Second), (:datepicker, :(Dates.Date), "date", Dates.Day) ]
    @eval begin
        function $func(theme::WidgetTheme, val=nothing; value=val, kwargs...)
            (value isa AbstractObservable) || (value = Observable{Union{$typ, Nothing}}(value))
            f = x -> x === nothing ? "" : _string(x)
            g = t -> _parse($typ, t)
            pair = ObservablePair(value, f=f, g=g)
            ui = input(theme, pair.second; typ=$str, kwargs...)
            Widget{$(Expr(:quote, func))}(ui, output = value)
        end

        function $func(theme::WidgetTheme, vals::AbstractRange, val=medianelement(vals); value=val, kwargs...)
            f = x -> x === nothing ? "" : _string(x)
            fs = x -> x === nothing ? "" : split(string(convert($unit, x)), ' ')[1]
            min, max = extrema(vals)
            $func(theme; value=value, min=f(min), max=f(max), step=fs(step(vals)), kwargs...)
        end
    end
end

"""
`colorpicker(value::Union{Color, Observable}=colorant"#000000")`

Create a widget to select colors.
"""
function colorpicker(theme::WidgetTheme, val=colorant"#000000"; value=val, kwargs...)
    (value isa AbstractObservable) || (value = Observable{Color}(value))
    f = t -> "#"*hex(t)
    g = t -> parse(Colorant,t)
    pair = ObservablePair(value, f=f, g=g)
    ui = input(theme, pair.second; typ="color", kwargs...)
    Widget{:colorpicker}(ui, output = value)
end

"""
`spinbox([range,] label=""; value=nothing)`

Create a widget to select numbers with placeholder `label`. An optional `range` first argument
specifies maximum and minimum value accepted as well as the step. Use `step="any"` to allow all
decimal numbers.
"""
function spinbox(theme::WidgetTheme, label=""; value=nothing, placeholder=label, isinteger=nothing, kwargs...)
    isinteger === nothing || @warn "`isinteger` is deprecated"
    if !isa(value, AbstractObservable)
        T = something(isinteger, isa(value, Integer)) ? Int : Float64
        value = Observable{Union{T, Nothing}}(value)
    end
    ui = input(theme, value; isnumeric=true, placeholder=placeholder, typ="number", kwargs...)
    Widget{:spinbox}(ui; output = value)
end

spinbox(theme::WidgetTheme, vals::AbstractRange, args...; value=first(vals), kwargs...) =
    spinbox(theme, args...; value=value, min=minimum(vals), max=maximum(vals), step=step(vals), kwargs...)

"""
`autocomplete(options, label=""; value="")`

Create a textbox input with autocomplete options specified by `options`, with `value`
as initial value and `label` as label.
"""
function autocomplete(::WidgetTheme, options, args...; attributes=PropDict(), kwargs...)
    (options isa AbstractObservable) || (options = Observable{Any}(options))
    option_array = _js_array(options)
    s = gensym()
    attributes = merge(attributes, PropDict(:list => s))
    t = textbox(args...; extra_obs=["options_js" => option_array], attributes=attributes, kwargs...)
    Widgets.scope(t).dom = node(:div,
        Widgets.scope(t).dom,
        node(:datalist, node(:option, attributes=Dict("data-bind"=>"value : key"));
            attributes = Dict("data-bind" => "foreach : options_js", "id" => s))
    )
    w = Widget{:autocomplete}(t)
    w[:options] = options
    w
end

"""
`input(o; typ="text")`

Create an HTML5 input element of type `type` (e.g. "text", "color", "number", "date") with `o`
as initial value.
"""
@noinline function input(theme::WidgetTheme, o; extra_js=js"", extra_obs=[], label=nothing, typ="text", wdgtyp=typ,
    className="", style=Dict(), isnumeric=Knockout.isnumeric(o),
    computed=[], attributes=Dict(), bind="value", bindto="value", valueUpdate="input", changes=0, kwargs...)

    (o isa AbstractObservable) || (o = Observable(o))
    (changes isa AbstractObservable) || (changes = Observable(changes))
    data = Pair{String, Any}["changes" => changes, bindto => o]
    if isnumeric && bind == "value"
        bind = "numericValue"
    end
    append!(data, (string(key) => val for (key, val) in extra_obs))
    countChanges = js_lambda("this.changes(this.changes()+1)")
    attrDict = merge(
        attributes,
        Dict(:type => typ,
            Symbol("data-bind") => "$bind: $bindto, valueUpdate: '$valueUpdate', event: {change: $countChanges}"
        )
    )
    className = mergeclasses(getclass(theme, :input, wdgtyp), className)
    template = node(:input; className=className, attributes=attrDict, style=style, kwargs...)()
    ui = knockout(template, data, extra_js; computed=computed)
    (label != nothing) && (ui.dom = flex_row(wdglabel(theme, label), ui.dom))
    slap_design!(ui, theme)
    Widget{:input}(data, scope = ui, output = ui[bindto], layout = node(:div, className = "field interact-widget")∘Widgets.scope)
end

@noinline function input(theme::WidgetTheme; typ="text", kwargs...)
    if typ in ["checkbox", "radio"]
        o = false
    elseif typ in ["number", "range"]
        o = 0.0
    else
        o = ""
    end
    input(theme, o; typ=typ, kwargs...)
end

"""
`button(content... = "Press me!"; value=0)`

A button. `content` goes inside the button.
Note the button `content` supports a special `clicks` variable, that gets incremented by `1`
with each click e.g.: `button("clicked {{clicks}} times")`.
The `clicks` variable is initialized at `value=0`. Given a button `b`, `b["is-loading"]` defines
whether the button is in a loading state (spinning wheel). Use `b["is-loading"][]=true` or
`b["is-loading"][]=false` respectively to display or take away the spinner.
"""
function button(theme::WidgetTheme, content...; label = "Press me!", value = 0, style = Dict{String, Any}(),
    className = getclass(theme, :button, "primary"), attributes=Dict(), kwargs...)
    isempty(content) && (content = (label,))
    (value isa AbstractObservable) || (value = Observable(value))
    loading = Observable(false)
    className = "delete" in split(className, ' ') ? className : mergeclasses(getclass(theme, :button), className)
    countClicks = js_lambda("this.clicks(this.clicks()+1)")
    attrdict = merge(
        Dict("data-bind"=>"click: $countClicks, css: {'is-loading' : loading}"),
        attributes
    )
    template = node(:button, content...; className=className, attributes=attrdict, style=style, kwargs...)
    button = knockout(template, ["clicks" => value, "loading" => loading])
    slap_design!(button, theme)
    Widget{:button}(["is-loading" => loading], scope = button, output = value,
        layout = node(:div, className = "field interact-widget")∘Widgets.scope)
end

for wdg in [:toggle, :checkbox]
    @eval begin
        $wdg(theme::WidgetTheme, value, lbl::AbstractString=""; label=lbl, kwargs...) =
            $wdg(theme; value=value, label=label, kwargs...)

        $wdg(theme::WidgetTheme, label::AbstractString, val=false; value=val, kwargs...) =
            $wdg(theme; value=value, label=label, kwargs...)

        $wdg(theme::WidgetTheme, value::AbstractString, label::AbstractString; kwargs...) =
            error("value cannot be a string")

        function $wdg(theme::WidgetTheme; bind="checked", valueUpdate="change", value=false, label="", labelclass="", kwargs...)
            s = gensym() |> string
            (label isa Tuple) || (label = (label,))
            widgettype = $(Expr(:quote, wdg))
            wdgtyp = string(widgettype)
            labelclass = mergeclasses(getclass(theme, :input, wdgtyp, "label"), labelclass)
            ui = input(theme, value; bind=bind, typ="checkbox", valueUpdate="change", wdgtyp=wdgtyp, id=s, kwargs...)
            Widgets.scope(ui).dom = node(:div, className = "field interact-widget")(Widgets.scope(ui).dom, dom"label[className=$labelclass, for=$s]"(label...))
            Widget{widgettype}(ui)
        end
    end
end

"""
`checkbox(value::Union{Bool, AbstractObservable}=false; label)`

A checkbox.
e.g. `checkbox(label="be my friend?")`
"""
function checkbox end

"""
`toggle(value::Union{Bool, AbstractObservable}=false; label)`

A toggle switch.
e.g. `toggle(label="be my friend?")`
"""
function toggle end

"""
`textbox(hint=""; value="")`

Create a text input area with an optional placeholder `hint`
e.g. `textbox("enter number:")`. Use `typ=...` to specify the type of text. For example
`typ="email"` or `typ="password"`. Use `multiline=true` to display a `textarea` spanning
several lines.
"""
function textbox(theme::WidgetTheme, hint=""; multiline=false, placeholder=hint, value="", typ="text", kwargs...)
    multiline && return textarea(theme; placeholder=placeholder, value=value, kwargs...)
    Widget{:textbox}(input(theme, value; typ=typ, placeholder=placeholder, kwargs...))
end

"""
`textarea(hint=""; value="")`

Create a textarea with an optional placeholder `hint`
e.g. `textarea("enter number:")`. Use `rows=...` to specify how many rows to display
"""
function textarea(theme::WidgetTheme, hint=""; label=nothing, className="",
    placeholder=hint, value="", attributes=Dict(), style=Dict(), bind="value", valueUpdate = "input", kwargs...)

    (value isa AbstractObservable) || (value = Observable(value))
    attrdict = convert(PropDict, attributes)
    attrdict[:placeholder] = placeholder
    attrdict["data-bind"] = "$bind: value, valueUpdate: '$valueUpdate'"
    className = mergeclasses(getclass(theme, :textarea), className)
    template = node(:textarea; className=className, attributes=attrdict, style=style, kwargs...)
    ui = knockout(template, ["value" => value])
    (label != nothing) && (ui.dom = flex_row(wdglabel(theme, label), ui.dom))
    slap_design!(ui, theme)
    Widget{:textarea}(scope = ui, output = ui["value"], layout = node(:div, className = "field interact-widget")∘Widgets.scope)
end

function wdglabel(theme::WidgetTheme, text; padt=5, padr=10, padb=0, padl=10,
    className="", style = Dict(), kwargs...)

    className = mergeclasses(getclass(theme, :wdglabel), className)
    padding = Dict(:padding=>"$(padt)px $(padr)px $(padb)px $(padl)px")
    node(:label, text; className=className, style = merge(padding, style), kwargs...)
end

function flex_row(a,b,c=dom"div"())
    node(
        :div,
        node(:div, a, className = "interact-flex-row-left"),
        node(:div, b, className = "interact-flex-row-center"),
        node(:div, c, className = "interact-flex-row-right"),
        className = "interact-flex-row interact-widget"
    )
end

flex_row(a) = node(:div, a, className = "interact-flex-row interact-widget")


text:
using ITensors,
      Test

@testset "ITensor broadcast syntax" begin

  i = Index(2,"i")
  A = randomITensor(i,i')
  B = randomITensor(i',i)
  α = 2
  β = 3

  @testset "Copy" begin
    Bc = copy(B)
    Bc .= A
    @test Bc[1,1] == A[1,1]
    @test Bc[2,1] == A[1,2]
    @test Bc[1,2] == A[2,1]
    @test Bc[2,2] == A[2,2]
  end

  @testset "Fill" begin
    Bc = copy(B)
    Bc .= α
    @test Bc[1,1] == α
    @test Bc[2,1] == α
    @test Bc[1,2] == α
    @test Bc[2,2] == α
  end

  @testset "Scaling" begin
    Bc = copy(B)
    Bc .*= α
    @test Bc[1,1] == α * B[1,1]
    @test Bc[2,1] == α * B[2,1]
    @test Bc[1,2] == α * B[1,2]
    @test Bc[2,2] == α * B[2,2]
  end

  @testset "Dividing" begin
    Bc = copy(B)
    Bc ./= α
    @test Bc[1,1] == B[1,1] / α
    @test Bc[2,1] == B[2,1] / α
    @test Bc[1,2] == B[1,2] / α
    @test Bc[2,2] == B[2,2] / α
  end

  @testset "Scalar multiplication (in-place)" begin
    Bc = copy(B)
    Bc .= α .* A
    @test Bc[1,1] == α * A[1,1]
    @test Bc[2,1] == α * A[1,2]
    @test Bc[1,2] == α * A[2,1]
    @test Bc[2,2] == α * A[2,2]
  end

  @testset "Dividing (in-place)" begin
    Bc = copy(B)
    Bc .= A ./ α
    @test Bc[1,1] == A[1,1] / α
    @test Bc[2,1] == A[1,2] / α
    @test Bc[1,2] == A[2,1] / α
    @test Bc[2,2] == A[2,2] / α
    @test_throws ErrorException Bc .= α ./ A
  end

  @testset "Scalar multiplication (out-of-place)" begin
    Bc = α .* A
    @test Bc[1,1] == α * A[1,1]
    @test Bc[2,1] == α * A[2,1]
    @test Bc[1,2] == α * A[1,2]
    @test Bc[2,2] == α * A[2,2]
  end

  @testset "Addition" begin
    Bc = copy(B)
    Bc .= A .+ Bc
    @test Bc[1,1] == A[1,1] + B[1,1]
    @test Bc[2,1] == A[1,2] + B[2,1]
    @test Bc[1,2] == A[2,1] + B[1,2]
    @test Bc[2,2] == A[2,2] + B[2,2]
  end

  @testset "Addition (with α)" begin
    Bc = copy(B)
    Bc .+= A .* α

    @test Bc[1,1] == α * A[1,1] + B[1,1]
    @test Bc[2,1] == α * A[1,2] + B[2,1]
    @test Bc[1,2] == α * A[2,1] + B[1,2]
    @test Bc[2,2] == α * A[2,2] + B[2,2]
  end

  @testset "Addition (with α and β)" begin
    Bc = copy(B)
    Bc .= α .* A .+ β .* Bc

    @test Bc[1,1] == α * A[1,1] + β * B[1,1]
    @test Bc[2,1] == α * A[1,2] + β * B[2,1]
    @test Bc[1,2] == α * A[2,1] + β * B[1,2]
    @test Bc[2,2] == α * A[2,2] + β * B[2,2]
  end

  @testset "Addition errors" begin
    C = randomITensor(i,i')
    @test_throws ErrorException C .= A .+ B
    @test_throws ErrorException C = A .+ B
    @test_throws ErrorException C .= A .* B
  end

  @testset "Contraction" begin
    ii = Index(2; tags="ii")
    jj = Index(2; tags="jj")
    kk = Index(2; tags="kk")

    AA = randomITensor(ii, jj)
    BB = randomITensor(kk, jj)
    CC = randomITensor(kk, ii)

    R = copy(CC)
    R .= AA .* BB
    @test AA*BB ≈ R

    R = copy(CC)
    R .= α .* AA .* BB
    @test α*AA*BB ≈ R

    R = copy(CC)
    R .= AA .* α .* BB
    @test α*AA*BB ≈ R

    R = copy(CC)
    R .= AA .* BB .* α
    @test α*AA*BB ≈ R

    R = copy(CC)
    R .+= α .* AA .* BB
    @test α*AA*BB+CC ≈ R

    R = copy(CC)
    R .= β .* R .+ AA .* BB .* α
    @test α*AA*BB+β*CC ≈ R
  end

  @testset "General functions" begin
    absA = abs.(A)

    @test absA[1,1] == abs(A[1,1])
    @test absA[2,1] == abs(A[2,1])

    Bc = copy(B)
    Bc .= sqrt.(absA)

    @test Bc[1,1] == sqrt(absA[1,1])
    @test Bc[2,1] == sqrt(absA[1,2])

    Bc2 = copy(B)
    Bc2 .+= sqrt.(absA)

    @test Bc2[1,1] == B[1,1]+sqrt(absA[1,1])
    @test Bc2[2,1] == B[2,1]+sqrt(absA[1,2])
  end

  @testset "Some other operations" begin
    i = Index(2)
    A = randomITensor(i)
    B = randomITensor(i)

    absA = abs.(A)

    @test absA[1] == abs(A[1])
    @test absA[2] == abs(A[2])

    Bc = copy(B)
    Bc .= sqrt.(absA)

    @test Bc[1] == sqrt(absA[1])
    @test Bc[2] == sqrt(absA[2])

    Bc2 = copy(B)
    Bc2 .+= sqrt.(absA)

    @test Bc2[1] == B[1]+sqrt(absA[1])
    @test Bc2[2] == B[2]+sqrt(absA[2])

    Bc3 = copy(B)
    Bc3 .= sqrt.(absA) .+ sin.(Bc3)

    @test Bc3[1] == sin(B[1])+sqrt(absA[1])
    @test Bc3[2] == sin(B[2])+sqrt(absA[2])

    sqrtabsA = sqrt.(abs.(A))

    @test sqrtabsA[1] == sqrt(abs(A[1]))
    @test sqrtabsA[2] == sqrt(abs(A[2]))

    sqrtabsA = cos.(sin.(sqrt.(abs.(A))))

    @test sqrtabsA[1] == cos(sin(sqrt(abs(A[1]))))
    @test sqrtabsA[2] == cos(sin(sqrt(abs(A[2]))))

    Ap = A .+ 3

    @test Ap[1] == A[1] + 3
    @test Ap[2] == A[2] + 3

    Apow1 = A .^ 2.0

    @test Apow1[1] == A[1]^2
    @test Apow1[2] == A[2]^2

    Apow2 = A .^ 3

    @test Apow2[1] == A[1]^3
    @test Apow2[2] == A[2]^3
  end

end


text:
module JuSpglib

using Reexport

export spglib_version

include("../deps/build.jl")

# Load library after build
if isfile(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl"))
    include(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl"))
else
    error("Spglib not properly installed. Please run Pkg.build(\"JuSpglib\")")
end

macro lazy_version(cfuncname)
    return :(ccall(($cfuncname, spglib), Cint, ()))
end

"""
This returns version number of spglib.
"""
const spglib_version = VersionNumber(@lazy_version(:spg_get_major_version), @lazy_version(:spg_get_minor_version), @lazy_version(:spg_get_micro_version))

include("DataStructure.jl")
@reexport using .DataStructure

include("CAPIs.jl")
@reexport using .CAPIs

end # module


text:
# univariate

@inline function enclose(f::Function, X::Interval, ba::BranchAndBoundEnclosure;
                         df=Base.Fix1(ForwardDiff.derivative, f))
    return _branch_bound(ba, f, X, df)
end

@inline function enclose(f::Function, X::IntervalBox, ba::BranchAndBoundEnclosure;
                         df=t->ForwardDiff.gradient(w->f(w...), t.v))
    return _branch_bound(ba, f, X, df)
end

function _branch_bound(ba::BranchAndBoundEnclosure, f::Function, X::Interval_or_IntervalBox, df;
                       initial=emptyinterval(first(X)),
                       cnt=1)

    dfX = df(X)
    range_extrema, flag = _monotonicity_check(f, X, dfX)
    flag && return hull(range_extrema, initial)

    fX = f(X...)  # TODO: allow user to choose how to evaluate this (mean value, natural enclosure)
    # if tolerance or maximum number of iteration is met, return current enclosure
    if diam(fX) <= ba.tol || cnt == ba.maxdepth
        return hull(fX, initial)
    end

    X1, X2 = bisect(X)
    y1 = _branch_bound(ba, f, X1, df; initial=initial, cnt=cnt+1)
    return _branch_bound(ba, f, X2, df; initial=y1, cnt=cnt+1)
end

function _monotonicity_check(f::Function, X::Interval, dfX::Interval)
    if inf(dfX) >= 0 || sup(dfX) <= 0  # monotone function, evaluate at extrema
        lo = Interval(X.lo)
        hi = Interval(X.hi)
        return hull(f(lo), f(hi)), true
    end

    return zero(eltype(dfX)), false
end

function _monotonicity_check(f::Function, X::IntervalBox{N}, ∇fX::AbstractVector) where {N}
    low = zeros(eltype(X), N)
    high = zeros(eltype(X), N)

    @inbounds for (i, di) in enumerate(∇fX)
        if di >=0  #  increasing
            high[i] = sup(X[i])
            low[i] = inf(X[i])
        elseif di <= 0  # decreasing
            high[i] = inf(X[i])
            low[i] = sup(X[i])
        else
            return di, false
        end
    end

    return hull(f(low...), f(high...)), true
end


text:
struct Forward end

struct Reverse end

# For some reason, this is necessary for printing. I don't really understand it.
Base.length(::Union{Forward, Reverse}) = 0

Base.reverse(::Forward) = Reverse()
Base.reverse(::Reverse) = Forward()

"""
    GaussMarkovModel

Specifies the following Gauss-Markov model.
```julia
x[0] ∼ x0
x[t] = A[t] * x[t-1] + a[t] + ε[t], ε[t] ∼ N(0, Q)
```
"""
struct GaussMarkovModel{
    Tordering<:Union{Forward, Reverse},
    TAs<:AbstractVector{<:AbstractMatrix{<:Real}},
    Tas<:AbstractVector{<:AbstractVector{<:Real}},
    TQs<:AbstractVector{<:AbstractMatrix{<:Real}},
    Tx0<:Gaussian,
} <: AbstractMvNormal
    ordering::Tordering
    As::TAs
    as::Tas
    Qs::TQs
    x0::Tx0
end

# Helps Zygote out with some type-stability issues. Why this helps is unclear.
function Zygote._pullback(::AContext, ::Type{<:GaussMarkovModel}, ordering, As, as, Qs, x0)
    function GaussMarkovModel_pullback(Δ)
        return (nothing, nothing, Δ.As, Δ.as, Δ.Qs, Δ.x0)
    end
    return GaussMarkovModel(ordering, As, as, Qs, x0), GaussMarkovModel_pullback
end

ordering(model::GaussMarkovModel) = model.ordering

Base.eltype(model::GaussMarkovModel) = eltype(first(model.As))

Base.eachindex(model::GaussMarkovModel{Forward}) = 1:length(model)

Base.eachindex(model::GaussMarkovModel{Reverse}) = reverse(1:length(model))

Base.length(model::GaussMarkovModel) = length(model.As)

function Base.getindex(model::GaussMarkovModel, n::Int)
    return SmallOutputLGC(model.As[n], model.as[n], model.Qs[n])
end

function Base.:(==)(x::GaussMarkovModel, y::GaussMarkovModel)
    return (x.As == y.As) && (x.as == y.as) && (x.Qs == y.Qs) && (x.x0 == y.x0)
end

dim(model::GaussMarkovModel) = length(first(model.as))

storage_type(model::GaussMarkovModel) = storage_type(first(model.As))

function is_of_storage_type(model::GaussMarkovModel, s::StorageType)
    return is_of_storage_type((model.As, model.as, model.Qs, model.x0), s)
end

x0(model::GaussMarkovModel) = Zygote.literal_getfield(model, Val(:x0))

function get_adjoint_storage(x::GaussMarkovModel, n::Int, Δx::NamedTuple{(:A, :a, :Q)})
    return (
        ordering = nothing,
        As = get_adjoint_storage(x.As, n, Δx.A),
        as = get_adjoint_storage(x.as, n, Δx.a),
        Qs = get_adjoint_storage(x.Qs, n, Δx.Q),
        x0 = nothing,
    )
end

function _accum_at(
    Δxs::NamedTuple{(:ordering, :As, :as, :Qs, :x0)},
    n::Int,
    Δx::NamedTuple{(:A, :a, :Q)},
)
    return (
        ordering = nothing,
        As = _accum_at(Δxs.As, n, Δx.A),
        as = _accum_at(Δxs.as, n, Δx.a),
        Qs = _accum_at(Δxs.Qs, n, Δx.Q),
        x0 = nothing,
    )
end


text:
function init(renderer::GRU.Renderer, def::Dict{Symbol, Any})
	if haskey(def, :clear_color)
		GRU.set_clear_color(renderer, (def[:clear_color]...,))
	end
	if haskey(def, :clear_stencil)
		GRU.set_clear_stencil(renderer, def[:clear_stencil])
	end
	if haskey(def, :clear_depth)
		GRU.set_clear_depth(renderer, def[:clear_depth])
	end
end

function init(engine::Engine, ::Type{GRU.Shader}, def::Dict{Symbol, Any})::GRU.Shader
	path = def[:shader]
	assetPath = asset_path(engine, path)
	id = Symbol(assetPath)
	if GRU.has_resource(engine.renderer, id)
		return GRU.get_resource(engine.renderer, id)
	end
	setupFn = get_typed!(def, :setupfn, identity, Function)
	GRU.init(GRU.Shader(), engine.renderer, assetPath, setupFn)
end

function init(engine::Engine, ::Type{GRU.Texture}, def::Dict{Symbol, Any})::GRU.Texture
	path = def[:image]
	assetPath = asset_path(engine, path)
	id = Symbol(assetPath)
	if GRU.has_resource(engine.renderer, id)
		return GRU.get_resource(engine.renderer, id)
	end
	GRU.init(GRU.Texture(), engine.renderer, assetPath)
end

function init(engine::Engine, ::Type{GRU.Mesh}, def::Dict{Symbol, Any})::GRU.Mesh
	local model, id
	if haskey(def, :obj)
		path = def[:obj]
		assetPath = asset_path(engine, path)
		id = Symbol(assetPath)
		if GRU.has_resource(engine.renderer, id)
			return GRU.get_resource(engine.renderer, id)
		end
		model = ObjGeom.load_obj(assetPath)
	else
		shape = def[:shape]
		sides = get(def, :sides, 4)
		smooth = get_typed!(def, :smoothing, ObjGeom.SmoothNone)
		id = Symbol(asset_id(shape, sides, Int(smooth)))
		if GRU.has_resource(engine.renderer, id)
			return GRU.get_resource(engine.renderer, id)
		end
		model =
			if shape == "prism"
				ObjGeom.prism(sides; smooth = smooth)
			elseif shape == "pyramid"
				ObjGeom.pyramid(sides; smooth = smooth)
			elseif shape == "sphere"
				ObjGeom.sphere(sides; smooth = smooth)
			elseif shape == "regularpoly"
				ObjGeom.regularpoly(sides)
			end
	end
	if haskey(def, :project_texcoord)
		direction = get_typed(def, :project_texcoord, Vector{Float32})
		ObjGeom.add_texcoord(model, direction)
	end
	shader = load_def(engine, def[:shader])
	streams, indices = ObjGeom.get_indexed(model)
	GRU.init(GRU.Mesh(), shader, streams, map(UInt16, indices), positionFunc = GRU.position_func(:position), id = id)
end

function uniform_type(t::DataType)
	if GRU.isvector(t) || GRU.ismatrix(t)
		return Vector{eltype(t)}
	end
	if t <: GRU.SamplerType
		return GRU.Texture
	end
	t
end

function init_material(engine::Engine, material::GRU.Material, def::Dict{Symbol, Any})
	if haskey(def, :uniforms)
		uniforms = def[:uniforms]
		get!(def, Symbol("uniforms#org")) do; deepcopy(uniforms) end
		for (u, v) in uniforms
			if isa(v, String) || isa(v, Dict)
				v = load_def(engine, v)
				uniforms[u] = v
			end
			varType = uniform_type(material.shader.uniforms[u].varType)
			GRU.setuniform(material, u, convert(varType, v))
		end
	end
	if haskey(def, :states)
		states = def[:states]
		get!(def, Symbol("states#org")) do; copy(states) end
		for i = 1:length(states)
			if !isa(states[i], GRU.RenderState)
				states[i] = eval(Meta.parse(states[i]))::GRU.RenderState
			end
			GRU.setstate(material, states[i])
		end
	end
	material
end

function init(engine::Engine, ::Type{GRU.Material}, def::Dict{Symbol, Any})::GRU.Material
	shader = load_def(engine, def[:shader])
	init_material(engine, GRU.Material(shader), def)
end

function init(engine::Engine, ::Type{GRU.Model}, def::Dict{Symbol, Any})::GRU.Model
	mesh = load_def(engine, def[:mesh])
	material = load_def(engine, def[:material])
	model = GRU.Model(mesh, material)
	GRU.settransform(model, get_transform(def))
	model
end

function init(engine::Engine, ::Type{GRU.Font}, def::Dict{Symbol, Any})::GRU.Font
	facename = def[:facename]
	sizeXY = get(def, :sizexy, (32, 32))
	faceIndex = get(def, :faceindex, 0)
	chars = get(def, :chars, '\u0000':'\u00ff')
	ftFont = FTFont.loadfont(asset_path(engine, facename), sizeXY = (sizeXY...,), faceIndex = faceIndex, chars = chars)

	shader = load_def(engine, def[:shader])
	positionFunc = get_typed!(def, :positionfn, GRU.position_func(:position))
	textureUniform = get_typed!(def, :texture_uniform, :diffuseTexture)
	maxCharacters = get(def, :maxchars, 2048)
	font = GRU.init(GRU.Font(), ftFont, shader, positionFunc = positionFunc, textureUniform = textureUniform, maxCharacters = maxCharacters)
	init_material(engine, font.model.material, def)
	font
end


text:
"""
    module Laplacians
"""
module Laplacians

export Laplacian

using ...Defaults
import ..Shapes
using ..Simulations
using SparseArrays

"""
	struct Laplacian

	Laplacian(sim,k) --> laplacian

fields:
	`D`: Laplacian operator with PMLs (total kinetic energy)
	`K`: part of the Laplacian due to ordinary coupling between sites
	`Σ`: part of the Laplacian due to the self-energy from the boundary conditions
	`M`: mass matrix (scaled dielectric)
	`Mf`: auxilliary mass matrix (scaled pump)
	`scaling`: scalings associated with PMLs and/or polar coordinates
	`k`: the frequency which was used to construct the PML's.
	`ka`, `kb`: the Bloch vectors used to construct the Laplacian.
"""
struct Laplacian{TK,TKX,TKY}
	D::SparseMatrixCSC{ComplexF64,Int64} # total kinetic operator
	K::SparseMatrixCSC{ComplexF64,Int64} # coupling term
	Σ::SparseMatrixCSC{ComplexF64,Int64} # self-energy
	M::SparseMatrixCSC{ComplexF64,Int64} # mass matrix
	Mf::SparseMatrixCSC{ComplexF64,Int64} # mass matrix
	scaling::SparseMatrixCSC{ComplexF64,Int64}
	k::Float64
	ka::TKX
	kb::TKY

	function Laplacian(sim::Simulation,k::Number,kx::Number=0,ky::Number=0)
		k = real(k)
		sx,sy = boundary_layer_site(sim)
		HX = spdiagm(0=>map(z->(1+1im*z/k),sx))
		HY = spdiagm(0=>map(z->(1+1im*z/k),sy))

		Kx,Σx = laplacian_x(sim,k)
		Ky,Σy = laplacian_y(sim,k)
		K = HY*Kx+HX*Ky
		Σ = HY*Σx+HX*Σy
		D = K + Σ
		scaling = HX*HY
		M = scaling*spdiagm(0=>sim.ε[sim.interior])
		Mf = scaling*spdiagm(0=>sim.F[sim.interior])
		return new{typeof(k),typeof(kx),typeof(ky)}(D, K, Σ, M, Mf, scaling, k, kx, ky)
	end

	function Base.show(io::IO,lap::Laplacian)
		println(io,"Laplacian:")
		println(io,"\toptimal frequency: k=$(lap.k)")
		println(io,"\tBloch phases: ka=$(lap.ka), kb=$(lap.ka)")
		println(io,"\tD (kinetic): $(nnz(lap.D)) nonzeros")
		println(io,"\tK (coupling): $(nnz(lap.K)) nonzeros")
		println(io,"\tΣ (self-energy): $(nnz(lap.Σ)) nonzeros")
		println(io,"\tM, Mf (mass): $(nnz(lap.M)) nonzeros")
	end
end
# r = lds.r0 + ij[i][1]*lds.dr
# r_half₋ = lds.r0 + (ij[i][1]-1/2)*lds.dr
# r_half₊ = lds.r0 + (ij[i][1]+1/2)*lds.dr

function laplacian_x(sim::Simulation,k::Real)
	N = sum(sim.interior)

	link_weight = complex.(sim.link_weight[sim.link_x_bool])
	s = boundary_layer_link_x(sim)
	h = @. 1/(1+1im*s/k)
	for i ∈ eachindex(link_weight)
		if !iszero(s[i])
			link_weight[i] *= h[i]
		end
	end
	start = sim.link_start[sim.link_x_bool]
	stop = sim.link_stop[sim.link_x_bool]
	rows = vcat(start, start)
	cols = vcat( stop, start)
	vals = vcat(link_weight,-link_weight)
	Kx = sparse(rows,cols,vals,N,N,+)

	self_weight = complex.(sim.self_weight[sim.self_x_bool])
	s = boundary_layer_self_x(sim)
	h = @. 1/(1+1im*s/k)
	for i ∈ eachindex(self_weight)
		if !iszero(s[i])
			self_weight[i] *= h[i]
		end
	end
	self = sim.self_index[sim.self_x_bool]
	Σx = sparse(self,self,self_weight,N,N,+)
	return Kx,Σx
end

function laplacian_y(sim::Simulation,k::Real)
	N = sum(sim.interior)

	link_weight = complex.(sim.link_weight[sim.link_y_bool])
	s = boundary_layer_link_y(sim)
	h = @. 1/(1+1im*s/k)
	for i ∈ eachindex(link_weight)
		if !iszero(s[i])
			link_weight[i] *= h[i]
		end
	end
	start = sim.link_start[sim.link_y_bool]
	stop = sim.link_stop[sim.link_y_bool]
	rows = vcat(start, start)
	cols = vcat( stop, start)
	vals = vcat(link_weight,-link_weight)
	Ky = sparse(rows,cols,vals,N,N,+)

	self_weight = complex.(sim.self_weight[sim.self_y_bool])
	s = boundary_layer_self_y(sim)
	h = @. 1/(1+1im*s/k)
	for i ∈ eachindex(self_weight)
		if !iszero(s[i])
			self_weight[i] *= h[i]
		end
	end
	self = sim.self_index[sim.self_y_bool]
	Σy = sparse(self,self,self_weight,N,N,+)
	return Ky,Σy
end

function boundary_layer_site(sim::Simulation)
	domains = sim.domains
	x = sim.x[sim.interior]
	y = sim.y[sim.interior]
	domain_index = sim.domain_index[sim.interior]
	return boundary_layer(domains,x,y,domain_index)
end

boundary_layer_link_x(sim::Simulation) = boundary_layer_link(sim,sim.link_x_bool)[1]
boundary_layer_link_y(sim::Simulation) = boundary_layer_link(sim,sim.link_y_bool)[2]
function boundary_layer_link(sim::Simulation,link_half_bool)
	domains = sim.domains
	x = sim.link_half_x[link_half_bool]
	y = sim.link_half_y[link_half_bool]
	domain_index = sim.domain_index[sim.link_start[link_half_bool]]
	return boundary_layer(domains,x,y,domain_index)
end

boundary_layer_self_x(sim::Simulation) = boundary_layer_self(sim,sim.self_x_bool)[1]
boundary_layer_self_y(sim::Simulation) = boundary_layer_self(sim,sim.self_y_bool)[2]
function boundary_layer_self(sim::Simulation,self_half_bool)
	domains = sim.domains
	x = sim.self_half_x[self_half_bool]
	y = sim.self_half_y[self_half_bool]
	domain_index = sim.domain_index[sim.self_index[self_half_bool]]
	return boundary_layer(domains,x,y,domain_index)
end

function boundary_layer(domains,x,y,domain_index)
	Σx = Array{ComplexF64}(undef,length(x))
	Σy = Array{ComplexF64}(undef,length(y))
	boundary_layer!(Σx,Σy,domains,x,y,domain_index)
	return Σx, Σy
end
function boundary_layer!(Σx,Σy,domains,x,y,domain_index)
	for i ∈ eachindex(x)
		Σx[i],Σy[i] = boundary_layer(domains[domain_index[i]],x[i],y[i])
	end
	return nothing
end
function boundary_layer(domain,x,y)
	x,y = Shapes.unrotate(x,y,domain.boundary.shape)
	N = Shapes.get_number_of_sides(domain.boundary.shape)
	if N==4
	    bl1 = domain.boundary.bls[1]
		bl2 = domain.boundary.bls[2]
		bl3 = domain.boundary.bls[3]
		bl4 = domain.boundary.bls[4]
		σx = bl1(x,y)+bl2(x,y)
		σy = bl3(x,y)+bl4(x,y)
	elseif N==2
		bl1 = domain.boundary.bls[1]
		bl2 = domain.boundary.bls[2]
		σx = bl1(x,y)+bl2(x,y)
        σy = complex(0.)
	elseif N==1
		bl1 = domain.boundary.bls[1]
		σx = bl1(x,y)
		σy = complex(0.)
	else
		σx, σy =complex(0.), complex(0.)
	end
    return σx,σy
end

end #module


text:
import Pkg;
import PkgMirrors;

Pkg.Registry.add(Pkg.RegistrySpec(url="https://mirrors.zju.edu.cn/julia/registries/General.git"))

## 杭州大学镜像
Pkg.add(Pkg.PackageSpec(url="https://mirrors.zju.edu.cn/julia/PkgMirrors.jl.git", rev="v1.2.0"))
PkgMirrors.setmirror("ZJU")

## 中国科学技术大学开源软件镜像
# Pkg.add(Pkg.PackageSpec(url="https://mirrors.ustc.edu.cn/julia/PkgMirrors.jl.git", rev="v1.2.0"))
# PkgMirrors.setmirror("USTC")

# 列出所有可用的镜像。
PkgMirrors.availables()

# 停用镜像
PkgMirrors.deactivate()

# 清空镜像缓存
PkgMirrors.clear()

max_stars_repo_name: huangzhen/aquila

max_issues_repo_name: huangzhen/aquila

max_forks_repo_name: huangzhen/aquila


text:
using XConv, LinearAlgebra, PyPlot, Flux, BenchmarkTools, Statistics
using MLDatasets

BLAS.set_num_threads(Sys.CPU_THREADS)

function getbatch_CIFAR10(bsize, nchin, nchout; gaussian=false)
    X, _ = CIFAR10.traindata();
    gaussian && (X0 = randn(Float32, 32, 32, 3*nchin, bsize))
    !gaussian && (X0 = reshape(X[:, :, :, rand(1:50000, bsize*nchin)], 32, 32, 3*nchin, bsize))
    Y0 = reshape(X[:, :, :, rand(1:50000, bsize*nchout)], 32, 32, 3*nchout, bsize)
    return Float32.(X0), Float32.(Y0)
end

function getbatch_MNIST(bsize, nchin, nchout;gaussian=false)
    X, _ = MNIST.traindata();
    gaussian && (X0 = randn(Float32, 28, 28, nchin, bsize))
    !gaussian && (X0 = reshape(X[:, :, rand(1:60000, 3*bsize*nchin)], 28, 28, 3*nchin, bsize))
    Y0 = reshape(X[:, :, rand(1:60000, 3*bsize*nchout)], 28, 28, 3*nchout, bsize)
    return Float32.(X0), Float32.(Y0)
end

function qi(grads, level=.95)
    qis = [quantile(abs.(grads[i, :].- mean(grads[i, :])), level) for i=1:size(grads, 1)]
    return mean(grads; dims=2)[:, 1] .- vec(qis), mean(grads; dims=2)[:, 1] .+ vec(qis)
end

n_in = 1
n_out = 1
nx = 32
ny = 32

diffgrad = Array{Any}(undef, 4, 6)

close("all")
for ps=1:4
    fig, axs = subplots(3, 2, figsize=(10, 5), sharex=true, sharey=true)
    fig.suptitle("Gradient comparisons \n N=$(nx)x$(ny), nchi=$(3*n_in), ncho=$(3*n_out), prob_size=$(2^(ps+1))")

    for i=1:6
        batchsize= 2^(i-1)
        stride = 1
        n_bench = 5
        nw   = 3;

        # X, Y0 = getbatch_MNIST(batchsize, n_in, n_out)
        X, Y0 = getbatch_CIFAR10(batchsize, n_in, n_out; gaussian=true)
        # Flux network
        C = Conv((nw, nw), 3*n_in=>3*n_out, identity;pad=1, stride=stride)
        w = C.weight

        XConv.initXConv(0, "TrueGrad")
        @time g1 = gradient(w->.5f0*norm(conv(X, w;pad=1) - Y0), w)[1]

        XConv.initXConv(2^(ps+1), "EVGrad")
        @time g2 = gradient(w->.5f0*norm(conv(X, w;pad=1) - Y0), w)[1]
        nn = 100
        all_g = hcat([vec(gradient(w->.5f0*norm(conv(X, w;pad=1)- Y0), w)[1]) for i=1:nn]...)
        gm2 = mean(all_g, dims=2)[:, 1]
        st = std(all_g; dims=2)[:, 1]
        qtl, qtr = qi(all_g, .95)

        t = 1:length(vec(g2))
        axsl = axs[i]
        axsl.set_title("batchsize=$(batchsize)")
        axsl.plot(t, vec(all_g[:, 1]), label="sample")
        axsl.plot(t, vec(gm2), label="mean")
        axsl.fill_between(t, gm2 - st, gm2 + st, facecolor="orange", alpha=0.75, label="std($nn samples)")
        axsl.fill_between(t, qtl, qtr, facecolor="cyan", alpha=0.5, label="95% of values")
        axsl.plot(t, vec(g1), label="true")
        diffgrad[ps, i] = 1 .- vec(gm2)./vec(g1)
    end
    lines, labels = fig.axes[end].get_legend_handles_labels()
    fig.legend(lines, labels, loc = "upper left")
    tight_layout()
    # savefig("./benchmark/var_grad$(ps)_CIFAR10-randX.png", bbox_inches="tight")
end

fig, axs = subplots(3, 2, figsize=(10, 5), sharex=true, sharey=true)
title("Errors 1 - mean/true")
for i=1:6
    for ps=1:4
        axs[i].plot(diffgrad[ps, i], label="probe_size=$(2^(ps+1))")
    end
    axs[i].set_title("batch_size=$(2^(i-1))")
    axs[i].set_ylim([-1, 1])
end
lines, labels = fig.axes[end].get_legend_handles_labels()
fig.legend(lines, labels, loc = "upper left")
tight_layout()
# savefig("./benchmark/err_CIFAR10-randX.png", bbox_inches="tight")

max_stars_repo_name: sheevy/XConv

max_issues_repo_name: sheevy/XConv

max_forks_repo_name: sheevy/XConv


text:
using DirectDetections, Distributions, Plots, MCMCChains

@named B = DirectDetections.Planet(
    Deterministic(
        e = (sys, pl) -> 10^pl.loge,
        a = (sys, pl) -> 10^pl.loga,
        mass = (sys, pl) -> 10^pl.logm,
    ),
    Priors(
        # a = Uniform(0, 25),
        # e = TruncatedNormal(0.0, 0.2, 0, 1.0),
        # Note: priors with sharp edges (e.g. Uniform priors) are challenging for HMC samplers.
        # An alternative could be wide Gaussians, for example.
        τ = Normal(0.5, 1),
        ω = Normal(pi, 2pi),
        i = Normal(pi, 2pi),
        Ω = Normal(pi, 2pi),
        # mass = Uniform(0, 1),

        # Reparameterize a few properties for better sampling
        loge = TruncatedNormal(-2, 1.5, -Inf, 0),
        loga = Uniform(-1, 1.5),
        logm = Uniform(-2, 0),

    ),
    Astrometry(
        (epoch=mjd("2016-12-15"), ra=0.133*1e3, dec=-0.174*1e3, σ_ra=0.007*1e3, σ_dec=0.007*1e3),
        (epoch=mjd("2017-03-12"), ra=0.126*1e3, dec=-0.176*1e3, σ_ra=0.004*1e3, σ_dec=0.004*1e3),
        (epoch=mjd("2017-03-13"), ra=0.127*1e3, dec=-0.172*1e3, σ_ra=0.004*1e3, σ_dec=0.004*1e3),
        (epoch=mjd("2018-02-08"), ra=0.083*1e3, dec=-0.133*1e3, σ_ra=0.010*1e3, σ_dec=0.010*1e3),
        (epoch=mjd("2018-11-28"), ra=0.058*1e3, dec=-0.122*1e3, σ_ra=0.010*1e3, σ_dec=0.020*1e3),
        (epoch=mjd("2018-12-15"), ra=0.056*1e3, dec=-0.104*1e3, σ_ra=0.008*1e3, σ_dec=0.008*1e3),
    )
)

@named HD82134 = System(
    # Deterministic(
    #     j = Returns(1),
    #     k = Returns(1)
    # ),
    Priors(
        μ = Normal(1.61, 0.05),
        plx = gaia_plx(gaia_id=756291174721509376),
    ),
    ProperMotionAnomHGCA(gaia_id=756291174721509376),
    B,
)

##
plot(B.astrometry)

## Sampling from the posterior
chain = DirectDetections.hmc(
    HD82134,  0.90,
    # MCMCThreads(),
    # num_chains=4,
    adaptation =    500,
    iterations =  2_000,
    tree_depth =     10,
)

##
gelmandiag(chain)

## Plot the input data and samples from the posterior
plotmodel(chain, clims=(5, 20))

## Create a corner plot / pair plot.
# We can access any property from the chain specified in Priors or in Deterministic.
using PairPlots
table = (;
    a=         chain["B[a]"],
    μ=         chain["μ"],
    m=         chain["B[mass]"],
    e=         chain["B[e]"],
    i=rad2deg.(rem2pi.(chain["B[i]"],RoundDown)),
    Ω=rad2deg.(rem2pi.(chain["B[Ω]"],RoundDown)),
    ω=rad2deg.(rem2pi.(chain["B[ω]"],RoundDown)),
    τ=     rem.(chain["B[τ]"],1,RoundDown),
)
labels=[
    "a",
    "\\mu",
    "m",
    "e",
    "i",
    "\\Omega",
    "\\omega",
    "\\tau",
]
units = [
    "(au)",
    "(_\\odot)",
    "(_\\odot)",
    "",
    "(\\degree)",
    "(\\degree)",
    "(\\degree)",
    "",
]
corner(table, labels, units);
nothing


text:
@testset "002" begin
    @test tdpe.e002(0) == 0
    @test tdpe.e002(1) == 0
    @test tdpe.e002(2) == 0
    @test tdpe.e002(3) == 2
    @test tdpe.e002(35) == 44
end

max_stars_repo_name: mithodin/ttpe

max_issues_repo_name: mithodin/ttpe

max_forks_repo_name: mithodin/ttpe


text:
__precompile__()

"""

A module for searching for eigenvalues of symmetric matrices.

See `Problem`, `Options`, and `solve` for more info.

"""
baremodule A3

using Base

# Load the internal logic
Base.include(A3, "Internal/Internal.jl")

# Export the second-level components
using .Internal

# Export the first-level components
export Problem, Options, solve

end