text:
# site chosen with DirectCR, rx or hop chosen with Direct

############################ DirectCRDirect ###################################
const MINJUMPRATE = 2.0^exponent(1e-12)

#NOTE state vector u is a matrix. u[i,j] is species i, site j
#NOTE hopping_constants is a matrix. hopping_constants[i,j] is species i, site j
mutable struct DirectCRDirectJumpAggregation{J,T,RX,HOP,RNG,DEPGR,VJMAP,JVMAP,SS,U<:PriorityTable,W<:Function} <: AbstractSSAJumpAggregator
    next_jump::SpatialJump{J} #some structure to identify the next event: reaction or hop
    prev_jump::SpatialJump{J} #some structure to identify the previous event: reaction or hop
    next_jump_time::T
    end_time::T
    rx_rates::RX
    hop_rates::HOP
    site_rates::Vector{T}
    # rates::F1 #rates for constant-rate jumps
    # affects!::F2 #affects! function determines the effect of constant-rate jumps
    save_positions::Tuple{Bool,Bool}
    rng::RNG
    dep_gr::DEPGR #dep graph is same for each site
    vartojumps_map::VJMAP #vartojumps_map is same for each site
    jumptovars_map::JVMAP #jumptovars_map is same for each site
    spatial_system::SS
    numspecies::Int #number of species
    rt::U
    ratetogroup::W
end

function DirectCRDirectJumpAggregation(nj::SpatialJump{J}, njt::T, et::T, rx_rates::RX, hop_rates::HOP, site_rates::Vector{T}, sps::Tuple{Bool,Bool}, rng::RNG, spatial_system::SS; num_specs, minrate=convert(T,MINJUMPRATE), vartojumps_map=nothing, jumptovars_map=nothing, dep_graph=nothing, kwargs...) where {J,T,RX,HOP,RNG,SS}

    # a dependency graph is needed
    if dep_graph === nothing
        dg = make_dependency_graph(num_specs, rx_rates.ma_jumps)
    else
        dg = dep_graph
        # make sure each jump depends on itself
        add_self_dependencies!(dg)
    end

    # a species-to-reactions graph is needed
    if vartojumps_map === nothing
        vtoj_map = var_to_jumps_map(num_specs, rx_rates.ma_jumps)
    else
        vtoj_map = vartojumps_map
    end

    if jumptovars_map === nothing
        jtov_map = jump_to_vars_map(rx_rates.ma_jumps)
    else
        jtov_map = jumptovars_map
    end

    # mapping from jump rate to group id
    minexponent = exponent(minrate)

    # use the largest power of two that is <= the passed in minrate
    minrate = 2.0^minexponent
    ratetogroup = rate -> priortogid(rate, minexponent)

    # construct an empty initial priority table -- we'll reset this in init
    rt = PriorityTable(ratetogroup, zeros(T, 1), minrate, 2*minrate)

    DirectCRDirectJumpAggregation{J,T,RX,HOP,RNG,typeof(dg),typeof(vtoj_map),typeof(jtov_map),SS,typeof(rt), typeof(ratetogroup)}(nj, nj, njt, et, rx_rates, hop_rates, site_rates, sps, rng, dg, vtoj_map, jtov_map, spatial_system, num_specs, rt, ratetogroup)
end

############################# Required Functions ##############################
# creating the JumpAggregation structure (function wrapper-based constant jumps)
function aggregate(aggregator::DirectCRDirect, starting_state, p, t, end_time, constant_jumps, ma_jumps, save_positions, rng; hopping_constants, spatial_system, kwargs...)
    num_species = size(starting_state,1)
    majumps = ma_jumps
    if majumps === nothing
        majumps = MassActionJump(Vector{typeof(end_time)}(), Vector{Vector{Pair{Int,Int}}}(), Vector{Vector{Pair{Int,Int}}}())
    end

    next_jump = SpatialJump{Int}(typemax(Int),typemax(Int),typemax(Int)) #a placeholder
    next_jump_time = typemax(typeof(end_time))
    rx_rates = RxRates(num_sites(spatial_system), majumps)
    hop_rates = HopRates(hopping_constants, spatial_system)
    site_rates = zeros(typeof(end_time), num_sites(spatial_system))

    DirectCRDirectJumpAggregation(next_jump, next_jump_time, end_time, rx_rates, hop_rates, site_rates, save_positions, rng, spatial_system; num_specs = num_species, kwargs...)
end

# set up a new simulation and calculate the first jump / jump time
function initialize!(p::DirectCRDirectJumpAggregation, integrator, u, params, t)
    p.end_time = integrator.sol.prob.tspan[2]
    fill_rates_and_get_times!(p, u, t)
    generate_jumps!(p, integrator, params, u, t)
    nothing
end

# calculate the next jump / jump time
function generate_jumps!(p::DirectCRDirectJumpAggregation, integrator, params, u, t)
    p.next_jump_time  = t + randexp(p.rng) / p.rt.gsum
    p.next_jump_time >= p.end_time && return nothing
    site = sample(p.rt, p.site_rates, p.rng)
    p.next_jump = sample_jump_direct(p, site)
    nothing
end

# execute one jump, changing the system state
function execute_jumps!(p::DirectCRDirectJumpAggregation, integrator, u, params, t)
    # execute jump
    update_state!(p, integrator)

    # update current jump rates and times
    update_dependent_rates_and_firing_times!(p, integrator.u, t)
    nothing
end

######################## SSA specific helper routines ########################
"""
    fill_rates_and_get_times!(aggregation::DirectCRDirectJumpAggregation, u, t)

reset all stucts, reevaluate all rates, repopulate the priority table
"""
function fill_rates_and_get_times!(aggregation::DirectCRDirectJumpAggregation, u, t)
    @unpack spatial_system, rx_rates, hop_rates, site_rates, rt = aggregation

    reset!(rx_rates)
    reset!(hop_rates)
    site_rates .= zero(typeof(t))

    rxs = 1:num_rxs(rx_rates)
    species = 1:aggregation.numspecies

    for site in 1:num_sites(spatial_system)
        update_rx_rates!(rx_rates, rxs, u, site)
        update_hop_rates!(hop_rates, species, u, site, spatial_system)
        site_rates[site] = total_site_rate(rx_rates, hop_rates, site)
    end
    # setup PriorityTable
    reset!(rt)
    for (pid,priority) in enumerate(site_rates)
        insert!(rt, pid, priority)
    end
    nothing
end

"""
    update_dependent_rates_and_firing_times!(p, u, t)

recalculate jump rates for jumps that depend on the just executed jump (p.prev_jump)
"""
function update_dependent_rates_and_firing_times!(p::DirectCRDirectJumpAggregation, u, t)
    site_rates = p.site_rates
    jump = p.prev_jump
    if is_hop(p, jump)
        source_site = jump.src
        target_site = jump.dst
        update_rates_after_hop!(p, u, source_site, target_site, jump.jidx)

        # update site rates
        oldrate = site_rates[source_site]
        site_rates[source_site] = total_site_rate(p.rx_rates, p.hop_rates, source_site)
        update!(p.rt, source_site, oldrate, site_rates[source_site])

        oldrate = site_rates[target_site]
        p.site_rates[target_site] = total_site_rate(p.rx_rates, p.hop_rates, target_site)
        update!(p.rt, target_site, oldrate, site_rates[target_site])
    else
        site = jump.src
        update_rates_after_reaction!(p, u, site, reaction_id_from_jump(p,jump))

        # update site rates
        oldrate = site_rates[site]
        site_rates[site] = total_site_rate(p.rx_rates, p.hop_rates, site)
        update!(p.rt, site, oldrate, site_rates[site])
    end
end

"""
    num_constant_rate_jumps(aggregator::DirectCRDirectJumpAggregation)

number of constant rate jumps
"""
num_constant_rate_jumps(aggregator::DirectCRDirectJumpAggregation) = 0


text:
module NodeGraphEditor

using Genie, Logging, LoggingExtras
using JLD2
using JSON3
using StructTypes

const CACHE_DIR = ".nodegraph_cache"
const NG_FILE_NAME = "data.jld2"
const NG_FILE = joinpath(CACHE_DIR, NG_FILE_NAME)
const PTYPES_FILE_NAME = "ptypes"
const PTYPES_FILE = joinpath(CACHE_DIR, PTYPES_FILE_NAME)
const NNID_FILE_NAME = "nnid"
const NNID_FILE = joinpath(CACHE_DIR, NNID_FILE_NAME)

include("graph.jl")
include("node_params_API.jl")
include("editor_formatting.jl")
include("file_interface.jl")
include("operations.jl")

export edit_graph

function __init__()
    Genie.genie(; context = @__MODULE__)
    Core.eval(Main, :(using Genie))
    isdir(CACHE_DIR) || mkdir(CACHE_DIR)
end

"""
    edit_graph(nodegraph::NodeGraph = NodeGraph(), param_type_symbols::Vector{Symbol} = [])

Edit the given `nodegraph`. Valid [`type_symbol`](@ref)s for node types should be provided
through `param_type_symbols`. These will be available in the "Add Node" menu in the editor.
This runs a webserver on `localhost:8000`. Navigating to this URL will open the editor.
"""
function edit_graph(
    ng::NodeGraph = NodeGraph(),
    param_type_symbols::Vector{Symbol} = []
)
    nn_id = 0
    for node in ng.nodes
        nn_id = max(nn_id, node.id + 1)
    end
    set_nodegraph(ng)
    set_param_types(param_type_symbols)
    set_new_node_id(nn_id)
    Genie.up()
end

end

ext: jl


text:
using DataStructures

"""
    alloc_half_mms(V::Additive, C::Counts; α=0.5)

Find a 1/2-approximate MMS allocation that obeys the constraints imposed by C.
First the instance is reduced to an ordered normalized instance where each good
is worth less than 1/2. While there are more than one agent remaining, the
algorithm creates a bundle with the ``⌊length(category)/n⌋`` lowest-valued
items in each category. Repeatedly, it converts each of these to the
highest-valued remaining item in the category until it either runs out of items
to convert or an agent values the bundle at least 1/2. If the procedure runs
out of items to convert, it adds the highest-valued remaining item in each
category, in order, to get ``⌈length(category)/n⌉`` items from each category.
After each such item is added, the value is again checked for each agent. Since
the instance was ordered normalized and without items worth 1/2 or more, the
bundle created will always be worth more than 1/2 to one of the remaining
agents before the procedure runs out of items to add to it or convert from low-
to high-valued.

Another approximation ratio, `α`, can be supplied. If `α ≤ 0.5` the algorithm is
guaranteed to succeed. Otherwise, the method will try to find an allocation with
an approximation ratio of `α`, but may fail. In the latter case, the results
will indicate a failure by setting `res.fail` to `true`.
"""
function alloc_half_mms(V::Additive, C::Counts; α=0.5)
    R = order(V, C)

    # Normalize and allocate items worth α or more.
    R′ = reduce(profile(R), constraint(R), α)

    R = chain(R, R′)

    V, C = profile(R), constraint(R)
    N, n = agents(V), na(V)

    while n > 1
        # Fill the bundle with the floor(k_h/n) least valuable items
        B = union(Set{Int}(), [c[end - floor_n(c, n) + 1:end] for c in C]...)

        # Convert lower value items to higher value items
        categories = Iterators.Stateful(C)
        c, converted = popfirst!(categories), 0
        while all(i -> value(V, i, B) < α, N)
            if converted == floor_n(c, n)
                isempty(categories) && break

                c, converted = popfirst!(categories), 0
                continue
            end

            # Convert by removing the lowest valued remaining item and adding
            # the  highest valued item not in the bundle.
            converted += 1
            setdiff!(B, Set(c[end-converted+1]))
            B = B ∪ c[converted]
        end

        # Add higher value items
        categories = Iterators.Stateful(C)
        while all((i) -> value(V, i, B) < α, N)
            # If there are no more categories, the supplied α does not work
            if isempty(categories)
                return (alloc=nothing, fail=true)
            end

            c = popfirst!(categories)

            if length(c) % n > 0
                B = B ∪ c[ceil_n(c, n)]
            end
        end

        R′ = reduce(V, C, findfirst(i -> value(V, i, B) ≥ α, N), B)
        R = chain(R, R′)
        V, C = profile(R), constraint(R)
        N, n = agents(V), na(V)
    end

    # Use an empty allocation if there are no more agents remaining
    if n == 0
        A = Allocation()
    else
        # If there is not enough value remaining for the last agent, the
        # supplied α does not work
        if value(V, 1, items(V)) < α
            return (alloc=nothing, fail=true)
        end

        A = Allocation(V)
        give!(A, 1, items(V))
    end

    # Create an allocation in the original instance
    return (alloc=transform(R, A), fail=false)

end

"""
    alloc_rand(V)

A straightforward lottery that allocates the items randomly to the agents. For
each item, its agent is selected uniformly at random. The valuation profile `V`
is not used, other than to determine the number of agents and items. The return
value is a named tuple with the field `alloc` (the `Allocation`).
"""
alloc_rand(V) = alloc_rand(na(V), ni(V))

"""
    alloc_rand(n::Int, m::Int)

Same as `alloc_rand(V)`, for `n` agents and `m` items.
"""
function alloc_rand(n::Int, m::Int)

    A = Allocation(n, m)

    N, M = agents(A), items(A)

    for g in items(A)
        give!(A, rand(N), g)
    end

    return (alloc=A,)

end

"""
    alloc_rand(V, C::Conflicts)

Allocate items to agents randomly, respecting the item conflicts. Uses the
[randomized coloring procedure with
symmetry-breaking](https://doi.org/10.1007/978-3-540-70575-8_26) of Pemmaraju
and Srinivasan, which works as follows:

1. Give the items random priorities, corresponding to a permutation selected
   uniformly at ramdom.
2. Tentatively allocate each item randomly to an agent, without concern for the
   item conflicts.
3. If an agent has received conflicting items, it keeps the highest-priority
   item (i.e., earliest in the permutation), and the others are reallocated
   arbitrarily.

This final arbitrary reallocation is also performed randomly in this
implementation, by going through the items in random order, allocating each to a
randomly selected agent among those able to receive it.

The valuation profile `V` is not used, other than to determine the number of
agents and items.

For this algorithm to function properly, the maximum degree of the conflict
graph should be strictly less than the number of agents.
"""
alloc_rand(V, C::Conflicts) = alloc_rand(na(V), ni(V), C)

"""
    alloc_rand(n::Int, m::Int, C::Conflicts)

Same as `alloc_rand(V, C)`, for `n` agents and `m` items.
"""
function alloc_rand(n::Int, m::Int, C::Conflicts)

    G = graph(C)

    @assert Δ(G) < n

    # Random permutation of the items, for priority:
    π = randperm(m)

    # Tentative allocation:
    A = alloc_rand(n, m).alloc

    removed = Set{Int}()

    # In every conflict, remove lower-priority items:
    for e in edges(G)

        g, h = src(e), dst(e)

        (owned(A, g) && owned(A, h)) || continue

        i, j = owner(A, g), owner(A, h)

        i == j || continue

        g = π[g] < π[h] ? h : g

        deny!(A, i, g)
        push!(removed, g)

    end

    # Allocate each item to some agent able to receive it:
    for g in shuffle(collect(removed))

        conflicted = neighbors(G, g)

        for i in randperm(n)

            isempty(conflicted ∩ bundle(A, i)) || continue

            give!(A, i, g)

            break

        end

    end

    return (alloc=A,)

end

"""
    alloc_bkv18_1(V; randpri=true)

The first algorithm (**Alg-Identical**) described by Barman, Krishnamurty and
Vaish in their 2018 paper [Greedy Algorithms for Maximizing Nash Social
Welfare](https://doi.org/10.1145/3355902). The algorithm finds a
1.061-approximate MNW allocation when agents have identical valuations, i.e.,
for any agents `i`, `j` and item `g`, `value(V, i, g) == value(V, j, g)`. (This
approximation ratio applies to the geometric mean of agent utilities, not the
raw product.) The result will also be envy-free up to any item (EFX).

The algorithm follows a straightforward greedy allocation procedure, where in
each iteration, the most valuable item is allocated to the agent with the lowest
utility. By default, ties are broken by giving the agents random priorities; if
`randpri` is set to false, they are instead broken lexicographically (as
specified by Barman et al.), so that the agent with the lower index is
preferred.
"""
function alloc_bkv18_1(V; randpri=true)

    N, M = agents(V), items(V)
    v = [value(V, 1, g) for g in M]

    π = randpri ? shuffle(N) : N

    pq = PriorityQueue(i => (0, p) for (i, p) in zip(N, π))

    A = Allocation(V)

    for g in sortperm(v, rev=true)
        i, (u, p) = peek(pq)
        give!(A, i, g)
        pq[i] = (u + v[g], p)
    end

    return (alloc = A,)

end

# The original descriptions of the following algorithm leaves a great latitude
# in how one actually implements it. The gist of the procedure is a local
# search, where an allocation is incrementally improved along paths in a graph,
# which has edges added or removed as items are reallocated. Here one might,
# e.g., use one of the structures developed for dynamically updating transitive
# closures. Because the graph in question is one with the agents as its
# vertices, which means it will probably be quite small, any asymptotic
# improvements could quite possibly be swamped by overhead. For now, things are
# kept simple, with a fresh traversal for each agent pair.

"""
    alloc_bkv18_2(V; randpri=true, complete=false)
    alloc_hpps20_1(V; randpri=true, complete=false) # alias

The second algorithm (**Alg-Binary**) described by Barman, Krishnamurty and
Vaish in their 2018 paper [Greedy Algorithms for Maximizing Nash Social
Welfare](https://doi.org/10.1145/3355902). The algorithm finds MNW allocations
in polynomial time for binary additive valuations, i.e., where each agent values
any given object at 0 or 1 (e.g., an `Additive{BitMatrix}`). It also works in a
more general setting, where `value(V, i, S)`, for any given `i`, is a concave
function of the number of items `g` in `S` for which `value(V, i, g) == 1`.

The original algorithm builds on an initial allocation, but does not specify
what this allocation should be. It also does not deal with the case where one or
more agents ends up with zero utility; in fact the procedure will not work even
if we start with two or more agents with zero utility in the intial allocation.
The strategy followed here is the same as that of Caragiannis et al.
(https://doi.org/10.1145/3355902), where a maximum cardinality set of agents
achieving positive utility is found using bipartite matching (with no fairness
considerations). The remaining items are randomly allocated to agents among
these that value them, if any. Remaining agents and items are ignored by the
procedure.

Following the algorithm of Barman et al., the tie-breaking procedure (Algorithm
1) of Halpern et al. (http://arxiv.org/abs/2007.06073) is used, where the MNW
allocation is transformed into the lexicographically greatest MNW, according to
some ordering of the agents, providing group-strategyproofness (GSP) in addition
to the EF1 and PO guarantees that follow from MNW. By default, the agent
ordering/priority is random; if this randomization is turned off (with `randpri`
set to false), the default ordering is used, with agent `1` receiving the
highest priority, etc.

!!! note

    Despite the use of randomization here, by default, this is the
    *deterministic* procedure of Halpern et al. They also describe a randomized
    procedure, which functions in an entirely different manner.

Finally, if the `complete` argument is set to `true`, the allocation is
completed with `fill_even!` (which means that some agents that must necessarily
get a utility of zero can still receive items valued zero, if that evens out the
bundle cardinalities). Note that this undermines the GSP guarantee, which
requires that these items be discarded. The return value is a named tuple with
the fields `alloc` (the `Allocation`) and `mnw` (the Nash welfare, ignoring
agents with zero utility).
"""
function alloc_bkv18_2(V; randpri=true, complete=false)

    n, m = na(V), ni(V)

    # Tentative allocation, where O[g] is the owner of g. Will assign positive
    # utility to as many agents as possible (the set N, found through maximum
    # bipartite matching betewen agents and items they value), and will allocate
    # items only to them.
    O = zeros(Int, m)

    N = Int[]
    for (i, g) in bipartite_matching(V)
        O[g] = i
        push!(N, i)
    end

    M = Int[]
    for g in 1:m
        # To simplify the logic below, and ideally to reduce the number of
        # iterations, we assign items only to agents who value them, if any. If
        # no agents value an item, it will not participate in the procedure
        # anyway, and is left unallocated until the end of the function.
        fans = [i for i in N if value(V, i, g) > 0]
        isempty(fans) && continue # Leave g unallocated for now
        push!(M, g)
        O[g] == 0 && (O[g] = rand(fans))
    end

    # U[i]: The utility of agent i
    U = zeros(Int, n)
    for g = M
        i = O[g]
        U[i] += value(V, i, g)
    end

    # Pre-allocation
    G = zeros(Int, n, n)
    Π = zeros(Int, n, n)

    # Update the graph G based on current ownership.
    function update()

        # A form of adjacency matrix, where any G[i, j] > 0 is an arbitrary item
        # we may transfer from j to i.
        G .= 0
        for i in N, g in M
            (value(V, i, g) == 0 || O[g] == i) && continue
            G[i, O[g]] = g
        end

    end

    # We should never need more than this (as shown by Barman, Krishnamurty and
    # Vaish). We'll add one iteration, which should never be reached, as a
    # failsafe (with an assertion, below the loop):
    maxiter = ceil(Int, 2m*(n+1)*log(n*m))

    it = 0
    for outer it = 1:maxiter + 1

        update()

        # Predecessor matrix, where any Π[i, j] > 0 is the predecessor of j
        # along the path (sequence of possible transfers) from i.
        Π .= 0
        for i in N, j in N
            G[i, j] > 0 && (Π[i, j] = i)
        end

        # Straightforward version of Warshall's algorithm, using only Π and
        # restricting nodes of interest to N, i.e., the agents with nonzero
        # utility.
        for k in N, i in N, j in N
            (i == j || Π[i, j] > 0) && continue
            if Π[i, k] > 0 && Π[k, j] > 0
                Π[i, j] = Π[k, j]
            end
        end

        # The augmentation procedure of Barman et al., except that we only care
        # about agents with nonzero utility. Only items that are valued by at
        # least one agent will be used in Π. Such items are initially allocated
        # to agents that value them, and will remain so after each update, which
        # means that transfering the item will necessarily increment U[i] and
        # decrement U[j].
        s = t = 0
        best = 0
        for i in N, j in N
            @assert U[i] ≥ 1 && U[j] ≥ 1
            (Π[i, j] > 0 && U[j] > 1) || continue
            change = (U[i] + 1)*(U[j] - 1)/(U[i]*U[j])
            if change > best
                best = change
                s, t = i, j
            end
        end

        # No improvement, so we're done
        best > 1 || break

        # There was an improvement, and the transfer of one unit of utility.
        U[t] -= 1; U[s] += 1

        # Pass items back along path from s to t
        while t ≠ s
            p = Π[s, t]
            @assert O[G[p, t]] == t
            O[G[p, t]] = p
            t = p
        end

    end

    @assert it ≤ maxiter # Make sure we found optimum

    # We now have an MNW allocation, but want to applie the tiebreaker from
    # HPPS20(1), ending up with the maximum MNW allocation in lexicographic
    # order, given some priority ordering of the agents.

    # Normally, the priorities are randomized (and function primarily to prevent
    # group strategyproofness), but that may be turned off, in which case agents
    # are prioritized by their index.
    if randpri
        shuffle!(N) # Random priorities
    else
        sort!(N)    # If i < j, i is prioritized
    end

    # Find a path (in the same graph as before) from i (with priority p) to some
    # node j with lower priority and utility U[i] + 1. If one is found, pass
    # items back to increment U[i] and decrement U[j]. (For simplicity, we
    # increment/decrement along the path as well, though that cancels out.)
    #
    # We could have use Warshall here as well, but we'll potentially be
    # modifying the graph between every start node, so a straight DFS is
    # probably less wasteful.
    function search(p, i, u, seen=falses(n))
        seen[i] = true
        for (q, j) in enumerate(N)
            (G[i, j] > 0 && !seen[j]) || continue
            target = (q > p && U[j] == u + 1)
            if target || search(q, j, u, seen)
                @assert O[G[i, j]] == j
                O[G[i, j]] = i
                U[j] -= 1; U[i] += 1
                return true
            end
        end
        return false
    end

    # In order of priority, p:
    for (p, i) in enumerate(N)
        update()
        search(p, i, U[i])
    end

    # We now have the appropriate MNW allocation, represented by the ownership
    # vector O, and create the resulting Allocation.

    A = Allocation(n, m)
    for g in M
        give!(A, O[g], g)
    end

    if complete
        # Distribute the unvalued items as well, somewhat evenly. Note that this
        # ruins the group strategyproofness guarantees of HPPS20.
        fill_even!(A)
    end

    return (alloc=A, mnw=nash_welfare(V, A))

end

# Rather than BKV18(2) with HPPS20(1) for tie-breaking, one can view the
# function as HPPS20(1) with BKV18(2) for finding the initial allocation.
const alloc_hpps20_1 = alloc_bkv18_2

"""
    alloc_ghss18_4(V::Submodular, MMSs)

The fourth algorithm (**Algorithm 4**) described by Ghodsi et al. in the 2018
paper [Fair allocation of Indivisible Goods: Improvements and
Generalizations](https://arxiv.org/abs/1704.00222). The algorithm finds a
1/3-approximate MMS allocation for a given submodular instance and corresponding
maximin shares for the agents (`MMSs[i]` should be the MMS of agent `i`). If the
supplied maximin shares, are higher than the actual maximin shares, the method
may fail. In that case, this will be indicated in the result, where `res.fail`
will be set to true and `res.agent` will be set to the agent last considered
when the method failed to improve. If the maximin shares are unknown, use
`alloc_ghss18_4b`.
"""
function alloc_ghss18_4(V::Submodular, MMSs)

    @assert na(V) == length(MMSs) "There must be one MMS per agent"

    # Scale valuations so that each agent's MMS is 1
    Vo = V
    V = Submodular([B -> (value(Vo, i, B) / MMSs[i]) for i in agents(Vo)], ni(Vo))

    # Allocate all items worth at least 1/3
    R = reduce(V, 1/3)
    V = profile(R)

    A = Allocation(V)

    # In case all agents have received an item worth at least 1/3
    if na(V) == 0
        A = transform(R, A)

        # The algorithm does not specify what to do in this situation. We
        # therefore simply allocate the goods randomly to the agents.
        fill_random!(A)

        return (alloc=A, fail=false, agent=0)
    end

    # Create an arbitrary allocation with the remaining items by for each item
    # allocating it to a random agent
    fill_random!(A)

    N, M = agents(V), items(V)

    # All valuations from now on will be wrapped by the ceiling function for 2/3
    V′ = Submodular([B -> min(2/3, value(V, i, B)) for i in N], ni(V))

    # Continue until all agents have a bundle valued at least 1/3
    while any(i -> value(V′, i, bundle(A, i)) < 1/3, N)
        # Want to improve the bundle of the worst of agent
        i = argmin([value(V′, i, bundle(A, i)) for i in N])

        found = false
        for g in M
            j = owner(A, g)
            i == j && continue

            Bᵢ, Bⱼ = bundle(A, i), bundle(A, j)

            # Only the contribution of agents `i` and `j` to `ex(A)` (the
            # function defined by Ghodsi et al.) will change. Thus, we do not
            # need to compute the entire value and can simply check the
            # difference in the contributions of agents `i` and `j` before and
            # after moving `g`.
            v = value(V′, i, Bᵢ) + value(V′, j, Bⱼ)
            v′ = value(V′, i, Bᵢ ∪ g) + value(V′, j, symdiff(Bⱼ, g))

            # If the change in total value is more than 1/3m when moving g from
            # i to j, perform the move.
            if v′ ≥ v + 1/(3*ni(V))
                found = true
                deny!(A, j, g)
                give!(A, i, g)
                break
            end
        end

        # If no sufficently large improvement can be made, then the MMS values
        # are incorrect (too large)
        if !found
            return (alloc=nothing, fail=true, agent=agent(R, i))
        end
    end

    # Convert the allocation to one for the original instance
    return (alloc=transform(R, A), fail=false, agent=0)

end

"""
    alloc_ghss18_4b(V::Submodular; a=3, x_warn=true)

A variation on the fourth algorithm (**Algorithm 4**) described by Ghodsi et al.
in the 2018 paper [Fair allocation of Indivisible Goods: Improvements and
Generalizations](https://arxiv.org/abs/1704.00222). The algorithm finds a
1/3-approximate MMS allocation for a given submodular instance. The method
starts by overestimating the MMS of each agent and slowly decreasing the MMS of
specific agents until `alloc_ghss18_4` returns an allocation.

The amount that the MMS of an agent should be reduced by in each iteration is
not specified by Ghodsi et al. One can show that if the factor is `1/(1 + 1/x)`,
where `x ≥ 3n - 1`, then the algorithm will successfully find a 1/3-approximate
MMS allocation. One way to show this, is to modify Lemma 4.6 in their paper to
assume that each of the bundles `Sᵢ` is valued at least `1/(1 + 1/x)`. Using
this modified version of Lemma 4.6, one can modify the proof of Theorem 4.7 to
show that as long as `x ≥ 3n - 1`, the change in expectance from moving an item
is at least `1/(3m)`. The value of `x` used in this implementation is `x = an`,
where the keyword argument `a` is set to `3` by default (i.e., `x = 3n`). If `a`
is set so that `x < 3n - 1` a warning will be given. The warning can be turned
off by setting `x_warn` to `false`.
"""
function alloc_ghss18_4b(V::Submodular; a=3, x_warn=true)

    if x_warn && a * na(V) < 3 * na(V) - 1
        @warn("The value of `a` may be too small")
    end

    MMSs = Vector{Float64}([value(V, i, items(V)) for i in agents(V)])

    res = alloc_ghss18_4(V, MMSs)
    while res.fail
        MMSs[res.agent] /= (1 + 1/(a * na(V)))
        res = alloc_ghss18_4(V, MMSs)
    end

    return (alloc=res.alloc,)

end

"""
    alloc_bb18_3(V::Additive, C::Counts; a=3, ghss18_4b_warn=true)

The 1/3-approximate MMS-allocation under cardinality constraints algorithm
(Section 5) described by Biswas and Barman in their 2018 paper [Fair Division
Under Cardinality Constraints] (https://doi.org/10.24963/ijcai.2018/13). Finds a
1/3-approximate MMS allocation for an instance of the fair allocation problem
under cardinality constraints by converting the additive instance under
cardinality constraints to a submodular instance without cardinality
constraints. The allocation is then found by using the method of Ghodsi et al.
(`alloc_ghhs18_4b`), with possible reallocation of items to satisfy the
constraints. Both keyword arguments, `a` and `ghss18_4b_warn`, are passed
directly to `alloc_ghhs18_4b` as respectively the keyword arguments `a` and
`x_warn`. See [`alloc_ghhs18_4b`](@ref) for documentation on how to use them.
"""
function alloc_bb18_3(V::Additive, C::Counts; a=3, ghss18_4b_warn=true)

    # The submodular valuation function construction for V and C
    function submodular_valuation(i, B)
        total = 0

        for c in C
            # Extract the overlap between B and the category. We do not care
            # about which items overlap, only their value.
            overlap = [value(V, i, g) for g in c.members ∩ B]

            # An agent cannot get value from more items than the threshold of
            # the given category.
            n_items = min(threshold(c), length(overlap))

            # Give the agent the `n_items` most valuable items in the overlap.
            total += sum(partialsort!(overlap, 1:n_items, rev=true))
        end

        return total
    end

    N = agents(V)

    # Create a submodular instance and use the method of Ghodsi et al.
    V′ = Submodular([B -> submodular_valuation(i, B) for i in N], ni(V))
    res = alloc_ghss18_4b(V′, a=a, x_warn=ghss18_4b_warn)
    A = res.alloc

    # Easy way to find overlap between agent's bundle and a category
    overlap(i, c) = bundle(A, i) ∩ c.members

    # It is possible that the 1/3-approximate MMS allocation produced by
    # `alloc_ghss18_4b` does not adhere to the cardinality constraints. If this
    # is the case, then we for each non-adhering bundle, give away the
    # corresponding agent's least-preferred items in the bundle from any
    # category for which the bundle breaks the category's threshold. This does
    # not result in a decrease in the agent's perceived value of its bundle
    # based on the submodular valuations.
    for i in N, c in C
        B = overlap(i, c)
        # If there are more items than allowed
        while length(B) > threshold(c)
            # Give away the item in c that `i` prefers the least
            g = minimum(g -> (value(V, i, g), g), B)[2]

            # Give the item away to any agent that is allowed more items from
            # the category
            i′ = findfirst(i′ -> length(overlap(i′, c)) < threshold(c), N)

            deny!(A, i, g)
            give!(A, i′, g)
            delete!(B, g)
        end
    end

    return (alloc=A,)

end


text:
function polar_kernel(n::Int, nlines::Int, line_start::Int, scale::Float64,
                     u_curr::CuDeviceArray{Float64,1}, v_curr::CuDeviceArray{Float64,1},
                     l_curr::CuDeviceArray{Float64,1}, rho::CuDeviceArray{Float64,1},
                     shift_lines::Int, param::CuDeviceArray{Float64,2},
                     _YffR::CuDeviceArray{Float64,1}, _YffI::CuDeviceArray{Float64,1},
                     _YftR::CuDeviceArray{Float64,1}, _YftI::CuDeviceArray{Float64,1},
                     _YttR::CuDeviceArray{Float64,1}, _YttI::CuDeviceArray{Float64,1},
                     _YtfR::CuDeviceArray{Float64,1}, _YtfI::CuDeviceArray{Float64,1},
                     frBound::CuDeviceArray{Float64,1}, toBound::CuDeviceArray{Float64,1})

    tx = threadIdx().x
    ty = threadIdx().y
    I = blockIdx().x
    id_line = I + shift_lines

    x = @cuDynamicSharedMem(Float64, n)
    xl = @cuDynamicSharedMem(Float64, n, n*sizeof(Float64))
    xu = @cuDynamicSharedMem(Float64, n, (2*n)*sizeof(Float64))

    @inbounds begin
        YffR = _YffR[id_line]; YffI = _YffI[id_line]
        YftR = _YftR[id_line]; YftI = _YftI[id_line]
        YttR = _YttR[id_line]; YttI = _YttI[id_line]
        YtfR = _YtfR[id_line]; YtfI = _YtfI[id_line]

        pij_idx = line_start + 8*(I-1)

        xl[1] = sqrt(frBound[2*(id_line-1)+1])
        xu[1] = sqrt(frBound[2*id_line])
        xl[2] = sqrt(toBound[2*(id_line-1)+1])
        xu[2] = sqrt(toBound[2*id_line])
        xl[3] = -2*pi
        xu[3] = 2*pi
        xl[4] = -2*pi
        xu[4] = 2*pi

        x[1] = min(xu[1], max(xl[1], sqrt(u_curr[pij_idx+4])))
        x[2] = min(xu[2], max(xl[2], sqrt(u_curr[pij_idx+5])))
        x[3] = min(xu[3], max(xl[3], u_curr[pij_idx+6]))
        x[4] = min(xu[4], max(xl[4], u_curr[pij_idx+7]))

        param[1,id_line] = l_curr[pij_idx]
        param[2,id_line] = l_curr[pij_idx+1]
        param[3,id_line] = l_curr[pij_idx+2]
        param[4,id_line] = l_curr[pij_idx+3]
        param[5,id_line] = l_curr[pij_idx+4]
        param[6,id_line] = l_curr[pij_idx+5]
        param[7,id_line] = l_curr[pij_idx+6]
        param[8,id_line] = l_curr[pij_idx+7]
        param[9,id_line] = rho[pij_idx]
        param[10,id_line] = rho[pij_idx+1]
        param[11,id_line] = rho[pij_idx+2]
        param[12,id_line] = rho[pij_idx+3]
        param[13,id_line] = rho[pij_idx+4]
        param[14,id_line] = rho[pij_idx+5]
        param[15,id_line] = rho[pij_idx+6]
        param[16,id_line] = rho[pij_idx+7]
        param[17,id_line] = v_curr[pij_idx]
        param[18,id_line] = v_curr[pij_idx+1]
        param[19,id_line] = v_curr[pij_idx+2]
        param[20,id_line] = v_curr[pij_idx+3]
        param[21,id_line] = v_curr[pij_idx+4]
        param[22,id_line] = v_curr[pij_idx+5]
        param[23,id_line] = v_curr[pij_idx+6]
        param[24,id_line] = v_curr[pij_idx+7]

        CUDA.sync_threads()

        status, minor_iter = tron_kernel(n, shift_lines, 500, 200, 1e-6, scale, true, x, xl, xu,
                                         param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)

        vi_vj_cos = x[1]*x[2]*cos(x[3] - x[4])
        vi_vj_sin = x[1]*x[2]*sin(x[3] - x[4])

        u_curr[pij_idx] = YffR*x[1]^2 + YftR*vi_vj_cos + YftI*vi_vj_sin
        u_curr[pij_idx+1] = -YffI*x[1]^2 - YftI*vi_vj_cos + YftR*vi_vj_sin
        u_curr[pij_idx+2] = YttR*x[2]^2 + YtfR*vi_vj_cos - YtfI*vi_vj_sin
        u_curr[pij_idx+3] = -YttI*x[2]^2 - YtfI*vi_vj_cos - YtfR*vi_vj_sin
        u_curr[pij_idx+4] = x[1]^2
        u_curr[pij_idx+5] = x[2]^2
        u_curr[pij_idx+6] = x[3]
        u_curr[pij_idx+7] = x[4]
    end

    return
end

function polar_kernel_cpu(n::Int, nline::Int, line_start::Int, scale::Float64,
                          u_curr::AbstractVector{Float64}, v_curr::AbstractVector{Float64},
                          l_curr::AbstractVector{Float64}, rho::AbstractVector{Float64},
                          shift::Int,
                          param::Array{Float64},
                          YffR::Array{Float64}, YffI::Array{Float64},
                          YftR::Array{Float64}, YftI::Array{Float64},
                          YttR::Array{Float64}, YttI::Array{Float64},
                          YtfR::Array{Float64}, YtfI::Array{Float64},
                          frBound::Array{Float64}, toBound::Array{Float64})
    avg_minor_it = 0

    x = zeros(n)
    xl = zeros(n)
    xu = zeros(n)

    xl[3] = -2*pi
    xu[3] = 2*pi
    xl[4] = -2*pi
    xu[4] = 2*pi

    @inbounds for I=1:nline
        pij_idx = line_start + 8*(I-1)
        id_line = shift + I

        xl[1] = sqrt(frBound[2*(id_line-1)+1])
        xu[1] = sqrt(frBound[2*id_line])
        xl[2] = sqrt(toBound[2*(id_line-1)+1])
        xu[2] = sqrt(toBound[2*id_line])

        x[1] = min(xu[1], max(xl[1], sqrt(u_curr[pij_idx+4])))
        x[2] = min(xu[2], max(xl[2], sqrt(u_curr[pij_idx+5])))
        x[3] = min(xu[3], max(xl[3], u_curr[pij_idx+6]))
        x[4] = min(xu[4], max(xl[4], u_curr[pij_idx+7]))

        param[1,id_line] = l_curr[pij_idx]
        param[2,id_line] = l_curr[pij_idx+1]
        param[3,id_line] = l_curr[pij_idx+2]
        param[4,id_line] = l_curr[pij_idx+3]
        param[5,id_line] = l_curr[pij_idx+4]
        param[6,id_line] = l_curr[pij_idx+5]
        param[7,id_line] = l_curr[pij_idx+6]
        param[8,id_line] = l_curr[pij_idx+7]
        param[9,id_line] = rho[pij_idx]
        param[10,id_line] = rho[pij_idx+1]
        param[11,id_line] = rho[pij_idx+2]
        param[12,id_line] = rho[pij_idx+3]
        param[13,id_line] = rho[pij_idx+4]
        param[14,id_line] = rho[pij_idx+5]
        param[15,id_line] = rho[pij_idx+6]
        param[16,id_line] = rho[pij_idx+7]
        param[17,id_line] = v_curr[pij_idx]
        param[18,id_line] = v_curr[pij_idx+1]
        param[19,id_line] = v_curr[pij_idx+2]
        param[20,id_line] = v_curr[pij_idx+3]
        param[21,id_line] = v_curr[pij_idx+4]
        param[22,id_line] = v_curr[pij_idx+5]
        param[23,id_line] = v_curr[pij_idx+6]
        param[24,id_line] = v_curr[pij_idx+7]

        function eval_f_cb(x)
            f = eval_f_polar_kernel_cpu(id_line, scale, x, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return f
        end

        function eval_g_cb(x, g)
            eval_grad_f_polar_kernel_cpu(id_line, scale, x, g, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        function eval_h_cb(x, mode, rows, cols, _scale, lambda, values)
            eval_h_polar_kernel_cpu(id_line, x, mode, scale, rows, cols, lambda, values,
                                    param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        nele_hess = 10
        tron = ExaTron.createProblem(4, xl, xu, nele_hess, eval_f_cb, eval_g_cb, eval_h_cb;
                                     :tol => 1e-6, :matrix_type => :Dense, :max_minor => 200,
                                     :frtol => 1e-12)

        tron.x .= x
        status = ExaTron.solveProblem(tron)
        x .= tron.x
        avg_minor_it += tron.minor_iter

        cos_ij = cos(x[3] - x[4])
        sin_ij = sin(x[3] - x[4])
        vi_vj_cos = x[1]*x[2]*cos_ij
        vi_vj_sin = x[1]*x[2]*sin_ij

        u_curr[pij_idx] = YffR[id_line]*x[1]^2 + YftR[id_line]*vi_vj_cos + YftI[id_line]*vi_vj_sin
        u_curr[pij_idx+1] = -YffI[id_line]*x[1]^2 - YftI[id_line]*vi_vj_cos + YftR[id_line]*vi_vj_sin
        u_curr[pij_idx+2] = YttR[id_line]*x[2]^2 + YtfR[id_line]*vi_vj_cos - YtfI[id_line]*vi_vj_sin
        u_curr[pij_idx+3] = -YttI[id_line]*x[2]^2 - YtfI[id_line]*vi_vj_cos - YtfR[id_line]*vi_vj_sin
        u_curr[pij_idx+4] = x[1]^2
        u_curr[pij_idx+5] = x[2]^2
        u_curr[pij_idx+6] = x[3]
        u_curr[pij_idx+7] = x[4]
    end

    return 0, avg_minor_it / nline
end

function polar_kernel_two_level(
    n::Int, nline::Int, line_start::Int, bus_start::Int, scale::Float64,
    u::CuDeviceArray{Float64}, xbar::CuDeviceArray{Float64},
    z::CuDeviceArray{Float64}, l::CuDeviceArray{Float64}, rho::CuDeviceArray{Float64},
    shift_lines::Int, param::CuDeviceArray{Float64},
    _YffR::CuDeviceArray{Float64}, _YffI::CuDeviceArray{Float64},
    _YftR::CuDeviceArray{Float64}, _YftI::CuDeviceArray{Float64},
    _YttR::CuDeviceArray{Float64}, _YttI::CuDeviceArray{Float64},
    _YtfR::CuDeviceArray{Float64}, _YtfI::CuDeviceArray{Float64},
    frBound::CuDeviceArray{Float64}, toBound::CuDeviceArray{Float64},
    brBusIdx::CuDeviceArray{Int}
)
    I = blockIdx().x
    id_line = I + shift_lines

    x = @cuDynamicSharedMem(Float64, n)
    xl = @cuDynamicSharedMem(Float64, n, n*sizeof(Float64))
    xu = @cuDynamicSharedMem(Float64, n, (2*n)*sizeof(Float64))

    @inbounds begin
        YffR = _YffR[id_line]; YffI = _YffI[id_line]
        YftR = _YftR[id_line]; YftI = _YftI[id_line]
        YttR = _YttR[id_line]; YttI = _YttI[id_line]
        YtfR = _YtfR[id_line]; YtfI = _YtfI[id_line]

        pij_idx = line_start + 8*(I-1)
        xbar_pij_idx = line_start + 4*(I-1)

        xl[1] = sqrt(frBound[2*(id_line-1)+1])
        xu[1] = sqrt(frBound[2*id_line])
        xl[2] = sqrt(toBound[2*(id_line-1)+1])
        xu[2] = sqrt(toBound[2*id_line])
        xl[3] = -2*pi
        xu[3] = 2*pi
        xl[4] = -2*pi
        xu[4] = 2*pi

        x[1] = min(xu[1], max(xl[1], sqrt(u[pij_idx+4])))
        x[2] = min(xu[2], max(xl[2], sqrt(u[pij_idx+5])))
        x[3] = min(xu[3], max(xl[3], u[pij_idx+6]))
        x[4] = min(xu[4], max(xl[4], u[pij_idx+7]))

        fr_idx = bus_start + 2*(brBusIdx[2*(id_line-1)+1] - 1)
        to_idx = bus_start + 2*(brBusIdx[2*id_line] - 1)

        param[1,id_line] = l[pij_idx]
        param[2,id_line] = l[pij_idx+1]
        param[3,id_line] = l[pij_idx+2]
        param[4,id_line] = l[pij_idx+3]
        param[5,id_line] = l[pij_idx+4]
        param[6,id_line] = l[pij_idx+5]
        param[7,id_line] = l[pij_idx+6]
        param[8,id_line] = l[pij_idx+7]
        param[9,id_line] = rho[pij_idx]
        param[10,id_line] = rho[pij_idx+1]
        param[11,id_line] = rho[pij_idx+2]
        param[12,id_line] = rho[pij_idx+3]
        param[13,id_line] = rho[pij_idx+4]
        param[14,id_line] = rho[pij_idx+5]
        param[15,id_line] = rho[pij_idx+6]
        param[16,id_line] = rho[pij_idx+7]
        param[17,id_line] = xbar[xbar_pij_idx] - z[pij_idx]
        param[18,id_line] = xbar[xbar_pij_idx+1] - z[pij_idx+1]
        param[19,id_line] = xbar[xbar_pij_idx+2] - z[pij_idx+2]
        param[20,id_line] = xbar[xbar_pij_idx+3] - z[pij_idx+3]
        param[21,id_line] = xbar[fr_idx] - z[pij_idx+4]
        param[22,id_line] = xbar[to_idx] - z[pij_idx+5]
        param[23,id_line] = xbar[fr_idx+1] - z[pij_idx+6]
        param[24,id_line] = xbar[to_idx+1] - z[pij_idx+7]

        CUDA.sync_threads()

        status, minor_iter = tron_kernel(n, shift_lines, 500, 200, 1e-6, scale, true, x, xl, xu,
                                         param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)

        vi_vj_cos = x[1]*x[2]*cos(x[3] - x[4])
        vi_vj_sin = x[1]*x[2]*sin(x[3] - x[4])

        u[pij_idx] = YffR*x[1]^2 + YftR*vi_vj_cos + YftI*vi_vj_sin
        u[pij_idx+1] = -YffI*x[1]^2 - YftI*vi_vj_cos + YftR*vi_vj_sin
        u[pij_idx+2] = YttR*x[2]^2 + YtfR*vi_vj_cos - YtfI*vi_vj_sin
        u[pij_idx+3] = -YttI*x[2]^2 - YtfI*vi_vj_cos - YtfR*vi_vj_sin
        u[pij_idx+4] = x[1]^2
        u[pij_idx+5] = x[2]^2
        u[pij_idx+6] = x[3]
        u[pij_idx+7] = x[4]
    end

    return
end

function polar_kernel_two_level_cpu(
    n::Int, nline::Int, line_start::Int, bus_start::Int, scale::Float64,
    u::AbstractVector{Float64}, xbar::AbstractVector{Float64},
    z::AbstractVector{Float64}, l::AbstractVector{Float64}, rho::AbstractVector{Float64},
    shift::Int,
    param::Array{Float64},
    YffR::Array{Float64}, YffI::Array{Float64},
    YftR::Array{Float64}, YftI::Array{Float64},
    YttR::Array{Float64}, YttI::Array{Float64},
    YtfR::Array{Float64}, YtfI::Array{Float64},
    frBound::Array{Float64}, toBound::Array{Float64},
    brBusIdx::Array{Int}
)
    avg_minor_it = 0

    x = zeros(n)
    xl = zeros(n)
    xu = zeros(n)

    xl[3] = -2*pi
    xu[3] = 2*pi
    xl[4] = -2*pi
    xu[4] = 2*pi

    @inbounds for I=1:nline
        pij_idx = line_start + 8*(I-1)
        xbar_pij_idx = line_start + 4*(I-1)
        id_line = shift + I

        xl[1] = sqrt(frBound[2*(id_line-1)+1])
        xu[1] = sqrt(frBound[2*id_line])
        xl[2] = sqrt(toBound[2*(id_line-1)+1])
        xu[2] = sqrt(toBound[2*id_line])

        x[1] = min(xu[1], max(xl[1], sqrt(u[pij_idx+4])))
        x[2] = min(xu[2], max(xl[2], sqrt(u[pij_idx+5])))
        x[3] = min(xu[3], max(xl[3], u[pij_idx+6]))
        x[4] = min(xu[4], max(xl[4], u[pij_idx+7]))

        fr_idx = bus_start + 2*(brBusIdx[2*(id_line-1)+1] - 1)
        to_idx = bus_start + 2*(brBusIdx[2*id_line] - 1)

        param[1,id_line] = l[pij_idx]
        param[2,id_line] = l[pij_idx+1]
        param[3,id_line] = l[pij_idx+2]
        param[4,id_line] = l[pij_idx+3]
        param[5,id_line] = l[pij_idx+4]
        param[6,id_line] = l[pij_idx+5]
        param[7,id_line] = l[pij_idx+6]
        param[8,id_line] = l[pij_idx+7]
        param[9,id_line] = rho[pij_idx]
        param[10,id_line] = rho[pij_idx+1]
        param[11,id_line] = rho[pij_idx+2]
        param[12,id_line] = rho[pij_idx+3]
        param[13,id_line] = rho[pij_idx+4]
        param[14,id_line] = rho[pij_idx+5]
        param[15,id_line] = rho[pij_idx+6]
        param[16,id_line] = rho[pij_idx+7]
        param[17,id_line] = xbar[xbar_pij_idx] - z[pij_idx]
        param[18,id_line] = xbar[xbar_pij_idx+1] - z[pij_idx+1]
        param[19,id_line] = xbar[xbar_pij_idx+2] - z[pij_idx+2]
        param[20,id_line] = xbar[xbar_pij_idx+3] - z[pij_idx+3]
        param[21,id_line] = xbar[fr_idx] - z[pij_idx+4]
        param[22,id_line] = xbar[to_idx] - z[pij_idx+5]
        param[23,id_line] = xbar[fr_idx+1] - z[pij_idx+6]
        param[24,id_line] = xbar[to_idx+1] - z[pij_idx+7]

        function eval_f_cb(x)
            f = eval_f_polar_kernel_cpu(id_line, scale, x, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return f
        end

        function eval_g_cb(x, g)
            eval_grad_f_polar_kernel_cpu(id_line, scale, x, g, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        function eval_h_cb(x, mode, rows, cols, _scale, lambda, values)
            eval_h_polar_kernel_cpu(id_line, x, mode, scale, rows, cols, lambda, values,
                                    param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        nele_hess = 10
        tron = ExaTron.createProblem(4, xl, xu, nele_hess, eval_f_cb, eval_g_cb, eval_h_cb;
                                     :tol => 1e-6, :matrix_type => :Dense, :max_minor => 200,
                                     :frtol => 1e-12)

        tron.x .= x
        status = ExaTron.solveProblem(tron)
        x .= tron.x
        avg_minor_it += tron.minor_iter

        vi_vj_cos = x[1]*x[2]*cos(x[3] - x[4])
        vi_vj_sin = x[1]*x[2]*sin(x[3] - x[4])

        u[pij_idx] = YffR[id_line]*x[1]^2 + YftR[id_line]*vi_vj_cos + YftI[id_line]*vi_vj_sin
        u[pij_idx+1] = -YffI[id_line]*x[1]^2 - YftI[id_line]*vi_vj_cos + YftR[id_line]*vi_vj_sin
        u[pij_idx+2] = YttR[id_line]*x[2]^2 + YtfR[id_line]*vi_vj_cos - YtfI[id_line]*vi_vj_sin
        u[pij_idx+3] = -YttI[id_line]*x[2]^2 - YtfI[id_line]*vi_vj_cos - YtfR[id_line]*vi_vj_sin
        u[pij_idx+4] = x[1]^2
        u[pij_idx+5] = x[2]^2
        u[pij_idx+6] = x[3]
        u[pij_idx+7] = x[4]
    end

    return 0, avg_minor_it / nline
end


text:
module TestIndirectImports

using Base: PkgId
using InteractiveUtils: code_llvm
using MacroTools
using Pkg
using Test
using UUIDs

using IndirectImports
using IndirectImports: IndirectFunction, IndirectPackage, isloaded, _uuidfor

macro test_thrown(ex)
    quote
        let err = nothing
            try
                $(esc(ex))
            catch err
            end
            @test err isa Exception
            err
        end
    end
end

@testset "Macros" begin
    @test isexpr((@eval @macroexpand @indirect import A), :const)
    @test isexpr(unblock((@eval @macroexpand @indirect import A: x)), :let)
    @test isexpr(unblock((@eval @macroexpand @indirect import A: x, y)), :let)

    let err = @test_thrown _uuidfor(Main, :dummy)
        @test occursin(
            "does not have associated source code file",
            sprint(showerror, err))
    end
end

function devtest(uuid, name)
    if Base.locate_package(PkgId(UUID(uuid), name)) === nothing
        Pkg.develop(PackageSpec(
            name = name,
            path = joinpath(@__DIR__, name),
        ))
    end
end

devtest("20db8cd4-68a4-11e9-2de0-29cd367489cf", "_TestIndirectImportsUpstream")
devtest("63e77324-6b0a-11e9-11e4-8be33209e5fa", "_TestIndirectImportsDownstream")
devtest("32ce5e6c-7227-11e9-3206-ad1ab32cb15a", "_TestIndirectImportsDownstream2")
devtest("f64c0b4a-788b-11e9-0aea-4d12ea983b28", "_TestIndirectImportsPong")
devtest("f40e7e44-788b-11e9-0103-df8b12c08352", "_TestIndirectImportsPing")
using _TestIndirectImportsUpstream
using _TestIndirectImportsDownstream
using _TestIndirectImportsDownstream2
using _TestIndirectImportsPing
using _TestIndirectImportsPong

const Upstream = _TestIndirectImportsUpstream
const Downstream = _TestIndirectImportsDownstream
const Downstream_Upstream =
    _TestIndirectImportsDownstream._TestIndirectImportsUpstream

@testset "Core" begin
    @test Upstream.fun(1) == 2
    @test Upstream.fun === Downstream_Upstream.fun
    @test Val(Upstream.fun) isa Val{Upstream.fun}

    @testset for name in [:f1, :f2, :f3, :f4, :f5, :f6]
        @test getproperty(Downstream, name) === getproperty(Downstream_Upstream, name)
    end

    # A method defined in Upstream:
    @test Upstream.fun(1im) == 1 + 2im

    @test Downstream.dispatch(Downstream_Upstream.fun) === :fun
    @test Downstream.dispatch(Downstream_Upstream.fun2) === :fun2

    @test_throws(
        ErrorException("Only the top-level module can be indirectly imported."),
        IndirectPackage(Module()),
    )

    let err = @test_thrown @eval @indirect sin() = nothing
        @test occursin(
            "Function name `sin` does not refer to an indirect function.",
            sprint(showerror, err))
    end

    let err = @test_thrown @eval @indirect non_existing_function_name() = nothing
    end

    let err = @test_thrown @eval @indirect struct Spam end
        @test occursin("Cannot handle:", sprint(showerror, err))
    end

    let err = @test_thrown @eval @indirect import A=x, y
        @test occursin("Cannot handle:", sprint(showerror, err))
    end
end

@testset "Accessors" begin
    pkg = _TestIndirectImportsDownstream._TestIndirectImportsUpstream
    @test IndirectPackage(pkg.fun) === pkg
    @test nameof(pkg) === :_TestIndirectImportsUpstream
    @test nameof(pkg.fun) === :fun
    @test PkgId(Test) === PkgId(IndirectPackage(Test))
end

@testset "Ping-Pong Fibonacci" begin
    desired = [1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
    xs = 1:length(desired)
    @test _TestIndirectImportsPing.ping === _TestIndirectImportsPong.ping
    @test _TestIndirectImportsPing.pong === _TestIndirectImportsPong.pong
    @testset for f in [
        _TestIndirectImportsPing.ping
        _TestIndirectImportsPong.pong
    ]
        @test f.(xs) == desired
    end
end

struct Voldemort end
Base.nameof(::Voldemort) = error("must not be named")
IndirectImports.IndirectPackage(pkg::Voldemort) = pkg

@testset "Printing" begin
    upstreamname = "_TestIndirectImportsUpstream"
    @test repr(Upstream.fun) == "$upstreamname.fun"
    @test repr(IndirectPackage(Test).fun) == "Test.fun"

    @testset "2-arg `show` MUST NOT fail" begin
        f = IndirectFunction(Voldemort(), :fun)
        @debug "repr(f) = $(repr(f))"
        @test match(r".*\bIndirectFunction\{.*Voldemort\(\), *:fun\}",
                    repr(f)) !== nothing
    end

    # A fake package that is not loaded:
    pkg = IndirectPackage(UUID("7f75c6e9-3b46-4b36-93ee-72f09c6fb1e2"),
                          :NotLoaded)
    @test !isloaded(pkg)
    @test sprint(show, pkg.fun; context=:color=>true) ==
        "\e[31mNotLoaded\e[39m.fun"  # `NotLoaded` in red

    # But `Test` is a genuine so it's loaded:
    @test isloaded(IndirectPackage(Test))
    @test sprint(show, IndirectPackage(Test).fun; context=:color=>true) ==
        "\e[32mTest\e[39m.fun"  # `Test` in green
end

onetoten(config) = Upstream.reduceop(config, 0, 1:10)

@testset "Code gen" begin
    @test onetoten(Downstream.Config1()) == 55
    @test onetoten(Downstream.Config2()) == -55
    llvm = sprint(code_llvm, onetoten, Tuple{Downstream.Config1})
    @test occursin(r"i(32|64) 55", llvm)
end

end  # module


text:
using GridapGmsh
using GridapGmsh: gmsh, GmshDiscreteModel

gmsh.initialize()
gmsh.model.add("pincell")

factory = gmsh.model.geo

# in cm
p = 1.6        # pitch
p_2 = p / 2
ri = 0.5       # internal radius
t = 0.1        # wall thickness
ro = ri + t    # external radius

lc = 4e-2

# inner circle
factory.addPoint(p_2, p_2, 0, lc, 1)       # center
factory.addPoint(p_2 + ri, p_2, 0, lc, 2)  # right
factory.addPoint(p_2, p_2 + ri, 0, lc, 3)  # up
factory.addPoint(p_2 - ri, p_2, 0, lc, 4)  # left
factory.addPoint(p_2, p_2 - ri, 0, lc, 5)  # down
factory.addCircleArc(2, 1, 3, 1)
factory.addCircleArc(3, 1, 4, 2)
factory.addCircleArc(4, 1, 5, 3)
factory.addCircleArc(5, 1, 2, 4)
factory.addCurveLoop([1, 2, 3, 4], 1)

# outer circle
# factory.addPoint(p_2, p_2, 0, lc, 1)     # center
factory.addPoint(p_2 + ro, p_2, 0, lc, 6)  # right
factory.addPoint(p_2, p_2 + ro, 0, lc, 7)  # up
factory.addPoint(p_2 - ro, p_2, 0, lc, 8)  # left
factory.addPoint(p_2, p_2 - ro, 0, lc, 9)  # down
factory.addCircleArc(6, 1, 7, 5)
factory.addCircleArc(7, 1, 8, 6)
factory.addCircleArc(8, 1, 9, 7)
factory.addCircleArc(9, 1, 6, 8)
factory.addCurveLoop([5, 6, 7, 8], 2)

# square cell
factory.addPoint(0, 0, 0, lc, 10)
factory.addPoint(p, 0, 0, lc, 11)
factory.addPoint(p, p, 0, lc, 12)
factory.addPoint(0, p, 0, lc, 13)
factory.addLine(10, 11,  9)
factory.addLine(11, 12, 10)
factory.addLine(12, 13, 11)
factory.addLine(13, 10, 12)
factory.addCurveLoop([9, 10, 11, 12], 3)

factory.addPlaneSurface([1], 1)
factory.addPlaneSurface([2, 1], 2) # le agrego el hole `1` (no estoy teniendo en cuenta nada de sentidos de giro, no se si hay que hacerlo)
factory.addPlaneSurface([3, 2, 1], 3)

# materials
factory.addPhysicalGroup(2, [1], 1)
factory.addPhysicalGroup(2, [2], 2)
factory.addPhysicalGroup(2, [3], 3)
GridapGmsh.gmsh.model.setPhysicalName(2, 1, "pin")
GridapGmsh.gmsh.model.setPhysicalName(2, 2, "cladding")
GridapGmsh.gmsh.model.setPhysicalName(2, 3, "water")

# boundaries
factory.addPhysicalGroup(1,  [9], 4)
factory.addPhysicalGroup(1, [10], 5)
factory.addPhysicalGroup(1, [11], 6)
factory.addPhysicalGroup(1, [12], 7)
GridapGmsh.gmsh.model.setPhysicalName(1, 4, "bottom")
GridapGmsh.gmsh.model.setPhysicalName(1, 5, "right")
GridapGmsh.gmsh.model.setPhysicalName(1, 6, "top")
GridapGmsh.gmsh.model.setPhysicalName(1, 7, "left")

factory.synchronize()

gmsh.model.mesh.generate(2)

gmsh.write("pincell.msh")

if !("-nopopup" in ARGS)
    gmsh.fltk.run()
end

gmsh.finalize

# move to json file format
using GridapGmsh: GmshDiscreteModel
mshfile = joinpath(@__DIR__,"pincell.msh")
model = GmshDiscreteModel(mshfile; renumber=true)
Gridap.Io.to_json_file(model, "pincell.json")


text:
println("\nRunning runtests.jl ...")

using FactCheck, DictFiles

shouldtest(f, a) = length(ARGS) == 0 || in(a, ARGS) ? facts(f, a) : nothing
shouldtestcontext(f, a) = length(ARGS) < 2 || a == ARGS[2] ? facts(f, a) : nothing

macro throws_pred(ex) FactCheck.throws_pred(ex) end

filename = tempname()
data = [1 2 3 4 5]

shouldtest("Helpers") do
    @fact DictFiles.sortkeys(Any[]) --> Any[]
    @fact DictFiles.sortkeys([3,2,1]) --> [1,2,3]
    @fact DictFiles.sortkeys([3,2,"1"]) --> ["1",2,3]
    @fact DictFiles.sortkeys([3,2,:a,"b"]) --> [2,3,:a,"b"]
end

shouldtest("basic") do
    dictopen(filename) do a
        @fact stat(filename).inode --> not(0)

        a["a"] = (1,2,"a")
        @fact a["a"] --> (1,2,"a")
        a["a"] = "aa"
        @fact a["a"] --> "aa"
        @fact a[] --> Dict("a"=>"aa")
        a["a",1] = 11
        @fact a["a",1] --> 11
        @fact a["a"] --> Dict(1 => 11)
        @fact a[] --> Dict("a" => Dict(1 => 11))
        a["a",2] = 22
        @fact a["a",2] --> 22
        @fact a["a"] --> Dict(1 => 11, 2 => 22)
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22))
        a["b"] = data
        @fact a["b"] --> data
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22), "b" => data)

        delete!(a, "a")
        @fact a[] --> Dict("b" => data)

        delete!(a, "b")
        @fact a[] --> Dict()

        a[] = Dict("a" => Dict(1 => 11, 2 => 22), "b" => data)
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22), "b" => data)
        a["a"] = Dict(1 => 11, 2 => 22)
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22), "b" => data)
        a[:c] = "c"
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22), "b" => data, :c => "c")
    end

    dictopen(filename) do a
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22), "b" => data, :c => "c")
        @fact haskey(a, "a") --> true
        @fact haskey(a, "a", 1) --> true
        @fact haskey(a, "a", :nope) --> false
        @fact haskey(a, "z") --> false

        @fact in("a",keys(a)) --> true
        @fact in("b",keys(a)) --> true
        @fact in(:c,keys(a)) --> true
        @fact in(1,keys(a, "a")) --> true
        @fact in(2,keys(a, "a")) --> true

        a[(1,)] = 1
        @fact a[(1,)] --> 1
        @fact in((1,),keys(a)) --> true

        @fact in(Dict(1 => 11, 2 => 22),values(a)) --> true
        @fact in(11,values(a,"a")) --> true
        @fact in(22,values(a,"a")) --> true
        @fact keys(a,"b") --> Any[]
        @fact values(a,"b") --> Any[]

        @fact get(a, 1, "a") --> Dict(1 => 11, 2 => 22)
        @fact get(a, 1, "z") --> 1

        @fact getkey(a, 1, "a") --> ("a",)
        @fact getkey(a, 1, "z") --> 1
    end

    context("overwrite fields") do
        dictopen(filename,"w") do a
            a[1,1]   = "hi1"
            a[1]     = "hi2"
            a[1,2,1] = "hi3"
            a[1,2]   = "hi4"
            a[1,2,1] = "hi5"
            a[1]     = "hi6"
            a[1,2,1] = "hi7"
            a[1,3]   = "hi8"
        end
    end
end

shouldtest("dictread") do
    filename = tempname()
    a = Dict(:1 => 1, :dict => Dict(:2 => 2, :3 => 3))
    dictwrite(a, filename)
    r = dictread(filename)
    @fact r --> a
    r = dictread(filename, :1)
    @fact r --> 1
    r = dictread(filename, :dict, :2)
    @fact r --> 2
end

shouldtest("Compacting") do
    rm(filename)
    dictopen(filename) do a
        a["a"] = rand(1000,1000)
    end
    oldsize = filesize(filename)
    dictopen(filename) do a
        a[] = Dict("a" => Dict(1 => 11, 2 => 22), "b" => data, :c => "c")
    end
    compact(filename)
    dictopen(filename) do a
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22), "b" => data, :c => "c")
    end
    @fact filesize(filename) < oldsize --> true
end

shouldtest("Error handling") do
    dictopen(filename) do a
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22), "b" => data, :c => "c")
        try
            a["asdf"]
            @fact "no exception" --> "exception"
        catch
            @fact "exception" --> "exception"
        end
        try
            a["a",123]
            @fact "no exception" --> "exception"
        catch
            @fact "exception" --> "exception"
        end
    end
end

shouldtest("Tuple handling") do
    dictopen(filename) do a
        a["a","b",("ID", "param", 0, 0x5bca7c69b794f8ce)] = 123
        @fact a["a","b",("ID", "param", 0, 0x5bca7c69b794f8ce)] --> 123
        @fact a["a"]["b"] --> Dict(("ID", "param", 0, 0x5bca7c69b794f8ce) => 123)
    end
end

shouldtest("mmap") do
    dictopen(filename) do a
        data = rand(2,3)
        setindex!(a, data, "m"; mmap = true)
        m = DictFiles.mmap(a, "m")
        @fact copy(m) --> data
    end
end

shouldtest("Subviews through DictFile(a, keys)") do
    rm(filename)
    dictopen(filename) do a
        a[] = Dict("a" => Dict(1 => 11, 2 => 22), "b" => data, :c => "c")
        b = DictFile(a, "a")
        @fact keys(b) --> [1,2]
        @fact values(b) --> [11,22]
        b[3] = 33
        @fact b[3] --> 33
        @fact a[] --> Dict("a" => Dict(1 => 11, 2 => 22, 3 => 33), "b" => data, :c => "c")
    end
end

shouldtest("makekey(a, k)") do
    dictopen(filename) do a
        @fact DictFiles.makekey(a, (1,)) --> "/1"
        @fact DictFiles.makekey(a, ('a',)) --> "/'a'"
        @fact DictFiles.makekey(a, ("a",)) --> "/\"a\""
        @fact DictFiles.makekey(a, (1,2)) --> "/1/2"
        @fact DictFiles.makekey(a, (1,2,3,4,5)) --> "/1/2/3/4/5"
        @fact DictFiles.makekey(a, (1,'a',3)) --> "/1/'a'/3"
        @fact DictFiles.makekey(a, (1,'a',"a")) --> "/1/'a'/\"a\""
        @fact DictFiles.makekey(a, ("abc",)) --> "/\"abc\""
        @fact DictFiles.makekey(a, ("abc",(1,2))) --> "/\"abc\"/(1,2)"
        @fact DictFiles.makekey(a, (1.,(1,2))) --> "/1.0/(1,2)"
        @fact DictFiles.makekey(a, (1.f0,(1,2))) --> "/1.0f0/(1,2)"
        @fact DictFiles.makekey(a, (1.f0,(1,2),"asdf",'b')) --> "/1.0f0/(1,2)/\"asdf\"/'b'"
        @fact DictFiles.makekey(a, ([1,2,3],)) --> "/[1,2,3]"
        @fact DictFiles.makekey(a, (1000,)) --> "/1000"
    end
end

try
    rm(filename)
end

shouldtest("stress test") do
    nopen = 10
    nwrites = 200
    types = [Int32,UInt32,Int64,UInt64,Float32,Float64]
    payload() = rand(types[rand(1:length(types))], tuple(1,tuple(map( x -> rand(1:5), 5)...)...)...)

    filename = tempname()

    for iopen = 1:nopen
        data = Array{Any}(nwrites)
        readdata = Array{Any}(nwrites)
        dictopen(filename) do a
            for i = 1:nwrites
                key = map( x -> rand(1:10), 1:rand(1:3))
                data[i] = payload()
                a[key...] = data[i]
                readdata[i] = a[key...]
            end
        end
        for i = 1:length(data) # FIXME remove this
            if readdata[i] != data[i]
                @show i readdata[i] data[i]
            end
        end
        @fact readdata  --> data
    end

    DictFiles.compact(filename)

    rm(filename)
end

shouldtest("parallel") do
    addprocs(3)
    @everywhere using DictFiles
    filename = tempname()
    a = @fetchfrom 2 DictFile(filename)
    a[1] = 10
    @fact (@fetchfrom 2 a[1]) --> 10
    @fact (@fetchfrom 3 a[1]) --> 10
end

shouldtest("blosc") do
    filename = tempname()
    dictopen(filename) do a
        data = rand(2,3)
        a["a"] = blosc(data)
        @fact a["a"] --> data
        data = Dict(1 => rand(2,3), 'c' => "asdf")
        a["a"] = blosc(data)
        @fact a["a"] --> data
    end
end

shouldtest("serialize") do
    filename = tempname()
    dictopen(filename) do a
        data = rand(2,3)
        b = serialized(data)
        a["a"] = b
        @fact a["a"] --> data
    end
end

shouldtest("snapshot") do
    data = rand(10)
    snapshot(data)
    @fact loadsnapshot() --> data
    snapshot("data", data)
    @fact loadsnapshot("data") --> data
    @fact loadsnapshot() --> Dict("data" => data)
    snapshot("data", data, "a", 1)
    @fact loadsnapshot() --> Dict("data" => data, "a" => 1)
    @fact loadsnapshot("a") --> 1
    @fact loadsnapshot("data","a") --> Dict("data" => data, "a" => 1)

end

println("runtests.jl done!")
FactCheck.exitstatus()


text:
function bs(O::Option)
    d1 = (log(O.s / O.k) + (O.r + O.σ * O.σ / 2) * O.t) / (O.σ * √O.t)
    d2 = d1 - O.σ * √O.t

    if O.call_or_put == "Put"
        return CND(-d2) * O.k * exp(-O.r * O.t) - CND(-d1) * O.s
    else
        return O.s * CND(d1) - exp(-O.r * O.t) * O.k * CND(d2)
    end
end

function CND(x::Float64)
    a1 = 0.31938153
    a2 = -0.356563782
    a3 = 1.781477937
    a4 = -1.821255978
    a5 = 1.330274429
    l = abs(x)
    k = 1 / (1 + 0.2316419 * l)

    CND = 1 - 1 / sqrt(2 * π) * exp(-l^2 / 2) * (a1 * k + a2 * k^2 + a3 * k^3 + a4 * k^4 + a5 * k^5)

    if x < 0
        return 1 - CND
    end

    return CND
end


text:
### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 91615676-4eb1-11eb-2599-9f94046c2a5d
begin
	let
		env = mktempdir()
		import Pkg
		Pkg.activate(env)
		Pkg.Registry.update()
		Pkg.add(Pkg.PackageSpec(;name="PlutoUI", version="0.6.7-0.6"))
	end
	using PlutoUI
end

# ╔═╡ f381d0ce-3faf-11eb-2586-ffe002ea8d9d
begin
	using Plots
	using LinearAlgebra
	using SparseArrays
end

# ╔═╡ 6ce54252-3fb0-11eb-34e6-978f07a94d99
md"""

# Simulation of a model PDE

Let's try to build some intuition for how a PDE behaves. Suppose we have a PDE on ``[-1,1]`` with ``u(x,t)`` as the solution

```math
\begin{gather*}
\frac{\partial u}{\partial t} + a\frac{\partial u}{\partial x} - b\frac{\partial^2 u}{\partial x^2} = f(x,t)
\end{gather*}
```
with left and right boundary conditions
```math
u(-1,t) = u_L(t), \qquad u'(1,t) = 0
```
and initial condition ``u(x,0) = u_0(x)``.

Different choices of ``a,b`` transition between different types of PDEs.

## Physical parameters

Try choosing different values for each parameter.

``a =`` $(@bind a NumberField(0:.1:10, default = 1))
``b =`` $(@bind b NumberField(0:.1:10, default = 1))

Different parameter values can yield different PDEs. For example:
- ``a=1, b = 0`` results in the advection equation
- ``a=0, b = 1`` results in the heat equation
"""

# ╔═╡ e6a60308-3fb8-11eb-090f-9d533d92759b
md"Let's now define a forcing function `f(x,t)`, left boundary data `uL(t)`, and an initial condition `u0(x)`. Feel free to change these (set them to zero to remove)."

# ╔═╡ eb538830-3fb8-11eb-1d93-230874209dc7
begin
	u0(x)  = 0*(.5*sin(pi*x) - .25*sin(2*pi*x) + .25*sin(4*pi*x))
	f(x,t) = 0*(exp(-100*(x+.75)^2))*(.5*t + sin(2*t))*exp(-.1*t)
	uL(t)  = .75*sin(1*pi*t)
end;

# ╔═╡ 6bab77f0-6006-11eb-08b4-6b0a33f38472
md"Let's plot the initial condition below"

# ╔═╡ 1afcd228-3fc1-11eb-127a-b5059e1a0029
md"""
# Discretization parameters
"""

# ╔═╡ 952ca744-3fb7-11eb-20ab-a90257510dfb
md"""
Number of nodes `N` and timestep `dt`

`N =` $(@bind N NumberField(3:2:200, default = 201))

`dt =` $(@bind dt NumberField(.0001:.0001:10, default = .01))

Final time = $(@bind FinalTime NumberField(0:.0001:10, default = .25))
"""

# ╔═╡ d323417c-3fb5-11eb-0cdb-655211612ff0
begin
	# Here is the code to construct discretization matrices for our simulation
	x = LinRange(-1,1,N)
	h = 2/(N-1) # mesh size
	K = spdiagm(-1=>-ones(N-1),1=>-ones(N-1),0=>2*ones(N))
	K[1,1] = -1
	K[N,N] = 1 # corrections at the left/right endpoints
	K = (@. K/h^2)

	Q = spdiagm(-1=>-ones(N-1),1=>ones(N-1))
	Q[end,end] = 1
	Q = (@. Q/h)

	# modify matrices for Dirichlet BCs
	K[1,:] .= 0
	K[1,1] = 1
	Q[1,:] .= 0
	Q[1,1] = 1
end;

# ╔═╡ 73588218-6006-11eb-095e-e9c90ff9bbb8
begin
	plot(x,u0.(x),xlim=(-1,1),ylim=(-1,1),lw=3)
	plot!(title="Initial condition",ratio=1,legend=false)
	scatter!(x,u0.(x))
end

# ╔═╡ 3b1b2ad0-3fbe-11eb-0745-2f8aeccfea54
md"Now we'll solve our PDE until final time $(FinalTime) and animate the results."

# ╔═╡ 4b1df9c6-3fbe-11eb-1dcf-3f883d7d92de
begin
	B = (a*Q + b*K)
	A = I + .5*dt*B # Crank-Nicholson in time
	Nsteps = ceil(Int,FinalTime/dt)
	u = u0.(x)

	anim = @animate for i = 1:Nsteps
		global u
		t = i*dt

		bb = u + dt*f.(x,t) - .5*dt*B*u
		bb[1] = uL(t) # impose BC at left endpoint
		u = A\bb

		plot(x,u,xlim=(-1,1),ylim=(-1,1),lw=3)
		scatter!(x,u,label="")
		plot!(title="Time=$t",ratio=1,legend=false)
	end	every 1
	gif(anim, "fem.gif", fps = 10)
end

# ╔═╡ Cell order:
# ╟─6ce54252-3fb0-11eb-34e6-978f07a94d99
# ╟─e6a60308-3fb8-11eb-090f-9d533d92759b
# ╠═eb538830-3fb8-11eb-1d93-230874209dc7
# ╟─6bab77f0-6006-11eb-08b4-6b0a33f38472
# ╟─73588218-6006-11eb-095e-e9c90ff9bbb8
# ╟─1afcd228-3fc1-11eb-127a-b5059e1a0029
# ╟─952ca744-3fb7-11eb-20ab-a90257510dfb
# ╟─d323417c-3fb5-11eb-0cdb-655211612ff0
# ╟─3b1b2ad0-3fbe-11eb-0745-2f8aeccfea54
# ╟─4b1df9c6-3fbe-11eb-1dcf-3f883d7d92de
# ╟─91615676-4eb1-11eb-2599-9f94046c2a5d
# ╟─f381d0ce-3faf-11eb-2586-ffe002ea8d9d


text:
using BAPOMCP
using Base.Test
using POMDPs
using POMDPModels
using POMDPToolbox
using NBInclude
using D3Trees

test_solver(POMCPSolver(), BabyPOMDP())

pomdp = BabyPOMDP()
solver = POMCPSolver(rng = MersenneTwister(1))
planner = solve(solver, pomdp)

tree = BasicPOMCP.POMCPTree(pomdp, solver.tree_queries)
node = BasicPOMCP.POMCPObsNode(tree, 1)

r = @inferred BasicPOMCP.simulate(planner, initial_state(pomdp, MersenneTwister(1)), node, 20)

sim = HistoryRecorder(max_steps=10)
simulate(sim, pomdp, planner, updater(pomdp))

solver = POMCPSolver(max_time=0.1, tree_queries=typemax(Int), rng = MersenneTwister(1))
planner = solve(solver, pomdp)
action_info(planner, initial_state_distribution(pomdp))
println("time below should be about 0.1 seconds")
@time a, info = action_info(planner, initial_state_distribution(pomdp))
@show info[:search_time_us]

solver = POMCPSolver(max_time=0.1, tree_queries=typemax(Int), rng = MersenneTwister(1))
planner = solve(solver, pomdp)
a, info = action_info(planner, initial_state_distribution(pomdp), tree_in_info=true)

d3t = D3Tree(planner, title="test")
d3t = D3Tree(info[:tree], title="test tree")
# inchrome(d3t)
show(STDOUT, MIME("text/plain"), d3t)

solver = POMCPSolver(max_time=0.1, tree_queries=typemax(Int), rng=MersenneTwister(1), tree_in_info=true)
planner = solve(solver, pomdp)
a, info = action_info(planner, initial_state_distribution(pomdp))

d3t = D3Tree(info[:tree], title="test tree (tree_in_info solver option)")

nbinclude(joinpath(dirname(@__FILE__), "..", "notebooks", "Minimal_Example.ipynb"))

d3t = D3Tree(planner, title="test")
# inchrome(d3t)

# test consistency when rng is specified
pomdp = BabyPOMDP()
solver = POMCPSolver(rng = MersenneTwister(1))
planner = solve(solver, pomdp)
hist1 = simulate(HistoryRecorder(max_steps=1000, rng=MersenneTwister(3)), pomdp, planner)

solver = POMCPSolver(rng = MersenneTwister(1))
planner = solve(solver, pomdp)
hist2 = simulate(HistoryRecorder(max_steps=1000, rng=MersenneTwister(3)), pomdp, planner)

@test discounted_reward(hist1) == discounted_reward(hist2)

# REQUIREMENTS
solver = POMCPSolver()
pomdp = TigerPOMDP()

println("============== @requirements_info with only solver:")
@requirements_info solver
println("============== @requirements_info with solver and pomdp:")
@requirements_info solver pomdp


text:
#####################################
# StaticArray Interface
#####################################
Base.propertynames(::SLArray{S,T,N,L,Syms}) where {S,T,N,L,Syms<:Tuple{Syms1,Syms2}} where {Syms1,Syms2} = Syms1,Syms2
symnames(::Type{SLArray{S,T,N,L,Syms}}) where {S,T,N,L,Syms<:Tuple{Syms1,Syms2}} where {Syms1,Syms2} = Syms1,Syms2

"""
    @SLSliced Size Names1, Names2
    @SLSlice Eltype Size Names1, Names2

Creates an anonymous function that builds a labelled static vector with eltype
`ElType` with names determined from the `Names` tuple with size `Size`. If no eltype
is given, then the eltype is determined from the arguments in the constructor.

For example:

```julia
ABC = @SLSliced (4,2) (:a,:b,:c,:d), (:x,:y)
x = ABC([1.0 2.5; 3.0 5.0; 9.0 11.4; 12.9 17.7])
x.a.x == 1.0
x.a.y == 2.5
x.c.x == x[3,1]
x.d.y == x[4,2]
```

"""
macro SLSliced(dims,syms)
  dims = esc(dims)
  syms = esc(syms)
  quote
    SLArray{Tuple{$dims...,},Tuple{$syms...,}}
  end
end

macro SLSliced(T,dims,syms)
  T = esc(T)
  dims = esc(dims)
  syms = esc(syms)
  quote
      SLArray{Tuple{$dims...},$T,length($dims),prod($dims),Tuple{$syms...}}
  end
end

# """
#     dimSymbols(::SLArray{S,T,N,L,Syms}, dim::Int) where {Syms<:Tuple}

# Returns the labels of the `SLArray` associated with dimension `dim`.

# For example:

#     A = @LSliced [1 2; 3 4; 5 6] (:a,:b,:c), (:x, :y)
#     dimSymbols(A,2)  # (:x, :y)
# """
# dimSymbols(::SLArray{S,T,N,L,Syms}, dim::Int) where
# {S,T,N,L,Syms<:Tuple} = Syms.parameters[dim]

# "returns dimSymbols(,1)"
# rowSymbols(::SLArray{S,T,N,L,Syms}) where
# {S,T,N,L,Syms<:Tuple} = Syms.parameters[1]

# "returns dimSymbols(,2)"
# colSymbols(::SLArray{S,T,N,L,Syms}) where
# {S,T,N,L,Syms<:Tuple} = Syms.parameters[2]

symbols(::SLArray{S,T,N,L,Syms}) where
{rows,cols,S,T,N,L,Syms<:Tuple{rows,cols}} = rows,cols


text:
### utility functions ###

function logsumexp(w)
    offset = maximum(w)
    exp_w = exp.(w .- offset)
    s = sum(exp_w)
    log(s) + offset
end


text:
#############################################################
#
# Custom Types
#
#############################################################

#Type to hold Twitter API credentials
type TWCRED
    consumer_key::AbstractString
    consumer_secret::AbstractString
    oauth_token::AbstractString
    oauth_secret::AbstractString
end

#Places response object
type PLACES
    attributes::@compat(Union{Dict, Void})
    bounding_box::@compat(Union{Dict, Void})
    country::@compat(Union{AbstractString, Void})
    country_code::@compat(Union{AbstractString, Void})
    full_name::@compat(Union{AbstractString, Void})
    id::@compat(Union{AbstractString, Void})
    name::@compat(Union{AbstractString, Void})
    place_type::@compat(Union{AbstractString, Void})
    url::@compat(Union{AbstractString, Void})
end

#Users response object
type USERS
    contributors_enabled::@compat(Union{Bool, Void})
    created_at::@compat(Union{AbstractString, Void})
    default_profile::@compat(Union{Bool, Void})
    default_profile_image::@compat(Union{Bool, Void})
    description::@compat(Union{AbstractString, Void})
    entities::@compat(Union{Dict, Void})
    favourites_count::@compat(Union{Int, Void})
    follow_request_sent::@compat(Union{Bool, Void})
    followers_count::@compat(Union{Int, Void})
    friends_count::@compat(Union{Int, Void})
    geo_enabled::@compat(Union{Bool, Void})
    id::@compat(Union{Int, Void})
    id_str::@compat(Union{AbstractString, Void})
    is_translator::@compat(Union{Bool, Void})
    listed_count::@compat(Union{Int, Void})
    location::@compat(Union{AbstractString, Void})
    name::@compat(Union{AbstractString, Void})
    profile_background_color::@compat(Union{AbstractString, Void})
    profile_background_image_url::@compat(Union{AbstractString, Void})
    profile_background_image_url_https::@compat(Union{AbstractString, Void})
    profile_backround_tile::@compat(Union{Bool, Void})
    profile_banner_url::@compat(Union{AbstractString, Void})
    profile_image_url::@compat(Union{AbstractString, Void})
    profile_image_url_https::@compat(Union{AbstractString, Void})
    profile_link_color::@compat(Union{AbstractString, Void})
    profile_sidebar_border_color::@compat(Union{AbstractString, Void})
    profile_sidebar_fill_color::@compat(Union{AbstractString, Void})
    profile_text_color::@compat(Union{AbstractString, Void})
    profile_use_background_image::@compat(Union{Bool, Void})
    protected::@compat(Union{Bool, Void})
    screen_name::@compat(Union{AbstractString, Void})
    show_all_inline_media::@compat(Union{Bool, Void})
    status::@compat(Union{Dict, Void})
    statuses_count::@compat(Union{Int, Void})
    time_zone::@compat(Union{AbstractString, Void})
    url::@compat(Union{AbstractString, Void})
    utc_offset::@compat(Union{Int, Void})
    verified::@compat(Union{Bool, Void})
    withheld_in_countries::@compat(Union{AbstractString, Void})
    withheld_scope::@compat(Union{AbstractString, Void})
end

#Tweets response object
type TWEETS
    contributors::@compat(Union{Dict, Void})
    coordinates::@compat(Union{Dict, Void})
    created_at::@compat(Union{AbstractString, Void})
    current_user_retweet::@compat(Union{Dict, Void})
    entities::@compat(Union{Dict, Void})
    favorite_count::@compat(Union{Int, Void})
    favorited::@compat(Union{Bool, Void})
    filter_level::@compat(Union{AbstractString, Void})
    full_text::@compat(Union{AbstractString, Void})
    id::@compat(Union{Int, Void})
    id_str::@compat(Union{AbstractString, Void})
    in_reply_to_screen_name::@compat(Union{AbstractString, Void})
    in_reply_to_status_id::@compat(Union{Int, Void})
    in_reply_to_status_id_str::@compat(Union{AbstractString, Void})
    in_reply_to_user_id::@compat(Union{Int, Void})
    in_reply_to_user_id_str::@compat(Union{AbstractString, Void})
    lang::@compat(Union{AbstractString, Void})
    place::@compat(Union{Dict, Void})
    possibly_sensitive::@compat(Union{Bool, Void})
    scopes::@compat(Union{Dict, Void})
    retweet_count::@compat(Union{Int, Void})
    retweeted::@compat(Union{Bool, Void})
    retweeted_status::@compat(Union{Dict, Void})
    source::@compat(Union{AbstractString, Void})
    text::@compat(Union{AbstractString, Void})
    truncated::@compat(Union{Bool, Void})
    user::@compat(Union{Dict, Void})
    withheld_copyright::@compat(Union{Bool, Void})
    withheld_in_countries::@compat(Union{Array, Void})
    withheld_scope::@compat(Union{AbstractString, Void})
end

#Lists response object
type LISTS
    created_at::@compat(Union{AbstractString, Void})
    description::@compat(Union{AbstractString, Void})
    following::@compat(Union{Bool, Void})
    full_name::@compat(Union{AbstractString, Void})
    id::@compat(Union{Int, Void})
    id_str::@compat(Union{AbstractString, Void})
    member_count::@compat(Union{Int, Void})
    mode::@compat(Union{AbstractString, Void})
    name::@compat(Union{AbstractString, Void})
    slug::@compat(Union{AbstractString, Void})
    subscriber_count::@compat(Union{Int, Void})
    uri::@compat(Union{AbstractString, Void})
    user::@compat(Union{Dict, Void})
end