text:
using JuMP
using Memento
using MathOptInterface

"""
$(SIGNATURES)

Checks whether a JuMPLinearType is constant (and thus has no model associated)
with it. This can only be true if it is an affine expression with no stored
variables.
"""
function is_constant(x::JuMP.AffExpr)
    # TODO (vtjeng): Determine whether there is a built-in function for this as of JuMP>=0.19
    all(values(x.terms) .== 0)
end

function is_constant(x::JuMP.VariableRef)
    false
end

function get_tightening_algorithm(
    x::JuMPLinearType,
    nta::Union{TighteningAlgorithm,Nothing},
)::TighteningAlgorithm
    if is_constant(x)
        return interval_arithmetic
    elseif !(nta === nothing)
        return nta
    else
        # x is not constant, and thus x must have an associated model
        model = owner_model(x)
        return !haskey(model.ext, :MIPVerify) ? DEFAULT_TIGHTENING_ALGORITHM :
               model.ext[:MIPVerify].tightening_algorithm
    end
end

@enum BoundType lower_bound_type = -1 upper_bound_type = 1
#! format: off
bound_f = Dict(
    lower_bound_type => lower_bound,
    upper_bound_type => upper_bound
)
bound_obj = Dict(
    lower_bound_type => MathOptInterface.MIN_SENSE,
    upper_bound_type => MathOptInterface.MAX_SENSE
)
bound_delta_f = Dict(
    lower_bound_type => (b, b_0) -> b - b_0,
    upper_bound_type => (b, b_0) -> b_0 - b
)
bound_operator = Dict(
    lower_bound_type => >=,
    upper_bound_type => <=
)
#! format: on

"""
$(SIGNATURES)

Context manager for running `f` on `model`. If `should_relax_integrality` is true, the
integrality constraints are relaxed before `f` is run and re-imposed after.
"""
function relax_integrality_context(f, model::Model, should_relax_integrality::Bool)
    if should_relax_integrality
        undo_relax = relax_integrality(model)
    end
    r = f(model)
    if should_relax_integrality
        undo_relax()
    end
    return r
end

"""
$(SIGNATURES)

Optimizes the value of `objective` based on `bound_type`, with `b_0`, computed via interval
arithmetic, as a backup.

- If an optimal solution is reached, we return the objective value. We also verify that the
  objective found is better than the bound `b_0` provided; if this is not the case, we throw an
  error.
- If we reach the user-defined time limit, we compute the best objective bound found. We compare
  this to `b_0` and return the better result.
- For all other solve statuses, we warn the user and report `b_0`.
"""
function tight_bound_helper(m::Model, bound_type::BoundType, objective::JuMPLinearType, b_0::Number)
    @objective(m, bound_obj[bound_type], objective)
    optimize!(m)
    status = JuMP.termination_status(m)
    if status == MathOptInterface.OPTIMAL
        b = JuMP.objective_value(m)
        db = bound_delta_f[bound_type](b, b_0)
        if db < -1e-8
            Memento.warn(MIPVerify.LOGGER, "Δb = $(db)")
            Memento.error(
                MIPVerify.LOGGER,
                "Δb = $(db). Tightening via interval arithmetic should not give a better result than an optimal optimization.",
            )
        end
        return b
    elseif status == MathOptInterface.TIME_LIMIT
        return b_0
    else
        Memento.warn(
            MIPVerify.LOGGER,
            "Unexpected solve status $(status); using interval_arithmetic to obtain bound.",
        )
        return b_0
    end
end

"""
Calculates a tight bound of type `bound_type` on the variable `x` using the specified
tightening algorithm `nta`.

If an upper bound is proven to be below cutoff, or a lower bound is proven to above cutoff,
the algorithm returns early with whatever value was found.
"""
function tight_bound(
    x::JuMPLinearType,
    nta::Union{TighteningAlgorithm,Nothing},
    bound_type::BoundType,
    cutoff::Real,
)
    tightening_algorithm = get_tightening_algorithm(x, nta)
    b_0 = bound_f[bound_type](x)
    if tightening_algorithm == interval_arithmetic ||
       is_constant(x) ||
       bound_operator[bound_type](b_0, cutoff)
       return b_0
    end
    should_relax_integrality = (tightening_algorithm == lp)
    # x is not constant, and thus x must have an associated model
    bound_value = return relax_integrality_context(owner_model(x), should_relax_integrality) do m
        tight_bound_helper(m, bound_type, x, b_0)
    end

    return bound_value
end

function tight_upperbound(
    x::JuMPLinearType;
    nta::Union{TighteningAlgorithm,Nothing} = nothing,
    cutoff::Real = -Inf,
)
    tight_bound(x, nta, upper_bound_type, cutoff)
end

function tight_lowerbound(
    x::JuMPLinearType;
    nta::Union{TighteningAlgorithm,Nothing} = nothing,
    cutoff::Real = Inf,
)
    tight_bound(x, nta, lower_bound_type, cutoff)
end

function log_gap(m::JuMP.Model)
    gap = abs(1 - JuMP.objective_bound(m) / JuMP.objective_value(m))
    Memento.info(
        MIPVerify.LOGGER,
        "Hit user limit during solve to determine bounds. Multiplicative gap was $gap.",
    )
end

function relu(x::T)::T where {T<:Real}
    return max(zero(T), x)
end

function relu(x::AbstractArray{T}) where {T<:Real}
    return relu.(x)
end

function relu(x::T, l::Real, u::Real)::JuMP.AffExpr where {T<:JuMPLinearType}
    if u < l
        # TODO (vtjeng): This check is in place in case of numerical error in the calculation of bounds.
        # See sample number 4872 (1-indexed) when verified on the lp0.4 network.
        Memento.warn(
            MIPVerify.LOGGER,
            "Inconsistent upper and lower bounds: u-l = $(u - l) is negative. Attempting to use interval arithmetic bounds instead ...",
        )
        u = upper_bound(x)
        l = lower_bound(x)
    end

    if u <= 0
        # rectified value is always 0
        return zero(T)
    elseif u == l
        return one(T) * l
    elseif u < l
        error(
            MIPVerify.LOGGER,
            "Inconsistent upper and lower bounds even after using only interval arithmetic: u-l = $(u - l) is negative",
        )
    elseif l >= 0
        # rectified value is always x
        return x
    elseif solver_algorithm == exact_mip
        # since we know that u!=l, x is not constant, and thus x must have an associated model
        model = owner_model(x)
        x_rect = @variable(model)
        a = @variable(model, binary = true)

        # refined big-M formulation that takes advantage of the knowledge
        # that lower and upper bounds  are different.
        @constraint(model, x_rect <= x + (-l) * (1 - a))
        @constraint(model, x_rect >= x)
        @constraint(model, x_rect <= u * a)
        @constraint(model, x_rect >= 0)

        # Manually set the bounds for x_rect so they can be used by downstream operations.
        set_lower_bound(x_rect, 0)
        set_upper_bound(x_rect, u)
        return x_rect
    elseif solver_algorithm == overapprox_deeppoly
        model = owner_model(x)
        x_rect = @variable(model)

        # add overaprroximated constraints (deep poly fig. 4(a))
        @constraint(model, x_rect >= x)
        @constraint(model, x_rect >= 0)
        @constraint(model, (x-l)*u/(u-l) >= x_rect)

        # Manually set the bounds for x_rect so they can be used by downstream operations.
        set_lower_bound(x_rect, 0)
        set_upper_bound(x_rect, u)
        return x_rect
    else # solver_algorithm == overapprox_mqp
        model = owner_model(x)
        x_rect = @variable(model)

        l_b = @variable(model, lower_bound = l, upper_bound = u)
        u_b =  @variable(model, lower_bound = l, upper_bound = u)
        @constraint(model, u_b >= x)
        @constraint(model, x >= l_b)

        # add overaprroximated constraints (deep poly fig. 4(c))
        @constraint(model, x_rect >= x)
        @constraint(model, u_b >= x_rect)
        @constraint(model, (x-l_b)*u_b >= x_rect*(u_b-l_b))

        # Manually set the bounds for x_rect so they can be used by downstream operations.
        set_lower_bound(x_rect, 0)
        set_upper_bound(x_rect, u)
        return x_rect
    end
end

@enum ReLUType split = 0 zero_output = -1 linear_in_input = 1 constant_output = 2

function get_relu_type(l::Real, u::Real)::ReLUType
    if u <= 0
        return zero_output
    elseif u == l
        return constant_output
    elseif l >= 0
        return linear_in_input
    else
        return split
    end
end

struct ReLUInfo
    lowerbounds::Array{Real}
    upperbounds::Array{Real}
end

function Base.show(io::IO, s::ReLUInfo)
    relutypes = get_relu_type.(s.lowerbounds, s.upperbounds)
    print(io, "  Behavior of ReLUs - ")
    for t in instances(ReLUType)
        n = count(x -> x == t, relutypes)
        print(io, "$t: $n")
        if t != last(instances(ReLUType))
            print(io, ", ")
        end
    end
end

"""
Calculates the lower_bound only if `u` is positive; otherwise, returns `u` (since we expect)
the ReLU to be fixed to zero anyway.
"""
function lazy_tight_lowerbound(
    x::JuMPLinearType,
    u::Real;
    nta::Union{TighteningAlgorithm,Nothing} = nothing,
    cutoff = 0,
)::Real
    (u <= cutoff) ? u : tight_lowerbound(x; nta = nta, cutoff = cutoff)
end

function relu(x::JuMPLinearType)::JuMP.AffExpr
    u = tight_upperbound(x, cutoff = 0)
    l = lazy_tight_lowerbound(x, u, cutoff = 0)
    relu(x, l, u)
end

"""
$(SIGNATURES)
Expresses a rectified-linearity constraint: output is constrained to be equal to
`max(x, 0)`.
"""
function relu(
    x::AbstractArray{T};
    nta::Union{TighteningAlgorithm,Nothing} = nothing,
)::Array{JuMP.AffExpr} where {T<:JuMPLinearType}
    show_progress_bar::Bool =
        MIPVerify.LOGGER.levels[MIPVerify.LOGGER.level] > MIPVerify.LOGGER.levels["debug"]
    if !show_progress_bar
        u = tight_upperbound.(x, nta = nta, cutoff = 0)
        l = lazy_tight_lowerbound.(x, u, nta = nta, cutoff = 0)
        return relu.(x, l, u)
    else
        p1 = Progress(length(x), desc = "  Calculating upper bounds: ")
        u = map(x_i -> (next!(p1); tight_upperbound(x_i, nta = nta, cutoff = 0)), x)
        p2 = Progress(length(x), desc = "  Calculating lower bounds: ")
        l = map(v -> (next!(p2); lazy_tight_lowerbound(v..., nta = nta, cutoff = 0)), zip(x, u))

        reluinfo = ReLUInfo(l, u)
        Memento.info(MIPVerify.LOGGER, "$reluinfo")

        p3 = Progress(length(x), desc = "  Imposing relu constraint: ")
        return x_r = map(v -> (next!(p3); relu(v...)), zip(x, l, u))
    end
end

function masked_relu(x::T, m::Real)::T where {T<:Real}
    if m < 0
        zero(T)
    elseif m > 0
        x
    else
        relu(x)
    end
end

function masked_relu(x::AbstractArray{<:Real}, m::AbstractArray{<:Real})
    masked_relu.(x, m)
end

function masked_relu(x::T, m::Real)::JuMP.AffExpr where {T<:JuMPLinearType}
    if m < 0
        zero(T)
    elseif m > 0
        x
    else
        relu(x)
    end
end

"""
$(SIGNATURES)
Expresses a masked rectified-linearity constraint, with three possibilities depending on
the value of the mask. Output is constrained to be:
```
1) max(x, 0) if m=0,
2) 0 if m<0
3) x if m>0
```
"""
function masked_relu(
    x::AbstractArray{<:JuMPLinearType},
    m::AbstractArray{<:Real};
    nta::Union{TighteningAlgorithm,Nothing} = nothing,
)::Array{JuMP.AffExpr}
    @assert(size(x) == size(m))
    s = size(m)
    # We add the constraints corresponding to the active ReLUs to the model
    zero_idx = Iterators.filter(i -> m[i] == 0, CartesianIndices(s)) |> collect
    d = Dict(zip(zero_idx, relu(x[zero_idx], nta = nta)))

    # We determine the output of the masked relu, which is either:
    #  1) the output of the relu that we have previously determined when adding the
    #     constraints to the model.
    #  2, 3) the result of applying the (elementwise) masked_relu function.
    return map(i -> m[i] == 0 ? d[i] : masked_relu(x[i], m[i]), CartesianIndices(s))
end

function maximum(xs::AbstractArray{T})::T where {T<:Real}
    return Base.maximum(xs)
end

function maximum_of_constants(xs::AbstractArray{T}) where {T<:JuMPLinearType}
    @assert all(is_constant.(xs))
    max_val = map(x -> x.constant, xs) |> maximum
    return one(JuMP.VariableRef) * max_val
end

"""
$(SIGNATURES)
Expresses a maximization constraint: output is constrained to be equal to `max(xs)`.
"""
function maximum(xs::AbstractArray{T})::JuMP.AffExpr where {T<:JuMPLinearType}
    if length(xs) == 1
        return xs[1]
    end

    if all(is_constant.(xs))
        return maximum_of_constants(xs)
    end
    # at least one of xs is not constant.
    model = owner_model(xs)

    # TODO (vtjeng): [PERF] skip calculating lower_bound for index if upper_bound is lower than
    # largest current lower_bound.
    p1 = Progress(length(xs), desc = "  Calculating upper bounds: ")
    us = map(x_i -> (next!(p1); tight_upperbound(x_i)), xs)
    p2 = Progress(length(xs), desc = "  Calculating lower bounds: ")
    ls = map(x_i -> (next!(p2); tight_lowerbound(x_i)), xs)

    l = Base.maximum(ls)
    u = Base.maximum(us)

    if l == u
        return one(T) * l
        Memento.info(MIPVerify.LOGGER, "Output of maximum is constant.")
    end
    # at least one index will satisfy this property because of check above.
    filtered_indexes = us .> l

    # TODO (vtjeng): Smarter log output if maximum function is being used more than once (for example, in a max-pooling layer).
    Memento.info(
        MIPVerify.LOGGER,
        "Number of inputs to maximum function possibly taking maximum value: $(filtered_indexes |> sum)",
    )

    return maximum(xs[filtered_indexes], ls[filtered_indexes], us[filtered_indexes])
end

function maximum(
    xs::AbstractArray{T,1},
    ls::AbstractArray{<:Real,1},
    us::AbstractArray{<:Real,1},
)::JuMP.AffExpr where {T<:JuMPLinearType}

    @assert length(xs) > 0
    @assert length(xs) == length(ls)
    @assert length(xs) == length(us)

    if all(is_constant.(xs))
        return maximum_of_constants(xs)
    end
    # at least one of xs is not constant.
    model = owner_model(xs)
    if length(xs) == 1
        return first(xs)
    else
        l = Base.maximum(ls)
        u = Base.maximum(us)
        x_max = @variable(model, lower_bound = l, upper_bound = u)
        a = @variable(model, [1:length(xs)], binary = true)
        @constraint(model, sum(a) == 1)
        for (i, x) in enumerate(xs)
            umaxi = Base.maximum(us[1:end.!=i])
            @constraint(model, x_max <= x + (1 - a[i]) * (umaxi - ls[i]))
            @constraint(model, x_max >= x)
        end
        return x_max
    end
end

"""
$(SIGNATURES)
Expresses a one-sided maximization constraint: output is constrained to be at least
`max(xs)`.

Only use when you are minimizing over the output in the objective.

NB: If all of xs are constant, we simply return the largest of them.
"""
function maximum_ge(xs::AbstractArray{T})::JuMPLinearType where {T<:JuMPLinearType}
    @assert length(xs) > 0
    if all(is_constant.(xs))
        return maximum_of_constants(xs)
    end
    if length(xs) == 1
        return first(xs)
    end
    # at least one of xs is not constant.
    model = owner_model(xs)
    x_max = @variable(model)
    @constraint(model, x_max .>= xs)
    return x_max
end

"""
$(SIGNATURES)
Expresses a one-sided absolute-value constraint: output is constrained to be at least as
large as `|x|`.

Only use when you are minimizing over the output in the objective.
"""
function abs_ge(x::JuMPLinearType)::JuMP.AffExpr
    if is_constant(x)
        return one(JuMP.VariableRef) * abs(x.constant)
    end
    model = owner_model(x)
    u = upper_bound(x)
    l = lower_bound(x)
    if u <= 0
        return -x
    elseif l >= 0
        return x
    else
        x_abs = @variable(model)
        @constraint(model, x_abs >= x)
        @constraint(model, x_abs >= -x)
        set_lower_bound(x_abs, 0)
        set_upper_bound(x_abs, max(-l, u))
        return x_abs
    end
end

function get_target_indexes(
    target_index::Integer,
    array_length::Integer;
    invert_target_selection::Bool = false,
)

    get_target_indexes(
        [target_index],
        array_length,
        invert_target_selection = invert_target_selection,
    )

end

function get_target_indexes(
    target_indexes::Array{<:Integer,1},
    array_length::Integer;
    invert_target_selection::Bool = false,
)::AbstractArray{<:Integer,1}

    @assert length(target_indexes) >= 1
    @assert all(target_indexes .>= 1) && all(target_indexes .<= array_length)

    invert_target_selection ? filter((x) -> x ∉ target_indexes, 1:array_length) : target_indexes
end

function get_vars_for_max_index(
    xs::Array{<:JuMPLinearType,1},
    target_indexes::Array{<:Integer,1},
)::Tuple{JuMPLinearType,Array{<:JuMPLinearType,1}}

    @assert length(xs) >= 1

    target_vars = xs[Bool[i ∈ target_indexes for i in 1:length(xs)]]
    nontarget_vars = xs[Bool[i ∉ target_indexes for i in 1:length(xs)]]

    maximum_target_var = length(target_vars) == 1 ? target_vars[1] : MIPVerify.maximum(target_vars)

    return (maximum_target_var, nontarget_vars)
end

"""
$(SIGNATURES)

Imposes constraints ensuring that one of the elements at the target_indexes is the
largest element of the array x. More specifically, we require `x[j] - x[i] ≥ margin` for
some `j ∈ target_indexes` and for all `i ∉ target_indexes`.
"""
function set_max_indexes(
    model::Model,
    xs::Array{<:JuMPLinearType,1},
    target_indexes::Array{<:Integer,1};
    margin::Real = 0,
)::Nothing

    (maximum_target_var, nontarget_vars) = get_vars_for_max_index(xs, target_indexes)

    @constraint(model, nontarget_vars .<= maximum_target_var - margin)
    return nothing
end


text:
function testTrees3(n,nruns)

st0 = Array(Float64,0)
st2 = Array(Float64,0)
st5 = Array(Float64,0)
stold = Array(Float64,0)
stp = Array(Float64,0)

t0 = Array(Float64,0)
t2 = Array(Float64,0)
t5 = Array(Float64,0)
told = Array(Float64,0)
tp = Array(Float64,0)

for i in 1:nruns
    a = wtedChimera(n+i,1)
    f(t) = sum(compStretches(t,a))/nnz(a)

    try

    tic()
    tr = akpwish(a)
    at0 = toq()
    ast0 = f(tr)

    tic()
    tr = akpwish(a,ver=2)
    at2 = toq()
    ast2 = f(tr)

    tic()
    tr = akpwish(a,ver=5)
    at5 = toq()
    ast5 = f(tr)

    tic()
    tr = akpw(a)
    atold = toq()
    astold = f(tr)

    tic()
    tr = randishPrim(a)
    atp = toq()
    astp = f(tr)

    push!(t0,at0)
    push!(st0,ast0)

    push!(t2,at2)
    push!(st2,ast2)

    push!(t5,at5)
    push!(st5,ast5)

    push!(tp,atp)
    push!(stp,astp)

    push!(told,atold)
    push!(stold,astold)

    catch err
        println("bad on wtedChimera ", n, " ", i)
        @show err
    end

end

vstat = function(x)
    x = sort(x)
    n = length(x)
    n10 = div(n,10)
    [mean(x) minimum(x) x[n10:n10:(9*n10)]' maximum(x)]
end

println("st0 : ", round(vstat(st0./st5),3))
println("st2 : ", round(vstat(st2./st5),3))
println("stp : ", round(vstat(stp./st5),3))
println("stold : ", round(vstat(stold./st5),3))

println("t0 : ", round(vstat(t0./t5),3))
println("t2 : ", round(vstat(t2./t5),3))
println("tp : ", round(vstat(tp./t5),3))
println("told : ", round(vstat(told./t5),3))

    return t0, t2, t5, tp, told, st0, st2, st5, stp, stold
end

function testTrees2(n,nruns)

st2 = Array(Float64,0)
st5 = Array(Float64,0)
stold = Array(Float64,0)
stp = Array(Float64,0)

t2 = Array(Float64,0)
t5 = Array(Float64,0)
told = Array(Float64,0)
tp = Array(Float64,0)

for i in 1:nruns
    a = wtedChimera(n,i)
    f(t) = sum(compStretches(t,a))/nnz(a)

    try

    tic()
    tr = akpwish(a,ver=2)
    at2 = toq()
    ast2 = f(tr)

    tic()
    tr = akpwish(a,ver=5)
    at5 = toq()
    ast5 = f(tr)

    tic()
    tr = akpw(a)
    atold = toq()
    astold = f(tr)

    tic()
    tr = randishPrim(a)
    atp = toq()
    astp = f(tr)

    push!(t2,at2)
    push!(st2,ast2)

    push!(t5,at5)
    push!(st5,ast5)

    push!(tp,atp)
    push!(stp,astp)

    push!(told,atold)
    push!(stold,astold)

    catch err
        println("bad on wtedChimera ", n, " ", i)
        @show err
    end

end

vstat = function(x)
    x = sort(x)
    n = length(x)
    n10 = div(n,10)
    [mean(x) minimum(x) x[n10:n10:(9*n10)]' maximum(x)]
end

println("st2vs5 : ", round(vstat(st2./st5),3))
println("st5 : ", round(vstat(st5./st2),3))
println("stp : ", round(vstat(stp./st2),3))
println("stold : ", round(vstat(stold./st2),3))

println("t2vs5 : ", round(vstat(t2./t5),3))
println("t5 : ", round(vstat(t5./t2),3))
println("tp : ", round(vstat(tp./t2),3))
println("told : ", round(vstat(told./t2),3))

    return t2, t5, tp, told, st2, st5, stp, stold
end

function testTrees(n,nruns,fn)

st0 = Array(Float64,0)
st2 = Array(Float64,0)
st4 = Array(Float64,0)
st5 = Array(Float64,0)
stp = Array(Float64,0)
stk = Array(Float64,0)
stold = Array(Float64,0)

t0 = Array(Float64,0)
t2 = Array(Float64,0)
t4 = Array(Float64,0)
t5 = Array(Float64,0)
tp = Array(Float64,0)
tk = Array(Float64,0)
told = Array(Float64,0)

h = open(fn,"w")

for i in 1:nruns
    a = wtedChimera(n,i)
    f(t) = sum(compStretches(t,a))/nnz(a)

    try

    tic()
    tr = akpwish(a,ver=0)
    at0 = toq()
    ast0 = f(tr)

    tic()
    tr = akpwish(a,ver=2)
    at2 = toq()
    ast2 = f(tr)

    tic()
    tr = akpwish(a,ver=4)
    at4 = toq()
    ast4 = f(tr)

        #=
    tic()
    tr = akpwish(a,ver=5)
    at5 = toq()
    ast5 = f(tr)
        =#

    tic()
    tr = randishKruskal(a)
    atk = toq()
    astk = f(tr)

    tic()
    tr = randishPrim(a)
    atp = toq()
    astp = f(tr)

    tic()
    tr = akpw(a)
    atold = toq()
    astold = f(tr)

    print(h, n, " ", i, " : ")

    print(h, ast0, " ")
    push!(t0,at0)
    push!(st0,ast0)

    print(h, ast2, " ")
    push!(t2,at2)
    push!(st2,ast2)

    print(h, ast4, " ")
    push!(t4,at4)
    push!(st4,ast4)

        #=
    print(h, ast5, " ")
    push!(t5,at5)
    push!(st5,ast5)
        =#

    print(h, astk, " ")
    push!(tk,atk)
    push!(stk,astk)

    print(h, astp, " ")
    push!(tp,atp)
    push!(stp,astp)

    print(h, astold, " ")
    push!(told,atold)
    push!(stold,astold)

    print(h, "\n")

    catch
        println("bad on wtedChimera ", n, " ", i)
    end

end

mist = minimum([st0 st2 st4 stk stp stold],2)
mint = minimum([t0 t2 t4  tk tp told],2)

f = 1.3

stats(v) = [mean(v./mist) median(v./mist) maximum(v./mist) mean(v .== mist) mean(v .> f*mist)]
statt(v) = [mean(v./mint) median(v./mint) maximum(v./mint) mean(v .== mint) mean(v .> f*mint)]

println(h, "st0 : ", stats(st0))
println(h, "st2 : ", stats(st2))
println(h, "st4 : ", stats(st4))
#println(h, "st5 : ", stats(st5))
println(h, "stp : ", stats(stp))
println(h, "stk : ", stats(stk))
println(h, "stold : ", stats(stold))

println(h, "t0 : ", statt(t0))
println(h, "t2 : ", statt(t2))
println(h, "t4 : ", statt(t4))
#println(h, "t5 : ", statt(t5))
println(h, "tp : ", statt(tp))
println(h, "tk : ", statt(tk))
println(h, "told : ", statt(told))

    close(h)

end

ext: jl


text:
include("../src/EDKit.jl")
using LinearAlgebra
using .EDKit
using Test
#-------------------------------------------------------------------------------------------------------------------------
# Test Translational PXP
#-------------------------------------------------------------------------------------------------------------------------
@testset "Multi-threads Translational PXP" begin
    L, k, p = 28, 0, 1
    mat = begin
        P = Diagonal([1, 1, 1, 0, 1, 1, 0, 0])
        X = [0 1; 1 0]
        P * kron(I(2), X, I(2)) * P
    end
    pxpf(v::Vector{<:Integer}) = all(v[i]==0 || v[mod(i, length(v))+1]==0 for i=1:length(v))
    println("--------------------------------------")
    print("Single-threads:")
    @time bs = TranslationalBasis(pxpf, k, L, threaded=false)
    print("Multi-threads :")
    @time bm = TranslationalBasis(pxpf, k, L, threaded=true)
    @test bs.I == bm.I
    @test norm(bs.R-bm.R) ≈ 0.0
end

#-------------------------------------------------------------------------------------------------------------------------
# Test Translational Parity PXP
#-------------------------------------------------------------------------------------------------------------------------
@testset "Multi-threads Translational Parity PXP" begin
    L, k, p = 28, 0, 1
    mat = begin
        P = Diagonal([1, 1, 1, 0, 1, 1, 0, 0])
        X = [0 1; 1 0]
        P * kron(I(2), X, I(2)) * P
    end
    pxpf(v::Vector{<:Integer}) = all(v[i]==0 || v[mod(i, length(v))+1]==0 for i=1:length(v))
    println("--------------------------------------")
    print("Single-threads:")
    @time bs = TranslationParityBasis(pxpf, k, p, L, threaded=false)
    print("Multi-threads :")
    @time bm = TranslationParityBasis(pxpf, k, p, L, threaded=true)
    @test bs.I == bm.I
    @test norm(bs.R-bm.R) ≈ 0.0
end


text:
module KroneckerProducts

using Base.Cartesian
using LinearAlgebra

export Kronecker, KroneckerProduct, KroneckerFactorization, ⊗

"""
A `Kronecker` object is a wrapper object that represents the applying a binary
operator to all combinations of the elements in two arrays in a fashion similar
to how the Kronecker Product works.

Creating a `Kronecker` object requires almost no time or storage, as it just
references the input arrays. The computation is performed each time an element
of the new array is accessed. Mutating the input arrays will mutate the result.
"""
struct Kronecker{T,N,F,A,B} <: AbstractArray{T,N}
    f::F
    a::A
    b::B
end

Kronecker{T,N}(f::Function,a::AbstractArray,b::AbstractArray) where {T,N} = Kronecker{T,N,typeof(f),typeof(a),typeof(b)}(f,a,b)

function Kronecker{T}(f::Function, a::AbstractArray{<:Any,NA}, b::AbstractArray{<:Any,NB}) where {T,NA,NB}
   N = max(NA,NB)
   Kronecker{T,N}(f,a,b)
end

function Kronecker(f::Function, a::AbstractArray{TA}, b::AbstractArray{TB}) where {TA,TB}
   T = typeof(f(zero(TA),zero(TB)))
   Kronecker{T}(f,a,b)
end

"""
A `KroneckerProduct` is a wrapper object that represents the Kronecker product
of two matrices (or the extension thereof for arrays in any dimension).

It is a special case of a `Kronecker` object.
"""
KroneckerProduct{T,N,A,B} = Kronecker{T,N,typeof(*),A,B}
KroneckerProduct{T,N}(a, b) where {T,N} = Kronecker{T,N}(*,a,b)
KroneckerProduct(a, b) = Kronecker(*,a,b)

"""
    KroneckerProducts.kron(a, b)

Form the Kronecker product of `a` and `b`, same as `Base.kron` but use a
KroneckerProduct object when possible.
"""
kron(a, b) = Base.kron(a, b)
kron(a::AbstractMatrix, b::AbstractMatrix) = KroneckerProduct(a, b)

"The ⊗ operator is an alias for KroneckerProducts.kron"
⊗(a,b) = kron(a,b)

Base.size(x::Kronecker{T,N}) where {T,N} = ntuple(i->size(x.a, i)*size(x.b, i), N)

@generated function Base.getindex(x::Kronecker{T,N,F}, idx::CartesianIndex{N})::T where {T,N,F}
    quote
        a = x.a
        b = x.b
        @nexprs $N i->(sz_i = size(b, i))
        x.f(@nref($N, a, i->(idx[i]-1)÷sz_i+1), @nref($N, b, i->(idx[i]-1)%sz_i+1))
    end
end

#Base.getindex(x::Kronecker, idx::Integer) = x[CartesianIndices(x)[idx]] # Not needed
Base.getindex(x::Kronecker, idxs::Integer...) = x[CartesianIndex(idxs)]

# Pass through wrappers from adjoint/transpose
Base.adjoint(x::KroneckerProduct{T,N}) where {T,N} = KroneckerProduct{T,N}(x.a', x.b')
Base.transpose(x::KroneckerProduct{T,N}) where {T,N} = KroneckerProduct{T,N}(transpose(x.a), transpose(x.b))

function LinearAlgebra.det(x::KroneckerProduct{T,2}) where {T}
    LinearAlgebra.checksquare(x)
    size(x.a,1) != size(x.a,2) && return zero(T)
    size(x.b,1) != size(x.b,2) && return zero(T)
    det(x.a)^size(x.b,1) * det(x.b)^size(x.a, 1)
end

function LinearAlgebra.tr(x::KroneckerProduct{T,2}) where {T}
    n = LinearAlgebra.checksquare(x)
    if size(x.a,1) == size(x.a,2)
        tr(x.a)*tr(x.b)
    else
        mapreduce(i->x[i,i], +, 1:n)
    end
end

gettmp(x::KroneckerProduct{Tx,2}, y::AbstractVector{Ty}) where {Tx, Ty} =
    Matrix{promote_type(Tx, Ty)}(undef, size(x.b,1), size(x.a,2))

function LinearAlgebra.mul!(y!::Vector, x::KroneckerProduct{<:Any,2}, y::AbstractVector;
        tmp=gettmp(x, y) )
    size(x,1) == length(y!) || throw(DimensionMismatch())
    size(x,2) == length(y) || throw(DimensionMismatch())
    mul!(tmp, x.b, reshape(y, size(x.b,2), size(x.a,2)))
    mul!(reshape(y!, size(x.b,1), size(x.a,1)), tmp, transpose(x.a))
    y!
end

"A KroneckerFactorization is a factorization of a KroneckerProduct."
struct KroneckerFactorization{T,A,B} <: Factorization{T}
    a::A
    b::B
end

KroneckerFactorization(a::A, b::B) where {A<:Factorization{T}, B<:Factorization{T}} where {T} =
   KroneckerFactorization{T,A,B}(a,b)

# TODO: handle case where factorizations have different element types
for f in (:factorize, :lu, :qr, :cholesky)
    @eval LinearAlgebra.$f(x::KroneckerProduct) = KroneckerFactorization($f(x.a), $f(x.b))
end

Base.size(x::KroneckerFactorization) = size(x.a).*size(x.b)
Base.size(x::KroneckerFactorization, i::Integer) = size(x.a, i)*size(x.b, i)

gettmp(x::KroneckerFactorization{T}, y::AbstractVector) where {T} =
    Matrix{T}(undef, size(x.a,2), size(x.b,1))

function LinearAlgebra.ldiv!(y!::Vector, x::KroneckerFactorization, y::AbstractVector;
        tmp=gettmp(x, y) )
    size(x,1) == length(y!) || throw(DimensionMismatch())
    size(x,2) == length(y) || throw(DimensionMismatch())
    tmp .= transpose(reshape(y, size(x.b,1), size(x.a,1)))
    ldiv!(x.a, tmp)
    y! .= vec(transpose(tmp))
    ldiv!(x.b, reshape(y!, size(x.b,2), size(x.a,2)))
    y!
end

LinearAlgebra.ldiv!(x::KroneckerFactorization, y!::Vector) = ldiv!(y!, x, y!)

end # module


text:
@testset "Zero" begin
    z = Zero()
    @test extern(z) === false
    @test z + z == z
    @test z + 1 == 1
    @test 1 + z == 1
    @test z * z == z
    @test z * 1 == z
    @test 1 * z == z
    for x in z
        @test x === z
    end
    @test broadcastable(z) isa Ref{Zero}
    @test conj(z) == z
end


text:
using Gadfly

plot(sin, 0, 25, Coord.cartesian(aspect_ratio=1.0))


text:
using Postgres
P = Postgres
naval = P.Types.naval

conn = connect(PostgresServer, db="julia_test", host="localhost")
curs = cursor(conn)
query(curs, "select 1")

f = x -> repr(P.Types.PostgresValue(naval(x)))
exprs = [(v, f(v)) for v in values(conn.pgtypes)]
proto = x -> "select $x from generate_series(1, 1000000)"
results = []
for (t, exp) in exprs
     @time query(curs, proto(exp))
     println("$exp : $(get(curs.query_time))")
     push!(results, (get(curs.query_time).fetch_time, t))
end

for (time, kind) in results
    println("$(kind.name)\t\t\t$(round(time, 5))")
end

max_stars_repo_name: NCarson/Postgres

max_issues_repo_name: NCarson/Postgres

max_forks_repo_name: NCarson/Postgres


text:
using Colors, ColorBrewer

function do_palette(name="Blues", size::Int=9)
  width, height = "1cm", "0.5cm"
  path = "$colorbrewer_git_path/images"
  svg_filename = "$path/$name.svg"
  png_filename = "$path/$name.png"
  f = open(svg_filename, "w")
  @printf f "%s\n" "<svg xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\" width=\"$(size)cm\" height=\"0.5cm\">"
  x = 0
  for c in palette(name, size)
    r, g, b = int(round(c.r*255)), int(round(c.g*255)), int(round(c.b*255))
    @printf f "%s\n" "<rect x=\"$(x)cm\" width=\"$width\" height=\"$height\" fill=\"rgb($r, $g, $b)\"/>"
    x += 1
  end
  @printf f "</svg>"
  close(f)
  run(`svg2png $svg_filename $png_filename`)
  run(`rm $svg_filename`)
end

function main()
  for name in ["Blues","Oranges","Greens","Reds","Purples","Greys","OrRd","GnBu","PuBu","PuRd","BuPu","BuGn","YlGn","RdPu","YlOrBr","YlGnBu","YlOrRd","PuBuGn"]
    do_palette(name, 9)
  end
  for name in ["Spectral","RdYlGn","RdBu","PiYG","PRGn","RdYlBu","BrBG","RdGy","PuOr"]
    do_palette(name, 11)
  end
  do_palette("Set1", 9)
  do_palette("Set2", 8)
  do_palette("Set3", 12)
  do_palette("Accent", 8)
  do_palette("Dark2", 8)
  do_palette("Paired", 12)
  do_palette("Pastel1", 9)
  do_palette("Pastel2", 8)
end

nothing


text:
## UC Model
function UC(Ω, ρ, RES; printing = 0)
    mUC = Model(() -> Gurobi.Optimizer(GUROBI_ENV))
    set_optimizer_attribute(mUC, "OutputFlag", printing)

        ### Variables
        # 1st stage
        @variables(mUC, begin
            v[G,T], (Bin) # Unit state
            0.0 ≤ c_su[G,T] # start-up cost
            0.0 ≤ c_rc[G,T] # reserve capacity cost
            0.0 ≤ obj1
            0. ≤ p_exp[g in G,T] ≤ GenData.Pmax[g] # Expected generation
            0. ≤ r_up[g in G,T] ≤ GenData.Rup[g] # Upward reserve capacity
            0. ≤ r_dw[g in G,T] ≤ GenData.Rdw[g] # Downward reserve capacity
            end)

            # 2nd stage
            @variables(mUC, begin
                0. ≤ p[g in G,Ω,T] ≤ GenData.Pmax[g] # Real-time generation
                0. ≤ r_p[g in G,Ω,T] #≤ GenData.Rplus[g] # Real-time up reserve
                0. ≤ r_m[g in G,Ω,T] #≤ GenData.Rminus[g] # Real-time down reserve
                0. ≤ l_s[t in T, Ω] ≤ sum(Pd,dims=1)[t] # load shedding
                0. ≤ w_s[t in T, ω in Ω] ≤ RES[t,ω] # wind spillage
                0. ≤ Θ[Ω] # operation costs

            end)

        ### Constraints
        #### 1st stage
            # Power limits
            @constraints(mUC, begin
                PMax[g in G, t in T], p_exp[g,t] + r_up[g,t]≤ v[g,t]*GenData.Pmax[g]
                PMin[g in G, t in T], p_exp[g,t] - r_dw[g,t]≥ v[g,t]*GenData.Pmin[g]
                RupLim[g in G, ω in Ω, t in T], r_up[g,t] ≤ v[g,t]*GenData.Rup[g]
                RdwLim[g in G, ω in Ω, t in T], r_dw[g,t] ≤ v[g,t]*GenData.Rdw[g]
                end)
            # Costs
            @constraint(mUC, obj1 == sum(GenData.Cg[g]*p_exp[g,t]
                                + c_su[g,t] + c_rc[g,t] for g in G, t in T))

            @constraints(mUC, begin
                CostReserveCap[g in G, t in T],
                        c_rc[g,t] == r_up[g,t]*GenData.Cup[g] + r_dw[g,t]*GenData.Cdw[g]
                #
                CostSU[g in G, t in T],
                        c_su[g,t] ≥ GenData.Cstart[g]*(v[g,t] - v[g,T[t-1]])
                            end)

            # Minimum Online & Offline Times
            @constraints(mUC, begin
                MinOn[g in G, t in T],
                    sum(v[g,n] for n in t:min(t+GenData.UT[g]-1,24)) ≥ GenData.UT[g]*(v[g,t]-v[g,T[t-1]])

                MinOff[g in G, t in T],
                    sum(1-v[g,n] for n in t:min(t+GenData.DT[g]-1,24)) ≥ GenData.DT[g]*(v[g,T[t-1]]-v[g,t])
                end)

            #### 2nd stage
            # Power balance
            @constraint(mUC, PowerBalance_Second[ω in Ω, t in T],
                                sum(Pd,dims=1)[t] - l_s[t,ω]== RES[t,ω] + sum(p[:,ω,t]) - w_s[t,ω])

            # Power limits
            @constraints(mUC, begin
                PMax2nd[g in G, ω in Ω, t in T], p[g,ω,t] ≤ v[g,t]*GenData.Pmax[g]
                PMin2nd[g in G, ω in Ω, t in T], p[g,ω,t] ≥ v[g,t]*GenData.Pmin[g]
                RplusLim[g in G, ω in Ω, t in T], r_p[g,ω,t] ≤ r_up[g,t]
                RminusLim[g in G, ω in Ω, t in T], r_m[g,ω,t] ≤ r_dw[g,t]
                Ramps_2nd[g in G, ω in Ω, t in T], -GenData.Rdw[g] ≤ p[g,ω,t] - p[g,ω,T[t-1]] ≤ GenData.Rup[g]

                end)

            # Generated Power
            @constraint(mUC, PGen_2nd[g in G, ω in Ω, t in T],
                    p[g,ω,t] == p_exp[g,t] + r_p[g,ω,t] - r_m[g,ω,t])

            # Total Cost
            @constraint(mUC, Cost_2nd[ω in Ω],
                    Θ[ω] == sum(C_ls*l_s[t,ω] + C_ws*w_s[t,ω] + sum(GenData.Cplus[g]*r_p[g,ω,t] - GenData.Cminus[g]*r_m[g,ω,t] for g in G) for t in T)
                    )

        ### Objective
        @objective(mUC, Min, obj1 + sum(ρ[ω]*Θ[ω] for ω in Ω))

        optimize!(mUC)
        # println("Gap: $(MOI.get(mUC, MOI.RelativeGap()))")

    V = JuMP.value.(v)
        V= collect(V[i,j] for i in G, j in T)
            Pexp = (JuMP.value.(p_exp))
            Pexp= collect(Pexp[i,j] for i in G, j in T)
            Rup = (JuMP.value.(r_up))
            Rup= collect(Rup[i,j] for i in G, j in T)
            Rdw = (JuMP.value.(r_dw))
            Rdw= collect(Rdw[i,j] for i in G, j in T)
            Obj1 = JuMP.value.(obj1)
            LB = JuMP.objective_value.(mUC)

            # println("l_s:")
            # println(C_ls*sum(JuMP.value.(l_s)))
            # println("w_s")
            # println(C_ws*sum(JuMP.value.(w_s)))
            # println("Reserves:")
            # println(sum(sum(GenData.Cplus[g]*JuMP.value.(r_p[g,ω,t]) - GenData.Cminus[g]*JuMP.value.(r_m[g,ω,t]) for g in G) for t in T, ω in Ω))

    return V, Pexp, Rup, Rdw, Obj1, LB, mUC

    end

max_stars_repo_name: gonzalez-alvaro/MultidimensionalGAPM

max_issues_repo_name: gonzalez-alvaro/MultidimensionalGAPM

max_forks_repo_name: gonzalez-alvaro/MultidimensionalGAPM


text:
using ImageTracking
using Test

include("farneback.jl")
include("lucas_kanade.jl")
include("haar.jl")
include("visualization.jl")
include("error_evaluation.jl")


text:
export SceneNil, is_nil

struct SceneNil <: Ranger.AbstractScene
    base::NodeData

    function SceneNil()
        new(NodeData(UInt32(0), "_Nil_"))
    end
end

function is_nil(node::Ranger.AbstractScene)
    node.base.id == 0
end

max_stars_repo_name: wdevore/Ranger-Julia-SDL

max_issues_repo_name: wdevore/Ranger-Julia-SDL

max_forks_repo_name: wdevore/Ranger-Julia-SDL


text:
module Mesh

using LinearAlgebra
using LinearAlgebra.BLAS
using jInv.Utils
using SparseArrays
using Printf
using Distributed

export AbstractMesh
export AbstractTensorMesh
abstract type AbstractMesh end
abstract type AbstractTensorMesh <: AbstractMesh end

import Distributed.clear!
function clear!(M::AbstractTensorMesh)
	M.Div  = clear!(M.Div )
	M.Grad = clear!(M.Grad)
	M.Curl = clear!(M.Curl)
	M.Af   = clear!(M.Af  )
	M.Ae   = clear!(M.Ae  )
	M.An   = clear!(M.An  )
	M.V    = clear!(M.V   )
	M.F    = clear!(M.F   )
	M.L    = clear!(M.L   )
	M.Vi   = clear!(M.Vi  )
	M.Fi   = clear!(M.Fi  )
	M.Li   = clear!(M.Li  )
	M.nLap = clear!(M.nLap  )
end

function speye(n)
	return sparse(1.0I,n,n);
end

function spdiagm(x::Vector)
	return sparse(Diagonal(x));
end

include("generic.jl")
include("tensor.jl")
include("regular.jl")
include("interpmat.jl")
include("display.jl")
include("regularVectorFaces.jl")

export getNodalConstraints, getEdgeConstraints, getFaceConstraints
getNodalConstraints(M::AbstractMesh) = (UniformScaling(1.0), UniformScaling(1.0), 1:prod(M.n+1))
getEdgeConstraints(M::AbstractMesh) = (UniformScaling(1.0), UniformScaling(1.0), 1:sum(M.ne))
getFaceConstraints(M::AbstractMesh) = (UniformScaling(1.0), UniformScaling(1.0))

export clear!

end


text:
using OrderedCollections: OrderedDict
using CUDAapi: has_cuda

using Oceananigans: AbstractOutputWriter, AbstractDiagnostic, TimeStepper
using Oceananigans.Architectures: AbstractArchitecture
using Oceananigans.Buoyancy: validate_buoyancy
using Oceananigans.TurbulenceClosures: ν₀, κ₀, with_tracers

mutable struct Model{TS, E, A<:AbstractArchitecture, G, T, B, R, SW, U, C, Φ, F,
                     BCS, S, K, OW, DI, Θ} <: AbstractModel

           architecture :: A         # Computer `Architecture` on which `Model` is run
                   grid :: G         # Grid of physical points on which `Model` is solved
                  clock :: Clock{T}  # Tracks iteration number and simulation time of `Model`
               buoyancy :: B         # Set of parameters for buoyancy model
               coriolis :: R         # Set of parameters for the background rotation rate of `Model`
          surface_waves :: SW        # Set of parameters for surfaces waves via the Craik-Leibovich approximation
             velocities :: U         # Container for velocity fields `u`, `v`, and `w`
                tracers :: C         # Container for tracer fields
              pressures :: Φ         # Container for hydrostatic and nonhydrostatic pressure
                forcing :: F         # Container for forcing functions defined by the user
                closure :: E         # Diffusive 'turbulence closure' for all model fields
    boundary_conditions :: BCS       # Container for 3d bcs on all fields
            timestepper :: TS        # Object containing timestepper fields and parameters
        pressure_solver :: S         # Pressure/Poisson solver
          diffusivities :: K         # Container for turbulent diffusivities
         output_writers :: OW        # Objects that write data to disk
            diagnostics :: DI        # Objects that calc diagnostics on-line during simulation
             parameters :: Θ         # Container for arbitrary user-defined parameters
end

"""
    Model(;
                   grid,
           architecture = CPU(),
             float_type = Float64,
                tracers = (:T, :S),
                closure = ConstantIsotropicDiffusivity(float_type, ν=ν₀, κ=κ₀),
                  clock = Clock{float_type}(0, 0),
               buoyancy = SeawaterBuoyancy(float_type),
               coriolis = nothing,
          surface_waves = nothing,
                forcing = ModelForcing(),
    boundary_conditions = HorizontallyPeriodicSolutionBCs(),
         output_writers = OrderedDict{Symbol, AbstractOutputWriter}(),
            diagnostics = OrderedDict{Symbol, AbstractDiagnostic}(),
             parameters = nothing,
             velocities = VelocityFields(architecture, grid),
              pressures = PressureFields(architecture, grid),
          diffusivities = TurbulentDiffusivities(architecture, grid, tracernames(tracers), closure),
            timestepper = :AdamsBashforth,
        pressure_solver = PressureSolver(architecture, grid, PressureBoundaryConditions(boundary_conditions.u))
    )

Construct an `Oceananigans.jl` model on `grid`.

Keyword arguments
=================
- `grid`: (required) The resolution and discrete geometry on which `model` is solved.
- `architecture`: `CPU()` or `GPU()`. The computer architecture used to time-step `model`.
- `float_type`: `Float32` or `Float64`. The floating point type used for `model` data.
- `closure`: The turbulence closure for `model`. See `TurbulenceClosures`.
- `buoyancy`: Buoyancy model parameters.
- `coriolis`: Parameters for the background rotation rate of the model.
- `forcing`: User-defined forcing functions that contribute to solution tendencies.
- `boundary_conditions`: User-defined boundary conditions for model fields. Can be either`SolutionBoundaryConditions`
  or `ModelBoundaryConditions`. See `BoundaryConditions`, `HorizontallyPeriodicSolutionBCs`, and `ChannelSolutionBCs`.
- `parameters`: User-defined parameters for use in user-defined forcing functions and boundary condition functions.
"""
function Model(;
                   grid,
           architecture = CPU(),
             float_type = Float64,
                tracers = (:T, :S),
                closure = ConstantIsotropicDiffusivity(float_type, ν=ν₀, κ=κ₀),
                  clock = Clock{float_type}(0, 0),
               buoyancy = SeawaterBuoyancy(float_type),
               coriolis = nothing,
          surface_waves = nothing,
                forcing = ModelForcing(),
    boundary_conditions = HorizontallyPeriodicSolutionBCs(),
         output_writers = OrderedDict{Symbol, AbstractOutputWriter}(),
            diagnostics = OrderedDict{Symbol, AbstractDiagnostic}(),
             parameters = nothing,
             velocities = VelocityFields(architecture, grid),
              pressures = PressureFields(architecture, grid),
          diffusivities = TurbulentDiffusivities(architecture, grid, tracernames(tracers), closure),
            timestepper = :AdamsBashforth,
        pressure_solver = PressureSolver(architecture, grid, PressureBoundaryConditions(boundary_conditions.u))
    )

    if architecture == GPU()
        !has_cuda() && throw(ArgumentError("Cannot create a GPU model. No CUDA-enabled GPU was detected!"))
        if mod(grid.Nx, 16) != 0 || mod(grid.Ny, 16) != 0
            throw(ArgumentError("For GPU models, Nx and Ny must be multiples of 16."))
        end
    end

    timestepper = TimeStepper(timestepper, float_type, architecture, grid, tracernames(tracers))

    tracers = TracerFields(architecture, grid, tracers)
    validate_buoyancy(buoyancy, tracernames(tracers))

    # Regularize forcing, boundary conditions, and closure for given tracer fields
    forcing = ModelForcing(tracernames(tracers), forcing)
    boundary_conditions = ModelBoundaryConditions(tracernames(tracers), boundary_conditions)
    closure = with_tracers(tracernames(tracers), closure)

    return Model(architecture, grid, clock, buoyancy, coriolis, surface_waves, velocities, tracers,
                 pressures, forcing, closure, boundary_conditions, timestepper,
                 pressure_solver, diffusivities, output_writers, diagnostics, parameters)
end


text:
lines = readlines("d07.in")

baggraph = Dict{String,Vector{Tuple{Int,String}}}()

function invertgraph(baggraph)
    invertedgraph = Dict{String,Set{String}}()
    for (outerbag, innerbags) in baggraph
        for (_, innerbag) in innerbags
            push!(invertedgraph, innerbag => union!(get(invertedgraph, innerbag, Set{String}()), [outerbag]))
        end
    end
    return invertedgraph
end

function containingbags(invertedgraph, desiredbag)
    bags = Set{String}()
    visit = Set([desiredbag])
    while !isempty(visit)
        bag = pop!(visit)
        if bag in keys(invertedgraph)
            outers = invertedgraph[bag]
            union!(bags, outers)
            union!(visit, outers)
            pop!(invertedgraph, bag)
        end
    end
    return bags
end

function countbags(baggraph, desiredbag)
    innerbags = baggraph[desiredbag]
    if isempty(innerbags)
        return 1
    else
        return 1 + sum(n * countbags(baggraph, innerbag) for (n, innerbag) in innerbags)
    end
end

for line in lines
    m = match(r"^([\w ]+) bags contain ", line)
    outerbag = m.captures[1]
    innerbags = eachmatch(r"(\d+) ([\w ]+) bags?[,.]", line[length(m.match)+1:end])
    push!(baggraph, outerbag => [Tuple([parse(Int, i.captures[1]), i.captures[2]]) for i in innerbags])
end

desiredbag = "shiny gold"

P1 = length(containingbags(invertgraph(baggraph),desiredbag))
P2 = countbags(baggraph, desiredbag) - 1

println("P1: ", P1)
println("P2: ", P2)

max_stars_repo_name: eduellery/adventofcode

max_issues_repo_name: eduellery/adventofcode

max_forks_repo_name: eduellery/adventofcode


text:
using Rayly
using StaticArrays
using Test

include("aabb.jl")


text:
using Pkg, Random
Pkg.activate(".")

module EnzymeMFABC
include("../example.jl")
end

using FileIO, JLD2

Random.seed!(444)

R = 5
θ_lo = (EnzymeMFABC.k1, EnzymeMFABC.k2, EnzymeMFABC.k3, EnzymeMFABC.e0)
t_lo, x_lo, c_lo = EnzymeMFABC.simulate(EnzymeMFABC.EnzymeModels.EnzymeMM(θ_lo...), [EnzymeMFABC.s0, EnzymeMFABC.p0], EnzymeMFABC.tspan)
θ_hi = (EnzymeMFABC.k1, EnzymeMFABC.k2, EnzymeMFABC.k3)

uncoupled = map(1:R) do r
    t_hi_uc, x_hi_uc, _ = EnzymeMFABC.simulate(EnzymeMFABC.EnzymeModels.Enzyme(θ_hi...), [EnzymeMFABC.s0, EnzymeMFABC.e0, EnzymeMFABC.c0, EnzymeMFABC.p0], EnzymeMFABC.tspan)
    (t_hi_uc, x_hi_uc)
end

coupled = map(1:R) do r
    t_hi_cp, x_hi_cp, _ = EnzymeMFABC.simulate(EnzymeMFABC.EnzymeModels.Enzyme(θ_hi...), [EnzymeMFABC.s0, EnzymeMFABC.e0, EnzymeMFABC.c0, EnzymeMFABC.p0], EnzymeMFABC.tspan, (nothing, nothing, c_lo[1]))
    (t_hi_cp, x_hi_cp)
end

save("eg_traj.jld2", Dict(
    "R" => R,
    "t_lo" => t_lo,
    "x_lo" => x_lo,
    "uncoupled" => uncoupled,
    "coupled" => coupled,
))

max_stars_repo_name: tpprescott/mf-lf

max_issues_repo_name: tpprescott/mf-lf

max_forks_repo_name: tpprescott/mf-lf


text:
#!/usr/bin/env julia
#=
performance is the same for isnan() and == nothing in Julia
=#
println("--> Julia $VERSION")

try
  using BenchmarkTools
catch
  import Pkg
  Pkg.add("BenchmarkTools")
  Pkg.update("BenchmarkTools")
  using BenchmarkTools
end

function fnan()
for i = 0:100
   x = isnan(i)
   if i > 1e100; break; end  # to break JIT
end
end

function fnothing()
for i = 0:100
   x = nothing == 0.
   if i > 1e100; break; end  # to break JIT
end
end

o=@benchmark fnan()
println("isnan", o)

println("is nothing")
o=@benchmark fnothing()
println(o)

max_stars_repo_name: scienceopen/numba-examples

max_issues_repo_name: scienceopen/numba-examples

max_forks_repo_name: scienceopen/numba-examples


text:
@testset "mapcube" begin
    @testset "Loop Axis permutation" begin
        x,y,z = (RangeAxis("x",1:4), RangeAxis("y", 1:5), RangeAxis("z", 1:6))
        a1 = YAXArray([x,y,z], rand(4,5,6))
        a2 = YAXArray([x,z,y], rand(4,6,5))
        a3 = YAXArray([x,y], rand(4,5))
        indims = InDims("x")
        outdims = OutDims("x")
        r = mapCube((a1, a2), indims=(indims, indims), outdims=outdims) do xout, x1, x2
            xout .= x1 .+ x2
        end
        @test r.data == a1.data .+ permutedims(a2.data,(1,3,2))
        r = mapCube((a2, a3), indims=(indims, indims), outdims=outdims) do xout, x1, x2
            xout .= x1 .+ x2
        end
        @test r.data == a2.data .+ reshape(a3.data,(4,1,5))
    end

    @testset "Reduction to empty outputs" begin
        yax = YAXArray([RangeAxis("axis1", 1:10)], rand(10))
        r = mapslices(sum, yax, dims="axis1")
        @test r.data[] == sum(yax.data)

    end
end

max_stars_repo_name: meggart/YAXArrays

max_forks_repo_name: meggart/YAXArrays


text:
const runtime_path = "/var/runtime"
const julia_depot_path = "/var/runtime/julia_depot"
const temp_path = "/tmp"
const jot_github_url = "https://github.com/harris-chris/Jot.jl#main"

function dockerfile_add_julia_image(julia_base_version::String)::String
  """
  FROM julia:$julia_base_version
  """
end

function dockerfile_add_additional_registries(additional_registries::Vector{String})::String
  using_pkg_script = "using Pkg; "
  add_registries_script = begin
    str = ""
    for reg in additional_registries
      str *= "Pkg.Registry.add(RegistrySpec(url = \\\"$(reg)\\\")); "
    end
    str *= "Pkg.Registry.add(\\\"General\\\")"
    str == "" ? str : str * "; "
  end
  """
  RUN julia -e \"$using_pkg_script$add_registries_script\"
  """
end

function dockerfile_add_utilities()::String
  """
  RUN apt-get update && apt-get install -y \\
    build-essential
  """
end

function dockerfile_add_runtime_directories()::String
  """
  RUN mkdir -p $julia_depot_path
  ENV JULIA_DEPOT_PATH=$julia_depot_path
  RUN mkdir -p $temp_path
  ENV TMPDIR=$temp_path
  RUN mkdir -p $runtime_path
  WORKDIR $runtime_path
  """
end

function dockerfile_copy_build_dir()::String
  """
  COPY . .
  """
end

function dockerfile_add_responder(res::LocalPackageResponder)::String
  using_pkg_script = "using Pkg; "
  add_module_script = "Pkg.develop(path=\\\"$runtime_path/$(res.package_name)\\\"); "
  instantiate_script = "Pkg.instantiate(); "
  """
  RUN julia -e \"$using_pkg_script$add_module_script$instantiate_script\"
  """
end

function dockerfile_add_jot()::String
  test_running = get(ENV, "JOT_TEST_RUNNING", nothing)
  jot_url = if isnothing(test_running) || test_running == "false"
    jot_github_url
  else
    this_branch = readchomp(`git branch --show-current`)
    jot_github_url * "#$this_branch"
  end
  @debug jot_url

  """
  RUN julia -e "using Pkg; Pkg.add([\\\"HTTP\\\", \\\"JSON3\\\"]); Pkg.add(url=\\\"$jot_url\\\")"
  """
end

function dockerfile_add_aws_rie()::String
  """
  RUN curl -Lo ./aws-lambda-rie \\
  https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie
  RUN chmod +x ./aws-lambda-rie
  """
end

function dockerfile_add_bootstrap(
    package_name::String,
    function_name::String,
    ::Type{IT}
  )::String where {IT}
  """
  ENV PKG_NAME=$(package_name)
  ENV FUNC_FULL_NAME=$package_name.$function_name
  ENV FUNC_PARAM_TYPE=$IT
  RUN chmod 775 . -R
  ENTRYPOINT ["/var/runtime/bootstrap"]
  """
end

function dockerfile_add_precompile(package_compile::Bool)::String
  run_init = """
  RUN julia init.jl
  """
end

function dockerfile_add_labels(labels::Labels)::String
  labels_str = to_docker_buildfile_format(labels)
  """
  LABEL $labels_str
  """
end

function get_dockerfile_build_cmd(
    dockerfile::String,
    image_full_name_plus_tag::String,
    no_cache::Bool,
    build_args::AbstractDict{String, String},
  )::Cmd
  options = ["--rm", "--iidfile", "id", "--tag", "$image_full_name_plus_tag"]
  if length(build_args) > 0
    build_args = foldl(build_args; init=" ") do acc, (opt, arg)
      acc * "--build arg $opt=$arg "
    end
    push!(options, build_args)
  end
  # [push!(options, e) for e in map(x -> [first(x), last(x)], build_args) |> x -> reduce(vcat, x)]
  no_cache && push!(options, "--no-cache")
  `docker build $options .`
end


text:
"""
    struct ScaledMonomialBasis{MT<:MP.AbstractMonomial, MV<:AbstractVector{MT}} <: AbstractPolynomialBasis
        monomials::MV
    end

*Scaled monomial basis* (see [Section 3.1.5, BPT12]) with the monomials of the vector `monomials`.
Given a monomial ``x^\\alpha = x_1^{\\alpha_1} \\cdots x_n^{\\alpha_n}`` of degree ``d = \\sum_{i=1}^n \\alpha_i``,
the corresponding polynomial of the basis is
```math
{d \\choose \\alpha}^{\\frac{1}{2}} x^{\\alpha} \\quad \\text{ where } \\quad
{d \\choose \\alpha} = \\frac{d!}{\\alpha_1! \\alpha_2! \\cdots \\alpha_n!}.
```

For instance, create a polynomial with the basis ``[xy^2, xy]`` creates the polynomial
``\\sqrt{3} a xy^2 + \\sqrt{2} b xy`` where `a` and `b` are new JuMP decision variables.
Constraining the polynomial ``axy^2 + bxy`` to be zero with the scaled monomial basis constrains
`a/√3` and `b/√2` to be zero.

[BPT12] Blekherman, G.; Parrilo, P. A. & Thomas, R. R.
*Semidefinite Optimization and Convex Algebraic Geometry*.
Society for Industrial and Applied Mathematics, **2012**.
"""
struct ScaledMonomialBasis{MT<:MP.AbstractMonomial, MV<:AbstractVector{MT}} <: AbstractMonomialBasis{MT, MV}
    monomials::MV
end
ScaledMonomialBasis(monomials) = ScaledMonomialBasis(monovec(monomials))

MP.polynomialtype(::ScaledMonomialBasis{MT}, T::Type) where MT = MP.polynomialtype(MT, promote_type(T, Float64))
MP.polynomial(f::Function, basis::ScaledMonomialBasis) = MP.polynomial(i -> scaling(basis.monomials[i]) * f(i), basis.monomials)

function Base.promote_rule(::Type{ScaledMonomialBasis{MT, MV}}, ::Type{MonomialBasis{MT, MV}}) where {MT, MV}
    return MonomialBasis{MT, MV}
end

function change_basis(Q::AbstractMatrix, basis::ScaledMonomialBasis{MT, MV}, B::Type{MonomialBasis{MT, MV}}) where {MT, MV}
    n = length(basis)
    scalings = map(scaling, basis.monomials)
    scaled_Q = [Q[i, j] * scalings[i] * scalings[j] for i in 1:n, j in 1:n]
    return scaled_Q, MonomialBasis(basis.monomials)
end

function MP.polynomial(Q::AbstractMatrix, basis::ScaledMonomialBasis{MT, MV}, T::Type) where {MT, MV}
    return MP.polynomial(change_basis(Q, basis, MonomialBasis{MT, MV})..., T)
end

scaling(m::MP.AbstractMonomial) = √(factorial(MP.degree(m)) / prod(factorial, MP.exponents(m)))
unscale_coef(t::MP.AbstractTerm) = coefficient(t) / scaling(monomial(t))
function MP.coefficients(p, ::Type{<:ScaledMonomialBasis})
    return unscale_coef.(MP.terms(p))
end


text:
module types

export ColumnSchema, TableSchema, Schema

using Dates
using Parsers

import Base.parse  # For extending Base.parse to Base.parse(s::ColumnSchema, val)

using ..handle_validvalues

mutable struct ColumnSchema
    name::Symbol
    description::String
    datatype::DataType    # The non-missing type of each value in the column
    iscategorical::Bool   # Specifies whether the values represent categories. If so, order is specified by valueorder.
    isrequired::Bool      # Is non-missing data required?
    isunique::Bool        # Is each value in the column unique?
    validvalues::Union{DataType, <:AbstractRange, <:Set}             # Either the full range of the data type or a user-supplied restriction.
    valueorder::Union{DataType, <:AbstractRange, <:Vector, Nothing}  # If iscategorical, valueorder specifies the ordering of categories. Else nothing.
    parser::Function      # outputvalue = parser(inputvalue)

    function ColumnSchema(name, description, datatype, iscategorical, isrequired, isunique, validvalues, valueorder, parser)
        # Ensure eltyp and validvalues are consistent with each other
        tp_validvals = get_datatype(validvalues)
        datatype != tp_validvals && error("Column :$(name). Type of valid values ($(tp_validvals)) does not match that of eltype ($(datatype)).")
        new(Symbol(name), description, datatype, iscategorical, isrequired, isunique, validvalues, valueorder, parser)
    end
end

function ColumnSchema(name, description, datatype, iscategorical, isrequired, isunique, validvalues)
    valueorder  = iscategorical ? validvalues : nothing
    validvalues = validvalues isa Vector ? Set(validvalues) : validvalues
    parser      = constructparser(nothing, nothing, nothing, datatype)
    ColumnSchema(name, description, datatype, iscategorical, isrequired, isunique, validvalues, valueorder, parser)
end

function ColumnSchema(d::Dict)
    name          = d["name"]
    description   = d["description"]
    datatype      = d["datatype"] isa DataType ? d["datatype"] : eval(Meta.parse(d["datatype"]))
    iscategorical = d["iscategorical"]
    isrequired    = d["isrequired"]
    isunique      = d["isunique"]
    if haskey(d, "parser")
        func   = d["parser"]["function"] isa String ? eval(Meta.parse(d["parser"]["function"])) : d["parser"]["function"]
        args   = haskey(d["parser"], "args") ? d["parser"]["args"] : nothing
        kwargs = haskey(d["parser"], "kwargs") ? Dict{Symbol, Any}(Symbol(k) => v for (k,v) in d["parser"]["kwargs"]) : nothing
        parser = constructparser(func, args, kwargs, datatype)
    else
        parser = constructparser(nothing, nothing, nothing, datatype)
    end
    valueorder  = parse_validvalues(parser, datatype, d["validvalues"])
    validvalues = valueorder isa Vector ? Set(valueorder) : valueorder
    valueorder  = iscategorical ? valueorder : nothing
    ColumnSchema(name, description, datatype, iscategorical, isrequired, isunique, validvalues, valueorder, parser)
end

function constructparser(func, args, kwargs, returntype)
    # Special cases
    if ((func == Date) || (isnothing(func) && returntype == Date)) && !isnothing(args) && length(args) == 1
        df = DateFormat(args[1])
        return (x) -> try Date(x, df) catch e missing end
    end

    # General cases
    if func isa DataType || isnothing(func)
        opts = isnothing(kwargs) ? Parsers.Options() : Parsers.Options(kwargs...)
        function closure(val)
            len = val isa IO ? 0 : sizeof(val)  # Use default pos=1
            x, code, vpos, vlen, tlen = Parsers.xparse(returntype, val isa AbstractString ? codeunits(val) : val, 1, len, opts)
            Parsers.ok(code) ? x : missing
        end
        return closure
    end
    !(func isa Function) && error("Parser function does have type Function.")
    isnothing(args)  && isnothing(kwargs)  && return (x) -> try func(x)                     catch e missing end
    !isnothing(args) && isnothing(kwargs)  && return (x) -> try func(x, args...)            catch e missing end
    isnothing(args)  && !isnothing(kwargs) && return (x) -> try func(x; kwargs...)          catch e missing end
    !isnothing(args) && !isnothing(kwargs) && return (x) -> try func(x, args...; kwargs...) catch e missing end
    error("Invalid specification of the parser.")
end

################################################################################
struct TableSchema
    name::Symbol
    description::String
    colname2colschema::Dict{Symbol, ColumnSchema}
    columnorder::Vector{Symbol}  # Determines the order of the columns
    primarykey::Vector{Symbol}   # Vector of column names
    intrarow_constraints::Vector{Tuple{String, Function}}  # (msg, testfunc). Constraints between columns within a row (e.g., marriage date > birth date)

    function TableSchema(name, description, colname2colschema, columnorder, primarykey, intrarow_constraints=Function[])
        for colname in primarykey
            !haskey(colname2colschema, colname) && error("Table: $(name). Primary key has a non-existent column ($(colname)).")
            colschema = colname2colschema[colname]
            !colschema.isrequired    && error("Table: $(name). Primary key has a column ($(colname)) that allows missing data.")
        end
        if length(primarykey) == 1 && colname2colschema[primarykey[1]].isunique == false
            error("Table: $(name). Primary key must have isunique == true.")
        end
        new(name, description, colname2colschema, columnorder, primarykey, intrarow_constraints)
    end
end

function TableSchema(name, description, colschemata::Vector{ColumnSchema}, primarykey, intrarow_constraints=Function[])
    columnorder       = [colschema.name for colschema in colschemata]
    colname2colschema = Dict(colschema.name => colschema for colschema in colschemata)
    TableSchema(name, description, colname2colschema, columnorder, primarykey, intrarow_constraints)
end

function TableSchema(d::Dict)
    name        = Symbol(d["name"])
    description = d["description"]
    pk          = d["primarykey"]  # String or Vector{String}
    primarykey  = typeof(pk) == String ? [Symbol(pk)] : [Symbol(colname) for colname in pk]
    columns     = d["columns"]
    columnorder = fill(Symbol("x"), size(columns, 1))
    colname2colschema = Dict{Symbol, ColumnSchema}()
    i = 0
    for colname2schema in columns
        for (colname, colschema) in colname2schema
            i += 1
            columnorder[i]    = Symbol(colname)
            colschema["name"] = columnorder[i]
            colname2colschema[columnorder[i]] = ColumnSchema(colschema)
        end
    end
    intrarow_constraints = construct_intrarow_constraints(d)
    TableSchema(name, description, colname2colschema, columnorder, primarykey, intrarow_constraints)
end

function construct_intrarow_constraints(d::Dict)
    !haskey(d, "intrarow_constraints") && return Tuple{String, Function}[]
    d = d["intrarow_constraints"]
    n = length(d)
    result = Vector{Tuple{String, Function}}(undef, n)
    i = 0
    for (msg, func) in d
      s  = "(r) -> $(func)"
      f  = eval(Meta.parse(s))
      i += 1
      result[i] = (msg, f)
    end
    result
end

################################################################################
struct Schema
    name::Symbol
    description::String
    tables::Dict{Symbol, TableSchema}  # table_name => table_schema

    function Schema(name, description, tables)
        new(name, description, tables)
    end
end

function Schema(d::Dict)
    name = Symbol(d["name"])
    description = d["description"]
    tables = Dict{Symbol, TableSchema}()
    for (tblname, tblschema) in d["tables"]
        tblschema["name"]       = tblname
        tables[Symbol(tblname)] = TableSchema(tblschema)
    end
    Schema(name, description, tables)
end

end


text:
module NextGP

#exporting run-time equivalent of functions
export ran
export PR

#
export makeA

using DataFrames
using CategoricalArrays
using StatsModels
using MixedModels
using Printf

include("equations.jl")
include("runTime.jl")
include("MCMC.jl")
include("misc.jl")
include("outFiles.jl")

using .MCMC
export runGibbs
using .IO
export summaryMCMC

end


text:
"""
    greenspline(cmd0::String="", arg1=nothing; kwargs...)

Reads randomly-spaced (x,y,z) triples and produces a binary grid file of gridded values z(x,y) by solving:

		(1 - T) * L (L (z)) + T * L (z) = 0

Full option list at [`greenspline`]($(GMTdoc)greenspline.html)

Parameters
----------

- $(GMT.opt_R)
- **I** | **inc** :: [Type => Str | Number]

    *x_inc* [and optionally *y_inc*] is the grid spacing.
    ($(GMTdoc)greenspline.html#i)
- **A** | **gradient** :: [Type => Str | Array]		``Arg = gradfile+f1|2|3|4|5 | (data=Array, format=x)``

    The solution will partly be constrained by surface gradients v = v*n, where v is the
    gradient magnitude and n its unit vector direction.
    ($(GMTdoc)greenspline.html#a)
- **C** | **approx** | **approximate** :: [Type => Str | Number]	``Arg = [n]value[+ffile]``

    Find an approximate surface fit: Solve the linear system for the spline coefficients by
    SVD and eliminate the contribution from all eigenvalues whose ratio to the largest
    eigenvalue is less than value.
    ($(GMTdoc)greenspline.html#c)
- **G** | **grid** :: [Type => Str | []]

    Output grid file name. Note that this is optional and to be used only when saving
    the result directly on disk. Otherwise, just use the G = greenspline(....) form.
    ($(GMTdoc)greenspline.html#g)
- **D** | **mode** :: [Type => Number]

    Sets the distance flag that determines how we calculate distances between data points.
    ($(GMTdoc)greenspline.html#d)
- **E** :|**misfit** :: [Type => Str | []]		``Arg = [misfitfile]``

    Evaluate the spline exactly at the input data locations and report statistics of
    the misfit (mean, standard deviation, and rms).
    ($(GMTdoc)greenspline.html#e)
- **L** | **leave_trend** :: [Type => Bool]

    Do not remove a linear (1-D) or planer (2-D) trend when -D selects mode 0-3.
    ($(GMTdoc)greenspline.html#l)
- **N** | **nodes** :: [Type => Number | Array]			``Arg = nodefile``

    ASCII file with coordinates of desired output locations x in the first column(s).
    ($(GMTdoc)greenspline.html#n)
- **Q** | **dir_derivative** :: [Type => Str]		``Arg = az|x/y/z``

    Rather than evaluate the surface, take the directional derivative in the az azimuth and
    return the magnitude of this derivative instead.
    ($(GMTdoc)greenspline.html#q)
- **S** | **splines** :: [Type => Str]				``Arg = c|t|l|r|p|q[pars]``

    Select one of six different splines. The first two are used for 1-D, 2-D, or 3-D Cartesian splines.
    ($(GMTdoc)greenspline.html#s)
- **T** | **mask** :: [Type => Str]					``Arg = maskgrid``

    For 2-D interpolation only. Only evaluate the solution at the nodes in the maskgrid that are
    not equal to NaN.
    ($(GMTdoc)greenspline.html#t)
- $(GMT.opt_V)
- **W** | **uncertainties** :: [Type => Str | []]	``Arg = [w]``

	Data one-sigma uncertainties are provided in the last column. We then compute weights that
	are inversely proportional to the uncertainties squared.
    ($(GMTdoc)greenspline.html#w)
- $(GMT.opt_b)
- $(GMT.opt_d)
- $(GMT.opt_e)
- $(GMT.opt_f)
- $(GMT.opt_h)
- $(GMT.opt_i)
- $(GMT.opt_o)
- $(GMT.opt_r)
- $(GMT.opt_x)
- $(GMT.opt_w)
- $(GMT.opt_swap_xy)
"""
function greenspline(cmd0::String="", arg1=nothing; kwargs...)

	length(kwargs) == 0 && return monolitic("greenspline", cmd0, arg1)

	d = init_module(false, kwargs...)[1]		# Also checks if the user wants ONLY the HELP mode
	arg2 = nothing;     arg3 = nothing

	cmd, = parse_common_opts(d, "", [:I :R :V_params :bi :d :e :f :h :i :o :r :x :w :yx])
	cmd  = parse_these_opts(cmd, d, [[:C :approx :approximate], [:D :mode], [:E :misfit],
	                                 [:G :grid], [:L :leave_trend], [:Q :dir_derivative], [:S :splines], [:T :mask], [:W :uncertainties]])
	cmd, args, n, = add_opt(d, cmd, "A", [:A :gradient], :data, Array{Any,1}([arg1, arg2]), (format="+f",))
	if (n > 0)  arg1, arg2 = args[:]  end
	cmd, args, n, = add_opt(d, cmd, "N", [:N :nodes], :data, Array{Any,1}([arg1, arg2, arg3]), (x="",))
	if (n > 0)  arg1, arg2, arg3 = args[:]  end

	common_grd(d, cmd0, cmd, "greenspline ", arg1, arg2, arg3)		# Finish build cmd and run it
end

# ---------------------------------------------------------------------------------------------------
greenspline(arg1, cmd0::String=""; kw...) = greenspline(cmd0, arg1; kw...)


text:
using CairoMakie, LsqFit

#=
Use tanh to approximate the thermodynamic constraint
flux_forward = kinetics * (1 - exp(ΔG/RT))
flux_backward = kinetics * (exp(-ΔG/RT) - 1)
=#

xs = [
    range(-10, stop = -3, length = 1000)
    range(-3, stop = 3, length = 10000)
    range(3, stop = 10, length = 1000)
];

ys = [x <= 0 ? 1.0 - exp(x) : exp(-x) - 1.0 for x in xs]

@. fitfunc(x, p) = tanh(p[1] * x)
p0 = [1.0]

fit = curve_fit(fitfunc, xs, ys, p0)
vs = round.(coef(fit), digits = 1)

ysfit = fitfunc.(xs, Ref(vs));

fig = Figure()
ax1 = Axis(fig[1, 1], xlabel = "ΔG/RT", ylabel = "Thermodynamic factor")
lines!(ax1, xs, ys, label = "True")
lines!(ax1, xs, ysfit, label = "Approximation")
axislegend()

ax2 = Axis(fig[1, 2], xlabel = "ΔG/RT")
lines!(ax2, xs, ys)
lines!(ax2, xs, ysfit)
xlims!(ax2, -0.3, 0.3)
ylims!(ax2, -0.3, 0.3)

ax3 = Axis(fig[2, 1:2], xlabel = "ΔG/RT", ylabel = "% Error")
lines!(ax3, xs, (ys .- ysfit) ./ ys .* 100)

fig
CairoMakie.FileIO.save(joinpath("docs", "imgs", "tanh_approximation.pdf"), fig)


text:
#!/usr/bin/env julia

data = Array{Int}(undef, 2, 3)

for i in 1:size(data, 1), j in 1:size(data, 2)
    data[i, j] = (i - 1)*size(data, 2) + j
end

println(data)


text:
module Day25

using LinearAlgebra

inp = Set()
for line in eachline("day25.txt")
    push!(inp, line |> strip |> x -> split(x, ',') |> x -> parse.(Int, x))
end

mandist(a, b) = norm(a - b, 1)

function solvea(v)
    #incomplete
    consts = Set()

    function addstar(star)
        for constellation in consts, existstar in constellation
            if mandist(existstar, star) <= 3
                push!(constellation, star)
                return
            end
        end

        push!(consts, [star])
    end

    for star in sort(collect(v))
        addstar(star)
    end

    length(consts)
end

end

max_stars_repo_name: AnAverageHuman/competitive

max_issues_repo_name: AnAverageHuman/competitive

max_forks_repo_name: AnAverageHuman/competitive


text:
module Maxwell3D

    using ..BEAST
    Mod = BEAST

    """
        singlelayer(;gamma, alpha, beta)
        singlelayer(;wavenumber, alpha, beta)

    Bilinear form given by:

    ```math
        α ∬_{Γ×Γ} j(x)⋅k(y) G_{γ}(x,y) + β ∬_{Γ×Γ} div j(x) div k(y) G_{γ}(x,y)
    ```

    with ``G_{γ} = e^{-γ|x-y|} / 4π|x-y|``.
    """
    function singlelayer(;
            gamma=nothing,
            wavenumber=nothing,
            alpha=nothing,
            beta=nothing)

        if (gamma == nothing) && (wavenumber == nothing)
            error("Supply one of (not both) gamma or wavenumber")
        end

        if (gamma != nothing) && (wavenumber != nothing)
            error("Supply one of (not both) gamma or wavenumber")
        end

        if gamma == nothing
            if iszero(real(wavenumber))
                gamma = -imag(wavenumber)
            else
                gamma = im*wavenumber
            end
        end

        @assert gamma != nothing

        alpha == nothing && (alpha = -gamma)
        beta  == nothing && (beta  = -1/gamma)

        Mod.MWSingleLayer3D(gamma, alpha, beta)
    end

    """
        doublelayer(;gamma)
        doublelaher(;wavenumber)

    Bilinear form given by:

    ```math
        ∬_{Γ^2} k(x) ⋅ (∇G_γ(x-y) × j(y))
    ```

    with ``G_γ = e^{-γ|x-y|} / 4π|x-y|``
    """
    function doublelayer(;
            gamma=nothing,
            wavenumber=nothing)

        if (gamma == nothing) && (wavenumber == nothing)
            error("Supply one of (not both) gamma or wavenumber")
        end

        if (gamma != nothing) && (wavenumber != nothing)
            error("Supply one of (not both) gamma or wavenumber")
        end

        if gamma == nothing
            if iszero(real(wavenumber))
                gamma = -imag(wavenumber)
            else
                gamma = im*wavenumber
            end
        end

        @assert gamma != nothing

        Mod.MWDoubleLayer3D(gamma)
    end

    planewave(;
            direction    = error("missing arguement `direction`"),
            polarization = error("missing arguement `polarization`"),
            wavenumber   = error("missing arguement `wavenumber`"),
            amplitude    = one(real(typeof(wavenumber)))) =
        Mod.PlaneWaveMW(direction, polarization, wavenumber, amplitude)

    farfield(;
        wavenumber = error("missing argument: `wavenumber`")) =
            Mod.MWFarField3D(wavenumber=wavenumber)
end


text:
#const tagfield = collect(fieldnames(tag));

function cyclic_check(atag::tag, xtaglist::Array{tag})
  return any(x->x===atag, xtaglist)
  end;

function verifytags(atag::tag)

  taglist = tag[];
  tagstack = [atag];
  taglim = [size(atag.attach, 1)];
  tagptr = [0];
  tp_ptr = 1;
  itercount = 1;

  while tp_ptr>0
    local ctag = tagptr[tp_ptr];
    local xtag;
    local zval;
    local subtag =0;

    if(ctag == 0)
      xtag = tagstack[tp_ptr];
      ctag += 1;
      ctag > taglim[1] && return tag[];
      end;

    if(ctag <= taglim[tp_ptr])
      xtag = tagstack[tp_ptr].attach[ctag];
      ctag += 1;
      tagptr[tp_ptr] = ctag;
      subtag = size(xtag.attach,1);

      if(subtag > 0 && cyclic_check(xtag, tagstack))
        append!(taglist, [tagstack[tp_ptr]; xtag]);
        subtag = 0;
        end;

      if(subtag > 0)
        push!(tagstack, xtag);
        push!(taglim, subtag);
        push!(tagptr, 1);
        tp_ptr += 1;
        ctag = 1;
        end;
      end;

    if(ctag > taglim[tp_ptr])
      pop!(tagstack);#, xtag);
      pop!(taglim); #, subtag);
      pop!(tagptr); #, 1);
      tp_ptr -= 1;
      end;
    end;
  return taglist;
  end;


text:
"""
Tracks the specific informations and statistics we need to use Var-El to coordinate_action
the joint action in Factored-Value MCTS.

Fields:
    coord_graph_components::Vector{Vector{Int64}}
        The list of coordination graph components, i.e., cliques, where each element is a list of agent IDs that are in a mutual clique.

    min_degree_ordering::Vector{Int64}
        Ordering of agent IDs in increasing CG degree. This ordering is the heuristic most typically used for the elimination order in Var-El.

    n_component_stats::Dict{AbstractVector{S},Vector{Vector{Int64}}}
        Maps each joint state in the tree (for which we need to compute the UCB action) to the frequency of each component's various local actions.

    q_component_stats::Dict{AbstractVector{S},Vector{Vector{Float64}}}
        Maps each joint state in the tree to the accumulated q-value of each component's various local actions.
"""
mutable struct VarElStatistics{S} <: CoordinationStatistics
    coord_graph_components::Vector{Vector{Int64}}
    min_degree_ordering::Vector{Int64}
    n_component_stats::Dict{S,Vector{Vector{Int64}}}
    q_component_stats::Dict{S,Vector{Vector{Float64}}}
end

function clear_statistics!(ve_stats::VarElStatistics)
    empty!(ve_stats.n_component_stats)
    empty!(ve_stats.q_component_stats)
end

"""
Runs variable elimination at the current state using the VarEl Statistics to compute the best joint action with the component-wise exploration bonus.
FYI: Rather complicated.
"""
function coordinate_action(mdp::JointMDP{S,A}, tree::FVMCTSTree{S,A,VarElStatistics{S}}, s::S,
                           exploration_constant::Float64=0.0, node_id::Int64=0) where {S,A}

    n_agents = length(s)
    best_action_idxs = MVector{n_agents}([-1 for i in 1:n_agents])

    # !Note: Acquire lock so as to avoid race
    state_q_stats = lock(tree.lock) do
        tree.coordination_stats.q_component_stats[s]
    end
    state_n_stats = lock(tree.lock) do
        tree.coordination_stats.n_component_stats[s]
    end
    state_total_n = lock(tree.lock) do
        (node_id > 0) ? tree.total_n[node_id] : 0
    end

    # Maintain set of potential functions
    # NOTE: Hashing a vector here
    potential_fns = Dict{Vector{Int64},Vector{Float64}}()
    for (comp, q_stats) in zip(tree.coordination_stats.coord_graph_components, state_q_stats)
        potential_fns[comp] = q_stats
    end

    # Need this for reverse process
    # Maps agent to other elements in best response functions and corresponding set of actions
    # E.g. Agent 2 -> (3,4) in its best response and corresponding vector of agent 2 best actions
    best_response_fns = Dict{Int64,Tuple{Vector{Int64},Vector{Int64}}}()

    state_dep_actions = [agent_actions(mdp, i, si) for (i, si) in enumerate(s)]

    # Iterate over variable ordering
    # Need to maintain intermediate tables
    for ag_idx in tree.coordination_stats.min_degree_ordering

        # Lookup factors with agent in them and simultaneously construct
        # members of new potential function, and delete old factors
        agent_factors = Vector{Vector{Int64}}(undef, 0)
        new_potential_members = Vector{Int64}(undef, 0)
        for k in collect(keys(potential_fns))
            if ag_idx in k

                # Agent to-be-eliminated is in factor
                push!(agent_factors, k)

                # Construct key for new potential as union of all others except ag_idx
                for ag in k
                    if ag != ag_idx && ~(ag in new_potential_members)
                        push!(new_potential_members, ag)
                    end
                end
            end
        end

        if isempty(new_potential_members) == true
            # No out neighbors..either at beginning or end of ordering
            @assert agent_factors == [[ag_idx]] "agent_factors $(agent_factors) is not just [ag_idx] $([ag_idx])!"
            best_action_idxs[ag_idx] = _best_actionindex_empty(potential_fns,
                                                               state_dep_actions,
                                                               tree.all_agent_actions,
                                                               ag_idx)

        else

            # Generate new potential function and the best response vector for eliminated agent
            n_comp_actions = prod([length(tree.all_agent_actions[c]) for c in new_potential_members])

            # NOTE: Tuples should ALWAYS use tree.all_agent_actions for indexing
            comp_tup = Tuple(1:length(tree.all_agent_actions[c]) for c in new_potential_members)

            # Initialize q-stats for new potential and best response action vector
            # will be inserted into corresponding dictionaries at the end
            new_potential_stats = Vector{Float64}(undef, n_comp_actions)
            best_response_vect = Vector{Int64}(undef, n_comp_actions)

            # Iterate over new potential joint actions and compute new payoff and best response
            for comp_ac_idx = 1:n_comp_actions

                # Get joint action for other members in potential
                ct_idx = CartesianIndices(comp_tup)[comp_ac_idx]

                # For maximizing over agent actions
                # As before, we now need to init with -Inf
                ag_ac_values = zeros(length(tree.all_agent_actions[ag_idx]))

                # TODO: Agent actions should already be in order
                # Only do anything if action legal
                for (ag_ac_idx, ag_ac) in enumerate(tree.all_agent_actions[ag_idx])

                    if ag_ac in state_dep_actions[ag_idx]

                        # Need to look up corresponding stats from agent_factors
                        for factor in agent_factors

                            # NOTE: Need to reconcile the ORDER of ag_idx in factor
                            factor_action_idxs = MVector{length(factor),Int64}(undef)

                            for (idx, f) in enumerate(factor)

                                # if f is ag_idx, set corresponding factor action to ag_ac
                                if f == ag_idx
                                    factor_action_idxs[idx] = ag_ac_idx
                                else
                                    # Lookup index for corresp. agent action in ct_idx
                                    new_pot_idx = findfirst(isequal(f), new_potential_members)
                                    factor_action_idxs[idx] = ct_idx[new_pot_idx]
                                end # f == ag_idx
                            end

                            # NOW we can look up the stats of the factor
                            factor_tup = Tuple(1:length(tree.all_agent_actions[c]) for c in factor)
                            factor_action_linidx = LinearIndices(factor_tup)[factor_action_idxs...]

                            ag_ac_values[ag_ac_idx] += potential_fns[factor][factor_action_linidx]

                            # Additionally add exploration stats if factor in original set
                            factor_comp_idx = findfirst(isequal(factor), tree.coordination_stats.coord_graph_components)
                            if state_total_n > 0 && ~(isnothing(factor_comp_idx)) # NOTE: Julia1.1
                                ag_ac_values[ag_ac_idx] += exploration_constant * sqrt((log(state_total_n+1.0))/(state_n_stats[factor_comp_idx][factor_action_linidx]+1.0))
                            end
                        end # factor in agent_factors
                    else
                        ag_ac_values[ag_ac_idx] = -Inf
                    end # ag_ac in state_dep_actions
                end # ag_ac_idx = 1:length(tree.all_agent_actions[ag_idx])

                # Now we lookup ag_ac_values for the best value to be put in new_potential_stats
                # and the best index to be put in best_response_vect
                # NOTE: The -Inf mask should ensure only legal idxs chosen
                # If all ag_ac_values equal, should we sample randomly?
                best_val, best_idx = findmax(ag_ac_values)

                new_potential_stats[comp_ac_idx] = best_val
                best_response_vect[comp_ac_idx] = best_idx
            end # comp_ac_idx in n_comp_actions

            # Finally, we enter new stats vector and best response vector back to dicts
            potential_fns[new_potential_members] = new_potential_stats
            best_response_fns[ag_idx] = (new_potential_members, best_response_vect)
        end # isempty(new_potential_members)

        # Delete keys in agent_factors from potential fns since variable has been eliminated
        for factor in agent_factors
            delete!(potential_fns, factor)
        end
    end # ag_idx in min_deg_ordering

    # NOTE: At this point, best_action_idxs has at least one entry...for the last action obtained
    @assert !all(isequal(-1), best_action_idxs) "best_action_idxs is still undefined!"

    # Message passing in reverse order to recover best action
    for ag_idx in Base.Iterators.reverse(tree.coordination_stats.min_degree_ordering)

        # Only do something if best action already not obtained
        if best_action_idxs[ag_idx] == -1

            # Should just be able to lookup best response function
            (agents, best_response_vect) = best_response_fns[ag_idx]

            # Members of agents should already have their best action defined
            agent_ac_tup = Tuple(1:length(tree.all_agent_actions[c]) for c in agents)
            best_agents_action_idxs = [best_action_idxs[ag] for ag in agents]
            best_response_idx = LinearIndices(agent_ac_tup)[best_agents_action_idxs...]

            # Assign best action for ag_idx
            best_action_idxs[ag_idx] = best_response_vect[best_response_idx]
        end # isdefined
    end

    # Finally, return best action by iterating over best action indices
    # NOTE: best_action should use state-dep actions to reverse index
    best_action = [tree.all_agent_actions[ag][idx] for (ag, idx) in enumerate(best_action_idxs)]

    return best_action
end

"""
Take the q-value from the MCTS step and distribute the updates across the component q-stats as per the formula in the Amato-Oliehoek paper.
"""
function update_statistics!(mdp::JointMDP{S,A}, tree::FVMCTSTree{S,A,VarElStatistics{S}},
                            s::S, ucb_action::A, q::AbstractVector{Float64}) where {S,A}

    n_agents = length(s)

    for (idx, comp) in enumerate(tree.coordination_stats.coord_graph_components)

        # Create cartesian index tuple
        comp_tup = Tuple(1:length(tree.all_agent_actions[c]) for c in comp)

        # RECOVER local action corresp. to ucb action
        # TODO: Review this carefully. Need @req for action index for agent.
        local_action = [ucb_action[c] for c in comp]
        local_action_idxs = [agent_actionindex(mdp, c, a) for (a, c) in zip(local_action, comp)]

        comp_ac_idx = LinearIndices(comp_tup)[local_action_idxs...]

        # NOTE: NOW we can update stats. Could generalize incremental update more here
        lock(tree.lock) do
            tree.coordination_stats.n_component_stats[s][idx][comp_ac_idx] += 1
            q_comp_value = sum(q[c] for c in comp)
            tree.coordination_stats.q_component_stats[s][idx][comp_ac_idx] +=
                (q_comp_value - tree.coordination_stats.q_component_stats[s][idx][comp_ac_idx]) / tree.coordination_stats.n_component_stats[s][idx][comp_ac_idx]
        end
    end
end

# TODO: is this the correct thing to do?
# My guess is no, but not sure.
function update_statistics!(mdp::JointMDP{S,A}, tree::FVMCTSTree{S,A,VarElStatistics{S}},
    s::S, ucb_action::A, q::Float64) where {S, A}
    n_agents = length(s)

    for (idx, comp) in enumerate(tree.coordination_stats.coord_graph_components)

        # Create cartesian index tuple
        comp_tup = Tuple(1:length(tree.all_agent_actions[c]) for c in comp)

        # RECOVER local action corresp. to ucb action
        # TODO: Review this carefully. Need @req for action index for agent.
        local_action = [ucb_action[c] for c in comp]
        local_action_idxs = [agent_actionindex(mdp, c, a) for (a, c) in zip(local_action, comp)]

        comp_ac_idx = LinearIndices(comp_tup)[local_action_idxs...]

        # NOTE: NOW we can update stats. Could generalize incremental update more here
        lock(tree.lock) do
            tree.coordination_stats.n_component_stats[s][idx][comp_ac_idx] += 1
            q_comp_value = q * length(comp) # Maintains equivalence with `sum(q[c] for c in comp)`
            tree.coordination_stats.q_component_stats[s][idx][comp_ac_idx] +=
                (q_comp_value - tree.coordination_stats.q_component_stats[s][idx][comp_ac_idx]) / tree.coordination_stats.n_component_stats[s][idx][comp_ac_idx]
        end
    end
end

function init_statistics!(tree::FVMCTSTree{S,A,VarElStatistics{S}}, planner::FVMCTSPlanner,
                          s::S) where {S,A}

    n_comps = length(tree.coordination_stats.coord_graph_components)
    n_component_stats = Vector{Vector{Int64}}(undef, n_comps)
    q_component_stats = Vector{Vector{Float64}}(undef, n_comps)

    n_agents = length(s)

    # TODO: Could actually make actions state-dependent if need be
    for (idx, comp) in enumerate(tree.coordination_stats.coord_graph_components)

        n_comp_actions = prod([length(tree.all_agent_actions[c]) for c in comp])

        n_component_stats[idx] = Vector{Int64}(undef, n_comp_actions)
        q_component_stats[idx] = Vector{Float64}(undef, n_comp_actions)

        comp_tup = Tuple(1:length(tree.all_agent_actions[c]) for c in comp)

        for comp_ac_idx = 1:n_comp_actions

            # Generate action subcomponent and call init_Q and init_N for it
            ct_idx = CartesianIndices(comp_tup)[comp_ac_idx] # Tuple corresp to
            local_action = [tree.all_agent_actions[c][ai] for (c, ai) in zip(comp, Tuple(ct_idx))]

            # NOTE: init_N and init_Q are functions of component AND local action
            # TODO(jkg): init_N and init_Q need to be defined
            n_component_stats[idx][comp_ac_idx] = init_N(planner.solver.init_N, planner.mdp, s, comp, local_action)
            q_component_stats[idx][comp_ac_idx] = init_Q(planner.solver.init_Q, planner.mdp, s, comp, local_action)
        end
    end

    # Update tree member
    lock(tree.lock) do
        tree.coordination_stats.n_component_stats[s] = n_component_stats
        tree.coordination_stats.q_component_stats[s] = q_component_stats
    end
end

@inline function _best_actionindex_empty(potential_fns, state_dep_actions, all_agent_actions, ag_idx)
    # NOTE: This is inefficient but necessary for state-dep actions?
    if length(state_dep_actions[ag_idx]) == length(all_agent_actions[ag_idx])
        _, best_ac_idx = findmax(potential_fns[[ag_idx]])
    else
        # Now we need to choose the best index from among legal actions
        # Create an array with illegal actions having -Inf and then fill legal vals
        # TODO: More efficient way to do this?
        masked_action_vals = fill(-Inf, length(all_agent_actions[ag_idx]))
        for (iac, ac) in enumerate(all_agent_actions[ag_idx])
            if ac in state_dep_actions[ag_idx]
                masked_action_vals[iac] = potential_fns[[ag_idx]][iac]
            end
        end
        _, best_ac_idx = findmax(masked_action_vals)
    end
    return best_ac_idx
end


text:
using CRCBS
using Documenter

makedocs(;
    modules=[CRCBS,CRCBS.BenchmarkInterface],
    authors="kylebrown <kylejbrown17@gmail.com> and contributors",
    repo="https://github.com/kylejbrown17/CRCBS.jl/blob/{commit}{path}#L{line}",
    sitename="CRCBS.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://kylejbrown17.github.io/CRCBS.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "Getting Started" => "getting_started.md",
        "Core Types and Methods" => "library.md",
        "API Reference" => "reference.md",
    ],
)

deploydocs(;
    repo="github.com/kylejbrown17/CRCBS.jl",
)


text:
"""
spin26.jl - This is a julia implementation of backwards diff
GRAPE with projected gradient descent via Adam.
"""

WDIR = get(ENV, "ROBUST_QOC_PATH", "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using ForwardDiff
using HDF5
using LinearAlgebra
using ReverseDiff
using StaticArrays
using Statistics
using Zygote

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin26"
const SAVE_PATH = abspath(joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME))

### PADE DUE TO HIGHAM 2005 ###

# Pade approximants from algorithm 2.3.
const B = (
    64764752532480000,
    32382376266240000,
    7771770303897600,
    1187353796428800,
    129060195264000,
    10559470521600,
    670442572800,
    33522128640,
    1323241920,
    40840800,
    960960,
    16380,
    182,
    1,
)

function one_norm(a)
    """
    Return the one-norm of the matrix.
    References:
    [0] https://www.mathworks.com/help/dsp/ref/matrix1norm.html
    Arguments:
    a :: ndarray(N x N) - The matrix to compute the one norm of.

    Returns:
    one_norm_a :: float - The one norm of a.
    """
    return maximum(sum(map(abs, a), dims=1))
end

function pade3(a, i)
    a2 = a * a
    u = a * (B[3] * a2) + B[2] * a
    v = B[3] * a2 + B[1] * i
    return u, v
end

function pade5(a, i)
    a2 = a * a
    a4 = a2 * a2
    u = a * (B[6] * a4 + B[4] * a2) + B[2] * a
    v = B[5] * a4 + B[3] * a2 + B[1] * i
    return u, v
end

function pade7(a, i)
    a2 = a * a
    a4 = a2 * a2
    a6 = a2 * a4
    u = a * (B[8] * a6 + B[6] * a4 + B[4] * a2) + B[2] * a
    v = B[7] * a6 + B[5] * a4 + B[3] * a2 + B[1] * i
    return u, v
end

function pade9(a, i)
    a2 = a * a
    a4 = a2 * a2
    a6 = a2 * a4
    a8 = a2 * a6
    u = a * (B[10] * a8 + B[8] * a6 + B[6] * a4 + B[4] * a2) + B[2] * a
    v = B[9] * a8 + B[7] * a6 + B[5] * a4 + B[3] * a2 + B[1] * i
    return u, v
end

function pade13(a, i)
    a2 = a * a
    a4 = a2 * a2
    a6 = a2 * a4
    u = a * (a6 * (B[14] * a6 + B[12] * a4 + B[10] * a2) + B[8] * a6 + B[6] * a4 + B[4] * a2) + B[2] * a
    v = a6* (B[13] * a6 + B[11] * a4 + B[9] * a2) + B[7] * a6 + B[5] * a4 + B[3] * a2 + B[1] * i
    return u, v
end

# Valid pade orders for algorithm 2.3.
const PADE_ORDERS = (
    3,
    5,
    7,
    9,
    13,
)

# Pade approximation functions.
const PADE = [
    nothing,
    nothing,
    nothing,
    pade3,
    nothing,
    pade5,
    nothing,
    pade7,
    nothing,
    pade9,
    nothing,
    nothing,
    nothing,
    pade13,
]

# Constants taken from table 2.3.
const THETA = (
    0,
    0,
    0,
    1.495585217958292e-2,
    0,
    2.539398330063230e-1,
    0,
    9.504178996162932e-1,
    0,
    2.097847961257068,
    0,
    0,
    0,
    5.371920351148152,
)

function expm_pade(a)
    """
    Compute the matrix exponential via pade approximation.
    References:
    [0] http://eprints.ma.man.ac.uk/634/1/high05e.pdf
    [1] https://github.com/scipy/scipy/blob/v0.14.0/scipy/linalg/_expm_frechet.py#L10
    [2] https://github.com/SchusterLab/qoc/blob/dev-robust/qoc/standard/functions/expm.py#L212
    Arguments:
    a :: ndarray(N x N) - The matrix to exponentiate.

    Returns:
    expm_a :: ndarray(N x N) - The exponential of a.
    """
    # If the one norm is sufficiently small,
    # pade orders up to 13 are well behaved.
    scale = 0
    size_ = size(a, 1)
    pade_order = nothing
    one_norm_ = one_norm(a)
    for pade_order_ in PADE_ORDERS
        if one_norm_ < THETA[pade_order_]
            pade_order = pade_order_
        end
    end

    # If the one norm is large, scaling and squaring
    # is required.
    if isnothing(pade_order)
        pade_order = 14
        scale = max(0, Int(ceil(log2(one_norm_ / THETA[14]))))
        a = a * 2^(-scale)
    end

    # Execute pade approximant.
    i = I(size_)
    u, v = PADE[pade_order](a, i)
    r = (u + v)/(-u + v)

    # Do squaring if necessary.
    for i = 1:scale
        r = r * r
    end

    return r
end

"""
gram_schmidt

Arguments:
    basis::Matrix - cols are basis vectors
Returns:
    basis_ ::Vector{Vector} - elements are
    orthogonal basis vectors
References:
[0] http://vmls-book.stanford.edu/vmls-julia-companion.pdf
"""
function gram_schmidt(basis)
    basis_ = Vector{Vector}(undef, size(basis, 2))
    for i = 1:size(basis, 2)
        q = basis[:, i]
        for j = 1:i - 1
            q -= basis_[j]'basis[:, i] * basis_[j]
        end
        basis_[i] = q / norm(q)
    end
    return basis_
end

function adam_update(params::Vector, grads::Vector, iter::Int, learning_rate::Float64,
                     gradient_moment::Vector, gradient_square_moment::Vector;
                     eps=1e-8, β1=0.9, β2=0.999)
    gradient_moment .= β1 * gradient_moment + (1 - β1) * grads
    gradient_square_moment .= β2 * gradient_square_moment + (1 - β2) * grads.^2
    gradient_moment_hat = gradient_moment / (1 - β1^iter)
    gradient_square_moment_hat = gradient_square_moment / (1 - β2^iter)
    params .= params .- learning_rate .* gradient_moment ./ (map(sqrt, gradient_square_moment) .+ eps)
end

"""
znf_basis - construct an orthonormal basis for the tangent space where all elements
of a vector sum to zero
"""
function znf_basis(knot_count)
    basis = zeros(knot_count, knot_count - 1)
    for i = 2:knot_count - 1
        basis[i, i - 1] = -1
    end
    basis[end, end] = -1
    basis[1, :] .= 1
    basis = gram_schmidt(basis)
    return basis
end

"""
project - project a vector onto an orthonormal basis
"""
function project(vec::Vector, basis::Vector{<:Vector})
    vec_ = zeros(size(vec, 1))
    for i = 1:size(basis, 1)
        vec_ = vec_ + vec'basis[i] * basis[i]
    end
    return vec_
end

"""
rollout - integrate the Schrodinger equation like it's hot
"""
function rollout(controls::AbstractVector, state1::SVector{HDIM_ISO},
                 state2::SVector{HDIM_ISO}, dt::Real, knot_count::Int64)
    for i = 1:knot_count
        h_prop = expm_pade(dt .* (FQ_NEGI_H0_ISO + (controls[i] .* NEGI_H1_ISO)))
        state1 = h_prop * state1
        state2 = h_prop * state2
    end

    return (state1, state2)
end

function objective(controls::AbstractVector, state1::SVector{HDIM_ISO},
                   state2::SVector{HDIM_ISO}, target1::SVector{HDIM_ISO}, target2::SVector{HDIM_ISO},
                   qs::SVector{3}, dt::Real, knot_count::Int64)
    # fidelity
    (fstate1, fstate2) = rollout(controls, state1, state2, dt, knot_count)
    d1 = fstate1 - target1
    d2 = fstate2 - target2
    state_cost = qs[1] * (d1'd1 + d2'd2)

    # control variation
    dcontrols_diff = diff(controls)
    dcontrols_cost = qs[2] * dcontrols_diff'dcontrols_diff
    d2controls_diff = diff(dcontrols_diff)
    d2controls_cost = qs[3] * d2controls_diff'd2controls_diff

    return state_cost + dcontrols_cost + d2controls_cost
end

@inline max_violation(state1::SVector{HDIM_ISO}, state2::SVector{HDIM_ISO},
                      target1::SVector{HDIM_ISO}, target2::SVector{HDIM_ISO}) = (
    maximum((1 - fidelity_vec_iso2(state1, target1), 1 - fidelity_vec_iso2(state2, target2)))
)

"""
run_pgrape

Arguments:
qs :: Array - (ψ, ∂a, ∂2a)
"""
function run_pgrape(;gate_type=zpiby2, evolution_time=18.,
                    ctol=1e-8, dt=1e-1, qs=SVector{3}([1e0, 1e-1, 1e-1]),
                    learning_rate=1e-2, max_iterations=Int64(1e3),
                    min_control=-0.5, max_control=0.5)
    # setup
    knot_count = Int(div(evolution_time, dt))
    controls = fill(1e-2, knot_count)
    znf_basis_ = znf_basis(knot_count)

    # initial state
    state1 = SVector{HDIM_ISO}(IS1_ISO_)
    state2 = SVector{HDIM_ISO}(IS2_ISO_)

    # final state
    if gate_type == xpiby2
        target1 = XPIBY2_ISO_1
        target2 = XPIBY2_ISO_2
    elseif gate_type == zpiby2
        target1 = ZPIBY2_ISO_1
        target2 = ZPIBY2_ISO_2
    end

    # initialize optimization
    max_viol = 1
    objective_(controls_) = objective(controls_, state1, state2,
                                      target1, target2, qs, dt, knot_count)
    grad_controls = zeros(knot_count)
    gradient_moment = zeros(knot_count)
    gradient_square_moment = zeros(knot_count)
    # perform optimization
    for i = 1:max_iterations
        # check if converged
        (fs1, fs2) = rollout(controls, state1, state2, dt, knot_count)
        max_viol = max_violation(fs1, fs2, target1, target2)
        println("max_viol: $(max_viol)")
        if max_viol < ctol
            return (fs1, fs2)
            break
        end
        # compute gradient
        (grad_controls,) = Zygote.gradient(objective_, controls)
        # update controls
        adam_update(controls, grad_controls, i, learning_rate,
                    gradient_moment, gradient_square_moment)
        # project controls onto zero net flux constraint manifold
        controls = project(controls, znf_basis_)
        # ensure controls satisfy bound constraints
        controls[1] = controls[end] = 0.
        for j = 2:knot_count - 1
            controls[j] = clamp(controls[j], min_control, max_control)
        end
    end
end

max_stars_repo_name: SchusterLab/rbqoc

max_issues_repo_name: SchusterLab/rbqoc

max_forks_repo_name: SchusterLab/rbqoc