text:
#using Distributed
using ClearStacktrace

using Plots

using SpikeNetOpt
#@show(varinfo(SNO))
using JLD
using SpikingNeuralNetworks
using Evolutionary
using SpikingNN
using Metaheuristics
using Test

using Revise

SNN = SpikingNeuralNetworks
SNN.@load_units
SNO = SpikeNetOpt

unicodeplots()
##
##

const MU = 10

##
# Ground truths
##
const weight_gain_factor = 0.77
#const GT = 26

#g, Cg = SpikeNetOpt.make_net_from_graph_structure(GT)

spkd_ground_dic,weights1,synpases1 = SNO.sim_net_darsnack_used(weight_gain_factor)#(Ne, Ni, σee = 0.5, pee = 0.8, σei = 0.5, pei = 0.8)
spkd_ground = SNO.get_trains_dars(spkd_ground_dic)
#@show(weights1)
@show(spkd_ground)
error1 = SNO.spike_train_difference(spkd_ground, spkd_ground)

@test error1==0
@show(error1)
#println(error==0)
#weight_gain_factor1 = 0.77
nspkd_ground_dic,weights2,synpases2 = SNO.sim_net_darsnack_used(weight_gain_factor)#(Ne, Ni, σee = 0.5, pee = 0.8, σei = 0.5, pei = 0.8)
nspkd_ground = SNO.get_trains_dars(nspkd_ground_dic)

nerror = SNO.spike_train_difference(nspkd_ground, nspkd_ground)
@test nerror==0
@show(spkd_ground)
@show(nspkd_ground)
nerror2 = SNO.spike_train_difference(spkd_ground, nspkd_ground)

@show(nerror2)
#@test weights1==weights2
#@test all(isapprox.(weights1, weights2, rtol = 0.0))
#@test spkd_ground==nspkd_ground
#@test all(isapprox.(spkd_ground, nspkd_ground, rtol = 0.0))

#@test nerror2==0
#@show(weights2)

#println(nerror)
#@show(nerror)

#@test weights1===weights2

META_HEUR_OPT = true
EVOLUTIONARY_OPT = false
function loss(model)

    weight_gain_factor = model[1]
    println("\n genes \n")
    @show(model)
    println(model[1])
    println("\n")

    train_dic,weights,synpases = SNO.sim_net_darsnack_used(weight_gain_factor)
    println("\n ground truth: \n")
    rasterplot(spkd_ground_dic)|>display#

    spkd_found = SNO.get_trains_dars(train_dic)

    l0 = length(spkd_ground)
    l1 = length(spkd_found)
    @show(spkd_ground)
    s0 = sum(spkd_ground[1][:,:])
    if !isempty(spkd_found[1][:,:])
        s1 = sum(spkd_found[1][:,:])
    else
        #@show(spkd_found)

        s1 = 0

    end
    temp = abs(l0-l1)+abs(s0-s1)
    #if temp == 0
    error = SNO.spike_train_difference(spkd_ground, spkd_found)
    error = sum(error)+temp

    #else
    #    error = temp
    #end
    @show(error)
    error
end

function eval_best(params)
    spkd_found,weights,synpases = SNO.sim_net_darsnack_used(params[1])#Ne, Ni, σee = σee, pee = pee, σei = σei, pei = pei)#,a=a)

    #title!("Raster Plot")
    #xlabel!("Time (sec)")
    spkd_found2 = SNO.get_trains_dars(spkd_found)
    @show(spkd_found2)
    #display(Plots.historam(spkd_found2))
    #spkd_found
    #rasterplot(spkd_found2)#, label = ["Input 1"])#, "Input 2"])
end

if EVOLUTIONARY_OPT

    lower = Float32[0.0]#0.0 0.0 0.0]# 0.03 4.0]
    upper = Float32[3.0]# 1.0 1.0 1.0]# 0.2 20.0]
    lower = vec(lower)
    upper = vec(upper)

    options = GA(
        populationSize = MU,
        ɛ = 4,
        mutationRate = 0.5,
        selection = ranklinear(1.5),#ranklinear(1.5),#ss,
        crossover = intermediate(0.5),#xovr,
        mutation = uniform(0.5),#(.015),#domainrange(fill(1.0,ts)),#ms
    )
    #Random.seed!(0);
    D = 1
    bounds = [0.0ones(D) 3.0ones(D)]'
    result = Evolutionary.optimize(
        loss,
        lower,
        upper,
        vec([1.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0 0.0]),
        options,
        Evolutionary.Options(
            iterations = 50,
            successive_f_tol = 1,
            show_trace = true,
            store_trace = true,
        ),
    )
    fitness = minimum(result)

    filename = string("GAsolution.jld")#, py"target_num_spikes")#,py"specimen_id)
    params = result.minimizer

    spkd_found = eval_best(params)
    save(
        filename,
        "spkd_ground",
        spkd_ground,
        "spkd_found",
        spkd_found,
        "Ne",
        Ne,
        "Ni",
        Ni,
        "sim_length",
        sim_length,
    )
    println("best result")
    loss(result.minimizer)
    #println("σee = 0.5,  pee= 0.8,σei = 0.5,  pei= 0.8")

    #println("σee = 0.45,  pee= 0.8,σei = 0.4,  pei= 0.9)")
    @show(result.minimizer)

    @show(fitness)

    @show(result)
    @show(result.trace)
    trace = result.trace
    #SNO.dir(trace[1, 1, 1])
    trace[1, 1, 1].metadata#["population"]
    spkd_found = eval_best(params)
    evo_loss = [t.value for t in trace[2:length(trace)]]
    display(plot(evo_loss))
    filename = string("PopulationScatter.jld")#, py"target_num_spikes")#,py"specimen_id)
    save(filename, "trace", trace)

end
dir(x) = fieldnames(typeof(x))

#if META_HEUR_OPT
logger(st) = begin
    A = st.population
    temp = [a.x for a in A]
    bs = st.best_sol
    scatter(temp, label="opt", title="Gen: $(st.iteration)") |> display
    plot!([bs.x], label="Pareto Front", lw=1) |> display
end
#function main()
       # Transliterated fortran main subroutine

D = 1

#bounds = [0.0ones(D);# lower bounds
#        1.0ones(D)]# upper bounds
bounds = Array([0.0ones(D) 1.0ones(D)]')
#pareto_set = [ generateChild(loss(bounds[i])) for i in 1:length(bounds) ]

#information = Information(f_optimum = 0.0)
options = Options( seed = 1, iterations=350, f_calls_limit =350)
#optimize(f, bounds, NSGA2())
information = Information(f_optimum =0.0)

algorithm = ECA(information = information, options = options)

result = optimize(loss, bounds, algorithm, logger=logger)#, algorithm)#, logger=logger)#, SMS_EMOA(N = 15, n_samples=1, options=options))
@show(result)
#end

#end

#main()

@show minimum(result)
@show minimizer(result)

f_calls, best_f_value = convergence(result)

animation = @animate for i in 1:length(result.convergence)
    l = @layout [a b]
    p = plot( layout=l)

    X = positions(result.convergence[i])
    scatter!(p[1], X[:,1], X[:,2], label="", xlim=(-5, 5), ylim=(-5,5), title="Population")
    x = minimizer(result.convergence[i])
    scatter!(p[1], x[1:1], x[2:2], label="")

    # convergence
    plot!(p[2], xlabel="Generation", ylabel="fitness", title="Gen: $i")
    plot!(p[2], 1:length(best_f_value), best_f_value, label=false)
    plot!(p[2], 1:i, best_f_value[1:i], lw=3, label=false)
    x = minimizer(result.convergence[i])
    scatter!(p[2], [i], [minimum(result.convergence[i])], label=false)
end

# save in different formats
# gif(animation, "anim-convergence.gif", fps=30)
mp4(animation, "anim-convergence.mp4", fps=30)
    #=
    result = optimize(loss, bounds, NSGA2(options=options))
    #result = optimize(SpikeNetOpt.loss, bounds)#, method)
    @show(result)

    methods = [
            SMS_EMOA(N = 5, n_samples=5, options=options),
            NSGA2(options=options)
            ]
            #MOEAD_DE(gen_ref_dirs(nobjectives, npartitions), options=Options( seed = 1, iterations = 500)),
            #NSGA3(options=options),

    for method in methods
        f_calls = 0
        result = optimize(loss, bounds, method)
        #result = optimize(SpikeNetOpt.loss, bounds)#, method)
        @show(result)

    end
    =#
#end

# ## Visualize the result

#contour(0:.1:1, 0:.1:1, (x, y) -> model([x,y])[].data, fill=true)


text:
# Implement Interval Tree using DataStructures's SortedDict
struct IntervalTreeOrder <: DataStructures.Ordering end
DataStructures.lt(::IntervalTreeOrder, a::UnitRange{Int}, b::UnitRange{Int}) =
    last(a) < first(b)
DataStructures.eq(::IntervalTreeOrder, a::UnitRange{Int}, b::UnitRange{Int}) =
    isequal(a, b) || first(a) in b || first(b) in a

function is_str_or_cmd(t::Tokens.Kind)
    t === Tokens.CMD && return true
    t === Tokens.TRIPLE_CMD && return true
    t === Tokens.STRING && return true
    t === Tokens.TRIPLE_STRING && return true
    return false
end

function is_str_or_cmd(cst::CSTParser.EXPR)
    CSTParser.isstring(cst) || CSTParser.iscmd(cst)
end

function tokenize(val::AbstractString)::Tokens.Kind
    toks = collect(Tokenize.tokenize(val))
    toks[1].kind
end

struct Document
    text::AbstractString

    range_to_line::SortedDict{UnitRange{Int},Int,IntervalTreeOrder}
    line_to_range::Dict{Int,UnitRange{Int}}

    # mapping the offset in the file to the raw literal
    # string and what lines it starts and ends at.
    lit_strings::Dict{Int,Tuple{Int,Int,String}}
    comments::Dict{Int,Tuple{Int,String}}

    # CSTParser does not detect semicolons.
    # It's useful to know where these are for a few node types.
    # The key is the line number and the value is list of semicolon counts.
    # Each count indicates a different collection of semicolons.
    #
    # These counts are used for the multi-dimensional arrays syntax introduced inline
    # v1.7. For other purposes simply checking if the line has a semicolon is sufficient.
    semicolons::Dict{Int,Vector{Int}}

    # List of tuples where a tuple contains
    # the start and end lines of regions in the
    # file formatting should be skipped.
    format_skips::Vector{Tuple{Int,Int,String}}
end

function Document(text::AbstractString)
    ranges = UnitRange{Int}[]
    lit_strings = Dict{Int,Tuple{Int,Int,String}}()
    comments = Dict{Int,Tuple{Int,String}}()
    semicolons = Dict{Int,Vector{Int}}()
    format_skips = Tuple{Int,Int,String}[]
    prev_tok = Tokens.Token() # dummy initial token
    stack = Int[]
    format_on = true
    str = ""

    goffset = 0
    for (idx, t) in enumerate(Tokenize.tokenize(text))
        if t.kind === Tokens.WHITESPACE
            offset = goffset
            for c in t.val
                if c == '\n'
                    s = length(ranges) > 0 ? last(ranges[end]) + 1 : 1
                    push!(ranges, s:offset+1)
                end
                offset += 1
            end
        elseif t.kind === Tokens.ENDMARKER
            s = length(ranges) > 0 ? last(ranges[end]) + 1 : 1
            push!(ranges, s:goffset)
        elseif is_str_or_cmd(t.kind)
            offset = goffset
            lit_strings[offset] = (t.startpos[1], t.endpos[1], t.val)
            if t.startpos[1] != t.endpos[1]
                nls = findall(x -> x == '\n', t.val)
                bidx = 1
                cidx = 1
                for nl in nls
                    s = length(ranges) > 0 ? last(ranges[end]) + 1 : 1

                    # newline position in character length instead
                    # of byte length.
                    nl2 = cidx + length(t.val[bidx:nl]) - 1
                    push!(ranges, s:offset+nl2)

                    bidx = nl + 1
                    cidx = nl2 + 1
                end
            end
        elseif t.kind === Tokens.COMMENT
            ws = 0
            if prev_tok.kind === Tokens.WHITESPACE
                # Handles the case where the value of the
                # WHITESPACE token is like " \n ".
                i = findlast(c -> c == '\n', prev_tok.val)
                i === nothing && (i = 1)
                ws = count(c -> c == ' ', prev_tok.val[i:end])
            end

            if t.startpos[1] == t.endpos[1]
                # Determine the number of spaces prior to a possible inline comment
                comments[t.startpos[1]] = (ws, t.val)
            else
                # multiline comment of the form
                # #=
                #
                # #=

                line = t.startpos[1]
                offset = goffset
                cs = ""
                for (i, c) in enumerate(t.val)
                    cs *= c
                    if c == '\n'
                        s = length(ranges) > 0 ? last(ranges[end]) + 1 : 1
                        push!(ranges, s:offset+1)
                        fc = findfirst(c -> !isspace(c), cs)
                        if fc === nothing
                            # comment is all whitespace
                            comments[line] = (ws, cs[end:end])
                        else
                            idx = min(fc, ws + 1)
                            comments[line] = (ws, cs[idx:end])
                        end
                        line += 1
                        cs = ""
                    end
                    offset += 1
                end
                # last comment
                idx = min(findfirst(c -> !isspace(c), cs), ws + 1)
                comments[line] = (ws, cs[idx:end])
            end

            if occursin(r"^#!\s*format\s*:\s*off\s*$", t.val) && length(stack) == 0
                # There should not be more than 1
                # "off" tag on the stack at a time.
                push!(stack, t.startpos[1])
                format_on = false
            elseif occursin(r"^#!\s*format\s*:\s*on\s*$", t.val) && length(stack) > 0
                # If "#! format: off" has not been seen
                # "#! format: on" is treated as a normal comment.
                idx1 = findfirst(c -> c == '\n', str)
                idx2 = findlast(c -> c == '\n', str)
                str = str[idx1:idx2]
                push!(format_skips, (pop!(stack), t.startpos[1], str))
                str = ""
                format_on = true
            end
        elseif t.kind === Tokens.SEMICOLON
            if haskey(semicolons, t.startpos[1])
                if prev_tok.kind === Tokens.SEMICOLON
                    semicolons[t.startpos[1]][end] += 1
                else
                    push!(semicolons[t.startpos[1]], 1)
                end
            else
                semicolons[t.startpos[1]] = Int[1]
            end
        end
        prev_tok = t

        if t.kind === Tokens.COMMENT
            goffset += (t.endbyte - t.startbyte + 1)
        elseif t.kind === Tokens.WHITESPACE
            goffset += (t.endbyte - t.startbyte + 1)
        else
            goffset += length(Tokenize.untokenize(t))
        end

        if !format_on
            str *= Tokenize.untokenize(t)
        end
    end

    range_to_line = SortedDict{UnitRange{Int},Int}(IntervalTreeOrder())
    line_to_range = Dict{Int,UnitRange{Int}}()
    for (l, r) in enumerate(ranges)
        insert!(range_to_line, r, l)
        line_to_range[l] = r
    end

    # If there is a SINGLE "#! format: off" tag
    # do not format from the "off" tag onwards.
    if length(stack) == 1 && length(format_skips) == 0
        # -1 signifies everything afterwards "#! format: off"
        # will not formatted.
        idx1 = findfirst(c -> c == '\n', str)
        str = str[idx1:end]
        push!(format_skips, (stack[1], -1, str))
    end

    return Document(
        text,
        range_to_line,
        line_to_range,
        lit_strings,
        comments,
        semicolons,
        format_skips,
    )
end

ext: jl


text:
cd(@__DIR__)
using Pkg
Pkg.activate(".")

# To load diana changes automatically, uncomment if you wish to try new things with Diana
# Pkg.add(PackageSpec(name="Diana", url="../../"))

function main()
    include(joinpath("src", "TbsServerJulia.jl"))
end; main()


text:
lib_name = "liblsl64"
lib_path = ["$(@__DIR__)/lib"]

const LIB = Libdl.find_library(lib_name,lib_path)

forever = 32000000.0

type LslError
    code::Int
end

function _lsl_get_channel_format(o::Ptr{Void})
    ccall((:lsl_get_channel_format,Lsl.LIB),Cint,(Ptr{Void},),o)
end

function _lsl_get_channel_count(o::Ptr{Void})
    ccall((:lsl_get_channel_count,Lsl.LIB),Cint,(Ptr{Void},),o)
end

function _lsl_open_stream(o::Ptr{Void};timeout::Float64=forever)
    er = Ref{Cint}(0)
    ccall((:lsl_open_stream,LIB),Void,(Ptr{Void},Cdouble,Ref{Cint}),o,timeout,er)
    er[] != 0 && throw(LslError(er[]))
end

function _lsl_time_correction(o::Ptr{Void};timeout::Float64=forever)
    er = Ref{Cint}(0)
    offset = ccall((:lsl_time_correction,LIB),Cdouble,(Ptr{Void},Cdouble,Ref{Cint}),o,timeout,er)
    er[] != 0 && throw(LslError(er[]))
    return offset
end


text:
# Labs 3 & 4

"""
    nwscore(base1, base2)

Your documentation here.
"""
function nwscore(base1::Char, base2::Char)

end

nwscore(base::Char, ::Nothing) = -1
nwscore(::Nothing, base::Char) = nwscore(base, nothing)

nwscore(::Nothing, ::Nothing) = throw(ArgumentError("Score for two gaps is not defined"))


text:
import Base.isequal
import Base.delete!, Base.union
import Base.join

export add!, delete!, contract!,  induce, add_edges!
export line_graph, complement, complement!, adjoint
export cartesian, lex, relabel, trim
export disjoint_union, union, join

# check for graph equality
function isequal(G::SimpleGraph, H::SimpleGraph)
    if G.V != H.V || NE(G) != NE(H)
        return false
    end
    try  # check vertex sets for sortability
        sort(collect(G.V))
        # sort(collect(H.V))
    catch  #if not sortable, resort to a slow method
        for e in G.E
            if ! has(H,e[1],e[2])
                return false
            end
        end
        return true
    end
    # otherwise faster to check edge sets this way
    return G.E==H.E
end

function ==(G::SimpleGraph, H::SimpleGraph)
    return isequal(G,H)
end

# adding vertices
"""
`add!(G,v)` adds a new vertex `v` to the graph.

`add!(G,v,w)` adds a new edge `(v,w)` to the graph. If one (or both)
of those vertices is not already in the graph, it is added to the
vertex set.
"""
function add!(G::SimpleGraph{T}, v) where {T}
    if has(G,v)
        return false
    end
    cache_clear(G)
    push!(G.V, v)
    if G.Nflag
        G.N[v] = Set{T}()
    end
    return true
end

# adding edges
function add!(G::SimpleGraph{T}, v, w) where {T}
    if v==w
        return false
    end

    try
        if v > w
            v,w = w,v
        end
    catch
    end

    if ~has(G,v)
        add!(G,v)
    end
    if ~has(G,w)
        add!(G,w)
    end
    if has(G,v,w)
        return false
    end
    cache_clear(G)
    push!(G.E, (v,w))
    if G.Nflag
        push!(G.N[v], w)
        push!(G.N[w], v)
    end
    return true
end

"""
`add_edges!(G,edge_table)` adds edges to the graph `G`. The argument
`edge_table` is an iterable list (array or set) of edges to be
added (two-tuples). Returns the number of edges added to the graph.

This works both when `G` is a `SimpleGraph` and when `G` is
a `SimpleDigraph`.
"""
function add_edges!(G::AbstractSimpleGraph, edge_list)::Int
    m0 = NE(G)
    for e in edge_list
        add!(G,e[1],e[2])
    end
    return NE(G)-m0
end

# edge deletion
"""
`delete!(G,v)` deletes vertex `v` (and any edges incident with `v`)
from the graph.

`delete!(G,v,w)` deletes the edge `(v,w)` from `G`.
"""
function delete!(G::SimpleGraph, v, w)
    flag = false
    if has(G,v,w)
        cache_clear(G)
        flag = true
        delete!(G.E,(v,w))
        delete!(G.E,(w,v))
        if G.Nflag
            delete!(G.N[v],w)
            delete!(G.N[w],v)
        end
    end
    return flag
end

# vertex deletion
function delete!(G::SimpleGraph, v)
    flag = false
    if has(G,v)
        cache_clear(G)
        flag = true
        Nv = G[v]
        for w in Nv
            delete!(G,v,w)
        end
        delete!(G.V, v)
        if G.Nflag
            delete!(G.N,v)
        end
    end
    return flag
end

# Contract an edge in a graph. If uv is an edge of G, we add all of
# v's neighbors to u's neighborhood and then delete v. This mutates
# the graph. Usually, vertices u and v are adjacent, but we don't
# require that. If either u or v is not a vertex of this graph, or if
# u==v, we return false. Otherwise we return true to indicate success.
"""
`contract!(G,u,v)` contracts the edge `(u,v)` in the graph. The merged
vertex is named `u` (hence `contract!(G,v,u)` results in a different,
albeit isomorphic, graph).

Note: The edge `(u,v)` need not be present in the graph. If missing,
this is equivalent to first adding the edge to the graph and then
contracting it.
"""
function contract!(G::SimpleGraph, u, v)
    if !has(G,u) || !has(G,v) || u==v
        return false
    end
    cache_clear(G)
    Nv = G[v]
    for x in Nv
        add!(G,u,x)
    end
    delete!(G,v)
    return true
end

# Given a simple graph G and a set of vertices A, form the induced
# subgraph G[A]. Note that A must be a subset of V(G).
"""
`induce(G,A)` creates the induced subgraph of `G` with vertices in the
set `A`.
"""
function induce(G::SimpleGraph{T}, A::Set) where {T}
    # Check that A is a subset of V(G)
    for v in A
        if ~has(G,v)
            error("The set A must be a subset of V(G)")
        end
    end
    H = SimpleGraph{T}()  # place to hold the answer

    # add all the vertices in A to H
    for v in A
        add!(H,v)
    end

    # The method we choose depends on the size of A. For small A, best
    # to iterate over pairs of elements of A. For large A, best to
    # iterate over G.E
    a = length(A)

    if a<2
        return H
    end

    # case of small A, iterate over pairs from A
    if a*(a-1) < 2*NE(G)
        alist = collect(A)
        for i=1:a-1
            u = alist[i]
            for j=i+1:a
                v = alist[j]
                if has(G,u,v)
                    add!(H,u,v)
                end
            end
        end

    else  # case of large A, iterate over G.E
        for e in G.E
            if in(e[1],A) && in(e[2],A)
                add!(H,e[1],e[2])
            end
        end
    end
    return H
end

# Create the line graph of a given graph
"""
`line_graph(G)` creates the line graph of `G`.
"""
function line_graph(G::SimpleGraph{T}) where {T}
    H = SimpleGraph{Tuple{T,T}}()

    m = NE(G)
    E = elist(G)
    for e in E
        add!(H,e)
    end

    for i=1:m-1
        e = E[i]
        for j=i+1:m
            f = E[j]
            if e[1]==f[1] || e[2]==f[1] || e[1]==f[2] || e[2]==f[2]
                add!(H,e,f)
            end
        end
    end
    return H
end

# Create the complement of a graph.
"""
`complement(G)` creates (as a new graph) the complement of `G`.
Note that `G'` is a short cut for `complement(G)`.
"""
function complement(G::SimpleGraph{T}) where {T}
    H = SimpleGraph{T}()
    V = vlist(G)
    n = NV(G)

    for i=1:n
        add!(H,V[i])
    end

    for i=1:n-1
        v = V[i]
        for j=i+1:n
            w = V[j]
            if ~has(G,v,w)
                add!(H,v,w)
            end
        end
    end
    return H
end

# We can use G' to mean complement(G)
"""
`G'` is equivalent to `complement(G)`.
"""
adjoint(G::SimpleGraph) = complement(G)

# complement in place. Returns None.
"""
`complement!(G)` replaces `G` with its complement.
"""
function complement!(G::SimpleGraph)
    cache_clear(G)
    n = NV(G)
    V = vlist(G)
    for i=1:n-1
        v = V[i]
        for j=i+1:n
            w = V[j]
            if has(G,v,w)
                delete!(G,v,w)
            else
                add!(G,v,w)
            end
        end
    end
    nothing
end

# Create the cartesian product of two graphs
"""
`cartesian(G,H)` creates the Cartesian product of the two graphs.
This can be abbreviated as `G*H`.
"""
function cartesian(G::SimpleGraph{S}, H::SimpleGraph{T}) where {S,T}
    K = SimpleGraph{Tuple{S,T}}()
    for v in G.V
        for w in H.V
            add!(K,(v,w))
        end
    end

    for e in G.E
        (u,v) = e
        for w in H.V
            add!(K, (u,w), (v,w))
        end
    end

    for e in H.E
        (u,v) = e
        for w in G.V
            add!(K, (w,u), (w,v))
        end
    end

    return K
end

# Use G*H for Cartesian product
"""
For `SimpleGraph`s: `G*H` is equivalent to `cartesian(G,H)`.
"""
function *(G::SimpleGraph{S},H::SimpleGraph{T}) where {S,T}
    return cartesian(G,H)
end

# The join of two graphs is formed by taking disjoint copies of the
# graphs and all possible edges between the two.

"""
`join(G,H)` is a new graph formed by taking disjoint copies of
`G` and `H` together with all possible edges between those copies.
"""
function join(G::SimpleGraph{S}, H::SimpleGraph{T}) where {S,T}
    K = disjoint_union(G,H)
    for v in G.V
        for w in H.V
            add!(K, (v,1), (w,2) )
        end
    end
    return K
end

# Create the union of two graphs. If they have vertices or edges in
# common, that's OK.

"""
`union(G,H)` creates the union of the graphs `G` and `H`. The graphs
may (and typically do) have common vertices or edges.
"""
function union(G::SimpleGraph{S}, H::SimpleGraph{T}) where {S,T}
    if S==T
        K = SimpleGraph{S}()
    else
        K = SimpleGraph{Any}()
    end

    for v in G.V
        add!(K,v)
    end
    for v in H.V
        add!(K,v)
    end

    for e in G.E
        add!(K,e[1],e[2])
    end
    for e in H.E
        add!(K,e[1],e[2])
    end
    return K
end

# This is an unexposed helper function that takes a graph and creates
# a new graph in which the name of each vertex has an integer
# appended. For example, if the vertex type is String in the original
# graph, the new vertices are type (String, Int).
function label_append(G::SimpleGraph{S}, a::Int) where {S}
    mapper = Dict{S,Tuple{S,Int}}()
    for v in G.V
        mapper[v] = (v,a)
    end
    return relabel(G,mapper)
end

# The disjoint_union of two graphs, G and H, is a new graph consisting of
# disjoint copies of G and H.

"""
`disjoint_union(G,H)` is a new graph formed by taking disjoint copies
of `G` and `H` (and no additional edges).
"""
function disjoint_union(G::SimpleGraph{S}, H::SimpleGraph{T}) where {S,T}
    GG = label_append(G,1)
    HH = label_append(H,2)
    if S==T
        K = SimpleGraph{Tuple{S,Int}}()
    else
        K = SimpleGraph{Any}()
    end

    for v in GG.V
        add!(K,v)
    end
    for v in HH.V
        add!(K,v)
    end

    for e in GG.E
        add!(K,e[1],e[2])
    end
    for e in HH.E
        add!(K,e[1],e[2])
    end
    return K
end

# Repeatedly remove vertices with the given degree or less until there
# are no such vertices remaining in the graph. The default trim(G)
# simply removes all isolated vertices.

"""
`trim(G)` returns a copy of `G` with all isolated vertices removed.

`trim(G,d)` returns a copy of `G` in which we iteratively remove all
vertices of degree `d` or smaller. For example, if `G` is a tree,
`trim(G,1)` will eventually remove all vertices.
"""
function trim(G::SimpleGraph, d::Int = 0)
    H = deepcopy(G)
    while NV(H) > 0 && minimum(deg(H)) <= d
        for v in H.V
            if deg(H,v) <= d
                delete!(H,v)
            end
        end
    end
    return H
end

# Relabel the vertics of a graph based on a dictionary mapping old
# vertex names to new

"""
`relabel(G)` returns a copy of `G` in which the vertices are renamed
`1:n`.

`relabel(G,d)` (where `d` is a `Dict`) returns a copy of `G` in which
vertex `v` is renamed `d[v]`.
"""
function relabel(G::SimpleGraph{S}, label::Dict{S,T}) where {S,T}
    H = SimpleGraph{T}()
    for v in G.V
        add!(H,label[v])
    end

    for e in G.E
        u = label[e[1]]
        v = label[e[2]]
        add!(H,u,v)
    end
    return H
end

# Relabel the vertices with the integers 1:n
function relabel(G::SimpleGraph{S}) where {S}
    verts = vlist(G)
    n = length(verts)
    label = Dict{S,Int}()
    sizehint!(label,n)

    for idx = 1:n
        label[verts[idx]] = idx
    end

    return relabel(G,label)
end

"""
`lex(G,H)` computes the lexicographic product of the graphs `G` and
`H`. The vertex set of this graph is the set of all ordered pairs
`(g,h)` (with `g` a vertex of `G` and `h` a vertex of `H` in which we
have the edge `(g,h)~(g',h')` if either `g~g'` in `G` or else `g=g`
and `h~h'`.

We can use the notation `G[H]` also to create `lex(G,H)`.
"""
function lex(G::SimpleGraph{S}, H::SimpleGraph{T}) where {S,T}
    VT = Tuple{S,T}
    K = SimpleGraph{VT}()
    # Create vertex set
    for a in G.V
        for b in H.V
            add!(K,(a,b))
        end
    end

    # Add edges of the form (a,x)~(b,y) where a~b
    for e in G.E
        a = e[1]
        b = e[2]

        for x in H.V
            for y in H.V
                add!(K, (a,x), (b,y))
            end
        end
    end

    # Add edges of the form (a,x)~(a,y) where x~y
    for a in G.V
        for e in H.E
            x = e[1]
            y = e[2]
            add!(K, (a,x), (a,y))
        end
    end

    return K
end

"""
Abbreviation for `lex(G,H)` for `SimpleGraph`s.
"""
getindex(G::SimpleGraph, H::SimpleGraph) = lex(G,H)


text:
using AccurateArithmetic
using ForwardDiff
using LinearAlgebra
using MuladdMacro

#this allows to bail out on other numeric types, but
#keep using the accurate arithmetic package when available

function accurate_dot(x::AbstractArray{T1},y::AbstractArray{T2}) where {T1<:Union{Float32,Float64},T2<:Union{Float32,Float64}}
    return dot_oro(x,y)
end

function accurate_dot(x,y)
    return dot(x,y)
end
@inline _isbad(x) = isnan(x) || isinf(x)
@inline _mod1(x, m) = (x - 1) % m + 1

Base.@kwdef mutable struct State{T<:Real}

    check_obj    :: Bool
    Lp           :: T
    maps_limit   :: T
    time_limit   :: T
    buffer       :: T
    t0           :: T
    tol          :: T

    go_on        :: Bool = true
    converged    :: Bool = false
    maps         :: Int64 = 0 # nb of maps or of g! calls if optim == true.
    k            :: Int64 = 0
    f_calls      :: Int64 = 0
    σ            :: T = zero(T)
    α            :: T = one(T)
    obj_now      :: T = T(Inf)
    norm_∇       :: T = T(Inf)

    ix           :: Int64 = 1
    ord_best     :: Int64 = 0
    i_ord        :: Int64 = 0
    α_best       :: T = one(T)
    σ_mult_fail  :: T = one(T)
    σ_mult_loop  :: T = one(T)
    α_mult       :: T = one(T)
    norm_best    :: T = T(Inf)
    obj_best     :: T = T(Inf)
    α_boost      :: T = one(T)

    σs_i         :: Int64 = 1
end

function Base.eltype(::State{F}) where F
    return F
end

#####
##### Core functions
#####

# Makes sure x_try[i] stays within boundaries with a gap buffer * (bound[i] - x_old[i]).
function bound!(
    extr, x_try :: T, x_old :: T, bound :: T, buffer
) where {T<:AbstractArray}

    for i ∈ eachindex(x_try)
        x_try[i] = extr(x_try[i], (1 - buffer) * bound[i] + buffer * x_old[i])
    end
    return nothing
end

@inline function descent!(
    x_out :: T, ∇ :: T, s :: State,
    x_in :: T, lower :: Union{T, Nothing}, upper :: Union{T, Nothing}
) where {T<:AbstractArray}
    @muladd @. x_out = x_in - s.α * ∇
    if lower ≠ nothing; bound!(max, x_out, x_in, lower, s.buffer) end
    if upper ≠ nothing; bound!(min, x_out, x_in, upper, s.buffer) end
    if lower ≠ nothing || upper ≠ nothing; @muladd @. ∇ = (x_in - x_out) / s.α end
    return nothing
end

@inline function update_x!(
    gm!, x_out :: T, ∇ :: T, s :: State, x_in :: T,
    lower :: Union{T, Nothing}, upper :: Union{T, Nothing}, is_map :: Bool
) where {T<:AbstractArray}

    if !is_map
        if s.k > 1 || s.ix > 1; gm!(∇, x_in) end   # If s.k == 1 && s.ix ≤ 1, ∇ and x were...
        descent!(x_out, ∇, s, x_in, lower, upper)  # ...already computed while initializing α.
    else
        ∇ .= x_in
        gm!(x_out, x_in)
        ∇ .-= x_out

        # If the algorithm has reached an infeasible value, it helps stability ...
        # ...if we also slow down the change of x from mappings.
        if s.α < 1; @muladd @. x_out += (1 - s.α) * ∇ end
    end
    s.norm_∇ =  norm(∇, s.Lp)
    return nothing
end

function check_convergence!(f, s :: State, x :: AbstractArray)
    if s.norm_∇ <= s.tol # only based on norm_∇ (for now)
        s.go_on = false

        if s.check_obj
            s.obj_now = f(x)
            s.f_calls += 1
        end
        s.converged = !s.check_obj || !_isbad(s.obj_now)
    end
    return nothing
end

@inline function mapping!(
    f, gm!, x_out :: T, ∇ :: T, s :: State, info, x_in :: T,
    lower :: Union{T, Nothing}, upper :: Union{T, Nothing}, is_map :: Bool, store_info :: Bool
) where {T<:AbstractArray}

    if s.go_on
        update_x!(gm!, x_out, ∇, s, x_in, lower, upper, is_map)
        s.ix += 1
        s.maps += 1
        !_isbad(s.norm_∇) ? check_convergence!(f, s, x_out) : s.go_on = false
        if store_info; store_info!(info, x_out, s, 1, false) end
    end
    return nothing
end

function prodΔ(∇s :: Vector{T}, tmp1 :: T, tmp2 :: T, p :: Int64) where {T<:AbstractArray}
    @. tmp1 = ∇s[2] - ∇s[1]
    if p == 2
        return accurate_dot(tmp1, ∇s[1]), accurate_dot(tmp1, tmp1)
    elseif p == 3
        @muladd @. tmp2 = ∇s[3] - 2∇s[2] + ∇s[1]
        return accurate_dot(tmp2, tmp1), accurate_dot(tmp2, tmp2)
    end
end

function extrapolate!(
    x_out :: T, x_in :: T, ∇s :: Vector{T}, s :: State, p :: Int64
) where {T<:AbstractArray}

    if p == 1
        @muladd @. x_out = x_in - s.α * ∇s[1] # Just a stabilizing step
    elseif p == 2
		c1 = s.α * (2s.σ - s.σ^2)
		c2 = s.α * s.σ^2
        @muladd @. x_out = x_in - c1 * ∇s[1] - c2 * ∇s[2]
    elseif p == 3
		c1 = s.α * (3s.σ - 3s.σ^2 + s.σ^3)
		c2 = s.α * (3s.σ^2 - 2s.σ^3)
		c3 = s.α * s.σ^3
        @muladd @. x_out = x_in - c1 * ∇s[1] - c2 * ∇s[2] - c3 * ∇s[3]
    end
    return nothing
end

@inline function update_α!(s :: State{F}, σ_new :: F, ΔᵇΔᵇ :: F) where F

    # Increasing / decreasing α to maintain σ ∈ [1, 2] as much as possible.
    s.α *= F(1.5)^((σ_new > 2) - (σ_new < 1))

    if ΔᵇΔᵇ < 1e-100 # if ΔᵇΔᵇ gets really small (fairly rare)...
        s.α_boost *= 2 # Increasing more aggressively each time
        s.α = min(s.α * s.α_boost, one(F)) # We boost α
    end
    return nothing
end

#####
##### Functions to handle failures
#####

# update_progress! is used to know which iterate to fall back on when
# backtracking and, optionally, which x was the smallest minimizer in case the
# problem has multiple minima.

@inline function update_progress!(f, s :: State{F}, x₀ :: AbstractArray, x_best) where F
    if s.go_on && s.check_obj
        s.obj_now = f(x₀)
        s.f_calls += 1
        s.go_on = !_isbad(s.obj_now)
    end

    if s.go_on && (( s.check_obj && s.obj_now < s.obj_best)
                || (!s.check_obj && s.norm_∇ < s.norm_best))
		x_best .= x₀
        s.ord_best, s.α_best = (s.i_ord, s.α)
        s.σ_mult_fail = s.α_mult = one(F)
        s.check_obj ? s.obj_best = s.obj_now : s.norm_best = s.norm_∇
    end
    return nothing
end

function backtrack!(s :: State, x₀ :: T, x_best :: T) where T <: AbstractArray

    x₀ .= x_best
    s.i_ord = s.ord_best
    s.σ_mult_fail /= 2
    s.α_mult /= 2
    s.α = s.α_best * s.α_mult
    if s.check_obj; s.obj_now = s.obj_best end
    return nothing
end

# The future of the algorithm is fully determined by x₀, α and i_ord. Since ACX
# non-monotonic, the same (x₀, α, i_ord) could in principle appear twice with
# non-zero probability, creating an infinite loop. The probability is very small,
# so little resources should be devoted to it. We verify this by comparing values
# of σ and norm(∇) over time. If (σ, norm(∇)) appear twice within a certain number
# of iterations, we divide α and σ by 2.

function check_∞_loop!(
    s :: State, last_order :: Bool, σs :: Vector{T}, norm_∇s :: Vector{T}
) where T<:Real

    if s.σ > σs[s.σs_i]
        σs[s.σs_i] = s.σ
        norm_∇s[s.σs_i] = s.norm_∇
    end

    l = length(σs)
    if last_order && σs[s.σs_i] > 0.001 && s.σ_mult_loop > 0.01 && s.σs_i == l
        loop = false
        for i ∈ 1:l - 1, j ∈ i+1:l
            if σs[i] > 1e-10 && abs(σs[i] - 1) > 1e-10 &&
                    abs(1 - σs[j]/σs[i]) < 1e-10 &&
                    abs(1 - norm_∇s[j]/norm_∇s[i]) < 1e-10
                loop = true
                s.σ_mult_loop /= 2
                s.α /= 2
                σs .= 0
                norm_∇s .= 0
                break
            end
        end
        if !loop; s.σ_mult_loop = 1.0 end
    end
    s.σs_i = _mod1(s.σs_i + last_order, l)
    return nothing
end

#####
##### Miscellaneous
#####

function store_info!(info, x, s :: State, p, extrapolating)

    push!(info.x, copy(x))
    push!(info.σ, s.σ)
    push!(info.α, s.α)
    push!(info.p, p)
    push!(info.extrapolating, extrapolating)
    return nothing
end

#####
##### Initialization funtions
#####

function check_arguments(
    f, x_in :: T, lower :: Union{T, Nothing}, upper :: Union{T, Nothing},
    s :: State
) where {T<:AbstractArray}

    if s.check_obj && f === nothing
        throw(ArgumentError("if check_obj == true, f must be provided."))
    end
    if (upper ≠ nothing && maximum(x_in .- upper) > 0) ||
       (lower ≠ nothing && maximum(lower .- x_in) > 0)
        throw(DomainError(x_in, "infeasible starting point"))
    end

    if !((eltype(x_in) <: Real) | (eltype(x_in) <: Integer))
        throw(ArgumentError("starting point must be of type Float"))
    end
    return nothing
end

# While ACX tolerates a wide range of starting descent step sizes (α), some
# difficult problems benefit from a good initial α. A good starting α has the
# largest possible value while respecting two conditions: i) The Armijo condition
# (if using_f == true) with constant set to 0.25 and ii) the gradient norm does
# not increase too fast (|∇f(x₁)|₂/|∇f(x₀)|₂ ≤ 2).

function α_too_large!(
    f, gm!, s :: State{F}, ∇s :: Vector{T}, xs :: Vector{T},
    x_in :: T, ∇_in :: T, lower :: Union{T, Nothing}, upper :: Union{T, Nothing},
    using_f :: Bool
) :: Bool where {F,T<:AbstractArray}

    ∇s[1] .= ∇_in # Because of box constraints, descent_x! may modify ∇s[1]
    descent!(xs[1], ∇s[1], s, x_in, lower, upper)
    ∇∇1 = ∇s[1] ⋅ ∇s[1]

    if using_f
        obj_new = f(xs[1])
        s.f_calls += 1
        return _isbad(obj_new) || (obj_new > s.obj_now - 0.25s.α * ∇∇1)
    else
        gm!(∇s[2], xs[1])
        s.maps += 1
        descent!(xs[2], ∇s[2], s, xs[1], lower, upper)
        ∇∇2 = ∇s[2] ⋅ ∇s[2]
        return _isbad(∇∇2) || ∇∇2 / ∇∇1 > 4
    end
end

# From an initial α = 1, we converge on the largest possible α within a factor of
# 2 that respects the Armijo condition and the gradient not increasing too fast.
@inline function initialize_α!(
    f, gm!, s :: State{F}, ∇s :: Vector{T}, xs :: Vector{T}, tmp1 :: T, tmp2 :: T,
    x_in :: T, ∇_in :: T, lower :: Union{T, Nothing}, upper :: Union{T, Nothing}
) where {F,T<:AbstractArray}

    max_α = 1e10
    min_mult = 2.0
    mult = min_mult * 64^2
    using_f = f ≠ nothing # Armijo condition first since computing f should be faster
    is_too_large = α_too_large!(f, gm!, s, ∇s, xs, x_in, ∇_in, lower, upper, using_f)
    for i ∈ 0:100
        s.α *= mult^(1 - 2is_too_large)
        was_too_large = is_too_large
        is_too_large = s.α > max_α || α_too_large!(f, gm!, s, ∇s, xs, x_in, ∇_in, lower, upper, using_f)
        if !is_too_large
            tmp1 .= ∇s[1]
            tmp2 .= ∇s[2]
        end
        if mult == min_mult && (is_too_large ⊻ was_too_large)
            if is_too_large # We are overshooting for the last time
                s.α /= min_mult  # The last good s.α
                ∇s[1] .= tmp1    # The last good ∇s[1], so no need to recompute it.
                ∇s[2] .= tmp2    # The last good ∇s[2], so no need to recompute it.
            end
            if using_f # Second condition to check: the gradient must decrease, or not increase too fast
                if s.α < max_α && !α_too_large!(f, gm!, s, ∇s, xs, x_in, ∇_in,
                    lower, upper, false)
                    break
                else
                    max_α = s.α # Updating max_α to make sure the 1st condition is still respected, just in case.
                    using_f = false
                    mult = min_mult * 64^2
                    is_too_large = true   # To make sure we restart reducing α.
                    was_too_large = false # To make sure we restart reducing α.
                end
            else
                break
            end
        end
        mult = max(mult / 64^(is_too_large ⊻ was_too_large), min_mult)
    end
    return nothing
end

#####
##### Main funtion
#####

# We use this internal _speedmapping function to remove the type
# instability associated with g! and m!.
function _speedmapping(
    x_in :: AbstractArray{F}, f, gm!, is_map :: Bool, orders :: Vector{Int64},
    σ_min, stabilize, check_obj, tol, Lp, maps_limit, time_limit,
    lower, upper, buffer, store_info :: Bool
) where F <: Real

    s = State{F}(; tol, buffer, Lp, maps_limit, time_limit, t0 = time(), check_obj)

    if lower !== nothing && eltype(lower) ≠ F; lower = F.(lower) end
    if upper !== nothing && eltype(upper) ≠ F; lower = F.(upper) end

    check_arguments(f, x_in, lower, upper, s)

    # Inserting stabilization steps
    if stabilize; orders = Int.(vec(hcat(ones(length(orders)),orders)')) end

    # Two x₀s to avoid copying at each improvement (maybe this is excessive optimization?)
    #using copy instead of similar because BigFloats have problems
    x₀ = copy(x_in)
    xs = [similar(x_in) for i ∈ 1:maximum(orders)]
    ∇s = [similar(x_in) for i ∈ 1:maximum(orders)] # Storing ∇s is equivalent to Δs
	x_best = similar(x_in)
    (tmp1, tmp2) = (similar(x_in),similar(x_in)) # temp storage

    if f ≠ nothing && (!is_map || check_obj)
        s.obj_now = s.obj_best = f(x_in) # Useful for initialize_α and tracking progress
        s.f_calls += 1
    end

    if !is_map
        gm!(x_best, x_in) # Here x_best acts purely as temp storage for the initial ∇ to avoid allocating a tmp3.
        s.maps = -1 # To avoid double counting since we'll save the 2 first gradient evaluations
        if x_best ⋅ x_best == 0
            throw(DomainError(x_in, "∇f(x_in) = 0 (extremum or saddle point)"))
        end
        initialize_α!(f, gm!, s, ∇s, xs, tmp1, tmp2, x_in, x_best, lower, upper)
        if abs(∇s[2] ⋅ ∇s[1]) / (∇s[2] ⋅ ∇s[2]) < 1; s.i_ord = length(orders) - 1 end
    end

    info = store_info ? (x = [x_in], σ = [s.σ], α = [s.α], p = [0], extrapolating = [false]) : nothing

    σs      = zeros(F,10)
    norm_∇s = zeros(F,10)
    while !s.converged && s.maps ≤ maps_limit && time() - s.t0 ≤ time_limit
        s.k += 1
        s.i_ord += 1
        s.ix = 0 # Which x is currently beeing updated

        # Avoiding a stabilization step if the last extrapolation was close to 1
        s.i_ord += stabilize && s.i_ord % 2 == 1 && abs(s.σ - 1) < 0.01 && s.k > 1

        io = _mod1(s.i_ord, length(orders))
        p = orders[io]
        s.go_on = true

        mapping!(f, gm!, xs[1], ∇s[1], s, info, x₀, lower, upper, is_map, store_info)
        update_progress!(f, s, x₀, x_best)
        for i ∈ 2:p
            mapping!(f, gm!, xs[i], ∇s[i], s, info, xs[i - 1], lower, upper, is_map, store_info)
        end

        if !s.converged && s.go_on
            if p > 1
                ΔᵃΔᵇ, ΔᵇΔᵇ = prodΔ(∇s, tmp1, tmp2, p)
                σ_new = F(abs(ΔᵃΔᵇ) > 1e-100 && ΔᵇΔᵇ > 1e-100 ? abs(ΔᵃΔᵇ) / ΔᵇΔᵇ : 1.0)
                s.σ = max(σ_min, σ_new) * s.σ_mult_fail * s.σ_mult_loop
            end

            extrapolate!(tmp1, x₀, ∇s, s, p)
            if lower ≠ nothing; bound!(max, tmp1, x₀, lower, s.buffer) end
            if upper ≠ nothing; bound!(min, tmp1, x₀, upper, s.buffer) end
            x₀ .= tmp1
            if store_info; store_info!(info, x₀, s, p, true) end

            if p > 1
                if !is_map; update_α!(s, σ_new, F(ΔᵇΔᵇ)) end
                check_∞_loop!(s, io == length(orders), σs, norm_∇s)
            end
        elseif !s.converged && !s.go_on
            backtrack!(s, x₀, x_best )
        end
    end

    if s.maps > maps_limit
        @warn "Maximum mappings exceeded."
    elseif time() - s.t0 > time_limit
        @warn "Exceeded time limit of $time_limit seconds."
    end

    minimizer = s.check_obj && s.obj_best < s.obj_now ? x_best : xs[s.ix]

    return (minimizer = minimizer, maps = s.maps, f_calls = s.f_calls,
        converged = s.converged, norm_∇ = s.norm_∇, info = info)
end

"""
    SpeedMapping
`speedmapping(x₀; m!, kwargs...)` accelerates the convergence of a mapping
`m!(x_out, x_in)` to a fixed point of `m!` by the Alternating cyclic
extrapolation algorithm (ACX). Since gradient descent is an example
of such mapping, `speedmapping(x0; g!, kwargs...)` can also perform multivariate
optimization based on the gradient function `g!(∇, x)`.

Reference:
N. Lepage-Saucier, _Alternating cyclic extrapolation methods for optimization
algorithms_, arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974

### Arguments

*   `x₀ :: AbstractArray`: The starting point; must be of eltype `Float`.

Main keyword arguments:
*   `m!(x_out, x_in)`: A map for which a fixed-point must be found.
*   `g!(∇, x)`: The gradient of a function to be minimized.
*   `f(x)`: The objective of the function to be minimized. It is useful to *i*)
    compute a good initial α (learning  rate) for the gradient descent, *ii*)
    optimize using autodiff or *iii*) track the progress of the algorithm. In
    case neither `m!` nor `g!` is provided, then `f` is used to compute `g`
    using `ForwardDiff`.
*   `lower, upper :: Union{AbstractArray, Nothing} = nothing`: Box constraints
    for the optimization. NOTE: When appropriate, it should also be used with
    `m!`. Even if `m!` always keeps `x_out` within bounds, an extrapolation
    step could throw `x` out of bounds.
*   `tol :: Float64 = 1e-8, Lp :: Real = 2` When using `m!`, the algorithm
    stops when `norm(F(xₖ) - xₖ, Lp) ≤ tol`. When using `g!`, the algorithm
    stops when `norm(∇f(xₖ), Lp) ≤ tol`.
*   `maps_limit :: Real = 1e6`: Maximum number of `m!` calls or `g!` calls.
*   `time_limit :: Real = 1000`: Maximum time in seconds.

Minor keyword arguments:
*   `orders :: Array{Int64} = [3, 3, 2]` determines `ACX`'s alternating order.
    Must be between 1 and 3 (where 1 means no extrapolation). The two recommended
    orders are [3, 2] and [3, 3, 2], the latter being *potentially* better for
    highly non-linear applications (see paper).
*   `check_obj :: Bool = false`: In case of `NaN` or `Inf` values, the algorithm
    restarts at the best past iterate. If `check_obj = true`, progress is
    monitored with the value of the objective (requires `f`).
    Otherwise, it is monitored with `norm(F(xₖ) - xₖ, Lp)`. Advantages of
    `check_obj = true`: more precise and if the algorithm converges on a bad
    local minimum, it can return the best of all past iterates. Advantages of
    `check_obj = false`: for well-behaved convex problems, it avoids the
    effort and time of providing `f` and calling it at every iteration.
*   `store_info :: Bool = false`: Stores `xₖ`, `σₖ` and `αₖ` (see paper).
*   `buffer :: Float64 = 0.01` If `xₖ` goes out of bounds, it is brought back in
    with a buffer. Ex. `xₖ = buffer * xₖ₋₁ + (1 - buffer) * upper`. Setting
    `buffer = 0.001` may speed-up box-constrained optimization.

Keyword arguments to fine-tune fixed-point mapping acceleration (using `m!`):
*   `σ_min :: Real = 0.0`: Setting to `1` may avoid stalling (see paper).
*   `stabilize :: Bool = false`: performs a stabilization mapping before
    extrapolating. Setting to `true` may improve the performance for
    applications like accelerating the EM or MM algorithms (see paper).

# Example: Finding a dominant eigenvalue
```jldoctest
julia> using LinearAlgebra

julia> using SpeedMapping

julia> A = ones(10) * ones(10)' + Diagonal(1:10);

julia> function power_iteration!(x_out, x_in, A)
           mul!(x_out, A, x_in)
           x_out ./= maximum(abs.(x_out))
       end;

julia> res = speedmapping(ones(10); m! = (x_out, x_in) -> power_iteration!(x_out, x_in, A))
(minimizer = [0.4121491412218099, 0.4409506073968953, 0.47407986465655094, 0.5125916147320677, 0.5579135738427361, 0.612027372716759, 0.6777660406970623, 0.7593262786058275, 0.8632012019116189, 1.0], maps = 16, f_calls = 0, converged = true, norm_∇ = 2.804263612262994e-9, info = nothing)

julia> V = res.minimizer;

julia> dominant_eigenvalue = V'A * V / V'V
16.3100056907922

```
# Example: Minimizing a multidimensional Rosenbrock
```
julia> f(x) = sum(100 * (x[i,1]^2 - x[i,2])^2 + (x[i,1] - 1)^2 for i ∈ 1:size(x,1));

julia> function g!(∇, x)
           ∇[:,1] .=  400(x[:,1].^2 .- x[:,2]) .* x[:,1] .+ 2(x[:,1] .- 1)
           ∇[:,2] .= -200(x[:,1].^2 .- x[:,2])
           return nothing
       end;

julia> x₀ = 1.0 * [-4 -3 -2 -1; 0 1 2 3]';
```

Optimizing, providing f and g!
```
julia> speedmapping(x₀; f, g!)
(minimizer = [0.999999999982878 0.9999999999656839; 0.9999999999732417 0.9999999999463607; 0.9999999998875755 0.9999999997746609; 0.9999999951927082 0.9999999903661674], maps = 180, f_calls = 11, converged = true, norm_∇ = 4.306438901515058e-9, info = nothing)
```

Optimizing without objective
```
julia> speedmapping(x₀; g!)
(minimizer = [1.000000000000002 1.000000000000004; 0.999999999999956 0.9999999999999117; 0.9999999999998761 0.9999999999997516; 0.999999999999863 0.9999999999997254], maps = 148, f_calls = 0, converged = true, norm_∇ = 2.7446698204458767e-13, info = nothing)
```

Optimizing without gradient
```
julia> speedmapping(x₀; f)
(minimizer = [0.9999999999957527 0.9999999999914151; 0.9999999999933037 0.9999999999865801; 0.9999999999716946 0.9999999999432596; 0.9999999987753265 0.999999997545751], maps = 172, f_calls = 11, converged = true, norm_∇ = 1.0971818937506587e-9, info = nothing)
```

Optimizing with a box constraint
```
julia> speedmapping(x₀; f, g!, upper = [0.5ones(4) Inf * ones(4)])
(minimizer = [0.5 0.25; 0.49999999999999434 0.24999999996939753; 0.5 0.24999999999999997; 0.4999999999999999 0.24999999999948902], maps = 71, f_calls = 7, converged = true, norm_∇ = 8.135561263867014e-9, info = nothing)
```
"""
function speedmapping(
    x_in :: AbstractArray{F}; f = nothing, g! = nothing, m! = nothing,
    orders :: Array{Int64} = [3,3,2], σ_min :: Real = zero(F), stabilize :: Bool = false,
    check_obj :: Bool = false, tol = F(1e-8), Lp :: Real = 2,
    maps_limit :: Real = 1e6, time_limit :: Real = 1000,
    lower :: Union{AbstractArray, Nothing} = nothing,
    upper :: Union{AbstractArray, Nothing} = nothing, buffer  = F(0.01),
    store_info :: Bool = false
) where F <: Real

    kwargs = (orders, σ_min, stabilize, check_obj, tol, Lp, maps_limit,
        time_limit, lower, upper, buffer, store_info)
    if m! ≠ nothing
        if g! ≠ nothing; throw(ArgumentError("must not provide both m! and g!")) end
        return _speedmapping(x_in, f, m!, true, kwargs...)
    elseif g! === nothing
        return _speedmapping(x_in, f, (∇, x) -> ForwardDiff.gradient!(∇, f, x),
            false, kwargs...)
    else
        return _speedmapping(x_in, f, g!, false, kwargs...)
    end
end


text:
using SafeTestsets

@safetestset "My f tests" begin
    include("my_f_tests.jl")
end


text:
using CameraCalibrations
using Test

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


text:
module DSS

using ..Grids
using ClimateMachine.MPIStateArrays
using CUDA
using KernelAbstractions
using DocStringExtensions

export dss!

"""
    dss3d(Q::MPIStateArray,
        grid::DiscontinuousSpectralElementGrid)

This function computes the 3D direct stiffness summation for all variables in the MPIStateArray.

# Fields
 - `Q`: MPIStateArray
 - `grid`: Discontinuous Spectral Element Grid
"""
function dss!(
    Q::MPIStateArray,
    grid::DiscontinuousSpectralElementGrid{FT, 3};
    max_threads = 256,
) where {FT}
    DA = arraytype(grid)                                      # device array
    device = arraytype(grid) <: Array ? CPU() : CUDADevice()  # device
    #----communication--------------------------
    event = MPIStateArrays.begin_ghost_exchange!(Q)
    event = MPIStateArrays.end_ghost_exchange!(Q, dependencies = event)
    wait(event)
    #----Direct Stiffness Summation-------------
    vertmap = grid.vertmap
    edgemap = grid.edgemap
    facemap = grid.facemap

    vtconn = grid.topology.vtconn
    fcconn = grid.topology.fcconn
    edgconn = grid.topology.edgconn

    vtconnoff = grid.topology.vtconnoff
    edgconnoff = grid.topology.edgconnoff

    nvt, nfc, nedg =
        length(vtconnoff) - 1, size(fcconn, 1), length(edgconnoff) - 1

    Nq = polynomialorders(grid) .+ 1
    Nqmax = maximum(Nq)
    Nemax = Nqmax - 2
    Nfmax = size(facemap, 1)
    args = (
        Q.data,
        vertmap,
        edgemap,
        facemap,
        vtconn,
        vtconnoff,
        edgconn,
        edgconnoff,
        fcconn,
        nvt,
        nedg,
        nfc,
        Nemax,
        Nfmax,
    )
    if device == CPU()
        dss3d_CPU!(args...)
    else
        n_items = nvt + nfc + nedg
        tx = max(n_items, max_threads)
        bx = cld(n_items, tx)
        @cuda threads = (tx) blocks = (bx) dss3d_CUDA!(args...)
    end
    return nothing
end

function dss3d_CUDA!(
    data,
    vertmap,
    edgemap,
    facemap,
    vtconn,
    vtconnoff,
    edgconn,
    edgconnoff,
    fcconn,
    nvt,
    nedg,
    nfc,
    Nemax,
    Nfmax,
)
    I = eltype(nvt)
    FT = eltype(data)

    tx = threadIdx().x        # threaid id
    bx = blockIdx().x         # block id
    bxdim = blockDim().x      # block dimension
    glx = tx + (bx - 1) * bxdim # global id
    nvars = size(data, 2)

    # A mesh node is either
    # - interior (no dss required)
    # - on an element corner / vertex,
    # - edge (excluding end point element corners / vertices)
    # - face (excluding edges and corners)
    if glx ≤ nvt #vertex DSS
        vx = glx
        for ivar in 1:nvars
            dss_vertex!(vtconn, vtconnoff, vertmap, data, ivar, vx, FT)
        end
    elseif glx > nvt && glx ≤ (nvt + nedg) # edge DSS
        ex = glx - nvt
        for ivar in 1:nvars
            dss_edge!(edgconn, edgconnoff, edgemap, Nemax, data, ivar, ex, FT)
        end
    elseif glx > (nvt + nedg) && glx ≤ (nvt + nedg + nfc) # face DSS
        fx = glx - (nvt + nedg)
        for ivar in 1:nvars
            dss_face!(fcconn, facemap, Nfmax, data, ivar, fx)
        end
    end

    return nothing
end

function dss3d_CPU!(
    data,
    vertmap,
    edgemap,
    facemap,
    vtconn,
    vtconnoff,
    edgconn,
    edgconnoff,
    fcconn,
    nvt,
    nedg,
    nfc,
    Nemax,
    Nfmax,
)
    I = eltype(nvt)
    FT = eltype(data)
    nvars = size(data, 2)

    # A mesh node is either
    # - interior (no dss required)
    # - on an element corner / vertex,
    # - edge (excluding end point element corners / vertices)
    # - face (excluding edges and corners)
    for ivar in 1:nvars
        for vx in 1:nvt # vertex DSS
            dss_vertex!(vtconn, vtconnoff, vertmap, data, ivar, vx, FT)
        end
        for ex in 1:nedg # edge DSS
            dss_edge!(edgconn, edgconnoff, edgemap, Nemax, data, ivar, ex, FT)
        end
        for fx in 1:nfc # face DSS
            dss_face!(fcconn, facemap, Nfmax, data, ivar, fx)
        end
    end
    return nothing
end

"""
    dss_vertex!(
    vtconn,
    vtconnoff,
    vertmap,
    data,
    ivar,
    vx,
    ::Type{FT},
) where {FT}

This function computes the direct stiffness summation for the vertex `vx`.

# Fields
 - `vtconn`: vertex connectivity array
 - `vtconnoff`: offsets for vertex connectivity array
 - `vertmap`: map to vertex degrees of freedom: `vertmap[vx]` contains the
    degree of freedom located at vertex `vx`.
 - `data`: data field of MPIStateArray
 - `ivar`: variable # in the MPIStateArray
 - `vx`: unique edge number
 - `::Type{FT}`: Floating point type
"""
function dss_vertex!(
    vtconn,
    vtconnoff,
    vertmap,
    data,
    ivar,
    vx,
    ::Type{FT},
) where {FT}
    @inbounds st = vtconnoff[vx]
    @inbounds nlvt = Int((vtconnoff[vx + 1] - vtconnoff[vx]) / 2)
    sumv = -FT(0)
    @inbounds for j in 1:nlvt
        lvt = vtconn[st + (j - 1) * 2]
        lelem = vtconn[st + (j - 1) * 2 + 1]
        loc = vertmap[lvt]
        sumv += data[loc, ivar, lelem]
    end
    @inbounds for j in 1:nlvt
        lvt = vtconn[st + (j - 1) * 2]
        lelem = vtconn[st + (j - 1) * 2 + 1]
        loc = vertmap[lvt]
        data[loc, ivar, lelem] = sumv
    end
end

"""
    dss_edge!(
    edgconn,
    edgconnoff,
    edgemap,
    Nemax,
    data,
    ivar,
    ex,
    ::Type{FT},
) where {FT}

This function computes the direct stiffness summation for
all degrees of freedom corresponding to edge `ex`. dss_edge!
applies only to interior (non-vertex) edge nodes.

# Fields
 - `edgconn`: edge connectivity array
 - `edgconnoff`: offsets for edge connectivity array
 - `edgemap`: map to edge degrees of freedom:
    `edgemap[i, edgno, orient]` contains the element node index of
    the `i`th interior node on edge `edgno`, under orientation `orient`.
 - `Nemax`: # of relevant degrees of freedom per edge (other dof are marked as -1)
 - `data`: data field of MPIStateArray
 - `ivar`: variable # in the MPIStateArray
 - `ex`: unique edge number
 - `::Type{FT}`: Floating point type
"""
function dss_edge!(
    edgconn,
    edgconnoff,
    edgemap,
    Nemax,
    data,
    ivar,
    ex,
    ::Type{FT},
) where {FT}
    @inbounds st = edgconnoff[ex]
    @inbounds nledg = Int((edgconnoff[ex + 1] - edgconnoff[ex]) / 3)

    @inbounds for k in 1:Nemax
        sume = -FT(0)
        @inbounds for j in 1:nledg
            ledg = edgconn[st + (j - 1) * 3]
            lor = edgconn[st + (j - 1) * 3 + 1]
            lelem = edgconn[st + (j - 1) * 3 + 2]
            loc = edgemap[k, ledg, lor]
            if loc ≠ -1
                sume += data[loc, ivar, lelem]
            end
        end
        @inbounds for j in 1:nledg
            ledg = edgconn[st + (j - 1) * 3]
            lor = edgconn[st + (j - 1) * 3 + 1]
            lelem = edgconn[st + (j - 1) * 3 + 2]
            loc = edgemap[k, ledg, lor]
            if loc ≠ -1
                data[loc, ivar, lelem] = sume
            end
        end
    end
end
"""
    dss_face!(fcconn, facemap, Nfmax, data, ivar, fx)

This function computes the direct stiffness summation for
all degrees of freedom corresponding to face `fx`. dss_face!
applies only to interior (non-vertex and non-edge) face nodes.

# Fields
 - `fcconn`: face connectivity array
 - `facemap`: map to face degrees of freedom: `facemap[ij, fcno, orient]`
    contains the element node index of the `ij`th
    interior node on face `fcno` under orientation `orient`
 - `Nfmax`: # of relevant degrees of freedom per face (other dof are marked as -1)
 - `data`: data field of MPIStateArray
 - `ivar`: variable # in the MPIStateArray
 - `fx`: unique face number
"""
function dss_face!(fcconn, facemap, Nfmax, data, ivar, fx)
    @inbounds lfc = fcconn[fx, 1]
    @inbounds lel = fcconn[fx, 2]
    @inbounds nabrlfc = fcconn[fx, 3]
    @inbounds nabrlel = fcconn[fx, 4]
    @inbounds ordr = fcconn[fx, 5]
    ordr = ordr == 3 ? 2 : 1
    # mesh orientation 3 is a flip along the horizontal edges.
    # This is the only orientation currently support by the mesh generator
    # see vertsortandorder in
    # src/Numerics/Mesh/BrickMesh.jl
    @inbounds for j in 1:Nfmax
        loc1 = facemap[j, lfc, ordr]
        loc2 = facemap[j, nabrlfc, 1]
        if loc1 ≠ -1 && loc2 ≠ -1
            sumf = data[loc1, ivar, lel] + data[loc2, ivar, nabrlel]
            data[loc1, ivar, lel] = sumf
            data[loc2, ivar, nabrlel] = sumf
        end
    end
end

end


text:
function autodiff_array_gradient(a,i_to_x)
  i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient),i_to_x)
  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient),i_to_x)
  i_to_ydual = a(i_to_xdual)
  i_to_result = lazy_map(AutoDiffMap(ForwardDiff.gradient),i_to_ydual,i_to_x,i_to_cfg)
  i_to_result
end

function autodiff_array_jacobian(a,i_to_x)
  i_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian),i_to_x)
  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian),i_to_x)
  i_to_ydual = a(i_to_xdual)
  i_to_result = lazy_map(AutoDiffMap(ForwardDiff.jacobian),i_to_ydual,i_to_x,i_to_cfg)
  i_to_result
end

function autodiff_array_hessian(a,i_to_x)
  agrad = i_to_y -> autodiff_array_gradient(a,i_to_y)
  autodiff_array_jacobian(agrad,i_to_x)
end

function autodiff_array_gradient(a,i_to_x,j_to_i)
  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient),i_to_x)
  j_to_ydual = a(i_to_xdual)
  j_to_x = lazy_map(Reindex(i_to_x),j_to_i)
  j_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient),j_to_x)
  j_to_result = lazy_map(AutoDiffMap(ForwardDiff.gradient),j_to_ydual,j_to_x,j_to_cfg)
  j_to_result
end

function autodiff_array_jacobian(a,i_to_x,j_to_i)
  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian),i_to_x)
  j_to_ydual = a(i_to_xdual)
  j_to_x = lazy_map(Reindex(i_to_x),j_to_i)
  j_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian),j_to_x)
  j_to_result = lazy_map(AutoDiffMap(ForwardDiff.jacobian),j_to_ydual,j_to_x,j_to_cfg)
  j_to_result
end

function autodiff_array_hessian(a,i_to_x,i_to_j)
  agrad = i_to_y -> autodiff_array_gradient(a,i_to_y,i_to_j)
  autodiff_array_jacobian(agrad,i_to_x,i_to_j)
end

struct ConfigMap{F} <: Map
 f::F
end

# TODO Prescribing long chunk size can lead to slow compilation times!
function return_cache(k::ConfigMap{typeof(ForwardDiff.gradient)},x)
  cfg = ForwardDiff.GradientConfig(nothing,x,ForwardDiff.Chunk{length(x)}())
  cfg
end

function return_cache(k::ConfigMap{typeof(ForwardDiff.jacobian)},x)
  cfg = ForwardDiff.JacobianConfig(nothing,x,ForwardDiff.Chunk{length(x)}())
  cfg
end

function evaluate!(cfg,k::ConfigMap,x)
  cfg
end

struct DualizeMap{F} <: Map
 f::F
end

function return_cache(k::DualizeMap,x)
  return_cache(ConfigMap(k.f),x)
end

function evaluate!(cfg,k::DualizeMap,x)
  xdual = cfg.duals
  ForwardDiff.seed!(xdual, x, cfg.seeds)
  xdual
end

struct AutoDiffMap{F} <: Map
  f::F
end

function return_cache(k::AutoDiffMap,ydual,x,cfg::ForwardDiff.GradientConfig{T}) where T
  ydual isa Real || throw(ForwardDiff.GRAD_ERROR)
  result = similar(x, ForwardDiff.valtype(ydual))
  result
end

function evaluate!(result,k::AutoDiffMap,ydual,x,cfg::ForwardDiff.GradientConfig{T}) where T
  @notimplementedif ForwardDiff.chunksize(cfg) != length(x)
  @notimplementedif length(result) != length(x)
  result = ForwardDiff.extract_gradient!(T, result, ydual)
  return result
end

function return_cache(k::AutoDiffMap,ydual,x,cfg::ForwardDiff.JacobianConfig{T,V,N}) where {F,T,V,N}
  ydual isa AbstractArray || throw(ForwardDiff.JACOBIAN_ERROR)
  result = similar(ydual, ForwardDiff.valtype(eltype(ydual)), length(ydual), N)
  result
end

function evaluate!(result,k::AutoDiffMap,ydual,x,cfg::ForwardDiff.JacobianConfig{T,V,N}) where {F,T,V,N}
  @notimplementedif ForwardDiff.chunksize(cfg) != length(x)
  @notimplementedif size(result,2) != length(x)
  ForwardDiff.extract_jacobian!(T, result, ydual, N)
  ForwardDiff.extract_value!(T, result, ydual)
  return result
end


text:
struct ZGroup{S<:AbstractStore}
    storage::S
    arrays::Dict{String, ZArray}
    groups::Dict{String, ZGroup}
    attrs::Dict
end

zname(g::ZGroup) = zname(g.storage)

#Open an existing ZGroup
function ZGroup(s::T,mode="r") where T <: AbstractStore
  arrays = Dict{String, ZArray}()
  groups = Dict{String, ZGroup}()

  for d in subdirs(s)
    dshort = splitpath(d)[end]
    m = zopen(getsub(s,dshort),mode)
    if isa(m, ZArray)
      arrays[dshort] = m
    else
      groups[dshort] = m
    end
  end
  attrs = getattrs(s)
  ZGroup(s, arrays, groups, attrs)
end

function Base.show(io::IO, g::ZGroup)
    print(io, "ZarrGroup at ", g.storage)
    !isempty(g.arrays) && print(io, "\nVariables: ", map(i -> string(zname(i), " "), values(g.arrays))...)
    !isempty(g.groups) && print(io, "\nGroups: ", map(i -> string(zname(i), " "), values(g.groups))...)
    nothing
end
Base.haskey(g::ZGroup,k)= haskey(g.groups,k) || haskey(g.arrays,k)

function Base.getindex(g::ZGroup, k)
    if haskey(g.groups, k)
        return g.groups[k]
    elseif haskey(g.arrays, k)
       return g.arrays[k]
    else
       throw(KeyError("Zarr Dataset does not contain $k"))
    end
end

"""
    zopen(s::AbstractStore, mode="r"; consolidated = false)

Opens a zarr Array or Group at Store `s`. If `consolidated` is set to "true",
Zarr will search for a consolidated metadata field as created by the python zarr
`consolidate_metadata` function. This can substantially speed up metadata parsing
of large zarr groups.
"""
function zopen(s::AbstractStore, mode="r"; consolidated = false)
    # add interfaces to Stores later
    consolidated && return zopen(ConsolidatedStore(s), mode)
    if is_zarray(s)
        return ZArray(s,mode)
    elseif is_zgroup(s)
        return ZGroup(s,mode)
    else
        x = path(s)
        throw(ArgumentError("Specified store ($x) is neither a ZArray nor a ZGroup"))
    end
end

"""
    zopen(p::String, mode="r")

Open a zarr Array or group at disc path p.
"""
function zopen(s::String, mode="r"; kwargs...)
  zopen(storefromstring(s), mode; kwargs...)
end

import AWSCore
function storefromstring(s)
  if startswith(s,"gs://")
    aws_google = AWSCore.aws_config(creds=nothing, region="", service_host="googleapis.com", service_name="storage")
    decomp = split(s,"/")
    bucket = decomp[3]
    path = join(decomp[4:end],"/")
    S3Store(String(bucket),path, aws=aws_google, listversion=1)
  elseif startswith(s,"http://")
    return ConsolidatedStore(HTTPStore(s))
  else
    return DirectoryStore(s)
  end
end

"""
    zgroup(s::AbstractStore; attrs=Dict())

Create a new zgroup in the store `s`
"""
function zgroup(s::AbstractStore; attrs=Dict())
    d = Dict("zarr_format"=>2)
    isempty(s) || error("Store is not empty")
    b = IOBuffer()
    JSON.print(b,d)
    s[".zgroup"]=take!(b)
    writeattrs(s,attrs)
    ZGroup(s, Dict{String,ZArray}(), Dict{String,ZGroup}(), attrs)
end

zgroup(s::String;kwargs...)=zgroup(DirectoryStore(s);kwargs...)

"Create a subgroup of the group g"
zgroup(g::ZGroup, name; attrs=Dict()) = g.groups[name] = zgroup(newsub(g.storage,name),attrs=attrs)

"Create a new subarray of the group g"
function zcreate(::Type{T},g::ZGroup, name::String, addargs...; kwargs...) where T
  newstore = newsub(g.storage,name)
  z = zcreate(T, newstore, addargs...; kwargs...)
  g.arrays[name] = z
  return z
end

HTTP.serve(s::Union{ZArray,ZGroup}, args...; kwargs...) = HTTP.serve(s.storage, args...; kwargs...)
consolidate_metadata(z::Union{ZArray,ZGroup}) = consolidate_metadata(z.storage)


text:
module FiniteDiff

using LinearAlgebra, SparseArrays, StaticArrays, ArrayInterface, Requires

import Base: resize!

include("iteration_utils.jl")
include("epsilons.jl")
include("derivatives.jl")
include("gradients.jl")
include("jacobians.jl")
include("hessians.jl")

end # module


text:
include("kepler_init.jl")

# Initializes N-body integration for a plane-parallel hierarchical system
# (see Hamers & Portugies-Zwart 2016 (HPZ16); Beust 2003).
# We want to define a "simplex" hierarchy which can be decomposed into N-1 Keplerian binaries.
# This can be diagramed with a "mobile" diagram, pioneered by Evans (1968).  Here's an example:
#   Level             |
#    4         _______|_______
#    3   _____|____           |
#    2  |          |      ____|_____
#    1  |          |     |       ___|____
#       |          |     |      |        |
#       5          4     3      2        1

# Number of levels:  n_level
# Number of bodies:  n_body
#  - Problem is divided up into N-1 Kepler problems:  there is a single Kepler problem at each level.
#  - For example, in the above "mobile" diagram, 1-2 could be a binary star,
#    while 3 is an interior planets orbiting the stars, and 4/5 are a planet/moon orbiting exterior.
#  - We compute the N-body positions with each Keplerian connection (one at each level), starting
#    at the bottom and moving upwards.
#

function init_nbody(elements::Array{T,2},t0::T,n_body::Int64) where {T <: Real}
# the "_plane" is to remind us that this is currently plane-parallel, so inclination & Omega are zero
n_level = n_body-1
# Input -
# elements: masses & orbital elements for each Keplerian (in this case, each planet plus star)
# Output -
# x: NDIM x n_body array of positions  of each planet.
# v: NDIM x n_body array of velocities "   "      "
#
# Read in the orbital elements:
# elements = readdlm("elements.txt",',')
# Initialize N-body for each Keplerian:
# Get the indices:
indices = get_indices_planetary(n_body)
# Set up "A" matrix (Hamers & Portegies-Zwart 2016) which transforms from
# cartesian coordinates to Keplerian orbits (we are using opposite sign
# convention of HPZ16, for example, r_1 = R_2-R_1).
amat = zeros(T,n_body,n_body)
# Mass vector:
mass = vcat(elements[:,1])
# Set up array for orbital positions of each Keplerian:
rkepler = zeros(T,n_body,NDIM)
rdotkepler = zeros(T,n_body,NDIM)
# Fill in the A matrix & compute the Keplerian elements:
for i=1:n_body-1
  # Sums of masses for two components of Keplerian:
  m1 = zero(T)
  m2 = zero(T)
  for j=1:n_body
    if indices[i,j] == 1
      m1 += mass[j]
    end
    if indices[i,j] == -1
      m2 += mass[j]
    end
  end
  # Compute Kepler problem: r is a vector of positions of "body" 2 with respect to "body" 1; rdot is velocity vector
  # For now set inclination to Inclination = pi/2 and longitude of nodes to Omega = pi:
#  r,rdot = kepler_init(t0,m1+m2,[elements[i+1,2:5];pi/2;pi])
  r,rdot = kepler_init(t0,m1+m2,elements[i+1,2:7])
#  rbig,rdotbig = kepler_init(big(t0),big(m1+m2),big.(elements[i+1,2:7]))
#  r = convert(Array{T,1},rbig); rdot = convert(Array{T,1},rdotbig)
  for j=1:NDIM
    rkepler[i,j] = r[j]
    rdotkepler[i,j] = rdot[j]
  end
  # Now, fill in the A matrix
  for j=1:n_body
    if indices[i,j] == 1
      amat[i,j] = -mass[j]/m1
    end
    if indices[i,j] == -1
      amat[i,j] =  mass[j]/m2
    end
  end
end
mtot = sum(mass)
# Final row is for center-of-mass of system:
for j=1:n_body
  amat[n_body,j] = mass[j]/mtot
end
# Compute inverse of A matrix to convert from Keplerian coordinates
# to Cartesian coordinates:
ainv = inv(amat)
# Now, compute the Cartesian coordinates (eqn A6 from HPZ16):
x = zeros(T,NDIM,n_body)
v = zeros(T,NDIM,n_body)
#for i=1:n_body
#  for j=1:NDIM
#    for k=1:n_body
#      x[j,i] += ainv[i,k]*rkepler[k,j]
#      v[j,i] += ainv[i,k]*rdotkepler[k,j]
#    end
#  end
#end
#x = transpose(*(ainv,rkepler))
x = permutedims(*(ainv,rkepler))
#v = transpose(*(ainv,rdotkepler))
v = permutedims(*(ainv,rdotkepler))
# Return the cartesian position & velocity matrices:
#return x,v,amat,ainv
return x,v
end

# Version including derivatives:
function init_nbody(elements::Array{T,2},t0::T,n_body::Int64,jac_init::Array{T,2}) where {T <: Real}

# the "_plane" is to remind us that this is currently plane-parallel, so inclination & Omega are zero
n_level = n_body-1
# Input -
# elements: masses & orbital elements for each Keplerian (in this case, each planet plus star)
# Output -
# x: NDIM x n_body array of positions  of each planet.
# v: NDIM x n_body array of velocities "   "      "
# jac_init: derivative of cartesian coordinates, (x,v,m) for each body, with respect to initial conditions
#  n_body x (period, t0, e*cos(omega), e*sin(omega), inclination, Omega, mass) for each Keplerian, with
#  the first body having orbital elements set to zero, so the first six derivatives are zero.
#jac_init = zeros(T,7*n_body,7*n_body)
fill!(jac_init,0.0)
#
# Read in the orbital elements:
# elements = readdlm("elements.txt",',')
# Initialize N-body for each Keplerian:
# Get the indices:
indices = get_indices_planetary(n_body)
#println("Indices: ",indices)
# Set up "A" matrix (Hamers & Portegies-Zwart 2016) which transforms from
# cartesian coordinates to Keplerian orbits (we are using opposite sign
# convention of HPZ16, for example, r_1 = R_2-R_1).
amat = zeros(T,n_body,n_body)
# Derivative of i,jth element of A matrix with respect to mass of body k:
damatdm = zeros(T,n_body,n_body,n_body)
# Mass vector:
mass = vcat(elements[:,1])
# Set up array for orbital positions of each Keplerian:
rkepler = zeros(T,n_body,NDIM)
rdotkepler = zeros(T,n_body,NDIM)
# Set up Jacobian for transformation from n_body-1 Keplerian elements & masses
# to (x,v,m) - the last is center-of-mass, which is taken to be zero.
jac_21 = zeros(T,7,7)
# jac_kepler saves jac_21 for each set of bodies:
jac_kepler = zeros(T,n_body*6,n_body*7)
# Fill in the A matrix & compute the Keplerian elements:
for i=1:n_body-1  # i labels the row of matrix, which weights masses in current Keplerian
  # Sums of masses for two components of Keplerian:
  m1 = 0.0 ; m2 = 0.0
  for j=1:n_body
    if indices[i,j] == 1
      m1 += mass[j]
    end
    if indices[i,j] == -1
      m2 += mass[j]
    end
  end
  # Compute Kepler problem: r is a vector of positions of "body" 2 with respect to "body" 1; rdot is velocity vector:
  r,rdot = kepler_init(t0,m1+m2,elements[i+1,2:7],jac_21)
#  rbig,rdotbig = kepler_init(big(t0),big(m1+m2),big.(elements[i+1,2:7]))
#  r = convert(Array{T,1},rbig); rdot = convert(Array{T,1},rdotbig)
  for j=1:NDIM
    rkepler[i,j] = r[j]
    rdotkepler[i,j] = rdot[j]
  end
  # Save Keplerian Jacobian to a matrix.  First, positions/velocities vs. elements:
  for j=1:6, k=1:6
    jac_kepler[(i-1)*6+j,i*7+k] = jac_21[j,k]
  end
  # Now add in mass derivatives:
  for j=1:n_body
    # Check which bodies participate in current Keplerian:
    if indices[i,j] != 0
      for k=1:6
        jac_kepler[(i-1)*6+k,j*7] = jac_21[k,7]
      end
    end
  end
  # Now, fill in the A matrix:
  for j=1:n_body
    if indices[i,j] == 1
      amat[i,j] = -mass[j]/m1
      damatdm[i,j,j] += -1.0/m1
      for k=1:n_body
        if indices[i,k] == 1
          damatdm[i,j,k] += mass[j]/m1^2
        end
      end
    end
    if indices[i,j] == -1
      amat[i,j] =  mass[j]/m2
      damatdm[i,j,j] +=  1.0/m2
      for k=1:n_body
        if indices[i,k] == -1
          damatdm[i,j,k] -= mass[j]/m2^2
        end
      end
    end
  end
end
mtot = sum(mass)
for j=1:n_body
  amat[n_body,j] = mass[j]/mtot
  damatdm[n_body,j,j] = 1.0/mtot
  for k=1:n_body
    damatdm[n_body,j,k] -= mass[j]/mtot^2
  end
end
ainv = inv(amat)
# Propagate mass uncertainties (11/17/17 notes):
dainvdm = zeros(T,n_body,n_body,n_body)
#dainvdm_num = zeros(T,n_body,n_body,n_body)
#damatdm_num = zeros(T,n_body,n_body,n_body)
#dlnq = 2e-3
#elements_tmp = zeros(T,n_body,7)
for k=1:n_body
  dainvdm[:,:,k]=-ainv*damatdm[:,:,k]*ainv
# Compute derivatives numerically:
#  elements_tmp = copy(elements)
#  elements_tmp .= elements
#  elements_tmp[k,1] *= (1.0-dlnq)
#  x_minus,v_minus,amat_minus,ainv_minus = init_nbody(elements_tmp,t0,n_body)
#  elements_tmp[k,1] *= (1.0+dlnq)/(1.0-dlnq)
#  x_plus,v_plus,amat_plus,ainv_plus = init_nbody(elements_tmp,t0,n_body)
#  dainvdm_num[:,:,k] = (ainv_plus.-ainv_minus)/(2dlnq*elements[k,1])
#  damatdm_num[:,:,k] = (amat_plus.-amat_minus)/(2dlnq*elements[k,1])
end
#println("damatdm: ",maximum(abs.(damatdm-damatdm_num)))
#println("dainvdm: ",maximum(abs.(dainvdm-dainvdm_num)))
#for k=1:n_body
#  println("damatdm: ",k," ",abs.(damatdm[k,:,:]-damatdm_num[k,:,:]))
##  println("damatdm_num: ",k," ",abs.(damatdm_num[k,:,:]))
#end
#println("dainvdm: ",abs.(dainvdm))
#println("dainvdm_num: ",abs.(dainvdm_num))
# Now, compute the Cartesian coordinates (eqn A6 from HPZ16):
x = zeros(T,NDIM,n_body)
v = zeros(T,NDIM,n_body)
#for i=1:n_body
#  for j=1:NDIM
#    for k=1:n_body
#      x[j,i] += ainv[i,k]*rkepler[k,j]
#      v[j,i] += ainv[i,k]*rdotkepler[k,j]
#    end
#  end
#end
#x = transpose(*(ainv,rkepler))
x = permutedims(*(ainv,rkepler))
#v = transpose(*(ainv,rdotkepler))
v = permutedims(*(ainv,rdotkepler))
# Finally, compute the overall Jacobian.
# First, compute it for the orbital elements:
dxdm = zeros(T,3,n_body); dvdm = zeros(T,3,n_body)
for i=1:n_body
  # Propagate derivatives of Kepler coordinates with respect to
  # orbital elements to the Cartesian coordinates:
  for k=1:n_body
    for j=1:3, l=1:7*n_body
      jac_init[(i-1)*7+j,l] += ainv[i,k]*jac_kepler[(k-1)*6+j,l]
      jac_init[(i-1)*7+3+j,l] += ainv[i,k]*jac_kepler[(k-1)*6+3+j,l]
    end
  end
  # Include the mass derivatives of the A matrix:
#  dainvdm .= -ainv*damatdm[:,:,i]*ainv  # derivative with respect to the ith body
  # Multiply the derivative time Kepler positions/velocities to convert to
  # Cartesian coordinates:
  # Cartesian coordinates of all of the bodies can depend on the masses
  # of all the others so we need to loop over indices of each body, k:
  for k=1:n_body
#    dxdm = transpose(dainvdm[:,:,k]*rkepler)
    dxdm = permutedims(dainvdm[:,:,k]*rkepler)
#    dvdm = transpose(dainvdm[:,:,k]*rdotkepler)
    dvdm = permutedims(dainvdm[:,:,k]*rdotkepler)
    jac_init[(i-1)*7+1:(i-1)*7+3,k*7] += dxdm[1:3,i]
    jac_init[(i-1)*7+4:(i-1)*7+6,k*7] += dvdm[1:3,i]
  end
  # Masses are conserved:
  jac_init[i*7,i*7] = 1.0
end

return x,v
end

function get_indices_planetary(n_body)
# This sets up a planetary-hierarchy index matrix
indices = zeros(Int64,n_body,n_body)
for i=1:n_body-1
 for j=1:i
  indices[i,j ]=-1
 end
 indices[i,i+1]= 1
 indices[n_body,i]=1
end
indices[n_body,n_body]=1
# This is an example for TRAPPIST-1
# indices = [[-1, 1, 0, 0, 0, 0, 0, 0],  # first two bodies orbit in a binary
#            [-1,-1, 1, 0, 0, 0, 0, 0],  # next planet orbits about these
#            [-1,-1,-1, 1, 0, 0, 0, 0],  # etc...
#            [-1,-1,-1,-1, 1, 0, 0, 0],
#            [-1,-1,-1,-1,-1, 1, 0, 0],
#            [-1,-1,-1,-1,-1,-1, 1, 0],
#            [-1,-1,-1,-1,-1,-1,-1, 1],
#            [ 1, 1, 1, 1, 1, 1, 1, 1]  # center of mass of the system
return indices
end


text:
export dis2spt,
       spt2dis,
       spt2wlet,
       power,
       softThresh,
       fista,
       alterMini,
       alterCG,
       CG_w,
       CG_dis,
       zeros,
       ones,
       dotSrcSpt,
       AddSrcSpt,
       CGjoint,
       CGjoint_GP,
       GPW,
       obtainS,
       CG_GP

include("dis2spt.jl")
include("fista.jl")
include("CGLS.jl")
include("CGjoint.jl")

max_stars_repo_name: WneleiGao/AcousticWaveLab

max_issues_repo_name: WneleiGao/AcousticWaveLab

max_forks_repo_name: WneleiGao/AcousticWaveLab


text:
# computes 2^k exp(x) - 1 where |x| < ln(2)/2
# this is simply a translation of the last bit of openlib/src/s_expm1.c
# the only change is the handling of k == 0 to handle x.lo != 0

function expm1_kernel(k::Int, x::Double64)
    hfx = 0.5*x.hi
    hxs = x.hi*hfx
    r1  = @horner(hxs,1.0,
                 -3.33333333333331316428e-02,
                 1.58730158725481460165e-03,
                 -7.93650757867487942473e-05,
                 4.00821782732936239552e-06,
                 -2.01099218183624371326e-07)

    t  = 3.0-r1*hfx
    d  = hxs*((r1-t)/(6.0 - x.hi*t))
    e  = (x.hi*(d-x.lo)-x.lo) - hxs

    k ==  0 && return x.hi-e
    k == -1 && return 0.5*(x.hi-e)-0.5
    k ==  1 && return x.hi < -0.25 ? -2.0*(e-(x.hi+0.5)) :
                                     1.0+2.0*(x.hi-e)

    twopk = ldexp(1.0,k)
    if k <= -2 || k > 56 # suffice to return exp(x)-1
	y = 1.0-(e-x.hi)
	s = k == 1024 ? (y*2.0)*0x1p1023 : y*twopk
	return s-1.0
    end

    if k < 20
        t = 1.0-ldexp(1.0,-k)
	y = (t-(e-x.hi))*twopk
    else
        t = ldexp(1.0,-k)
	y = ((x.hi-(e+t))+1.0)*twopk
    end
    y
end


text:
"""
$(TYPEDEF)

A system of stochastic differential equations.

# Fields
$(FIELDS)

# Example

```
using ModelingToolkit

@parameters t σ ρ β
@variables x(t) y(t) z(t)
D = Differential(t)

eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]

noiseeqs = [0.1*x,
            0.1*y,
            0.1*z]

de = SDESystem(eqs,noiseeqs,t,[x,y,z],[σ,ρ,β])
```
"""
struct SDESystem <: AbstractODESystem
    """The expressions defining the drift term."""
    eqs::Vector{Equation}
    """The expressions defining the diffusion term."""
    noiseeqs::AbstractArray
    """Independent variable."""
    iv::Sym
    """Dependent (state) variables."""
    states::Vector
    """Parameter variables."""
    ps::Vector
    observed::Vector
    """
    Time-derivative matrix. Note: this field will not be defined until
    [`calculate_tgrad`](@ref) is called on the system.
    """
    tgrad::RefValue
    """
    Jacobian matrix. Note: this field will not be defined until
    [`calculate_jacobian`](@ref) is called on the system.
    """
    jac::RefValue
    """
    `Wfact` matrix. Note: this field will not be defined until
    [`generate_factorized_W`](@ref) is called on the system.
    """
    Wfact::RefValue
    """
    `Wfact_t` matrix. Note: this field will not be defined until
    [`generate_factorized_W`](@ref) is called on the system.
    """
    Wfact_t::RefValue
    """
    Name: the name of the system
    """
    name::Symbol
    """
    Systems: the internal systems
    """
    systems::Vector{SDESystem}
    """
    default_u0: The default initial conditions to use when initial conditions
    are not supplied in `ODEProblem`.
    """
    default_u0::Dict
    """
    default_p: The default parameters to use when parameters are not supplied
    in `ODEProblem`.
    """
    default_p::Dict
end

function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
                   observed = [],
                   systems = SDESystem[],
                   default_u0=Dict(),
                   default_p=Dict(),
                   name = gensym(:SDESystem))
    iv′ = value(iv)
    dvs′ = value.(dvs)
    ps′ = value.(ps)

    default_u0 isa Dict || (default_u0 = Dict(default_u0))
    default_p isa Dict || (default_p = Dict(default_p))
    default_u0 = Dict(value(k) => value(default_u0[k]) for k in keys(default_u0))
    default_p = Dict(value(k) => value(default_p[k]) for k in keys(default_p))

    tgrad = RefValue(Vector{Num}(undef, 0))
    jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
    Wfact   = RefValue(Matrix{Num}(undef, 0, 0))
    Wfact_t = RefValue(Matrix{Num}(undef, 0, 0))
    SDESystem(deqs, neqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, default_u0, default_p)
end

function generate_diffusion_function(sys::SDESystem, dvs = states(sys), ps = parameters(sys); kwargs...)
    return build_function(get_noiseeqs(sys),
                          map(x->time_varying_as_func(value(x), sys), dvs),
                          map(x->time_varying_as_func(value(x), sys), ps),
                          independent_variable(sys); kwargs...)
end

"""
$(TYPEDSIGNATURES)

Choose correction_factor=-1//2 (1//2) to converte Ito -> Stratonovich (Stratonovich->Ito).
"""
function stochastic_integral_transform(sys::SDESystem, correction_factor)
    # use the general interface
    if typeof(get_noiseeqs(sys)) <: Vector
        eqs = vcat([equations(sys)[i].lhs ~ get_noiseeqs(sys)[i] for i in eachindex(states(sys))]...)
        de = ODESystem(eqs,get_iv(sys),states(sys),parameters(sys))

        jac = calculate_jacobian(de, sparse=false, simplify=false)
        ∇σσ′ = simplify.(jac*get_noiseeqs(sys))

        deqs = vcat([equations(sys)[i].lhs ~ equations(sys)[i].rhs+ correction_factor*∇σσ′[i] for i in eachindex(states(sys))]...)
    else
        dimstate, m = size(get_noiseeqs(sys))
        eqs = vcat([equations(sys)[i].lhs ~ get_noiseeqs(sys)[i] for i in eachindex(states(sys))]...)
        de = ODESystem(eqs,get_iv(sys),states(sys),parameters(sys))

        jac = calculate_jacobian(de, sparse=false, simplify=false)
        ∇σσ′ = simplify.(jac*get_noiseeqs(sys)[:,1])
        for k = 2:m
            eqs = vcat([equations(sys)[i].lhs ~ get_noiseeqs(sys)[Int(i+(k-1)*dimstate)] for i in eachindex(states(sys))]...)
            de = ODESystem(eqs,get_iv(sys),states(sys),parameters(sys))

            jac = calculate_jacobian(de, sparse=false, simplify=false)
            ∇σσ′ = ∇σσ′ + simplify.(jac*get_noiseeqs(sys)[:,k])
        end

        deqs = vcat([equations(sys)[i].lhs ~ equations(sys)[i].rhs + correction_factor*∇σσ′[i] for i in eachindex(states(sys))]...)
    end

    SDESystem(deqs,get_noiseeqs(sys),get_iv(sys),states(sys),parameters(sys))
end

"""
```julia
function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = sys.states, ps = sys.ps;
                                     version = nothing, tgrad=false, sparse = false,
                                     jac = false, Wfact = false, kwargs...) where {iip}
```

Create an `SDEFunction` from the [`SDESystem`](@ref). The arguments `dvs` and `ps`
are used to set the order of the dependent variable and parameter vectors,
respectively.
"""
function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = states(sys), ps = parameters(sys),
                                     u0 = nothing;
                                     version = nothing, tgrad=false, sparse = false,
                                     jac = false, Wfact = false, eval_expression = true, kwargs...) where {iip}
    f_gen = generate_function(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
    f_oop,f_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in f_gen) : f_gen
    g_gen = generate_diffusion_function(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
    g_oop,g_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in g_gen) : g_gen

    f(u,p,t) = f_oop(u,p,t)
    f(du,u,p,t) = f_iip(du,u,p,t)
    g(u,p,t) = g_oop(u,p,t)
    g(du,u,p,t) = g_iip(du,u,p,t)

    if tgrad
        tgrad_gen = generate_tgrad(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
        tgrad_oop,tgrad_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in tgrad_gen) : tgrad_gen
        _tgrad(u,p,t) = tgrad_oop(u,p,t)
        _tgrad(J,u,p,t) = tgrad_iip(J,u,p,t)
    else
        _tgrad = nothing
    end

    if jac
        jac_gen = generate_jacobian(sys, dvs, ps; expression=Val{eval_expression}, sparse=sparse, kwargs...)
        jac_oop,jac_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in jac_gen) : jac_gen
        _jac(u,p,t) = jac_oop(u,p,t)
        _jac(J,u,p,t) = jac_iip(J,u,p,t)
    else
        _jac = nothing
    end

    if Wfact
        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true; expression=Val{true}, kwargs...)
        Wfact_oop, Wfact_iip = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in tmp_Wfact) : tmp_Wfact
        Wfact_oop_t, Wfact_iip_t = eval_expression ? (@RuntimeGeneratedFunction(ex) for ex in tmp_Wfact_t) : tmp_Wfact_t
        _Wfact(u,p,dtgamma,t) = Wfact_oop(u,p,dtgamma,t)
        _Wfact(W,u,p,dtgamma,t) = Wfact_iip(W,u,p,dtgamma,t)
        _Wfact_t(u,p,dtgamma,t) = Wfact_oop_t(u,p,dtgamma,t)
        _Wfact_t(W,u,p,dtgamma,t) = Wfact_iip_t(W,u,p,dtgamma,t)
    else
        _Wfact,_Wfact_t = nothing,nothing
    end

    M = calculate_massmatrix(sys)
    _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0',M)

    sts = states(sys)
    SDEFunction{iip}(f,g,
                     jac = _jac === nothing ? nothing : _jac,
                     tgrad = _tgrad === nothing ? nothing : _tgrad,
                     Wfact = _Wfact === nothing ? nothing : _Wfact,
                     Wfact_t = _Wfact_t === nothing ? nothing : _Wfact_t,
                     mass_matrix = _M,
                     syms = Symbol.(states(sys)))
end

function DiffEqBase.SDEFunction(sys::SDESystem, args...; kwargs...)
    SDEFunction{true}(sys, args...; kwargs...)
end

"""
```julia
function DiffEqBase.SDEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
                                     ps = parameters(sys);
                                     version = nothing, tgrad=false,
                                     jac = false, Wfact = false,
                                     skipzeros = true, fillzeros = true,
                                     sparse = false,
                                     kwargs...) where {iip}
```

Create a Julia expression for an `SDEFunction` from the [`SDESystem`](@ref).
The arguments `dvs` and `ps` are used to set the order of the dependent
variable and parameter vectors, respectively.
"""
struct SDEFunctionExpr{iip} end

function SDEFunctionExpr{iip}(sys::SDESystem, dvs = states(sys),
                                     ps = parameters(sys), u0 = nothing;
                                     version = nothing, tgrad=false,
                                     jac = false, Wfact = false,
                                     sparse = false,linenumbers = false,
                                     kwargs...) where {iip}

    idx = iip ? 2 : 1
    f = generate_function(sys, dvs, ps; expression=Val{true}, kwargs...)[idx]
    g = generate_diffusion_function(sys, dvs, ps; expression=Val{true}, kwargs...)[idx]
    if tgrad
        _tgrad = generate_tgrad(sys, dvs, ps; expression=Val{true}, kwargs...)[idx]
    else
        _tgrad = :nothing
    end

    if jac
        _jac = generate_jacobian(sys, dvs, ps; sparse = sparse, expression=Val{true}, kwargs...)[idx]
    else
        _jac = :nothing
    end

    if Wfact
        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps; expression=Val{true}, kwargs...)
        _Wfact =  tmp_Wfact[idx]
        _Wfact_t =  tmp_Wfact_t[idx]
    else
        _Wfact,_Wfact_t = :nothing,:nothing
    end

    M = calculate_massmatrix(sys)

    _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0',M)

    ex = quote
        f = $f
        g = $g
        tgrad = $_tgrad
        jac = $_jac
        Wfact = $_Wfact
        Wfact_t = $_Wfact_t
        M = $_M
        SDEFunction{$iip}(f,g,
                         jac = jac,
                         tgrad = tgrad,
                         Wfact = Wfact,
                         Wfact_t = Wfact_t,
                         mass_matrix = M,
                         syms = $(Symbol.(states(sys))))
    end
    !linenumbers ? striplines(ex) : ex
end

function SDEFunctionExpr(sys::SDESystem, args...; kwargs...)
    SDEFunctionExpr{true}(sys, args...; kwargs...)
end

function rename(sys::SDESystem,name)
    SDESystem(sys.eqs, sys.noiseeqs, sys.iv, sys.states, sys.ps, sys.tgrad, sys.jac, sys.Wfact, sys.Wfact_t, name, sys.systems)
end

"""
```julia
function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,p=parammap;
                                    version = nothing, tgrad=false,
                                    jac = false, Wfact = false,
                                    checkbounds = false, sparse = false,
                                    sparsenoise = sparse,
                                    skipzeros = true, fillzeros = true,
                                    linenumbers = true, parallel=SerialForm(),
                                    kwargs...)
```

Generates an SDEProblem from an SDESystem and allows for automatically
symbolically calculating numerical enhancements.
"""
function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,parammap=DiffEqBase.NullParameters();
                                    sparsenoise = nothing,
                                    kwargs...) where iip
    f, u0, p = process_DEProblem(SDEFunction{iip}, sys, u0map, parammap; kwargs...)
    sparsenoise === nothing && (sparsenoise = get(kwargs, :sparse, false))

    noiseeqs = get_noiseeqs(sys)
    if noiseeqs isa AbstractVector
        noise_rate_prototype = nothing
    elseif sparsenoise
        I,J,V = findnz(SparseArrays.sparse(noiseeqs))
        noise_rate_prototype = SparseArrays.sparse(I,J,zero(eltype(u0)))
    else
        noise_rate_prototype = zeros(eltype(u0),size(noiseeqs))
    end

    SDEProblem{iip}(f,f.g,u0,tspan,p;noise_rate_prototype=noise_rate_prototype,kwargs...)
end

function DiffEqBase.SDEProblem(sys::SDESystem, args...; kwargs...)
    SDEProblem{true}(sys, args...; kwargs...)
end

"""
```julia
function DiffEqBase.SDEProblemExpr{iip}(sys::AbstractODESystem,u0map,tspan,
                                    parammap=DiffEqBase.NullParameters();
                                    version = nothing, tgrad=false,
                                    jac = false, Wfact = false,
                                    checkbounds = false, sparse = false,
                                    linenumbers = true, parallel=SerialForm(),
                                    kwargs...) where iip
```

Generates a Julia expression for constructing an ODEProblem from an
ODESystem and allows for automatically symbolically calculating
numerical enhancements.
"""
struct SDEProblemExpr{iip} end

function SDEProblemExpr{iip}(sys::SDESystem,u0map,tspan,
                             parammap=DiffEqBase.NullParameters();
                             sparsenoise = nothing,
                             kwargs...) where iip
    f, u0, p = process_DEProblem(SDEFunctionExpr{iip}, sys, u0map, parammap; kwargs...)
    linenumbers = get(kwargs, :linenumbers, true)
    sparsenoise === nothing && (sparsenoise = get(kwargs, :sparse, false))

    noiseeqs = get_noiseeqs(sys)
    if noiseeqs isa AbstractVector
        noise_rate_prototype = nothing
    elseif sparsenoise
        I,J,V = findnz(SparseArrays.sparse(noiseeqs))
        noise_rate_prototype = SparseArrays.sparse(I,J,zero(eltype(u0)))
    else
        T = u0 === nothing ? Float64 : eltype(u0)
        noise_rate_prototype = zeros(T,size(get_noiseeqs(sys)))
    end
    ex = quote
        f = $f
        u0 = $u0
        tspan = $tspan
        p = $p
        noise_rate_prototype = $noise_rate_prototype
        SDEProblem(f,f.g,u0,tspan,p;noise_rate_prototype=noise_rate_prototype,$(kwargs...))
    end
    !linenumbers ? striplines(ex) : ex
end

function SDEProblemExpr(sys::SDESystem, args...; kwargs...)
    SDEProblemExpr{true}(sys, args...; kwargs...)
end


text:
struct AnalyticalMapCubedSphereTriangulation{T} <: Triangulation{2,3}
  cell_map::T
  btrian::Triangulation{2,3}
end

# Triangulation API

# Delegating to the underlying face Triangulation

Gridap.Geometry.get_cell_coordinates(trian::AnalyticalMapCubedSphereTriangulation) = Gridap.Geometry.get_cell_coordinates(trian.btrian)

Gridap.Geometry.get_reffes(trian::AnalyticalMapCubedSphereTriangulation) = Gridap.Geometry.get_reffes(trian.btrian)

Gridap.Geometry.get_cell_type(trian::AnalyticalMapCubedSphereTriangulation) = Gridap.Geometry.get_cell_type(trian.btrian)

Gridap.Geometry.get_node_coordinates(trian::AnalyticalMapCubedSphereTriangulation) = Gridap.Helpers.get_node_coordinates(trian.btrian)

Gridap.Geometry.get_cell_node_ids(trian::AnalyticalMapCubedSphereTriangulation) = Gridap.Helpers.get_cell_node_ids(trian.btrian)

Gridap.Geometry.get_cell_map(trian::AnalyticalMapCubedSphereTriangulation) = trian.cell_map

# Genuine methods

Gridap.Geometry.TriangulationStyle(::Type{<:AnalyticalMapCubedSphereTriangulation}) = SubTriangulation()

Gridap.Geometry.get_background_triangulation(trian::AnalyticalMapCubedSphereTriangulation) =
      Gridap.Geometry.get_background_triangulation(trian.btrian)

Gridap.Geometry.get_cell_to_bgcell(trian::AnalyticalMapCubedSphereTriangulation) = get_cell_to_bgcell(trian.btrian)

function Gridap.Geometry.get_cell_to_bgcell(
  trian_in::AnalyticalMapCubedSphereTriangulation,
  trian_out::AnalyticalMapCubedSphereTriangulation)
  Gridap.Helpers.@notimplemented
end

function Gridap.Geometry.is_included(
  trian_in::AnalyticalMapCubedSphereTriangulation,
  trian_out::AnalyticalMapCubedSphereTriangulation)
  Gridap.Helpers.@notimplemented
end

function Gridap.Geometry.get_cell_ref_map(trian::AnalyticalMapCubedSphereTriangulation)
  get_cell_ref_map(trian.btrian)
end

function Gridap.Geometry.get_cell_ref_map(
  trian_in::AnalyticalMapCubedSphereTriangulation,
  trian_out::AnalyticalMapCubedSphereTriangulation)
  Gridap.Helpers.@notimplemented
end


text:
include("imports.jl")

include("squared_euclidean.jl")


text:
function plotfdistr(rs, f; kwargs...)
    m = mean(f)
    v = sqrt.(var(f))
    plot(;kwargs...)
    plot!(rs, m, label="μ")
    plot!(rs, m + v, label="μ + σ")
    plot!(rs, m - v, label="μ - σ")
end

plotfdistr(f; kwargs...) = plotfdistr(length(f[1]), f; kwargs...)

plot_point_process(ϕ; label=nothing, xlims=(-0.1, 1.1), ylims=(-0.1, 1.1), kwargs...) =
    scatter(Tuple.(ϕ); label, xlims, ylims, kwargs...)


text:
#JuliaRL/src/show_stats.jl

using BSON
using PyPlot

"""
    show_stats(save_dir)

Plots the stats saved in the `stats.bson` file inside of `save_dir`.
Stats are total reward per episode, and time per step.
"""
function show_stats(save_dir::T) where {T<:AbstractString}
    BSON.@load joinpath(save_dir,"stats.bson") total_reward_per_episode time_per_step

    figure(figsize=(10, 5))
    subplot(121)
    title(save_dir)
    ylabel("Total reward")
    xlabel("Episode")
    plot(total_reward_per_episode.rewards)

    subplot(122)
    ylabel("Time")
    xlabel("Step")
    x = (1:length(time_per_step.times)) * 100
    y = time_per_step.times / 100
end

max_stars_repo_name: FeLusiani/JuliaRL

max_issues_repo_name: FeLusiani/JuliaRL

max_forks_repo_name: FeLusiani/JuliaRL