text:
module RegularExpressions

using Base: SubstitutionString, Generator
using Base.Iterators: Filter

struct Not
end
"""
    not

Use to negate a pattern.

Use with [`short`](@ref), [`option`](@ref), [`class`](@ref), [`one_of`](@ref), [`property`](@ref), or [`script`](@ref).
"""
const not = Not()
export not

function group(them...)
    joined = join(them)
    if length(joined) == 1
        joined
    else
        kind(:group, joined)
    end
end

"""
    raw(it)

Escape punctuation.

```jldoctest
julia> using RegularExpressions

julia> p = pattern(raw("1.0"))
r"1\\.0"

julia> occursin(p, "v1.0")
true
```
"""
function raw(it)
    replace(it, r"([[:punct:]])" => s"\\\g<1>")
end
export raw

"""
    CONSTANTS

Plain commands.

```jldoctest
julia> using RegularExpressions

julia> p = pattern(CONSTANTS.any)
r"."

julia> occursin(p, "a")
true
```
"""
const CONSTANTS = (
    any = ".",
    start = "^",
    stop = "\$",
    define = "DEFINE",
    recur = "(?R)",
    recurred = "R"
)
export CONSTANTS

"""
    SHORTS

Access with [`short`](@ref).
"""
const SHORTS = (
    unit = "C",
    digit = "d",
    horizontal = "h",
    newline = "n",
    newlines = "R",
    space = "s",
    vertical = "v",
    word = "w",
    cluster = "X",
    boundary = "b",
    strict_start = "A",
    strict_stop = "z",
    chomp_stop = "Z",
    match_start = "G",
    reset = "K"
)
export SHORTS
"""
    short([::Not], it)

A short command. Access [`SHORTS`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(short(:space))
r"\\s"

julia> occursin(p, " ")
true

julia> p = pattern(short(not, :space))
r"\\S"

julia> occursin(p, "a")
true
```
"""
short(it) = "\\$(SHORTS[it])"
short(::Not, it) = "\\$(uppercase(SHORTS[it]))"
export short

"""
    PROPERTIES

Access with [`property`](@ref).
"""
const PROPERTIES = (
    other = ("C", (
        control = "c",
        format = "f",
        unassigned = "n",
        private = "o",
        surrogate = "s"
    )),
    letter = ("L", (
        lowercase = "l",
        modifier = "m",
        other = "o",
        titlecase = "t",
        uppercase = "u",
        cased = "&"
    )),
    mark = ("M", (
        spacing = "c",
        enclosing = "e",
        non_spacing = "n"
    )),
    number = ("N", (
        decimal = "d",
        letter = "l",
        other = "o"
    )),
    punctuation = ("P", (
        connector = "c",
        dash = "d",
        close = "e",
        final = "f",
        initial = "i",
        other = "o",
        open = "s"
    )),
    symbol = ("S", (
        currency = "c",
        modifier = "k",
        math = "m",
        other = "o"
    )),
    seperator = ("Z", (
        line = "l",
        paragraph = "p",
        space = "s"
    )),
    special = ("X", (
        letter_or_digit = "an",
        space = "sp",
        exists = "uc",
        word = "wd"
    ))
)
export PROPERTIES
"""
    property([::Not], general, [specific])

A character property. Access [`PROPERTIES`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(property(:seperator))
r"\\p{Z}"

julia> occursin(p, " ")
true

julia> p = pattern(property(not, :seperator))
r"\\P{Z}"

julia> occursin(p, "a")
true

julia> p = pattern(property(:seperator, :space))
r"\\p{Zs}"

julia> occursin(p, " ")
true

julia> p = pattern(property(not, :seperator, :space))
r"\\P{Zs}"

julia> occursin(p, "a")
true
```
"""
property(general) = "\\p{$(PROPERTIES[general][1])}"
property(::Not, general) = "\\P{$(PROPERTIES[general][1])}"
function specific_property(general, specific)
    general_code, specifics = PROPERTIES[general]
    "$(general_code)$(specifics[specific])"
end
property(general, specific) = "\\p{$(specific_property(general, specific))}"
property(::Not, general, specific) =
    "\\P{$(specific_property(general, specific))}"
export property

"""
    script([::Not], it)

A character from a script.

```jldoctest
julia> using RegularExpressions

julia> p = pattern(script(:Han))
r"\\p{Han}"

julia> occursin(p, "中")
true

julia> p = pattern(script(not, :Han))
r"\\P{Han}"

julia> occursin(p, "a")
true
```
"""
script(it) = "\\p{$it}"
script(::Not, it) = "\\P{$it}"
export script

"""
    one_of([::Not], them...)

Create a character class.

```jldoctest
julia> using RegularExpressions

julia> p = pattern(one_of('a', 'b'))
r"[ab]"

julia> occursin(p, "b")
true

julia> p = pattern(one_of(not, 'a', 'b'))
r"[^ab]"

julia> occursin(p, "c")
true
```
"""
one_of(them...) = "[$(them...)]"
one_of(::Not, them...) = "[^$(them...)]"
export one_of

"""
    through(start, stop)

A range of characters

```jldoctest
julia> using RegularExpressions

julia> p = pattern(one_of(through('a', 'c')))
r"[a-c]"

julia> occursin(p, "b")
true
```
"""
through(start, stop) = "$start-$stop"
export through

"""
    CLASSES

Access with [`class`](@ref).
"""
const CLASSES = (
    letter_or_digit = "alnum",
    letter = "alpha",
    standard = "ascii",
    blank = "blank",
    control = "cntrl",
    digit = "digit",
    prints = "graph",
    lowercase = "lower",
    prints_or_space = "print",
    punctuation = "punct",
    space = "space",
    uppercase = "upper",
    word = "word",
    hexadecimal = "xdigit"
)
export CLASSES
"""
    class([::Not], it)

Character classes. Access [`CLASSES`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(one_of(class(:space)))
r"[[:space:]]"

julia> occursin(p, " ")
true

julia> p = pattern(one_of(class(not, :space)))
r"[[:^space:]]"

julia> occursin(p, "a")
true
```
"""
class(it) = "[:$(CLASSES[it]):]"
class(::Not, it) = "[:^$(CLASSES[it]):]"
export class

"""
    GREEDS

Access with [`of`](@ref) or [`between`](@ref).
"""
const GREEDS = (
    possessive = "+",
    lazy = "?",
    greedy = ""
)
export GREEDS
"""
    QUANTITIES

Access with [`of`](@ref).
"""
const QUANTITIES = (
    maybe = "?",
    none_or_some = "*",
    some = "+"
)
export QUANTITIES
"""
    of(quantity::Symbol, them...; greed = :greedy)
    of(quantity::Number, them...)

A `quantity` `of` `it` with a certain `greed`. Acccess [`QUANTITIES`](@ref) and [`GREEDS`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(of(:some, "a"))
r"a+"

julia> occursin(p, "aa")
true

julia> p = pattern(of(2, "a"))
r"a{2}"

julia> occursin(p, "aa")
true
```
"""
of(quantity::Symbol, them...; greed = :greedy) =
    "$(group(them...))$(QUANTITIES[quantity])$(GREEDS[greed])"
of(quantity::Number, them...) = "$(group(them...)){$quantity}"
export of

"""
    between(low, high, them...; greed = :greedy)

Between `low` and `high` of `it` with a certain `greed`. Access [`GREEDS`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(between(1, 3, "a"))
r"a{1,3}"

julia> occursin(p, "aa")
true

julia> p = pattern(between(2, Inf, "a"))
r"a{2,}"

julia> occursin(p, "aaa")
true
```
"""
between(low, high, them...; greed = :greedy) =
    if isinf(high)
        "$(group(them...)){$low,}$(GREEDS[greed])"
    else
        "$(group(them...)){$low,$high}$(GREEDS[greed])"
    end
export between

"""
    or(them...)

One of `them`.

```jldoctest
julia> using RegularExpressions

julia> p = pattern(or("a", "b"))
r"a|b"

julia> occursin(p, "b")
true
```
"""
or(them...) = join(them, "|")
export or

"""
    capture(them...; name = nothing)

Capture `them` with optional `name`. See examples in [`captured`](@ref).
"""
capture(them...; name = nothing) =
    if name === nothing
        "($(them...))"
    else
        "(?<$name>$(them...))"
    end
export capture

"""
    KINDS

Access via [`kind`](@ref).
"""
const KINDS = (
    group = ":",
    reset = "|",
    atomic = ">",
    comment = "#",
    after = "=",
    not_after = "!",
    before = "<=",
    not_before = "<!",
    callout = "C"
)

export KINDS
"""
    kind(a_kind, them...)

A variety of syntaxes: `a_kind` of `them`. Access [`KINDS`](@ref). Use `repr` to
pass strings to callouts.

```jldoctest
julia> using RegularExpressions

julia> p = pattern(kind(:before, "a"), "b")
r"(?<=a)b"

julia> occursin(p, "ab")
true
```
"""
kind(kind, them...) = "(?$(KINDS[kind])$(them...))"
export kind

"""
    OPTIONS

Access with [`option`](@ref)
"""
const OPTIONS = (
    caseless = "i",
    duplicate_names = "J",
    multi_line = "m",
    no_auto_capture = "n",
    single_line = "s",
    lazy = "U",
    ignore_space = "x",
    ignore_all_space = "xx",
    unset = "^",
    recur = "R",
    callout = "C"
)
export OPTIONS
make_flags(options, OPTIONS) =
    join(Generator(
        pair -> OPTIONS[pair[1]],
        Filter(pair -> pair[2], options)
    ))
"""
    option([::Not]; options...)

`option`. Access [`OPTIONS`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(option(caseless = true, ignore_space = true), "a ")
r"(?ix)a "

julia> occursin(p, "A")
true

julia> p = pattern(option(caseless = true), option(not, caseless = true), "a")
r"(?i)(?-i)a"

julia> occursin(p, "A")
false
```
"""
option(; options...) = "(?$(make_flags(options, OPTIONS)))"
option(::Not; options...) = "(?-$(make_flags(options, OPTIONS)))"
export option

"""
    EXTRAS

Access with [`extra`](@ref).
"""
const EXTRAS = (
    limit_depth = "LIMIT_DEPTH",
    limit_heap = "LIMIT_HEAP",
    limit_match = "LIMIT_MATCH",
    not_empty = "NOTEMPTY",
    not_empty_at_start = "NOTEMPTY_ATSTART",
    no_auto_possess = "NO_AUTO_POSSESS",
    no_dot_star_anchor = "NO_DOTSTAR_ANCHOR",
    no_just_in_time = "NO_JIT",
    no_start_optimization = "NO_START_OPT",
    UTF = "UTF",
    unicode_properties = "UCP",
    carriage_return = "CR",
    linefeed = "LF",
    carriage_return_linefeed = "CRLF",
    standard_newline = "ANYCRLF",
    unicode_newline = "ANY",
    nul = "NUL",
    standard_boundary = "BSR_ANYCRLF",
    unicode_boundary = "BSR_UNICODE",
    accept = "ACCEPT",
    fail = "FAIL",
    mark = "MARK",
    commit = "COMMIT",
    prune = "PRUNE",
    skip = "SKIP",
    then = "THEN"
)
export EXTRAS
"""
    extra(it)
    extra(it, name)
    extra(it, value::Number)

`extra` command. Access [`EXTRAS`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(extra(:standard_newline), "a")
r"(*ANYCRLF)a"

julia> occursin(p, "a\\r")
true

julia> extra(:limit_match, 0)
"(*LIMIT_MATCH=0)"

julia> extra(:mark, "name")
"(*MARK:name)"
```
"""
extra(it) = "(*$(EXTRAS[it]))"
extra(it, name) = "(*$(EXTRAS[it]):$name)"
extra(it, value::Number) = "(*$(EXTRAS[it])=$value)"
export extra

"""
    relative(it)

Mark a reference as relative. For use with [`captured`](@ref) or [`whether`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(captured(relative(1)), capture("a"))
r"\\g<+1>(a)"

julia> occursin(p, "aa")
true

julia> p = pattern(capture("a"), captured(relative(-1)))
r"(a)\\g<-1>"

julia> occursin(p, "aa")
true
```
"""
relative(it) =
    if it < 0
        it
    else
        "+$it"
    end
export relative

"""
    captured(it)

Refer to a [`capture`](@ref)d group. See [`relative`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(capture("a"), captured(1))
r"(a)\\g<1>"

julia> occursin(p, "aa")
true
```
"""
captured(it) = "\\g<$it>"
export captured

"""
    pattern(them..., options...)

Splat of `Regex`. Options can be in [`OPTIONS`](@ref)

```jldoctest
julia> using RegularExpressions

julia> p = pattern("a", "b")
r"ab"

julia> occursin(p, "ab")
true

julia> p = pattern("A", caseless = true)
r"A"i

julia> occursin(p, "a")
true
```
"""
pattern(them...; options...) =
    Regex(string(them...), make_flags(options, OPTIONS))
export pattern

"""
    exists(it)

Check whether a capture group. For use with [`whether`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(
            CONSTANTS.start,
            of(:maybe, capture("a", name = "first")),
            whether(exists("first"), "b", "c")
        )
r"^(?:(?<first>a))?(?(<first>)b|c)"

julia> occursin(p, "ab")
true
```
"""
exists(it) = "<$it>"
export exists

"""
    template(them...)

Splat of `SubstitutionString`.

```jldoctest
julia> using RegularExpressions

julia> p = pattern(capture("a"))
r"(a)"

julia> t = template(captured(1), "b")
s"\\\\g<1>b"

julia> replace("a", p => t)
"ab"
```
"""
template(them...) = SubstitutionString(string(them...))
export template

"""
    version(it; at_least = false)

Check whether the version of PCRE2 is `it`, (or, `at_least` `it`). For use with [`whether`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(whether(version(1), "new", "old"))
r"(?(VERSION=1)new|old)"

julia> occursin(p, "new")
false

julia> p = pattern(whether(version(1, at_least = true), "new", "old"))
r"(?(VERSION>=1)new|old)"

julia> occursin(p, "new")
true
```
"""
version(it; at_least = false) = "VERSION$(if at_least
        ">"
    else
        ""
    end)=$it"
export version

"""
    whether(condition, yes, no = "")

Test for a condition. See [`relative`](@ref), [`exists`](@ref), [`recurred`](@ref), and [`version`](@ref).

```jldoctest
julia> using RegularExpressions

julia> p = pattern(
            CONSTANTS.start,
            of(:maybe, capture("a")),
            whether(1, "b", "c")
        )
r"^(?:(a))?(?(1)b|c)"

julia> occursin(p, "ab")
true
```
"""
whether(condition, yes, no = "") = "(?($condition)$yes|$no)"
export whether

"""
    recurred(it::Number)
    recurred(it)

Check for recursion. Use with [`whether`](@ref).

```jldoctest
julia> using RegularExpressions

julia> recurred(1)
"R1"

julia> recurred("name")
"R&name"
```
"""
recurred(it::Number) = "R$it"
recurred(it) = "R&$it"
export recurred

end


text:
# Contains all streams in simulation. This way we can step each stream on each pass
# The first design of the simulation only has one cell so the cell output isn't included.
const SEED_SCALER = 10000.0
# For debugging or learning we usually want a known sequence.
const RESET_SEED = 13163

mutable struct Streams
    bit_streams::Array{AbstractBitStream,1}
    stim_streams::Array{AbstractBitStream,1}
    merger_streams::Array{AbstractBitStream,1}

    id::Int64
    stim_ids::Int64

    function Streams()
        o = new()
        o.id = 1
        o.stim_ids = 1
        o
    end
end

function id(stream::AbstractBitStream)
    stream.base.id
end

function add_stream!(streams::Streams, stream::AbstractBitStream)
    push!(streams.bit_streams, stream)

    # id is used to align to synapses. "id" can be thought of as
    # synapse-id.
    stream.base.id = streams.id

    streams.id += 1
end

function add_stimulus_stream!(streams::Streams, stream::AbstractBitStream)
    push!(streams.stim_streams, stream)

    # id is used to align to synapses. "id" can be thought of as
    # synapse-id.
    stream.base.id = streams.stim_ids

    streams.stim_ids += 1
end

function add_merger_stream!(streams::Streams, stream::AbstractBitStream)
    push!(streams.merger_streams, stream)
end

function config_streams!(streams::Streams)
    streams.bit_streams = Array{AbstractBitStream,1}()
    streams.stim_streams = Array{AbstractBitStream,1}()
    streams.merger_streams = Array{AbstractBitStream,1}()
end

function config_poi_streams!(streams::Streams, synapses::Int64, firing_rate::Float64)
    Random.seed!(RESET_SEED)

    # Create a stream for each synapse
    for syn_id in 1:synapses
        seed = Int64(round(rand(1)[1] * SEED_SCALER))

        stream = Simulation.PoissonStream(UInt64(seed), firing_rate)
        stream.base.id = syn_id
        add_stream!(streams, stream)
    end
end

function get_poisson_stream(streams::Streams, id::Int64)
    for stream in streams.bit_streams
        if stream.base.id == id
            return stream
        end
    end

    nothing
end

function get_stimulus_stream(streams::Streams, id::Int64)
    for stream in streams.stim_streams
        if stream.base.id == id
            return stream
        end
    end

    nothing
end

function config_stimulus_streams!(streams::Streams, model::Model.ModelData)
    # Load stimulus pattern
    pattern_file_prefix = Model.source_stimulus(model)

    working_dir = pwd()

    src_range = findlast("src/", working_dir)
    root_path = working_dir[1:src_range[end]]

    data_file = root_path * "data/" * pattern_file_prefix * ".data"

    println("Loading stimulus: ", data_file)

    # Create stream for loading
    stream = Simulation.RegularPatternStream()

    # Collect it for later iteration.
    add_stimulus_stream!(streams, stream)

    frequency = Model.hertz(model)
    expand_scaler = Model.stimulus_scaler(model)
    Simulation.load!(stream, data_file, frequency, expand_scaler)
end

function exercise!(streams::Streams)
    # After integration all streams need to be exercised.
    # This causes any internal state to move through their
    # respective activites.
    for stream in streams.bit_streams # Typically Poisson
        step!(stream)
    end

    for stream in streams.stim_streams # Typically Patterns
        step!(stream)
    end

    for stream in streams.merger_streams # Typically Mergers
        step!(stream)
    end
end

function collect!(streams::Streams, samples::Model.Samples, t::Int64)
    # Iterate all streams and collect data from stream outputs.

    # These streams are generally poission types
    for stream in streams.bit_streams
        Model.store_poi_sample!(samples, id(stream), t, output(stream))
    end

    # Interate through each stimulus stream and store the current output
    # into the samples collection.
    for stream in streams.stim_streams
        # Pattern streams have multiple internal streams, one for each target synapse.
        for syn_id in 1:stream.synapses
            Model.store_stimulus_sample!(samples, syn_id, t, output(stream, syn_id))
        end
    end
end

function collect!(samples::Model.Samples, soma::AbstractSoma, t::Int64)
    # Capture soma data
    Model.store_apFast_sample!(samples, soma.apFast, t)

    Model.store_apSlow_sample!(samples, soma.apSlow, t)

    Model.store_cell_sample!(samples, output(soma.axon), t)

    Model.store_soma_psp_sample!(samples, soma.psp, t)
end

function collect_synapse!(samples::Model.Samples, syn::AbstractSynapse, t::Int64)
    Model.store_syn_weight_sample!(samples, syn.id, syn.w, t)
    Model.store_syn_surge_sample!(samples, syn.id, syn.surge, t)
    Model.store_syn_psp_sample!(samples, syn.id, syn.psp, t)
    Model.store_syn_input_sample!(samples, syn.id, Float64(output(syn.stream)), t)
end

function collect_dendrite!(samples::Model.Samples, dendrite::AbstractDendrite, t::Int64)
    Model.store_dendrite_avg_sample!(samples, dendrite.average, t)
end

ext: jl


text:
module SystemStructures

using ..ModelingToolkit
import ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!, flatten
using SymbolicUtils: arguments
using ..BipartiteGraphs
using UnPack
using Setfield
using SparseArrays

#=
When we don't do subsitution, variable information is split into two different
places, i.e. `states` and the right-hand-side of `observed`.

eqs = [0 ~ z + x; 0 ~ y + z^2]
states = [y, z]
observed = [x ~ sin(y) + z]
struct Reduced
    var
    expr
    idxs
end
fullvars = [Reduced(x, sin(y) + z, [2, 3]), y, z]
active_𝑣vertices = [false, true, true]
      x   y   z
eq1:  1       1
eq2:      1   1

      x   y   z
eq1:      1   1
eq2:      1   1

for v in 𝑣vertices(graph); active_𝑣vertices[v] || continue

end
=#

export SystemStructure, initialize_system_structure
export diffvars_range, dervars_range, algvars_range
export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq
export DIFFERENTIAL_VARIABLE, ALGEBRAIC_VARIABLE, DERIVATIVE_VARIABLE
export DIFFERENTIAL_EQUATION, ALGEBRAIC_EQUATION
export vartype, eqtype

struct SystemStructure
    dxvar_offset::Int
    fullvars::Vector # [xvar; dxvars; algvars]
    varassoc::Vector{Int}
    algeqs::BitVector
    graph::BipartiteGraph{Int}
    solvable_graph::BipartiteGraph{Int}
    assign::Vector{Int}
    inv_assign::Vector{Int}
    scc::Vector{Vector{Int}}
    partitions::Vector{NTuple{4, Vector{Int}}}
end

diffvars_range(s::SystemStructure) = 1:s.dxvar_offset
dervars_range(s::SystemStructure) = s.dxvar_offset+1:2s.dxvar_offset
algvars_range(s::SystemStructure) = 2s.dxvar_offset+1:length(s.fullvars)

isdiffvar(s::SystemStructure, var::Integer) = var in diffvars_range(s)
isdervar(s::SystemStructure, var::Integer) = var in dervars_range(s)
isalgvar(s::SystemStructure, var::Integer) = var in algvars_range(s)

@enum VariableType DIFFERENTIAL_VARIABLE ALGEBRAIC_VARIABLE DERIVATIVE_VARIABLE

function vartype(s::SystemStructure, var::Integer)::VariableType
    isdiffvar(s, var) ? DIFFERENTIAL_VARIABLE :
    isdervar(s, var)  ? DERIVATIVE_VARIABLE :
    isalgvar(s, var)  ? ALGEBRAIC_VARIABLE : error("Variable $var out of bounds")
end

@enum EquationType DIFFERENTIAL_EQUATION ALGEBRAIC_EQUATION

isalgeq(s::SystemStructure, eq::Integer) = s.algeqs[eq]
isdiffeq(s::SystemStructure, eq::Integer) = !isalgeq(s, eq)
eqtype(s::SystemStructure, eq::Integer)::EquationType = isalgeq(s, eq) ? ALGEBRAIC_EQUATION : DIFFERENTIAL_EQUATION

function initialize_system_structure(sys)
    sys, dxvar_offset, fullvars, varassoc, algeqs, graph, solvable_graph = init_graph(flatten(sys))
    @set sys.structure = SystemStructure(
                                         dxvar_offset,
                                         fullvars,
                                         varassoc,
                                         algeqs,
                                         graph,
                                         solvable_graph,
                                         Int[],
                                         Int[],
                                         Vector{Int}[],
                                         NTuple{4, Vector{Int}}[]
                                        )
end

function Base.show(io::IO, s::SystemStructure)
    @unpack fullvars, dxvar_offset, solvable_graph, graph = s
    algvar_offset = 2dxvar_offset
    print(io, "xvars: ")
    print(io, fullvars[1:dxvar_offset])
    print(io, "\ndxvars: ")
    print(io, fullvars[dxvar_offset+1:algvar_offset])
    print(io, "\nalgvars: ")
    print(io, fullvars[algvar_offset+1:end], '\n')

    S = incidence_matrix(graph, Num(Sym{Real}(:×)))
    print(io, "Incidence matrix:")
    show(io, S)
end

# V-nodes `[x_1, x_2, x_3, ..., dx_1, dx_2, ..., y_1, y_2, ...]` where `x`s are
# differential variables and `y`s are algebraic variables.
function collect_variables(sys)
    dxvars = []
    eqs = equations(sys)
    algeqs = trues(length(eqs))
    for (i, eq) in enumerate(eqs)
        if isdiffeq(eq)
            algeqs[i] = false
            lhs = eq.lhs
            # Make sure that the LHS is a first order derivative of a var.
            @assert !(arguments(lhs)[1] isa Differential) "The equation $eq is not first order"

            push!(dxvars, lhs)
        end
    end

    xvars = (first ∘ var_from_nested_derivative).(dxvars)
    algvars  = setdiff(states(sys), xvars)
    return xvars, dxvars, algvars, algeqs
end

function init_graph(sys)
    xvars, dxvars, algvars, algeqs = collect_variables(sys)
    dxvar_offset = length(xvars)
    algvar_offset = 2dxvar_offset

    fullvars = [xvars; dxvars; algvars]
    eqs = equations(sys)
    idxmap = Dict(fullvars .=> 1:length(fullvars))
    graph = BipartiteGraph(length(eqs), length(fullvars))
    solvable_graph = BipartiteGraph(length(eqs), length(fullvars))

    vs = Set()
    for (i, eq) in enumerate(eqs)
        # TODO: custom vars that handles D(x)
        # TODO: add checks here
        lhs = eq.lhs
        if isdiffeq(eq)
            v = lhs
            haskey(idxmap, v) && add_edge!(graph, i, idxmap[v])
        else
            vars!(vs, lhs)
        end
        vars!(vs, eq.rhs)
        for v in vs
            haskey(idxmap, v) && add_edge!(graph, i, idxmap[v])
        end
        empty!(vs)
    end

    varassoc = Int[(1:dxvar_offset) .+ dxvar_offset; zeros(Int, length(fullvars) - dxvar_offset)] # variable association list
    sys, dxvar_offset, fullvars, varassoc, algeqs, graph, solvable_graph
end

end # module


text:
#================================================

Basic typedefs, function / macro defintions, etc. for the module

=================================================#

using Base.Threads

#===========================
Constants
===========================#

const DEFAULT_P = 0.5
const DEFAULT_MAX_HEIGHT = 64
const DEFAULT_NODE_CAPACITY = 256

const FLAG_IS_LEFT_SENTINEL = 0x80
const FLAG_IS_RIGHT_SENTINEL = 0x40
const IS_SENTINEL = FLAG_IS_LEFT_SENTINEL | FLAG_IS_RIGHT_SENTINEL

const NODE_NOT_FOUND = -1

#===========================
Abstract typedefs
===========================#

abstract type AbstractNode{T,M} end
abstract type AbstractSkipList{T,M} end

abstract type LeftSentinel{T,M} end
abstract type RightSentinel{T,M} end

#===========================
Helper functions
===========================#

"""
    random_height(p, args...)

Samples a number from a geometric distribution with parameter ``p`` and uses it
for the height of a new node in a Skiplist.

# Arguments

# Examples
"""
function random_height(p, args...; max_height = DEFAULT_MAX_HEIGHT)
    # This function uses the fact that the c.d.f. of a geometric distribution
    # is 1 - (1 - p)^k. To generate the height for a new node in a skip list,
    # we want it to be distributed as a geometric RV plus one.
    #
    # To perform this sampling, we randomly sample X ∈ [0,1], and find the
    # smallest value of k for which cdf(k) > X. We observe that
    #
    #           1 - (1 - p)^k   ≥    X                          =>
    #           (1 - p)^k       ≤    1 - X                      =>
    #           k log(1 - p)    ≤    log(1 - X)                 =>
    #           k               ≥    log(1 - X) / log(1 - p)
    #
    # (The inequality is flipped in the last step since log(1 - p) is necessarily
    # negative.) We can simplify this further by observing that Y = 1 - X has
    # the same distribution as X (i.e., Uniform([0,1])). As a result, to sample
    # a new random number, all we need to do is find the smallest integer k
    # satisfying k ≥ log(Y) / log(1 - p) for some Y ~ Uniform([0,1]), which
    # implies that
    #
    #           k = ⌈log(Y) / log(1 - p)⌉
    #

    p_scaler = 1 / log(1 - p)
    Y = rand(args...)
    @.(ceil(Int64, log(Y) * p_scaler) |> x -> min(max_height, x))
end

function _check_mode(M)
    if M != :List && M != :Set
        "SkipList mode $M is not recognized. Valid options are :List and :Set" |>
        ErrorException |>
        throw
    end
end


text:
using Test, NiLang, NiLang.AD
import NiLang: ⊕, ⊖

const add = ⊕(identity)

@testset "NGrad" begin
    @test NGrad{3}(exp) isa NGrad{3,typeof(exp)}
end

@testset "instr" begin
    x, y = 3.0, 4.0
    @instr Grad(add)(x, y; iloss=1)
    @test grad(x) == 1.0
    @test grad(y) == 1.0
    @test check_inv(Grad(add), (3.0, 4.0); verbose=true, atol=1e-5, iloss=1)
    x, y = 3.0, 4.0
    @test check_grad(add, (x, y); iloss=1)

    x, y = 3.0, 4.0
    Grad(add)(x, NoGrad(y); iloss=1)
    @test grad(y) === 0.0

    @test check_inv(⊕(*), (0.4, 0.4, 0.5))
    @test ⊖(*)(GVar(0.0, 1.0), GVar(0.4), GVar(0.6)) == (GVar(-0.24, 1.0), GVar(0.4, 0.6), GVar(0.6, 0.4))
    @test check_grad(⊕(*), (0.4, 0.4, 0.5); iloss=1)
    @test check_grad(⊖(*), (0.4, 0.4, 0.5); iloss=1)
end

@testset "i" begin
    @i function test1(a, b, out)
        a += b
        out += a * b
    end

    @i function tt(a, b)
        out ← 0.0
        test1(a, b, out)
        (~test1)(a, b, out)
        a += b
    end

    # compute (a+b)*b -> out
    x = 3.0
    y = 4.0
    out = 0.0
    @test check_grad(test1, (x, y, out); iloss=3)
    @test check_grad(tt, (x, y); iloss=1)
end

@testset "broadcast" begin
    # compute (a+b)*b -> out
    @i function test1(a, b)
        a .+= b
    end
    @i function test2(a, b, out, loss)
        a .+= b
        out .+= (a .* b)
        loss += out[1]
    end

    x = [3, 1.0]
    y = [4, 2.0]
    out = [0.0, 1.0]
    loss = 0.0
    # gradients
    @test check_grad(test2, (x, y, out, loss); iloss=4)
end

@testset "broadcast 2" begin
    # compute (a+b)*b -> out
    @i function test1(a, b)
        a += b
    end
    @i function test2(a, b, out)
        a += b
        out += (a * b)
    end

    # gradients
    a = 1.0
    b = 1.3
    c = 1.9
    @test check_grad(test2, (a,b,c); iloss=3)

    x = GVar([3, 1.0])
    y = GVar([4, 2.0])
    lout = GVar.([0.0, 1.0], [0.0, 2.0])
    @instr (~test2).(x, y, lout)
    @test grad.(lout) == [0,2.0]
    @test grad.(x) == [0, 4.0]
    @test grad.(y) == [0, 6.0]
end

@testset "function call function" begin
    # compute (a+b)*b -> out
    @i function test1(a, b)
        a += b
    end

    @i function test2(a, b, out)
        test1(a, out)
        (~test1)(a, out)
        out += (a * b)
    end

    a = 1.0
    b = 1.3
    c = 1.9
    @test check_grad(test2, (a,b,c); iloss=3)
end

@testset "neg sign" begin
    @i function test(out, x, y)
        out += x * (-y)
    end
    @test check_grad(test, (0.1, 2.0, -2.5); verbose=true, iloss=1)
end

@testset "i" begin
    @i function test1(a::T, b, out) where T<:Number
        add(a, b)
        out += a * b
    end
    @test isreversible(Grad(test1), Tuple{Number, Any,Any})
    @test isreversible(~Grad(test1), Tuple{Number, Any,Any})
    @test Grad(~test1) != ~(Grad(test1)) # this is not true
end

@testset "gradient" begin
    @test gradient((⊕(*)), (0.0, 2.0, 3.0); iloss=1) == (1.0, 3.0, 2.0)
end


text:
using Pkg
Pkg.activate(".")
Pkg.instantiate()

using RadMod2D

include("./models2D.jl")
include("./plot2D.jl")

using CairoMakie
CairoMakie.activate!(type = "svg")
filetype = ".svg"
# CairoMakie.activate!(type = "png")
# filetype = ".png"

function fig_c1_mesh2D()
    m = model_circle_in_circle(0.8, 1.6, 0.05)
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    plot_model(fig, ax, m, shownvec = true, shownodes = false, showcom = false, showleg = false, colors = [:blue, :red], linewidth = linewidth)
    # display(fig)
    save("fig_c1_mesh2D"*filetype, fig)
end

function fig_c1_mesh2D_legend()
    # settings for leg
    # leg = fig[end+1, 1] = Legend(fig, plt_elem, leg_string, orientation = :horizontal, tellwidth = false, tellheight = true)
    # leg.padding = (5.0f0, 5.0f0, 4.0f0, 4.0f0)
    # leg.linewidth = 1.0
    # leg.patchsize = (5, 3)
    # leg.rowgap = 0.1
    # leg.colgap = 8
    # for copy in mesh2D()
    m = model_circle_in_circle(0.8, 1.6, 0.05)
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    # legend horizontal at the bottom
    plot_model(fig, ax, m, shownvec = true, shownodes = false, showcom = false, showleg = true, colors = [:blue, :red], linewidth = linewidth)
    # display(fig)
    save("fig_c1_mesh2D_legend"*filetype, fig)
end

function fig_c1_view2D_existing_2elem()
    m = model_circle_in_circle(0.8, 1.6, 0.1)
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    # plot all
    plot_model(fig, ax, m, shownvec = false, shownodes = false, showcom = false, showleg = false, colors = [:black, :black], linewidth = linewidth)
    # elem pair check -> existing
    i1 = 47
    i2 = 33
    plot_model_elements(fig, ax, m, [i1,i2], shownvec = true, shownodes = false, showcom = false, color = :blue, linewidth = linewidth)
    linesegments!(ax, [Point2f0(m.elem[i1].com.x, m.elem[i1].com.y) => Point2f0(m.elem[i2].com.x, m.elem[i2].com.y)], color = :red, linewidth = linewidth)
    # elem pair check -> not existing
    i2 = 20
    plot_model_elements(fig, ax, m, [i1,i2], shownvec = true, shownodes = false, showcom = false, color = :blue, linewidth = linewidth)
    linesegments!(ax, [Point2f0(m.elem[i1].com.x, m.elem[i1].com.y) => Point2f0(m.elem[i2].com.x, m.elem[i2].com.y)], color = :red, linewidth = linewidth, linestyle = :dash)
    # display(fig)
    save("fig_c1_view2D_existing_2elem"*filetype, fig)
end

function fig_c1_view2D_existing1()
    m = model_circle_in_circle(0.8, 1.6, 0.1)
    #### vfmat calculation
    vfmat = zeros(Float64, m.nelem, m.nelem)
    @time existing_vf!(m, vfmat)
    #### 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    plot_model(fig, ax, m, shownvec = true, shownodes = false, showcom = false, showleg = false, colors = [:black, :black], linewidth = linewidth)
    plot_existing(fig, ax, m, 75, vfmat, color = :red, linewidth = 0.5*linewidth)
    # display(fig)
    save("fig_c1_view2D_existing1"*filetype, fig)
end

function fig_c1_view2D_existing2()
    m = model_circle_in_circle(0.8, 1.6, 0.1)
    #### vfmat calculation
    vfmat = zeros(Float64, m.nelem, m.nelem)
    @time existing_vf!(m, vfmat)
    n = 10
    dx, dy = create_tiles(m, n)
    @time t_occ = check_tile_occupation(m, dx, dy, n)
    @time blocking_vf_with_tiles!(m, vfmat, dx, dy, n, t_occ)
    #### 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    plot_model(fig, ax, m, shownvec = true, shownodes = false, showcom = false, showleg = false, colors = [:black, :black], linewidth = linewidth)
    plot_existing(fig, ax, m, 75, vfmat, color = :red, linewidth = 0.5*linewidth)
    # display(fig)
    save("fig_c1_view2D_existing2"*filetype, fig)
end

function fig_c1_view2D_blocking_occ()
    m = model_circle_in_circle(0.8, 1.6, 0.05)
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    # create tiles
    n = 15
    dx, dy = create_tiles(m, n)
    t_occ = check_tile_occupation(m, dx, dy, n)
    plot_occupied_tiles(fig, ax, dx, dy, n, t_occ)
    plot_empty_tiles(fig, ax, dx, dy, n, linewidth = linewidth)
    plot_model(fig, ax, m, shownvec = true, shownodes = false, showcom = false, showleg = false, colors = [:black, :black], linewidth = linewidth)
    # display(fig)
    save("fig_c1_view2D_bl_occ_nv"*filetype, fig)
end

function fig_c1_view2D_blocking_2elem_tiles()
    m = model_circle_in_circle(0.8, 1.6, 0.05)
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    # create tiles
    n = 15
    dx, dy = create_tiles(m, n)
    t_occ = check_tile_occupation(m, dx, dy, n)
    # plot_occupied_tiles(fig, ax, dx, dy, n, t_occ)
    # elem pair check
    i1 = 90
    i2 = 150
    isexisting = existing_vf_2elem(m, i1, i2)
    if isexisting
        blocking_vf_with_tiles_2elem_tiles(fig, ax, i1, i2, m, dx, dy, n, t_occ)
    end
    plot_model(fig, ax, m, shownvec = false, shownodes = false, showcom = false, showleg = false, colors = [:black, :black], linewidth = linewidth)
    plot_empty_tiles(fig, ax, dx, dy, n, linewidth = linewidth)
    plot_model_elements(fig, ax, m, [i1,i2], shownvec = true, shownodes = false, showcom = false, color = :blue, linewidth = linewidth)
    if isexisting
        linesegments!(ax, [Point2f0(m.elem[i1].com.x, m.elem[i1].com.y) => Point2f0(m.elem[i2].com.x, m.elem[i2].com.y)], color = :red, linewidth = linewidth)
    end
    # display(fig)
    save("fig_c1_view2D_bl_2elem_tiles"*filetype, fig)
end

function fig_c1_view2D_blocking_2elem()
    m = model_circle_in_circle(0.8, 1.6, 0.05)
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    # create tiles
    n = 15
    dx, dy = create_tiles(m, n)
    t_occ = check_tile_occupation(m, dx, dy, n)
    plot_occupied_tiles(fig, ax, dx, dy, n, t_occ)
    # elem pair check
    i1 = 90
    i2 = 150
    isexisting = existing_vf_2elem(m, i1, i2)
    if isexisting
        blocking_vf_with_tiles_2elem(fig, ax, i1, i2, m, dx, dy, n, t_occ)
    end
    plot_model(fig, ax, m, shownvec = false, shownodes = false, showcom = false, showleg = false, colors = [:black, :black], linewidth = linewidth)
    plot_empty_tiles(fig, ax, dx, dy, n, linewidth = linewidth)
    plot_model_elements(fig, ax, m, [i1,i2], shownvec = true, shownodes = false, showcom = false, color = :blue, linewidth = linewidth)
	# elements in bucket
	plot_model_elements(fig, ax, m, t_occ[6,11], shownvec = true, shownodes = false, showcom = false, color = :orange, linewidth = linewidth)
	# hitten element
	plot_model_elements(fig, ax, m, [17], shownvec = true, shownodes = false, showcom = false, color = :green, linewidth = linewidth)
    if isexisting
        linesegments!(ax, [Point2f0(m.elem[i1].com.x, m.elem[i1].com.y) => Point2f0(m.elem[i2].com.x, m.elem[i2].com.y)], color = :red, linewidth = linewidth)
    end
    # display(fig)
    save("fig_c1_view2D_bl_2elem"*filetype, fig)
end

function fig_c1_therm2D_temp()
    m = model_circle_in_circle(0.8, 1.6, 0.02)
    temp = zeros(m.nelem,1)
    set_bc_part!(m, temp, 1, 300)
    set_bc_part!(m, temp, 2, 400)
    # bc for some elements
    temp[189:313,1] .= 600
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    cbar = plot_model_with_value(fig, ax, m, temp[:,1], "Temperature in K", cmap = :viridis, linewidth = linewidth*1.5, showcbar = false)
    # cbar.ticks = 300:100:600
    # display(fig)
    save("fig_c1_therm2D_temp"*filetype, fig)
    # for i = 1:m.nnodes
    #     println(m.nodes[i].x, ";", m.nodes[i].y)
    # end
    # for i = 1:m.nelem
    #     println(m.elem[i].node1, ";", m.elem[i].node2, ";", m.elem[i].com, ";", m.elem[i].nvec, ";", m.elem[i].area)
    # end
    # for i = 1:m.nelem
    #     println(temp[i,:])
    # end
end

function fig_c1_therm2D_temp_legend()
    m = model_circle_in_circle(0.8, 1.6, 0.02)
    temp = zeros(m.nelem,1)
    set_bc_part!(m, temp, 1, 300)
    set_bc_part!(m, temp, 2, 400)
    # bc for some elements
    temp[189:313,1] .= 600
    # 2D plot
    fig = Figure(resolution = (500, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    cbar = plot_model_with_value(fig, ax, m, temp[:,1], "Temperature in K", cmap = :viridis, linewidth = linewidth*1.5, showcbar = true)
    cbar.ticks = 300:100:600
    # cbar.flipaxis = true
    cbar.vertical = false
    cbar.width = 120
    cbar.height = 12
    cbar.labelpadding = 1.0 # default 5.0
    cbar.ticklabelpad = 1.0 # default 3.0
    cbar.ticksize = 2.0 # default 6.0
    # display(fig)
    save("fig_c1_therm2D_temp_legend"*filetype, fig)
end

function fig_c1_therm2D_qrad()
    m = model_circle_in_circle(0.8, 1.6, 0.02)
    # vfmat calculation
    vfmat = zeros(Float64, m.nelem, m.nelem)
    existing_vf!(m, vfmat)
    n = 15
    dx, dy = create_tiles(m, n)
    t_occ = check_tile_occupation(m, dx, dy, n)
    blocking_vf_with_tiles!(m, vfmat, dx, dy, n, t_occ)
    calculating_vf!(m, vfmat, normit = true)
    # solve Qp
    epsilon = zeros(m.nelem,1)
    set_bc_part!(m, epsilon, 1, 0.9)
    set_bc_part!(m, epsilon, 2, 0.5)
    temp = zeros(m.nelem,1)
    set_bc_part!(m, temp, 1, 300)
    set_bc_part!(m, temp, 2, 400)
    # bc for some elements
    temp[189:313,1] .= 600
    @time Qp, G = tempsolver(m, vfmat, temp, epsilon)
    Qp_parts = [sum(Qp[m.elem2par[i].first:m.elem2par[i].last,1]) for i = 1:m.npar]
    area = [m.elem[i].area for i = m.elem2par[1].first:m.elem2par[end].last]
    qp_area = Qp[:] ./ area[:]
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    cbar = plot_model_with_value(fig, ax, m, qp_area, "Heat flux density in W/m", cmap = :viridis, linewidth = linewidth*1.5, showcbar = false)
    # cbar.ticks = -80:20:60
    # display(fig)
    save("fig_c1_therm2D_qrad"*filetype, fig)
    # for i = 1:m.nelem
    #     println(qp_area[i,:])
    # end
end

function fig_c1_therm2D_qrad_legend()
    m = model_circle_in_circle(0.8, 1.6, 0.02)
    # vfmat calculation
    vfmat = zeros(Float64, m.nelem, m.nelem)
    existing_vf!(m, vfmat)
    n = 15
    dx, dy = create_tiles(m, n)
    t_occ = check_tile_occupation(m, dx, dy, n)
    blocking_vf_with_tiles!(m, vfmat, dx, dy, n, t_occ)
    calculating_vf!(m, vfmat, normit = true)
    # solve Qp
    epsilon = zeros(m.nelem,1)
    set_bc_part!(m, epsilon, 1, 0.9)
    set_bc_part!(m, epsilon, 2, 0.5)
    temp = zeros(m.nelem,1)
    set_bc_part!(m, temp, 1, 300)
    set_bc_part!(m, temp, 2, 400)
    # bc for some elements
    temp[189:313,1] .= 600
    @time Qp, G = tempsolver(m, vfmat, temp, epsilon)
    Qp_parts = [sum(Qp[m.elem2par[i].first:m.elem2par[i].last,1]) for i = 1:m.npar]
    area = [m.elem[i].area for i = m.elem2par[1].first:m.elem2par[end].last]
    qp_area = Qp[:] ./ area[:]
    qp_area_kw = qp_area ./ 1000
    # 2D plot
    fig = Figure(resolution = (270, 270), font = "Arial", fontsize = 10)
    ax = fig[1, 1] = Axis(fig)
    linewidth = 1.0
    setup_axis(ax, linewidth)
    ax.xlabel = "X in m"
    ax.ylabel = "Y in m"
    ax.aspect = DataAspect()
    cbar = plot_model_with_value(fig, ax, m, qp_area_kw, "Heat flux density in kW/m", cmap = :viridis, linewidth = linewidth*1.5, showcbar = true)
    cbar.limits = (-4, 3)
    cbar.ticks = -4:1:3
    cbar.vertical = false
    cbar.width = 120
    cbar.height = 12
    cbar.labelpadding = 1.0 # default 5.0
    cbar.ticklabelpad = 1.0 # default 3.0
    cbar.ticksize = 2.0 # default 6.0
    # display(fig)
    save("fig_c1_therm2D_qrad_legend"*filetype, fig)
end

function calc_c1_view2D(;elemsize = 0.05)
    m = model_circle_in_circle(0.8, 1.6, elemsize)
    #### vfmat calculation
    vfmat = zeros(Float64, m.nelem, m.nelem)
    @time existing_vf!(m, vfmat)
    n = 15
    dx, dy = create_tiles(m, n)
    @time t_occ = check_tile_occupation(m, dx, dy, n)
    @time blocking_vf_with_tiles!(m, vfmat, dx, dy, n, t_occ)
    @time calculating_vf!(m, vfmat, normit = false)
    vfmatp = vfmat_to_parts(m, vfmat, normit = false)
    vfmatp
end

function calc_c1_analytical()
    d1 = 0.8
    d2 = 1.6
    r1 = d1 / 2
    r2 = d2 / 2
    vfmat = zeros(Float64, 2, 2)
    vfmat[1,1] = 0
    vfmat[1,2] = 1
    vfmat[2,1] = r1 / r2
    vfmat[2,2] = 1- vfmat[2,1]
    vfmat
end

##### cyl in cyl as example for algorithm
# fig_c1_mesh2D()
# fig_c1_mesh2D_legend()
# fig_c1_view2D_existing_2elem()
# fig_c1_view2D_existing1()
# fig_c1_view2D_existing2()
# fig_c1_view2D_blocking_occ()
# fig_c1_view2D_blocking_2elem_tiles()
# fig_c1_view2D_blocking_2elem()
# fig_c1_therm2D_temp()
# fig_c1_therm2D_temp_legend()
# fig_c1_therm2D_qrad()
# fig_c1_therm2D_qrad_legend()

##### vf calculation vs analytical solution
# calc_c1_analytical()
# calc_c1_view2D(elemsize = 0.2)
# calc_c1_view2D(elemsize = 0.1)
# calc_c1_view2D(elemsize = 0.05)
# calc_c1_view2D(elemsize = 0.01)
# calc_c1_view2D(elemsize = 0.005)

# include("./test/run_cyl_in_cyl.jl")


text:
using Flux

export BasicBlock

struct BasicBlock
    pre_block::Any
    downsample_block::Any
    drop_block::Any
end

function BasicBlock(
    chs::Pair{<:Integer,<:Integer};
    init = Flux.kaiming_normal,
    drop_prob = 0.0,
    use_dropblock = false,
    block_size = 3,
    pool_stride = 2,
)
    in_channel, out_channel = chs
    pre_block = Chain(
        Conv((3, 3), in_channel => out_channel, init = init, stride = 1, pad = 1),
        BatchNorm(out_channel, (x) -> leakyrelu(x, Float32(0.1))),
        Conv((3, 3), out_channel => out_channel, init = init, stride = 1, pad = 1),
        BatchNorm(out_channel, (x) -> leakyrelu(x, Float32(0.1))),
        Conv((3, 3), out_channel => out_channel, init = init, stride = 1, pad = 1),
        BatchNorm(out_channel),
        (x) -> leakyrelu.(x, Float32(0.1)),
        MaxPool((pool_stride, pool_stride), pad = SamePad(), stride = pool_stride),
    )
    downsample_block = in_channel != out_channel || pool_stride != 1 ?
        Chain(
        Conv((1, 1), in_channel => out_channel, stride = pool_stride, pad = SamePad()),
        BatchNorm(out_channel),
    ) :
        identity
    drop_block = use_dropblock ? DropBlock(block_size, Float32(1 - drop_prob)) : Dropout(Float32(drop_prob))
    if drop_prob > 0
        Flux.trainmode!(drop_block)
    else
        Flux.testmode!(drop_block)
    end
    BasicBlock(pre_block, downsample_block, drop_block)
end

function (block::BasicBlock)(x)
    residule = block.downsample_block(x)
    out = block.pre_block(x)
    out = out .+ residule
    out = block.drop_block(out)
    out
end

Flux.testmode!(m::BasicBlock, mode = true) = (
    map(
        x -> Flux.testmode!(x, mode),
        (m.pre_block, m.downsample_block, m.drop_block),
    );
    m
)

Flux.functor(::Type{<:BasicBlock}, m) =
    (m.pre_block, m.downsample_block, m.drop_block),
    blocks -> BasicBlock(blocks...)


text:
using Flatten

abstract type AgComponent end

macro def(name, definition)
    return quote
        macro $(esc(name))()
            esc($(Expr(:quote, definition)))
        end
    end
end

@def add_preprefix begin
    if !isnothing(id_prefix)
        prefix = id_prefix * prefix
    end
end

function create(spec::Dict)
    dt = deepcopy(spec)
    cls_name = pop!(dt, :component)

    cls = nothing
    try
        cls = eval(Symbol(cls_name))
    catch e
        if e isa UndefVarError
            println("$(cls_name) is not a valid Agtor component name")
            println("Error in specification for '$(dt[:name])'")
        end

        throw(e)
    end

    if cls_name == "FarmZone"
        cd_fn::Union{String, Nothing} = get(dt, :climate_data, nothing)
        if isnothing(cd_fn)
            throw(ArgumentError("Climate data not found in Zone spec."))
        end

        climate::Climate = load_climate(cd_fn)

        return create(dt, climate)
    end

    return cls(; dt...)
end

function create(spec::Dict, climate_data::String)
    dt = deepcopy(spec)
    cls_name = pop!(dt, :component)
    cls = eval(Symbol(cls_name))

    if cls_name == "Basin"
        return cls(; name=dt[:name], zone_specs=dt[:zone_spec],
                   climate_data=dt[:climate_data])
    end

    throw(ArgumentError("Unknown component: $(cls_name), with additional parameter '$(climate_data)'"))
end

function Base.getproperty(A::AgComponent, v::Symbol)
    field = getfield(A, v)
    if field isa AgParameter
        return field.value
    end

    return field
end

max_stars_repo_name: ConnectedSystems/agtor-jl

max_issues_repo_name: ConnectedSystems/agtor-jl

max_forks_repo_name: ConnectedSystems/agtor-jl


text:
@polly function kernel_jacobi_2d(tsteps, A, B)
    n = size(A,1)

    for t = 1:tsteps
        for i = 2:(n-1), j = 2:(n-1)
            B[i,j] = 0.2 * (A[i,j] + A[i,j-1] + A[i,1+j] + A[1+i,j] + A[i-1,j])
        end
        for i = 2:(n-1), j = 2:(n-1)
            A[i,j] = 0.2 * (B[i,j] + B[i,j-1] + B[i,1+j] + B[1+i,j] + B[i-1,j])
        end
    end
end

let
    tsteps = 500
    n = 1300

    A = zeros(Float32, n, n)
    B = zeros(Float32, n, n)

    for i = 1:n, j = 1:n
        A[i,j] = (i*(j+2) + 2) / n
        B[i,j] = (i*(j+3) + 3) / n
    end

    SUITE["jacobi-2d"] = @benchmarkable kernel_jacobi_2d($tsteps, A, B) setup = (A = copy($A); B = copy($B))
end


text:
module anagrams
    export storeanagrams
    export readanagrams

    using JLD2, FileIO

    JLDBFILENAME = ""::String
    const JLDBKEY = "ANAGRAMS"

    function readwords()::Vector{String}
        wordslist = []
        for line in eachline("words.txt")
            push!(wordslist, line)
        end
        return wordslist
    end

    function get_anagrams(wordslist::Vector{String})::Dict{String, Vector{String}}
        wordsset = Dict{String, Vector{String}}()

        for word in wordslist
            key = join(sort([l for l in word]))
            wordsset[key] = push!(get!(wordsset, key, []), word)
        end
        return filter!(kv -> length(kv[2]) >= 2, wordsset)
    end

    function JLDSAVE(filename::String, dict::Dict{String, Vector{String}})
        global JLDBKEY
        save(filename, Dict(JLDBKEY=>dict))
    end

    function JLDLOAD(filename::String)
        global JLDBKEY
        return load(filename)[JLDBKEY]
    end

    function storeanagrams(name::String)
        global JLDBFILENAME
        basewords::Vector{String} = readwords()
        anagramsbook::Dict{String, Vector{String}} = get_anagrams(basewords)
        name = name*".jld2"
        try
            JLDSAVE(name, anagramsbook)
            JLDBFILENAME = name
        catch exc
            println("(!) Error: $exc")
        end
    end

    function readanagrams(word::String)::Vector{String}
        RET::Vector{String} = []
        word = join(sort!([w for w in word]))
        global JLDBFILENAME
        try
            if isfile(JLDBFILENAME)
                RET = JLDLOAD(JLDBFILENAME)[word]
            else
                error("(!) Error: DB File Doesn't Exist, create one by calling storeanagrams(filename)")
            end
        catch exc
            println("(!) Error: $exc")
        finally
            return RET
        end
    end
end

max_stars_repo_name: planetnoob/LearnJulia

max_issues_repo_name: planetnoob/LearnJulia

max_forks_repo_name: planetnoob/LearnJulia


text:
#------------------------
#--------  Functions

#--------  random generation

using Random
function GenSymArrow(n::Integer,i::Integer)
    # generates symmetric n x n arrowhad matrix with arrow top at (i,i)
    SymArrow(rand(n-1),rand(n-1),rand(),i)
end

function GenSymDPR1(n::Integer)
    # generates symmetric n x n DPR1 matrix
    SymDPR1(rand(n),rand(n),rand())
end

#-------- Inverses
function inv(A::SymArrow{T},i::Integer,τ::Vector{Float64}=[1e3;10.0*length(A.D)]) where T
    # COMPUTES: inverse of a SymArrow matrix A, inv(A-A.D[i]*I) which is again SymArrow
    # uses higher precision to compute top of the arrow element accurately, if
    # needed.
    # τ=[tolb,tolz] are tolerances, usually [1e3, 10*n]
    # [0.0,0.0] forces DoubleDouble, [1e50,1e50] would never use it
    # RETURNS:  SymArrow(D,z,b,i), Kb, Kz, Qout
    # Kb - condition Kb, Kz - condition Kz, Qout = 1 / 0 - double was / was not used

    n=length(A.D)
    D=Array{T}(undef,n)
    z=Array{T}(undef,n)
    wz=one(T)/A.z[i]
    σ=A.D[i]
    for k=1:i-1
        D[k]=one(T)/(A.D[k]-σ)
        z[k]=-A.z[k]*D[k]*wz
    end
    for k=i+1:n
        D[k-1]=one(T)/(A.D[k]-σ)
        z[k-1]=-A.z[k]*D[k-1]*wz
    end
    # D=[1./(A.D[1:i-1]-shift),1./(A.D[i+1:end]-σ)]
    # wz=1/A.z[i]
    # z=-[A.z[1:i-1], A.z[i+1:end]].*D1*wz1
    a=A.a-A.D[i]
    # compute the sum in a plain loop
    P=zero(T)
    Q=zero(T)
    Qout=0
    nn=n-1
    for k=1:i-1
        D[k]>0.0 ? P+=A.z[k]^2*D[k] : Q+=A.z[k]^2*D[k]
    end
    for k=i:nn
        D[k]>0.0 ? P+=A.z[k+1]^2*D[k] : Q+=A.z[k+1]^2*D[k]
    end
    a<0 ? P-=a : Q-=a
    Kb=(P-Q)/abs(P+Q)
    Kz=maximum(abs,A.z)*abs(wz)
    if Kb<τ[1] ||  Kz<τ[2]
        b=(P+Q)*wz*wz
    else  # recompute in Double or BigFloat
        if Kz<1.0/eps()^2
            Type=Double
            Qout=1
        else
        # Example of a matrix where BigFloat is neeed, courtesy of Stan Eisenstat, is:
        # A=SymArrow([1+eps(), 1-eps(), 0,-1+2*eps(), -1-2*eps()],[2,2,eps()^4,1,1.0],6.0,5)
            Type=BigFloat
            Qout=50
        end
        σ₁=map(Type,A.D[i])
        Dd=[[Type(A.D[k])-σ₁ for k=1:i-1];
            [Type(A.D[k])-σ₁ for k=i+1:length(A.D)]]
        wzd=Type(A.z[i])
        ad=Type(A.a)-σ₁

        Pd,Qd=map(Type,(0.0,0.0))
        for k=1:i-1
            convert(Float64,Dd[k])>0.0 ? Pd+=Type(A.z[k])^2/Dd[k] :
            Qd+=Type(A.z[k])^2/Dd[k]
        end
        for k=i:nn
            convert(Float64,Dd[k])>0.0 ? Pd+=Type(A.z[k+1])^2/Dd[k] :
            Qd+=Type(A.z[k+1])^2/Dd[k]
        end
        convert(Float64,ad)<0 ?   Pd=Pd-ad : Qd=Qd-ad
        bd=(Pd+Qd)/(wzd*wzd)
        b=convert(Float64,bd)
    end

    if i<A.i
        # This reduces memory allocation (4 vectors) and is 30-40% faster
        for k=n-1:-1:A.i-1
            D[k+1]=D[k]
            z[k+1]=z[k]
        end
        D[A.i-1]=zero(T)
        z[A.i-1]=wz
        SymArrow(D,z,b,i),Kb,Kz,Qout
    else
        for k=n-1:-1:A.i
            D[k+1]=D[k]
            z[k+1]=z[k]
        end
        D[A.i]=zero(T)
        z[A.i]=wz
        SymArrow(D,z,b,i+1),Kb,Kz,Qout
    end
end # inv

function inv!(B::SymArrow{T},A::SymArrow{T},i::Integer,τ::Vector{Float64}=[1e3;10.0*length(A.D)]) where T
    # In-place version!
    # COMPUTES: inverse of a SymArrow matrix A, inv(A-A.D[i]*I) which is again SymArrow
    # uses higher precision to compute top of the arrow element accurately, if
    # needed.
    # τ=[tolb,tolz] are tolerances, usually [1e3, 10*n]
    # [0.0,0.0] forces DoubleDouble, [1e50,1e50] would never use it
    # RETURNS:  B=SymArrow(D,z,b,i), Kb, Kz, Qout
    # Kb - condition Kb, Kz - condition Kz, Qout = 1 / 0 - double was / was not used

    n=length(A.D)
    # D=Array{T}(undef,n)
    # z=Array{T}(undef,n)
    wz=one(T)/A.z[i]
    σ=A.D[i]
    for k=1:i-1
        B.D[k]=one(T)/(A.D[k]-σ)
        B.z[k]=-A.z[k]*B.D[k]*wz
    end
    for k=i+1:n
        B.D[k-1]=one(T)/(A.D[k]-σ)
        B.z[k-1]=-A.z[k]*B.D[k-1]*wz
    end

    # D=[1./(A.D[1:i-1]-shift),1./(A.D[i+1:end]-σ)]
    # wz=1/A.z[i]
    # z=-[A.z[1:i-1], A.z[i+1:end]].*D1*wz1
    a=A.a-A.D[i]

    # compute the sum in a plain loop
    P=zero(T)
    Q=zero(T)
    Qout=0
    nn=n-1
    for k=1:i-1
        B.D[k]>0.0 ? P+=A.z[k]^2*B.D[k] : Q+=A.z[k]^2*B.D[k]
    end
    for k=i:nn
        B.D[k]>0.0 ? P+=A.z[k+1]^2*B.D[k] : Q+=A.z[k+1]^2*B.D[k]
    end
    a<0 ? P-=a : Q-=a

    Kb=(P-Q)/abs(P+Q)
    Kz=maximum(abs,A.z)*abs(wz)

    if Kb<τ[1] ||  Kz<τ[2]
        b=(P+Q)*wz*wz
    else  # recompute in Double or BigFloat
        if Kz<1.0/eps()^2
            Type=Double
            Qout=1
        else
        # Example of a matrix where BigFloat is neeed, courtesy of Stan Eisenstat, is:
        # A=SymArrow([1+eps(), 1-eps(), 0,-1+2*eps(), -1-2*eps()],[2,2,eps()^4,1,1.0],6.0,5)
            Type=BigFloat
            Qout=50
        end
        σ₁=map(Type,A.D[i])
        Dd=[[Type(A.D[k])-σ₁ for k=1:i-1];
            [Type(A.D[k])-σ₁ for k=i+1:length(A.D)]]
        wzd=Type(A.z[i])
        ad=Type(A.a)-σ₁

        Pd,Qd=map(Type,(0.0,0.0))
        for k=1:i-1
            convert(Float64,Dd[k])>0.0 ? Pd+=Type(A.z[k])^2/Dd[k] :
            Qd+=Type(A.z[k])^2/Dd[k]
        end
        for k=i:nn
            convert(Float64,Dd[k])>0.0 ? Pd+=Type(A.z[k+1])^2/Dd[k] :
            Qd+=Type(A.z[k+1])^2/Dd[k]
        end
        convert(Float64,ad)<0 ?   Pd-=ad : Qd-=ad
        bd=(Pd+Qd)/(wzd*wzd)
        b=convert(Float64,bd)
    end

    if i<A.i
        # This reduces memory allocation (4 vectors) and is 30-40% faster
        for k=n-1:-1:A.i-1
            B.D[k+1]=B.D[k]
            B.z[k+1]=B.z[k]
        end
        B.D[A.i-1]=zero(T)
        B.z[A.i-1]=wz
        B.i=i
        B.a=b
        # SymArrow(D,z,b,i),Kb,Kz,Qout
        return Kb,Kz,Qout
    else
        for k=n-1:-1:A.i
            B.D[k+1]=B.D[k]
            B.z[k+1]=B.z[k]
        end
        B.D[A.i]=zero(T)
        B.z[A.i]=wz
        B.i=i+1
        B.a=b
        # SymArrow(D,z,b,i+1),Kb,Kz,Qout
        return Kb,Kz,Qout
    end
end # inv!

function inv(A::SymArrow{T}, σ::Float64, τ::Float64=1.0e3) where T
    # COMPUTES: inverse of the shifted SymArrow A, inv(A-σ*I) which is SymDPR1
    # uses DoubleDouble to compute rho accurately, if needed.
    # τ is tolerance, usually 1e3,  0.0 forces Double, 1e50 would never use it
    # RETURNS: SymDPR1(D,u,ρ), Kρ, Qout
    # Kρ - condition Kρ, Qout = 1 / 0 - Double was / was not used

    n=length(A.D)
    D=Array{T}(undef,n+1)
    u=Array{T}(undef,n+1)
    for k=1:A.i-1
        D[k]=one(T)/(A.D[k]-σ)
        u[k]=A.z[k]*D[k]
    end
    D[A.i]=zero(T)
    u[A.i]=-one(T)
    for k=A.i:n
        D[k+1]=one(T)/(A.D[k]-σ)
        u[k+1]=A.z[k]*D[k+1]
    end

    # compute rho and Krho
    #--- compute the sum in a plain loop
    a=A.a-σ
    P=zero(T)
    Q=zero(T)
    Qout=0
    for k=1:A.i-1
        D[k]>0.0 ? P+=A.z[k]^2*D[k] : Q+=A.z[k]^2*D[k]
    end
    for k=A.i:n
        D[k+1]>0.0 ? P+=A.z[k]^2*D[k+1] : Q+=A.z[k]^2*D[k+1]
    end
    a<0 ? P-=a : Q-=a

    # Condition of rho
    Kρ=(P-Q)/abs(P+Q)
    if Kρ<τ
        ρ=-1.0/(P+Q)
    else  # recompute in Double
        Qout=1
        Pd,Qd=map(Double,(0.0,0.0))
        σ₁=Double(σ)
        ad=Double(A.a)-σ₁
        for k=1:A.i-1
            D[k]>0.0 ? Pd+=Double(A.z[k])^2/(Double(A.D[k])-σ₁) : Qd+=Double(A.z[k])^2/(Double(A.D[k])-σ₁)
        end
        for k=A.i:n
            D[k+1]>0.0 ? Pd+=Double(A.z[k])^2/(Double(A.D[k])-σ₁) : Qd+=Double(A.z[k])^2/(Double(A.D[k])-σ₁)
        end
        ad.hi+ad.lo<0 ? Pd-=ad : Qd-=ad
        if Pd+Qd!=Double(0.0,0.0)
            Kρ=Float64((Pd-Qd)/abs(Pd+Qd))
            r=Double(1.0)/(Pd+Qd)
            ρ=-(r.hi+r.lo)
        else
        # Here we need quadruple working precision. We are using BigFloat.
        # Example of a matrix where this is neeed, courtesy of Stan Eisenstat, is:
        # A=SymArrow([1+eps(), 1-eps(), -1+2*eps(), -1-2*eps()],[2,2,1,1.0],6.0,5)
            Qout=100
            Pd,Qd=map(BigFloat,(0.0,0.0))
            σd=BigFloat(σ)
            ad=BigFloat(A.a)-σ₁
            for k=1:A.i-1
                D[k]>0.0 ? Pd+=BigFloat(A.z[k])^2/(BigFloat(A.D[k])-σ₁) : Qd+=BigFloat(A.z[k])^2/(BigFloat(A.D[k])-σ₁)
            end
            for k=A.i:n
                D[k]>0.0 ? Pd+=BigFloat(A.z[k])^2/(BigFloat(A.D[k])-σ₁) : Qd+=BigFloat(A.z[k])^2/(BigFloat(A.D[k])-σ₁)
            end
            ad<0 ? Pd-=ad : Qd-=ad
            Kρ=Float64((Pd-Qd)/abs(Pd+Qd))
            r=BigFloat(-1.0)/(Pd+Qd)
            ρ=Float64(r)
        end
    end

    # returns the following
    SymDPR1(D,u,ρ), Kρ, Qout
end # inv

function inv(A::SymArrow{T}, σ::Double) where T
    # COMPUTES: inverse of the shifted SymArrow A, inv(A-shift*I), which is a SymDPR1
    # here shift is Double so it uses Double to compute everything
    # RETURNS: SymDPR1(D₁,u₁,ρ), Qout
    # Qout = 1 on exit meaning Double was used

    n=length(A.D)
    D=Array{Double}(undef,n+1)
    u=Array{Double}(undef,n+1)
    for k=1:n
        D[k]=Double(A.D[k])-σ
    end
    u[1:n]=map(Double,A.z)
    a=map(Double,A.a)-σ
    i=A.i
    oned=Double(1.0)
    zerod=Double(0.0)
    for k=1:i-1
        u[k]=u[k]/D[k]
        D[k]=oned/D[k]
    end
    for k=i:n
        u[k+1]=u[k]/D[k]
        D[k+1]=oned/D[k+1]
    end
    D[i]=zerod
    u[i]=Double(-1.0)

    # compute rho and Kρ
    P,Q=zerod,zerod
    Qout=1
    for k=1:i-1
        D[k].hi > 0.0 ? P+=Double(A.z[k])*u[k] : Q+=Double(A.z[k])*u[k]
    end
    for k=i:n
        D[k+1].hi >0.0 ? P+=Double(A.z[k])*u[k+1] : Q+=Double(A.z[k])*u[k+1]
    end
    a.hi+a.lo<0.0  ? P-=a : Q-=a
    r=oned/(P+Q)
    ρ=-(r.hi+r.lo)
    # returns the following
    D₁=Array{T}(undef,n+1)
    u₁=Array{T}(undef,n+1)
    for k=1:n+1
        D₁[k]=D[k].hi+D[k].lo
    end
    for k=1:n+1
        u₁[k]=u[k].hi+u[k].lo
    end
    SymDPR1(D₁,u₁,ρ), Qout
end # inv

mutable struct Info
    Sind::Int
    Kb::Float64
    Kz::Float64
    Kν::Float64
    Kρ::Float64
    Qout::Int
end

function  eigen( A::SymArrow{T}, Ainv::SymArrow{T}, k::Integer,
    τ::Vector{Float64}=[1e3,10.0*length(A.D),1e3,1e3,1e3]) where T
    # COMPUTES: k-th eigenpair of an ordered irreducible SymArrow
    # A = [Diagonal(D) z; z' alpha]
    # τ=[tolb,tolz,tolnu,tolrho,tollambda] = [1e3,10.0*n,1e3,1e3,1e3]
    # RETURNS: λ,v, Info
    # where
    # λ - k-th eigenvalue in descending order
    # v - λ's normalized eigenvector
    # Info = Sind, Kb, Kz, Kν, Kρ, Qout
    # Sind = shift index i for the k-th eigenvalue
    # Kb, Kz, Kν, Kρ - condition numbers
    # Qout = 1 / 0 - Double was / was not used

    # Set the dimension
    n = length(A.D) + 1
    # Set all conditions initially to zero
    Kb,Kz,Kν,Kρ=0.0,0.0,0.0,0.0
    Qout=0
    v=zeros(T,n)
    # Kz is former κ_ν
    # Determine the shift σ, the shift index i, and whether λ
    # is on the left or the right side of the nearest pole
    # Exterior eigenvalues (k = 1 or k = n):
    if k == 1
        σ,i,side = A.D[1],1,'R'
    elseif k==n
        σ,i,side = A.D[n-1],n-1,'L'
    else
        # Interior eigenvalues (k in (2,...n-1) ):
        Dtemp = A.D.-A.D[k]
        atemp = A.a-A.D[k]
        middle = Dtemp[k-1]/2.0
        Fmiddle = (atemp-middle)-sum(A.z.^2 ./(Dtemp.-middle))
        σ,i,side = Fmiddle < 0.0 ? (A.D[k],k,'R') : (A.D[k-1],k-1,'L')
    end

    # Compute the inverse of the shifted matrix, A_i^(-1), Kb and Kz
    # Ainv,Kb,Kz,Qout = inv(A,i)
    Kb,Kz,Qout = inv!(Ainv,A,i)

    # Compute the eigenvalue of the inverse shifted matrix
    ν = bisect( Ainv,side )

    if abs(ν)==Inf
        # this is nonstandard
        # Deflation in aheig (ν=Inf)
        v[i]=1.0
        λ=σ
    else
        # standard case, full computation
        # ν₁ is the F- or 1-norm of the inverse of the shifted matrix
        # ν₁=maximum([sum(abs.(Ainv.z))+abs(Ainv.a), maximum(abs,Ainv.D)+abs.(Ainv.z))])
        ν₁=0.0
        for l=1:n-1
            ν₁=max(ν₁,abs(Ainv.D[l])+abs(Ainv.z[l]))
        end
        ν₁=max(sum(abs,Ainv.z)+abs(Ainv.a), ν₁)
        Kν=ν₁/abs(ν)
        while  Kν>τ[3]
            # Remedies according to Remark 3 - we shift between original
            # eigenvalues and compute DPR1 matrix
            # 1/ν₁+σ, 1/ν+σ
            # println("Remedy 3 ")
            ν = side=='R' ? abs(ν) : -abs(ν)
            ν₁=-sign(ν)*ν₁
            σ₁=(ν₁+ν)/(2.0*ν*ν₁)+σ
            if count(isequal(σ₁),A.D)>0  # we came back with a pole
                # recompute σᵥ more accurately with Dekker
                σᵥ=(Double(ν₁)+Double(ν))/(Double(2.0)*Double(ν)*Double(ν₁))+Double(σ)
                # Compute the inverse of the shifted arrowhead (DPR1)
                Ainv,Qout₁=inv(A,σᵥ) # Ainv is Float64
                ν₁=bisect(Ainv,side)
                μ₁ = 1.0/ν₁
                D0=map(Double,A.D)-σᵥ
                D1=[D0[l].hi+D0[l].lo for l=1:length(D0)]
                if count(isequal(μ₁),D1)>0
                    ind=findall(isequal(μ₁),D1)
                    v=zeros(T,n)
                    v[ind]=1.0
                else
                    v=[ A.z./(D1.-μ₁);-1.0]
                end
                # Shift the eigenvalue back in Double
                lam = Double(1.0)/Double(ν₁)+σᵥ
                # Return this
                λ=lam.hi+lam.lo
                normalize!(v)
                κ=Info(i,Kb,Kz,Kν,Kρ,Qout)
                return λ,v,κ
            else
                # Compute the inverse of the shifted arrowhead (DPR1)
                Ainv₁, Kρ,Qout₁=inv(A,σ₁,τ[4]) # Ainv is Float64
                # Compute the eigenvalue by bisect for DPR1
	            # Note: instead of bisection could use dlaed4 (need a wrapper) but
	            # it is not faster. There norm(u)==1
                ν= bisect(Ainv₁,side)
                if side=='R' && A.D[1]>0.0 && ν<0.0
                    ν=bisect(Ainv₁,'L')
                end
                μ=1.0/ν
                ν₁=maximum(abs,Ainv₁.D)+abs(Ainv₁.r)*dot(Ainv₁.u,Ainv₁.u)
                Kν=ν₁/abs(ν)
                σ=σ₁
            end
            Qout=Qout+2*Qout₁
        end

        # Accuracy is fine, compute the eigenvector
        μ = 1.0/ν
        for l=1:n-1
            v[l] = A.z[l]/((A.D[l]-σ)-μ)
        end
        v[n]=-1.0
        λ=μ+σ
        normalize!(v)

        # Remedy according to Remark 1 - we recompute the the eigenvalue
        # near zero from the inverse of the original matrix (a DPR1 matrix).
        if (abs(A.D[i])+abs(1.0/ν))/abs(λ)>τ[5]
            if k==1 && A.D[1]<0.0 || k==n && A.D[n-1]>0.0 || i<n-1 && side=='L' &&
                sign(A.D[i])+sign(A.D[i+1])==0 || i>1 &&
                side=='R' && sign(A.D[i])+sign(A.D[i-1])==0
                # println("Remedy 1 ")
                # Compute the inverse of the original arrowhead (DPR1)
                Ainv₁,Kρ,Qout₁ = inv(A,0.0,τ[4]) # Ainv is Float64
                Qout=Qout+4*Qout₁
                if abs(Ainv₁.r)==Inf
                    λ=0.0
                else
                    # Here we do not need bisection. We compute the Rayleigh
                    # quotient by using already computed vectors which is
                    # componentwise accurate
                    ν₁=sum(v.^2 .*Ainv₁.D)+Ainv₁.r*sum(v.*Ainv₁.u)^2;
                    λ=1.0/ν₁
                end
            end
        end
    end

    # Return this
    κ=Info(i,Kb,Kz,Kν,Kρ,Qout)
    λ,v,κ
end # eigen(k)

function eigen(A::SymArrow{T}, τ::Vector{Float64}=[1e3,10.0*length(A.D),1e3,1e3,1e3]) where T
    # COMPUTES: all eigenvalues and eigenvectors of a real symmetric SymArrow
    # A = [diag(D) z;z' alpha] (notice, here we assume A.i==n)
    # τ = [tolb,tolz,tolnu,tolrho,tollambda] = [1e3,10.0*n,1e3,1e3,1e3] or similar
    # RETURNS: Eigen(Λ,U), κ
    # where
    # Λ = eigenvalues in decreasing order, U = eigenvectors,
    # κ[k]=Info(Sind[k], Kb[k], Kz[k], Kν[k], Kρ[k], Qout[k])
    # Sind[k] - shift index i for the k-th eigenvalue
    # Kb, Kz, Kν, Kρ [k] - respective conditions for the k-th eigenvalue
    # Qout[k] = 1 / 0 - Double was / was not used when computing k-th eigenvalue

    n=length(A.D)+1
    n0=n
    # Ordering the matrix
    is=sortperm(A.D,rev=true)
    D=A.D[is]
    z=A.z[is]
    U=Array{T}(I,n,n)
    Λ=zeros(T,n)
    κ=[Info(0, 0.0, 0.0, 0.0, 0.0, 0) for i=1:n]
    # Quick return for 1x1
    if n==1
        return Eigen([A.a],U),κ
    end

    #  test for deflation in z
    z0=findall(iszero,z)
    zx=findall(!iszero,z)
    if isempty(zx)  # nothing to do
        Λ=[A.D;A.a]
        isΛ=sortperm(Λ,rev=true)
        Λ=Λ[isΛ]
        U=view(U,:,isΛ) # U[:,isΛ]
        return Eigen(Λ,U),κ
    end
    if !isempty(z0)
        Λ[z0]=D[z0]
        D=D[zx]
        z=z[zx]
        if isempty(z)
            Λ[n]=A.a
        else
            n=length(z)+1
        end
    end
    zx=[zx;n0]

    #  Test for deflation in D
    g=D[1:n-2]-D[2:n-1]
    # Can play with inexact deflation
    # g0=find(abs(g)<eps)
    # gx=find(abs(g)>=eps)
    # Only exact deflation !!
    g0=findall(iszero,g)
    gx=findall(!iszero,g)
    if !isempty(g0)
        # Deflation
        Dgx=D[gx]; zgx=z[gx]
        lg0=length(g0)
        R=Array{Tuple{LinearAlgebra.Givens{Float64},Float64}}(undef,lg0)
        for l=lg0:-1:1
            R[l]=givens(z[g0[l]],z[g0[l]+1],zx[g0[l]],zx[g0[l]+1])
            z[g0[l]]=R[l][2]; z[g0[l]+1]=0.0
            # A_mul_Bc!(U,R) # multiply R'*U later
            Λ[zx[g0[l]+1]]=D[g0[l]+1]
        end
        # remains
        gx=[0;gx].+1
        nn=length(gx)
        zxx=zx[[gx;n]]
        Axx=SymArrow(D[gx],z[gx],A.a,nn+1)
        Bxx=[deepcopy(Axx) for i=1:Threads.nthreads()]
        Threads.@threads for k=1:nn+1
            tid=Threads.threadid()
            Λ[zxx[k]],U[zxx,zxx[k]],κ[zxx[k]] = eigen(Axx,Bxx[tid],k)
        end
        for l=1:lg0
            # U=R[l][1]'*U
            lmul!(R[l][1],U)
        end
    else
        # No deflation in D
        Ax=SymArrow(D,z,A.a,n)
        Bx=[deepcopy(Ax) for i=1:Threads.nthreads()]
        Threads.@threads for k=1:n
            tid=Threads.threadid()
            Λ[zx[k]],U[zx,zx[k]],κ[zx[k]] = eigen(Ax,Bx[tid],k)
        end
    end
    # end
    # back premutation of vectors
    isi=sortperm(is)
    # must sort Λ once more
    es=sortperm(Λ,rev=true)
    Λ.=Λ[es]
    U=view(U,[isi[1:A.i-1];n0;isi[A.i:n0-1]],es)
    # U.=U[[isi[1:A.i-1];n0;isi[A.i:n0-1]],es]
    κ.=κ[es]
    # Return this
    return Eigen(Λ,U),κ
end # eigen (all)

function bisect(A::SymArrow{T}, side::Char) where T
    # COMPUTES: the leftmost (for side='L') or the rightmost (for side='R') eigenvalue
    # of a SymArrow A = [diag (D) z; z'] by bisection.
    # RETURNS: the eigenvalue

    # Determine the starting interval for bisection, [left; right]
    # left, right = side == 'L' ? {minimum([A.D-abs(A.z),A.a-sum(abs(A.z))]), minimum(A.D)} :
    #   {maximum(A.D),maximum([A.D+abs.(A.z),A.a+sum(abs,A.z)])}
    z = abs.(A.z)
    if side == 'L'
        left  = minimum(A.D .- z)
        left  = min(left, A.a - sum(z))
        right = minimum(A.D)
    else
        left  = maximum(A.D)
        right = maximum(A.D .+ z)
        right = max(right, A.a + sum(z))
    end

    # Bisection
    middle = (left + right) / convert(T,2)
    z.^= 2
    count, n = 0, length(A.D)
    while (right-left) > 2.0 * eps() * max(abs(left), abs(right))
        # in @time 50% of the time was garbage collection. The following line
        # assigns new vector every time it is called, so it is much better in the
        # loop?? Original timing were 30 secs for n=4000, 2.17 sec for n=1000

        # Fmiddle = A.a-middle-sum(z./(A.D-middle))
        Fmiddle = zero(T)
        for k=1:n
            Fmiddle += z[k] / (A.D[k] - middle)
        end
        Fmiddle = A.a - middle - Fmiddle

        if Fmiddle > zero(T)
            left = middle
        else
            right = middle
        end
        middle = (left + right) / convert(T,2)
    end
    # Return the eigenvalue
    right
end # bisect

function bisect( A::SymDPR1, side::Char )
    # COMPUTES: the leftmost (for side='L') or the rightmost (for side='R') eigenvalue
    # of a SymDPR1 matrix A = diagm(A.D) + A.r*A.u*(A.u)' by bisection.
    # RETURNS: the eigenvalue

    n=length(A.D)
    # Determine the starting interval for bisection, [left; right]
    indD=sortperm(A.D,rev=true)
    if A.r>0.0
        left, right = side == 'L' ? (A.D[indD[n]], A.D[indD[n-1]]) :  (A.D[indD[1]], A.D[indD[1]]+A.r*dot(A.u,A.u))
    else # rho<=0
        left, right = side == 'L' ? (A.D[indD[n]]+A.r*dot(A.u,A.u), A.D[indD[n]]) : (A.D[indD[2]], A.D[indD[1]])
    end

    # Bisection
    middle = (left + right)/2.0
    u2=A.u.^2
    n = length(A.D)
    while (right-left) > 2.0*eps()*maximum([abs(left),abs(right)])
        # Fmiddle = 1.0+A.r*sum(u2./(A.D-middle))
        Fmiddle=0.0
        for k=1:n
            Fmiddle=Fmiddle+u2[k]/(A.D[k]-middle)
        end
        Fmiddle=1.0+A.r*Fmiddle
        sign(A.r)*Fmiddle < 0.0 ? left = middle : right = middle
        middle = (left + right)/2.0
    end
    # Return the eigenvalue
    right
end # bisect


text:
using Distributions
using JSON, ForwardDiff, Calculus, PDMats # test dependencies
using Test
using Distributed
using Random
using StatsBase

const tests = [
    "truncate",
    "truncnormal",
    "truncated_exponential",
    "normal",
    "mvnormal",
    "mvlognormal",
    "types",
    "utils",
    "samplers",
    "categorical",
    "univariates",
    "continuous",
    "fit",
    "multinomial",
    "binomial",
    "poissonbinomial",
    "dirichlet",
    "dirichletmultinomial",
    "logitnormal",
    "mvtdist",
    "kolmogorov",
    "edgeworth",
    "matrix",
    "vonmisesfisher",
    "conversion",
    "mixture",
    "gradlogpdf",
    "noncentralt",
    "locationscale",
    "quantile_newton",
    "semicircle",
    "qq",
    "pgeneralizedgaussian",
    "product",
    "discretenonparametric",
    "functionals",
    "chernoff",
]

printstyled("Running tests:\n", color=:blue)

Random.seed!(345679)

# to reduce redundancy, we might break this file down into seperate `$t * "_utils.jl"` files
include("testutils.jl")

for t in tests
    @testset "Test $t" begin
        Random.seed!(345679)
        include("$t.jl")
    end
end

# print method ambiguities
println("Potentially stale exports: ")
display(Test.detect_ambiguities(Distributions))
println()


text:
export SG_Filter, maxDerivativeOrder, polynomialOrder, apply_SG_filter, apply_SG_filter2D

import Base: filter, length

function _Vandermonde(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)::Array{T,2}

    @assert halfWidth>=0
    @assert degree>=0

    x=T[i for i in -halfWidth:halfWidth]

    n = degree+1
    m = length(x)

    V = Array{T}(undef,m, n)

    for i = 1:m
        V[i,1] = T(1)
    end
    for j = 2:n
        for i = 1:m
            V[i,j] = x[i] * V[i,j-1]
        end
    end

    return V
end

# ================================================================

struct SG_Filter{T<:AbstractFloat,N}
    _filter_set::Array{LinearFilter_DefaultCentered{T,N},1}
end

"""
    function SG_Filter(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)

Creates a `SG_Filter` structure used to store Savitzky-Golay filters.

* filter length is 2*`halfWidth`+1
* polynomial degree is `degree`, which defines `maxDerivativeOrder`

You can apply these filters using the
* `apply_SG_filter`
* `apply_SG_filter2D`
functions.

Example:

```jldoctest
julia> sg = SG_Filter(halfWidth=5,degree=3);

julia> maxDerivativeOrder(sg)
3

julia> length(sg)
11

julia> filter(sg,derivativeOrder=2)
Filter(r=-5:5,c=[0.03497, 0.01399, -0.002331, -0.01399, -0.02098, -0.02331, -0.02098, -0.01399, -0.002331, 0.01399, 0.03497])

```
"""
function SG_Filter(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)::SG_Filter
    @assert degree>=0
    @assert halfWidth>=1
    @assert 2*halfWidth>degree

    V=_Vandermonde(T,halfWidth=halfWidth,degree=degree)
    Q,R=qr(V)
    # breaking change in Julia V1.0,
    # see https://github.com/JuliaLang/julia/issues/27397
    #
    # before Q was a "plain" matrix, now stored in compact form
    #
    # SG=R\Q'
    #
    # must be replaced by
    #
    # Q=Q*Matrix(I, size(V))
    # SG=R\Q'
    #
    Q=Q*Matrix(I, size(V))
    SG=R\Q'

    n_filter,n_coef = size(SG)

    buffer=Array{LinearFilter_DefaultCentered{T,n_coef},1}()

    for i in 1:n_filter
        SG[i,:]*=factorial(i-1)
        push!(buffer,LinearFilter(SG[i,:]))
    end

# Returns filters set
    return SG_Filter{T,n_coef}(buffer)
end

# ================================================================

"""
    function filter(sg::SG_Filter{T,N};derivativeOrder::Int=0)

Returns the filter to be used to compute the smoothed derivatives of order `derivativeOrder`.

See: `SG_Filter`
"""
function filter(sg::SG_Filter{T,N};derivativeOrder::Int=0) where {T<:AbstractFloat,N}
    @assert 0<= derivativeOrder <= maxDerivativeOrder(sg)
    return sg._filter_set[derivativeOrder+1]
end

"""
    function length(sg::SG_Filter{T,N})

Returns filter length, this is an odd number.

See: `SG_Filter`
"""
Base.length(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = length(filter(sg))

"""
    function maxDerivativeOrder(sg::SG_Filter{T,N})

Maximum order of the smoothed derivatives we can compute using `sg` filters.

See: `SG_Filter`
"""
maxDerivativeOrder(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = size(sg._filter_set,1)-1

"""
    function polynomialOrder(sg::SG_Filter{T,N})

Returns the degree of the polynomial used to construct the
Savitzky-Golay filters. This is mainly a 'convenience' function, as it
is equivalent to `maxDerivativeOrder`

See: `SG_Filter`
"""
polynomialOrder(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = maxDerivativeOrder(sg)

"""
    function apply_SG_filter(signal::Array{T,1},
                             sg::SG_Filter{T};
                             derivativeOrder::Int=0,
                             left_BE::Type{<:BoundaryExtension}=ConstantBE,
                             right_BE::Type{<:BoundaryExtension}=ConstantBE)

Applies an 1D Savitzky-Golay and returns the smoothed signal.
"""
function apply_SG_filter(signal::AbstractArray{T,1},
                         sg::SG_Filter{T};
                         derivativeOrder::Int=0,
                         left_BE::Type{<:BoundaryExtension}=ConstantBE,
                         right_BE::Type{<:BoundaryExtension}=ConstantBE) where {T<:AbstractFloat}

    return directCrossCorrelation(filter(sg,derivativeOrder=derivativeOrder),
                                  signal,
                                  left_BE,
                                  right_BE)
end

"""
    function apply_SG_filter2D(signal::Array{T,2},
                               sg_I::SG_Filter{T},
                               sg_J::SG_Filter{T};
                               derivativeOrder_I::Int=0,
                               derivativeOrder_J::Int=0,
                               min_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                               max_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                               min_J_BE::Type{<:BoundaryExtension}=ConstantBE,
                               max_J_BE::Type{<:BoundaryExtension}=ConstantBE)

Applies an 2D Savitzky-Golay and returns the smoothed signal.
"""
function apply_SG_filter2D(signal::AbstractArray{T,2},
                           sg_I::SG_Filter{T},
                           sg_J::SG_Filter{T};
                           derivativeOrder_I::Int=0,
                           derivativeOrder_J::Int=0,
                           min_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                           max_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                           min_J_BE::Type{<:BoundaryExtension}=ConstantBE,
                           max_J_BE::Type{<:BoundaryExtension}=ConstantBE) where {T<:AbstractFloat}

    return directCrossCorrelation2D(filter(sg_I,derivativeOrder=derivativeOrder_I),
                                    filter(sg_J,derivativeOrder=derivativeOrder_J),
                                    signal,
                                    min_I_BE,
                                    max_I_BE,
                                    min_J_BE,
                                    max_J_BE)
end

max_stars_repo_name: vincent-picaud/DirectConvolution

max_issues_repo_name: vincent-picaud/DirectConvolution

max_forks_repo_name: vincent-picaud/DirectConvolution


text:
include("../gcc_common.jl")
build_and_upload_gcc(v"11.1.0")

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
# Code to reproduce figure 1 numerical experiments for the paper "Understanding the Origins of Information-Seeking Exploratin in Probabilistic Objectives for Control"
# Essentially what we do here is demonstrate the difference between an evidence "maximizing" objective vs a divergence "probability matching" objective in a simple univariate scenario where we have two overlapping Gaussians
# as our "desire distribution". The evidence objective finds the peak of the distribution and concentrates all probability mass there. The divergence objective explicitly tries to match the desire distribution in all its modes.
# Even though both the desire and evidence objective can fit two Gaussians, the evidence objective nevertheless only utilizes one Gaussian and collapses the weight of the other to zero.

using Plots
using LinearAlgebra
using Statistics
using Distributions
using Optim

e = MathConstants.e

### Utility Function ###

function sigmoid(x)
    return 1 / (1 + exp(-x))
end
function Gaussian_cdf(x, mu, var)
    return (1 / sqrt(2 * pi * var)) * exp(-(x - mu)^2 / (2 * var))
end

function Gaussian_KL(mu1, mu2, var1, var2)
    log(sqrt(var2) / sqrt(var1)) + (var1 + (mu1 - mu2)^2) / (2 * var2) - 0.5
end

function Gaussian_entropy(mu, var)
    return log(sqrt(var * 2 * pi * e))
end

function approximate_KL(cdfs1, cdfs2)
    sum = 0
    for i in 1:length(cdfs1)
        sum += cdfs1[i] * log(cdfs1[i] / cdfs2[i])
    end
    return sum
end

function approximate_evidence(cdfs1, cdfs2)
    return sum(cdfs1 .* -log.(cdfs2))
end
function approximate_entropy(cdfs)
    return sum(cdfs .* log.(cdfs))
end

# Create our desire distribution cdf
function get_desired_cdf(mu_a=1, mu_b=4, var_1=1, var_b=0.4)
    xs = [i*0.01 for i in -500:1000]
    cdfs = (0.5 * Gaussian_cdf.(xs, mu_a, var_a)) + (0.5 * Gaussian_cdf.(xs, mu_b, var_b))
    return cdfs
end

### loss functions for evidence and divergence objective ###

function compute_divergence_loss(mu1, mu2, logvar1, logvar2, alpha,beta)
    lambda = 1e6
    desired_cdfs = get_desired_cdf()
    xs = [i*0.01 for i in -500:1000]
    cdfs = (sigmoid(alpha) .* Gaussian_cdf.(xs, mu1, exp(logvar1))) + (sigmoid(beta) .* Gaussian_cdf.(xs, mu2, exp(logvar2)))
    loss = approximate_KL(cdfs / sum(cdfs), desired_cdfs / sum(desired_cdfs))
    # ensure that constraint is respected
    loss += lambda * ((sigmoid(alpha) + sigmoid(beta)) - 1)^2
    return loss
end
function compute_evidence_loss(mu1, mu2, logvar1, logvar2, alpha,beta)
    lambda = 1e6
    desired_cdfs = get_desired_cdf()
    xs = [i*0.01 for i in -500:1000]
    cdfs = (sigmoid(alpha) .* Gaussian_cdf.(xs, mu1, exp(logvar1))) + (sigmoid(beta) .* Gaussian_cdf.(xs, mu2, exp(logvar2)))
    loss = approximate_evidence(cdfs / sum(cdfs), desired_cdfs / sum(desired_cdfs))
    # ensure that constraint is respected
    loss += lambda * ((sigmoid(alpha) + sigmoid(beta)) - 1)^2
    return loss
end

### Optimize using nonlinear optimization ###

function optimize_loss(loss_fn)
    init_params = zeros(6)
    init_params[1:4] = rand(4)
    print("$init_params \n")
    if loss_fn == "divergence"
        ll_func = compute_divergence_loss
    elseif loss_fn == "evidence"
        ll_func = compute_evidence_loss
    else
        error("Loss function not recognized")
    end
    opt = optimize(init_params -> ll_func(init_params...), init_params)
    output_params = Optim.minimizer(opt)
    final_loss = Optim.minimum(opt)
    return output_params,final_loss
end

### and plot ###

function plot_cdf(params,loss_fn)
    params, final_loss = optimize_loss(loss_fn)
    mu1, mu2, logvar1, logvar2, alpha,beta = params
    xs = [i*0.01 for i in -500:1000]
    desired_cdf = get_desired_cdf()
    cdfs = (sigmoid(alpha) .* Gaussian_cdf.(xs, mu1, exp(logvar1))) + (sigmoid(beta) .* Gaussian_cdf.(xs, mu2, exp(logvar2)))
    plot(xs ./ 10, cdfs/sum(cdfs), xaxis="X value", yaxis="Probability Density", label="Predicted Density")
    if loss_fn == "evidence"
        plot!(xs ./10, desired_cdf, title="Peak Finding with Evidence Objectives",linestyle=:dash,label="Desired Density")
    else
        plot!(xs ./10, desired_cdf / sum(desired_cdf),title="Probability Matching with Divergence objective",linestyle=:dash,label="Desired Density") # this is just for visibility
    end
    savefig("figures/$loss_fn")
    #return cdfs
end

plot_cdf(params,"evidence")
plot_cdf(params, "divergence")


text:
using NonconvexUtils, ForwardDiff, ReverseDiff, Tracker, Zygote
using Test, LinearAlgebra, SparseArrays, NLsolve, IterativeSolvers
using StableRNGs, ChainRulesCore, NonconvexCore

@testset "AbstractDiffFunction" begin
    global T = Nothing
    f = function (x)
        global T = eltype(x)
        return sum(x)
    end
    _f = ForwardDiffFunction(f)
    Zygote.gradient(_f, [1.0, 1.0])
    @test T <: ForwardDiff.Dual

    _f = AbstractDiffFunction(f, AD.ReverseDiffBackend())
    Zygote.gradient(_f, [1.0, 1.0])
    @test T <: ReverseDiff.TrackedReal

    _f = AbstractDiffFunction(f, AD.TrackerBackend())
    Zygote.gradient(_f, [1.0, 1.0])
    @test T <: Tracker.TrackedReal
end

@testset "TraceFunction" begin
    f = TraceFunction(sum, on_call = true)
    @test f.on_call == true
    @test f.on_grad == false
    f([2.0])
    @test f.trace == [(input = [2.0], output = 2.0)]
    Zygote.gradient(f, [3.0])
    @test f.trace == [(input = [2.0], output = 2.0)]

    f = TraceFunction(sum, on_grad = true)
    @test f.on_call == false
    @test f.on_grad == true
    f([2.0])
    @test f.trace == []
    Zygote.gradient(f, [3.0])
    @test f.trace == [(input = [3.0], output = 3.0, grad = [1.0])]

    f = TraceFunction(sum)
    @test f.on_call == true
    @test f.on_grad == true
    f = TraceFunction(sum, on_call = true, on_grad = true)
    @test f.on_call == true
    @test f.on_grad == true
    f([2.0])
    @test f.trace == [(input = [2.0], output = 2.0)]
    Zygote.gradient(f, [3.0])
    @test f.trace == [
        (input = [2.0], output = 2.0),
        (input = [3.0], output = 3.0, grad = [1.0]),
    ]

    f = TraceFunction(sum, on_call = false, on_grad = false)
    @test f.on_call == false
    @test f.on_grad == false
    f([2.0])
    @test f.trace == []
    Zygote.gradient(f, [3.0])
    @test f.trace == []
end

@testset "CustomGradFunction" begin
    fakeg = [2.0]
    f = CustomGradFunction(sum, x -> fakeg)
    @test Zygote.gradient(f, [1.0]) == (fakeg,)

    fakeJ = [1.0 2.0; 0.0 -1.0]
    f = CustomGradFunction(identity, x -> fakeJ)
    @test Zygote.jacobian(f, [1.0, 1.0]) == (fakeJ,)
end

@testset "CustomHessianFunction" begin
    fakeH = [3.0 -1.0; -1.0 2.0]
    f = CustomHessianFunction(sum, x -> fakeg, x -> fakeH)
    fakeg = [2.0, 2.0]
    @test Zygote.gradient(f, [1.0, 1.0]) == (fakeg,)
    @test Zygote.jacobian(x -> Zygote.gradient(f, x)[1], [1.0, 1.0]) == (fakeH,)

    hvp = (x, v) -> fakeH * v
    f = CustomHessianFunction(sum, x -> fakeg, hvp; hvp = true)
    H = Zygote.jacobian(x -> Zygote.gradient(f, x)[1], [1.0, 1.0])[1]
    @test norm(H - fakeH) < 1e-6
end

@testset "Implicit functions" begin
    @testset "Non-closure conditions" begin
        @testset "Vector input and output - Jacobian $jac - matrixfree $matrixfree" for jac in (false, true), matrixfree in (false, true)
            # Adapted from https://github.com/JuliaNLSolvers/NLsolve.jl/issues/205
            rng = StableRNG(123)
            nonlin = 0.1
            get_info_vec = N -> begin
                A = spdiagm(0 => fill(10.0, N), 1 => fill(-1.0, N-1), -1 => fill(-1.0, N-1))
                p0 = randn(rng, N)
                f = (p, x) -> A*x + nonlin*x.^2 - p
                solve_x = (p) -> begin
                    xstar = nlsolve(x -> f(p, x), zeros(N), method=:anderson, m=10).zero
                    return xstar, jac ? Zygote.jacobian(x -> f(p, x), xstar)[1] : nothing
                end
                g_analytic = gmres((A + Diagonal(2*nonlin*solve_x(p0)[1]))', ones(N))
                return solve_x, f, p0, g_analytic
            end

            solve_x, f, p0, g_analytic = get_info_vec(10)
            imf = ImplicitFunction(solve_x, f; matrixfree)
            obj = p -> sum(imf(p))
            g_auto = Zygote.gradient(obj, p0)[1]
            @test norm(g_analytic - g_auto) < 1e-6

            imf = ImplicitFunction(solve_x, f; matrixfree, tol = 0.0)
            obj = p -> sum(imf(p))
            g_auto = Zygote.gradient(obj, p0)[1]
            @test all(isnan.(g_auto))

            imf = ImplicitFunction(solve_x, f; matrixfree, tol = 0.0, error_on_tol_violation = true)
            obj = p -> sum(imf(p))
            @test_throws ArgumentError Zygote.gradient(obj, p0)[1]
        end

        @testset "Non-vector input and output" begin
            rng = StableRNG(123)
            nonlin = 0.1
            get_info_nonvec = N -> begin
                A = spdiagm(0 => fill(10.0, N), 1 => fill(-1.0, N-1), -1 => fill(-1.0, N-1))
                p0 = (a = randn(rng, N),)
                f = (p, x) -> A*x.a + nonlin*x.a.^2 - p.a
                solve_x = (p) -> begin
                    return (a = nlsolve(x -> f(p, (a = x,)), zeros(N), method=:anderson, m=10).zero,), nothing
                end
                g_analytic = (a = gmres((A + Diagonal(2*nonlin*solve_x(p0)[1].a))', ones(N)),)
                return solve_x, f, p0, g_analytic
            end

            for matrixfree in (false, true)
                solve_x, f, p0, g_analytic = get_info_nonvec(10)
                imf = ImplicitFunction(solve_x, f; matrixfree)
                obj = p -> sum(imf(p).a)
                g_auto = Zygote.gradient(obj, p0)[1]
                @test norm(g_analytic.a - g_auto.a) < 1e-6

                imf = ImplicitFunction(solve_x, f; matrixfree, tol = 0.0)
                obj = p -> sum(imf(p).a)
                g_auto = Zygote.gradient(obj, p0)[1]
                @test all(isnan.(g_auto.a))

                imf = ImplicitFunction(solve_x, f; matrixfree, tol = 0.0, error_on_tol_violation = true)
                obj = p -> sum(imf(p).a)
                @test_throws ArgumentError Zygote.gradient(obj, p0)[1]
            end
        end
    end

    @testset "Closure conditions" begin
        @testset "Vector input and output - Jacobian $jac - matrixfree $matrixfree" for jac in (false, true), matrixfree in (false, true)
            rng = StableRNG(123)
            nonlin = 0.1
            get_info_closure_vec = N -> begin
                A = spdiagm(0 => fill(10.0, N), 1 => fill(-1.0, N-1), -1 => fill(-1.0, N-1))
                p0 = randn(rng, N)
                f = (x) -> A*x + nonlin*x.^2 - p0
                solve_x = () -> begin
                    xstar = nlsolve(f, zeros(N), method=:anderson, m=10).zero
                    xstar, jac ? Zygote.jacobian(f, xstar)[1] : nothing
                end
                g_analytic = gmres((A + Diagonal(2*nonlin*solve_x()[1]))', ones(N))
                return p0, A, g_analytic
            end

            p0, A, g_analytic = get_info_closure_vec(10)
            obj = p -> begin
                N = 10
                # f closes over p
                f = (x) -> A*x + nonlin*x.^2 - p
                solve_x = () -> begin
                    return nlsolve(f, zeros(N), method=:anderson, m=10).zero, nothing
                end
                imf = ImplicitFunction(solve_x, f; matrixfree)
                return sum(imf())
            end
            g_auto = Zygote.gradient(obj, p0)[1]
            @test norm(g_analytic - g_auto) < 1e-6

            obj = p -> begin
                N = 10
                # f closes over p
                f = (x) -> A*x + nonlin*x.^2 - p
                solve_x = () -> begin
                    return nlsolve(f, zeros(N), method=:anderson, m=10).zero, nothing
                end
                imf = ImplicitFunction(solve_x, f; matrixfree, tol = 0.0)
                return sum(imf())
            end
            g_auto = Zygote.gradient(obj, p0)[1]
            @test all(isnan.(g_auto))

            obj = p -> begin
                N = 10
                # f closes over p
                f = (x) -> A*x + nonlin*x.^2 - p
                solve_x = () -> begin
                    return nlsolve(f, zeros(N), method=:anderson, m=10).zero, nothing
                end
                imf = ImplicitFunction(solve_x, f; matrixfree, tol = 0.0, error_on_tol_violation = true)
                return sum(imf())
            end
            @test_throws ArgumentError Zygote.gradient(obj, p0)[1]
        end

        @testset "Non-vector input and output" begin
            rng = StableRNG(123)
            nonlin = 0.1
            get_info_closure_nonvec = N -> begin
                N = 10
                A = spdiagm(0 => fill(10.0, N), 1 => fill(-1.0, N-1), -1 => fill(-1.0, N-1))
                p0 = (a = randn(rng, N),)
                _f = (x) -> A*x.a + nonlin*x.a.^2 - p0.a
                _solve_x = () -> begin
                    return (a = nlsolve(x -> _f((a = x,)), zeros(N), method=:anderson, m=10).zero,), nothing
                end
                return p0, A, (a = gmres((A + Diagonal(2*nonlin*_solve_x()[1].a))', ones(N)),)
            end

            for matrixfree in (false, true)
                p0, A, g_analytic = get_info_closure_nonvec(10)
                obj = p -> begin
                    N = 10
                    # f closes over p
                    f = (x) -> A*x.a + nonlin*x.a.^2 - p.a
                    solve_x = () -> begin
                        return (a = nlsolve(x -> f((a = x,)), zeros(N), method=:anderson, m=10).zero,), nothing
                    end
                    imf = ImplicitFunction(solve_x, f; matrixfree)
                    return sum(imf().a)
                end
                g_auto = Zygote.gradient(obj, p0)[1]
                @test norm(g_analytic.a - g_auto.a) < 1e-6

                obj = p -> begin
                    N = 10
                    # f closes over p
                    f = (x) -> A*x.a + nonlin*x.a.^2 - p.a
                    solve_x = () -> begin
                        return (a = nlsolve(x -> f((a = x,)), zeros(N), method=:anderson, m=10).zero,), nothing
                    end
                    imf = ImplicitFunction(solve_x, f; matrixfree, tol = 0.0)
                    return sum(imf().a)
                end
                g_auto = Zygote.gradient(obj, p0)[1]
                @test all(isnan.(g_auto.a))

                obj = p -> begin
                    N = 10
                    # f closes over p
                    f = (x) -> A*x.a + nonlin*x.a.^2 - p.a
                    solve_x = () -> begin
                        return (a = nlsolve(x -> f((a = x,)), zeros(N), method=:anderson, m=10).zero,), nothing
                    end
                    imf = ImplicitFunction(solve_x, f; matrixfree, tol = 0.0, error_on_tol_violation = true)
                    return sum(imf().a)
                end
                @test_throws ArgumentError Zygote.gradient(obj, p0)[1]
            end
        end
    end
end

# TODO tests:
## Multiple outputs
## Struct output
## Functor f - fix first
@testset "ForwardDiff frule" begin
    @eval begin
        f1(x, y) = x + y
        global frule_count = 0
        function ChainRulesCore.frule((_, Δx1, Δx2), ::typeof(f1), x1, x2)
            global frule_count += 1
            println("frule was called")
            return f1(x1, x2), Δx1 + Δx2
        end
        NonconvexUtils.@ForwardDiff_frule f1(x1::ForwardDiff.Dual, x2::ForwardDiff.Dual)
        NonconvexUtils.@ForwardDiff_frule f1(x1::AbstractVector{<:ForwardDiff.Dual}, x2::AbstractVector{<:ForwardDiff.Dual})
        NonconvexUtils.@ForwardDiff_frule f1(x1::AbstractMatrix{<:ForwardDiff.Dual}, x2::AbstractMatrix{<:ForwardDiff.Dual})

        f2(x::NamedTuple, y::NamedTuple) = (a = x.a + y.a, b = x.b + y.b)
        f2(x::AbstractVector, y::AbstractVector) = f2.(x, y)
        function ChainRulesCore.frule((_, Δx1, Δx2), ::typeof(f2), x1::NamedTuple, x2::NamedTuple)
            global frule_count += 1
            println("frule was called")
            return f2(x1, x2), (a = Δx1.a + Δx2.a, b = Δx1.b + Δx2.b)
        end
        NonconvexUtils.@ForwardDiff_frule f2(x1::NamedTuple{<:Any, <:Tuple{Vararg{<:ForwardDiff.Dual}}}, x2::NamedTuple{<:Any, <:Tuple{Vararg{<:ForwardDiff.Dual}}})

        struct MyStruct{T, T1, T2}
            a::T1
            b::T2
        end
        MyStruct(a, b) = MyStruct{typeof(a), typeof(a), typeof(b)}(a, b)

        # The @constructor macro takes the type (first) and constructor function (second)
        # The constructor function takes input the fields generated from ntfromstruct (as multiple positional arguments)
        # The ntfromstruct function can be overloaded for your type
        NonconvexCore.@constructor MyStruct MyStruct

        f2(x::MyStruct, y::MyStruct) = MyStruct(x.a + y.a, x.b + y.b)
        function ChainRulesCore.frule((_, Δx1, Δx2), ::typeof(f2), x1::MyStruct, x2::MyStruct)
            global frule_count += 1
            println("frule was called")
            return f2(x1, x2), MyStruct(Δx1.a + Δx2.a, Δx1.b + Δx2.b)
        end
        NonconvexUtils.@ForwardDiff_frule f2(x1::MyStruct{<:ForwardDiff.Dual}, x2::MyStruct{<:ForwardDiff.Dual})
        Base.sum(s::MyStruct) = s.a + s.b

        # I recommend creating your own type to avoid piracy
        _eigvals!(x) = eigvals!(x)
        function ChainRulesCore.frule((_, Δx), ::typeof(_eigvals!), x::Symmetric{<:Real})
            global frule_count += 1
            println("frule was called")
            return frule((NoTangent(), Δx), eigvals!, x)
        end

        # I recommend creating your own type to avoid piracy
        NonconvexCore.@constructor Symmetric Symmetric
        import NamedTupleTools: ntfromstruct, structfromnt
        ntfromstruct(a::Symmetric) = (data = a.data,)
        structfromnt(::Type{Symmetric}, x::NamedTuple) = Symmetric(x.data, :U)
        NonconvexUtils.@ForwardDiff_frule _eigvals!(A::Symmetric{<:ForwardDiff.Dual})
    end
    global frule_count = 0
    @testset "2 real inputs - 1 real output" begin
        _f = x -> f1(x[1], x[2])
        _f(rand(2))
        g1 = ForwardDiff.gradient(_f, rand(2))
        @test frule_count == 2
        cfg = ForwardDiff.GradientConfig(_f, rand(2), ForwardDiff.Chunk{2}())
        g2 = ForwardDiff.gradient(_f, rand(2), cfg)
        @test frule_count == 4
        g3 = Zygote.gradient(_f, rand(2))[1]
        @test g1 == g2 == g3
    end
    frule_count = 0
    @testset "2 vector inputs - 1 real output" begin
        _f = x -> sum(f1(x[1:2], x[3:4]))
        g1 = ForwardDiff.gradient(_f, rand(4))
        @test frule_count == 4
        cfg = ForwardDiff.GradientConfig(_f, rand(4), ForwardDiff.Chunk{2}())
        g2 = ForwardDiff.gradient(_f, rand(4), cfg)
        @test frule_count == 8
        g3 = Zygote.gradient(_f, rand(4))[1]
        @test g1 == g2 == g3
    end
    frule_count = 0
    @testset "2 vector inputs - 1 vector output" begin
        _f = x -> f1(x[1:2], x[3:4])
        j1 = ForwardDiff.jacobian(_f, rand(4))
        @test frule_count == 4
        cfg = ForwardDiff.JacobianConfig(_f, rand(4), ForwardDiff.Chunk{2}())
        j2 = ForwardDiff.jacobian(_f, rand(4), cfg)
        @test frule_count == 8
        j3 = Zygote.jacobian(_f, rand(4))[1]
        @test j1 == j2 == j3
    end
    frule_count = 0
    @testset "2 matrix inputs - 1 real output" begin
        _f = x -> sum(f1(x[1:2,1:2], x[3:4,3:4]))
        g1 = ForwardDiff.gradient(_f, rand(4, 4))
        @test frule_count == 16
        cfg = ForwardDiff.GradientConfig(_f, rand(4, 4), ForwardDiff.Chunk{2}())
        g2 = ForwardDiff.gradient(_f, rand(4, 4), cfg)
        @test frule_count == 32
        g3 = Zygote.gradient(_f, rand(4, 4))[1]
        @test g1 == g2 == g3
    end
    frule_count = 0
    @testset "2 NamedTuple inputs - 1 real output" begin
        _f = x -> sum(f2((a = x[1], b = x[2]), (a = x[3], b = x[4])))
        g1 = ForwardDiff.gradient(_f, rand(4))
        @test frule_count == 4
        cfg = ForwardDiff.GradientConfig(_f, rand(4), ForwardDiff.Chunk{2}())
        g2 = ForwardDiff.gradient(_f, rand(4))
        @test frule_count == 8
        g3 = Zygote.gradient(_f, rand(4))[1]
        @test g1 == g2 == g3
    end
    frule_count = 0
    @testset "2 struct inputs - 1 real output" begin
        _f = x -> sum(f2(MyStruct(x[1], x[2]), MyStruct(x[3], x[4])))
        g1 = ForwardDiff.gradient(_f, rand(4))
        @test frule_count == 4
        cfg = ForwardDiff.GradientConfig(_f, rand(4), ForwardDiff.Chunk{2}())
        g2 = ForwardDiff.gradient(_f, rand(4))
        @test frule_count == 8
        g3 = Zygote.gradient(_f, rand(4))[1]
        @test g1 == g2 == g3
    end
    frule_count = 0
    @testset "eigvals" begin
        # Gradient of trace
        g = ForwardDiff.gradient(x -> sum(_eigvals!(Symmetric(x))), rand(4, 4))
        @test frule_count == 16
        @test norm(g - I) < 1e-6
    end
end