text:
using DataStructures: DefaultDict

struct Reagent
    name::String
    quantity::Int64
end

struct Reaction
    inputs::Array{Reagent}
    output::Reagent
end

function parse_reagent(s)::Reagent
    components = split(s, " ")
    return Reagent(last(components), parse(Int32, first(components)))
end

function parse_reaction(s)::Reaction
    components = split(s, " => ")
    return Reaction(parse_reagent.(split(first(components), ", ")), parse_reagent(last(components)))
end

function min_ore_required(outputs, output_name, n, leftover = DefaultDict(0))
    if output_name == "ORE" return n end

    # Use leftovers before making more
    if output_name ∈ keys(leftover)
        if leftover[output_name] >= n
            leftover[output_name] -= n
            return 0
        else
            n -= leftover[output_name]
            leftover[output_name] = 0
        end
    end

    # Produce at least n of the output, putting the rest in leftover dict
    r = outputs[output_name]
    multiplier = Int64(ceil(n / r.output.quantity))
    leftover[r.output.name] += r.output.quantity * multiplier - n
    return sum([min_ore_required(outputs, i.name, i.quantity * multiplier, leftover) for i in r.inputs])
end

"""Return the largest value x in [lower, upper] where f(x) is true, assuming f(lower) and !f(upper)"""
function binary_search(lower, upper, f)
    midpoint = (lower + upper) ÷ 2
    if midpoint == lower return lower end
    return f(midpoint) ? binary_search(midpoint, upper, f) : binary_search(lower, midpoint, f)
end

function main()
    # Part 1: read reaction spec and compute min ore requirement by walking up the tree recursively
    reactions = parse_reaction.(open(f -> readlines(f), "day14_input.txt"))
    outputs = Dict(r.output.name => r for r in reactions)
    ore = min_ore_required(outputs, "FUEL", 1)
    println("Min ore required for 1 FUEL: $ore")

    # Part 2: Fuel available with 1 trillion ore. Quick binary search?
    trillion = 1000000000000
    fuel = binary_search(trillion ÷ ore, 2trillion ÷ ore, x -> (min_ore_required(outputs, "FUEL", x) <= trillion))
    println("Fuel for 1 trillion ore: $fuel")
end
main()

max_stars_repo_name: ericgreveson/adventofcode

max_issues_repo_name: ericgreveson/adventofcode

max_forks_repo_name: ericgreveson/adventofcode


text:
"""
Generate bindings for SDL from the C headers

## Outline

Download and extract source
Generate bindings with Clang.jl
    Running on linux gives fewer warnings and better output than windows
Remove SDL_ prefix
Compare with original
    Some of these are manual changes, some will be due to Clang.jl changes. Some might be
    deliberate configuration differences.

    Differences in SDL.jl
        They manually changed a bunch of types from Cint to more specific things.
        I don't know whether that was required to get things to run or whether they
        just wanted the documentation.

            Probably required. Running on an OS where Clang can find the
            stdlibs (linux) has defined them correctly.

        They have type annotations in the functions signature most of the time and I don't

            TODO: check if ccall helps us out more if we don't have specific
            types in the funcs.

        They manually added a macro

        There are a few more functions in mine

        There are a load of _mm_* functions in theirs

        Some functions in mine take and return no arguments when they definitely should
        (iconv, writeu8, utf8strlcpy)

            This was fixed by running on linux

    Differences in SDL_h.jl
        Mine uses @cenum, and that makes the diff a huge mess. I'll work on it another time.

        They've made some manual changes to Events (mostly just adding AbstractEvent as a parent)

        They've made empty structs for some CVoid types (mine just sets typename = CVoid

        They've capitalised some type names (mutex, sem) - seems like a bad idea? Names should be consistent, probably.

## Next

Done: Run the tests with these bindings.

Done: Port the remaining macros that NHDaly did:
 - https://github.com/jonathanBieler/SimpleDirectMediaLayer.jl/commit/ee594583aec7def27fa76ce9eee3dbce5cb9dd23#diff-16524f24c16a5c6939df897d14902bf7
 - https://github.com/jonathanBieler/SimpleDirectMediaLayer.jl/commit/f5ca6b95de65e514a86d6d6ae2bdc9ce14cbb990#diff-16524f24c16a5c6939df897d14902bf7

Try without removing the leading SDL_ ?

Filter the shit macros with clang

## Future work

Once the experimentation is over, the downloading and extracting code can go
and we can just vendor the includes from the right SDL2 version, maybe keeping
a script for downloading and extracting them for the future. Maybe not.

Decide if the opengl, opengles, vulkan, metal, headers should be excluded or not.

Try running your examples with these bindings.

Maybe parse and rewrite the files rather than string mangling.

Find out why `using SimpleDirectMediaLayer` is slow:

    `julia -e using pkgname`

    time(s) pkgname
    6.7     SimpleDirectMediaLayer
    6.6     SDL2_jll, SDL2_image_jll, SDL2_mixer_jll, SDL2_ttf_jll, ColorTypes
    5.7     SDL2_jll, SDL2_image_jll, SDL2_mixer_jll, SDL2_ttf_jll
    3.5     SDL2_jll
    4.5     SDL2_image_jll
    1.2     ColorTypes

    3.0     Strs
    1.5     ThreadsX

    Suggests we can only save time by loading the jlls lazily or by modifying
    the dependencies directly.

    Don't bother: the jlls are known to load slowly, but decent improvements are incoming.

### How to integrate manual changes

e.g. I'd want to edit SDL_Event and rename WindowShowMode

If the diffs are stable enough across versions I could just manually merge. Let's test that...
    Functions move up and down in an unpredictable way between 2.0.5 and 2.0.12, so lexical sorting at least within each file would probably be sensible.

### SDL_ prefix issue

Literally every function has an SDL_ prefix, so we can remove those fine.

With the structures and constants, however, there's a mix and we'll get (rare)
collisions if we just automatically remove all prefixes (WindowShowMode and
SDL_WindowShowMode are both defined, for example).

### Low priority aesthetic issues

Some types have lowercase names (mutex, sem, semaphore) (Maybe shouldn't fix this?)

Some types are aliased (Uint64 - UInt64, sem - semaphore, ...)

There are docstrings in the headers that we could extract. This site has them, but it's not very ergonomic.
https://happiness_follows.gitee.io/doc_translate/SDL2-2.0.7/html/annotated.html

doxygen is happy to output XML, but the XML is in a shit format (attrs, not nodes, plz), so I've just scraped the strings direct from the source code and added them after the fact.

Clang.jl wrap!() is the method that adds expressions to the buffer, but it's probably smarter to use the `rewrite` hook. rewrite accepts a vector of expressions.

"""

const dir = @__DIR__
cd(dir)

# Latest
sdl2_version = "2.0.12"

# The version in the comment of SDL.jl
# sdl2_version = "2.0.5"

urls = [
    :sdl2 => "http://www.libsdl.org/release/SDL2-$sdl2_version.tar.gz"
    # :sdl2_image => "https://www.libsdl.org/projects/SDL_image/release/SDL2_image-2.0.5.tar.gz"
    # :sdl2_mixer => "https://www.libsdl.org/projects/SDL_mixer/release/SDL2_mixer-2.0.4.tar.gz"
    # :sdl2_ttf => "https://www.libsdl.org/projects/SDL_ttf/release/SDL2_ttf-2.0.15.tar.gz"
]

# Get sources
mkpath("downloads")
for (_, url) in urls
    if !isfile(joinpath("downloads", basename(url)))
        download(url, joinpath("downloads", basename(url)))
    end
end

# Unpack their include dirs
using DataDeps: unpack

# Let's just do this one at a time for now.

isdir("SDL2-$sdl2_version") ||
    unpack("downloads/SDL2-$sdl2_version.tar.gz"; keep_originals=true)

### Generate bindings

# This is copied from the Clang.jl readme and then I changed the obvious bits.
using Clang: CLANG_INCLUDE, InternalOptions, init

# On windows, Clang will complain that it can't find a load of headers and this
# genuinely causes problems. Set a proper include or compile on a system where
# clang is more likely to find the headers.
LIBSDL_INCLUDE = "SDL2-$sdl2_version/include/"

# readdir() sorts the names, and Clang seems to generate the bindinds
# sequentially in this order, so that should make our diffs easy to read
# between versions
LIBSDL_HEADERS = [joinpath(LIBSDL_INCLUDE, header) for header in readdir(LIBSDL_INCLUDE) if endswith(header, ".h")]

# I think Jonathan excluded these, too.
filter!(!fn -> occursin("test", fn), LIBSDL_HEADERS)
filter!(!fn -> occursin("opengl", fn), LIBSDL_HEADERS)
filter!(!fn -> occursin("config", fn), LIBSDL_HEADERS)
filter!(!fn -> occursin("SDL_log", fn), LIBSDL_HEADERS)

bindings_dir = joinpath(dir, "bindings")
outputf = joinpath(bindings_dir, "libsdl2_funcs.jl")
commonf = joinpath(bindings_dir, "libsdl2_types.jl")

mkpath(bindings_dir)

wc = init(; headers = LIBSDL_HEADERS,
            output_dir = bindings_dir,
            output_file = outputf,
            common_file = commonf,
            clang_includes = vcat(LIBSDL_INCLUDE, CLANG_INCLUDE),
            #= clang_args = ["-I", joinpath(STDLIB_INCLUDE, "..")], =#
            # I don't know what this does, but deleting it gives us an error.
            header_wrapped = (root, current) -> root == current,
            header_library = x -> "libsdl2",
            # Generate mutable structs
            options = InternalOptions(true, true),
            )

run(wc)

### Remove SDL_ prefix

function rm_SDL_(fn)
    str = String(read(fn))
    str = replace(str, r"(::|[^:])SDL_" => s"\1");
    write(fn, str)
end

rm_SDL_(outputf)
rm_SDL_(commonf)

#=

### Compare with original

"""
Give diff a chance to work by sorting the blocks

Pretty useless with the `common` output because of the enums.

"""
function normjl(filename)
    lines = readlines(filename)
    lines = filter(s -> !startswith(s, "#"), lines)

    funcs = split(join(lines, "\n"), "\n\n")
    funcs = strip.(funcs)
    funcs = filter(!isempty, funcs)
    # funcs = replace.(funcs, r"([^:])SDL_" => s"\1");
    funcs = sort(funcs)

    str = join(funcs, "\n\n")
    return str
end

write("bindings/mine.jl", normjl(SDL_jl))
write("bindings/theirs.jl", normjl("../src/lib/SDL.jl"))

write("bindings/mine_h.jl", normjl(SDL_h_jl))
write("bindings/theirs_h.jl", normjl("../src/lib/SDL_h.jl"))

=#

ext: jl


text:
import DifferentialEquations as DE
import Sundials
using DifferentialEquations.DiffEqBase: ODEProblem, AbstractODEAlgorithm
import DifferentialEquations.DiffEqBase: solve

export FixedSparseFsp, AdaptiveFspSparse, solve
include("fspoutput.jl")
include("spaceadapters.jl")

function solve(model::CmeModel,
    initial_distribution::FspVectorSparse{NS,IntT,RealT},
    tspan::Union{Vector,Tuple},
    ode_method::Union{Nothing,AbstractODEAlgorithm};
    saveat = [], fsptol::AbstractFloat = 1.0E-6, odeatol::AbstractFloat = 1.0E-10, odertol::AbstractFloat = 1.0E-4) where {NS,IntT<:Integer,RealT<:AbstractFloat}

    statespace = StateSpaceSparse(model.stoich_matrix, initial_distribution.states)
    sink_count = get_sink_count(statespace)
    p0 = initial_distribution.values

    A = FspMatrixSparse(statespace, model.propensities, parameters = model.parameters)
    u0 = [p0; zeros(sink_count)]
    function odefun!(du, u, θ, t)
        matvec!(du, t, A, u)
        nothing
    end
    fspprob = ODEProblem(odefun!, u0, tspan, p = model.parameters)
    solutions = DE.solve(fspprob, ode_method, atol = odeatol, rtol = odertol, saveat = saveat)

    output = FspOutputSparse{NS,IntT,RealT}(
        t = Vector{RealT}(),
        p = Vector{FspVectorSparse{NS,IntT,RealT}}(),
        sinks = Vector{Vector{RealT}}()
    )
    for (t, u) in zip(solutions.t, solutions.u)
        push!(output.t, t)
        push!(output.p, FspVectorSparse(statespace.states, u[1:end-sink_count]))
        push!(output.sinks, u[end-sink_count+1:end])
    end

    return output
end

"""
   $(TYPEDEF)

Struct to store adaptive Finite State Projection algorithmic options. This type is intended to work with `solve()` methods that output CME solutions as sparse vectors of type `FspVectorSparse`.

# Fields

    ode_method
Instance of an ODE algorithm from the `DifferentialEquations.jl` package.

    space_adapter
Method to adapt the FSP state space when truncation error exceeds user-specified tolerance.

# See also
[`RStepAdapter`](@ref), [`FspVectorSparse`](@ref), [`CmeModel`](@ref)
"""
Base.@kwdef mutable struct AdaptiveFspSparse <: TransientCmeAlgorithm
    ode_method::Union{Nothing,AbstractODEAlgorithm}
    space_adapter::AbstractSpaceAdapterSparse
end

"""
    $(TYPEDSIGNATURES)

Numerical integration for the chemical master equation using an adaptive Finite State Projection (FSP) algorithm. This method is adaptive, meaning that states will be added during the course of integration to ensure the approximation error is below the user-specified tolerance.
In addition, depending on the input `fspalgorithm`, the space adapter may delete states with low probabilities before expanding the state space. This ensures not only accuracy but also efficiency of the CME integration.

# Arguments

    model::CmeModel
Chemical Master Equation model.

    initial_distribution::FspVectorSparse
Initial distribution.

    tspan::Tuple{AbstractFloat, AbstractFloat}
Timespan for the integration. The first element is the starting time, the second element is the end time of the integration.

    fspalgorithm::AdaptiveFspSparse
An instance of `AdaptiveFspSparse`, storing information about the specific adaptive FSP method.

    saveat::Vector{AbstractFloat} (optional)
Time points to store the FSP solutions at. Default is `[]`, which means all solutions at every timestep will be stored.

    fsptol (optional)
FSP truncation error tolerance. The FSP state space is adapted so that the output solution vectors sum to greater than or equal to `1-fsptol`. Default: 1.0E-6.

    odertol (optional)
Relative error tolerance for the ODE solver of the ODEs resulting from FSP truncation. Default: 1.0E-4.

    odeatol (optional)
Relative error tolerance for the ODE solver of the ODEs resulting from FSP truncation. Default: 1.0E-10.

    verbose (optional)
Whether to output status when updating the state space. Default: false.

# Returns
An instance of `FspOutputSparse`.

# See also
[`CmeModel`](@ref), [`AdaptiveFspSparse`](@ref), [`FspOutputSparse`](@ref).
"""
function solve(model::CmeModel,
    initial_distribution::FspVectorSparse{NS,IntT,RealT},
    tspan::Tuple{AbstractFloat,AbstractFloat},
    fspalgorithm::AdaptiveFspSparse; saveat = [], fsptol::AbstractFloat = 1.0E-6,
    odeatol::AbstractFloat = 1.0E-10,
    odertol::AbstractFloat = 1.0E-4,
    verbose::Bool = false) where {NS,IntT<:Integer,RealT<:AbstractFloat}

    tstart = min(tspan...)
    tend = max(tspan...)
    adapter = fspalgorithm.space_adapter

    p0 = deepcopy(initial_distribution.values)
    statespace = StateSpaceSparse(model.stoich_matrix, initial_distribution.states)
    sink_count = get_sink_count(statespace)
    init!(statespace, adapter, p0, tstart, fsptol)

    tnow = tstart
    unow = [p0; zeros(sink_count)]
    A = FspMatrixSparse{RealT}(statespace, model.propensities, parameters = model.parameters)

    output = FspOutputSparse{NS,IntT,RealT}(
        t = Vector{RealT}(),
            p = Vector{FspVectorSparse{NS,IntT,RealT}}(),
            sinks = Vector{Vector{RealT}}()
        )
    while tnow < tend
        # Set up callback for checking the growth of FSP error over time
        function fsprhs!(du, u, θ, t)
            matvec!(du, t, A, u)
        end
        function affect!(integrator)
            DE.terminate!(integrator)
        end
        # Error constraint for the intermediate solutions. The intermediate sinks must not grow beyond a linear function of time that reaches `fsptol` at the end time
        function fsp_error_constraint(u, t, integrator)
            sinks = u[end-sink_count+1:end]
            return sum(sinks) - fsptol * t / tend
        end

        fsp_cb = DE.ContinuousCallback(
            fsp_error_constraint,
            affect!,
            save_positions = (false, false),
            interp_points = 100,
            abstol = eps()
            )

        fspprob = DE.ODEProblem(fsprhs!, unow, (tnow, tend), p = model.parameters)
        integrator = DE.init(fspprob, fspalgorithm.ode_method, atol = odeatol, rtol = odertol, callback = fsp_cb, saveat = saveat)

        DE.step!(integrator, tend - tnow, true)

        for (t, u) in zip(integrator.sol.t, integrator.sol.u)
            push!(output.t, t)
            push!(output.p, FspVectorSparse(statespace.states, u[1:end-sink_count]))
            push!(output.sinks, u[end-sink_count+1:end])
        end

        tnow = integrator.t
        # If the ODE integrator exits before reaching `tend`, it means the intermediate solution with current state space has reached error tolerance threshold
        if tnow < tend
            p = integrator.u[1:end-sink_count]
            sinks = integrator.u[end-sink_count+1:end]
            # Call adapter to modify the state space and shrink/expand the solution vector as appropriate
            adapt!(statespace, adapter, p, sinks, tnow, tend, fsptol;integrator=integrator)
            A = FspMatrixSparse{RealT}(statespace, model.propensities, parameters = get_parameters(model))

            # This numerical trick ensures that DifferentialEquations.jl's event detection does not miss the event the updated state space is insufficient
            if sum(sinks) >= tnow*fsptol/tend
                sinks .-= eps()
            end

            # Extended solution vector to start the ODE integrator again
            unow = [p; sinks]

            if verbose
                println("t = $(round(tnow, digits=2)). Update state space. New size: $(get_state_count(statespace)).")
            end
        else # Otherwise, add the final slice to the output
            u = integrator.u
            push!(output.t, tnow)
            push!(output.p, FspVectorSparse(statespace.states, u[1:end-sink_count]))
            push!(output.sinks, u[end-sink_count+1:end])
        end
    end
    return output
end


text:
@testset "5 Inverse trig functions" begin
    include("5.1 Inverse sine/.jl")
    include("5.2 Inverse cosine/.jl")
    include("5.3 Inverse tangent/.jl")
    include("5.4 Inverse cotangent/.jl")
    include("5.5 Inverse secant/.jl")
    include("5.6 Inverse cosecant/.jl")
end

max_stars_repo_path: test/5 Inverse trig functions/.jl

max_issues_repo_path: test/5 Inverse trig functions/.jl

max_forks_repo_path: test/5 Inverse trig functions/.jl


text:
###
### DynamicHMC backend - https://github.com/tpapp/DynamicHMC.jl
###

"""
    DynamicNUTS

Dynamic No U-Turn Sampling algorithm provided by the DynamicHMC package.

To use it, make sure you have DynamicHMC package (version >= 2) loaded:
```julia
using DynamicHMC
```
"""
struct DynamicNUTS{AD,space} <: Hamiltonian{AD} end

DynamicNUTS(args...) = DynamicNUTS{ADBackend()}(args...)
DynamicNUTS{AD}(space::Symbol...) where AD = DynamicNUTS{AD, space}()

DynamicPPL.getspace(::DynamicNUTS{<:Any, space}) where {space} = space

struct DynamicHMCLogDensity{M<:Model,S<:Sampler{<:DynamicNUTS},V<:AbstractVarInfo}
    model::M
    sampler::S
    varinfo::V
end

function DynamicHMC.dimension(ℓ::DynamicHMCLogDensity)
    return length(ℓ.varinfo[ℓ.sampler])
end

function DynamicHMC.capabilities(::Type{<:DynamicHMCLogDensity})
    return DynamicHMC.LogDensityOrder{1}()
end

function DynamicHMC.logdensity_and_gradient(
    ℓ::DynamicHMCLogDensity,
    x::AbstractVector,
)
    return gradient_logp(x, ℓ.varinfo, ℓ.model, ℓ.sampler)
end

"""
    DynamicNUTSState

State of the [`DynamicNUTS`](@ref) sampler.

# Fields
$(TYPEDFIELDS)
"""
struct DynamicNUTSState{V<:AbstractVarInfo,C,M,S}
    vi::V
    "Cache of sample, log density, and gradient of log density."
    cache::C
    metric::M
    stepsize::S
end

function gibbs_update_state(state::DynamicNUTSState, varinfo::AbstractVarInfo)
    return DynamicNUTSState(varinfo, state.cache, state.metric, state.stepsize)
end

DynamicPPL.initialsampler(::Sampler{<:DynamicNUTS}) = SampleFromUniform()

function DynamicPPL.initialstep(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:DynamicNUTS},
    vi::AbstractVarInfo;
    kwargs...
)
    # Ensure that initial sample is in unconstrained space.
    if !DynamicPPL.islinked(vi, spl)
        DynamicPPL.link!(vi, spl)
        model(rng, vi, spl)
    end

    # Perform initial step.
    results = DynamicHMC.mcmc_keep_warmup(
        rng,
        DynamicHMCLogDensity(model, spl, vi),
        0;
        initialization = (q = vi[spl],),
        reporter = DynamicHMC.NoProgressReport(),
    )
    steps = DynamicHMC.mcmc_steps(results.sampling_logdensity, results.final_warmup_state)
    Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q)

    # Update the variables.
    vi[spl] = Q.q
    DynamicPPL.setlogp!(vi, Q.ℓq)

    # If a Gibbs component, transform the values back to the constrained space.
    if spl.selector.tag !== :default
        DynamicPPL.invlink!(vi, spl)
    end

    # Create first sample and state.
    sample = Transition(vi)
    state = DynamicNUTSState(vi, Q, steps.H.κ, steps.ϵ)

    return sample, state
end

function AbstractMCMC.step(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:DynamicNUTS},
    state::DynamicNUTSState;
    kwargs...
)
    # Compute next sample.
    vi = state.vi
    ℓ = DynamicHMCLogDensity(model, spl, vi)
    steps = DynamicHMC.mcmc_steps(
        rng,
        DynamicHMC.NUTS(),
        state.metric,
        ℓ,
        state.stepsize,
    )
    Q = if spl.selector.tag !== :default
        # When a Gibbs component, transform values to the unconstrained space
        # and update the previous evaluation.
        DynamicPPL.link!(vi, spl)
        DynamicHMC.evaluate_ℓ(ℓ, vi[spl])
    else
        state.cache
    end
    newQ, _ = DynamicHMC.mcmc_next_step(steps, Q)

    # Update the variables.
    vi[spl] = newQ.q
    DynamicPPL.setlogp!(vi, newQ.ℓq)

    # If a Gibbs component, transform the values back to the constrained space.
    if spl.selector.tag !== :default
        DynamicPPL.invlink!(vi, spl)
    end

    # Create next sample and state.
    sample = Transition(vi)
    newstate = DynamicNUTSState(vi, newQ, state.metric, state.stepsize)

    return sample, newstate
end


text:
using DaemonMode
using Test
using Sockets

@testset "Start Server" begin
    port = 3001
    @test_throws Base.IOError connect(port)
    task = @async serve(port)
    sleep(1)
    sendExitCode(port)
    wait(task)
end

@testset "runFile" begin
    port = 3002
    task = @async serve(port)
    sleep(1)
    buffer = IOBuffer()
    files = ["hello.jl", "hello2.jl", "print.jl"]
    outputs = ["Hello, World!\n", "Hello, World\n\nBye, World!\n", "Hello, World"]

    for (file, out) in zip(files, outputs)
        @test isfile(file)
        runfile(file, output=buffer, port=port)
        output = String(take!(buffer))
        @test output == out
    end

    sendExitCode(port)
    wait(task)
end

@testset "runExpr" begin
    port = 3003
    task = @async serve(port)
    sleep(1)

    buffer = IOBuffer()
    expr = "x = 3 ; for i = 1:x ; println(i) ; end"
    runexpr(expr, output=buffer, port=port)
    output = String(take!(buffer))
    @test output == "1\n2\n3\n"

    buffer = IOBuffer()
    expr = "begin
        x = 2
        for i = 1:2
            println(i)
        end
    end"
    runexpr(expr, output=buffer, port=port)
    output = String(take!(buffer))
    @test output == "1\n2\n"

    expr = "begin
        x = 2
        for i = 1:2
            println(i)
        end
    end

    "

    runexpr(expr, output=buffer, port=port)
    output = String(take!(buffer))
    @test output == "1\n2\n"
    @test istaskdone(task) == false

    sendExitCode(port)
    wait(task)
end

@testset "runConflict" begin
    port = 3004
    task = @async serve(port)
    sleep(1)
    buffer = IOBuffer()
    files = ["conflict1.jl", "conflict2.jl"]
    outputs = ["f(1) = 2\n", "f + 2 = 3\n"]

    for (file, out) in zip(files, outputs)
        @test isfile(file)
        runfile(file, output=buffer, port=port)
        output = String(take!(buffer))
        @test output == out
    end

    sendExitCode(port)
    wait(task)
end

@testset "testInclude" begin
    port = 3005
    task = @async serve(port)
    sleep(1)
    buffer = IOBuffer()
    files = ["conflict1.jl", "conflict2.jl"]
    runfile("include_test.jl", output=buffer, port=port)
    output = String(take!(buffer))
    @test output == "6\n"
end


text:
module UnitCells

using LinearAlgebra

export UnitCell, monkhorst_pack_mesh, calc_cell_pos!, calc_cell_pos, calc_site_pos!, calc_site_pos

"""
Represents a specfied lattice geometry.
"""
struct UnitCell{T<:AbstractFloat}

    "ndim: the number of dimensions that the geometry lives in."
    ndim::Int

    "norbits: the number of sites per unit cell."
    norbits::Int

    "lvecs:: (3 x 3) matrix where the columns give the lattice vectors."
    lvecs::Matrix{T}

    "rlvecs: (3 x 3) matrix where the columns give the reciprocal lattice vectors."
    rlvecs::Matrix{T}

    "bvecs: (3 x norbits) matrix where the columns give the basis vectors."
    bvecs::Matrix{T}

    #######################
    ## INNER CONSTRUCTOR ##
    #######################
    """
        UnitCell(ndim::Int,norbits::Int,lvecs::Matrix{T},bvecs::Matrix{T}) where {T<:AbstractFloat}

    Constructor for UnitCell type.
    """
    function UnitCell(ndim::Int,norbits::Int,lvecs::Matrix{T},bvecs::Matrix{T}) where {T<:AbstractFloat}

        # constructing matrix containing lattice vectors
        nrows = size(lvecs,1)
        ncols = size(lvecs,2)
        Lvecs = Matrix{T}(I,3,3) # intialized as identity matri
        Lvecs[1:nrows,1:ncols] = lvecs

        # calculating reciprocal lattice vectors
        RLvecs = 2*π*inv(Lvecs)

        # constucting matrix to contain basis vectors
        nrows = size(bvecs,1)
        ncols = size(bvecs,2)
        Bvecs = zeros(T,3,ncols)
        Bvecs[1:nrows,:] = bvecs

        new{T}(ndim,norbits,Lvecs,RLvecs,Bvecs)
    end
end

########################
## OUTER CONSTRUCTORS ##
########################

function UnitCell(ndim::Int, norbits::Int, lvecs::Matrix{T}, bvecs::Vector{Vector{T}}) where {T<:AbstractFloat}
    UnitCell(ndim, norbits, lvecs, hcat(bvecs...))
end

function UnitCell(ndim::Int, norbits::Int, lvecs::Vector{Vector{T}}, bvecs::Matrix{T}) where {T<:AbstractFloat}
    UnitCell(ndim, norbits, hcat(lvecs...), bvecs)
end

function UnitCell(ndim::Int, norbits::Int, lvecs::Vector{Vector{T}}, bvecs::Vector{Vector{T}}) where {T<:AbstractFloat}
    UnitCell(ndim, norbits, hcat(lvecs...), hcat(bvecs...))
end

#############################################
## DEFINING METHODS THAT USE GEOMETRY TYPE ##
#############################################

"""
Calculates the position of a unit cell in a lattice.
"""
function calc_cell_pos!(pos::AbstractVector{T},geom::UnitCell,l1::Int,l2::Int=0,l3::Int=0) where {T<:AbstractFloat}

    @assert length(pos)==3
    lv1 = @view geom.lvecs[:,1] # first lattice vector
    lv2 = @view geom.lvecs[:,2] # second lattice vector
    lv3 = @view geom.lvecs[:,3] # third lattice vector
    @. pos = l1*lv1 + l2*lv2 + l3*lv3 # calculating position
    return nothing
end

function calc_cell_pos(geom::UnitCell,l1::Int,l2::Int=0,l3::Int=0)::Vector{T} where {T<:AbstractFloat}

    pos = zeros(3)
    calc_cell_pos!(pos,geom,l1,l2,l3)
    return pos
end

"""
Calculates the position of a site in a lattice.
"""
function calc_site_pos!(pos::AbstractVector{T},geom::UnitCell,orbit::Int,l1::Int,l2::Int=0,l3::Int=0) where {T<:AbstractFloat}

    @assert orbit>0
    # calculating position of unit cell that site lives in
    calc_cell_pos!(pos,geom,l1,l2,l3)
    # adding basis vector for given orbital to position
    pos .+= @view geom.bvecs[:,orbit]
    return nothing
end

function calc_site_pos(geom::UnitCell{T},orbit::Int,l1::Int,l2::Int=0,l3::Int=0)::Vector{T}  where {T<:AbstractFloat}

    pos = zeros(3)
    calc_site_pos!(pos,geom,orbit,l1,l2,l3)
    return pos
end

"""
    monkhorst_pack_mesh(geom::UnitCell, L1::Int, L2::Int=1, L3::Int=1)::Matrix{AbstractFloat}

Returns a matrix where each column is a k-point vector in a Monkhort-Pack meshgrid over the full Brillouin Zone.
"""
function monkhorst_pack_mesh(geom::UnitCell{T}, L1::Int, L2::Int=1, L3::Int=1)::Matrix{T}  where {T<:AbstractFloat}

    kpoints = zeros(T,3,L1*L2*L3)
    v1 = @view geom.rlvecs[:,1]
    v2 = @view geom.rlvecs[:,2]
    v3 = @view geom.rlvecs[:,3]
    i = 1
    for l3=0:L3-1
        for l2=0:L2-1
            for l1=0:L1-1
                kpoints[:,i] = @. (l1/L1)*v1 + (l2/L2)*v2 + (l3/L3)*v3
                i += 1
            end
        end
    end
    return kpoints
end

end

max_stars_repo_name: cohensbw/Langevin

max_issues_repo_name: cohensbw/ElPhDynamics

max_forks_repo_name: cohensbw/ElPhDynamics


text:
"""
    function load_timeseries_data(data_path::String;
                              region::String="none",
                              T::Int=24,
                              years::Array{Int,1}=[2016],
                              att::Array{String,1}=Array{String,1}())

Return all time series as ClustData struct that are stored as csv files in the specified path.

- Loads `*.csv` files in the folder or the file `data_path`
- Loads all attributes (all `*.csv` files) if the `att`-Array is empty or only the files specified in `att`
- The `*.csv` files shall have the following structure and must have the same length:
|Timestamp |Year  |[column names...]|
|----------|------|-----------------|
|[iterator]|[year]|[values]         |
- The first column of a `.csv` file should be called `Timestamp` if it contains a time iterator
- The second column should be called `Year` and contains the corresponding year
- Each other column should contain the time series data. For one node systems, only one column is used; for an N-node system, N columns need to be used. In an N-node system, each column specifies time series data at a specific geolocation.
- Returns time series as ClustData struct
- The `.data` field of the ClustData struct is a Dictionary where each column in `[file name].csv` file is the key (called `"[file name]-[column name]"`). `file name` should correspond to the attribute name, and `column name` should correspond to the node name.

Optional inputs to `load_timeseries_data`:
- region-region descriptor
- T- Number of Segments
- years::Array{Int,1}= The years to be selected from the csv file as specified in `years column`
- att::Array{String,1}= The attributes to be loaded. If left empty, all attributes will be loaded.
"""
function load_timeseries_data(data_path::String;
                              region::String="none",
                              T::Int=24,
                              years::Array{Int,1}=[2016],
                              att::Array{String,1}=Array{String,1}())
  dt = Dict{String,Array}()
  num=0
  K=0
  #Check if data_path is directory or file
  if isdir(data_path)
      for full_data_name in readdir(data_path)
          if split(full_data_name,".")[end]=="csv"
              data_name=split(full_data_name,".")[1]
              if isempty(att) || data_name in att
                  # Add the
                  K=add_timeseries_data!(dt, data_name, data_path; K=K, T=T, years=years)
              end
          end
      end
  elseif isfile(data_path)
      full_data_name=splitdir(data_path)[end]
      data_name=split(full_data_name,".")[1]
      K=add_timeseries_data!(dt, data_name, dirname(data_path); K=K, T=T, years=years)
  else
      error("The path $data_path is neither recognized as a directory nor as a file")
  end
  # Store the data
  ts_input_data =  ClustData(FullInputData(region, years, num, dt),K,T)
  return ts_input_data
end #load_timeseries_data

"""
    function load_timeseries_data(existing_data::Symbol;
                              region::String="none",
                              T::Int=24,
                              years::Array{Int,1}=[2016],
                              att::Array{String,1}=Array{String,1}())

Return time series of example data sets as ClustData struct.

The choice of example data set is given by e.g. existing_data=:CEP-GER1. Example data sets are:
- `:DAM_CA` : Hourly Day Ahead Market Electricity prices for California-Stanford 2015
- `:DAM_GER` : Hourly Day Ahead Market Electricity prices for Germany 2015
- `:CEP_GER1` : Hourly Wind, Solar, Demand data Germany one node
- `:CEP_GER18`: Hourly Wind, Solar, Demand data Germany 18 nodes

Optional inputs to `load_timeseries_data`:
- region-region descriptor
- T- Number of Segments
- years::Array{Int,1}= The years to be selected from the csv file as specified in `years column`
- att::Array{String,1}= The attributes to be loaded. If left empty, all attributes will be loaded.
"""
function load_timeseries_data(existing_data::Symbol;
                              region::String="none",
                              T::Int=24,
                              years::Array{Int,1}=[2016],
                              att::Array{String,1}=Array{String,1}())
  data_path = normpath(joinpath(dirname(@__FILE__),"..","..","data"))
  if existing_data == :DAM_CA
      data_path=joinpath(data_path,"DAM","CA")
      years=[2015]
  elseif existing_data == :DAM_GER
      data_path=joinpath(data_path,"DAM","GER")
      years=[2015]
  elseif existing_data == :CEP_GER1
      data_path=joinpath(data_path,"TS_GER_1")
  elseif existing_data == :CEP_GER18
      data_path=joinpath(data_path,"TS_GER_18")
  else
      error("The symbol - $existing_data - does not exist")
  end
  return load_timeseries_data(data_path;region=region,T=T,years=years,att=att)
end

"""
    add_timeseries_data!(dt::Dict{String,Array}, data::DataFrame; K::Int=0, T::Int=24, years::Array{Int,1}=[2016])
selects first the years and second the data_points so that their number is a multiple of T and same with the other timeseries
"""
function add_timeseries_data!(dt::Dict{String,Array},
                            data_name::SubString,
                            data_path::String;
                            K::Int=0,
                            T::Int=24,
                            years::Array{Int,1}=[2016])
    #Load the data
    data_df=CSV.read(joinpath(data_path,data_name*".csv");strict=true)
    # Add it to the dictionary
    return add_timeseries_data!(dt,data_name, data_df; K=K, T=T, years=years)
end

"""
    add_timeseries_data!(dt::Dict{String,Array}, data::DataFrame; K::Int=0, T::Int=24, years::Array{Int,1}=[2016])
selects first the years and second the data_points so that their number is a multiple of T and same with the other timeseries
"""
function add_timeseries_data!(dt::Dict{String,Array},
                            data_name::SubString,
                            data::DataFrame;
                            K::Int=0,
                            T::Int=24,
                            years::Array{Int,1}=[2016])
    # find the right years to select
    time_name=find_column_name(data, [:Timestamp, :timestamp, :Time, :time, :Zeit, :zeit, :Date, :date, :Datum, :datum]; error=false)
    year_name=find_column_name(data, [:year, :Year, :jahr, :Jahr])
    data_selected=data[in.(data[year_name],[years]),:]
    for column in eachcol(data_selected, true)
        # check that this column isn't time or year
        if !(column[1] in [time_name, year_name])
            K_calc=Int(floor(length(column[2])/T))
            if K_calc!=K && K!=0
                error("The time_series $(column[1]) has K=$K_calc != K=$K of the previous")
            else
                K=K_calc
            end
            dt[data_name*"-"*string(column[1])]=Float64.(column[2][1:(Int(T*K))])
        end
    end
    return K
end

"""
        find_column_name(df::DataFrame, name_itr::Arrray{Symbol,1})
find wich of the supported name in `name_itr` is used as an
"""
function find_column_name(df::DataFrame, name_itr::Array{Symbol,1}; error::Bool=true)
    col_name=:none
    for name in name_itr
        if name in names(df)
            col_name=name
            break
        end
    end
    if error
        col_name!=:none || error("No $(name_itr) in $(repr(df)).")
    else
        col_name!=:none || @warn "No $(name_itr) in $(repr(df))."
    end
    return col_name
end

"""
        combine_timeseries_weather_data(ts::ClustData,ts_weather::ClustData)
-`ts` is the shorter timeseries with e.g. the demand
-`ts_weather` is the longer timeseries with the weather information
The `ts`-timeseries is repeated to match the number of periods of the longer `ts_weather`-timeseries.
If the number of periods of the `ts_weather` data isn't a multiple of the `ts`-timeseries, the necessary number of the `ts`-timeseries periods 1 to x are attached to the end of the new combined timeseries.
"""
function combine_timeseries_weather_data(ts::ClustData,
                                        ts_weather::ClustData)
    ts.T==ts_weather.T || error("The number of timesteps per period is not the same: `ts.T=$(ts.T)≢$(ts_weather.T)=ts_weather.T`")
    ts.K<=ts_weather.K || error("The number of timesteps in the `ts`-timeseries isn't shorter or equal to the ones in the `ts_weather`-timeseries.")
    ts_weather.K%ts.K==0 || @warn "The number of periods of the `ts_weather` data isn't a multiple of the other `ts`-timeseries: periods 1 to $(ts_weather.K%ts.K) are attached to the end of the new combined timeseries."
    ts_data=deepcopy(ts_weather.data)
    ts_mean=deepcopy(ts_weather.mean)
    ts_sdv=deepcopy(ts_weather.sdv)
    for (k,v) in ts.data
        ts_data[k]=repeat(v, 1, ceil(Int,ts_weather.K/ts.K))[:,1:ts_weather.K]
    end
    for (k,v) in ts.mean
        ts_mean[k]=v
    end
    for (k,v) in ts.sdv
        ts_sdv[k]=v
    end

    return ClustData(ts.region, ts_weather.years, ts_weather.K, ts_weather.T, ts_data, ts_weather.weights, ts_mean, ts_sdv, ts_weather.delta_t, ts_weather.k_ids)
end


text:
is_proximable(term::Term) = is_AAc_diagonal(term)

function is_proximable(terms::Tuple)
	# Check that each term is proximable
	if any(is_proximable.(terms) .== false)
		return false
	end
	# Construct the set of occurring variables
	vars = Set()
	for term in terms
		union!(vars, variables(term))
	end
	# Check that each variable occurs in only one term
	for v in vars
		tv = [t for t in terms if v in variables(t)]
		if length(tv) != 1
            if all( is_sliced.(tv) ) && all( is_proximable.(tv) )
				return true
			else
				return false
			end
		end
	end
	return true
end


text:
using JuMP
using Combinatorics
using DataFrames
using CSV
using ExcelFiles
using DelimitedFiles
using XLSX
using BenchmarkTools
using GLPK
using Gurobi

i = 0

println("Scenario $i initiated")

##### Sets #####
# read sets
sets_df = load("Scenario$i.xlsx", "sets") |> DataFrame;
sets_df;
set_K = sets_df[!,:K];
set_S = sets_df[!,:S];
set_L = sets_df[!,:L];
set_N = sets_df[!,:N];

set_K = findall(!ismissing, set_K);
set_S = findall(!ismissing, set_S);
set_L = findall(!ismissing, set_L);
set_N = findall(!ismissing, set_N);

set_A = [(i,j) for (i,j) in permutations(set_N,2) if i != j]
set_IL = [(i,l) for i in set_N for l in set_L];
set_KN = [(k,i) for k in set_K for i in set_N];
set_KSA = [(k,s,i,j) for k in set_K for s in set_S for (i,j) in set_A];
set_KSN = [(k,s,i) for k in set_K for s in set_S for i in set_N];

##### Parameters with multiple indices #####
# Read from input data file
param_c_df = load("Scenario$i.xlsx", "param_cc") |> DataFrame;
# Convert and create parameters
param_c_df[!,[:k, :s, :i, :j]] = convert.(Int,param_c_df[!,[:k, :s, :i, :j]]) #convert to int
param_c_df[!,:index] = zip(param_c_df[!,:k], param_c_df[!,:s],param_c_df[!,:i],param_c_df[!,:j]) |> collect #create index col
param_c_df = DataFrames.select!(param_c_df,[:index,:value]);
param_c = eachrow(param_c_df) |> Dict # convert to dict type

param_U_df = load("Scenario$i.xlsx", "param_U") |> DataFrame;
param_U_df[!,[:s, :i, :j]] = convert.(Int,param_U_df[!,[:s, :i, :j]]);
param_U_df[!,:index] = zip(param_U_df[!,:s],param_U_df[!,:i],param_U_df[!,:j]) |> collect;
param_U_df = DataFrames.select!(param_U_df,[:index,:value]);
param_U = eachrow(param_U_df) |> Dict;

param_v_df = load("Scenario$i.xlsx", "param_vv") |> DataFrame;
param_v_df[!,[:k,:s,:i]] = convert.(Int,param_v_df[!,[:k,:s,:i]]);
param_v_df[!,:index] = zip(param_v_df[!,:k], param_v_df[!,:s],param_v_df[!,:i]) |> collect;
param_v_df = DataFrames.select!(param_v_df,[:index,:value]);
param_v = eachrow(param_v_df) |> Dict;

param_prop_df = load("Scenario$i.xlsx", "param_prop") |> DataFrame;
param_prop_df[!,[:k,:s,:i]] = convert.(Int,param_prop_df[!,[:k,:s,:i]]);
param_prop_df[!,:index] = zip(param_prop_df[!,:k], param_prop_df[!,:s],param_prop_df[!,:i]) |> collect;
param_prop_df = DataFrames.select!(param_prop_df,[:index,:value]);
param_prop = eachrow(param_prop_df) |> Dict;

param_F_df = load("Scenario$i.xlsx", "param_F") |> DataFrame;
param_F_df[!,[:i,:l]] = convert.(Int,param_F_df[!,[:i,:l]]);
param_F_df[!,:index] = zip(param_F_df[!,:i], param_F_df[!,:l]) |> collect;
param_F_df = DataFrames.select!(param_F_df,[:index,:value]);
param_F = eachrow(param_F_df) |> Dict;

param_d_df = load("Scenario$i.xlsx", "param_dd") |> DataFrame;
param_d_df[!,[:i,:j]] = convert.(Int,param_d_df[!,[:i,:j]]);
param_d_df[!,:index] = zip(param_d_df[!,:i], param_d_df[!,:j]) |> collect;
param_d_df = DataFrames.select!(param_d_df,[:index,:value]);
param_d = eachrow(param_d_df) |> Dict;

##### Parameters with single index #####
param_M_df = load("Scenario$i.xlsx", "param_M") |> DataFrame;
param_M_df[!,[:index]] = convert.(Int,param_M_df[!,[:index]]);
param_M = eachrow(param_M_df) |> Dict;

param_b_df = load("Scenario$i.xlsx", "param_bb") |> DataFrame;
param_b_df[!,[:index]] = convert.(Int,param_b_df[!,[:index]]);
param_b = eachrow(param_b_df) |> Dict;

param_q_df = load("Scenario$i.xlsx", "param_qq") |> DataFrame;
param_q_df[!,[:index]] = convert.(Int,param_q_df[!,[:index]]);
param_q = eachrow(param_q_df) |> Dict;

param_h_df = load("Scenario$i.xlsx", "param_hh") |> DataFrame;
param_h_df[!,[:index]] = convert.(Int,param_h_df[!,[:index]]);
param_h = eachrow(param_h_df) |> Dict;

param_p_df = load("Scenario$i.xlsx", "param_pp") |> DataFrame;
param_p_df[!,[:index]] = convert.(Int,param_p_df[!,[:index]]);
param_p = eachrow(param_p_df) |> Dict;

param_P_df = load("Scenario$i.xlsx", "param_P") |> DataFrame;
param_P_df[!,[:index]] = convert.(Int,param_P_df[!,[:index]]);
param_P = eachrow(param_P_df) |> Dict;

param_u_df = load("Scenario$i.xlsx", "param_uu") |> DataFrame;
param_u_df[!,[:index]] = convert.(Int,param_u_df[!,[:index]]);
param_u = eachrow(param_u_df) |> Dict;

param_bigM_df = load("Scenario$i.xlsx", "param_bigM") |> DataFrame;
param_bigM_df[!,[:index]] = convert.(Int,param_bigM_df[!,[:index]]);
param_bigM = eachrow(param_bigM_df) |> Dict;

# Adjusting parameters
param_D_df = load("Scenario$i.xlsx", "param_D") |> DataFrame;
param_D_df[!,[:index]] = convert.(Int,param_D_df[!,[:index]]);
param_D = eachrow(param_D_df) |> Dict;

param_alpha_df = load("Scenario$i.xlsx", "param_alpha") |> DataFrame;
param_alpha_df[!,[:index]] = convert.(Int,param_alpha_df[!,[:index]]);
param_alpha = eachrow(param_alpha_df) |> Dict;

println("Parameters set up successfully for scenario $i")

####################################################################################################################
#### Model Prep ####
# Time the model
TimeTrack_start = time();

# Preparing an optimization model
m = Model(Gurobi.Optimizer);
# m = Model(GLPK.Optimizer);

##Declaring variables

# variables
@variable(m, y[index in set_IL], Bin);
@variable(m, r[index in set_KN] >= 0, Int);
@variable(m, x[index in set_KSA] >= 0, Int);
@variable(m, z[index in set_KSN] >= 0, Int);
@variable(m, w[index in set_KSN] >= 0, Int);
@variable(m, gamma[s in set_S], Bin);

##Objective function
@objective(m, Min, sum(param_F[(i,l)] * y[(i,l)] for (i,l) in set_IL) # Fix Cost
                    + sum(param_q[k] * r[(k,n)] for (k,n) in set_KN)  # Preposition Cost
                    + sum(param_P[s] * param_c[(k,s,i,j)] * x[(k,s,i,j)] for (k,s,i,j) in set_KSA) # Scenario Specific Cost 1
                    + sum(param_P[s] * param_h[k] * z[(k,s,i)] + param_P[s] * param_p[k] * w[(k,s,i)] for (k,s,i) in set_KSN) # Scenario Specific Cost 2
            );

##Constraints
# Flow Conservation
fc = @constraint(m, fc[k in set_K, s in set_S, i in set_N],
                        sum(x[(k,s,j,i)] for j in set_N if j != i) + param_prop[(k,s,i)]* r[(k,i)] - z[(k,s,i)]
                        ==
                        sum(x[(k,s,i,j)] for j in set_N if j != i) + param_v[(k,s,i)] - w[(k,s,i)]
            );

# Open Facility Capacity
ofc = @constraint(m, ofc[i in set_N],
                    sum(param_b[k] * r[(k,i)] for k in set_K)
                    <=
                    sum(param_M[l] * y[(i,l)] for l in set_L)
            );
# Facility Number per Node
fpn = @constraint(m, fpn[i in set_N],
                sum(y[(i,l)] for l in set_L)
                <=
                1
            );
# Arc Capacity
ac = @constraint(m, ac[s in set_S, (i,j) in set_A],
                sum(param_u[k] * x[(k,s,i,j)] for k in set_K)
                <=
                param_U[(s,i,j)]
            );

# Average Distance Limit
adl = @constraint(m, adl[k in set_K, s in set_S],
                    sum(param_d[(i,j)] * x[(k,s,i,j)] for (i,j) in set_A)
                    <=
                    param_D[k] * sum(param_v[(k,s,i)] for i in set_N) + param_bigM[1] * (1 - gamma[s])
            );

# Reliability Set Definition
rsd = @constraint(m, rsd,
                sum(param_P[s] * gamma[s] for s in set_S)
                >=
                param_alpha[1]
            );

# Demand requirements for scenarios in reliable set
dmrs = @constraint(m, dmrs[k in set_K, s in set_S, i in set_N],
                    w[(k,s,i)]
                    <=
                    param_v[(k,s,i)] * (1 - gamma[s])
            );

println("Model successfully set up for scenario $i")
####################################################################################################################
#### Model solving ####

# Printing the prepared optimization model
# print(m)

# Solving the optimization problem
JuMP.optimize!(m)
TimeTrack_elapsed = time() - TimeTrack_start;

####################################################################################################################
#### Post-processing ####

# Print result
println("Model running time is: ", round(TimeTrack_elapsed,digits = 5),"s")
println("Model solving time:",round(JuMP.solve_time(m),digits = 5), "s")

# save the optimized variables to *_star
y_star = JuMP.value.(y);
r_star = JuMP.value.(r);
x_star = JuMP.value.(x);
z_star = JuMP.value.(z);
w_star = JuMP.value.(w);
gamma_star = JuMP.value.(gamma);

# Print objective value
println("Objective value: ", JuMP.objective_value(m))

# Print the non-zero variables
println("If a variable is not printed, the value is 0!")
for (i,l) in set_IL
    if y_star[(i,l)] == 1
        println("facility $l located in node $i")
    else end
end;

for (k,i) in set_KN
    if r_star[(k,i)] > 0
        println("Amount of commodity $k pre-position at node $i: ", r_star[(k,i)])
    else end
end;

for (k,s,i,j) in set_KSA
    if x_star[(k,s,i,j)] > 0
    println("Amount of commodity $k shipped across link ($i,$j) in scenario $s: ", x_star[(k,s,i,j)])
    else end
end;

for (k,s,i) in set_KSN
    if z_star[(k,s,i)] > 0
        println("Commodity $k not used in scenario $s at node $i: ", z_star[(k,s,i)] )
    else end
end;

for (k,s,i) in set_KSN
    if w_star[(k,s,i)] > 0
        println("Shortage of commodity $k at node $i in scenario $s: ", w_star[(k,s,i)])
    else end
end;

for s in set_S
    if gamma_star[s] > 0
        println("Scenario $s belong to the reliable set")
    else end
end;


text:
module Replace

using Unicode

include("replace.jl")
include("kctest.jl")
include("sjon.jl")
include("sjreg.jl")

end # module


text:
abstract type PropODE end

struct GenerativePropODE{T <: Real, R <: Real, S <: Real, M <: Function} <: PropODE
    x₀ :: Vector{T}
    F :: Vector{M}
    p :: Matrix{R}
    g :: SimpleGraph
    τ :: Vector{S}
end

abstract type InferablePropODE <: PropODE end

struct ModelPropODE{T <: Real, R <: Real, S <: Real, M <: Function} <: InferablePropODE
    x₀ :: Vector{T}
    F :: Vector{M}
    p :: Matrix{R}
    prob :: GRNProblem
    τ :: Vector{S}
end

struct DeepPropODE{T <: Real, R <: Real, S <: Real} <: InferablePropODE
    x₀ :: Vector{T}
    F :: Vector{FastChain}
    p :: Matrix{R}
    prob :: GRNProblem
    τ :: Vector{S}
end

### Generative Portion ###

get_graph(prop :: GenerativePropODE) = prop.g
get_graph(prop :: T) where {T <: InferablePropODE} = prop.prob.g

function generate(g :: PropODE)
    function dynamic_system!(du, u, p, t)
        for (i, f) ∈ enumerate(g.F)
            n = neighbors(get_graph(g), i)
            du[i] = f(u[n], p[i, n], t)
        end
    end

    tspan = (g.τ[1], g.τ[end])
    prob = ODEProblem(dynamic_system!, g.x₀, tspan, g.p, saveat = g.τ)
    return solve(prob)
end

### Inference Portion ###


text:
using EconPDEs

Base.@kwdef struct LelandModel
    # consumption process parameters
    r::Float64 = 0.02
    μ::Float64 = 0.0
    σ::Float64 = 0.1
    C::Float64 = 0.05 # coupon rate
    τ::Float64 = 0.2 # tax rate
end

function (m::LelandModel)(state::NamedTuple, y::NamedTuple)
    (; r, μ, σ, C, τ) = m
    (; δ) = state
    (; E, Eδ_up, Eδ_down, Eδδ) = y
    μδ, σδ = μ * δ, σ * δ
    Eδ = (μδ >= 0) ? Eδ_up : Eδ_down
    f = δ - (1 - τ) * C
    Et = - (f + Eδ * μδ + 0.5 * Eδδ * σδ^2 - r * E)
    return (; Et)
end

m = LelandModel()
stategrid = OrderedDict(:δ => [0:0.005:0.2;])
yend = OrderedDict(:E => [max(δ / (m.r - m.μ) - (1 - m.τ) * m.C / m.r, 0.0) for δ in stategrid[:δ]])
y̲ = zeros(length(stategrid[:δ]))
bc = OrderedDict(:Eδ => (0.0, 1 / (m.r - m.μ),))
result = pdesolve(m, stategrid, yend; y̲ = y̲, bc = bc, reformulation = :smooth)
@assert result.residual_norm <= 1e-5

"""
    Analytical solution to Leland model. Given current cashflow δ, return the equity value.
"""
function analytical_solution(m::LelandModel, δ)
    (; r, μ, σ, C, τ) = m
    Δ = (1/2 * σ^2 - μ)^2 + 2 * σ^2 * r
    γ = (μ - 1/2 * σ^2 + sqrt(Δ))/σ^2
    η = -(μ - 1/2 * σ^2 - sqrt(Δ))/σ^2
    δb = (1-τ) * C * (r-μ)/r * γ/(1+γ)
    if δ < δb
        return zero(δ)
    else
        return δ / (r - μ) - (1 - τ) * C / r + ((1 - τ) * C / r - δb / (r - μ)) * (δ / δb)^(-γ)
    end
end


text:
using NumericalMethodsforEngineers
using Test

b = [16. -24. 18.; 3. -2. 0.; -9. 18. -17.]
x0 = [1., 1., 1.]

tol = 1.0e-5
limit = 100

(iters, eigval, eigvec) = nmex(b, x0; tol=tol, limit=100)

#eb = eig(b)
ebvals, ebvecs = (eigen(b)...,)

@test round.(eigvec[:,1], digits=5) == round.(ebvecs[:,1], digits=5)


text:
# -----------------------------------------------------------
#
# Implementing a nodal basis for DG sparse grids
#
# -----------------------------------------------------------

#=

We begin by defining analogues of our DG functions
to satisfy the following conditions:

	1) There are a set of collocation points respecting the hierarchical structure
	2) The nodal basis functions at a given level are those that vanish
		on all nodal points at the above levels and vanish on all but one
		of the cells on the level below
	3) Third constraint?

To do this, we _restrict_ to the continuous part of the DG basis, and focus
on piecewise polynomials that vanish on their boundaries. We then pick the
collocation points to be exactly at these boundaries.

For this, it is useful to define lagrange polynomials, and moreover
use k=2, 3, 5.
The k=2 case has our basis is exactly the hat function basis constructed
previously.

For k=2, the coarsest level has nodal points at {0, 1},
	the next level (the first hat) has nodal point at 0.5
	the next level has two cells with points at 0.25, 0.75 etc.

For k=3 the coarsest level has nodal points at {0, 0.5, 1},
	the next level has two nodal functions with points at 0.25, 0.75,
	the next level has two cells and two nodes per cell at all odd multiples of 0.125 etc.

For k=5 the coarsest level has nodal points at {0, 0.25, 0.5, 0.75, 1}
	the next level has four nodal functions with points at all multiples of 0.125
	the next level has two cells and four nodes per cell at all odd multiples of 0.0625 etc.

Notice that k=2,3,5 work because at each cell of the multiresolution basis the space of
	picewise continuous polynomials is k-1 = 2^i for i=0,1,2 in this case.
	A construction for other k would be harder, if not impossible.

=#

# Lagrange polynomials interpolating at the coarsest level
function lag_nodal(k::Int, mode::Int, x::Real)
	(k < 2 || k > 5) && throw(ArgumentError("k is not between 2 and 5"))
	(mode > k || mode < 1) && throw(ArgumentError("mode is not in [1, k]"))
	if k == 2
		mode == 1 && return max(zero(x), x)
		mode == 2 && return max(zero(x), 1-x)
	elseif k == 3
		mode == 1 && return (x < 0 || x > 1) ? zero(x) : 2 * (x-1) * (x - 1//2)
		mode == 2 && return (x < 0 || x > 1) ? zero(x) : 2 * x * (x - 1//2)
		mode == 3 && return (x < 0 || x > 1) ? zero(x) : -4 * x * (x-1)
	elseif k == 4
		throw(MethodError("k = 4 not implemented"))
	elseif k == 5
		mode == 1 && return (x > 1 || x < 0) ? zero(x) : 32//3 * (x - 1//4) * (x - 1//2) * (x - 3//4) * (x - 1)
		mode == 2 && return (x > 1 || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 3//4)
		mode == 3 && return (x > 1 || x < 0) ? zero(x) : 64 * x * (x - 1) * (x - 1//4) * (x - 3//4)
		mode == 4 && return (x > 1 || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 1//4)
		mode == 5 && return (x > 1 || x < 0) ? zero(x) : 32//3 * x * (x - 1//4) * (x - 1//2) * (x - 3//4)
	end
end

# Multiresolution functions for the levels below
function h_nodal(k::Int, mode::Int, x::Real)
	(k < 2 || k > 5) && throw(ArgumentError("k is not between 2 and 5"))
	(mode >= k) && return zero(x)
	#(abs(x) >= 1) && return zero(T)
	if k == 2
		return max(zero(x), 1-abs(x))
	elseif k == 3
		mode == 1 && return max(zero(x), -4 * x * (x + 1))
		mode == 2 && return max(zero(x), -4 * x * (x - 1))
	elseif k == 4
		throw(MethodError("k = 4 not implemented"))
	elseif k == 5
		mode == 1 && return (x < -1 || x > 0) ? zero(x) : -128//3 * x * (x + 1) * (x + 1//2) * (x + 1//4)
		mode == 2 && return (x < -1 || x > 0) ? zero(x) : -128//3 * x * (x + 1) * (x + 1//2) * (x + 3//4)
		mode == 3 && return (x > 1  || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 3//4)
		mode == 4 && return (x > 1  || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 1//4)
	end
end

# This yields our nodal basis:
function v_nodal(k::Int, level::Int, cell::Int, mode::Int, x::Real)
	if level==0
		return lag_nodal(k, mode, x)
	else
		return h_nodal(k, mode, (1<<level)*x - (2*cell-1))
	end
end

function v_nodal(k::Int, level::Int, cell::Int, mode::Int)
	return x->v_nodal(k, level, cell, mode, x)
end

# -----------------------------------------------------------
# Next we define the sparse matrices transforming between
# the heirarchical, nodal, and point bases
# -----------------------------------------------------------

# Evaluation of a given nodal function along all relevant gridpoints
function eval_points_1D(k::Int, max_level::Int, level::Int, cell::Int, mode::Int)
	i = 1
	I = Int[]
	V = Float64[]
	for l in 0:max_level
		l == 0 ? (node_range = 0:(1//(k-1)):1) :
				 (node_range = 1//(2*(k-1)):1//(k-1):(1-1//(2*(k-1))))

		for c in 1:1<<max(0, l-1)
			for node in node_range
				x = (c - 1 + node)//(1<<max(0, l-1))
				val = v_nodal(k, level, cell, mode, x)
				(val != zero(x)) && (push!(I, i); push!(V, val))
				i += 1
			end
		end
	end
	return dropzeros!(sparsevec(I, V))
end

# Sparse matrix converting nodal basis element to its value
# on the collocation points
function nodal2points_1D(k::Int, max_level::Int)
	I = Int[]
	J = Int[]
	V = Float64[]
	j = 1

	for level in 0:max_level
		level == 0 ? (mode_range = 1:k) : (mode_range = 1:(k-1))

		for cell in 1:1<<max(0, level-1)
			for mode in mode_range
				coeffs = eval_points_1D(k, max_level, level, cell, mode)
				for (i, val) in zip(findnz(coeffs)...)
					push!(I, i)
					push!(J, j)
					push!(V, val)
				end
				j += 1
			end
		end
	end
    return dropzeros!(sparse(I, J, V))
end

# Sparse matrix inverting the above construction
# Converting a set of values on the collocation points to the DG space
# in terms of the (for now continuous) nodal basis
# TODO: Allow for this to handle discontinuity across cells
function points2nodal_1D(k::Int, max_level::Int)
    Q = nodal2points_1D(k, max_level)
    return threshold(inv(full(Q)), abs_tol=1e-15)
end

# Converts a given nodal basis element to a sum in the standard
# hieararchical DG scheme
function nodal2modal_1D(k::Int, level::Int, cell::Int, mode::Int; rel_tol=1e-10, abs_tol=1e-15, max_evals=1500)
	I = Int[]
	V = Float64[]
	i = 1

	for hier_level in 0:level
		hier_cells = 1<<max(0, hier_level-1)
		for hier_cell in 1:hier_cells
			for hier_mode in 1:k
				x_min = (hier_cell-1)//hier_cells
				x_max = (hier_cell)//hier_cells
				val = hquadrature(x->v_nodal(k, level, cell, mode, x)*
									v(k, hier_level, hier_cell, hier_mode, x),
									x_min, x_max,
									reltol=rel_tol, abstol=abs_tol, maxevals=max_evals)[1]
				(abs(val) > eps(abs_tol)) && (push!(I, i); push!(V, val))
				i += 1
			end
		end
	end
	return dropzeros!(sparsevec(I, V))
end

function nodal2modal_1D(k::Int, max_level::Int; rel_tol=1e-10, abs_tol=1e-15, max_evals=1500)
	I = Int[]
	J = Int[]
	V = Float64[]
	j = 1

	for nodal_level in 0:max_level
		nodal_level == 0 ? num_nodes = k : num_nodes = k-1
		nodal_cells = 1<<max(0, nodal_level-1)
		for nodal_cell in 1:nodal_cells
			for nodal_mode in 1:num_nodes
				coeffs = nodal2modal_1D(k, nodal_level, nodal_cell, nodal_mode;
										rel_tol=rel_tol,
										abs_tol=abs_tol,
										max_evals=max_evals)
				for (i, val) in zip(findnz(coeffs)...)
					push!(I, i)
					push!(J, j)
					push!(V, val)
				end
				j += 1
			end
		end
	end
	return sparse(I, J, V)
end

# Accounting for ambiguity in the definition of the hierarchical
# basis across regions
function eval_v(k, level, cell, mode, x)
	xl = max(x - 1e-16, 0)
	xr = min(x + 1e-16, 1)
	(xl == 0) && (xl = 1 - 1e-16)
	(xr == 1) && (xr = 1e-16)
	return 0.5 * (v(k, level, cell, mode, xl) + v(k, level, cell, mode, xr))
end

# Evaluates a given hierarchical basis function at
# all collocation points
function modal2points_1D(k::Int, max_level::Int, level::Int, cell::Int, mode::Int)
	i = 1
	I = Int[]
	V = Float64[]

	# subsequent levels:
	for l in 0:max_level
		l == 0 ? (node_range = 0:(1/(k-1)):1) :
				 (node_range = 1/(2*(k-1)):1/(k-1):(1-1/(2*(k-1))))

		for c in 1:1<<max(0, l-1)
			for node in node_range
				x = (c - 1 + node)/(1<<max(0, l-1))
				val = eval_v(k, level, cell, mode, x)
				(abs(val) != 0) && (push!(I, i); push!(V, val))
				i += 1
			end
		end
	end
	return dropzeros!(sparsevec(I, V))
end

# Constructs the sparse matrix to convert from the
# hierarchical basis to the point basis
function modal2points_1D(k::Int, max_level::Int)
	I = Int[]
	J = Int[]
	V = Float64[]
	j = 1

	for level in 0:max_level
		for cell in 1:1<<max(0, level-1)
			for mode in 1:k
				coeffs = modal2points_1D(k, max_level, level, cell, mode)
				for (i, val) in zip(findnz(coeffs)...)
					push!(I, i)
					push!(J, j)
					push!(V, val)
				end
				j += 1
			end
		end
	end
	return dropzeros!(sparse(I, J, V))
end

function nodal2pos_1D(k, max_level, nodal_level, nodal_cell, nodal_mode;
				   rel_tol=1e-10, abs_tol=1e-15, max_evals=1500)
	I = Int[]
	V = Float64[]
	i = 1

	nodal_min = (nodal_cell-1)/(1<<max(0, nodal_level-1))
	nodal_max = (nodal_cell)/(1<<max(0, nodal_level-1))
	for c in 1:1<<max_level
		pos_min = (c-1)/(1<<max_level)
		pos_max = c/(1<<max_level)

		if (pos_min > nodal_max) || (pos_max < nodal_min)
			i += k
			continue
		end
		for m in 1:k
			val = hquadrature(x->v_nodal(k, nodal_level, nodal_cell, nodal_mode, x)*basis(max_level, c, m, x),
							  pos_min, pos_max,
							  reltol=rel_tol, abstol=abs_tol, maxevals=max_evals)[1]
			(abs(val) > eps(abs_tol)) && (push!(I, i); push!(V, val))
			i += 1
		end
	end
	return sparsevec(I, V)
end

function nodal2pos_1D(k, max_level; rel_tol=1e-10, abs_tol=1e-15, max_evals=1500)
	I = Int[]
	J = Int[]
	V = Float64[]
	j = 1

	for nodal_level in 0:max_level
		nodal_level == 0 ? num_nodes = k : num_nodes = k-1
		nodal_cells = 1<<max(0, nodal_level-1)
		for nodal_cell in 1:nodal_cells
			for nodal_mode in 1:num_nodes
				coeffs = nodal2pos_1D(k, max_level, nodal_level, nodal_cell, nodal_mode;
										rel_tol=rel_tol,
										abs_tol=abs_tol,
										max_evals=max_evals)
				for (i, val) in zip(findnz(coeffs)...)
					push!(I, i)
					push!(J, j)
					push!(V, val)
				end
				j += 1
			end
		end
	end
	return sparse(I, J, V)
end

# Inverts the above construction by factoring through
# Essentially all the different spaces
# TODO: Make this more direct!!
function points2modal_1D(k::Int, max_level::Int)
	p2n = points2nodal_1D(k, max_level)
	n2pos = nodal2pos_1D(k, max_level)
	pos2m = pos2hier(k, max_level)
	return threshold(pos2m*n2pos*p2n)
end


text:
using Test, Mill, Flux
using Random

function ngradient(f, xs::AbstractArray...)
  grads = zero.(xs)
  for (x, Δ) in zip(xs, grads), i in 1:length(x)
    δ = sqrt(eps())
    tmp = x[i]
    x[i] = tmp - δ/2
    y1 = f(xs...)
    x[i] = tmp + δ/2
    y2 = f(xs...)
    x[i] = tmp
    Δ[i] = (y2-y1)/δ
  end
  return grads
end

gradcheck(f, xs...) =
  all(isapprox.(ngradient(f, xs...),
                gradient(f, xs...), rtol = 1e-5, atol = 1e-5))

gradtest(f, xs::AbstractArray...) = gradcheck((xs...) -> sum(sin.(f(xs...))), xs...)
gradtest(f, dims...) = gradtest(f, rand.(Float64, dims)...)

function mngradient(f, xs::AbstractArray...)
    grads = zero.(xs)
    for (x, Δ) in zip(xs, grads), i in eachindex(x)
        δ = sqrt(eps())
        tmp = x[i]
        x[i] = tmp - δ/2
        y1 = f(xs...)
        x[i] = tmp + δ/2
        y2 = f(xs...)
        x[i] = tmp
        Δ[i] = (y2-y1)/δ
    end
    return grads
end

function mgradcheck(f, xs...)
    correct = true
    for (ng, ag) in zip(mngradient(f, xs...), gradient(f, xs...))
        if isnothing(ag)
            if any(ng .!= 0)
                correct = false
                @show ng
            end
        elseif !all(isapprox.(ng, ag, rtol = 1e-5, atol = 1e-5))
            correct = false
            grad_dif = [abs.(x) for x in (ng .- ag)]
            @show ng
            @show ag
            @show grad_dif
        end
    end
    correct
end

mgradtest(f, xs::AbstractArray...) = mgradcheck((xs...) -> sum(sin.(f(xs...))), xs...)

Random.seed!(42)

include("datanode.jl")
include("modelnode.jl")
include("missing.jl")
include("aggregation.jl")
include("gradtests.jl")
include("conv.jl")
include("bags.jl")
include("ngrams.jl")
include("activations.jl")


text:
using HexEdit

# construct type
hexedit = HexEd("test.bin")
# dump whole file
dump!(hexedit, 0x00)
# write foobar beginning at offset 0x04 using hex string as arg
edit!(hexedit, "0x666f6f626172", 4)
# dump 16 bytes beginning at offset 0x04
dump!(hexedit, 4, 16)
# write foobar beginning at offset 0x04 using ascii as arg
edit!(hexedit, "foobar", 0x04)
# dump 16 bytes beginning at offset 0x04
dump!(hexedit, 0x04, 16)
# find foobar
fbl = find!(hexedit, "foobar", 0x00)
if fbl != nothing
    println("found foobar at $(uppercase(string(fbl, base=16, pad=8)))")
end
# find hex signature
fbl = find!(hexedit, "0x07ae0971")
if fbl != nothing
    println("found hex signature at $(uppercase(string(fbl, base=16, pad=8)))")
end


text:
# This file is not an active part of the package. This is the code
# that uses the Clang.jl package to wrap Librsync using the headers.

# Find all headers
# To run the script from Julia console:
# include(joinpath(Pkg.dir("Librsync"), "src", "wrap.jl"));
const llvmpath = "/usr/local/Cellar/llvm/5.0.0"
push!(Libdl.DL_LOAD_PATH, joinpath(llvmpath, "lib"))
using Clang.wrap_c
const wdir = dirname(@__FILE__)
const pkgbasedir = joinpath(wdir, "..")

# `outpath` specifies, where the julian wrappers would be generated.
# If the generated .jl files are ok, they have to be copied to the "src" folder
# overwriting the old ones
const outpath = joinpath(pkgbasedir, "new")
rm(outpath, recursive = true, force = true)
mkpath(outpath)

const incpath = joinpath(pkgbasedir, "deps", "usr", "include")
if !isdir(incpath)
  error("Run Pkg.build(\"Librsync\") before trying to wrap C headers.")
end

info("Scanning Librsync headers in $incpath...")
const Librsync_header_files = ["librsync"]
const Librsync_headers =[joinpath(incpath, y) for y in (string(x, ".h") for x in Librsync_header_files)]

const clang_path = joinpath(llvmpath, "lib", "clang", "5.0.0") # change to your clang location
const includes = [
    joinpath(clang_path, "include"),
    incpath
]

function find_outfile(s)
    joinpath(outpath, string(first(splitext(basename(s))), ".jl"))
end
find_libfile(s) = "librsync"

const context = wrap_c.init(
    headers = Librsync_headers,
    common_file = joinpath(outpath, "Librsync_types.jl"),
    clang_args = [
        "-D", "__STDC_LIMIT_MACROS",
        "-D", "__STDC_CONSTANT_MACROS",
        # "-v"
    ],
    clang_diagnostics = true,
    header_library = find_libfile,
    header_outputfile = find_outfile,
    clang_includes = includes,
)

info("Generating .jl wrappers for Librsync in $outpath...")
run(context)
info("Done generating .jl wrappers for Librsync in $outpath")


text:
using Documenter
import Pkg
using RootedTrees

# Define module-wide setups such that the respective modules are available in doctests
DocMeta.setdocmeta!(RootedTrees,
  :DocTestSetup, :(using RootedTrees); recursive=true)

# Copy some files from the top level directory to the docs and modify them
# as necessary
open(joinpath(@__DIR__, "src", "license.md"), "w") do io
  # Point to source license file
  println(io, """
  ```@meta
  EditURL = "https://github.com/SciML/RootedTrees.jl/blob/main/LICENSE.md"
  ```
  """)
  # Write the modified contents
  println(io, "# License")
  println(io, "")
  for line in eachline(joinpath(dirname(@__DIR__), "LICENSE.md"))
    line = replace(line, "[LICENSE.md](LICENSE.md)" => "[License](@ref)")
    println(io, "> ", line)
  end
end

open(joinpath(@__DIR__, "src", "contributing.md"), "w") do io
  # Point to source license file
  println(io, """
  ```@meta
  EditURL = "https://github.com/SciML/RootedTrees.jl/blob/main/CONTRIBUTING.md"
  ```
  """)
  # Write the modified contents
  println(io, "# Contributing")
  println(io, "")
  for line in eachline(joinpath(dirname(@__DIR__), "CONTRIBUTING.md"))
    line = replace(line, "[LICENSE.md](LICENSE.md)" => "[License](@ref)")
    println(io, "> ", line)
  end
end

# Make documentation
makedocs(
  modules = [RootedTrees],
  sitename="RootedTrees.jl",
  format = Documenter.HTML(
    prettyurls = get(ENV, "CI", nothing) == "true",
    canonical = "https://SciML.github.io/RootedTrees.jl/stable",
    ansicolor = true
  ),
  # Explicitly specify documentation structure
  pages = [
    "Home" => "index.md",
    "Introduction" => "introduction.md",
    "Tutorials" => [
      "tutorials/RK_order_conditions.md",
      "tutorials/ARK_order_conditions.md",
    ],
    # "Benchmarks" => "benchmarks.md",
    "API reference" => "api_reference.md",
    "Contributing" => "contributing.md",
    "License" => "license.md"
  ],
  strict = true # to make the GitHub action fail when doctests fail
)

deploydocs(
  repo = "github.com/SciML/RootedTrees.jl",
  devbranch = "main",
  push_preview = true
)


text:
#=
= Energy and derivative plots.
=#

include("hcp.jl")

using HCP: Constants, Auxiliary, readvec
using LaTeXStrings
using StatPlots

function loc(data::Vector{Float64};
    s::Int=1, R::Float64=40907.,
    alpha::Float64=1., beta::Float64=200.)

    Y = sort(data)
    C = Constants(3,40907.,1.,200.,0.,0.,0.,0.)
    X = Auxiliary(3)
    X.locs[1], X.locs[5] = 0., 40907.
    X.locs[2] = 14_500 + 1000randn()
    X.locs[3] = 24_000 + 3000randn()
    X.locs[4] = 35_500 + 1000randn()
    X.hts[1] = abs(.0085 + .001randn())
    X.hts[2] = abs(.0025 + .0005randn())
    X.hts[3] = abs(.003 + .001randn())
    X.hts[4] = abs(.001 + .0005randn())
    HCP.loc_metrics(Y,X)
    HCP.loc_unconstrained(C,X)
    HCP.loc_potential(C,X)
    HCP.loc_momenta(C,X)
    HCP.ht_unconstrained(X)
    HCP.ht_momenta(C,X)
    HCP.loc_println(X)

    function e(x)
        X.locs[s+1] = x
        HCP.loc_metrics(Y,X)
        HCP.loc_unconstrained(C,X)
        return HCP.loc_potential(C,X)
    end

    function d(x)
        X.locs[s+1] = x
        HCP.loc_metrics(Y,X)
        HCP.loc_unconstrained(C,X)
        HCP.loc_potential(C,X)
        return X.∇tlocs[s]
    end

    r = X.locs[s+2]-X.locs[s]
    xp,xn = X.locs[s]+r/50, X.locs[s+2]-r/50
    ts = linspace(log(xp/(R-xp)), log(xn/(R-xn)), 1000)
    xs = map(t -> R/(1+exp(-t)), ts)
    es,ds = e.(xs), d.(xs)

    # Approximate the integral of the derivative to compare it to the
    # discontinuous function.

    fs = similar(es)
    fs[1] = 0
    for i = 2:length(ts)
        fs[i] = fs[i-1] + (ts[i]-ts[i-1])ds[i-1]
    end
    j = div(length(ts),2)
    fs += es[j]-fs[j]

    # Compute the derivatives numerically to verify the analyitic solution.
    # Sampled outliers arising from discontinuities in the energy function are
    # removed (hopefully).

    tmids = @. @view(ts[1:end-1]) + .5(@view(ts[2:end]) - @view(ts[1:end-1]))
    dmids = @. begin
        (@view(es[2:end]) - @view(es[1:end-1])) /
        (@view(ts[2:end]) - @view(ts[1:end-1]))
    end

    L,U = extrema(ds)
    for (i,d) in enumerate(dmids)
        (d < L || d > U) && (dmids[i] = NaN)
    end

    # Plot the numerical and analytical derivatives.

    p = plot(ts,es,lab="energy")
    plot!(ts,fs,lab="integral",ls=:dash)
    q = plot(ts,ds,lw=8,lab="analytic")
    plot!(q,tmids,dmids,lw=4,lab="numeric")
    plot(p,q,layout=(2,1),size=(1200,900))

    # p = plot(ts,es,lc=:black,lab=false,ylabel=L"H")
    # plot!(ts,fs,lc=:black,lab=false,ls=:dash)
    # q = plot(ts,ds,lc=:black,lab=false,xlabel=L"\hat{s}_2",ylabel=L"\frac{\partial H}{\partial \hat{s}_2}")
    # pq = plot(p,q,layout=(2,1),size=(480,360))
    # savefig(pq, "loc.pdf")
    # savefig(pq, "loc.png")
    # savefig(pq, "loc.tex")

end

function ht(data::Vector{Float64};
    h::Int=0, R::Float64=40907.,
    alpha::Float64=1., beta::Float64=200.)

    Y = sort(data)
    C = Constants(3,40907.,1.,200.,0.,0.,0.,0.)
    X = Auxiliary(3)
    X.locs[1], X.locs[5] = 0., 40907.
    X.locs[2] = 14_500 + 1000randn()
    X.locs[3] = 24_000 + 3000randn()
    X.locs[4] = 35_500 + 1000randn()
    X.hts[1] = abs(.0085 + .001randn())
    X.hts[2] = abs(.0025 + .0005randn())
    X.hts[3] = abs(.003 + .001randn())
    X.hts[4] = abs(.001 + .0005randn())
    HCP.loc_metrics(Y,X)
    HCP.loc_unconstrained(C,X)
    HCP.loc_potential(C,X)
    HCP.loc_momenta(C,X)
    HCP.ht_unconstrained(X)
    HCP.ht_momenta(C,X)
    HCP.ht_println(X)

    function e(x)
        X.hts[h+1] = x
        HCP.ht_unconstrained(X)
        return HCP.ht_potential(C,X)
    end

    function d(x)
        X.hts[h+1] = x
        HCP.ht_unconstrained(X)
        HCP.ht_potential(C,X)
        return X.∇lhts[h+1]
    end

    js = linspace(log(.0001),log(.025),1000)
    xs = exp.(js)
    es,ds = e.(xs), d.(xs)
    jmids = @. @view(js[1:end-1]) + .5(@view(js[2:end]) - @view(js[1:end-1]))
    dmids = @. begin
        (@view(es[2:end]) - @view(es[1:end-1])) /
        (@view(js[2:end]) - @view(js[1:end-1]))
    end

    p = plot(js,es,lab="energy")
    q = plot(js,ds,lw=8,lab="analytic")
    plot!(q,jmids,dmids,lw=4,lab="numeric")
    plot(p,q,layout=(2,1),size=(1200,900))

end

max_stars_repo_name: kevdur/rjcp

max_issues_repo_name: kevdur/rjcp

max_forks_repo_name: kevdur/rjcp


text:
# eventually boundary conditions should be a subtype of this
# we don't enforce it currently to make the transition easier
abstract type BoundaryCondition end

"""
    boundary_conditions(::BL)

Define the set of boundary conditions for the balance law `BL`. This should
return a tuple, where a boundary tagged with the integer `i` will use the `i`th
element of the tuple.
"""
function boundary_conditions end

"""
    boundary_state!(
        ::NumericalFluxGradient,
        ::BC
        ::BL,
        state_prognostic⁺::Vars,
        state_auxiliary⁺::Vars,
        normal⁻,
        state_prognostic⁻::Vars,
        state_auxiliary⁻::Vars,
        t
    )
    boundary_state!(
        ::NumericalFluxFirstOrder,
        ::BC
        ::BL,
        state_prognostic⁺::Vars,
        state_auxiliary⁺::Vars,
        normal⁻,
        state_prognostic⁻::Vars,
        state_auxiliary⁻::Vars,
        t
    )
    boundary_state!(
        ::NumericalFluxSecondOrder,
        ::BC
        ::BL,
        state_prognostic⁺::Vars,
        state_gradient_flux⁺::Vars,
        state_auxiliary⁺:
        Vars, normal⁻,
        state_prognostic⁻::Vars,
        state_gradient_flux⁻::Vars,
        state_auxiliary⁻::Vars,
        t
    )

Specify the opposite (+ side) face for the boundary condition type `BC` with balance law `BL`.

 - `NumericalFluxGradient` numerical flux (internal method)
 - `NumericalFluxFirstOrder` first-order unknowns
 - `NumericalFluxSecondOrder` second-order unknowns

"""
function boundary_state! end


text:
### Source
###
### [3] Moré JJ, Garbow BS, Hillstrom KE: Testing unconstrained optimization software. ACM T Math Software. 1981
###

##########################
### Extended Rosenbrock
###
### Problem (21) from [3]
##########################

function extrosenbrock(x::AbstractArray, param::MatVecHolder)
    # TODO: we could do this without the xt storage holder
    n = length(x)
    jodd = 1:2:n-1
    jeven = 2:2:n

    xt = param.vec
    @. xt[jodd] = 10.0 * (x[jeven] - x[jodd]^2)
    @. xt[jeven] = 1.0 - x[jodd]

    return 0.5*sum(abs2, xt)
end

function extrosenbrock_gradient!(storage::AbstractArray,
                                 x::AbstractArray, param::MatVecHolder)
    n = length(x)
    jodd = 1:2:n-1
    jeven = 2:2:n
    xt = param.vec
    @. xt[jodd] = 10.0 * (x[jeven] - x[jodd]^2)
    @. xt[jeven] = 1.0 - x[jodd]

    @. storage[jodd] = -20.0 * x[jodd] * xt[jodd] - xt[jeven]
    @. storage[jeven] = 10.0 * xt[jodd]
end

function extrosenbrock_fun_gradient!(storage::AbstractArray,
                                     x::AbstractArray, param::MatVecHolder)
    n = length(x)
    jodd = 1:2:n-1
    jeven = 2:2:n
    xt = param.vec
    @. xt[jodd] = 10.0 * (x[jeven] - x[jodd]^2)
    @. xt[jeven] = 1.0 - x[jodd]

    @. storage[jodd] = -20.0 * x[jodd] * xt[jodd] - xt[jeven]
    @. storage[jeven] = 10.0 * xt[jodd]
    return 0.5*sum(abs2, xt)
end

function extrosenbrock_hessian!(storage,x,param)
    error("Hessian not implemented for Extended Rosenbrock")
end

function _extrosenbrockproblem(N::Int;
                               initial_x::AbstractArray{T} = repeat([-1.2,1]; inner = round(Int, N/2)),
                               name::AbstractString = "Extended Rosenbrock ($N)") where T
    @assert mod(N,2) == 0
    OptimizationProblem(name,
                        extrosenbrock,
                        extrosenbrock_gradient!,
                        extrosenbrock_fun_gradient!,
                        extrosenbrock_hessian!,
                        nothing, # Constraints
                        initial_x,
                        fill(T(1), length(initial_x)),
                        zero(T),
                        true,
                        false,
                        MatVecHolder(Array{T}(undef, 0,0),similar(initial_x)))
end

examples["Extended Rosenbrock"] = _extrosenbrockproblem(100)

##########################
### Extended Powell
###
### Problem (22) from [3]
##########################

function extpowell(x::AbstractArray, param::MatVecHolder)
    # TODO: we could do this without the xt storage holder
    n = length(x)
    j1 = 1:4:n-3;
    j2 = 2:4:n-2;
    j3 = 3:4:n-1;
    j4 = 4:4:n;

    xt = param.vec
    @. xt[j1] = x[j1] + 10*x[j2]
    @. xt[j2] = sqrt(5)*(x[j3]-x[j4])
    @. xt[j3] = (x[j2] - 2*x[j3])^2;
    @. xt[j4] = sqrt(10)*(x[j1]-x[j4])^2;

    return 0.5*sum(abs2, xt)
end

function extpowell_gradient!(storage::AbstractArray,
                             x::AbstractArray, param::MatVecHolder)
    # TODO: we could do this without the xt storage holder
    n = length(x)
    j1 = 1:4:n-3;
    j2 = 2:4:n-2;
    j3 = 3:4:n-1;
    j4 = 4:4:n;

    xt = param.vec
    @. xt[j1] = x[j1] + 10*x[j2]
    @. xt[j2] = sqrt(5)*(x[j3]-x[j4])
    @. xt[j3] = (x[j2] - 2*x[j3])^2;
    @. xt[j4] = sqrt(10)*(x[j1]-x[j4])^2;

    @. storage[j1] = xt[j1] + 2*sqrt(10)*(x[j1]-x[j4]).*xt[j4];
    @. storage[j2] = 10*xt[j1] + 2*(x[j2]-2*x[j3]).*xt[j3];
    @. storage[j3] = sqrt(5)*xt[j2] - 4*(x[j2]-2*x[j3]).*xt[j3];
    @. storage[j4] = -sqrt(5)*xt[j2] - 2*sqrt(10)*(x[j1]-x[j4]).*xt[j4];
end

function extpowell_fun_gradient!(storage::AbstractArray,
                                 x::AbstractArray, param::MatVecHolder)
    # TODO: we could do this without the xt storage holder
    n = length(x)
    j1 = 1:4:n-3;
    j2 = 2:4:n-2;
    j3 = 3:4:n-1;
    j4 = 4:4:n;

    xt = param.vec
    @. xt[j1] = x[j1] + 10*x[j2]
    @. xt[j2] = sqrt(5)*(x[j3]-x[j4])
    @. xt[j3] = (x[j2] - 2*x[j3])^2;
    @. xt[j4] = sqrt(10)*(x[j1]-x[j4])^2;

    @. storage[j1] = xt[j1] + 2*sqrt(10)*(x[j1]-x[j4]).*xt[j4];
    @. storage[j2] = 10*xt[j1] + 2*(x[j2]-2*x[j3]).*xt[j3];
    @. storage[j3] = sqrt(5)*xt[j2] - 4*(x[j2]-2*x[j3]).*xt[j3];
    @. storage[j4] = -sqrt(5)*xt[j2] - 2*sqrt(10)*(x[j1]-x[j4]).*xt[j4];
    return 0.5*sum(abs2, xt)
end

function extpowell_hessian!(storage,x,param)
    error("Hessian not implemented for Extended Powell")
end

function _extpowellproblem(N::Int;
                           initial_x::AbstractArray{T} = repeat(float([3,-1,0,1]); inner = round(Int, N/4)),
                           name::AbstractString = "Extended Powell ($N)") where T
    @assert mod(N,4) == 0
    OptimizationProblem(name,
                        extpowell,
                        extpowell_gradient!,
                        extpowell_fun_gradient!,
                        extpowell_hessian!,
                        nothing, # Constraints
                        initial_x,
                        fill(zero(T), length(initial_x)),
                        zero(T),
                        true,
                        false,
                        MatVecHolder(Array{T}(undef, 0,0),similar(initial_x)))
end

examples["Extended Powell"] = _extpowellproblem(100)

##########################
### Penalty function I
###
### Problem (23) from [3]
### Default alpha = sqrt(1e-5)
##########################

function penfunI(x::AbstractArray, param)
    # TODO: we could do this without the xt storage holder
    xt = param.xt
    @. xt = param.alpha*(x-one(eltype(x)))

    xtend = sum(abs2,x)-0.25
    return 0.5*(sum(abs2, xt) + abs2(xtend)) # TODO: make 0.25 a parameter
end

function penfunI_gradient!(storage::AbstractArray,
                           x::AbstractArray, param)
    # TODO: we could do this without the xt storage holder
    xt = param.xt
    @. xt = param.alpha*(x-one(eltype(x)))

    xtend = sum(abs2,x)-0.25
    @. storage = param.alpha*xt + 2.0*xtend*x
end

function penfunI_fun_gradient!(storage::AbstractArray,
                               x::AbstractArray, param)
    # TODO: we could do this without the xt storage holder
    xt = param.xt
    @. xt = param.alpha*(x-one(eltype(x)))

    xtend = sum(abs2,x)-0.25
    @. storage = param.alpha*xt + 2.0*xtend*x
    return 0.5*(sum(abs2, xt) + abs2(xtend)) # TODO: make 0.25 a parameter
end

function penfunI_hessian!(storage,x,param)
    error("Hessian not implemented for Penalty Function I")
end

function _penfunIproblem(N::Int;
                         initial_x::AbstractArray{T} = collect(float(1:N)),
                         alpha::T = sqrt(1e-5),
                         name::AbstractString = "Penalty Function I ($N)") where T
    if N == 100
        # Calculated numerically to high precision |g(x)|<1e-15
        fsol = 0.00045124548840214817
        xsol = fill(0.05000949719895305, N)
    elseif N == 200
        # Calculated numerically to high precision |g(x)|<1e-15
        fsol = 0.000930530019118627
        xsol = fill(0.03536498146451683, N)
    else
        fsol  = NaN
        xsol = fill(NaN, N)
    end

    OptimizationProblem(name,
                        penfunI,
                        penfunI_gradient!,
                        penfunI_fun_gradient!,
                        penfunI_hessian!,
                        nothing, # Constraints
                        initial_x,
                        xsol,
                        fsol,
                        true,
                        false,
                        ParaboloidStruct(Array{T}(undef, 0,0),Array{T}(undef, 0),
                                         similar(initial_x), alpha))
end

examples["Penalty Function I"] = _penfunIproblem(100)

##########################
### Trigonometric function
###
### Problem (26) from [3]
##########################

function trigonometric(x::AbstractArray, param)
    # TODO: we could do this without the xt storage holder
    n = length(x)
    xt = param.vec
    scos = sum(cos,x)
    @. xt = n + (1:n)*(one(eltype(x)) - cos(x)) - sin(x) - scos

    return 0.5*sum(abs2, xt)
end

function trigonometric_gradient!(storage::AbstractArray,
                                 x::AbstractArray, param)
    # TODO: we could do this without the xt storage holder
    n = length(x)
    xt = param.vec
    scos = sum(cos,x)
    @. xt = n + (1:n)*(one(eltype(x)) - cos(x)) - sin(x) - scos

    sxt = sum(xt)
    @. storage = sxt*sin(x) + xt * ((1:n)*sin(x) - cos(x))
end

function trigonometric_fun_gradient!(storage::AbstractArray,
                                     x::AbstractArray, param)
    # TODO: we could do this without the xt storage holder
    n = length(x)
    xt = param.vec
    scos = sum(cos,x)
    @. xt = n + (1:n)*(one(eltype(x)) - cos(x)) - sin(x) - scos

    sxt = sum(xt)
    @. storage = sxt*sin(x) + xt * ((1:n)*sin(x) - cos(x))
    return 0.5*sum(abs2, xt)
end

function trigonometric_hessian!(storage,x,param)
    error("Hessian not implemented for Trigonometric Function")
end

function _trigonometricproblem(N::Int;
                               initial_x::AbstractArray{T} = ones(N)/N,
                               alpha::T = sqrt(1e-5),
                               name::AbstractString = "Trigonometric ($N)") where T
    OptimizationProblem(name,
                        trigonometric,
                        trigonometric_gradient!,
                        trigonometric_fun_gradient!,
                        trigonometric_hessian!,
                        nothing, # Constraints
                        initial_x,
                        fill(T(0), length(initial_x)),
                        zero(T),
                        true,
                        false,
                        MatVecHolder(Array{T}(undef, 0,0),similar(initial_x)))
end

examples["Trigonometric"] = _trigonometricproblem(100)

##########################################################################
###
### Beale (2D)
###
### Problem 5 in [3]
###
### Sum-of-squares objective, non-convex with g'*inv(H)*g == 0 at the
### initial position.
###
##########################################################################

### General utilities for sum-of-squares functions
# TODO: Update the other problems that are not Beale to use sumsq as well?

# Requires f(x) and J(x) computes the values and jacobian at x of a set of functions, and
# that H(x, i) computes the hessian of the ith function

sumsq_obj(f, x) = sum(f(x).^2)

function sumsq_gradient!(g::AbstractVector, f, J, x::AbstractVector)
    copyto!(g, sum((2.0 .* f(x)) .* J(x), dims = 1))
end

function sumsq_hessian!(h::AbstractMatrix, f, J, H, x::AbstractVector)
    fx = f(x)
    Jx = J(x)
    htmp = 2.0 .* (Jx' * Jx)
    for i = 1:length(fx)
        htmp += (2.0 * fx[i]) * H(x, i)
    end
    copyto!(h, htmp)
end

const beale_y = [1.5, 2.25, 2.625]

beale_f(x) = [beale_y[i] - x[1]*(1-x[2]^i) for i = 1:3]
beale_J(x) = hcat([-(1-x[2]^i) for i = 1:3],
                  [i*x[1]*x[2]^(i-1) for i = 1:3])
function beale_H(x, i)
    od = i*x[2]^(i-1)
    d2 = i > 1 ? i*(i-1)*x[1]*x[2]^(i-2) : zero(x[2])
    [0 od; od d2]
end

beale(x::AbstractVector) = sumsq_obj(beale_f, x)

function beale_gradient!(g::AbstractVector, x::AbstractVector)
    sumsq_gradient!(g, beale_f, beale_J, x)
end

function beale_hessian!(h::AbstractMatrix, x::AbstractVector)
    sumsq_hessian!(h, beale_f, beale_J, beale_H, x)
end

examples["Beale"] = OptimizationProblem("Beale",
                                        beale,
                                        beale_gradient!,
                                        nothing,
                                        beale_hessian!,
                                        nothing, # Constraints
                                        [1.0, 1.0],
                                        [3.0, 0.5],
                                        beale([3.0, 0.5]),
                                        true,
                                        true)


text:
#=
    Implement the vectorized OPF model proposed in:

    Feasible Path Identification in Optimal Power Flow with Sequential Convex Restriction
    Dongchan Lee, Konstantin Turitsyn, Daniel K. Molzahn, Line A. Roald

    Preprint: https://arxiv.org/abs/1906.09483
    Code:     https://github.com/dclee131/PowerFlowCVXRS

    The vectorized version of the OPF leads to a compact
    expression tree, with all the nonlinearities factorized inside
    a single nonlinear basis ψ. Comparing to the scalar version,
    the vectorized formulation allows to streamline the evaluation
    of the derivatives.

    The expression tree is given directly as:

                                    /------ (power flow equations)
                                   /
    (vm, va, pg) --->  ψ(vm, va) -/-------- (reactive power generations qg)
                                  \
                                   \------- (line flow)

    In addition, this model considers only a subset of the power
    flow equations, and lead to a more compact KKT system than
    the formulation adopted in PowerModels.jl.

=#

time_start = time()

using LinearAlgebra
using SparseArrays
using JuMP, Ipopt
using PowerModels

pkg_load_time = time() - time_start

time_start = time()

file_name = "../data/pglib_opf_case5_pjm.m"
# file_name = "../data/pglib_opf_case118_ieee.m"

#=
    Import data with PowerModels
=#
data = PowerModels.parse_file(file_name)
PowerModels.standardize_cost_terms!(data, order=2)
PowerModels.calc_thermal_limits!(data)
ref = PowerModels.build_ref(data)[:it][:pm][:nw][0]

data_load_time = time() - time_start

gen = data["gen"]
bus = data["bus"]
branch = data["branch"]
shunt = data["shunt"]
load = data["load"]
nbus = length(data["bus"])
ngen = length(data["gen"])
nlines = length(data["branch"])
nloads = length(data["load"])
nshunts = length(data["shunt"])
## GENERATORS
pg0 = Float64[gen["$i"]["pg"] for i in 1:ngen]
pgmin = Float64[gen["$i"]["pmin"] for i in 1:ngen]
pgmax = Float64[gen["$i"]["pmax"] for i in 1:ngen]
qgmin = Float64[gen["$i"]["qmin"] for i in 1:ngen]
qgmax = Float64[gen["$i"]["qmax"] for i in 1:ngen]
gen2bus = Int[gen["$i"]["gen_bus"] for i in 1:ngen]
## BUSES
busid = Int[bus["$i"]["bus_i"] for i in 1:nbus]
bustype = Int[bus["$i"]["bus_type"] for i in 1:nbus]
vm0 = Float64[bus["$i"]["vm"] for i in 1:nbus]
va0 = Float64[bus["$i"]["va"] for i in 1:nbus]
vmin = Float64[bus["$i"]["vmin"] for i in 1:nbus]
vmax = Float64[bus["$i"]["vmax"] for i in 1:nbus]
## SHUNTS
b_sh = Int[shunt["$i"]["shunt_bus"] for i in 1:nshunts]
gs_sh = Float64[shunt["$i"]["gs"] for i in 1:nshunts]
bs_sh = Float64[shunt["$i"]["bs"] for i in 1:nshunts]
shunts = gs_sh .+ 1im .* bs_sh
## LOADS
load2bus = Int[load["$i"]["load_bus"] for i in 1:nloads]
pd = Float64[load["$i"]["pd"] for i in 1:nloads]
qd = Float64[load["$i"]["qd"] for i in 1:nloads]
## BRANCHES
f_bus = Int[branch["$i"]["f_bus"] for i in 1:nlines]
t_bus = Int[branch["$i"]["t_bus"] for i in 1:nlines]
tap = Float64[branch["$i"]["tap"] for i in 1:nlines]
stat = Float64[branch["$i"]["br_status"] for i in 1:nlines]
br_r = Float64[branch["$i"]["br_r"] for i in 1:nlines]
br_x = Float64[branch["$i"]["br_x"] for i in 1:nlines]
b_fr = Float64[branch["$i"]["b_fr"] for i in 1:nlines]
b_to = Float64[branch["$i"]["b_to"] for i in 1:nlines]
br_b = b_fr + b_to
shift = Float64[branch["$i"]["shift"] for i in 1:nlines]
sline_max = Float64[branch["$i"]["rate_a"] for i in 1:nlines]
## COSTS
c2 = Float64[gen["$i"]["cost"][1] for i in 1:ngen]
c1 = Float64[gen["$i"]["cost"][2] for i in 1:ngen]
c0 = Float64[gen["$i"]["cost"][3] for i in 1:ngen]

#=
    Buses' classification
=#
pv = findall(isequal(2), bustype)
pq = findall(isequal(1), bustype)
ref = findall(isequal(3), bustype)
refgen = findfirst(isequal(ref[1]), gen2bus)
npv = length(pv)
npq = length(pq)
nref = length(ref)

#=
    Build incidence matrices encoding the problem's topology
=#
Cg = sparse(gen2bus, 1:ngen, ones(ngen), nbus, ngen)
Cl = sparse(load2bus, 1:nloads, ones(nloads), nbus, nloads)
Cf = sparse(f_bus, 1:nlines, ones(nlines), nbus, nlines)
Ct = sparse(t_bus, 1:nlines, ones(nlines), nbus, nlines)
E = Cf - Ct
Ysh = sparse(b_sh, b_sh, shunts, nbus, nbus)

#=
    Build admittance matrices
    (take expressions from MATPOWER)
=#
Ys = stat ./ (br_r + 1im * br_x)       # series admittance
Bc = stat .* br_b                      # line charging susceptance
tap = tap .* exp.(1im*pi/180 .* shift) # add phase shifters
Ytt = Ys .+ 1im .* Bc ./ 2.0
Yff = Ytt ./ (tap .* conj.(tap))
Yft = - Ys ./ conj.(tap)
Ytf = - Ys ./ tap

i = [1:nlines; 1:nlines]
Yf = sparse(i, [f_bus; t_bus], [Yff; Yft], nlines, nbus)
Yt = sparse(i, [f_bus; t_bus], [Ytf; Ytt], nlines, nbus)
Ybus = Cf * Yf + Ct * Yt + Ysh

Yc = Cf * Diagonal(Yft) + Ct * Diagonal(Ytf)
Ys = Cf * Diagonal(Yft) - Ct * Diagonal(Ytf)
Yd = Cf * Diagonal(Yff) * Cf' + Ct * Diagonal(Ytt) * Ct' + Ysh

yff = Diagonal(Yff)
yft = Diagonal(Yft)
ytf = Diagonal(Ytf)
ytt = Diagonal(Ytt)

#=
    Compact matrices
=#
# Eq (5): power flow equations
M = [ real(Yc)  imag(Ys)  real(Yd);
     -imag(Yc)  real(Ys) -imag(Yd)]
# Eqs (9-10): line flow constraints
Lfp = [real(yft)  imag(yft)  real(yff) * Cf']
Lfq = [-imag(yft) real(yft) -imag(yff) * Cf']
Ltp = [real(ytf)  -imag(ytf)  real(ytt) * Ct']
Ltq = [-imag(ytf) -real(ytf) -imag(ytt) * Ct']

# Eq (7): Power flow (equality)
M_eq = M[[pv; pq; nbus .+ pq], :]
Cg_eq = [Cg[pv, :] ; spzeros(2 * npq, ngen)]
τ_eq = [Cl[[pv; pq], :] * pd; Cl[pq, :] * qd]
# Eq (8): Power flow (inequality)
M_ineq = M[[ref; nbus .+ ref; nbus .+ pv], :]
τ_ineq = [Cl[ref, :] * pd; Cl[[ref; pv], :] * qd]
pmin = [Cg[ref, :] * pgmin; Cg[[ref;pv], :] *qgmin]
pmax = [Cg[ref, :] * pgmax; Cg[[ref;pv], :] *qgmax]

#=
    Build JuMP model
=#
time_start = time()

acopf = Model(Ipopt.Optimizer)

## 1. Variables

@variable(acopf, vmin[i] <= vm[i=1:nbus] <= vmax[i], start=vm0[i])
@variable(acopf, va[i=1:nbus], start=va0[i])
@variable(acopf, pgmin[i] <= pg[i=1:ngen] <= pgmax[i], start=pg0[i])
# nonlinear basis (encodes all the nonlinearities in the problem)
@variable(acopf, ψsin[i=1:nlines])
@variable(acopf, ψcos[i=1:nlines])
@variable(acopf, ψq[i=1:nbus])
# line flow
@variable(acopf, sfp[i = 1:nlines])
@variable(acopf, stp[i = 1:nlines])
@variable(acopf, sfq[i = 1:nlines])
@variable(acopf, stq[i = 1:nlines])

## 2. Constraints

@constraint(acopf, va[ref] .== 0)
# Angle difference
@expression(acopf, φ, E' * va)
# The nonlinear basis encodes all the nonlinearities in the problem.
@NLconstraint(acopf, [i=1:nlines], ψsin[i] == vm[f_bus[i]] * vm[t_bus[i]] * sin(φ[i]))
@NLconstraint(acopf, [i=1:nlines], ψcos[i] == vm[f_bus[i]] * vm[t_bus[i]] * cos(φ[i]))
@NLconstraint(acopf, [i=1:nbus],  ψq[i] == vm[i]^2)

# eq(11b): recover power flow equations with sparse operations
@constraint(acopf, τ_eq - Cg_eq * pg + M_eq * [ψcos; ψsin; ψq] .== 0)

# eq(11c) - eq(20): active power bounds on slack + reactive power bounds
@constraint(acopf, pmin .<= M_ineq * [ψcos; ψsin; ψq] .+ τ_ineq .<= pmax)

# Line-flow constraints
# eq(11e) - eq(20)
@constraint(acopf, sfp .== Lfp * [ψcos; ψsin; ψq])
@constraint(acopf, sfq .== Lfq * [ψcos; ψsin; ψq])
@constraint(acopf, stp .== Ltp * [ψcos; ψsin; ψq])
@constraint(acopf, stq .== Ltq * [ψcos; ψsin; ψq])
@constraint(acopf, sfp.^2 .+ sfq.^2  .<= sline_max.^2)
@constraint(acopf, stp.^2 .+ stq.^2  .<= sline_max.^2)

# Recover active power generation at slack node
@constraint(acopf, pg[refgen] .== M_ineq[1:1, :] * [ψcos; ψsin; ψq] .+ τ_ineq[1:1])

## 3. Objective

@objective(acopf, Min, dot(c2, pg.^2) + dot(c1, pg) + sum(c0))

model_build_time = time() - time_start

#=
    Resolution
=#

time_start = time()

optimize!(acopf)

solve_time = time() - time_start

#=
    Analysis
=#

nlp_block = MOI.get(acopf, MOI.NLPBlock())

println("")
println("\033[1mSummary\033[0m")
println("   case......: $(file_name)")
println("   cost......: $(round(Int, JuMP.objective_value(acopf)))")
println("   pkg time..: $(pkg_load_time)")
println("   data time.: $(data_load_time)")
println("   build time: $(model_build_time)")
println("   solve time: $(solve_time)")
println("   callbacks time:")
println("   * obj.....: $(nlp_block.evaluator.eval_objective_timer)")
println("   * grad....: $(nlp_block.evaluator.eval_objective_gradient_timer)")
println("   * cons....: $(nlp_block.evaluator.eval_constraint_timer)")
println("   * jac.....: $(nlp_block.evaluator.eval_constraint_jacobian_timer)")
println("   * hesslag.: $(nlp_block.evaluator.eval_hessian_lagrangian_timer)")
println("")

max_stars_repo_name: lanl-ansi/rosetta-opf

max_issues_repo_name: lanl-ansi/rosetta-opf

max_forks_repo_name: lanl-ansi/rosetta-opf


text:
# TODO :
#   - design optimizer functionality

abstract type Layer end

abstract type Loss <: Layer end

abstract type Activation <: Layer end

abstract type AbstractRNN <: Layer end


text:
using Mdp
using Test

@testset "VI.jl" begin
    include("test_vi.jl")
end

@testset "FH.jl" begin
    include("test_fh.jl")
end

include("test_pricing_problem.jl")

@testset "PricingAlgorithmFH.jl" begin
    include("test_pricing_algorithm_fh.jl")
end


text:
using Test
using ARMSampling
using Statistics

@testset "TrapezoidalProposal" begin
    logf(x) = -(x-1)^2
    x = [-3.0, 0.0, 1.0, 2.0, 5.0]
    y = [logf(x1) for x1 in x]

    proposal = TrapezoidalProposal(x, y)
    @test length(proposal) == 3
    @test proposal.w[2] ≈ (ℯ - 1.0)/ℯ
    @test proposal.w[3] ≈ (ℯ - 1.0)/ℯ
    @test x[1] <= sample(proposal) <= x[end]

    insert!(proposal, 0.5, logf(0.5))
    @test length(proposal) == 4
    @test sum(proposal.w[2:3]) > proposal.w[4]

    insert!(proposal, 1.5, logf(1.5))
    @test length(proposal) == 5
    @test sum(proposal.w[2:3]) ≈ sum(proposal.w[4:5])
    @test x[1] <= sample!(proposal, logf, 1.1) <= x[end]

    mysample = [1.1]
    sample_size = 20000
    sample!(mysample, sample_size, proposal, logf)
    @test length(mysample) == sample_size + 1
    @test all( proposal.x[1] .<= mysample .<= proposal.x[end] )
    @test abs(std(mysample) - sqrt(0.5)) < 1e-2
    @test abs(mean(mysample) - 1.0) < 1e-2
end


text:
using Documenter
using MetaGraphsNext
using Test

@testset "MetaGraphsNext" begin
    doctest(MetaGraphsNext)
end


text:
using OrdinaryDiffEq, HyperbolicDiffEq

function compute(K, N, fnum, compute_entropy_var, diss)
    xmin = -1.0
    xmax = +1.0
    tspan = (0., 1.0)
    ode_solver = SSPRK104()

    mesh = UniformPeriodicMesh1D(xmin, xmax, N)
    balance_law = Cubic()
    reconstruction = ModifiedENO(K)
    semidisc = ScalarUniformPeriodicTecno1D(balance_law, mesh, fnum, diss, reconstruction, compute_entropy_var)
    ode = semidiscretise(semidisc, sinpi, tspan)

    solve(ode, ode_solver, dt=1/N, save_everystep=false);
end

let K = 1
    N = 2^10
    fnum = CentralFlux()
    compute_entropy_var = (u,model) -> u^3
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * (wr-wl) # TeCNO
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = CentralFlux()
    compute_entropy_var = (u,model) -> u^3
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * abs(wr-wl)^(K-1) * (wr-wl) #ELW
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = EnergyConservativeFlux()
    compute_entropy_var = (u,model) -> u
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * (wr-wl) # TeCNO
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = EnergyConservativeFlux()
    compute_entropy_var = (u,model) -> u
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * abs(wr-wl)^(K-1) * (wr-wl) #ELW
    compute(K, N, fnum, compute_entropy_var, diss)
end

let K = 2
    N = 2^10
    fnum = CentralFlux()
    compute_entropy_var = (u,model) -> u^3
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * (wr-wl) # TeCNO
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = CentralFlux()
    compute_entropy_var = (u,model) -> u^3
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * abs(wr-wl)^(K-1) * (wr-wl) #ELW
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = EnergyConservativeFlux()
    compute_entropy_var = (u,model) -> u
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * (wr-wl) # TeCNO
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = EnergyConservativeFlux()
    compute_entropy_var = (u,model) -> u
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * abs(wr-wl)^(K-1) * (wr-wl) #ELW
    compute(K, N, fnum, compute_entropy_var, diss)
end

let K = 3
    N = 2^10
    fnum = CentralFlux()
    compute_entropy_var = (u,model) -> u^3
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * (wr-wl) # TeCNO
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = CentralFlux()
    compute_entropy_var = (u,model) -> u^3
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * abs(wr-wl)^(K-1) * (wr-wl) #ELW
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = EnergyConservativeFlux()
    compute_entropy_var = (u,model) -> u
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * (wr-wl) # TeCNO
    compute(K, N, fnum, compute_entropy_var, diss)

    fnum = EnergyConservativeFlux()
    compute_entropy_var = (u,model) -> u
    diss = (wl, wr, model) -> -0.5*max_abs_speed(wl, wr, model) * abs(wr-wl)^(K-1) * (wr-wl) #ELW
    compute(K, N, fnum, compute_entropy_var, diss)
end