text:
export
  generateCanonicalFG_Helix2D!,
  generateCanonicalFG_Helix2DSlew!,
  generateCanonicalFG_Helix2DSpiral!

"""
    $SIGNATURES

Generalized canonical graph generator function for helix patterns.

Notes
- assumes poses are labeled according to r"x\\d+"
- Gradient (i.e. angle) calculations are on the order of 1e-8.
- Use callback `spine_t(t)::Complex` to modify how the helix pattern is moved in x, y along the progression of `t`,
  - See related wrapper functions for convenient generators of helix patterns in 2D,
  - Real valued `xr_t(t)` and `yr_t(t)` can be modified (and will override) complex valued `spine_t` instead.
- use `postpose_cb = (fg_, lastestpose) -> ...` for additional user features after each new pose
- can be used to grow a graph with repeated calls, but keyword parameters are assumed identical between calls.

See also: [`generateCanonicalFG_Helix2DSlew!`](@ref), [`generateCanonicalFG_Helix2DSpiral!`](@ref), [`generateCanonicalFG_Beehive!`](@ref)
"""
function generateCanonicalFG_Helix2D!(numposes::Integer = 40;
                                      posesperturn::Integer = 20,
                                      graphinit = nothing,
                                      useMsgLikelihoods = nothing,
                                      solverParams::SolverParams = SolverParams(;graphinit=false),
                                      dfg::AbstractDFG = LightDFG{SolverParams}(;solverParams),
                                      radius::Real = 10,
                                      spine_t = (t)->0 + im*0,
                                      xr_t::Function = (t)->real(spine_t(t)),
                                      yr_t::Function = (t)->imag(spine_t(t)),
                                      poseRegex::Regex = r"x\d+",
                                      μ0 = [0;0;pi/2],
                                      refKey::Symbol = :simulated,
                                      Qd::AbstractMatrix{<:Real} = diagm( [0.1;0.1;0.05].^2 ),
                                      postpose_cb::Function = (fg_,latestpose)->()   )
  #
  (graphinit isa Nothing) ? nothing : @error("generateCanonicalFG_Helix2D! keyword graphinit obsolete, use solverParams=SolverParams(graphinit=..) instead.")
  (useMsgLikelihoods isa Nothing) ? nothing : @error("generateCanonicalFG_Helix2D! keyword useMsgLikelihoods obsolete, use solverParams=SolverParams(useMsgLikelihoods=..) instead.")

  # add first pose if not already exists
  _initpose = Symbol(match(r"[A-Za-z]+", poseRegex.pattern).match, 0)
  if !exists( dfg, _initpose )
    generateCanonicalFG_ZeroPose(;dfg, μ0, solverParams, postpose_cb) # , μ0=[0;0;1e-5] # tried for fix NLsolve on wrap issue
    # getSolverParams(dfg).useMsgLikelihoods = useMsgLikelihoods
    # reference ppe on :x0
    ppe = DFG.MeanMaxPPE(refKey, μ0, μ0, μ0)
    setPPE!(dfg[:x0], refKey, DFG.MeanMaxPPE, ppe)
  end

  # start from existsing poses
  _poses = ls(dfg, poseRegex) |> sortDFG
  # what is the last pose and posecount number
  lastpose = _poses[end]
  posecount = match(r"\d+", string(lastpose)).match |> x->parse(Int,x)
  # init how many poses at the beginning of a new turn
  bidx = length(_poses) # 1

  # fractional number of turns to make in total (after all graph generation is done)
  turns = numposes/posesperturn
  # generate helix pattern algebraically
  tmp = calcHelix_T(0, turns, posesperturn, radius=radius, spine_t=spine_t, xr_t=xr_t, yr_t=yr_t)
  # TODO, dont always start from 0 -- i.e. chop first repeat elements from deterministic helix

  # select the starting point
  _μ0 = μ0
  # @show _μ0 = 1 == bidx ? μ0 : getPPE(dfg, lastpose, refKey).suggested
  Tμ = SE2(_μ0-[0;0;pi/2])

  # current end pose count for number of turns
  eidx = 1
  for tn in 0:(ceil(Int, turns)-1)
    eidx += posesperturn
    # skip out early if extending a previous existing graph
    eidx = minimum( [eidx, length(tmp[1])] )
    # tmp_ = _calcHelix2DApprox(N_ppt=posesperturn, radius=radius, runback=runback)
    tmp_ = hcat(tmp[2][bidx:eidx,:],tmp[3][bidx:eidx])'
    # adjust for turn progression in x
    # tmp_[1,:] .+= tn*(2radius*(1-runback))
    oldpose = Tμ*SE2(tmp_[:,1])
    eidx < bidx ? continue : nothing

    # add each new pose (skippin the first element)
    for ps in 2:size(tmp_,2)
      # check exit condition
      numposes-1 <= posecount && break
      # add a new pose
      newpose = Tμ*TU.SE2(tmp_[:,ps])
      deltaodo = se2vee(oldpose \ newpose)
      factor = Pose2Pose2( MvNormal(deltaodo, Qd) )
      posecount += 1
      v_n = _addPoseCanonical!(dfg, lastpose, posecount, factor, poseRegex=poseRegex, refKey=refKey, overridePPE=se2vee(newpose), postpose_cb=postpose_cb)
      lastpose = getLabel(v_n)
      oldpose = newpose
    end

    bidx = eidx
  end

  #
  return dfg
end

"""
    $SIGNATURES

Generate canonical slewed helix graph (like a flattened slinky).

Notes
- Use `slew_x` and `slew_y` to pull the "slinky" out in different directions at constant rate.
- See generalized helix generator for more details.
- Defaults are choosen to slew along x and have multple trajectory intersects between consecutive loops of the helix.

Related

[`generateCanonicalFG_Helix2D!`](@ref), [`generateCanonicalFG_Helix2DSpiral!`](@ref)
"""
generateCanonicalFG_Helix2DSlew!( numposes::Integer=40;
                                  slew_x::Real=2/3,
                                  slew_y::Real=0,
                                  spine_t=(t)->slew_x*t + im*slew_y*t,
                                  kwargs...  ) = generateCanonicalFG_Helix2D!(numposes; spine_t=spine_t, kwargs...)
#

"""
    $SIGNATURES

Generate canonical helix graph that expands along a spiral pattern, analogous flower petals.

Notes
- This function wraps the complex `spine_t(t)` function to generate the spiral pattern.
  - `rate_a` and `rate_r` can be varied for different spiral behavior.
- See generalized helix generator for more details.
- Defaults are choosen to slewto have multple trajectory intersects between consecutive loops of the helix and do a decent job of moving around coverage area with a relative balance of encircled area sizes.

Related

[`generateCanonicalFG_Helix2D!`](@ref), [`generateCanonicalFG_Helix2DSlew!`](@ref)
"""
generateCanonicalFG_Helix2DSpiral!( numposes::Integer=100;
                                    rate_r=0.6,
                                    rate_a=6,
                                    spine_t=(t)->rate_r*(t^0.5)*cis(rate_a*(t^0.4)),
                                    kwargs...  ) = generateCanonicalFG_Helix2D!(numposes; spine_t=spine_t, kwargs...)
#

#


text:
# script to parse oneAPI headers and generate Julia wrappers

#
# Parsing
#

using Clang.Generators
using JuliaFormatter

function wrap(name, headers...; library="lib$name", defines=[], include_dirs=[])
    options =  load_options(joinpath(@__DIR__, "wrap.toml"))
    options["general"]["library_name"] = library
    options["general"]["output_file_path"] = "lib$(name).jl"

    args = get_default_args()
    ctx = create_context([headers...], args, options)

    build!(ctx, BUILDSTAGE_NO_PRINTING)

    # tweak exprs
    for node in get_nodes(ctx.dag)
        exprs = get_exprs(node)
        for (i, expr) in enumerate(exprs)
            exprs[i] = expr |> change_argument_types |> add_check_pass
        end
    end

    build!(ctx, BUILDSTAGE_PRINTING_ONLY)

    return options["general"]["output_file_path"]
end

## rewrite passes

# insert `@checked` before each function with a `ccall` returning a checked type`
checked_types = [
    "ze_result_t",
]
function add_check_pass(x::Expr)
    Meta.isexpr(x, :function) || return x
    body = x.args[2].args[1]
    @assert Meta.isexpr(body, :macrocall) # `@ccall`
    ret_type = string(body.args[3].args[2])
    if ret_type in checked_types
        return Expr(:macrocall, Symbol("@checked"), nothing, x)
    end
end

# change certain Ptr inputs to ZePtr
argument_types = Dict(
    :zeCommandListAppendMemoryCopy => Dict(
        :dstptr => :(PtrOrZePtr{Cvoid}),
        :srcptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeCommandListAppendMemoryFill => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
        :pattern => :(PtrOrZePtr{Cvoid}),
    ),
    :zeCommandListAppendMemoryCopyRegion => Dict(
        :dstptr => :(PtrOrZePtr{Cvoid}),
        :srcptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeCommandListAppendMemoryCopyFromContext => Dict(
        :dstptr => :(PtrOrZePtr{Cvoid}),
        :srcptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeCommandListAppendMemoryPrefetch => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeCommandListAppendMemAdvise => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeMemFree => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeMemFreeExt => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeMemGetAllocProperties => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeMemGetAddressRange => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeMemGetIpcHandle => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeMemOpenIpcHandle => Dict(
        :pptr => :(Ptr{PtrOrZePtr{Cvoid}}),
    ),
    :zeMemCloseIpcHandle => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeContextEvictMemory => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeContextMakeMemoryResident => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
    :zeVirtualMemFree => Dict(
        :ptr => :(PtrOrZePtr{Cvoid}),
    ),
)
function change_argument_types(x::Expr)
    global argument_types
    Meta.isexpr(x, :function) || return x
    body = x.args[2].args[1]
    @assert Meta.isexpr(body, :macrocall) # `@ccall`
    the_ccall = body.args[3]
    @assert Meta.isexpr(the_ccall, :(::))
    call, ret = the_ccall.args
    @assert Meta.isexpr(call, :call)
    target, argexprs... = call.args
    @assert Meta.isexpr(target, :(.))
    lib, f = target.args
    f = f.value
    for argexpr in argexprs
        @assert Meta.isexpr(argexpr, :(::))
        arg, argtyp = argexpr.args
        if haskey(argument_types, f) && haskey(argument_types[f], arg)
            new_argtyp = argument_types[f][arg]
            @info "Changing argument $f($arg::$argtyp) to $new_argtyp"
            argexpr.args[2] = new_argtyp
        end
    end
    return x
end

#
# Main application
#

using oneAPI_Level_Zero_Headers_jll

function process(name, headers...; modname=name, kwargs...)
    output_file = wrap(name, headers...; kwargs...)

    let file = output_file
        text = read(file, String)

        ## header

        squeezed = replace(text, "\n\n\n"=>"\n\n")
        while length(text) != length(squeezed)
            text = squeezed
            squeezed = replace(text, "\n\n\n"=>"\n\n")
        end
        text = squeezed

        write(file, text)
    end

    ## manual patches

    patchdir = joinpath(@__DIR__, "patches", name)
    if isdir(patchdir)
        for entry in readdir(patchdir)
            if endswith(entry, ".patch")
                path = joinpath(patchdir, entry)
                run(`patch -p1 -i $path`)
            end
        end
    end

    ## move to destination and format

    let src = output_file
        dst = joinpath(dirname(@__DIR__), "lib", modname, src)
        cp(src, dst; force=true)
        format(dst, YASStyle(), always_use_return=false)
    end

    return
end

function main()
    process("ze", oneAPI_Level_Zero_Headers_jll.ze_api;
            library="libze_loader", modname="level-zero")
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end

ext: jl


text:
using GameZero
game_include("PointQuad.jl")

const WIDTH = 800
const HEIGHT = 800
const BACKGROUND = colorant"black"

mutable struct P
    x::Real
    y::Real
end

const qt = PointQuad.QuadTree{P}(PointQuad.Sqr(WIDTH / 2, HEIGHT / 2, WIDTH), 4)

PointQuad.position(p::P) = (trunc(Int, p.x), trunc(Int, p.y))

function init()
    for _ in 1:1200
        PointQuad.insert!(qt, P(rand(1:WIDTH - 1), rand(1:HEIGHT - 1)))
    end
end

init()

function draw()
    for p in PointQuad.get_points(qt)
        p.x += rand(-1:1)
        p.y += rand(-1:1)
        if length(PointQuad.query(qt, PointQuad.Sqr(p.x, p.y, 16))) > 1
            draw(Circle(trunc(Int, p.x), trunc(Int, p.y), 4), colorant"green", fill = true)
        else
            draw(Circle(trunc(Int, p.x), trunc(Int, p.y), 4), colorant"white", fill = true)
        end
    end
    agents = PointQuad.get_points(qt)
    PointQuad.clear!(qt)
    for lol in agents
        PointQuad.insert!(qt, lol)
    end
    agents = nothing
end


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

module LittleEndianBase128

export encode, decodesigned, decodeunsigned, decode

const version = v"0.1.1"

movesign(x::T) where T = signed(xor(x >>> one(T), -(x & one(T))))

encode(input::Array{Bool,N}) where N = encode(Array(reinterpret(UInt8, input)))

function encode(input::Array{T,N}) where T<:Unsigned where N
  # Encode array of unsigned integers using LEB128
  maxbytes = ceil(Int, 8*sizeof(T)/ 7)
  output = Array{UInt8}(undef, maxbytes*length(input))  # maximum possible length
  k = 1  # position of next write in output array
  for j in 1:length(input)
    if input[j] == 0
      output[k] = zero(T)
      k += 1
    else
      x = input[j]
      while x != 0
        output[k] = x & 0x7F | 0x80  # can simplify?
        k += 1
        x >>= 7
      end
      output[k-1] &= 0x7F
    end
  end
  return output[1:k-1]
end

encode(n::T) where T<:Unsigned = encode([n])

encode(n::Bool) = encode([n])

encode(n::T) where T<:Signed = encode(unsigned(xor(n << 1, n >> (8*sizeof(T)-1))))

encode(input::Array{T,N}) where T<:Signed where N = encode(map(n -> unsigned(xor(n << 1, n >> 63)), input))

function decodeunsigned(input::Array{UInt8,1}, dtype::DataType=UInt64, outsize::Integer=0)
  # Decode unsigned integer using LEB128
  if outsize == 0
    outsize = length(input) # if don't know, just allocate largest possible array
  end
  output = Array{dtype}(undef, outsize)
  j = 1  # position of next read in input array
  k = 1  # position of current write in output array
  while k <= length(output)
    output[k] = 0
    shift = 0
    while true
        byte = input[j]
        j += 1
        output[k] |= (dtype(byte & 0x7F) << shift)
        if (byte & 0x80 == 0) # replace with right shift 7?
          break
        end
        shift += 7
    end
    k += 1
    if j > length(input)
      break
    end
  end
  return output[1:k-1]
end

function decodesigned(input::Array{UInt8,1}, dtype::DataType=Int64, outsize::Integer=0)
  n = decodeunsigned(input, dtype, outsize)
  # undo zigzag encoding to retrieve the sign
  return movesign.(n)
end

function decode(input::Array{UInt8,1}, dtype::DataType=UInt64, outsize::Integer=0)
  udtypes = Dict(1 => UInt8, 2 => UInt16, 4 => UInt32, 8 => UInt64, 16 => UInt128)
  if dtype <: Signed
    udtype = udtypes[sizeof(dtype)]
    n = decodeunsigned(input, udtype, outsize)
    return movesign.(n)
  else
    n = decodeunsigned(input, dtype, outsize)
    return map(dtype, n)
  end
end

decode(input::UInt8, dtype::DataType=UInt64, outsize::Integer=0) = decode([input], dtype, outsize)

end


text:
using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges

include("../utilities.jl")

mock = MOIU.MockOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()))
config = MOIT.TestConfig()

@testset "NonposToNonneg" begin
    bridged_mock = MOIB.Variable.NonposToNonneg{Float64}(mock)

    @testset "lin2v" begin
        mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [-4, 3, 16, 0],
            (MOI.VectorAffineFunction{Float64}, MOI.Zeros) => [[7, 2, -4]])
        MOIT.lin2vtest(bridged_mock, config)

        @test MOI.get(mock, MOI.NumberOfVariables()) == 4
        @test MOI.get(bridged_mock, MOI.NumberOfVariables()) == 4
        vis = MOI.get(bridged_mock, MOI.ListOfVariableIndices())
        y = vis[4]
        @test y.value == -1

        @test MOI.supports(bridged_mock, MOI.VariablePrimalStart(), MOI.VariableIndex)
        @test MOI.supports(bridged_mock, MOI.VariablePrimalStart(), typeof(MOIB.bridge(bridged_mock, y)))
        MOI.set(bridged_mock, MOI.VariablePrimalStart(), y, 1.0)
        x, y_flipped, z, s = MOI.get(mock, MOI.ListOfVariableIndices())
        @test MOI.get(mock, MOI.VariablePrimalStart(), y_flipped) == -1
        @test MOI.get(bridged_mock, MOI.VariablePrimalStart(), y) == 1

        var_names = ["x", "y", "z", "w"]
        MOI.set(mock, MOI.VariableName(), MOI.get(mock, MOI.ListOfVariableIndices()), var_names)
        con_w = MOI.get(mock, MOI.ListOfConstraintIndices{MathOptInterface.VectorOfVariables, MathOptInterface.Zeros}())[1]
        con_yz = MOI.get(mock, MOI.ListOfConstraintIndices{MathOptInterface.VectorOfVariables, MathOptInterface.Nonnegatives}())
        con_ex = MOI.get(mock, MOI.ListOfConstraintIndices{MathOptInterface.VectorAffineFunction{Float64}, MathOptInterface.Zeros}())[1]

        MOI.set(mock, MOI.ConstraintName(), con_w, "cw")
        MOI.set(mock, MOI.ConstraintName(), con_yz[1], "cy")
        MOI.set(mock, MOI.ConstraintName(), con_yz[2], "cz")
        MOI.set(mock, MOI.ConstraintName(), con_ex, "cex")

        MOI.set(bridged_mock, MOI.VariableName(), MOI.get(bridged_mock, MOI.ListOfVariableIndices())[4], "v")
        con_v = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MathOptInterface.VectorOfVariables, MathOptInterface.Nonpositives}())[1]
        MOI.set(bridged_mock, MOI.ConstraintName(), con_v, "cv")

        @testset "Test Mock model" begin
            s = """
            variables: x, y, z, w
            cw: [w] in MathOptInterface.Zeros(1)
            cy: [y] in MathOptInterface.Nonnegatives(1)
            cz: [z] in MathOptInterface.Nonnegatives(1)
            cex: [1*x + -1*w + 4.0, -1*y + 3.0, 1*x + 1*z + -12.0] in MathOptInterface.Zeros(3)
            minobjective: 3*x + -2*y + -4*z
            """
            model = MOIU.Model{Float64}()
            MOIU.loadfromstring!(model, s)
            MOIU.test_models_equal(mock, model, var_names, ["cw", "cy", "cz", "cex"])
        end

        @testset "Test Bridged model" begin
            s = """
            variables: x, z, w, v
            cv: [v] in MathOptInterface.Nonpositives(1)
            cw: [w] in MathOptInterface.Zeros(1)
            cz: [z] in MathOptInterface.Nonnegatives(1)
            cex: [1*x + -1*w + 4.0, 1*v + 3.0, 1*x + 1*z + -12.0] in MathOptInterface.Zeros(3)
            minobjective: 3*x + 2*v + -4*z
            """
            model = MOIU.Model{Float64}()
            MOIU.loadfromstring!(model, s)
            MOIU.test_models_equal(bridged_mock, model, ["x", "z", "w", "v"], ["cv", "cw", "cz", "cex"])
        end
    end

    @testset "lin4" begin
        MOIU.set_mock_optimize!(mock,
            (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(
                mock, MOI.INFEASIBLE, MOI.INFEASIBLE_POINT,
                MOI.INFEASIBILITY_CERTIFICATE)
        )
        MOIT.lin4test(bridged_mock, config)

        @test MOI.get(mock, MOI.NumberOfVariables()) == 1
        @test length(MOI.get(mock, MOI.ListOfVariableIndices())) == 1
        @test first(MOI.get(mock, MOI.ListOfVariableIndices())).value ≥ 0
        @test MOI.get(bridged_mock, MOI.NumberOfVariables()) == 1
        vis = MOI.get(bridged_mock, MOI.ListOfVariableIndices())
        @test vis == [MOI.VariableIndex(-1)]
        test_delete_bridged_variable(bridged_mock, vis[1], MOI.Nonpositives, 1, (
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
        ))
    end

    @testset "Delete in vector" begin
        MOI.empty!(bridged_mock)
        vis, ci = MOI.add_constrained_variables(bridged_mock, MOI.Nonpositives(4))
        test_delete_bridged_variable(bridged_mock, vis[2], MOI.Nonpositives, 4, (
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
        ), used_bridges = 0, used_constraints = 0)
    end
end


text:
"A PostgreSQL prepared statement"
struct Statement
    """
    A `Connection` for which this statement is valid.
    It may become invalid if the connection is reset.
    """
    jl_conn::Connection

    "An autogenerated neame for the prepared statement (using [`unique_id`](@ref)"
    name::String

    "The query string of the prepared statement"
    query::String

    "A `Result` containing a description of the prepared statement"
    description::Result

    "The number of parameters accepted by this statement according to `description`"
    num_params::Int
end

Base.broadcastable(stmt::Statement) = Ref(stmt)

"""
    prepare(jl_conn::Connection, query::AbstractString) -> Statement

Create a prepared statement on the PostgreSQL server using libpq.
The statement is given an generated unique name using [`unique_id`](@ref).

!!! note

    Currently the statement is not explicitly deallocated, but it is deallocated at the end
    of session per the [PostgreSQL documentation on DEALLOCATE](https://www.postgresql.org/docs/10/sql-deallocate.html).
"""
function prepare(jl_conn::Connection, query::AbstractString)
    uid = unique_id(jl_conn, "stmt")

    result = lock(jl_conn) do
        libpq_c.PQprepare(
            jl_conn.conn,
            uid,
            query,
            0,  # infer all parameters from the query string
            C_NULL,
        )
    end

    close(handle_result(Result(result, jl_conn); throw_error=true))

    result = lock(jl_conn) do
        libpq_c.PQdescribePrepared(jl_conn.conn, uid)
    end

    description = handle_result(Result(result, jl_conn); throw_error=true)

    Statement(jl_conn, uid, query, description, num_params(description))
end

"""
    show(io::IO, jl_result::Statement)

Show a PostgreSQL prepared statement and its query.
"""
function Base.show(io::IO, stmt::Statement)
    print(
        io,
        "PostgreSQL prepared statement named ",
        stmt.name,
        " with query ",
        stmt.query,
    )
end

"""
    num_params(stmt::Statement) -> Int

Return the number of parameters in the prepared statement.
"""
num_params(stmt::Statement) = num_params(stmt.description)

"""
    num_columns(stmt::Statement) -> Int

Return the number of columns that would be returned by executing the prepared statement.
"""
num_columns(stmt::Statement) = num_columns(stmt.description)

"""
    column_name(stmt::Statement, column_number::Integer) -> String

Return the name of the column at index `column_number` (1-based) that would be returned by
executing the prepared statement.
"""
function column_name(stmt::Statement, column_number::Integer)
    column_name(stmt.description, column_number)
end

"""
    column_names(stmt::Statement) -> Vector{String}

Return the names of all the columns in the query result that would be returned by executing
the prepared statement.
"""
column_names(stmt::Statement) = column_names(stmt.description)

"""
    column_number(stmt::Statement, column_name::AbstractString) -> Int

Return the index (1-based) of the column named `column_name` that would be returned by
executing the prepared statement.
"""
function column_number(stmt::Statement, column_name::AbstractString)
    column_number(stmt.description, column_name)
end

function execute(
    stmt::Statement,
    parameters::Union{AbstractVector, Tuple};
    throw_error::Bool=true,
    kwargs...
)
    num_params = length(parameters)
    string_params = string_parameters(parameters)
    pointer_params = parameter_pointers(string_params)

    result = lock(stmt.jl_conn) do
        _execute_prepared(stmt.jl_conn.conn, stmt.name, pointer_params)
    end

    return handle_result(Result(result, stmt.jl_conn; kwargs...); throw_error=throw_error)
end

function execute(
    stmt::Statement;
    throw_error::Bool=true,
    kwargs...
)
    result = lock(stmt.jl_conn) do
        _execute_prepared(stmt.jl_conn.conn, stmt.name)
    end

    return handle_result(Result(result, stmt.jl_conn; kwargs...); throw_error=throw_error)
end

function _execute_prepared(
    conn_ptr::Ptr{libpq_c.PGconn},
    stmt_name::AbstractString,
    parameters::Vector{Ptr{UInt8}}=Ptr{UInt8}[],
)
    num_params = length(parameters)

    return libpq_c.PQexecPrepared(
        conn_ptr,
        stmt_name,
        num_params,
        num_params == 0 ? C_NULL : parameters,
        C_NULL,  # paramLengths is ignored for text format parameters
        num_params == 0 ? C_NULL : zeros(Cint, num_params),  # all parameters in text format
        zero(Cint),  # return result in text format
    )
end


text:
# core.jl - single place to load all package definitions.
#
# If you are willing to learn the package structure, just go through
# included files one by one, read header notes and other comments

using IterTools
using Einsum
using Espresso
importall Espresso
import Espresso: canonical, canonical_calls, to_context, OpName, without_indices
import Espresso: reduce_equalities, flatten, @get, @get_or_create, isspecial

include("utils.jl")
include("ops.jl")
include("deriv.jl")
include("rules.jl")
include("tderiv.jl")
include("trules.jl")
include("special.jl")
include("specialrules.jl")
include("rdiff.jl")
include("tdiff.jl")


text:
"""
This function calculates the field components given two wire center position coordinates
    It assumes you want to fild the fields AT C₂ FROM C₁ as well as the (fully ANTIsymmetric) fields from the wire at C₂ on itself given some wire diameter WireD and current I₂ (assume 1 Amp)

    This will be some constant term, which will be ⟂ to the vector between the two(C₂-C₁), and have some gradient that decays in the direction of C₂-C₁.

This is a 2D function and returns the field
"""
function WireDist2Field(C₁::Vector,C₂::Vector,WireD=0.001,I₁=1,I₂=1)
    μ₀ = 4π*1e-7
    r = WireD/2
    DistVec = C₂ .- C₁
    Dist = √(sum(DistVec.^2))
    # D̂ = Dist ./ sum(Dist.^2) #the unit vector between the two coordinates
    # D′ = D̂ * [0 -1;1 0] #the unit vector rotated 90°

    MagH₂₁_Near = WireNearField_B(Dist-r,r;I=I₁)/μ₀ # The field from 1 on 2 is ⟂ to the vector between them
    MagH₂₁_Far = WireNearField_B(Dist+r,r;I=I₁)/μ₀ # The field from 1 on 2 is ⟂ to the vector between them
    GradH₂₁ = DerivField(Dist,r;I=I₁)/μ₀

    MagH₂₂ = WireNearField_B(r,r,I=I₂)/μ₀

    ## There are assumed to be four boundary conditions on wire 2, K₁,₂,₃,₄ where 1 is the point farthest from wire 1, 4 is the closest, and 2 is the point between 1,4 in the direction of D′ and 4 is the opposite of 2

    HK₁ =      MagH₂₂ + (MagH₂₁_Far)
    HK₃ = -1 * MagH₂₂ + (MagH₂₁_Near)

    HK₂ = -1 * MagH₂₂
    HK₄ =      MagH₂₂

    return HK₁,HK₂,HK₃,HK₄
end

"""
This function calculates the H field distribution in a Conductor (1D) where the width of the wire is W(meters),The skin depth is δ, the external field is ANTIsymmetric (skin effect) and directed in the Z direction (current is in Y) with a magnitude of K. X is the dimension that is ⟂ to I,H, and extends from -W/2, to W/2

    Boundary conditions are ±K A/m at x=±W/2
"""
function HAntiSym(x,δ,W = 1., K=1.)
    if (x)>(W/2)
        return -K
    elseif (x)<(-1*W/2)
        return K
    else
        return  K*(exp.(-1*(1+im)*x/δ) - exp.((1+im)*x/δ)) / (exp.((1+im)*W/(2*δ)) - exp.(-1*(1+im)*W/(2*δ)))
    end

end

"""
This resultant current distribution that corresponds to HAntiSym(...). J = -dH/dx

With anti-symmetric fields, the current on opposite sides of the conductor is in phase(non-zero net current)
"""
function JAntiSym(x,δ,W = 1., K=1.)
    if abs(x)>(W/2)
        return 0
    else
        return  K*((1+im)/δ)*(exp.(-1*(1+im)*x/δ) + exp.((1+im)*x/δ)) /
                      (exp.((1+im)*W/(2*δ)) - exp.(-1*(1+im)*W/(2*δ)))
    end

 end

 function HSym(x,δ,W = 1, K=1)
    if (x)>(W/2)
        return K
    elseif (x)<(-1*W/2)
        return K
    else
        return  K*(exp.(-1*(1+im)*x/δ) + exp.((1+im)*x/δ)) / (exp.((1+im)*W/(2*δ)) + exp.(-1*(1+im)*W/(2*δ)))
    end

end

"""
This resultant current distribution that corresponds to HSym(...). J = -dH/dx
With symmetric fields (even, as in proximity effect), the current on opposite sides of the conductor is in phase(zero net current)
"""
function JSym(x,δ,W = 1, K=1)
    if abs(x)>(W/2)
        return 0
    else
        return  K*((1+im)/δ)*(exp.(-1*(1+im)*x/δ) - exp.((1+im)*x/δ)) /
                      (exp.((1+im)*W/(2*δ)) - exp.(-1*(1+im)*W/(2*δ)))
    end

 end

 function SkinDepth(f,ρ = 1.67e-8)
    μ₀ = 4π*1e-7
    SkinDepth = sqrt.(ρ ./ (π * f * μ₀))
 end

 function plot1DMagDiffusion(K₋,K₊,W,δ,n=1000)

    HTot, JTot = MagDiffusion1D(K₋,K₊,W,δ,1000)
    δX = W/n
    x = (-1*W/2):δX:(W/2)
    figure()
    pygui(true)

    subplot(221)
    plot(x,abs.(abs.(HTot)))
    title("|H|")
    subplot(223)
    plot(x,rad2deg.(angle.(HTot)))
    title("∠H")
    subplot(222)
    plot(x,abs.(JTot))
    plot(x,real.(JTot))

    subplot(224)
    plot(x,rad2deg.(angle.(JTot)))

    println(sum(JTot)/length(JTot)*pi/4*.001^2)
    return HTot, JTot
 end

 """
Calculates the proximity effect losses in one wire at C₁ due to many others in an Nx2 array, Cₐ
 """
 function MultiWireProximity(C₁::Vector,Cₐ::Array,f,WireD=0.001,I₁=1,Iₐ=1)

    Nₐ = length(Cₐ[:,1])
    Losses = 0.
    Losses += OneWireLoss(f,WireD,I₁)
    for i in 1:Nₐ
        Losses += sum(TwoWireLoss(Cₐ[i,:],C₁,f,WireD,1,0))*2
    end
    println(Losses)

    return Losses
 end

 """
Calculates the losses IN WIRE 2 FROM WIRE 1. By default only proximity losses are considered so the default is to have zero current in wire 2

 """
 function TwoWireLoss(C₁::Vector,C₂::Vector,f,WireD=0.001,I₁=1,I₂=0)

    δ = SkinDepth(f)
    HK₁,HK₂,HK₃,HK₄ = WireDist2Field(C₁,C₂,WireD,I₁,I₂)
    println("$([HK₁,HK₂,HK₃,HK₄])")
    HTot_LR, JTot_LR = MagDiffusion1D(HK₁,HK₃,WireD,δ,100000)
    HTot_UD, JTot_UD = MagDiffusion1D(HK₂,HK₄,WireD,δ,100000)
    J_LR² =  abs.(JTot_LR).^2
    J_UD² =  abs.(JTot_UD).^2
    # JTot_LR = WeightRadial!(JTot_LR)
    # JTot_UD = WeightRadial!(JTot_UD)
    σ=58e6
    LossLR = 1/2*(sum(J_LR²)/length(JTot_LR)*pi/4*WireD^2)/σ
    LossUD = 1/2*(sum(J_UD²)/length(JTot_UD)*pi/4*WireD^2)/σ
    # println(LossLR)
    # println(LossUD)
    # θ = LinRange(-π,π,  100)
    # TotLoss =√(sum( (sin.(θ)*LossUD).^2 .+ (cos.(θ)*LossLR).^2))/length(θ)
    return LossLR, LossUD

 end

 function OneWireLoss(f,WireD=0.001,I=1)

    Area = pi/4*WireD^2
    μ₀ = 4π*1e-7
    r = WireD/2
    MagH = WireNearField_B(r,r,I=I)/μ₀
    δ = SkinDepth(f)
    HTot, JTot = MagDiffusion1D(-MagH,MagH,WireD,δ,100000)
    JTot = 2 .* JTot ##account for X and Y
    σ=58e6
    # JTot = WeightRadial!(JTot)
    # J² =  WeightRadial!(abs.(JTot).^2)
    J²RMS =  WeightRadial!(abs.(JTot ./ (√2)).^2)
    print("Net current:")
    println(round( (sum(JTot)/length(JTot)*Area);sigdigits=3))
    Losses = (sum(J²RMS)/length(JTot)*Area)/σ

    #println(Losses)

    return Losses
 end

function MagDiffusion1D(HK₋,HK₊,W,δ,n=1000)
    δX = W/n
    x = (-1*W/2):δX:(W/2)

    meanK = (HK₊+HK₋)
    ΔK = (HK₋-HK₊)/2

    HTot = HAntiSym.(x,δ,W, ΔK) .+ HSym.(x,δ,W, meanK)
    JTot = JAntiSym.(x,δ,W, ΔK) .+ JSym.(x,δ,W, meanK)

    return HTot, JTot
end

function SeriesWireBundleLosses(Cᵦ::Array, f,I,WireD)

    Nₐ = length(Cᵦ[:,1])
    Losses = 0.

    for i in 1:Nₐ
        Cₐ = Cᵦ[setdiff(1:end, i), :]
        C₁ = Cᵦ[i,:]

        Losses += MultiWireProximity(C₁,Cₐ,f,WireD,I,I)
    end
    return Losses
end

"""
This function takes in a vector and weights it such that the values are weighted by distance from the middle (for example when integrating a radial system)
"""
function WeightRadial!(V::Vector)
    Weights = abs.(collect(LinRange(-2,2,length(V))))
    Weights = Weights ./(sum(Weights))*length(V)
    V = V.*Weights
    return V
end

"""
This function takes in a vector and weights it such that the values are weighted by distance from the middle except with more weight towards the center.
"""
function WeightCenter(V::Vector)
    Weights = abs.(abs.(collect(LinRange(-1,1,length(V)))) .- 2)
    # Weights = Weights ./(sum(Weights))*length(V)
    V = V.*Weights
    return V
end


text:
# module for representing agent
# decide control order to robot
# get action based on policy
# by dynamic programming
# using probabilistic flow control

include(joinpath(split(@__FILE__, "src")[1], "src/decision_making/markov_decision_process/dynamic_programming.jl"))

mutable struct PfcAgent
  speed
  yaw_rate
  delta_time
  estimator
  prev_spd
  prev_yr
  puddle_coef
  puddle_depth
  total_reward
  in_goal
  final_value
  goal
  pose_min
  pose_max
  widths
  index_nums
  policy_data
  dp
  evaluations
  current_value
  stop_timer
  magnitude

  # init
  function PfcAgent(;delta_time::Float64=0.1,
                    estimator=nothing,
                    goal=nothing,
                    puddles=nothing,
                    sampling_num=10,
                    puddle_coef=100,
                    widths=[0.2, 0.2, pi/18],
                    lower_left=[-4.0, -4.0],
                    upper_right=[4.0, 4.0],
                    magnitude=2)
    self = new()
    self.speed = 0.0
    self.yaw_rate = 0.0
    self.delta_time = delta_time
    self.estimator = estimator
    self.prev_spd = 0.0
    self.prev_yr = 0.0
    self.puddle_coef = puddle_coef
    self.puddle_depth = 0.0
    self.total_reward = 0.0
    self.in_goal = false
    self.final_value = 0.0
    self.goal = goal

    self.pose_min = [lower_left[1], lower_left[2], 0.0]
    self.pose_max = [upper_right[1], upper_right[2], 2*pi]
    self.widths = widths
    self.index_nums = round.(Int64, (self.pose_max - self.pose_min)./self.widths)
    self.policy_data = init_policy(self)

    self.dp = DynamicProgramming(widths, goal, puddles, delta_time, sampling_num)
    self.dp.value_function = init_value(self)
    self.evaluations = [0.0, 0.0, 0.0] # to store q-mdp value
    self.current_value = 0.0 # to store average of current state value

    self.stop_timer = 0.0
    self.magnitude = magnitude

    return self
  end
end

function init_value(self::PfcAgent)
  tmp = zeros(Tuple(self.dp.index_nums))

  txt_path = "src/decision_making/partially_observable_mdp/value.txt"
  open(joinpath(split(@__FILE__, "src")[1], txt_path), "r") do fp
    for line in eachline(fp)
      d = split(line) # [i_x, i_y, i_theta, value]
      tmp[parse(Int64, d[1])+1, parse(Int64, d[2])+1, parse(Int64, d[3])+1] = parse(Float64, d[4])
    end
  end

  return tmp
end

function init_policy(self::PfcAgent)
  tmp = zeros(Tuple([self.index_nums[1], self.index_nums[2], self.index_nums[3], 2]))

  txt_path = "src/decision_making/partially_observable_mdp/policy.txt"
  open(joinpath(split(@__FILE__, "src")[1], txt_path), "r") do fp
    for line in eachline(fp)
      d = split(line) # [i_x, i_y, i_theta, speed, yaw_rate]
      tmp[parse(Int64, d[1])+1, parse(Int64, d[2])+1, parse(Int64, d[3])+1, :] .= [parse(Float64, d[4]), parse(Float64, d[5])]
    end
  end

  return tmp
end

function reward_per_sec(self::PfcAgent)
  return (-1.0 - self.puddle_depth*self.puddle_coef)
end

function to_index(self::PfcAgent, pose)
  index = Int64.(floor.((pose - self.pose_min)./self.widths))

  # normalize direction index
  index[3] = (index[3] + self.index_nums[3]*1000)%self.index_nums[3]

  for i in 1:2
    if index[i] < 0
      index[i] = 0
    elseif index[i] >= self.index_nums[i]
      index[i] = self.index_nums[i] - 1
    end
  end

  return index
end

function evaluation(self::PfcAgent, action, indexes)
  v = self.dp.value_function
  vs = []
  for i in indexes
    if abs(self.dp.value_function[i[1]+1, i[2]+1, i[3]+1]) > 0.0
      push!(vs, abs(self.dp.value_function[i[1]+1, i[2]+1, i[3]+1]))
    else
      push!(vs, 1e-10)
    end
  end

  qs = [action_value(self.dp, action, i, out_penalty=false) for i in indexes]

  ps = self.estimator.particles

  return sum([(p.weight/(v^self.magnitude))*q for (v, q, p) in zip(vs, qs, ps)])
end

function policy(self::PfcAgent, pose)
  # decrease weight of particle reached at goal
  for p in self.estimator.particles
    if inside(self.goal, p.pose) == true
      p.weight *= 1e-10
    end
  end
  resampling(self.estimator)

  indexes = [to_index(self, p.pose) for p in self.estimator.particles]
  self.current_value = sum([self.dp.value_function[i[1]+1, i[2]+1, i[3]+1] for i in indexes])/length(indexes)
  self.evaluations = [evaluation(self, a, indexes) for a in self.dp.actions]
  action = self.dp.actions[argmax(self.evaluations)]

  return action[1], action[2]
end

function draw_decision!(self::PfcAgent, observation)
  # reached goal
  if self.in_goal == true
    return 0.0, 0.0 # stop
  end

  if self.estimator !== nothing
    motion_update(self.estimator, self.prev_spd, self.prev_yr, self.delta_time)
    observation_update(self.estimator, observation)

    self.total_reward += self.delta_time * reward_per_sec(self)

    spd_tmp, yr_tmp = policy(self, self.estimator.estimated_pose)
    if spd_tmp == 0.0
      self.stop_timer += self.delta_time
    else
      self.stop_timer = 0.0
    end

    if self.stop_timer > 1.0
      self.speed, self.yaw_rate = 1.0, 0.0
    else
      self.speed, self.yaw_rate = spd_tmp, yr_tmp
    end
    self.prev_spd, self.prev_yr = self.speed, self.yaw_rate

    draw!(self.estimator)

    annotate!(-4.5, -4.6, text("$(round(self.current_value, digits=3)) => $(round.(self.evaluations, digits=3))", :left, 8))
  end

  return self.speed, self.yaw_rate
end

max_stars_repo_name: ShisatoYano/JuliaAutonomy

max_issues_repo_name: ShisatoYano/JuliaAutonomy

max_forks_repo_name: ShisatoYano/JuliaAutonomy


text:
#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
#  Released under the modified BSD license. See COPYING.md for more details.

using Statistics
using DataStructures
import Base.Threads: threadid

function take(
    pool::NodePool;
    suggestions::Array{Node} = [],
    time_remaining::Float64,
    gap_limit::Float64,
    node_limit::Int,
)::Union{Symbol,Node}
    t = threadid()
    lock(pool.lock) do
        n_processing = length(pool.processing)
        if (
            (pool.gap < gap_limit) ||
            (n_processing + pool.processed >= node_limit) ||
            (time_remaining < 0)
        )
            return :END
        end
        if isempty(pool.pending)
            if isempty(pool.processing)
                return :END
            else
                return :WAIT
            end
        else
            # If one of the suggested nodes is still pending, return it.
            # This is known in the literature as plunging.
            for s in suggestions
                if s in keys(pool.pending)
                    delete!(pool.pending, s)
                    pool.processing[s] = s.obj
                    return s
                end
            end

            # If all suggestions have already been processed
            # or pruned, find another node based on best bound.
            node = dequeue!(pool.pending)
            pool.processing[node] = node.obj
            return node
        end
    end
end

function offer(
    pool::NodePool;
    parent_node::Union{Nothing,Node},
    child_nodes::Vector{Node},
    time_elapsed::Float64 = 0.0,
    print_interval::Int = 100,
    detailed_output::Bool = false,
)::Nothing
    lock(pool.lock) do
        primal_update = false

        # Update node.processing and node.processed
        pool.processed += 1
        if parent_node !== nothing
            delete!(pool.processing, parent_node)
        end

        # Queue child nodes
        for node in child_nodes
            if node.status == :Infeasible
                continue
            end
            if node.obj >= pool.primal_bound
                continue
            end
            if isempty(node.fractional_variables)
                pool.primal_bound = min(pool.primal_bound, node.obj)
                primal_update = true
                continue
            end
            pool.pending[node] = node.obj
        end

        # Update dual bound
        pool.dual_bound = pool.primal_bound
        if !isempty(pool.pending)
            pool.dual_bound = min(pool.dual_bound, peek(pool.pending)[2])
        end
        if !isempty(pool.processing)
            pool.dual_bound = min(pool.dual_bound, peek(pool.processing)[2])
        end

        # Update gap
        if pool.primal_bound == pool.dual_bound
            pool.gap = 0
        else
            pool.gap = abs((pool.primal_bound - pool.dual_bound) / pool.primal_bound)
        end

        if parent_node !== nothing
            # Update branching variable history
            branch_var = child_nodes[1].branch_variables[end]
            offset = findfirst(isequal(branch_var), parent_node.fractional_variables)
            x = parent_node.fractional_values[offset]
            obj_change_up = child_nodes[1].obj - parent_node.obj
            obj_change_down = child_nodes[2].obj - parent_node.obj
            _update_var_history(
                pool = pool,
                var = branch_var,
                x = x,
                obj_change_down = obj_change_down,
                obj_change_up = obj_change_up,
            )

            # Update global history
            pool.history.avg_pseudocost_up =
                mean(vh.pseudocost_up for vh in values(pool.var_history))
            pool.history.avg_pseudocost_down =
                mean(vh.pseudocost_down for vh in values(pool.var_history))

            # Print progress
            print_progress(
                pool,
                parent_node,
                time_elapsed = time_elapsed,
                print_interval = print_interval,
                detailed_output = detailed_output,
                primal_update = primal_update,
            )
        end
    end
    return
end

function _update_var_history(;
    pool::NodePool,
    var::Variable,
    x::Float64,
    obj_change_down::Float64,
    obj_change_up::Float64,
)::Nothing
    # Create new history entry
    if var ∉ keys(pool.var_history)
        pool.var_history[var] = VariableHistory()
    end
    varhist = pool.var_history[var]

    # Push fractional value
    push!(varhist.fractional_values, x)

    # Push objective value changes
    push!(varhist.obj_change_up, obj_change_up)
    push!(varhist.obj_change_down, obj_change_down)

    # Push objective change ratios
    f_up = x - floor(x)
    f_down = ceil(x) - x
    if isfinite(obj_change_up)
        push!(varhist.obj_ratio_up, obj_change_up / f_up)
    end
    if isfinite(obj_change_down)
        push!(varhist.obj_ratio_down, obj_change_down / f_down)
    end

    # Update variable pseudocosts
    varhist.pseudocost_up = 0.0
    varhist.pseudocost_down = 0.0
    if !isempty(varhist.obj_ratio_up)
        varhist.pseudocost_up = sum(varhist.obj_ratio_up) / length(varhist.obj_ratio_up)
    end
    if !isempty(varhist.obj_ratio_down)
        varhist.pseudocost_down =
            sum(varhist.obj_ratio_down) / length(varhist.obj_ratio_down)
    end
    return
end

function generate_indices(pool::NodePool, n::Int)::Vector{Int}
    lock(pool.lock) do
        result = Int[]
        for i = 1:n
            push!(result, pool.next_index)
            pool.next_index += 1
        end
        return result
    end
end


text:
module Mid

using JuMP
using Cbc

include("Computation.jl")
using .Computation

include("Format.jl")
using .Format

export mid, vmid

# Determines upper bound alpha with Midpoint Method, optionally outputs proof
# Author: Jason Liu
function mid(m::Int64, s::Int64; output::Int64 = 2)
    output > 0 && printHeader(center("MIDPOINT METHOD"))

    if m <= 0 || s <= 0
        if output > 0
            printf("Midpoint Method does not apply", line = true)
            printEnd()
        end
        return 1 // 1
    elseif m < s
        alpha = m // s * mid(s, m, output = 0)
        if output > 1
            printfT(
                "Duality Theorem",
                "muffins($m,$s) = $m/$s · muffins($s,$m)",
                "Bounding muffins($s,$m) instead",
            )
            mid(s, m, output = output)
        elseif output > 0
            printfT(
                "Midpoint Method",
                "Upper bound α of muffins($m,$s) is $(formatFrac(alpha))",
            )
            printEnd()
        end
        return alpha
    elseif m % s == 0
        if output > 0
            printfT("Midpoint Method", "Since $m % $s = 0, muffins($m,$s) = 1")
            printEnd()
        end
        return 1 // 1
    end

    (V, W, sV, sW) = sv(m, s)
    numV = (V)sV
    numW = (W)sW

    # Compute potential alpha candidates
    alphas = []
    if numW > numV
        for (a, b, c) in comboTup(W, 3)
            append!(alphas, (a // 2 + (W)b + c - (1 + b)m // s) // ((W - 1)b + c))                  # Value for alpha derived by solving a(1/2) + b(1-y) + c(1-alpha) = m/s
            append!(alphas, ((1 - a + c)m // s + (W - 1)a - b // 2 - c) // ((W - 1)a + (V - 1)c))   # Value for alpha derived by solving a(y) + b(1/2) + c(1-x) = m/s
        end
    elseif numW < numV
        for (a, b, c) in comboTup(V, 3)
            append!(alphas, ((W)a + b // 2 - (1 + a - c)m // s) // ((W - 1)a + (V - 1)c))           # Value for alpha derived by solving a(1-y) + b(1/2) + c(x) = m/s
            append!(alphas, ((1 + b)m // s - b - c // 2) // (a + (V - 1)b))                         # Value for alpha derived by solving a(alpha) + b(1-x) + c(1/2) = m/s
        end
    end

    # Test alpha candidates
    for alpha in sort(unique(alphas))
        if denominator(alpha) != 0 && vmid(m, s, alpha, output = 0)
            vmid(m, s, alpha, output = output)
            return alpha
        end
    end

    if output > 0
        printf(
            "All α candidates were inconclusive, Midpoint Method inconclusive",
            line = true,
        )
        printEnd()
    end
    1 // 1
end

# Helper function for mid -- verifies whether mid(m, s, alpha) is conclusive
function vmid(m::Int64, s::Int64, alpha::Rational{Int64}; output::Int64 = 2)
    if !(m > s > 0) || m % s == 0
        output > 0 && printf("VMid does not apply", line = true)
        false
    elseif alpha < 1 / 3
        output > 0 && printfT("No piece size < 1/3", "For m ≥ s, α must be ≥ 1/3")
        false
    elseif alpha > 1
        output > 0 && printfT("No piece size > 1", "α must be ≤ 1")
        false
    else
        (V, W, sV, sW) = sv(m, s)
        numV = (V)sV
        numW = (W)sW

        # Initialize Midpoint Method proof
        if output > 1
            # Define and format variables for proof
            alphaF = formatFrac(alpha)
            size = formatFrac(m // s)
            a = formatFrac(m // (V + 1)s)
            b = max(1 - m // (W - 1)s, 0)
            b1 = formatFrac(1 - b)
            b = formatFrac(b)

            # Establish assumptions and premises
            printHeader("CLAIM")
            printfT(
                "Claim",
                "For all ($m, $s) procedures where $s students each receive $size muffins, the smallest piece size is bounded above by α = $alphaF",
                "Proven by contradicting the assumption",
            )

            printHeader("ASSUMPTION")
            printfT("Assumption", "Assume muffins($m,$s) > α ≥ 1/3")
            printfT(
                "Property of Buddies",
                "Since muffins($m,$s) > 1/3, each muffin must be cut into 2 shs, totaling $(2m) shs",
                "",
                "Since each muffin is cut into 2 shs that are buddies, each sh of size x must imply the existence of a sh of size 1-x",
            )

            # Begin casework
            printHeader("CASEWORK")

            printfT(
                "Case 1",
                "Alice has ≥ $(V+1) shs",
                "One of them is ≤ ($size × 1/$(V+1)) = $a",
                "",
                "Contradicts assumption if α ≥ $a",
            )

            if W - 1 <= 1 && m / s > 1
                printfT(
                    "Case 2",
                    "Bob cannot have ≤ $(W-1) shs since $size > 1, so this case is impossible",
                )
            else
                printfT(
                    "Case 2",
                    "Bob has ≤ $(W-1) shs",
                    "One of them is ≥ ($size × 1/$(W-1)) = $b1",
                    "Its buddy is ≤ (1 - $b1) = $b",
                    "",
                    "Contradicts assumption if α ≥ $b",
                )
            end
        end

        # Check if V-Conjecture applies
        if m // s * 1 // (V + 1) > alpha || 1 - m // s * 1 // (W - 1) > alpha
            output > 0 && printf("V-Conjecture does not apply, VMid failed", line = true)
            return false
        end

        ((_, x), (y, _)) = findend(m, s, alpha, V)
        xOriginal = m // s - alpha * (V - 1)
        yOriginal = m // s - (1 - alpha) * (V - 2)

        # Check if FindEnd works
        if x != xOriginal && y != yOriginal
            output > 0 && printf("FindEnd inconclusive, VMid failed", line = true)
            return false
        end

        # Define and format variables for proof
        cd = lcm(s, denominator(x), denominator(y), 2)
        alphaF = formatFrac(alpha, cd)
        alpha1 = formatFrac(1 - alpha, cd)
        size = formatFrac(m // s)
        a = formatFrac(m // s - x, cd)
        b = formatFrac(m // s - y, cd)
        xF = formatFrac(x, cd)
        x1 = formatFrac(1 - x, cd)
        yF = formatFrac(y, cd)
        y1 = formatFrac(1 - y, cd)
        diff = abs(numV - numW)

        # Continue proof
        if output > 1
            # Continue casework
            printfT(
                "Note",
                "The remaining cases deal with everyone having either $W or $V shs, so:",
                "",
                "$(W)·s_$W + $(V)·s_$V = $(2m)  (total shs)",
                "s_$W + s_$V = $s  (total students)",
                "where s_N = # of students with N shs",
                "",
                "The solution to the system is:",
                "s_$W = $sW, s_$V = $sV",
                "So there are $numW $W-shs and $numV $V-shs",
            )

            if x == xOriginal
                printfT(
                    "Case 3",
                    "Alice has a $V-sh ≥ $xF",
                    "Her other $(V-1) $V-shs sum to ≤ ($size - $xF) = $a",
                    "One of them is ≤ ($a × 1/$(V-1)) = $alphaF",
                    "",
                    "Contradicts assumption if α ≥ $alphaF",
                )
            else
                printfT("Case 3", "FindEnd did not produce a conclusive bound for $V-shs")
            end

            if y == yOriginal
                printfT(
                    "Case 4",
                    "Bob has a $W-sh ≤ $yF",
                    "His other $(W-1) $W-shs sum to ≥ ($size - $yF) = $b",
                    "One of them is ≥ ($b × 1/$(W-1)) = $alpha1",
                    "Its buddy is ≤ (1 - $alpha1) = $alphaF",
                    "",
                    "Contradicts assumption if α ≥ $alphaF",
                )
            else
                printfT("Case 4", "FindEnd did not produce a conclusive bound for $W-shs")
            end

            printHeader("CASE 5: INTERVAL ANALYSIS")
            println()
            printf("The following intervals capture the negation of the previous cases:")
            if alpha < x <= y < 1 - alpha
                println(
                    "\n",
                    interval(
                        ["(", alphaF],
                        [")[", xF],
                        ["](", yF],
                        [")", alpha1],
                        labels = ["$numV $V-shs", "0", "$numW $W-shs"],
                    ),
                )
                printLine()
            else
                if xF != alphaF
                    println(
                        "\n",
                        interval(
                            ["(", alphaF],
                            [")[", xF],
                            [")", alpha1],
                            labels = ["$numV $V-shs", "0 $V-shs"],
                        ),
                    )
                else
                    println("\n", "No negation interval for $V-shs")
                end

                if yF != alpha1
                    println(
                        "\n",
                        interval(
                            ["(", alphaF],
                            ["](", yF],
                            [")", alpha1],
                            labels = ["0 $W-shs", "$numW $W-shs"],
                        ),
                    )
                else
                    println("\n", "No negation interval for $W-shs")
                end

                printfT(
                    "Case 5",
                    "The Midpoint Method is inconclusive on these intervals, VMid failed",
                )
            end
        end

        # Fail if interval inconclusive
        if x > y || x == alpha || y == 1 - alpha
            output == 1 &&
                printf("Could not generate disjoint intervals, VMid failed", line = true)
            return false
        end

        # Define and format variables for proof
        (vMin, vMax, sMin, sMax) = (V, W, sV, sW)
        (i, j, k, l) = (xF, yF, y1, x1)
        (rngMin, rngMax, rngS, rngL) = ((alphaF, xF), (x1, alpha1), (yF, y1), (x1, alpha1))
        if numV > numW
            (vMin, vMax, sMin, sMax) = (W, V, sW, sV)
            (i, j, k, l) = (y1, x1, xF, yF)
            (rngMin, rngMax, rngS, rngL) =
                ((yF, alpha1), (alphaF, y1), (alphaF, y1), (x1, xF))
        end
        numMin = (vMin)sMin

        # Conclude interval case with case by case analysis
        if output > 1
            printfT(
                "Property of Buddies",
                "Because 0 shs lie in [$xF,$yF], 0 shs must also lie in [$y1,$x1]",
                "",
                "Similary, the $numMin $vMin-shs in ($(rngMin[1]),$(rngMin[2])) implies $numMin $vMax-shs in ($(rngMax[1]),$(rngMax[2]))",
            )
            println()
            printf("The following intervals capture the previous statements:")
            println(
                "\n",
                interval(
                    ["(", alphaF],
                    [")[", i],
                    ["](", j],
                    [")[", k],
                    ["](", l],
                    [")", alpha1],
                    labels = [
                        "$numMin $V-shs",
                        "0",
                        "$diff $vMax-shs",
                        "0",
                        "$numMin $W-shs",
                    ],
                ),
            )
            printLine()

            printfT(
                "Note",
                "Let the $vMax-shs that lie in ($(rngS[1]),$(rngS[2])) be small $vMax-shs",
                "Let the $vMax-shs that lie in ($(rngL[1]),$(rngL[2])) be large $vMax-shs",
            )
        end

        # Determine possible small-sh, large-sh combinations
        posCombo = []
        for k = 0:vMax
            upperB = (vMax - k) * toFrac(rngL[2]) + (k) * toFrac(rngS[2])
            lowerB = (vMax - k) * toFrac(rngL[1]) + (k) * toFrac(rngS[1])
            if upperB > m // s > lowerB
                append!(posCombo, k)
            end
        end

        # Diverge from Interval Method
        if length(posCombo) > 0
            mid = formatFrac(1 // 2, cd)
            if output > 1
                combos =
                    ["$k small $vMax-shs and $(vMax-k) large $vMax-shs" for k in posCombo]
                printfT(
                    "Case 5",
                    "All combinations of small and large $vMax-shs are impossible except for:",
                    "",
                    combos...,
                    "",
                    "because the range of muffin amounts for the others does not include $size",
                )

                printHeader("CASE 5: MIDPOINT ANALYSIS")
                printfT(
                    "Property of Buddies",
                    "The # of $vMax-shs in ($j, $mid) is equal to the # of $vMax-shs in ($mid, $k)",
                )
            end

            bounds = [[toFrac(j), 1 // 2], [1 // 2, toFrac(k)], [toFrac(l), toFrac(alpha1)]]
            smallBounds = [toFrac(j), toFrac(k)]
            intervals = [["(", j], ["|", mid], [")[", k], ["](", l], [")", alpha1]]
            labels = ["z $W-shs", "z $W-shs", "0", "$numV $W-shs"]
            intervalDefs = [
                "A = ($j,$mid)",
                "B = ($mid,$k) (|A| = |B| = z)",
                "C = ($l,$alpha1) (|C| = $numV)",
            ]
            if numW < numV
                bounds = [[alpha, toFrac(i)], [toFrac(j), 1 // 2], [1 // 2, toFrac(k)]]
                smallBounds = [alpha, toFrac(i)]
                intervals = [["(", alphaF], [")[", i], ["](", j], ["|", mid], [")", k]]
                labels = ["$numW $V-shs", "0", "z $V-shs", "z $V-shs"]
                intervalDefs = [
                    "A = ($alphaF, $i) (|A| = $numW)",
                    "B = ($j, $mid)",
                    "C = ($mid, $k) (|B| = |C| = z)",
                ]
            end

            if output > 1
                println()
                printf("The following intervals capture what we know about the $V-shs:")
                println("\n", interval(intervals..., labels = labels))
                printLine()
                printfT("Note", "We define the following intervals:", intervalDefs...)
            end

            # Determine possible interval combinations
            posInt = intCombo(m, s, vMax, bounds, smallBounds, posCombo)

            if length(posInt) > 0
                if output > 1
                    printfT(
                        "Note",
                        "Let the notation (a, b, c, ...) denote a combination of $vMax-shs where there are:",
                        "a shs in interval A",
                        "b shs in interval B",
                        "c shs in interval C",
                        "...",
                        "",
                        "A (a, b, c, ...)-student denotes a student with that combination of $vMax-shs",
                    )
                    printfT(
                        "Midpoint Analysis I",
                        "Given the possible small and large $vMax-sh combinations,",
                        "all combinations of $vMax-shs from intervals A, B, and C are impossible except for:",
                        "",
                        posInt...,
                        "",
                        "because the range of muffin amounts for the others does not include $size",
                    )
                end

                # Check if posInt is too large
                if length(posInt) > 5
                    output > 0 && printf(
                        "The # of possible interval combinations is too large, VMid failed",
                        line = true,
                    )
                    return false
                end

                # Set up and solve system of equations
                (a, b, c) =
                    (getindex.(posInt, 1), getindex.(posInt, 2), getindex.(posInt, 3))
                (A, B, C) = ("A", "B", "C")
                if numW < numV
                    (a, b, c) = (b, c, a)
                    (A, B, C) = (B, C, A)
                end

                iter = 1:length(posInt)
                equations = [
                    join(["$(a[i])·x_$i" for i in iter], " + ") *
                    " = " *
                    join(["$(b[i])·x_$i" for i in iter], " + "),
                    join(["$(c[i])·x_$i" for i in iter], " + ") * " = |$C| = $numMin",
                    join(["x_$i" for i in iter], " + ") * " = s_$vMax = $sMax",
                    "",
                    ["where x_$i is the # of $(posInt[i])-students" for i in iter]...,
                ]

                # Test all possible non-negative integer solutions with system
                model = Model(Cbc.Optimizer)
                set_optimizer_attribute(model, "logLevel", 0)
                set_optimizer_attribute(model, "cuts", "off")

                @variable(model, 0 <= x[i = 1:length(iter)], Int)

                @constraint(model, sum(x .* a) == sum(x .* b))
                @constraint(model, sum(x .* c) == numMin)
                @constraint(model, sum(x) == sMax)

                @objective(model, Max, x[1])
                optimize!(model)
                if termination_status(model) == MOI.OPTIMAL
                    x = convert.(Int64, round.(value.(x)))
                    solution = [
                        "One solution is " * join(["x_$i = $(x[i])" for i in iter], ", "),
                        "The system has a non-negative integer solution, so Case 5 is still possible",
                    ]
                else
                    solution =
                        ["The system has no non-negative integer solutions, so Case 5 is impossible"]
                end

                if output > 1
                    printfT(
                        "Midpoint Analysis II",
                        "Since |$A| = |$B|, we can set up the following system of equations:",
                        "",
                        equations...,
                        "",
                        solution...,
                    )
                end

                if length(solution) > 1
                    output == 1 && printf(
                        "Could not generate conclusive system of equations, VMid failed",
                        line = true,
                    )
                    return false
                end

            else
                if output > 1
                    printfT(
                        "Midpoint Analysis",
                        "Given the possible small and large $vMax-sh combinations,",
                        "all combinations of $vMax-shs from intervals A, B, and C are impossible since their total amount can not be $size",
                    )
                end
            end

        else
            if output > 1
                printfT(
                    "Case 5",
                    "All combinations of small and large $vMax-shs are impossible since their total amount can not be $size",
                )
            end
        end

        # Conclude proof
        alpha = formatFrac(alpha)
        if output > 1
            # Conclude with alpha's value
            printHeader("CONCLUSION")
            printfT(
                "Conclusion",
                "All possible cases contradict the assumption iff. α ≥ $alpha",
                "muffins($m,$s) ≤ α, ∀ α ≥ $alpha",
                "",
                "muffins($m,$s) ≤ $alpha",
            )
            printEnd()
        elseif output > 0
            printfT("Midpoint Method", "Upper bound α of muffins($m,$s) is $alpha")
            printEnd()
        end

        true
    end
end

end


text:
export DelayNetwork

"""
# DelayNetwork
Network with axonal conduction delays.
# Arguments
- `neurons::NeuronType`
- `batch_size::Integer=1`
...
- `weight::AbstractArray`: Final size= (#neurons, #neurons)
- `adjacency::AbstractArray`: Final size= (#neurons, #neurons)
- `spikes::AbstractArray`: Final size= (Batch size, #neurons)
- `voltage::AbstractArray`: Final size= (Batch size, #neurons)
- `recovery::AbstractArray`: Membrane recovery variable, only applicable to Izhikevich neurons
- `refractory::AbstractArray`: Refractory period, only applicable to LIF neurons. Final size= (Batch size, #neurons)
...
- `delay::AbstractArray`: Axonal conduction delays
- `_delay::AbstractArray`: Auxiliary delay array for calculations
- `delayed_voltages::AbstractArray`: Keeps track of upcomig voltages
...
"""
@with_kw mutable struct DelayNetwork{T<:NeuronType} <: SpikingNetwork{T}
    neurons::T
    batch_size::Int = 1
    learning_rule::Union{LearningRule,Nothing} = nothing
    weight::AbstractArray = Array{Float64}(undef, (0, 0))
    adjacency::AbstractArray = Array{Bool}(undef, (0, 0))
    spikes::AbstractArray = zeros(Bool, batch_size, 0)
    voltage::AbstractArray = ones(Float32, batch_size, 0)
    recovery::AbstractArray = ones(Float32, batch_size, 0)
    refractory::AbstractArray = zeros(Int32, batch_size, 0)
    learning::Bool = isnothing(learning_rule) ? false : true
    groups::Dict{String,NeuronGroup} = Dict{String,NeuronGroup}()
    delay::AbstractArray = Array{Float64}(undef, (0, 0))
    _delay::AbstractArray = Array{Float64}(undef, (0, 0))
    delayed_voltages::AbstractArray = Array{Float64}(undef, (0, 0))
end

DelayNetwork(neurons) = DelayNetwork(neurons=neurons)

function add_group!(network::DelayNetwork, N::Int; name::Union{String,Nothing}=nothing)
    network.delay = pad2D(network.delay, N)
    invoke(add_group!, Tuple{SpikingNetwork,Int}, network, N; name=name)
end

"""
Add a connection between two neuron groups in the network
"""
function connect!(
    network::DelayNetwork,
    source::NeuronGroup,
    target::NeuronGroup,
    weight::AbstractArray;
    adjacency::AbstractArray=ones(Bool, source.n, target.n),
    delay::AbstractArray=zeros(Bool, source.n, target.n)
)
    network.delay[source.idx, target.idx] = delay
    network._delay = zeros(size(network.delay)..., Int(maximum(network.delay) + 1))
    for i in 1:size(delay)[1]
        for j in 1:size(delay)[2]
            network._delay[i, j, Int(network.delay[i, j] + 1)] = 1
        end
    end
    network.delayed_voltages = zeros(size(network.delay)[1], Int(maximum(network.delay) + 1))
    invoke(
        connect!,
        Tuple{SpikingNetwork,NeuronGroup,NeuronGroup,AbstractArray},
        network, source, target, weight, ; adjacency=adjacency
    )
    return
end

function connect!(
    network::DelayNetwork,
    source::NeuronGroup,
    target::NeuronGroup,
    connection::Connection;
    delay::AbstractArray=zeros(Bool, source.n, target.n)
)
    connect!(network, source, target, connection.weight;
        adjacency=connection.adjacency, delay=delay)
end

function run!(
    network::DelayNetwork;
    input_spikes::Union{AbstractMatrix{Bool},Nothing}=nothing,
    input_voltage::Union{AbstractMatrix,Nothing}=nothing
)
    if network.batch_size > 1 #TODO: implement batch_size>1 for DelayNetwork
        error("`batch_size > 1` is not supported for DelayNetwork yet!")
    end
    # Evoke spikes
    network.spikes = network.voltage .>= network.neurons.v_thresh
    # External spikes
    if !isnothing(input_spikes)
        network.spikes = network.spikes .| input_spikes
    end
    # Supress neurons #TODO
    _update!(network, input_voltage)
    # update delayed voltages
    network.delayed_voltages[:, 1:end-1] = network.delayed_voltages[:, 2:end]
    network.delayed_voltages[:, end] .= 0
    # Learning process
    if network.learning # Apply the learning rule and update weights
        train!(network)
    end
    return network.spikes, network.voltage
end

function _update!(
    network::DelayNetwork{LIF},
    input_voltage::Union{AbstractMatrix,Nothing}
)
    # Decay voltages.
    network.voltage = (
        network.neurons.voltage_decay_factor .*
        (network.voltage .- network.neurons.v_rest)
        .+
        network.neurons.v_rest
    )
    # Get current
    current = network.spikes * network.weight  # + network.bias #TODO
    # External voltage:
    if !isnothing(input_voltage)
        input_voltage[network.refractory.>0] .= 0.0
        network.voltage += input_voltage
    end
    # update delayed voltages
    network.delayed_voltages += sum(
        network.spikes .* network.weight .* network._delay,
        dims=1
    )[1, :, :]
    current += transpose(network.delayed_voltages[:, 1])
    # update voltages
    network.voltage += current
    network.voltage[network.refractory.>0] .= network.neurons.v_rest # reset the voltage of the neurons in the refractory period
    network.voltage[network.spikes] .= network.neurons.v_reset  # change the voltage of spiked neurons to v_reset
    # Update refractory timepoints
    network.refractory .-= network.neurons.dt
    network.refractory[network.spikes] .= network.neurons.refractory_period
end

function _update!(
    network::DelayNetwork{Izhikevich},
    input_voltage::Union{AbstractMatrix,Nothing}
)

    # Get current
    current = network.spikes * network.weight
    # External voltage:
    if !isnothing(input_voltage)
        current += input_voltage
    end
    # update delayed voltages and current
    network.delayed_voltages += sum(
        transpose(network.spikes) .* network.weight .* network._delay,
        dims=1
    )[1, :, :]
    current += transpose(network.delayed_voltages[:, 1])
    # Update voltages (spiked neurons)
    network.voltage[network.spikes] .= (network.spikes.*network.neurons.c)[network.spikes]  # change the voltage of spiked neurons to c #TODO: optimize for scalar values
    network.recovery[network.spikes] .+= (network.spikes.*network.neurons.d)[network.spikes]  # add d to the recovery parameter of spiked neurons #TODO: optimize for scalar values
    # Update voltages
    for _ in 1:network.neurons.runge_kutta_order
        network.voltage += (network.neurons.dt / network.neurons.runge_kutta_order) .* (
            0.04 .* network.voltage .^ 2 + 5 .* network.voltage .+ 140 - network.recovery + current
        )
    end
    network.recovery += network.neurons.dt .* (
        network.neurons.a .* (network.neurons.b .* network.voltage - network.recovery)
    )
end


text:
function rxn_str(metnet::MetNet, ider::IDER_TYPE)
    ridx = rxnindex(metnet, ider)
    arrow_str = isblock(metnet, ridx) ? " >< " :
                isbkwd_bounded(metnet, ridx) ? " <== " :
                isfwd_bounded(metnet, ridx) ? " ==> " : " <==> "

    reacts = rxn_reacts(metnet, ridx)
    react_str = join([string("(", S(metnet, react, ridx), ") ",
        mets(metnet, react))  for react in reacts], " + ")

    prods = rxn_prods(metnet, ridx)
    prods_str = join([string("(", S(metnet, prod, ridx), ") ",
        mets(metnet, prod))  for prod in prods], " + ")
    return react_str * arrow_str * prods_str
end
rxn_str(metnet::MetNet, iders::Vector) = [rxn_str(metnet, ider) for ider in iders]
rxn_str(metnet::MetNet) = rxn_str(metnet, metnet.rxns)


text:
# Test that the way we specify boundary conditions works as expected
using MPI
using OrderedCollections
using StaticArrays
using Statistics
using Test

using CLIMAParameters
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

using ClimateMachine
using ClimateMachine.Land
using ClimateMachine.Land.SoilWaterParameterizations
using ClimateMachine.Land.Runoff
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.DGMethods: BalanceLaw, LocalGeometry
using ClimateMachine.MPIStateArrays
using ClimateMachine.GenericCallbacks
using ClimateMachine.ODESolvers
using ClimateMachine.VariableTemplates
using ClimateMachine.SingleStackUtils
using ClimateMachine.BalanceLaws:
    BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state

import ClimateMachine.DGMethods.FVReconstructions: FVLinear

@testset "NoRunoff" begin
    ClimateMachine.init()

    FT = Float64

    function init_soil_water!(land, state, aux, localgeo, time)
        myfloat = eltype(state)
        state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))
        state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))
    end

    soil_param_functions =
        SoilParamFunctions{FT}(porosity = 0.75, Ksat = 1e-7, S_s = 1e-3)
    surface_precip_amplitude = FT(3e-8)
    f = FT(pi * 2.0 / 300.0)
    precip = (t) -> surface_precip_amplitude * sin(f * t)
    ϑ_l0 = (aux) -> eltype(aux)(0.2)
    bc = LandDomainBC(
        bottom_bc = LandComponentBC(
            soil_water = Neumann((aux, t) -> eltype(aux)(0.0)),
        ),
        surface_bc = LandComponentBC(
            soil_water = SurfaceDrivenWaterBoundaryConditions(
                FT;
                precip_model = DrivenConstantPrecip{FT}(precip),
                runoff_model = NoRunoff(),
            ),
        ),
    )

    soil_water_model = SoilWaterModel(FT; initialϑ_l = ϑ_l0)
    soil_heat_model = PrescribedTemperatureModel()

    m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)
    sources = ()
    m = LandModel(
        param_set,
        m_soil;
        boundary_conditions = bc,
        source = sources,
        init_state_prognostic = init_soil_water!,
    )

    N_poly = 5
    nelem_vert = 50

    # Specify the domain boundaries
    zmax = FT(0)
    zmin = FT(-1)

    driver_config = ClimateMachine.SingleStackConfiguration(
        "LandModel",
        N_poly,
        nelem_vert,
        zmax,
        param_set,
        m;
        zmin = zmin,
        numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
    )

    t0 = FT(0)
    timeend = FT(150)
    dt = FT(0.05)

    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        ode_dt = dt,
    )
    n_outputs = 60
    mygrid = solver_config.dg.grid
    every_x_simulation_time = ceil(Int, timeend / n_outputs)
    state_types = (Prognostic(), Auxiliary(), GradientFlux())
    dons_arr =
        Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]
    time_data = FT[0] # store time data

    # We specify a function which evaluates `every_x_simulation_time` and returns
    # the state vector, appending the variables we are interested in into
    # `dons_arr`.

    callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do
        dons = dict_of_nodal_states(solver_config, state_types; interp = true)
        push!(dons_arr, dons)
        push!(time_data, gettime(solver_config.solver))
        nothing
    end

    # # Run the integration
    ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))

    dons = dict_of_nodal_states(solver_config, state_types; interp = true)
    push!(dons_arr, dons)
    push!(time_data, gettime(solver_config.solver))

    # Get z-coordinate
    z = get_z(solver_config.dg.grid; rm_dupes = true)
    N = length(dons_arr)
    # Note - we take the indices 2:N here to avoid the t = 0 spot. Gradients
    # are not calculated before the invoke! command, so we cannot compare at t=0.
    computed_surface_flux = [dons_arr[k]["soil.water.K∇h[3]"][end] for k in 2:N]
    t = time_data[2:N]
    prescribed_surface_flux = t -> FT(-1) * FT(3e-8 * sin(pi * 2.0 * t / 300.0))
    MSE = mean((prescribed_surface_flux.(t) .- computed_surface_flux) .^ 2.0)
    @test MSE < 5e-7
end

@testset "Explicit Dirichlet BC Comparison" begin
    # This tests the "if" branch of compute_surface_grad_bc
    ClimateMachine.init()
    FT = Float64
    function init_soil_water!(land, state, aux, localgeo, time)
        myfloat = eltype(state)
        state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))
        state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))
    end

    soil_param_functions = SoilParamFunctions{FT}(
        porosity = 0.495,
        Ksat = 0.0443 / (3600 * 100),
        S_s = 1e-4,
    )
    surface_precip_amplitude = FT(-5e-4)

    precip = (t) -> surface_precip_amplitude
    hydraulics = vanGenuchten{FT}(n = 2.0)
    function hydrostatic_profile(z, zm, porosity, n, α)
        myf = eltype(z)
        m = FT(1 - 1 / n)
        S = FT((FT(1) + (α * (z - zm))^n)^(-m))
        return FT(S * porosity)
    end
    N_poly = 2
    nelem_vert = 50

    # Specify the domain boundaries
    zmax = FT(0)
    zmin = FT(-0.35)
    Δz = abs(FT((zmin - zmax) / nelem_vert / 2))

    ϑ_l0 =
        (aux) -> eltype(aux)(hydrostatic_profile(
            aux.z,
            -0.35,
            0.495,
            hydraulics.n,
            hydraulics.α,
        ))

    soil_water_model = SoilWaterModel(
        FT;
        moisture_factor = MoistureDependent{FT}(),
        hydraulics = hydraulics,
        initialϑ_l = ϑ_l0,
    )
    soil_heat_model = PrescribedTemperatureModel()

    m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)
    sources = ()

    bc = LandDomainBC(
        bottom_bc = LandComponentBC(
            soil_water = Neumann((aux, t) -> eltype(aux)(0.0)),
        ),
        surface_bc = LandComponentBC(
            soil_water = SurfaceDrivenWaterBoundaryConditions(
                FT;
                precip_model = DrivenConstantPrecip{FT}(precip),
                runoff_model = CoarseGridRunoff{FT}(Δz),
            ),
        ),
    )

    m = LandModel(
        param_set,
        m_soil;
        boundary_conditions = bc,
        source = sources,
        init_state_prognostic = init_soil_water!,
    )

    driver_config = ClimateMachine.SingleStackConfiguration(
        "LandModel",
        N_poly,
        nelem_vert,
        zmax,
        param_set,
        m;
        zmin = zmin,
        numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
        #fv_reconstruction = FVLinear(),
    )

    t0 = FT(0)
    timeend = FT(60 * 60)
    dt = FT(0.1)

    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        ode_dt = dt,
    )
    state_types = (Prognostic(), Auxiliary(), GradientFlux())

    ClimateMachine.invoke!(solver_config;)# user_callbacks = (callback,))
    srf_dons = dict_of_nodal_states(solver_config, state_types; interp = true)

    ###### Repeat with explicit dirichlet BC
    soil_water_model2 = SoilWaterModel(
        FT;
        moisture_factor = MoistureDependent{FT}(),
        hydraulics = hydraulics,
        initialϑ_l = ϑ_l0,
    )
    soil_heat_model2 = PrescribedTemperatureModel()

    m_soil2 =
        SoilModel(soil_param_functions, soil_water_model2, soil_heat_model2)
    bc2 = LandDomainBC(
        bottom_bc = LandComponentBC(
            soil_water = Neumann((aux, t) -> eltype(aux)(0.0)),
        ),
        surface_bc = LandComponentBC(
            soil_water = Dirichlet((aux, t) -> eltype(aux)(0.495)),
        ),
    )

    m2 = LandModel(
        param_set,
        m_soil2;
        boundary_conditions = bc2,
        source = sources,
        init_state_prognostic = init_soil_water!,
    )

    driver_config2 = ClimateMachine.SingleStackConfiguration(
        "LandModel",
        N_poly,
        nelem_vert,
        zmax,
        param_set,
        m2;
        zmin = zmin,
        numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
        #  fv_reconstruction = FVLinear(),
    )

    solver_config2 = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config2,
        ode_dt = dt,
    )

    ClimateMachine.invoke!(solver_config2;)
    dir_dons = dict_of_nodal_states(solver_config2, state_types; interp = true)
    # Here we take the solution near the surface, where changes between the two
    # would occur.
    @test mean(
        (
            srf_dons["soil.water.ϑ_l"][(end - 15):end] .-
            dir_dons["soil.water.ϑ_l"][(end - 15):end]
        ) .^ 2.0,
    ) < 5e-5
end

@testset "Explicit Flux BC Comparison" begin
    # This tests the else branch of compute_surface_grad_bc
    ClimateMachine.init()
    FT = Float64
    function init_soil_water!(land, state, aux, localgeo, time)
        myfloat = eltype(state)
        state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))
        state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))
    end

    soil_param_functions = SoilParamFunctions{FT}(
        porosity = 0.495,
        Ksat = 0.0443 / (3600 * 100),
        S_s = 1e-4,
    )
    surface_precip_amplitude = FT(0)

    precip = (t) -> surface_precip_amplitude
    hydraulics = vanGenuchten{FT}(n = 2.0)
    function hydrostatic_profile(z, zm, porosity, n, α)
        myf = eltype(z)
        m = FT(1 - 1 / n)
        S = FT((FT(1) + (α * (z - zm))^n)^(-m))
        return FT(S * porosity)
    end

    ϑ_l0 =
        (aux) -> eltype(aux)(hydrostatic_profile(
            aux.z,
            -1,
            0.495,
            hydraulics.n,
            hydraulics.α,
        ))

    soil_water_model = SoilWaterModel(
        FT;
        moisture_factor = MoistureDependent{FT}(),
        hydraulics = hydraulics,
        initialϑ_l = ϑ_l0,
    )
    soil_heat_model = PrescribedTemperatureModel()

    m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)
    sources = ()
    N_poly = (1, 0)
    nelem_vert = 200

    # Specify the domain boundaries
    zmax = FT(0)
    zmin = FT(-1)
    Δz = FT(1 / 200 / 2)
    bc = LandDomainBC(
        bottom_bc = LandComponentBC(
            soil_water = Neumann((aux, t) -> eltype(aux)(0.0)),
        ),
        surface_bc = LandComponentBC(
            soil_water = SurfaceDrivenWaterBoundaryConditions(
                FT;
                precip_model = DrivenConstantPrecip{FT}(precip),
                runoff_model = CoarseGridRunoff{FT}(Δz),
            ),
        ),
    )

    m = LandModel(
        param_set,
        m_soil;
        boundary_conditions = bc,
        source = sources,
        init_state_prognostic = init_soil_water!,
    )

    driver_config = ClimateMachine.SingleStackConfiguration(
        "LandModel",
        N_poly,
        nelem_vert,
        zmax,
        param_set,
        m;
        zmin = zmin,
        numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
        fv_reconstruction = FVLinear(),
    )

    t0 = FT(0)
    timeend = FT(100000)
    dt = FT(4)

    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        ode_dt = dt,
    )
    state_types = (Prognostic(), Auxiliary(), GradientFlux())

    ClimateMachine.invoke!(solver_config;)

    srf_dons = dict_of_nodal_states(solver_config, state_types; interp = true)
    # Note - we only look at differences between the solutions at the surface, because
    # near the bottom they will agree.
    z = srf_dons["z"][150:end]
    error1 = sqrt(mean(
        (
            srf_dons["soil.water.ϑ_l"][150:end] .-
            hydrostatic_profile.(z, -1, 0.495, hydraulics.n, hydraulics.α)
        ) .^ 2.0,
    ))
    error2 = sqrt(mean(srf_dons["soil.water.K∇h[3]"][150:end] .^ 2.0))
    @test error1 < 1e-5
    @test error2 < eps(FT)
end


text:
using ThinningAndShift
using Documenter

DocMeta.setdocmeta!(ThinningAndShift, :DocTestSetup, :(using ThinningAndShift); recursive=true)

makedocs(;
    modules=[ThinningAndShift],
    authors="Dylan Festa",
    repo="https://github.com/dylanfesta/ThinningAndShift.jl/blob/{commit}{path}#{line}",
    sitename="ThinningAndShift.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://dylanfesta.github.io/ThinningAndShift.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/dylanfesta/ThinningAndShift.jl",
    devbranch="main",
)


text: abstract type DDMPreconditioner{T} <: DDMMatrix{T} end


text:
##############################################################################
##
## AbstractFixedEffectSolver
##
## this type must defined solve_residuals!, solve_coefficients!
##
##############################################################################
abstract type AbstractFixedEffectSolver{T} end

"""
Solve a least square problem for a set of FixedEffects

`solve_residuals!(y, fes, weights; method = :lsmr, maxiter = 10000, double_precision = true, tol = 1e-8)`

### Arguments
* `y` : A `AbstractVector` or an `AbstractMatrix`
* `fes`: A `Vector{<:FixedEffect}`
* `weights`: A `AbstractWeights`
* `method` : A `Symbol` for the method. Choices are :lsmr, :lsmr_threads, :lsmr_parallel, :lsmr_gpu (requires `CuArrays`. Use the option `double_precision = false` to use `Float32` on the GPU).
* `maxiter` : Maximum number of iterations
* `double_precision::Bool`: Should the demeaning operation use Float64 rather than Float32? Default to true.
* `tol` : Tolerance. Default to 1e-8 if `double_precision = true`, 1e-6 otherwise.

### Returns
* `res` :  Residual of the least square problem
* `iterations`: Number of iterations
* `converged`: Did the algorithm converge?

### Examples
```julia
using  FixedEffects
p1 = repeat(1:5, inner = 2)
p2 = repeat(1:5, outer = 2)
solve_residuals!(rand(10), [FixedEffect(p1), FixedEffect(p2)])
```
"""
function solve_residuals!(y::Union{AbstractVector{<: Number}, AbstractMatrix{<: Number}}, fes::AbstractVector{<: FixedEffect}, weights::AbstractWeights = Weights(Ones{eltype(y)}(size(y, 1)));
	method::Symbol = :lsmr, maxiter::Integer = 10000,
	double_precision::Bool = eltype(y) == Float64, tol::Real = double_precision ? 1e-8 : 1e-6)
	any((length(fe) != size(y, 1) for fe in fes)) && throw("FixedEffects must have the same length as y")
	any(ismissing.(fes)) && throw("FixedEffects must not have missing values")
	feM = AbstractFixedEffectSolver{double_precision ? Float64 : Float32}(fes, weights, Val{method})
	solve_residuals!(y, feM; maxiter = maxiter, tol = tol)
end

"""
Solve a least square problem for a set of FixedEffects

`solve_coefficients!(y, fes, weights; method = :lsmr, maxiter = 10000, double_precision = true, tol = 1e-8)`
d
### Arguments
* `y` : A `AbstractVector`
* `fes`: A `Vector{<:FixedEffect}`
* `weights`: A `AbstractWeights`
* `method` : A `Symbol` for the method. Choices are :lsmr, :lsmr_threads, :lsmr_parallel, :lsmr_gpu (requires `CuArrays`. Use the option `double_precision = false` to use `Float32` on the GPU).
* `maxiter` : Maximum number of iterations
* `double_precision::Bool`: Should the demeaning operation use Float64 rather than Float32? Default to true.
* `tol` : Tolerance. Default to 1e-8 if `double_precision = true`, 1e-6 otherwise.

### Returns
* `b` : Solution of the least square problem
* `iterations`: Number of iterations
* `converged`: Did the algorithm converge?
Fixed effects are generally not unique. We standardize the solution
in the following way: the mean of fixed effects within connected components is zero
(except for the first).
This gives the unique solution in the case of two fixed effects.

### Examples
```julia
using  FixedEffects
p1 = repeat(1:5, inner = 2)
p2 = repeat(1:5, outer = 2)
x = rand(10)
solve_coefficients!(rand(10), [FixedEffect(p1), FixedEffect(p2)])
```
"""
function solve_coefficients!(y::AbstractVector{<: Number}, fes::AbstractVector{<: FixedEffect}, weights::AbstractWeights  = Weights(Ones{eltype(y)}(length(y)));
	method::Symbol = :lsmr, maxiter::Integer = 10000,
	double_precision::Bool = eltype(y) == Float64, tol::Real = double_precision ? 1e-8 : 1e-6)
	any(ismissing.(fes)) && throw("Some FixedEffect has a missing value for reference or interaction")
	any((length(fe) != length(y) for fe in fes))  && throw("FixedEffects must have the same length as y")
	feM = AbstractFixedEffectSolver{double_precision ? Float64 : Float32}(fes, weights, Val{method})
	solve_coefficients!(y, feM; maxiter = maxiter, tol = tol)
end


text:
"""
Check if the matrix is square...
"""
function checkmatrix(m::Matrix)
    (size(m,1) != size(m,2) || !(eltype(m) <: Number)) && error("Only square matrices of Numbers accepted, sorry!!!");
    for r in 1:size(m,1)
        for c in 1:size(m,2)
            if r == c
                m[r,c] != 0 && error("The diagonal must be filled by 0s");
            else
                m[r,c] != m[c,r] && error("The matrix must be symmetrical");
            end
        end
    end
    return m;
end

#################################################################################

"""
Return all the possibles combinations of these items
"""
function all_possible_rutes(items::Vector{T})::Vector{Vector{T}} where T

    #combs
    combs = Vector{Vector{T}}();

    #base case
    if length(items) == 2
        return push!(combs, [items[2],items[1]],[items[1],items[2]]);
    end

    #recursion block
    for i in 1:length(items)
        push!(combs, pushfirst!.(all_possible_rutes(copy_and_remove(items,i)),items[i])...);
    end

    return combs;
end

#################################################################################

"""
get the total cost of a given rute. The cost include the travel from the last city back to home.
"""
function rute_total_cost(cmatrix::Matrix{T}, rute::Vector{Int}) where T

    cost = 0;
    #add travel cost
    for t in 1:length(rute)-1
        cost += cmatrix[rute[t],rute[t+1]]; # central part of the trip
    end
    return cost + cmatrix[rute[end], 1]; # the way home
end

#################################################################################

"""
get the total costs of the given rutes. The cost of each rute include the travel from the last city back to home.
"""
function rutes_total_costs(cmatrix::Matrix{T}, rutes::Vector{Vector{Int}})::Vector where T
    costs = Vector();
    for rute in rutes
       push!(costs,rute_total_cost(cmatrix, rute)); # pushing all the rutes costs
    end
    return costs;
end

#################################################################################
"""
return the costs of each stretch of the rute...
"""
function rute_stretches_costs(cmatrix::Matrix{T}, rute::Vector{Int}) where T
   costs = zeros(T,length(rute));
    for t in 1:length(rute)-1
        costs[t] = cmatrix[rute[t],rute[t+1]]; # central part of the trip
    end
    costs[end] = cmatrix[rute[end],1];
    return costs;
end

#################################################################################

"""
return the rute of minimal cost.
"""
function pick_cheaper(cmatrix::Matrix{T}, rutes::Vector{Vector{Int}})::Vector{Int} where T
    return rutes[findmin(rutes_total_costs(cmatrix, rutes))[2]];
end

#################################################################################

function copy_and_remove(a::AbstractVector,i::Int)
    return lazy_remove!(copy(a),i);
end

#################################################################################

"""
Remove an element from the array. Warning, it is very probable that this
methos unsort the array!!!
"""
function lazy_remove!(n::AbstractArray, i::Int)
    if i == length(n)
        pop!(n);
    else
        n[i] = pop!(n);
    end
    return n;
end

#################################################################################

"""
Pick the rute from the bit matrix
"""
function pick_rute!(bm)
        r = Vector{Int}()
        bml = size(bm,1)
        d1 = 1;
        for i in 1:bml
            for d2 in 1:bml
                if bm[d1,d2] != 0
                    # println("Action in ", d1, " , ", d2);

                    #Save d1
                    push!(r,d1);
                    # println(d1, " pushed ", r);

                    #Clear
                    # println("Clearing")
                    # display(bm);
                    bm[d1,d2] = 0;
                    bm[d2,d1] = 0;
                    # println("Cleared")
                    # display(bm);

                    #Jump to d2
                    # println("Jumping to ", d2)
                    d1 = d2;

                    #break
                    # println("Break")
                    # println();
                    break;
                end
            end
        end
        return r;
end

#################################################################################

"""
Get the number of possible rutes given a number of cities to visit...
"""
function rutescount(ncities::T)::T where {T <: Integer}
    if ncities <= 1 error("ncities must be greater than 1")
    elseif ncities == 2 return 1
    else
        return (ncities - 1) * rutescount(ncities - 1)
    end
end

#################################################################################

function reverserute(r)
    rr = reverse(r);
    po = pop!(rr);
    return pushfirst!(rr,po);
end


text:
struct BinnedMarginalDensity # this encodes the estimated density..
    marginal::Vector{Float64} # measure after truncation & binning
    marginal_grid::Vector{Float64}
    marginal_h::Float64
end

function BinnedMarginalDensity(d::NormalConvolutionProblem)
    BinnedMarginalDensity(d.marginal, d.marginal_grid, d.marginal_h)
end

function BinnedMarginalDensity(M_bd::Float64, marginal_grid, marginal_h)
    BinnedMarginalDensity(fill(M_bd, size(marginal_grid)), marginal_grid, marginal_h);
end


text:
#module LinearPredict

  export linPred,linPred2

  function Toeplitz(a::Array{Complex{Float64},1})
         m=length(a)
     T=[(k>=l) ? a[1+mod(k-l,m)] : conj(a[1+mod(l-k,m)]) for k=1:m,l=1:m]
     return(T)
  end

  function autoCorr(a::Array{Complex{Float64},1},p::Int64)
      c=zeros(Complex{Float64},p)
       for l=0:(p-1),k=1:(length(a)-l)
          c[l+1]+=a[k]*conj(a[k+l])
      end
      return( c )
  end

  function linPred(a::Array{Complex{Float64},1},p::Int,n::Int)
    lena=length(a)
    r=autoCorr(a,p+1)
    coeff = conj(reverse(Toeplitz(r[1:p,1]) \ r[2:p+1,1]))
    # print(size(coeff))
    # print(coeff)
    ap=zeros(Complex{Float64},lena+n) ; ap[1:lena]=a
    for k=(lena+1):(lena+n)
        ap[k] = sum(coeff.*ap[(k-p):(k-1)])
    end
    return(ap)
  end

  function linPred2(a::Array{Complex{Float64},2},p::Int,n::Int)
    (r,s)=size(a)
    ap=zeros(Complex{Float64},r,s+n)
    for k=1:r
        ap[k,:]=linPred(a[k,:],p,n)
    end
    return(ap)
  end

#end

max_stars_repo_name: marcel-utz/NMR

max_issues_repo_name: marcel-utz/NMR


text:
export schedule
export schedule_randomly, schedule_by_id, schedule_fastest, schedule_partially, schedule_by_property, schedule_by_type
export random_activation, by_id, fastest, partial_activation, property_activation, by_type
"""
    schedule(model)
Return an iterator over the scheduled IDs using the model's scheduler.
"""
schedule(model::ABM) = model.scheduler(model)

####################################
# Schedulers
####################################
"""
    schedule_fastest
A scheduler that activates all agents once per step in the order dictated by the
agent's container, which is arbitrary (the keys sequence of a dictionary).
This is the fastest way to activate all agents once per step.
"""
schedule_fastest(model::ABM) = keys(model.agents)

@deprecate fastest schedule_fastest

"""
    schedule_by_id
A scheduler that activates all agents agents at each step according to their id.
"""
function schedule_by_id(model::ABM)
    agent_ids = sort(collect(keys(model.agents)))
    return agent_ids
end

@deprecate by_id schedule_by_id
@deprecate as_added schedule_by_id

"""
    schedule_randomly
A scheduler that activates all agents once per step in a random order.
Different random ordering is used at each different step.
"""
function schedule_randomly(model::ABM)
    order = shuffle(model.rng, collect(keys(model.agents)))
end

@deprecate random_activation schedule_randomly

"""
    schedule_partially(p)
A scheduler that at each step activates only `p` percentage of randomly chosen agents.
"""
function schedule_partially(p::Real)
    function partial(model::ABM)
        ids = collect(keys(model.agents))
        return randsubseq(model.rng, ids, p)
    end
    return partial
end

@deprecate partial_activation schedule_partially

"""
    schedule_by_property(property)
A scheduler that at each step activates the agents in an order dictated by their `property`,
with agents with greater `property` acting first. `property` can be a `Symbol`, which
just dictates which field of the agents to compare, or a function which inputs an agent
and outputs a real number.
"""
function schedule_by_property(p)
    function by_property(model::ABM)
        ids = collect(keys(model.agents))
        properties = [Agents.get_data(model[id], p) for id in ids]
        s = sortperm(properties)
        return ids[s]
    end
end

@deprecate property_activation schedule_by_property

"""
    schedule_by_type(shuffle_types::Bool, shuffle_agents::Bool)
A scheduler useful only for mixed agent models using `Union` types.
- Setting `shuffle_types = true` groups by agent type, but randomizes the type order.
Otherwise returns agents grouped in order of appearance in the `Union`.
- `shuffle_agents = true` randomizes the order of agents within each group, `false` returns
the default order of the container (equivalent to [`schedule_fastest`](@ref)).
"""
function schedule_by_type(shuffle_types::Bool, shuffle_agents::Bool)
    function by_union(model::ABM{S,A}) where {S,A}
        types = union_types(A)
        sets = [Integer[] for _ in types]
        for agent in allagents(model)
            idx = findfirst(t -> t == typeof(agent), types)
            push!(sets[idx], agent.id)
        end
        shuffle_types && shuffle!(model.rng, sets)
        shuffle_agents && [shuffle!(model.rng, set) for set in sets]
        vcat(sets...)
    end
end

@deprecate by_type schedule_by_type

"""
    schedule_by_type((C, B, A), shuffle_agents::Bool)
A scheduler that activates agents by type in specified order (since `Union`s are not order
preserving). `shuffle_agents = true` randomizes the order of agents within each group.
"""
function schedule_by_type(order::Tuple{Type,Vararg{Type}}, shuffle_agents::Bool)
    function by_ordered_union(model::ABM{S,A}) where {S,A}
        types = union_types(A)
        if order !== nothing
            @assert length(types) == length(order) "Invalid dimension for `order`"
            types = order
        end
        sets = [Integer[] for _ in types]
        for agent in allagents(model)
            idx = findfirst(t -> t == typeof(agent), types)
            push!(sets[idx], agent.id)
        end
        shuffle_agents && [shuffle!(model.rng, set) for set in sets]
        vcat(sets...)
    end
end