text:
# Functions related to an individual's genome:
# - meiosis
# - mutation
# - genome creation
# - auxiliary functions

"""
    meiosis(genome, maternal)

Carry out meiosis on a genome (marked as maternal or not). Returns a haploid
gamete genome. (genome => array of chromosomes)
"""
function meiosis(genome::Array{Chromosome,1}, maternal::Bool)
    gametelength = Int(length(genome)/2)
    gamete = Chromosome[]
    m, p = 1, 1
    while length(gamete) < gametelength
        # find the next maternal/paternal chromosome
        while !genome[m].maternal
            m += 1
        end
        while genome[p].maternal
            p += 1
        end
        # then choose one at random and use it to create a new chromosome for the gamete
        push!(gamete, Chromosome(genome[rand([m,p])].genes, maternal))
        m += 1; p += 1
    end
    gamete
end

"""
    getmeantraitvalue(traits, traitidx)

Take an array of traits and return the mean value of the indexed trait.
"""
function getmeantraitvalue(traits::Array{Trait, 1}, traitidx::Integer)
    mean(skipmissing(map(x -> x.value, filter(x -> x.nameindex == traitidx, traits))))
end

"""
    getstdtraitvalue(traits, traitidx)

Take an array of traits and return the standard deviation of the indexed trait.
"""
function getstdtraitvalue(traits::Array{Trait, 1}, traitidx::Integer)
    if length(traits) <= 1
        0.0
    else
        std(skipmissing(map(x -> x.value, filter(x -> x.nameindex == traitidx, traits))))
    end
end

"""
    gettraitdict(chromosomes, traitnames)

Convert a genome (an array of chromosomes) into a dict of traits and their values.
"""
function gettraitdict(chrms::Array{Chromosome, 1}, traitnames::Array{String, 1})
    #XXX can this be made more efficient? It's called really often...
    genes = AbstractGene[]
    nchrms = 0
    ngenes = 0
    for chrm in chrms
        append!(genes, chrm.genes)
        nchrms += 1
        ngenes += length(chrm.genes)
    end
    traits = vcat(map(g -> g.codes, genes)...)
    traitdict = Dict{String, Float64}()
    for traitidx in eachindex(traitnames)
        traitdict[traitnames[traitidx]] = getmeantraitvalue(traits, traitidx)
        traitdict[string(traitnames[traitidx]) * "sd"] = getstdtraitvalue(traits, traitidx)
    end
    traitdict["ngenes"] = ngenes
    traitdict["nlnkgunits"] = nchrms
    traitdict
end

"""
    gettraitdict(traits, traitnames)

Construct a trait dict from a list of Trait objects.
"""
function gettraitdict(traits::Array{Trait, 1}, traitnames::Array{String, 1})
    traitdict = Dict{String, Float64}()
    for traitidx in unique(map(x -> x.nameindex, traits))
        traitdict[traitnames[traitidx]] = getmeantraitvalue(traits, traitidx)
    end
    traitdict
end

"""
    getseqsimilarity(seqone, seqtwo)

Compare two strings and return similarity.
"""
function getseqsimilarity(indgene::AbstractString, mategene::AbstractString)
    basediffs = 0
    for i in eachindex(indgene) # this is actually faster than `sum(collect(indgene) .== collect(mategene))`
        try
            indgene[i] != mategene[i] && (basediffs += 1) # alternatively use bioinformatics tools
        catch # e.g., in case of differently long genes
            basediffs += 1
        end
    end
    seqidentity = 1 - (basediffs / length(indgene))
end

"""
    getseq(genome, traitidx)

Find and return the sequence of one gene that codes for the given trait `traitidx`.
"""
function getseq(genome::Array{Chromosome, 1}, traitidx::Integer, compressgenes::Bool)
    seq = ""
    for chrm in genome
        for gene in chrm.genes
            # use one compatibility gene randomly #XXX what's that comment supposed to mean?
            #TODO have another look at this
            if any(x -> x.nameindex == traitidx, gene.codes)
                compressgenes ? seq = num2seq(gene.sequence) : seq = gene.sequence
            end
        end
    end
    seq
end

"""
    iscompatible(mate, individual, settings)

Check to see whether two individual organisms are reproductively compatible.
"""
function iscompatible(mate::Individual, ind::Individual, settings::Dict{String, Any})
    mate.lineage != ind.lineage && return false ##TODO must be possible under some conditions
    compatidx = findfirst(x -> x == "compat", settings["traitnames"])
    indgene = getseq(ind.genome, compatidx, settings["compressgenes"])
    mategene = getseq(mate.genome, compatidx, settings["compressgenes"])
    seqidentity = getseqsimilarity(indgene, mategene)
    seqidentity >= ind.traits["seqsimilarity"]
end

"""
    seq2num(sequence)

Convert a DNA base sequence (a string) into binary and then into an integer.
This saves memory.
"""
#XXX Actually, this hardly seems to make a difference :-( So using `compressgenes = false`
# seems to be a pretty good choice right now. What we should really do is create gene
# sequences as int right away and save ourselves the intermediate string allocation.
function seq2num(sequence::String)
    # This function effectively uses the same technique as seq2bignum, but it skips the
    # intermediate allocations and is therefore more efficient.
    num::Int64 = 0  # Int64 allows for max length of 21bp
    for b in eachindex(sequence)
        if sequence[end+1-b] == 'a'
            num += 2^(3*(b-1)) * 4 # b'100'
        elseif sequence[end+1-b] == 'c'
            num += 2^(3*(b-1)) * 5 # b'101'
        elseif sequence[end+1-b] == 'g'
            num += 2^(3*(b-1)) * 6 # b'110'
        elseif sequence[end+1-b] == 't'
            num += 2^(3*(b-1)) * 7 # b'111'
        end
    end
    num
end

function intseq(n::Int)
    #XXX doesn't seem to work either? I'm starting to think I'm optimising in the wrong place :-(
    num::Int64 = 0  # Int64 allows for max length of 21bp
    for b in 0:(n-1)
        # basically like seq2bignum(), but skips all intermediate allocations
        num += 2^(3*b) * rand(4:7)
    end
    num
end

"""
    seq2bignum(sequence)

Convert a DNA base sequence (a string) into binary and then into a BigInt (for
larger genes). This saves memory.
"""
function seq2bignum(sequence::String)
    bases = "acgt"
    binary = ""
    for base in sequence
        binary *= string(findfirst(x -> x == base, bases) + 3, base = 2)
    end
    parse(BigInt, binary, base = 2) # BigInt allows arbitrary-length sequences
end

"""
    num2seq(n)

Convert an integer into binary and then into a DNA base sequence string.
"""
function num2seq(n::Integer)
    bases = "acgt"
    binary = string(n, base = 2)
    sequence = ""
    for i in 1:3:(length(binary) - 2)
        sequence *= string(bases[parse(Int, binary[i:(i + 2)], base = 2) - 3])
    end
    sequence
end

"""
    createtraits(settings)

Create an array of trait objects generated from the default trait values (with a
random offset).
"""
function createtraits(settings::Dict{String, Any})
    #TODO: this is all very ugly. (case/switch w/ v. 2.0+?)
    traitnames = settings["traitnames"]
    traits = Trait[]
    # exponential distributions of body sizes:
    repoffset = settings["maxrepsize"] - settings["minrepsize"]
    seedoffset = settings["maxseedsize"] - settings["minseedsize"]
    tempoffset = settings["maxtemp"] - settings["mintemp"]
    sizes = Vector{Float64}(undef, 2)
    while true
        sizes[1] = exp(settings["minrepsize"] + repoffset * rand())
        sizes[2] = exp(settings["minseedsize"] + seedoffset * rand())
        sizes[1] > sizes[2] && break
    end
    repsize, seedsize = sizes
    for idx in eachindex(traitnames)
        if occursin("dispshape", traitnames[idx])
            push!(traits, Trait(idx, rand() * settings["dispshape"]))
        elseif occursin("dispmean", traitnames[idx])
            push!(traits, Trait(idx, rand() * settings["dispmean"]))
        elseif occursin("precopt", traitnames[idx])
            push!(traits, Trait(idx, rand() * settings["precrange"]))
        elseif occursin("prectol", traitnames[idx])
            push!(traits, Trait(idx, rand() * settings["maxbreadth"]))
        elseif occursin("repsize", traitnames[idx])
            push!(traits, Trait(idx, repsize))
        elseif occursin("seqsimilarity", traitnames[idx]) && settings["fixtol"]
            push!(traits, Trait(idx, settings["tolerance"]))
        elseif occursin("seqsimilarity", traitnames[idx]) && !settings["fixtol"]
            push!(traits, Trait(idx, rand())) # assortative mating might evolve
        elseif occursin("seedsize", traitnames[idx])
            push!(traits, Trait(idx, seedsize))
        elseif occursin("tempopt", traitnames[idx])
            push!(traits, Trait(idx, settings["mintemp"] + rand() * tempoffset))
        elseif occursin("temptol", traitnames[idx])
            push!(traits, Trait(idx, rand() * settings["maxbreadth"]))
        else
            push!(traits, Trait(idx, rand()))
        end
    end
    traits
end

"""
    creategenes(ngenes, traits, settings)

Randomly create a given number of gene objects, with their base sequence and
associated traits. Returns the result as an array of AbstractGenes.
"""
function creategenes(ngenes::Int, traits::Array{Trait,1}, settings::Dict{String, Any})
    genes = AbstractGene[]
    bases = ['a','c','g','t']
    compatidx = findfirst(x -> x == "compat", settings["traitnames"])
    # initialise each gene with an arbitrary sequence
    for i in 1:ngenes
        if settings["compressgenes"] #default
            push!(genes, Gene(intseq(settings["smallgenelength"]), Trait[]))
        else
            sequence = String(rand(bases, settings["smallgenelength"]))
            push!(genes, StringGene(sequence, Trait[]))
        end
    end
    # assign traits to the genes (allows for pleiotropy as well as polygenic inheritance)
    for trait in traits
        (trait.nameindex == compatidx) && continue # the compatibility gene is treated separately
        if settings["degpleiotropy"] == 0
            # Disable polygenic inheritance and pleiotropy: make sure one gene codes for one trait.
            for gene in genes
                if isempty(gene.codes)
                    push!(gene.codes, trait)
                    break
                end
            end
        else
            # We're actually setting polygenic inheritance here, rather than pleiotropy.
            # Pleiotropy is introduced indirectly, because a higher number of coding genes
            # increases the likelihood of picking a gene that already codes for other traits.
            ncodinggenes = rand(Geometric(1 - settings["degpleiotropy"])) + 1
            codinggenes = rand(genes,ncodinggenes)
            for gene in codinggenes
                push!(gene.codes,trait)
            end
        end
    end
    # append a gene that will be used to determine mating compatibility
    # Note: we only need big genes for the compatibility gene, because we need a longer base
    # sequence than offered by `smallgenelength` if we want to do phylogenetic analyses.
    settings["usebiggenes"] ? seql = settings["biggenelength"] : seql = settings["smallgenelength"]
    cseq = String(rand(bases, seql))
    ctrait = [Trait(compatidx, 0.5)]
    if settings["compressgenes"]
        if settings["usebiggenes"]
            push!(genes, BigGene(seq2bignum(cseq), ctrait))
        else
            push!(genes, Gene(seq2num(cseq), ctrait))
        end
    else
        push!(genes, StringGene(cseq, ctrait))
    end
    genes
end

"""
    createchrms(nchrms, genes)

Randomly distribute the passed genes into the given number of chromosomes.
Returns an array of chromosome objects.
"""
function createchrms(nchrms::Int,genes::Array{AbstractGene,1})
    ngenes=size(genes,1)
    if nchrms>1
        chrmsplits = sort(rand(1:ngenes,nchrms-1))
        chromosomes = Chromosome[]
        for chr in 1:nchrms
            if chr==1 # first chromosome
                push!(chromosomes, Chromosome(genes[1:chrmsplits[chr]], true))
            elseif chr==nchrms # last chromosome
                push!(chromosomes, Chromosome(genes[(chrmsplits[chr-1]+1):end], true))
            else
                push!(chromosomes, Chromosome(genes[(chrmsplits[chr-1]+1):chrmsplits[chr]], true))
            end
        end
    else # only one chromosome
        chromosomes = [Chromosome(genes, true)]
    end
    secondset = deepcopy(chromosomes)
    for chrm in secondset
        chrm.maternal = !chrm.maternal
    end
    append!(chromosomes,secondset)
    chromosomes
end

"""
    varyalleles!(genes, settings, locivar)

Mutate gene traits in the passed array of genes.
"""
function varyalleles!(genes::Array{AbstractGene, 1}, settings::Dict{String, Any}, locivar::Float64)
    locivar == 0 && return
    for gene in genes
        mutate!(gene.codes, settings, locivar)
    end
end

"""
    varyalleles!(chromosomes, settings, locivar)

Mutate gene traits in the passed array of chromosomes.
"""
function varyalleles!(chrms::Array{Chromosome, 1}, settings::Dict{String, Any}, locivar::Float64)
    locivar == 0 && return
    for chrm in chrms
        varyalleles!(chrm.genes, settings, locivar)
    end
end

"""
    mutate!(traits, settings, locivar)

Loop over an array of traits, mutating each value in place along a normal distribution.
`locivar` can be used to scale the variance of the normal distribution used to draw new
trait values (together with `settings[phylconstr]`).
"""
function mutate!(traits::Array{Trait, 1}, settings::Dict{String, Any}, locivar::Float64 = 1.0)
    settings["phylconstr"] * locivar == 0 && return
    for trait in traits
        traitname = settings["traitnames"][trait.nameindex]
        occursin("seqsimilarity", traitname) && settings["fixtol"] && continue
        oldvalue = trait.value
        occursin("tempopt", traitname) && (oldvalue -= 273)
        while oldvalue <= 0 # make sure sd of Normal dist != 0
            oldvalue = abs(rand(Normal(0,0.01)))
        end
        newvalue = rand(Normal(oldvalue, oldvalue * settings["phylconstr"] * locivar))
        (newvalue > 1 && occursin("prob", traitname)) && (newvalue=1.0)
        while newvalue <= 0
            newvalue = rand(Normal(oldvalue, oldvalue * settings["phylconstr"] * locivar))
        end
        occursin("tempopt", traitname) && (newvalue += 273)
        trait.value = newvalue
    end
end

"""
    mutate!(individual, temp, settings)

Mutate an individual's genome (sequence and traits) in place.
"""
function mutate!(ind::Individual, temp::Float64, settings::Dict{String, Any})
    muts = settings["mutationrate"] * exp(-act/(boltz*temp))
    nmuts = rand(Poisson(muts))
    nmuts == 0 && return
    chrmidcs = rand(eachindex(ind.genome), nmuts)
    for c in chrmidcs
        ind.genome[c] = deepcopy(ind.genome[c])
        length(ind.genome[c].genes) == 0 && continue
        g = rand(eachindex(ind.genome[c].genes))
        gseq = ind.genome[c].genes[g].sequence
        settings["compressgenes"] ? charseq = collect(num2seq(gseq)) : charseq = collect(gseq)
        i = rand(eachindex(charseq))
        newbase = rand(collect("acgt"),1)[1]
        while newbase == charseq[i]
            newbase = rand(collect("acgt"),1)[1]
        end
        charseq[i] = newbase
        mutate!(ind.genome[c].genes[g].codes, settings)
        ind.genome[c].genes[g].sequence = deepcopy(ind.genome[c].genes[g].sequence)
        if settings["compressgenes"]
            if length(charseq) > 21
                ind.genome[c].genes[g].sequence = seq2bignum(String(charseq))
            else
                ind.genome[c].genes[g].sequence = seq2num(String(charseq))
            end
        else
            ind.genome[c].genes[g].sequence = String(charseq)
        end
    end
    ind.traits = gettraitdict(ind.genome, settings["traitnames"])
end

"""
    mutate!(patch, setting)

Mutate all seed individuals in a patch.
"""
function mutate!(patch::Patch, settings::Dict{String, Any})
    for ind in patch.seedbank
        mutate!(ind, patch.temp, settings)
    end
end

"""
    mutate!(world, settings)

Mutate the world. (That sounds scary!)
"""
function mutate!(world::Array{Patch, 1}, settings::Dict{String, Any})
    for patch in world
        (patch.isisland || !settings["static"]) && mutate!(patch, settings)
    end
end

max_stars_repo_name: lleiding/gemm

max_issues_repo_name: lleiding/gemm

max_forks_repo_name: lleiding/gemm


text:
#!/usr/bin/env julia
using MuxPkgServer

# Host/Port
try
    server = get(ENV, "JULIA_MUX_PKG_SERVER", "http://0.0.0.0:8080")
    m = match(r"(https?://)?(.+):(\d+)", server)
    global host = String(m.captures[2])
    global port = parse(Int, m.captures[3])
catch
    @warn("Invalid JULIA_MUX_PKG_SERVER setting, ignoring and using default of 0.0.0.0:8080!")
    global host = "0.0.0.0"
    global port = 8080
end

# Upstream PkgServers
upstreams = get(ENV, "JULIA_MUX_PKG_SERVER_UPSTREAMS", "https://pkg.julialang.org")
upstreams = split(upstreams, ','; keepempty=false)

# Logging
log_dir = get(ENV, "JULIA_MUX_PKG_SERVER_LOGS") do
    log_dir = mktempdir(; cleanup=false)
    @warn "JULIA_MUX_PKG_SERVER_LOGS is not set, using temp directory for logs" log_dir
    return log_dir
end
mkpath(log_dir)

using Dates, Logging, LoggingExtras
const date_format = dateformat"yyyy-mm-dd\THH:MM:SS\Z"
timestamp_logger(logger) = TransformerLogger(logger) do log
    merge(log, (; message = "[$(Dates.format(unix2datetime(time()), date_format))] $(log.message)"))
end
global_logger(TeeLogger(
    timestamp_logger(
        MinLevelLogger(
            DatetimeRotatingFileLogger(
                log_dir,
                string(raw"yyyy-mm-dd-\m\u\x\p\k\g\s\e\r\v\e\r.\l\o\g"),
            ),
            Logging.Info,
        ),
    ),
    current_logger(),
))

MuxPkgServer.start(;
    host=host, port=port,
    upstreams = upstreams,
)

ext: jl


text:
"""
    AMRCallback(semi, controller [,adaptor=AdaptorAMR(semi)];
                interval,
                adapt_initial_condition=true,
                adapt_initial_condition_only_refine=true)

Performs adaptive mesh refinement (AMR) every `interval` time steps
for a given semidiscretization `semi` using the chosen `controller`.
"""
struct AMRCallback{Controller, Adaptor, Cache}
  controller::Controller
  interval::Int
  adapt_initial_condition::Bool
  adapt_initial_condition_only_refine::Bool
  dynamic_load_balancing::Bool
  adaptor::Adaptor
  amr_cache::Cache
end

function AMRCallback(semi, controller, adaptor;
                     interval,
                     adapt_initial_condition=true,
                     adapt_initial_condition_only_refine=true,
                     dynamic_load_balancing=true)
  # check arguments
  if !(interval isa Integer && interval >= 0)
    throw(ArgumentError("`interval` must be a non-negative integer (provided `interval = $interval`)"))
  end

  # AMR every `interval` time steps, but not after the final step
  if interval > 0
    condition = (u, t, integrator) -> integrator.iter % interval == 0 && !isfinished(integrator)
  else # disable the AMR callback except possibly for initial refinement during initialization
    condition = (u, t, integrator) -> false
  end

  to_refine  = Int[]
  to_coarsen = Int[]
  amr_cache = (; to_refine, to_coarsen)

  amr_callback = AMRCallback{typeof(controller), typeof(adaptor), typeof(amr_cache)}(
    controller, interval, adapt_initial_condition, adapt_initial_condition_only_refine,
    dynamic_load_balancing, adaptor, amr_cache)

  DiscreteCallback(condition, amr_callback,
                   save_positions=(false,false),
                   initialize=initialize!)
end

function AMRCallback(semi, controller; kwargs...)
  adaptor = AdaptorAMR(semi)
  AMRCallback(semi, controller, adaptor; kwargs...)
end

function AdaptorAMR(semi; kwargs...)
  mesh, _, solver, _ = mesh_equations_solver_cache(semi)
  AdaptorAMR(mesh, solver; kwargs...)
end

# TODO: Taal bikeshedding, implement a method with less information and the signature
# function Base.show(io::IO, cb::DiscreteCallback{Condition,Affect!}) where {Condition, Affect!<:AMRCallback}
#   amr_callback = cb.affect!
#   print(io, "AMRCallback")
# end
function Base.show(io::IO, mime::MIME"text/plain", cb::DiscreteCallback{Condition,Affect!}) where {Condition, Affect!<:AMRCallback}
  if get(io, :compact, false)
    show(io, cb)
  else
    amr_callback = cb.affect!

    summary_header(io, "AMRCallback")
    summary_line(io, "controller", typeof(amr_callback.controller).name)
    show(increment_indent(io), mime, amr_callback.controller)
    summary_line(io, "interval", amr_callback.interval)
    summary_line(io, "adapt IC", amr_callback.adapt_initial_condition ? "yes" : "no",)
    if amr_callback.adapt_initial_condition
      summary_line(io, "│ only refine", amr_callback.adapt_initial_condition_only_refine ? "yes" : "no")
    end
    summary_footer(io)
  end
end

# The function below is used to control the output depending on whether or not AMR is enabled.
"""
    uses_amr(callback)

Checks whether the provided callback or `CallbackSet` is an [`AMRCallback`](@ref)
or contains one.
"""
uses_amr(cb) = false
uses_amr(cb::DiscreteCallback{Condition,Affect!}) where {Condition, Affect!<:AMRCallback} = true
uses_amr(callbacks::CallbackSet) = mapreduce(uses_amr, |, callbacks.discrete_callbacks)

function get_element_variables!(element_variables, u, mesh, equations, solver, cache,
                                amr_callback::AMRCallback; kwargs...)
  get_element_variables!(element_variables, u, mesh, equations, solver, cache,
                         amr_callback.controller, amr_callback; kwargs...)
end

function initialize!(cb::DiscreteCallback{Condition,Affect!}, u, t, integrator) where {Condition, Affect!<:AMRCallback}
  amr_callback = cb.affect!
  semi = integrator.p

  @timeit_debug timer() "initial condition AMR" if amr_callback.adapt_initial_condition
    # iterate until mesh does not change anymore
    has_changed = amr_callback(integrator,
                               only_refine=amr_callback.adapt_initial_condition_only_refine)
    while has_changed
      compute_coefficients!(integrator.u, t, semi)
      u_modified!(integrator, true)
      has_changed = amr_callback(integrator,
                                 only_refine=amr_callback.adapt_initial_condition_only_refine)
    end
  end

  return nothing
end

# TODO: Taal remove?
# function (cb::DiscreteCallback{Condition,Affect!})(ode::ODEProblem) where {Condition, Affect!<:AMRCallback}
#   amr_callback = cb.affect!
#   semi = ode.p

#   @timeit_debug timer() "initial condition AMR" if amr_callback.adapt_initial_condition
#     # iterate until mesh does not change anymore
#     has_changed = true
#     while has_changed
#       has_changed = amr_callback(ode.u0, semi,
#                                  only_refine=amr_callback.adapt_initial_condition_only_refine)
#       compute_coefficients!(ode.u0, ode.tspan[1], semi)
#     end
#   end

#   return nothing
# end

function (amr_callback::AMRCallback)(integrator; kwargs...)
  u_ode = integrator.u
  semi = integrator.p

  @timeit_debug timer() "AMR" begin
    has_changed = amr_callback(u_ode, semi,
                               integrator.t, integrator.iter; kwargs...)
    if has_changed
      resize!(integrator, length(u_ode))
      u_modified!(integrator, true)
    end
  end

  return has_changed
end

@inline function (amr_callback::AMRCallback)(u_ode::AbstractVector,
                                             semi::SemidiscretizationHyperbolic,
                                             t, iter;
                                             kwargs...)
  amr_callback(u_ode, mesh_equations_solver_cache(semi)..., t, iter; kwargs...)
end

# `passive_args` is currently used for Euler with self-gravity to adapt the gravity solver
# passively without querying its indicator, based on the assumption that both solvers use
# the same mesh. That's a hack and should be improved in the future once we have more examples
# and a better understandin of such a coupling.
# `passive_args` is expected to be an iterable of `Tuple`s of the form
# `(p_u_ode, p_mesh, p_equations, p_dg, p_cache)`.
function (amr_callback::AMRCallback)(u_ode::AbstractVector, mesh::TreeMesh,
                                     equations, dg::DG, cache,
                                     t, iter;
                                     only_refine=false, only_coarsen=false,
                                     passive_args=())
  @unpack controller, adaptor = amr_callback

  u = wrap_array(u_ode, mesh, equations, dg, cache)
  lambda = @timeit_debug timer() "indicator" controller(u, mesh, equations, dg, cache,
                                                        t=t, iter=iter)

  if mpi_isparallel()
    # Collect lambda for all elements
    lambda_global = Vector{eltype(lambda)}(undef, nelementsglobal(dg, cache))
    # Use parent because n_elements_by_rank is an OffsetArray
    recvbuf = MPI.VBuffer(lambda_global, parent(cache.mpi_cache.n_elements_by_rank))
    MPI.Allgatherv!(lambda, recvbuf, mpi_comm())
    lambda = lambda_global
  end

  leaf_cell_ids = leaf_cells(mesh.tree)
  @boundscheck begin
   @assert axes(lambda) == axes(leaf_cell_ids) ("Indicator (axes = $(axes(lambda))) and leaf cell (axes = $(axes(leaf_cell_ids))) arrays have different axes")
  end

  @unpack to_refine, to_coarsen = amr_callback.amr_cache
  empty!(to_refine)
  empty!(to_coarsen)
  for element in 1:length(lambda)
    controller_value = lambda[element]
    if controller_value > 0
      push!(to_refine, leaf_cell_ids[element])
    elseif controller_value < 0
      push!(to_coarsen, leaf_cell_ids[element])
    end
  end

  @timeit_debug timer() "refine" if !only_coarsen && !isempty(to_refine)
    # refine mesh
    refined_original_cells = @timeit_debug timer() "mesh" refine!(mesh.tree, to_refine)

    # refine solver
    @timeit_debug timer() "solver" refine!(u_ode, adaptor, mesh, equations, dg, cache, refined_original_cells)
    for (p_u_ode, p_mesh, p_equations, p_dg, p_cache) in passive_args
      @timeit_debug timer() "passive solver" refine!(p_u_ode, adaptor, p_mesh, p_equations, p_dg, p_cache, refined_original_cells)
    end
  else
    # If there is nothing to refine, create empty array for later use
    refined_original_cells = Int[]
  end

  @timeit_debug timer() "coarsen" if !only_refine && !isempty(to_coarsen)
    # Since the cells may have been shifted due to refinement, first we need to
    # translate the old cell ids to the new cell ids
    if !isempty(to_coarsen)
      to_coarsen = original2refined(to_coarsen, refined_original_cells, mesh)
    end

    # Next, determine the parent cells from which the fine cells are to be
    # removed, since these are needed for the coarsen! function. However, since
    # we only want to coarsen if *all* child cells are marked for coarsening,
    # we count the coarsening indicators for each parent cell and only coarsen
    # if all children are marked as such (i.e., where the count is 2^ndims). At
    # the same time, check if a cell is marked for coarsening even though it is
    # *not* a leaf cell -> this can only happen if it was refined due to 2:1
    # smoothing during the preceding refinement operation.
    parents_to_coarsen = zeros(Int, length(mesh.tree))
    for cell_id in to_coarsen
      # If cell has no parent, it cannot be coarsened
      if !has_parent(mesh.tree, cell_id)
        continue
      end

      # If cell is not leaf (anymore), it cannot be coarsened
      if !is_leaf(mesh.tree, cell_id)
        continue
      end

      # Increase count for parent cell
      parent_id = mesh.tree.parent_ids[cell_id]
      parents_to_coarsen[parent_id] += 1
    end

    # Extract only those parent cells for which all children should be coarsened
    to_coarsen = collect(1:length(parents_to_coarsen))[parents_to_coarsen .== 2^ndims(mesh)]

    # Finally, coarsen mesh
    coarsened_original_cells = @timeit_debug timer() "mesh" coarsen!(mesh.tree, to_coarsen)

    # Convert coarsened parent cell ids to the list of child cell ids that have
    # been removed, since this is the information that is expected by the solver
    removed_child_cells = zeros(Int, n_children_per_cell(mesh.tree) * length(coarsened_original_cells))
    for (index, coarse_cell_id) in enumerate(coarsened_original_cells)
      for child in 1:n_children_per_cell(mesh.tree)
        removed_child_cells[n_children_per_cell(mesh.tree) * (index-1) + child] = coarse_cell_id + child
      end
    end

    # coarsen solver
    @timeit_debug timer() "solver" coarsen!(u_ode, adaptor, mesh, equations, dg, cache, removed_child_cells)
    for (p_u_ode, p_mesh, p_equations, p_dg, p_cache) in passive_args
      @timeit_debug timer() "passive solver" coarsen!(p_u_ode, adaptor, p_mesh, p_equations, p_dg, p_cache, removed_child_cells)
    end
  else
    # If there is nothing to coarsen, create empty array for later use
    coarsened_original_cells = Int[]
  end

  # Store whether there were any cells coarsened or refined
  has_changed = !isempty(refined_original_cells) || !isempty(coarsened_original_cells)
  if has_changed # TODO: Taal decide, where shall we set this?
    # don't set it to has_changed since there can be changes from earlier calls
    mesh.unsaved_changes = true
  end

  # Dynamically balance computational load by first repartitioning the mesh and then redistributing the cells/elements
  if has_changed && mpi_isparallel() && amr_callback.dynamic_load_balancing
    @timeit_debug timer() "dynamic load balancing" begin
      old_mpi_ranks_per_cell = copy(mesh.tree.mpi_ranks)

      partition!(mesh)

      rebalance_solver!(u_ode, mesh, equations, dg, cache, old_mpi_ranks_per_cell)
    end
  end

  # Return true if there were any cells coarsened or refined, otherwise false
  return has_changed
end

# After refining cells, shift original cell ids to match new locations
# Note: Assumes sorted lists of original and refined cell ids!
# Note: `mesh` is only required to extract ndims
function original2refined(original_cell_ids, refined_original_cells, mesh)
  # Sanity check
  @assert issorted(original_cell_ids) "`original_cell_ids` not sorted"
  @assert issorted(refined_original_cells) "`refined_cell_ids` not sorted"

  # Create array with original cell ids (not yet shifted)
  shifted_cell_ids = collect(1:original_cell_ids[end])

  # Loop over refined original cells and apply shift for all following cells
  for cell_id in refined_original_cells
    # Only calculate shifts for cell ids that are relevant
    if cell_id > length(shifted_cell_ids)
      break
    end

    # Shift all subsequent cells by 2^ndims ids
    shifted_cell_ids[(cell_id + 1):end] .+= 2^ndims(mesh)
  end

  # Convert original cell ids to their shifted values
  return shifted_cell_ids[original_cell_ids]
end

"""
    ControllerThreeLevel(semi, indicator; base_level=1,
                                          med_level=base_level, med_threshold=0.0,
                                          max_level=base_level, max_threshold=1.0)

An AMR controller based on three levels (in decending order of precedence):
- set the target level to `max_level` if `indicator > max_threshold`
- set the target level to `med_level` if `indicator > med_threshold`;
  if `med_level < 0`, set the target level to the current level
- set the target level to `base_level` otherwise
"""
struct ControllerThreeLevel{RealT<:Real, Indicator, Cache}
  base_level::Int
  med_level ::Int
  max_level ::Int
  med_threshold::RealT
  max_threshold::RealT
  indicator::Indicator
  cache::Cache
end

function ControllerThreeLevel(semi, indicator; base_level=1,
                                               med_level=base_level, med_threshold=0.0,
                                               max_level=base_level, max_threshold=1.0)
  med_threshold, max_threshold = promote(med_threshold, max_threshold)
  cache = create_cache(ControllerThreeLevel, semi)
  ControllerThreeLevel{typeof(max_threshold), typeof(indicator), typeof(cache)}(
    base_level, med_level, max_level, med_threshold, max_threshold, indicator, cache)
end

create_cache(indicator_type::Type{ControllerThreeLevel}, semi) = create_cache(indicator_type, mesh_equations_solver_cache(semi)...)

function Base.show(io::IO, controller::ControllerThreeLevel)
  print(io, "ControllerThreeLevel(")
  print(io, controller.indicator)
  print(io, ", base_level=", controller.base_level)
  print(io, ", med_level=",  controller.med_level)
  print(io, ", max_level=",  controller.max_level)
  print(io, ", med_threshold=", controller.med_threshold)
  print(io, ", max_threshold=", controller.max_threshold)
  print(io, ")")
end

function Base.show(io::IO, mime::MIME"text/plain", controller::ControllerThreeLevel)
  if get(io, :compact, false)
    show(io, controller)
  else
    summary_header(io, "ControllerThreeLevel")
    summary_line(io, "indicator", typeof(controller.indicator).name)
    show(increment_indent(io), mime, controller.indicator)
    summary_line(io, "base_level", controller.base_level)
    summary_line(io, "med_level", controller.med_level)
    summary_line(io, "max_level", controller.max_level)
    summary_line(io, "med_threshold", controller.med_threshold)
    summary_line(io, "max_threshold", controller.max_threshold)
    summary_footer(io)
  end
end

function get_element_variables!(element_variables, u, mesh, equations, solver, cache,
                                controller::ControllerThreeLevel, amr_callback::AMRCallback;
                                kwargs...)
  # call the indicator to get up-to-date values for IO
  controller.indicator(u, equations, solver, cache; kwargs...)
  get_element_variables!(element_variables, controller.indicator, amr_callback)
end

function get_element_variables!(element_variables, indicator::AbstractIndicator, ::AMRCallback)
  element_variables[:indicator_amr] = indicator.cache.alpha
  return nothing
end

# TODO: Taal refactor, merge the two loops of ControllerThreeLevel and IndicatorLöhner etc.?
#       But that would remove the simplest possibility to write that stuff to a file...
#       We could of course implement some additional logic and workarounds, but is it worth the effort?
function (controller::ControllerThreeLevel)(u::AbstractArray{<:Any},
                                            mesh::TreeMesh, equations, dg::DG, cache;
                                            kwargs...)

  @unpack controller_value = controller.cache
  resize!(controller_value, nelements(dg, cache))

  alpha = controller.indicator(u, equations, dg, cache; kwargs...)

  Threads.@threads for element in eachelement(dg, cache)
    cell_id = cache.elements.cell_ids[element]
    current_level = mesh.tree.levels[cell_id]

    # set target level
    target_level = current_level
    if alpha[element] > controller.max_threshold
      target_level = controller.max_level
    elseif alpha[element] > controller.med_threshold
      if controller.med_level > 0
        target_level = controller.med_level
        # otherwise, target_level = current_level
        # set med_level = -1 to implicitly use med_level = current_level
      end
    else
      target_level = controller.base_level
    end

    # compare target level with actual level to set controller
    if current_level < target_level
      controller_value[element] = 1 # refine!
    elseif current_level > target_level
      controller_value[element] = -1 # coarsen!
    else
      controller_value[element] = 0 # we're good
    end
  end

  return controller_value
end

"""
    ControllerThreeLevelCombined(semi, indicator_primary, indicator_secondary;
                                 base_level=1,
                                 med_level=base_level, med_threshold=0.0,
                                 max_level=base_level, max_threshold=1.0,
                                 max_threshold_secondary=1.0)

An AMR controller based on three levels (in decending order of precedence):
- set the target level to `max_level` if `indicator_primary > max_threshold`
- set the target level to `med_level` if `indicator_primary > med_threshold`;
  if `med_level < 0`, set the target level to the current level
- set the target level to `base_level` otherwise
If `indicator_secondary >= max_threshold_secondary`,
set the target level to `max_level`.
"""
struct ControllerThreeLevelCombined{RealT<:Real, IndicatorPrimary, IndicatorSecondary, Cache}
  base_level::Int
  med_level ::Int
  max_level ::Int
  med_threshold::RealT
  max_threshold::RealT
  max_threshold_secondary::RealT
  indicator_primary::IndicatorPrimary
  indicator_secondary::IndicatorSecondary
  cache::Cache
end

function ControllerThreeLevelCombined(semi, indicator_primary, indicator_secondary;
                                      base_level=1,
                                      med_level=base_level, med_threshold=0.0,
                                      max_level=base_level, max_threshold=1.0,
                                      max_threshold_secondary=1.0)
  med_threshold, max_threshold, max_threshold_secondary = promote(med_threshold, max_threshold, max_threshold_secondary)
  cache = create_cache(ControllerThreeLevelCombined, semi)
  ControllerThreeLevelCombined{typeof(max_threshold), typeof(indicator_primary), typeof(indicator_secondary), typeof(cache)}(
    base_level, med_level, max_level, med_threshold, max_threshold,
    max_threshold_secondary, indicator_primary, indicator_secondary, cache)
end

create_cache(indicator_type::Type{ControllerThreeLevelCombined}, semi) = create_cache(indicator_type, mesh_equations_solver_cache(semi)...)

function Base.show(io::IO, controller::ControllerThreeLevelCombined)
  print(io, "ControllerThreeLevelCombined(")
  print(io, controller.indicator_primary)
  print(io, ", ", controller.indicator_secondary)
  print(io, ", base_level=", controller.base_level)
  print(io, ", med_level=",  controller.med_level)
  print(io, ", max_level=",  controller.max_level)
  print(io, ", med_threshold=", controller.med_threshold)
  print(io, ", max_threshold_secondary=", controller.max_threshold_secondary)
  print(io, ")")
end

function Base.show(io::IO, mime::MIME"text/plain", controller::ControllerThreeLevelCombined)
  if get(io, :compact, false)
    show(io, controller)
  else
    summary_header(io, "ControllerThreeLevelCombined")
    summary_line(io, "primary indicator", typeof(controller.indicator_primary).name)
    show(increment_indent(io), mime, controller.indicator_primary)
    summary_line(io, "secondary indicator", typeof(controller.indicator_secondary).name)
    show(increment_indent(io), mime, controller.indicator_secondary)
    summary_line(io, "base_level", controller.base_level)
    summary_line(io, "med_level", controller.med_level)
    summary_line(io, "max_level", controller.max_level)
    summary_line(io, "med_threshold", controller.med_threshold)
    summary_line(io, "max_threshold", controller.max_threshold)
    summary_line(io, "max_threshold_secondary", controller.max_threshold_secondary)
    summary_footer(io)
  end
end

function get_element_variables!(element_variables, u, mesh, equations, solver, cache,
                                controller::ControllerThreeLevelCombined, amr_callback::AMRCallback;
                                kwargs...)
  # call the indicator to get up-to-date values for IO
  controller.indicator_primary(u, equations, solver, cache; kwargs...)
  get_element_variables!(element_variables, controller.indicator_primary, amr_callback)
end

function (controller::ControllerThreeLevelCombined)(u::AbstractArray{<:Any},
                                                    mesh::TreeMesh, equations, dg::DG, cache;
                                                    kwargs...)

  @unpack controller_value = controller.cache
  resize!(controller_value, nelements(dg, cache))

  alpha = controller.indicator_primary(u, equations, dg, cache; kwargs...)
  alpha_secondary = controller.indicator_secondary(u, equations, dg, cache)

  Threads.@threads for element in eachelement(dg, cache)
    cell_id = cache.elements.cell_ids[element]
    current_level = mesh.tree.levels[cell_id]

    # set target level
    target_level = current_level
    if alpha[element] > controller.max_threshold
      target_level = controller.max_level
    elseif alpha[element] > controller.med_threshold
      if controller.med_level > 0
        target_level = controller.med_level
        # otherwise, target_level = current_level
        # set med_level = -1 to implicitly use med_level = current_level
      end
    else
      target_level = controller.base_level
    end

    if alpha_secondary[element] >= controller.max_threshold_secondary
      target_level = controller.max_level
    end

    # compare target level with actual level to set controller
    if current_level < target_level
      controller_value[element] = 1 # refine!
    elseif current_level > target_level
      controller_value[element] = -1 # coarsen!
    else
      controller_value[element] = 0 # we're good
    end
  end

  return controller_value
end

include("amr_dg1d.jl")
include("amr_dg2d.jl")
include("amr_dg3d.jl")


text:
"""
`Manopt.jl` – Optimization on Manifolds in Julia.
"""
module Manopt
    using Colors
    using ColorSchemes
    using ColorTypes
    using Markdown
    using LinearAlgebra
    using Dates: Period, Nanosecond, value
    import Random:
        rand,
        randperm
    import Base:
        copy,
        identity
    import ManifoldsBase:
        ℝ,
        ℂ,
        ×,
        ^,
        AbstractVectorTransportMethod,
        AbstractRetractionMethod,
        ExponentialRetraction,
        ParallelTransport,
        Manifold,
        allocate_result,
        allocate_result_type,
        distance,
        exp,
        exp!,
        log,
        log!,
        injectivity_radius,
        inner,
        geodesic,
        manifold_dimension,
        norm,
        project,
        project!,
        retract,
        retract!,
        shortest_geodesic,
        vector_transport_to,
        vector_transport_to!,
        zero_tangent_vector,
        zero_tangent_vector!,
        DiagonalizingOrthonormalBasis,
        get_basis,
        get_coordinates,
        get_vector,
        get_vectors,
        representation_size
    using Manifolds:
        AbstractPowerManifold,
        PowerManifold
    using Manifolds: #temporary for random
        Circle,
        Euclidean,
        Grassmann,
        Hyperbolic,
        ProductManifold,
        Rotations,
        SymmetricPositiveDefinite,
        Stiefel,
        Sphere
    using Manifolds: # Wishlist for Base
        NestedPowerRepresentation,
        ArrayPowerRepresentation,
        mean,
        median,
        get_iterator,
        _read,
        _write,
        power_dimensions,
        ArrayReshaper,
        prod_point,
        ShapeSpecification,
        ProductRepr,
        submanifold_components,
        submanifold_component,
        get_component,
        set_component!

    include("plans/plan.jl")
    # Functions
    include("functions/manifold.jl")
    include("functions/adjointDifferentials.jl")
    include("functions/costs.jl")
    include("functions/differentials.jl")
    include("functions/gradients.jl")
    include("functions/jacobiFields.jl")
    include("functions/proximalMaps.jl")
    # solvers general framework
    include("solvers/solver.jl")
    # specific solvers
    include("solvers/conjugate_gradient_descent.jl")
    include("solvers/cyclic_proximal_point.jl")
    include("solvers/DouglasRachford.jl")
    include("solvers/NelderMead.jl")
    include("solvers/steepest_descent.jl")
    include("solvers/truncated_conjugate_gradient_descent.jl")
    include("solvers/trust_regions.jl")
    include("solvers/subgradient.jl")
    include("solvers/debug_solver.jl")
    include("solvers/record_solver.jl")
    include("plots/SpherePlots.jl")
    include("helpers/errorMeasures.jl")
    include("helpers/exports/Asymptote.jl")
    include("data/artificialDataFunctions.jl")

    include("random.jl")

export ×, ^, ℝ, ℂ
export AbstractOptionsAction, StoreOptionsAction
export has_storage, get_storage, update_storage!

export random_point, random_tangent, mid_point, mid_point!, reflect, sym_rem

export βdifferential_geodesic_startpoint, βdifferential_exp_basepoint
export βdifferential_exp_argument, βdifferential_log_basepoint, βdifferential_log_argument
export adjoint_differential_geodesic_startpoint, adjoint_differential_geodesic_endpoint
export adjoint_differential_exp_basepoint, adjoint_differential_exp_argument
export adjoint_differential_log_basepoint, adjoint_differential_log_argument
export adjoint_differential_forward_logs
export asymptote_export_S2_signals, asymptote_export_S2_data, asymptote_export_SPD
export render_asymptote
export costL2TV, costL2TVTV2, costL2TV2, costTV, costTV2, costIntrICTV12
export differential_geodesic_startpoint, differential_geodesic_endpoint
export differential_exp_basepoint, differential_exp_argument
export differential_log_basepoint, differential_log_argument, differential_forward_logs
export jacobi_field, adjoint_Jacobi_field
export ∇TV, ∇TV2, ∇intrinsic_infimal_convolution_TV12, forward_logs, ∇distance
export get_cost, get_gradient, get_subgradient, getProximalMap, get_options, get_initial_stepsize
export getHessian, approxHessianFD
export meanSquaredError, meanAverageError
export prox_distance, prox_TV, prox_parallel_TV, prox_TV2, prox_collaborative_TV
export random_point, random_tangent
export stopIfResidualIsReducedByFactor, stopIfResidualIsReducedByPower, StopWhenCurvatureIsNegative, StopWhenTrustRegionIsExceeded
export StopAfterIteration, StopWhenChangeLess, StopWhenGradientNormLess, StopWhenCostLess
export StopAfter, StopWhenAll, StopWhenAny
export get_active_stopping_criteria, get_stopping_criteria, get_reason

export DebugOptions, DebugAction, DebugGroup, DebugEntry, DebugEntryChange, DebugEvery
export DebugChange, DebugIterate, DebugIteration, DebugDivider
export DebugCost, DebugStoppingCriterion, DebugFactory, DebugActionFactory
export DebugGradient, DebugGradientNorm, DebugStepsize

export RecordGradient, RecordGradientNorm, RecordStepsize

export CostProblem, Problem, SubGradientProblem, GradientProblem, HessianProblem

export NelderMead, steepest_descent, subgradient_method, truncated_conjugate_gradient_descent, trust_regions
export cyclic_proximal_point, conjugate_gradient_descent

export DebugGradient, DebugGradientNorm, DebugStepsize

export Options, get_options
export is_options_decorator, dispatch_options_decorator

export
    ConjugateGradientDescentOptions,
    GradientDescentOptions,
    HessianOptions,
    SubGradientMethodOptions,
    NelderMeadOptions,
    TruncatedConjugateGradientOptions,
    TrustRegionsOptions

export DirectionUpdateRule,
    SteepestDirectionUpdateRule,
    HeestenesStiefelCoefficient,
    FletcherReevesCoefficient,
    PolakRibiereCoefficient,
    ConjugateDescentCoefficient,
    LiuStoreyCoefficient,
    DaiYuanCoefficient,
    HagerZhangCoefficient

export StoppingCriterion, StoppingCriterionSet, Stepsize
export EvalOrder, LinearEvalOrder, RandomEvalOrder, FixedRandomEvalOrder

export decorate_options
export initialize_solver!, step_solver!, get_solver_result, stop_solver!
export solve

export ConstantStepsize, DecreasingStepsize
export Linesearch, ArmijoLinesearch
export get_stepsize, get_initial_stepsize, get_last_stepsize

end


text:
using DeltaSugar
using Documenter

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

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

deploydocs(;
    repo="github.com/caseykneale/DeltaSugar.jl",
)


text:
using Test, MakieCore

# Main tests live in Makie.jl, but we should write some unit tests going forward!
using MakieCore: @recipe, Attributes, Plot
import MakieCore: plot!, convert_arguments, used_attributes, plot
import MakieCore: Observable, PointBased

struct AbstractTimeseriesSolution
    results::Vector{Float32}
end

function plot!(plot::Plot(AbstractTimeseriesSolution))
    # plot contains any keyword arguments that you pass to plot(series; kw...)
    var = get(plot, :var, Observable(5))
    density!(plot, map((v, r)-> v .* r.results, var, plot[1]))
end

struct Test2
    series::Any
end

struct Solution
    data::Any
end

function plot!(plot::Plot(Test2))
    arg1 = plot[1]
    scatter!(plot, arg1[].series)
    ser = AbstractTimeseriesSolution(arg1[].series)
    sol = Solution(arg1[].series)
    plot!(plot, ser, var = 10)
    scatter!(plot, sol, attribute = 3, color=:red)
end

used_attributes(::Any, x::Solution) = (:attribute,)

# Convert for all point based types (lines, scatter)
function convert_arguments(p::MakieCore.PointBased, x::Solution; attribute = 1.0)
    return convert_arguments(p, x.data .* attribute)
end

# using GLMakie
# using CairoMakie
# plot(Test2(rand(Float32, 10)))

# function MakieCore.plot!(myplot::MyPlot{<:Tuple{<:AbstractVector{<:Number}}})
#     lines!(myplot, rand(10), color = myplot[:plot_color])
#     plot!(myplot, myplot[:x])
#     myplot
# end

# myplot(1:3)

# function MakieCore.plot(P::Type{<: AbstractPlot}, fig::Makie.FigurePosition, arg::Solution; axis = NamedTuple(), kwargs...)

#     menu = Menu(fig, options = ["viridis", "heat", "blues"])

#     funcs = [sqrt, x->x^2, sin, cos]

#     menu2 = Menu(fig, options = zip(["Square Root", "Square", "Sine", "Cosine"], funcs))

#     fig[1, 1] = vgrid!(
#         Label(fig, "Colormap", width = nothing),
#         menu,
#         Label(fig, "Function", width = nothing),
#         menu2;
#         tellheight = false, width = 200)

#     ax = Axis(fig[1, 2]; axis...)

#     func = Node{Any}(funcs[1])

#     ys = @lift($func.(arg.data))

#     scat = plot!(ax, P, Attributes(color = ys), ys)

#     cb = Colorbar(fig[1, 3], scat)

#     on(menu.selection) do s
#         scat.colormap = s
#     end

#     on(menu2.selection) do s
#         func[] = s
#         autolimits!(ax)
#     end

#     menu2.is_open = true

#     return Makie.AxisPlot(ax, scat)
# end

# f = Figure();
# lines(f[1, 1], Solution(0:0.3:10))
# scatter(f[1, 2], Solution(0:0.3:10))
# f |> display

# @recipe(MyPlot, x) do scene
#     Theme(
#         plot_color = :red
#     )
# end

# function Makie.plot!(p::MyPlot)
#     @show p.transformation.transform_func[]
#     scatter!(p, p[1])
# end

# myplot(rand(4), axis=(xscale=log10,))
# using CairoMakie
# plot(Test2(rand(Float32, 10)))

# scatter(Solution(rand(4)), attribute = 10)


text:
#! format: off

abstract type AbstractLoadFormulation <: AbstractDeviceFormulation end
abstract type AbstractControllablePowerLoadFormulation <: AbstractLoadFormulation end
struct StaticPowerLoad <: AbstractLoadFormulation end
struct InterruptiblePowerLoad <: AbstractControllablePowerLoadFormulation end
struct DispatchablePowerLoad <: AbstractControllablePowerLoadFormulation end

########################### ElectricLoad ####################################

get_variable_sign(_, ::Type{<:PSY.ElectricLoad}, ::AbstractLoadFormulation) = -1.0

########################### ActivePowerVariable, ElectricLoad ####################################

get_variable_binary(::ActivePowerVariable, ::Type{<:PSY.ElectricLoad}, ::AbstractLoadFormulation) = false

get_variable_expression_name(::ActivePowerVariable, ::Type{<:PSY.ElectricLoad}) = :nodal_balance_active

get_variable_lower_bound(::ActivePowerVariable, d::PSY.ElectricLoad, ::AbstractLoadFormulation) = 0.0
get_variable_upper_bound(::ActivePowerVariable, d::PSY.ElectricLoad, ::AbstractLoadFormulation) = PSY.get_active_power(d)

########################### ReactivePowerVariable, ElectricLoad ####################################

get_variable_binary(::ReactivePowerVariable, ::Type{<:PSY.ElectricLoad}, ::AbstractLoadFormulation) = false

get_variable_expression_name(::ReactivePowerVariable, ::Type{<:PSY.ElectricLoad}) = :nodal_balance_reactive

get_variable_lower_bound(::ReactivePowerVariable, d::PSY.ElectricLoad, ::AbstractLoadFormulation) = 0.0
get_variable_upper_bound(::ReactivePowerVariable, d::PSY.ElectricLoad, ::AbstractLoadFormulation) = PSY.get_reactive_power(d)

########################### ReactivePowerVariable, ElectricLoad ####################################

get_variable_binary(::OnVariable, ::Type{<:PSY.ElectricLoad}, ::AbstractLoadFormulation) = true

#! format: on

####################################### Reactive Power Constraints #########################
"""
Reactive Power Constraints on Controllable Loads Assume Constant power_factor
"""
function DeviceRangeConstraintSpec(
    ::Type{<:RangeConstraint},
    ::Type{ReactivePowerVariable},
    ::Type{<:PSY.ElectricLoad},
    ::Type{<:AbstractControllablePowerLoadFormulation},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
    use_parameters::Bool,
    use_forecasts::Bool,
)
    return DeviceRangeConstraintSpec(;
        custom_optimization_container_func = custom_reactive_power_constraints!,
    )
end

function custom_reactive_power_constraints!(
    optimization_container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::Type{<:AbstractControllablePowerLoadFormulation},
) where {T <: PSY.ElectricLoad}
    time_steps = model_time_steps(optimization_container)
    constraint = JuMPConstraintArray(undef, [PSY.get_name(d) for d in devices], time_steps)
    assign_constraint!(optimization_container, REACTIVE, T, constraint)

    for t in time_steps, d in devices
        name = PSY.get_name(d)
        pf = sin(atan((PSY.get_max_reactive_power(d) / PSY.get_max_active_power(d))))
        reactive = get_variable(optimization_container, REACTIVE_POWER, T)[name, t]
        real = get_variable(optimization_container, ACTIVE_POWER, T)[name, t] * pf
        constraint[name, t] =
            JuMP.@constraint(optimization_container.JuMPmodel, reactive == real)
    end
end

function DeviceRangeConstraintSpec(
    ::Type{<:RangeConstraint},
    ::Type{ActivePowerVariable},
    ::Type{T},
    ::Type{<:DispatchablePowerLoad},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
    use_parameters::Bool,
    use_forecasts::Bool,
) where {T <: PSY.ElectricLoad}
    if (!use_parameters && !use_forecasts)
        return DeviceRangeConstraintSpec(;
            range_constraint_spec = RangeConstraintSpec(;
                constraint_name = make_constraint_name(
                    RangeConstraint,
                    ActivePowerVariable,
                    T,
                ),
                variable_name = make_variable_name(ActivePowerVariable, T),
                limits_func = x -> (min = 0.0, max = PSY.get_active_power(x)),
                constraint_func = device_range!,
                constraint_struct = DeviceRangeConstraintInfo,
            ),
        )
    end

    return DeviceRangeConstraintSpec(;
        timeseries_range_constraint_spec = TimeSeriesConstraintSpec(
            constraint_name = make_constraint_name(RangeConstraint, ActivePowerVariable, T),
            variable_name = make_variable_name(ActivePowerVariable, T),
            parameter_name = use_parameters ? ACTIVE_POWER : nothing,
            forecast_label = "max_active_power",
            multiplier_func = x -> PSY.get_max_active_power(x),
            constraint_func = use_parameters ? device_timeseries_param_ub! :
                              device_timeseries_ub!,
        ),
    )
end

function DeviceRangeConstraintSpec(
    ::Type{<:RangeConstraint},
    ::Type{ActivePowerVariable},
    ::Type{T},
    ::Type{<:InterruptiblePowerLoad},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
    use_parameters::Bool,
    use_forecasts::Bool,
) where {T <: PSY.ElectricLoad}
    if (!use_parameters && !use_forecasts)
        return DeviceRangeConstraintSpec(;
            range_constraint_spec = RangeConstraintSpec(;
                constraint_name = make_constraint_name(
                    RangeConstraint,
                    ActivePowerVariable,
                    T,
                ),
                variable_name = make_variable_name(ActivePowerVariable, T),
                bin_variable_names = [make_variable_name(OnVariable, T)],
                limits_func = x -> (min = 0.0, max = PSY.get_active_power(x)),
                constraint_func = device_semicontinuousrange!,
                constraint_struct = DeviceRangeConstraintInfo,
            ),
        )
    end

    return DeviceRangeConstraintSpec(;
        timeseries_range_constraint_spec = TimeSeriesConstraintSpec(
            constraint_name = make_constraint_name(RangeConstraint, ActivePowerVariable, T),
            variable_name = make_variable_name(ActivePowerVariable, T),
            bin_variable_name = make_variable_name(OnVariable, T),
            parameter_name = use_parameters ? ON : nothing,
            forecast_label = "max_active_power",
            multiplier_func = x -> PSY.get_max_active_power(x),
            constraint_func = use_parameters ? device_timeseries_ub_bigM! :
                              device_timeseries_ub_bin!,
        ),
    )
end

########################## Addition to the nodal balances ##################################

function NodalExpressionSpec(
    ::Type{T},
    ::Type{<:PM.AbstractPowerModel},
    use_forecasts::Bool,
) where {T <: PSY.ElectricLoad}
    return NodalExpressionSpec(
        "max_active_power",
        REACTIVE_POWER,
        use_forecasts ? x -> PSY.get_max_reactive_power(x) : x -> PSY.get_reactive_power(x),
        -1.0,
        T,
    )
end

function NodalExpressionSpec(
    ::Type{T},
    ::Type{<:PM.AbstractActivePowerModel},
    use_forecasts::Bool,
) where {T <: PSY.ElectricLoad}
    return NodalExpressionSpec(
        "max_active_power",
        ACTIVE_POWER,
        use_forecasts ? x -> PSY.get_max_active_power(x) : x -> PSY.get_active_power(x),
        -1.0,
        T,
    )
end

############################## FormulationControllable Load Cost ###########################
function AddCostSpec(
    ::Type{T},
    ::Type{DispatchablePowerLoad},
    ::OptimizationContainer,
) where {T <: PSY.ControllableLoad}
    cost_function = x -> (x === nothing ? 1.0 : PSY.get_variable(x))
    return AddCostSpec(;
        variable_type = ActivePowerVariable,
        component_type = T,
        variable_cost = cost_function,
        multiplier = OBJECTIVE_FUNCTION_NEGATIVE,
    )
end

function AddCostSpec(
    ::Type{T},
    ::Type{InterruptiblePowerLoad},
    ::OptimizationContainer,
) where {T <: PSY.ControllableLoad}
    cost_function = x -> (x === nothing ? 1.0 : PSY.get_fixed(x))
    return AddCostSpec(;
        variable_type = OnVariable,
        component_type = T,
        fixed_cost = cost_function,
        multiplier = OBJECTIVE_FUNCTION_NEGATIVE,
    )
end


text:
# Copyright (c) 2014: SCS.jl contributors
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

struct DirectSolver <: LinearSolver end

scsint_t(::Type{DirectSolver}) = Clonglong

function scs_set_default_settings(
    ::Type{DirectSolver},
    stgs::ScsSettings{I},
) where {I<:Clonglong}
    return @ccall(
        direct.scs_set_default_settings(stgs::Ref{ScsSettings{I}})::Cvoid,
    )
end

function scs_init(
    ::Type{DirectSolver},
    data::ScsData{I},
    cone::ScsCone{I},
    stgs::ScsSettings{I},
) where {I<:Clonglong}
    return @ccall direct.scs_init(
        data::Ref{ScsData{I}},
        cone::Ref{ScsCone{I}},
        stgs::Ref{ScsSettings{I}},
    )::Ptr{Cvoid}
end

function scs_update(
    ::Type{DirectSolver},
    work::Ptr{Cvoid},
    b::Vector{Float64},
    c::Vector{Float64},
)
    return @ccall direct.scs_update(
        work::Ptr{Cvoid},
        b::Ref{Float64},
        c::Ref{Float64},
    )::Clonglong
end

function scs_solve(
    ::Type{DirectSolver},
    work::Ptr{Cvoid},
    solution::ScsSolution,
    info::ScsInfo{I},
    warm_start::Integer,
) where {I<:Clonglong}
    return @ccall direct.scs_solve(
        work::Ptr{Cvoid},
        solution::Ref{ScsSolution},
        info::Ref{ScsInfo{I}},
        warm_start::Clonglong,
    )::Clonglong
end

function scs_finish(::Type{DirectSolver}, work::Ptr{Cvoid})
    return @ccall direct.scs_finish(work::Ptr{Cvoid})::Cvoid
end

function scs_version()
    return unsafe_string(@ccall direct.scs_version()::Cstring)
end


text:
include("../network_functions.jl")

using Plots: plot, plot!, display
using MLDatasets: MNIST
using Random: seed!
using CuArrays
using Images
# set random seed
seed!(1301)
# disallow scalar operations
CuArrays.allowscalar(false)

# ::: DATA
train_x, _ = MNIST.traindata()
train_x = permutedims(Array{Float32}(train_x), [3,2,1])
train_x = reshape(train_x, (length(train_x[:, 1, 1]), length(train_x[1, :, 1]) * length(train_x[1, :, 1])))
get_MNIST_input_batch(x) = begin
    return CuArray(x[rand(1:length(x[:, 1])), :])
end

# ::: PARAMS
INPUT_SIZE = length(train_x[1, :, :])
LATENT_SIZE = 70
BATCH_SIZE = 1
LEARNING_RATE = 0.04
TRAIN_EPISODES = 10000

# optimizer
opt = SGD(LEARNING_RATE)

# training data
train_data_x = [get_MNIST_input_batch(train_x) for _ in 1:TRAIN_EPISODES]

# save test image for comparison
function generate_images(number_of_images, path)
    for n in 1:number_of_images
        all_test_image_ref = get_MNIST_batch(train_x, BATCH_SIZE)
        all_test_images =  all_test_image_ref |> encoder |> decoder |> Tracker.data

        random_index = rand(1:BATCH_SIZE)
        test_image = all_test_images[:, random_index]
        test_image_ref = all_test_image_ref[:, random_index]

        test_image_ref_r = reshape(test_image_ref, (28, 28))
        test_image_r = clamp.(reshape(test_image, (28, 28)), 0, 1)
        save("$path/$(n)_reference_.png", colorview(Gray, test_image_ref_r))
        save("$path/$(n)_reconstruction.png", colorview(Gray, test_image_r))
    end
end


text:
@testset "Geometric operations" begin
  @testset "cat" begin
    d₁ = RegularGrid(10,10)
    d₂ = PointSet(rand(2,10))
    d = vcat(d₁, d₂)
    @test nelms(d) == 110

    d₁ = PointSet(rand(2,3))
    d₂ = PointSet(rand(2,2))
    s₁ = georef((a=[1,2,3],b=[4,5,6]), d₁)
    s₂ = georef((a=[7.,8.],c=["foo","bar"]), d₂)
    s = vcat(s₁, s₂)
    @test nelms(s) == 5
    @test isequal(s[:a], [1.,2.,3.,7.,8.])
    @test isequal(s[:b], [4,5,6,missing,missing])
    @test isequal(s[:c], [missing,missing,missing,"foo","bar"])

    coords = rand(2,3)
    d₁ = georef((a=[1,2,3],), coords)
    d₂ = georef((b=[4.,5.,6.],), coords)
    d = hcat(d₁, d₂)
    s = Tables.schema(values(d))
    @test s.names == (:a, :b)
    @test s.types == (Int, Float64)
    @test d[:a] == [1,2,3]
    @test d[:b] == [4.,5.,6.]
  end

  @testset "unique" begin
    X = [i*j for i in 1:2, j in 1:1_000_000]
    z = rand(1_000_000)
    d = georef((z=[z;z],), [X X])
    u = uniquecoords(d)
    U = coordinates(u)
    @test nelms(u) == 1_000_000
    @test Set(eachcol(U)) == Set(eachcol(X))

    X = rand(3,100)
    z = rand(100)
    n = [string(i) for i in 1:100]
    Xd = hcat(X, X[:,1:10])
    zd = vcat(z, z[1:10])
    nd = vcat(n, n[1:10])
    sdata = georef(DataFrame(z=zd, n=nd), PointSet(Xd))
    ndata = uniquecoords(sdata)
    @test nelms(ndata) == 100
  end

  @testset "inside" begin
    # point set + rectangle
    𝒫 = PointSet([0. 2. 5. 7. 10.; 0. 3. 5. 6. 11.])
    𝒮 = georef((z=[1,2,3,4,5],), 𝒫)
    R1 = Rectangle((0.,0.), (5.,5.))
    R2 = Rectangle((5.,5.), (10.,10.))
    I = inside(𝒫, R1)
    @test coordinates(I) == [0. 2. 5.; 0. 3. 5.]
    I = inside(𝒫, R2)
    @test coordinates(I) == [5. 7.; 5. 6.]
    I = inside(𝒮, R1)
    I[:z] == [1,2,3]
    I = inside(𝒮, R2)
    I[:z] == [3,4]

    # regular grid + rectangle
    𝒢 = RegularGrid(3,3)
    𝒮 = georef((z=1:9,), 𝒢)
    R1 = Rectangle((0.,0.),(1.,1.))
    R2 = Rectangle((1.,1.),(2.,2.))
    R3 = Rectangle((0.,0.),(2.,2.))
    I = inside(𝒢, R1)
    @test I isa RegularGrid
    @test origin(I) == [0.,0.]
    @test spacing(I) == [1.,1.]
    @test size(I) == (2,2)
    I = inside(𝒢, R2)
    @test I isa RegularGrid
    @test origin(I) == [1.,1.]
    @test spacing(I) == [1.,1.]
    @test size(I) == (2,2)
    I = inside(𝒢, R3)
    @test I isa RegularGrid
    @test origin(I) == [0.,0.]
    @test spacing(I) == [1.,1.]
    @test size(I) == (3,3)
    for R in [R1,R2,R3]
      Ig = inside(𝒢, R)
      Is = inside(𝒮, R)
      Ds = domain(Is)
      @test Ds isa RegularGrid
      @test coordinates(Ds) == coordinates(Ig)
    end
  end

  @testset "split" begin
    d = RegularGrid(10,10)
    l, r = split(d, 0.5)
    @test nelms(l) == 50
    @test nelms(r) == 50
    l, r = split(d, 0.5, (1.,0.))
    @test nelms(l) == 50
    @test nelms(r) == 50
    cl = mean(coordinates(l), dims=2)
    cr = mean(coordinates(r), dims=2)
    @test cl[1] < cr[1]
    @test cl[2] == cr[2]
    l, r = split(d, 0.5, (0.,1.))
    @test nelms(l) == 50
    @test nelms(r) == 50
    cl = mean(coordinates(l), dims=2)
    cr = mean(coordinates(r), dims=2)
    @test cl[1] == cr[1]
    @test cl[2] < cr[2]
  end

  @testset "groupby" begin
    d = georef((z=[1,2,3],x=[4,5,6]), rand(2,3))
    g = groupby(d, :z)
    @test all(nelms.(g) .== 1)
    rows = [[1 4], [2 5], [3 6]]
    for i in 1:3
      @test Tables.matrix(values(g[i])) ∈ rows
    end
  end

  @testset "boundbox" begin
    d = RegularGrid((10,10), (1.,1.), (1.,1.))
    @test extrema(boundbox(d)) == ([1.,1.], [10.,10.])

    r = boundbox(RegularGrid(100,200))
    @test r == Rectangle((0.,0.), (99.,199.))

    r = boundbox(PointSet([0. 1. 2.; 0. 2. 1.]))
    @test r == Rectangle((0.,0.), (2.,2.))

    r = boundbox(PointSet([1. 2.; 2. 1.]))
    @test r == Rectangle((1.,1.), (2.,2.))
  end

  @testset "sample" begin
    d = georef((z=rand(10,10),))
    s = sample(d, 50)
    @test nelms(s) == 50
    s = sample(d, 50, rand([1,2], 100))
    @test nelms(s) == 50
  end

  @testset "filter" begin
    𝒟 = georef((a=[1,2,3], b=[1,1,missing]))
    𝒫 = filter(s -> !ismissing(s.b), 𝒟)
    @test 𝒫[:a] == [1,2]
    @test 𝒫[:b] == [1,1]

    𝒟 = georef((a=[1,2,3],b=[3,2,1]))
    𝒫ₐ = filter(s -> s.a > 1, 𝒟)
    𝒫ᵦ = filter(s -> s.b > 1, 𝒟)
    𝒫ₐᵦ = filter(s -> s.a > 1 && s.b > 1, 𝒟)
    @test nelms(𝒫ₐ) == 2
    @test nelms(𝒫ᵦ) == 2
    @test nelms(𝒫ₐᵦ) == 1
    @test 𝒫ₐ[:a] == [2,3]
    @test 𝒫ₐ[:b] == [2,1]
    @test 𝒫ᵦ[:a] == [1,2]
    @test 𝒫ᵦ[:b] == [3,2]
    @test 𝒫ₐᵦ[:a] == [2]
    @test 𝒫ₐᵦ[:b] == [2]
  end

  @testset "slice" begin
    d = RegularGrid((10,10), (1.,1.), (1.,1.))
    s = slice(d, 0.:5., 0.:5.)
    @test s isa RegularGrid
    @test extrema(coordinates(s)) == (1.,5.)
    s = slice(d, 5.:20.5, 6.:21.0)
    @test coordinates(s, 1) == [5.,6.]
    @test coordinates(s, nelms(s)) == [10.,10.]

    d = RegularGrid(10,10,10)
    s = slice(d, 5.5:10.0, 2.3:4.2, -1.2:2.0)
    @test s isa RegularGrid
    @test origin(s) == [6.,3.,0.]
    @test spacing(s) == [1.,1.,1.]
    @test size(s) == (4,1,2)

    d = PointSet([1. 5. 7.; 2. 3. 6.])
    s = slice(d, 2.:6., 2.:6.)
    @test nelms(s) == 1
    @test coordinates(s, 1) == [5.,3.]
  end
end


text:
using SnakeGames
using Test

@testset "SnakeGames.jl" begin
    # Write your tests here.
end


text:
"""
    Assigns layers using the longest path method.
    Arguments:
    adj_list        Directed graph in adjacency list format
"""

function layerAssmtLongestpathRec(adj_list, layers, j)
    # Look for all children of vertex i, try to bump their layer
    n = length(adj_list)
    for i in adj_list[j]
        if layers[i] == -1 || layers[i] <= layers[j] #if the layer of immediate precedor node == -1 or lower than current node, add 1 to precedor layer
            layers[i] = layers[j] + 1
            layerAssmtLongestpathRec(adj_list, layers, i) # recursive on precedor node
        end
    end
end

function layerAssmtLongestpath(adj_list)
    n = length(adj_list)
    layers = fill(-1, n) #assign each node to a layer. each with value -1

    for j in 1:n #
        out_deg = 0
        for i in 1:n
            if j in adj_list[i]
                out_deg += 1
            end
        end
        if out_deg == 0 # node without out-edge will be assigned to layer 1
            # Start recursive walk from this vertex
            layers[j] = 1
            layerAssmtLongestpathRec(adj_list, layers, j)
        end
    end

    return layers
end

"""
    Given a layer assignment, introduce dummy vertices to break up
    long edges (more than one layer)
    Arguments:
    orig_adj_list   Original directed graph in adjacency list format
    layers          Assignment of original vertices
"""
function layerAssmtDummy(orig_adj_list, layers)
    adj_list = deepcopy(orig_adj_list)

    # This is essentially
    # for i in 1:length(adj_list)
    # but the adj_list is growing in the loop
    i = 1
    while i <= length(adj_list)
        for (k,j) in enumerate(adj_list[i]) # k is the counter and j is the value of adj_list[i] at kth position
            if layers[j] - layers[i] > 1 # if the difference between the layer is larger than 1 => long edge
                # Need to add a dummy vertex
                new_v = length(adj_list) + 1
                adj_list[i][k] = new_v  # Replace the node of the long-edge with the dummy node
                push!(adj_list, Int[j])  # Make the node of the long-edge precedor of the dummy node
                push!(layers, layers[i]+1)  # Layer for the dummy node
            end
        end
        i += 1
    end

    return adj_list, layers
end

"""
    Given a layer assignment, decide a permutation for each layer
    that attempts to minimizes edge crossings using the barycenter
    method proposed in the Sugiyama paper.
    Arguments:
    adj_list        Directed graph in adjacency list format
    layers          Assignment of vertices
    layer_verts     Dictionary of layer => vertices
"""
function orderingBarycentric(adj_list, layers, layer_verts)
    num_layers = maximum(layers)
    n = length(adj_list)

    for iter in 1:5
        # DOWN
        for L in 1:num_layers-1
            # Calculate barycenter for every vertex in next layer
            cur_layer = layer_verts[L]
            next_layer = layer_verts[L+1]
            barys = zeros(n)
            out_deg = zeros(n)

            for (p,i) in enumerate(cur_layer)
                # Because of the dummy vertices we know that
                # all vertices in adj list for i are in next layer
                for (q,j) in enumerate(adj_list[i])
                    barys[j] = barys[j] .+ p
                    out_deg[j] = out_deg[j] .+ 1
                end
            end

            barys = barys ./ out_deg

            # Arrange next layer by barys, in ascending order
            next_layer_barys = [barys[j] for j in next_layer]

            layer_verts[L+1] = next_layer[sortperm(next_layer_barys)]
        end

        # UP
        for L in num_layers-1:-1:1
            # Calculate barycenters for every vertex in cur layer
            cur_layer = layer_verts[L]
            next_layer = layer_verts[L+1]
            barys = zeros(n)
            out_deg = zeros(n)
            for (p,i) in enumerate(cur_layer)
                # Because of the dummy vertices we know that
                # all vertices in adj list for i are in next layer
                # We need to know their positions in next layer
                # though, unfortunately. Probably a smarter way
                # to do this step
                for (q,j) in enumerate(adj_list[i])
                    # Find position in next layer
                    for (r,k) in enumerate(next_layer)
                        if k == j
                            barys[i] = barys[i] .+ r
                            out_deg[i] = out_deg[i] .+ 1
                            break
                        end
                    end
                end
            end
            barys = barys ./ out_deg
            # Arrange cur layer by barys, in ascending order
            cur_layer_barys = [barys[j] for j in cur_layer]
            #println("UP $L")
            #println(cur_layer)
            #println(cur_layer_barys)
            layer_verts[L] = cur_layer[sortperm(cur_layer_barys)]
        end
        # Do something with phase 2 here - don't really understand
    end

    return layer_verts
end

"""
    Given a layer assignment, decide a permutation for each layer
    that minimizes edge crossings using integer programming.
    Based on the IP described in
      M. Junger, E. K. Lee, P. Mutzel, and T. Odenthal.
      A polyhedral approach to the multi-layer crossing minimization problem.
      In G. Di Battista, editor, Graph Drawing: 5th International Symposium,
      GD ’97, volume 1353 of Lecture Notes in Computer Science, pages 13–24,
      Rome, Italy, September 1997. Springer-Verlag.
    Arguments:
    adj_list        Directed graph in adjacency list format
    layers          Assignment of vertices
    layer_verts     Dictionary of layer => vertices (initial perm.)
    Returns:
    new_layer_verts An improved dictionary of layer => vertices (opt. perm.)
"""
function orderingIp(adj_list, layers, layer_verts)
    num_layers = maximum(layers)

    m = Model(GLPK.Optimizer)

    # Define crossing binary variables
    @variable(m, c[L=1:num_layers,        # for each layer
                 i=layer_verts[L],      # for each vertex in this layer
                 j=adj_list[i],         # and vertex in the next layer
                 k=layer_verts[L],      # for each vertex in this layer
                 l=adj_list[k]], Bin)   # and vertex in the next layer

    # Objective: minimize crossings
    @objective(m, Min, sum(c))

    # Define permutation variables for each layer
    # We'll define for both (i,j) and (j,i), and ensure they consistency
    # by adding constraints. Presolve in the IP solver will simplify
    # the problem for us by removing one of the variables.
    @variable(m, x[L=1:num_layers, i=layer_verts[L], j=layer_verts[L]], Bin)
    for L in 1:num_layers
        for i in layer_verts[L]
            for j in layer_verts[L]
                j <= i && continue  # Don't double-add
                # Ensure x[i,j] and x[j,i] are consistent
                @constraint(m, x[L,i,j] == 1 - x[L,j,i])
                # And ensure that triples are consistent
                for k in layer_verts[L]
                    k <= j && continue
                    @constraint(m, 0 <= x[L,i,j] + x[L,j,k] - x[L,i,k] <= 1)
                end
            end
        end
    end

    # Link permutations to crossings
    for L in 1:num_layers-1
        # For all (i,j)
        for i in layer_verts[L]
        for j in adj_list[i]
            # For all (k,l)
            for k in layer_verts[L]
            k == i && continue  # Can't cross if starting from same vertex!
            for l in adj_list[k]
                @constraint(m, -c[L,i,j,k,l] <= x[L+1,j,l] - x[L,i,k])
                @constraint(m,  c[L,i,j,k,l] >= x[L+1,j,l] - x[L,i,k])
            end
            end
        end
        end
    end

    # Solve the IP
    optimize!(m)

    # Extract permutation from solution
    x_sol = getvalue(x)
    new_layer_verts = [L => Int[] for L in 1:num_layers]
    for L in 1:num_layers
        old_perm = layer_verts[L]
        # For each vertex, count the number of times it is "in front"
        # of the other vertices. The higher this number, the earlier
        # the vertex appears in the layer.
        scores = zeros(length(old_perm))
        for (p,i) in enumerate(old_perm)
            for j in old_perm
                i == j && continue
                if round(Integer, x_sol[L,i,j]) == 1
                    # i appears before j
                    scores[p] += 1
                end
            end
        end
        new_layer_verts[L] = old_perm[sortperm(scores,rev=true)]
    end

    return new_layer_verts
end

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

"""
    Given a layer assignment and permutation, decide the coordinates for
    each vertex. The objective is to encourage straight edges, especially
    for longer edges. This function uses an integer program to decide the
    coordinates (although it is solved as a linear program), as described in
      Gansner, Emden R., et al.
      A technique for drawing directed graphs.
      Software Engineering, IEEE Transactions on 19.3 (1993): 214-230.
    Arguments:
    adj_list        Directed graph in adjacency list format
    layers          Assignment of vertices
    layer_verts     Dictionary of layer => vertices (final perm.)
    orig_n          Number of original (non-dummy) vertices
    widths          Width of each vertex
    xsep            Minimum seperation between each vertex
    Returns:
    layer_coords    For each layer and vertex, the x-coord
"""
function coordIp(adj_list, layers, layer_verts, orig_n, widths, xsep)
    num_layers = maximum(layers)

    m = Model(GLPK.Optimizer)

    # One variable for each vertex
    @variable(m, x[L=1:num_layers, i=layer_verts[L]] >= 0)

    # Constraint: must respect permutation, and spacing constraint
    for L in 1:num_layers
        for i in 1:length(layer_verts[L])-1
            a = layer_verts[L][i]
            b = layer_verts[L][i+1]
            @constraint(m, x[L,b] - x[L,a] >=
                (widths[a] + widths[b])/2 + xsep) #the distance between two
                #adjacent node must be at least xsep
        end
    end

    # Objective: minimize total misalignment
    # Use the weights from the Ganser paper:
    #   1 if both nodes "real"
    #   2 if one of the nodes is "real"
    #   8 if neither node is "real"
    # We use absolute distance in the objective so we'll need
    # auxilary variables for each pair of edges
    obj = AffExpr()

    @variable(m, absdiff[L=1:num_layers-1,
                        i=layer_verts[L], j=adj_list[i]] >= 0)

    for L in 1:num_layers-1
        for i in layer_verts[L]
            for j in adj_list[i]
                #add 2 constraints to deal with absolute value in objective
                @constraint(m, absdiff[L,i,j] >= x[L,i] - x[L+1,j])
                @constraint(m, absdiff[L,i,j] >= x[L+1,j] - x[L,i])
                if i > orig_n && j > orig_n
                    # Both dummy vertices
                    obj += 8*absdiff[L,i,j]
                elseif (i <= orig_n && j >  orig_n) ||
                       (i >  orig_n && j <= orig_n)
                    # Only one dummy vertix
                    obj += 2*absdiff[L,i,j]
                else
                    # Both real
                    obj += absdiff[L,i,j]
                end
            end
        end
    end
    @objective(m, Min, obj)

    # Solve it...
    optimize!(m)

    # ... and mangle the solution into shape
    x_sol = getvalue.(x)
    locs_x = zeros(length(layers))
    for L in 1:num_layers
        for i in layer_verts[L]
            locs_x[i] = x_sol[L,i]
        end
    end
    return locs_x
end

########################################################################
function layoutTree(adj_list,
    labels::Vector;
    filename    = "",

    cycles      = true,
    ordering    = :optimal,
    coord       = :optimal,
    xsep        = 50,
    x_start       = 100,
    ysep        = 100,
    scale       = 0.05,
    labelpad    = 1.2,
    background  = "white")
# Calculate the original number of vertices
n = length(adj_list)

# 1     Cycle removal
if cycles
# Need to remove cycles first
error("Cycle removal not implemented!")
end

# 2     Layering
# 2.1   Assign a layer to each vertex
layers = layerAssmtLongestpath(adj_list)
num_layers = maximum(layers)
# 2.2  Create dummy vertices for long edges
adj_list, layers = layerAssmtDummy(adj_list, layers)
orig_n, n = n, length(adj_list)

# 3     Vertex ordering [to reduce crossings]
# 3.1   Build initial permutation vectors
layer_verts = Dict([L => Int[] for L in 1:num_layers])

for i in 1:n
push!(layer_verts[layers[i]], i)
end

# 3.2  Reorder permutations to reduce crossings
if ordering == :barycentric
    layer_verts = orderingBarycentric(adj_list, layers, layer_verts)
elseif ordering == :optimal
    layer_verts = orderingIp(adj_list, layers, layer_verts)
end

# 4     Vertex coordinates [to straighten edges]
# 4.1   Place y coordinates in layers
locs_y = zeros(n)

# Luxor origin (0,0) is the upper-left corner. Lower layer has higher y location.
rev_y = sort(1:num_layers, rev = true) #reverse the layer
#ysep = 100

for L in 1:num_layers
    for (x,v) in enumerate(layer_verts[L])
        locs_y[v] = (rev_y[L])*ysep
    end
end

# 4.2   Get widths of each label, if there are any
widths  = ones(n)
widths[orig_n+1:n]  .= 0 #dummy node has 0 width
heights = ones(n)
heights[orig_n+1:n] .= 0 #dummy node has 0 height

# Note that we will convert these sizes into "absolute" units
# and then work in these same units throughout. The font size used
# here is just arbitrary, and unchanging. This hack arises because it
# is meaningless to ask for the size of the font in "relative" units
# but we don't want to collapse to absolute units until the end.

if length(labels) == orig_n
    extents = textextents.(labels)
    for (i,(width,height)) in enumerate(extents)
        widths[i]  = reduce(hcat, extents)[3,i]
        heights[i] = reduce(hcat, extents)[4,i]
    end
end

locs_x = coordIp(adj_list, layers, layer_verts, orig_n, widths, xsep) .+ x_start #= add x_start to shift all x location
to the right by 100 =#

# 4.3   Summarize vertex info
max_x, max_y = maximum(locs_x), maximum(locs_y)
max_w, max_h = maximum(widths), maximum(heights)

## find x_pad to the right and add the pad to the border of the most right node
x_pad = minimum(locs_x - widths/2)
w = ceil.(max_x .+ widths[locs_x .== max_x]/2 .+ x_pad)[1] # width of the drawing is the max_x + width of the node + x_pad
h = ceil(max_y + max_h/2 + ysep) # height of the drawing is the max_y + max_height + y_pad

# 5     Draw the tree

#=
@draw begin
    Luxor.fontsize(20)
    fontface("Noto-Sans")

    #widths = reduce(hcat, textextents.(labels))[3,:]
    widths = ceil.(widths .* 0.6)
    #push!(widths, 0)

    #heights = reduce(hcat, textextents.(labels))[4,:]
    #push!(heights, 0)
    heights = repeat([max_h],n)
    heights = ceil.(heights)

    for i in 1:orig_n
        #origin()
        point = Luxor.Point(locs_x[i], locs_y[i])

        squircle(point, widths[i], heights[i], rt = 0.3, :stroke)
        #textcentered(labels[i], point)
        Luxor.text(labels[i], point, halign=:center, valign=:middle)
    end

    for L in 1:num_layers, i in layer_verts[L], j in adj_list[i]
        if i <= orig_n #if the destination node is a dummy then dont add arrow
            Luxor.arrow(Point(locs_x[j], locs_y[j] + heights[j]), Point(locs_x[i], locs_y[i] - heights[i]), linewidth=2)
        else
            Luxor.line(Point(locs_x[j], locs_y[j] + heights[j]), Point(locs_x[i], locs_y[i] + heights[i]), :stroke)
        end
    end

    end
    =#

Drawing(w, h, "test.svg")

Luxor.background("$background")
#sethue("red")
Luxor.fontsize(20)
fontface("Noto-Sans")

for i in 1:orig_n

    point = Luxor.Point(locs_x[i], locs_y[i])

    squircle(point, widths[i]/1.5, ceil(max_h), rt = 0.3, :stroke)

    Luxor.text(labels[i], point, halign=:center, valign=:middle)
end

for L in 1:num_layers, i in layer_verts[L], j in adj_list[i]
    if i <= orig_n #if the destination node is a dummy then dont add arrow
        Luxor.arrow(Point(locs_x[j], locs_y[j] + ceil(max_h)), Point(locs_x[i], locs_y[i] - ceil(max_h)), linewidth=2)
    else
        Luxor.line(Point(locs_x[j], locs_y[j] + ceil(max_h)), Point(locs_x[i], locs_y[i] + ceil(max_h)), :stroke)
    end
end

finish()
preview()

end


text:
using CellSegmenter
using Test

@testset "CellSegmenter.jl" begin
    # Write your tests here.
end

max_stars_repo_name: neuro-myoung/CellSegmenter

max_issues_repo_name: neuro-myoung/CellSegmenter

max_forks_repo_name: neuro-myoung/CellSegmenter


text:
argtypes(x::Combined{T, A}) where {T, A} = A
argtypes(x) = Any

function data_limits(x)
    error("No datalimits for $(typeof(x)) and $(argtypes(x))")
end

function data_limits(x::Atomic)
    isempty(x.plots) ? atomic_limits(x) : data_limits(x.plots)
end

"""
Data limits calculate a minimal boundingbox from the data points in a plot.
This doesn't include any transformations, markers etc.
"""
function atomic_limits(x::Atomic{<: Tuple{Arg1}}) where Arg1
    return xyz_boundingbox(transform_func(x), to_value(x[1]))
end

function atomic_limits(x::Atomic{<: Tuple{X, Y, Z}}) where {X, Y, Z}
    return xyz_boundingbox(transform_func(x), to_value.(x[1:3])...)
end

function atomic_limits(x::Atomic{<: Tuple{X, Y}}) where {X, Y}
    return xyz_boundingbox(transform_func(x), to_value.(x[1:2])...)
end

_isfinite(x) = isfinite(x)
_isfinite(x::VecTypes) = all(isfinite, x)
scalarmax(x::Union{Tuple, AbstractArray}, y::Union{Tuple, AbstractArray}) = max.(x, y)
scalarmax(x, y) = max(x, y)
scalarmin(x::Union{Tuple, AbstractArray}, y::Union{Tuple, AbstractArray}) = min.(x, y)
scalarmin(x, y) = min(x, y)

extrema_nan(itr::Pair) = (itr[1], itr[2])
extrema_nan(itr::ClosedInterval) = (minimum(itr), maximum(itr))

function extrema_nan(itr)
    vs = iterate(itr)
    vs === nothing && return (NaN, NaN)
    v, s = vs
    vmin = vmax = v
    # find first finite value
    while vs !== nothing && !_isfinite(v)
        v, s = vs
        vmin = vmax = v
        vs = iterate(itr, s)
    end
    while vs !== nothing
        x, s = vs
        vs = iterate(itr, s)
        _isfinite(x) || continue
        vmax = scalarmax(x, vmax)
        vmin = scalarmin(x, vmin)
    end
    return (vmin, vmax)
end

function xyz_boundingbox(transform_func, mesh::GeometryBasics.Mesh)
    xyz_boundingbox(transform_func, coordinates(mesh))
end

function xyz_boundingbox(transform_func, xyz)
    isempty(xyz) && return FRect3D()
    mini, maxi = extrema_nan((apply_transform(transform_func, point) for point in xyz))
    w = maxi .- mini
    return FRect3D(to_ndim(Vec3f0, mini, 0), to_ndim(Vec3f0, w, 0))
end

const NumOrArray = Union{AbstractArray, Number}

function xyz_boundingbox(transform_func, x::AbstractVector, y::AbstractVector, z::NumOrArray=0)
    # use lazy variant of broadcast!
    points = Base.broadcasted(Point3, x, y', z)
    return xyz_boundingbox(transform_func, points)
end

function xyz_boundingbox(transform_func, x::NumOrArray, y::NumOrArray, z::NumOrArray = 0)
    # use lazy variant of broadcast!
    points = Base.broadcasted(Point3, x, y, z)
    return xyz_boundingbox(transform_func, points)
end

function xyz_boundingbox(transform_func, x, y, z = 0)
    isempty(x) && return FRect3D()
    minmax = extrema_nan.(apply_transform.((transform_func,), (x, y, z)))
    mini, maxi = Vec(first.(minmax)), Vec(last.(minmax))
    w = maxi .- mini
    return FRect3D(to_ndim(Vec3f0, mini, 0), to_ndim(Vec3f0, w, 0))
end

const ImageLike{Arg} = Union{Heatmap{Arg}, Image{Arg}}
function data_limits(x::ImageLike{<: Tuple{X, Y, Z}}) where {X, Y, Z}
    xyz_boundingbox(transform_func(x), to_value.((x[1], x[2]))...)
end

function data_limits(x::Volume)
    xyz_boundingbox(transform_func(x), to_value.((x[1], x[2], x[3]))...)
end

function text_limits(x::VecTypes)
    p = to_ndim(Vec3f0, x, 0.0)
    return FRect3D(p, p)
end

function text_limits(x::AbstractVector)
    return FRect3D(x)
end

function atomic_limits(x::Text{<: Tuple{Arg1}}) where Arg1
    return boundingbox(x)
end

function data_limits(x::Annotations)
    # data limits is supposed to not include any transformation.
    # for the annotation, we use the model matrix directly, so we need to
    # to inverse that transformation for the correct limits
    bb = data_limits(x.plots[1])
    return inv(modelmatrix(x)) * bb
end

Base.isfinite(x::Rect) = all(isfinite.(minimum(x))) &&  all(isfinite.(maximum(x)))

function data_limits(plots::Vector)
    isempty(plots) && return
    bb = FRect3D()
    plot_idx = iterate(plots)
    while plot_idx !== nothing
        plot, idx = plot_idx
        plot_idx = iterate(plots, idx)
        # axis shouldn't be part of the data limit
        isaxis(plot) && continue
        isa(plot, Legend) && continue
        bb2 = data_limits(plot)::FRect3D
        isfinite(bb) || (bb = bb2)
        isfinite(bb2) || continue
        bb = union(bb, bb2)
    end
    bb
end

data_limits(s::Scene) = data_limits(plots_from_camera(s))
data_limits(plot::Combined) = data_limits(plot.plots)


text:
# store current variables
const stored_init_rho = [0.5, 0.5]
const stored_init_lambda = [[0.5,], [0.5,]]

# univariate LBFGSB optimizer
const optimizer_lbfgs = L_BFGS_B(1, 17)
const stored_bounds = zeros(3, 1)
stored_bounds[1,1] = 1  # represents the type of bounds imposed on the variables:
                 #  0->unbounded, 1->only lower bound, 2-> both lower and upper bounds, 3->only upper bound
stored_bounds[2,1] = 0.0
stored_bounds[3,1] = Inf

# LBFGSB optimizer for marginal_projection_with_constraint
const optimizer_lbfgs_cs = L_BFGS_B(1024, 17)

function solve_p_given_abk(a::AbstractVector, b, k::Integer, ids::AbstractVector{<:Integer})

    n = length(a)
    T = eltype(a)

    raw = a .- b

    sorted_raw = raw[ids]
    neg_cut_id = findfirst(x -> x < 0, sorted_raw)      # first negative
    if isnothing(neg_cut_id)
        neg_cut_id = n + 1
        sraw_neg = T[]
        n_neg = 0
    else
        sraw_neg = @views sorted_raw[neg_cut_id:end]
        n_neg = length(sraw_neg)
    end

    sum_p = sum(@views sorted_raw[1:neg_cut_id-1])
    if  sorted_raw[1] <= sum_p / k
        return max.(0, raw)
    end

    # starting
    i = 1                   # #x in sum_p_else
    j = neg_cut_id - 2      # #x in sum_p_edge
    l = 0                   # #x in sum_p_neg

    sum_p_edge = sorted_raw[1]
    sum_p_else = sum_p - sorted_raw[1]
    sum_p_neg = zero(T)

    # solution
    sol_r = zero(T)
    p_add = zero(T)

    while true

        next_r = (sum_p_else + sum_p_neg + sum_p_edge * (j + l) / (k - i)) / (k - i + (i * (j + l)) / (k - i)) |> T
        p_else_add = (sum_p_edge - i * next_r) / (k - i) |> T

        if l < n_neg && p_else_add + sraw_neg[l + 1] > 0
            sum_p_neg += sraw_neg[l + 1]
            l += 1
            continue
        end

        if  i == k-1  || next_r + 1e-6 >= sorted_raw[i + 1] + p_else_add     # 1e-6 for floating point error
            sol_r = next_r
            p_add = p_else_add

            break
        end

        sum_p_edge += sorted_raw[i + 1]
        sum_p_else -= sorted_raw[i + 1]
        i += 1
        j -= 1

    end

    res = min.(max.(raw .+ p_add, 0), sol_r)
    return res

end

function solve_p_given_abk(a::AbstractVector, b, k::Integer)
    raw = a .- b
    ids = sortperm(raw, rev = true)

    return solve_p_given_abk(a, b, k, ids)
end

function obj_rho(rho_arr::AbstractVector, A::AbstractMatrix, IDS::AbstractMatrix)
    n = size(A, 1)
    T = eltype(rho_arr)
    rho = rho_arr[1]

    # to store best P
    best_P = zeros(T, n, n)

    obj = zero(T)
    grad = zero(T)

    for k = 1:n
        a = A[:, k]
        b = rho / k
        ids = IDS[:, k]

        opt_p = solve_p_given_abk(a, b, k, ids)
        best_P[:, k] = opt_p

        obj += sum(opt_p) * rho / k + sum((opt_p - a) .^ 2) / 2
        grad += sum(opt_p) / k
    end

    obj -= rho
    grad -= 1

    return -obj, [-grad], best_P
end

function marginal_projection(A::AbstractMatrix; init_storage_id = 1)

    n = size(A, 1)
    T = eltype(A)

    IDS = zeros(Int, n, n)
    for i = 1:n
        IDS[:, i] = sortperm(A[:,i], rev = true)
    end

    init_rho = stored_init_rho[init_storage_id]

    _, xout = optimizer_lbfgs(rho -> obj_rho(rho, A, IDS)[1:2], [init_rho], stored_bounds, maxiter=20, factr=1e8, pgtol=1e-3)

    opt_P = obj_rho(xout, A, IDS)[3]
    stored_init_rho[init_storage_id] = xout[1]

    opt_P = min.(max.(opt_P, zero(T)), one(T))

    return opt_P
end

## marginal projection with constraints

function obj_rho_lambda(rho_lambda::AbstractVector, A::AbstractMatrix,
    B_list::AbstractVector{<:AbstractMatrix}, c_list::AbstractVector{<:Number}, tau_list::AbstractVector{<:Number})

    n = size(A, 1)
    ncs = length(B_list)
    T = eltype(rho_lambda)

    rho = rho_lambda[1]
    lda = rho_lambda[2:end]

    B_wsum = zeros(n, n)
    for j = 1:ncs
        B_wsum += lda[j] * B_list[j]
    end

    # to store best P
    best_P = zeros(T, n, n)

    obj = zero(T)
    grad = zeros(T, length(rho_lambda))

    for k = 1:n
        a = A[:, k]
        b = (rho / k) .- B_wsum[:, k]

        opt_p = solve_p_given_abk(a, b, k)
        best_P[:, k] = opt_p

        obj += sum(opt_p .* b) + sum((opt_p - a) .^ 2) / 2
        grad[1] += sum(opt_p) / k
    end

    obj -= rho
    for j = 1:ncs
        obj += lda[j] * (tau_list[j] - c_list[j])
    end

    grad[1] -= 1
    for j = 1:ncs
        grad[j+1] += tau_list[j] - c_list[j] - sum(B_list[j] .* best_P)
    end

    return -obj, -grad, best_P
end

function marginal_projection_with_constraint(A::AbstractMatrix, B_list::AbstractVector{<:AbstractMatrix},
    c_list::AbstractVector{<:Number}, tau_list::AbstractVector{<:Number}, alg = :lbfgs; args...)

    if alg == :lbfgs
        return lbfgs_projection_with_constraint(A, B_list, c_list, tau_list; args...)
    elseif alg == :dykstra
        return dykstra_projection_with_constraint(A, B_list, c_list, tau_list; args...)
    else
        return lbfgs_projection_with_constraint(A, B_list, c_list, tau_list; args...)
    end
end

function lbfgs_projection_with_constraint(A::AbstractMatrix, B_list::AbstractVector{<:AbstractMatrix},
    c_list::AbstractVector{<:Number}, tau_list::AbstractVector{<:Number}; init_storage_id = 1)

    n = size(A, 1)
    T = eltype(A)
    ncs = length(B_list)

    # pull stored parameters
    init_rho = stored_init_rho[init_storage_id]
    if length(stored_init_lambda[init_storage_id]) != ncs
        stored_init_lambda[init_storage_id] = ones(ncs) * 0.5
    end
    init_lambda = stored_init_lambda[init_storage_id]
    init_pars = [init_rho; init_lambda]

    # bounds
    bounds = zeros(3, ncs + 1)
    bounds[1,:] .= 1  # represents the type of bounds imposed on the variables:
                      #  0->unbounded, 1->only lower bound, 2-> both lower and upper bounds, 3->only upper bound
    bounds[2,:] .= 0.0
    bounds[3,:] .= Inf

    _, xout = optimizer_lbfgs_cs(rho_lambda -> obj_rho_lambda(rho_lambda, A, B_list, c_list, tau_list)[1:2],
                    init_pars, bounds, maxiter=20, factr=1e8, pgtol=1e-3)

    opt_P = obj_rho_lambda(xout, A, B_list, c_list, tau_list)[3]
    opt_P = min.(max.(opt_P, zero(T)), one(T))

    # store pars
    stored_init_rho[init_storage_id] = xout[1]
    stored_init_lambda[init_storage_id] = xout[2:end]

    return opt_P
end

function exact_with_constraint(A::AbstractMatrix, B::AbstractMatrix,
    c::Number, tau::Number)

    needed = tau - c - sum(A .* B)
    if needed > 0
        z = needed / sum(B .^ 2)
        return z * B + A
    else
        return A
    end
end

function dykstra_projection_with_constraint(A::AbstractMatrix, B_list::AbstractVector{<:AbstractMatrix},
    c_list::AbstractVector{<:Number}, tau_list::AbstractVector{<:Number}; args...)

    n = size(A, 1)
    T = eltype(A)
    ncs = length(B_list)

    Zs = [zeros(T, n, n) for _ = 1:ncs+1]
    P = copy(A)

    niter = 10

    for iter = 1:niter
        for ic = 1:ncs+1
            if ic == ncs+1
                P_next = marginal_projection(P + Zs[ic]; args...)
            else
                P_next = exact_with_constraint(P + Zs[ic], B_list[ic], c_list[ic], tau_list[ic])
            end
            Zs[ic] = P + Zs[ic] - P_next
            P = P_next
        end
    end

    return P
end

## max sum k largest, non neg
# compute value
function sumlargest(a::AbstractVector, k::Integer)
    sa = sort(a, rev=true)
    return sum(sa[1:k])
end

function sumlargest(A::AbstractMatrix)
    n = size(A, 1)
    sl = zeros(n)
    for i = 1:n
        sl[i] = sumlargest(A[:,i], i)
    end
    return sl
end

function max_sumlargest(A::AbstractMatrix, non_neg = true)
    sl = sumlargest(A)
    m = maximum(sl)
    if non_neg
        m = max(0, m)
    end
    return m
end

## Proximal functions
function prox_max_sumlargest(A::AbstractMatrix; args...)
    return A - marginal_projection(A; args...)
end

function prox_max_sumlargest(A::AbstractMatrix, rho::Real; args...)
    return A - marginal_projection(A .* rho; args...) ./ rho
end

function prox_max_sumlargest_with_constraint(A::AbstractMatrix, B_list::AbstractVector{<:AbstractMatrix},
    c_list::AbstractVector{<:Number}, tau_list::AbstractVector{<:Number}; args...)

    return A - marginal_projection_with_constraint(A, B_list, c_list, tau_list; args...)
end

function prox_max_sumlargest_with_constraint(A::AbstractMatrix, rho::Real, B_list::AbstractVector{<:AbstractMatrix},
    c_list::AbstractVector{<:Number}, tau_list::AbstractVector{<:Number}; args...)

    return A - marginal_projection_with_constraint(A .* rho, B_list, c_list, tau_list; args...) ./ rho
end

# Best P: min_P <B, P>,   rev =true for max_P
function find_opt_p(B, rev = false)
    n = size(B, 1)
    T = eltype(B)

    SB = zeros(Int, n, n)
    for i = 1:n
        SB[:, i] = sortperm(B[:,i], rev = rev)
    end

    best_i = 0
    best_sum = zero(T)
    for i = 1:n
        sum_i = sum(B[SB[1:i, i], i])
        if rev && sum_i > best_sum
            best_i = i
            best_sum = sum_i
        elseif !rev && sum_i < best_sum
            best_i = i
            best_sum = sum_i
        end
    end

    P = zeros(T, n, n)
    if best_i > 0
        P[SB[1:best_i, best_i], best_i] .= one(T)
    end

    return P, best_sum
end

## reset storage
function reset_projection_storage()
    stored_init_rho[1] = 0.5
    stored_init_rho[2] = 0.5
    stored_init_lambda[1] = [0.5,]
    stored_init_lambda[2] = [0.5,]
end