text:
struct FunkyConnectivityModel{Tv} <: AbstractConnectivityModel
    α::Vector{Tv}
    β_vertex::Tv
    β_pin::Tv
    β_net::Tv
end

(mdl::FunkyConnectivityModel)(n_vertices, n_pins, n_nets, k) = mdl.α[k] + n_vertices * mdl.β_vertex + n_pins * mdl.β_pin + n_nets * mdl.β_net

@inline ChainPartitioners.cost_type(::Type{FunkyConnectivityModel{Tv}}) where {Tv} = Tv

struct FunkyMonotonizedSymmetricConnectivityModel{Tv} <: AbstractMonotonizedSymmetricConnectivityModel
    α::Vector{Tv}
    β_vertex::Tv
    β_pin::Tv
    β_net::Tv
    Δ_pins::Tv
end

@inline ChainPartitioners.cost_type(::Type{FunkyMonotonizedSymmetricConnectivityModel{Tv}}) where {Tv} = Tv

(mdl::FunkyMonotonizedSymmetricConnectivityModel)(n_vertices, n_pins, n_nets, k) = mdl.α[k] + n_vertices * mdl.β_vertex + n_pins * mdl.β_pin + n_nets * mdl.β_net

struct FunkyPrimaryConnectivityModel{Tv} <: AbstractPrimaryConnectivityModel
    α::Vector{Tv}
    β_vertex::Tv
    β_pin::Tv
    β_local_net::Tv
    β_remote_net::Tv
end

@inline ChainPartitioners.cost_type(::Type{FunkyPrimaryConnectivityModel{Tv}}) where {Tv} = Tv

(mdl::FunkyPrimaryConnectivityModel)(n_vertices, n_pins, n_local_nets, n_remote_nets, k) = mdl.α[k] + n_vertices * mdl.β_vertex + n_pins * mdl.β_pin + n_local_nets * mdl.β_local_net + n_remote_nets * mdl.β_remote_net

function ChainPartitioners.bound_stripe(A::SparseMatrixCSC, K, mdl::Union{FunkyConnectivityModel, FunkyMonotonizedSymmetricConnectivityModel})
    ocl = oracle_stripe(mdl, A)
    m, n = size(A)
    args = (minimum(mdl.α), maximum(mdl.α), maximum(ocl(1, n + 1, k) for k = 1:K))
    return (minimum(args), maximum(args))
end

function ChainPartitioners.bound_stripe(A::SparseMatrixCSC, K, Π, mdl::Union{FunkyPrimaryConnectivityModel})
    ocl = oracle_stripe(mdl, A, Π)
    m, n = size(A)
    args = (minimum(mdl.α), maximum(mdl.α), maximum(ocl(1, n + 1, k) for k = 1:K))
    return (minimum(args), maximum(args))
end

LazyBisectCost = Union{AbstractConnectivityModel, AbstractMonotonizedSymmetricConnectivityModel, AbstractPrimaryConnectivityModel}

struct ConvexWorkModel <: AbstractWorkCost
    α::Float64
    β_vertex::Float64
    β_pin::Float64
end

(mdl::ConvexWorkModel)(n_vertices, n_pins) = mdl.α + (n_vertices * mdl.β_vertex + n_pins * mdl.β_pin)^0.8

@inline ChainPartitioners.cost_type(::Type{ConvexWorkModel}) = Float64

struct ConcaveWorkModel <: AbstractWorkCost
    α::Float64
    β_vertex::Float64
    β_pin::Float64
end

(mdl::ConcaveWorkModel)(n_vertices, n_pins) = mdl.α + (n_vertices * mdl.β_vertex + n_pins * mdl.β_pin)^2

@inline ChainPartitioners.cost_type(::Type{ConcaveWorkModel}) = Float64

@testset "Partitioners" begin
    for A in [
        [
            matrices["LPnetlib/lpi_itest6"],
            SparseMatrixCSC(matrices["HB/can_292"]),
        ];
        reshape([sprand(m, n, 0.1) for m = [1, 2, 3, 4, 8], n = [1, 2, 3, 4, 8], trial = 1:4], :);
    ]
        (m, n) = size(A)
        for K = [1, 2, 3, 4, 8]
            for f = [
                AffineWorkModel(0, 10, 1);
                AffineConnectivityModel(0, 3, 1, 3);
                AffinePrimaryConnectivityModel(0, 2, 1, 3, 6);
                m == n ? AffineMonotonizedSymmetricConnectivityModel(0, 3, 1, 3, 5) : [];
                FunkyConnectivityModel(rand(1:10, K), 3, 1, 3);
                FunkyPrimaryConnectivityModel(rand(1:10, K), 2, 1, 3, 6);
                m == n ? FunkyMonotonizedSymmetricConnectivityModel(rand(1:10, K), 3, 1, 3, 5) : [];
            ]
                Π = partition_stripe(A', K, EquiSplitter())
                Φ = partition_stripe(A, K, ReferenceBottleneckSplitter(f), Π)
                c = bottleneck_value(A, Π, Φ, f)
                for (method, ϵ) = [
                    (ReferenceBottleneckSplitter(f), 0);
                    (DynamicBottleneckSplitter(f), 0);
                    (BisectIndexBottleneckSplitter(f), 0);
                    (BisectCostBottleneckSplitter(f, 0.1), 0.1);
                    (BisectCostBottleneckSplitter(f, 0.01), 0.01);
                    (LazyBisectCostBottleneckSplitter(f, 0.1), 0.1);
                    (LazyBisectCostBottleneckSplitter(f, 0.01), 0.01);
                ]
                    Φ′ = partition_stripe(A, K, method, Π)
                    @test issorted(Φ′.spl)
                    @test Φ′.spl[1] == 1
                    @test Φ′.spl[end] == n + 1
                    @test Φ′.K == K
                    @test bottleneck_value(A, Π, Φ′, f) <= bottleneck_value(A, Π, Φ, f) * (1 + ϵ)
                end
            end

            # Creating nonnegative monotonic decreasing cost functions requires some care.
            # These complicated affine terms are designed to ensure nonnegativity.
            # Affine cost functions are commented out because their upper and lower bounds assume nonnegative coefficients.
            for f = [
                #AffineWorkModel(1 + nnz(A), 0, -1);
                #AffineConnectivityModel(1 + nnz(A) + 3n + 3m, -3, -1, -3);
                #AffinePrimaryConnectivityModel(1 + nnz(A) + 3n + 6m, -2, -1, -3, -6);
                AffineSecondaryConnectivityModel(0, 2, 1, 3, 6);
                FunkyConnectivityModel(1 + nnz(A) + 3n + 3m .+ rand(1:10, K), -3, -1, -3);
                FunkyPrimaryConnectivityModel(1 + nnz(A) + 3n + 6m .+ rand(1:10, K), -2, -1, -3, -6);
                #m == n ? AffineMonotonizedSymmetricConnectivityModel(1 + nnz(A) + 18n + 3m, -3, -1, -3, 5) : [];
                m == n ? FunkyMonotonizedSymmetricConnectivityModel(1 + nnz(A) + 18n + 3m .+ rand(1:10, K), -3, -1, -3, 5) : [];
            ]
                Π = partition_stripe(A', K, EquiSplitter())
                Φ = partition_stripe(A, K, ReferenceBottleneckSplitter(f), Π)
                c = bottleneck_value(A, Π, Φ, f)
                for (method, ϵ) = [
                    (ReferenceBottleneckSplitter(f), 0);
                    (DynamicBottleneckSplitter(f), 0);
                    (FlipBisectIndexBottleneckSplitter(f), 0);
                    (FlipBisectCostBottleneckSplitter(f, 0.1), 0.1);
                    (FlipBisectCostBottleneckSplitter(f, 0.01), 0.01);
                    (FlipBisectCostBottleneckSplitter(f, 0.001), 0.001);
                    (FlipBisectCostBottleneckSplitter(f, 0.0001), 0.0001);
                    (LazyFlipBisectCostBottleneckSplitter(f, 0.1), 0.1);
                    (LazyFlipBisectCostBottleneckSplitter(f, 0.01), 0.01);
                    (LazyFlipBisectCostBottleneckSplitter(f, 0.001), 0.001);
                    (LazyFlipBisectCostBottleneckSplitter(f, 0.0001), 0.0001);
                ]
                    Φ′ = partition_stripe(A, K, method, Π)
                    @test issorted(Φ′.spl)
                    @test Φ′.spl[1] == 1
                    @test Φ′.spl[end] == n + 1
                    @test Φ′.K == K
                    @test bottleneck_value(A, Π, Φ′, f) <= bottleneck_value(A, Π, Φ, f) * (1 + ϵ)
                end
            end

            for f = [
                AffineConnectivityModel(0, 3, 1, 3);
                m == n ? AffineMonotonizedSymmetricConnectivityModel(0, 3, 1, 3, 5) : [];
            ]
                Π = partition_stripe(A', K, EquiSplitter())
                Φ = partition_stripe(A, K, ReferenceTotalSplitter(f), Π)
                c = total_value(A, Π, Φ, f)
                for (method, ϵ) = [
                    (ReferenceTotalSplitter(f), 0);
                    (DynamicTotalSplitter(f), 0);
                ]
                    Φ′ = partition_stripe(A, K, method, Π)
                    @test issorted(Φ′.spl)
                    @test Φ′.spl[1] == 1
                    @test Φ′.spl[end] == n + 1
                    @test Φ′.K == K
                    @test total_value(A, Π, Φ′, f) <= total_value(A, Π, Φ, f) * (1 + ϵ)
                end
            end

            for (f,) = [
                (ConvexWorkModel(0, 0, 1),);
                (AffineConnectivityModel(0.0, 0.0, 0.0, 1.0),);
                (AffineWorkModel(0, 0, 0),);
                (ConstrainedCost(AffineConnectivityModel(0, 0, 0, 1), AffineWorkModel(0, 1, 0), 2),);
                (ConstrainedCost(AffineConnectivityModel(0, 0, 0, 1), AffineWorkModel(0, 1, 0), 4),);
                (ConstrainedCost(AffineConnectivityModel(0, 0, 0, 1), AffineWorkModel(0, 1, 0), 8),);
                (ConstrainedCost(ConvexWorkModel(0, 1, 0), AffineWorkModel(0, 1, 0), 2),);
                (ConstrainedCost(ConvexWorkModel(0, 1, 0), AffineWorkModel(0, 1, 0), 4),);
                (ConstrainedCost(ConvexWorkModel(0, 0, 1), AffineWorkModel(0, 1, 0), 8),);
            ]
                Φ = partition_stripe(A, K, ReferenceTotalSplitter(f))
                c = total_value(A, Φ, f)
                for method = [
                    ReferenceTotalSplitter(f);
                    DynamicTotalSplitter(f);
                    DynamicTotalChunker(f);
                    ConvexTotalSplitter(f);
                ]
                    Φ′ = partition_stripe(A, K, method)
                    @test issorted(Φ′.spl)
                    @test Φ′.spl[1] == 1
                    @test Φ′.spl[end] == n + 1
                    @test total_value(A, Φ′, f) ≈ c
                end
            end

            for (f,) = [
                (ConcaveWorkModel(0, 0, 1),);
                (AffineWorkModel(0, 0, 0),);
                (ConstrainedCost(ConcaveWorkModel(0, 1, 0), AffineWorkModel(0, 1, 0), 2),);
                (ConstrainedCost(ConcaveWorkModel(0, 1, 0), AffineWorkModel(0, 1, 0), 4),);
                (ConstrainedCost(ConcaveWorkModel(0, 0, 1), AffineWorkModel(0, 1, 0), 8),);
            ]
                Φ = partition_stripe(A, K, ReferenceTotalSplitter(f))
                c = total_value(A, Φ, f)
                for method = [
                    ReferenceTotalSplitter(f);
                    DynamicTotalSplitter(f);
                    DynamicTotalChunker(f);
                    ConcaveTotalSplitter(f);
                ]
                    Φ′ = partition_stripe(A, K, method)
                    @test issorted(Φ′.spl)
                    @test Φ′.spl[1] == 1
                    @test Φ′.spl[end] == n + 1
                    @test total_value(A, Φ′, f) ≈ c
                end
            end
        end

        for (f, w_max) = [
            (ConstrainedCost(AffineConnectivityModel(0, 3, 1, 3), VertexCount(), 4), 4);
            (ConstrainedCost(BlockComponentCostModel{Int64}(0, 0, (10, identity), (2, x->2x)), VertexCount(), 4), 4);
            (ConstrainedCost(BlockComponentCostModel{Int64}(identity, x->3x, (10, identity), (2, x->2x)), VertexCount(), 4), 4);
        ]
            Π = pack_stripe(A', EquiChunker(2))
            Φ = pack_stripe(A, DynamicTotalChunker(f), Π) #The dynamic total chunker should be at least as good as optimal
            c = total_value(A, Π, Φ, f)
            for method = [
                ReferenceTotalChunker(f);
                DynamicTotalChunker(f)
                StrictChunker(w_max);
                OverlapChunker(0.9, w_max);
                OverlapChunker(0.8, w_max);
                OverlapChunker(0.7, w_max);
            ]
                Φ′ = pack_stripe(A, method, Π)
                @test issorted(Φ′.spl)
                @test all(Φ′.spl[2:end] .- Φ′.spl[1:end - 1] .<= w_max)
                @test Φ′.spl[1] == 1
                @test Φ′.spl[end] == n + 1
                @test total_value(A, Π, Φ′, f) >= total_value(A, Π, Φ, f)
            end
        end

        for (f,) = [
            (ConvexWorkModel(0.0, 0, 1),);
            (ConvexWorkModel(-0.7, 0, 1),);
            (AffineConnectivityModel(-0.5, 0.0, 0.0, 1.0),);
            (AffineConnectivityModel(0, 0, 0, 1),);
            (AffineWorkModel(0, 0, 0),);
            (ConstrainedCost(AffineConnectivityModel(0, 0, 0, 1), AffineWorkModel(0, 1, 0), 2),);
            (ConstrainedCost(AffineConnectivityModel(0, 0, 0, 1), AffineWorkModel(0, 1, 0), 4),);
            (ConstrainedCost(AffineConnectivityModel(0, 0, 0, 1), AffineWorkModel(0, 1, 0), 8),);
            (ConstrainedCost(ConvexWorkModel(0, 1, 0), AffineWorkModel(0, 1, 0), 2),);
            (ConstrainedCost(ConvexWorkModel(0, 1, 0), AffineWorkModel(0, 1, 0), 4),);
            (ConstrainedCost(ConvexWorkModel(0, 0, 1), AffineWorkModel(0, 1, 0), 8),);
        ]
            Φ = pack_stripe(A, ReferenceTotalChunker(f))
            c = total_value(A, Φ, f)
            for method = [
                ReferenceTotalChunker(f);
                DynamicTotalChunker(f);
                ConvexTotalChunker(f);
            ]
                Φ′ = pack_stripe(A, method)
                @test issorted(Φ′.spl)
                @test Φ′.spl[1] == 1
                @test Φ′.spl[end] == n + 1
                @test total_value(A, Φ′, f) ≈ c
            end
        end

        for (f,) = [
            (ConcaveWorkModel(0.0, 0, 1),);
            (ConcaveWorkModel(-0.7, 0, 1),);
            (AffineWorkModel(0, 0, 0),);
            (ConstrainedCost(ConcaveWorkModel(0, 1, 0), AffineWorkModel(0, 1, 0), 2),);
            (ConstrainedCost(ConcaveWorkModel(0, 1, 0), AffineWorkModel(0, 1, 0), 4),);
            (ConstrainedCost(ConcaveWorkModel(0, 0, 1), AffineWorkModel(0, 1, 0), 8),);
        ]
            Φ = pack_stripe(A, ReferenceTotalChunker(f))
            c = total_value(A, Φ, f)
            for method = [
                ReferenceTotalChunker(f);
                DynamicTotalChunker(f);
                ConcaveTotalChunker(f);
            ]
                Φ′ = pack_stripe(A, method)
                @test issorted(Φ′.spl)
                @test Φ′.spl[1] == 1
                @test Φ′.spl[end] == n + 1
                @test total_value(A, Φ′, f) ≈ c
            end
        end

    end
end


text:
const Bulma = Style{:bulma}
class(::Bulma) = "content"

# Block elements

function Markdown.html(io::IO, md::Styled{Markdown.Admonition,Bulma})
    Markdown.withtag(io, :article, :class => "message is-$(md.category)") do
        Markdown.withtag(io, :div, :class => "message-header") do
            Markdown.withtag(io, :p) do
                print(io, md.title)
            end
        end
        Markdown.withtag(io, :div, :class => "message-body") do
            html(io, style(md)(md.content))
        end
    end
end

# Inline elements

function Markdown.htmlinline(io::IO, md::Styled{Markdown.Image,Bulma})
    Markdown.withtag(io, :figure, :class => "image") do
        Markdown.tag(io, :img, :src => md.url, :alt => md.alt)
    end
end

# Extra elements

function Markdown.html(io::IO, md::Styled{Markdown.Table,Bulma})
    Markdown.withtag(io, :table, :class => "table") do
        Markdown.withtag(io, :thead) do
            htmltablerow(io, md, 1)
        end
        Markdown.withtag(io, :tbody) do
            for i ∈ 2:length(md.rows)
                htmltablerow(io, md, i)
            end
        end
    end
end

ext: jl


text:
@testset "Non_negative" begin

    p, p_inv = Non_negative()

    x = abs.(randn((10, 10)))
    x2 = 100 .*randn((10, 10))

    @test x ≈ p(p_inv(x))
    @test x ≈ p_inv(p(x))

    @test all(p(x2) .>= 0)
end

@testset "Map_0_1" begin

    p, p_inv = Map_0_1()

    x = abs.(randn((10, 10)))
    x2 = 100 .*randn((10, 10))

    @test x ≈ p_inv(p(x))

    @test all(p(x2) .>= 0)
    @test all(p(x2) .<= 1)
end


text:
"""
	voronoiarea(tess::Tessellation) -> Vector

Compute the area of all Voronoi cells in a tesselation.
"""
function voronoiarea(tess::Tessellation)
	map(polyarea, tess.Cells)
end

"""
	polyarea(p::Vector)

Compute the area of the polygon with vertices `p` using the shoelace formula.
It is assumed that `p` is sorted.
"""
function polyarea(pts)
	n_pts = length(pts)
	A = getx(pts[1]) * (gety(pts[2]) - gety(pts[n_pts])) + getx(pts[n_pts]) * (gety(pts[1]) - gety(pts[n_pts - 1]))

	for n in 2:n_pts - 1
		A += getx(pts[n])*(gety(pts[n + 1]) - gety(pts[n - 1]))
	end

	return 0.5 * abs(A)
end


text:
# Model definition
N = latin_generator.N
p = latin_generator.p
approximator_GNN = GeometricFlux.GraphConv(32 => 32, Flux.leakyrelu)
target_approximator_GNN = GeometricFlux.GraphConv(32 => 32, Flux.leakyrelu)
gnnlayers = 2

approximator_model = SeaPearl.CPNN(
    graphChain = Flux.Chain(
        GeometricFlux.GraphConv(numInFeatures => 32, Flux.leakyrelu),
        [approximator_GNN for i = 1:gnnlayers]...
    ),
    nodeChain = Flux.Chain(
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, 16, Flux.leakyrelu),
    ),
    globalChain = Flux.Chain(
        Flux.Dense(numGlobalFeature, 32, Flux.leakyrelu),
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, 16, Flux.leakyrelu),
    ),
    outputChain = Flux.Chain(
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, N),
    )
)
 #|> gpu
target_approximator_model = SeaPearl.CPNN(
    graphChain = Flux.Chain(
        GeometricFlux.GraphConv(numInFeatures => 32, Flux.leakyrelu),
        [approximator_GNN for i = 1:gnnlayers]...
    ),
    nodeChain = Flux.Chain(
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, 16, Flux.leakyrelu),
    ),
    globalChain = Flux.Chain(
        Flux.Dense(numGlobalFeature, 32, Flux.leakyrelu),
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, 16, Flux.leakyrelu),
    ),
    outputChain = Flux.Chain(
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, N),
    )
) #|> gpu

agent = RL.Agent(
    policy = RL.QBasedPolicy(
        learner = RL.DQNLearner(
            approximator = RL.NeuralNetworkApproximator(
                model = approximator_model,
                optimizer = ADAM()
            ),
            target_approximator = RL.NeuralNetworkApproximator(
                model = target_approximator_model,
                optimizer = ADAM()
            ),
            loss_func = Flux.Losses.huber_loss,
            batch_size = 16, #32,
            update_horizon = 3, #what if the number of nodes in a episode is smaller
            min_replay_history = 10,
            update_freq = 8,
            target_update_freq = 100,
            #rng = rng,
        ),
        explorer = RL.EpsilonGreedyExplorer(
            ϵ_stable = 0.1,
            #kind = :exp,
            decay_steps = 2000,
            step = 1,
            #rng = rng
        )
    ),
    trajectory = RL.CircularArraySLARTTrajectory(
        capacity = 1000,
        state = SeaPearl.DefaultTrajectoryState[] => (),
        legal_actions_mask = Vector{Bool} => (N, ),
    )
)

max_stars_repo_name: corail-research/SeaPearlZoo

max_issues_repo_name: corail-research/SeaPearlZoo

max_forks_repo_name: corail-research/SeaPearlZoo


text:
# This script analyzes the behavior of spatial smoothing by varying the Hankel
# parameter m which determines the trade off between the number of ambient
# measurements and the number of ambient snapshots of the transformed data
# matrix which is constructed in the respective algorithms.
#
# Inputs:
# N: number of measurements
# K: number of sinusoids
# L: number of snapshots
# SNR: SNR in [dB]
# NIters: number of Monte-Carlo iterations
# m: Hankel parameter
# f: vector of the frequencies of the sinusoids

include("MUSIC.jl")
include("rootMUSIC.jl")
include("ESPRIT.jl")
include("TLSesprit.jl")
include("MPM.jl")
include("ANM.jl")
include("EMAC.jl")
include("periodogram.jl")
include("findMinWrapDist.jl")
include("findMinWrapDistUnique.jl")
include("WrapDist.jl")
using Plots, Random

N = 30
K = 3
L = 1
SNR = 30
Niters = 100
m = range(20,stop=N)
f_true = [0.3; 0.35;0.4]

MSE = zeros(length(m))

a(f) = exp.(2*pi*1im*f*(0:N-1))
A(f) = hcat(a.(f')...)

for ii = 1:Niters, jj = 1:length(m)
    noise = (randn(N,L)+1im*randn(N,L))/ (sqrt(2)*10^(SNR/10))

    X = exp.(2*pi*rand(K,L)*1im)

    Y = A(f_true)*X + noise

    f_rootMUSIC = TLSesprit(Y,K,m[jj])

    MinWrapDist = findMinWrapDistUnique(f_true,f_rootMUSIC)
    MSE[jj] += norm(MinWrapDist)^2
end

MSE /= Niters*2

plot(m, MSE,yaxis=:log,tickfontsize = 6,
label = "root MUSIC",
xlabel = "window length",
ylabel = "MSE"
)
#savefig("/home/thomas/atom-amd64/usr/share/atom/example/spatial_smoothing.png")


text: include("sparse/sensoutputsparse.jl")


text:
#### InitialConditions

"""
    init_state_vecs!

Defines initial conditions for state vectors `q` and `aux` for all sub-domains.
"""
function init_state_vecs! end

"""
    initialize_updrafts!

Defines initial conditions for state vectors `q` and `aux` for the updraft sub-domains.
"""
function initialize_updrafts! end

"""
    init_forcing!

Defines initial conditions for forcing terms in the state vector `aux` for all sub-domains.
"""
function init_forcing! end

function initialize_updrafts!(
    q::StateVec,
    aux::StateVec,
    grid::Grid,
    params,
    output_dir::String,
    ::BOMEX,
)
    gm, en, ud, sd, al = allcombinations(q)
    @unpack SurfaceModel = params
    k_1 = first_interior(grid, Zmin())
    n_updrafts = length(ud)
    for i in ud
        for k in over_elems(grid)
            q[:w, k, i] = 0
            q[:a, k, i] = bound(0.0, params[:a_bounds])
        end
        q[:a, k_1, i] = bound(SurfaceModel.area / n_updrafts, params[:a_bounds])
    end
    return
end

function init_state_vecs!(
    q::StateVec,
    aux::StateVec,
    grid::Grid,
    params,
    output_dir::String,
    case::BOMEX,
)
    @unpack a_bounds, SurfaceModel, param_set = params
    qtg = SurfaceModel.q_tot
    Tg = SurfaceModel.T
    Pg = SurfaceModel.P
    z = grid.zc

    gm, en, ud, sd, al = allcombinations(q)
    k_1 = first_interior(grid, Zmin())

    @inbounds for k in over_elems(grid)
        @inbounds for ϕ in var_names(q)
            @inbounds for i in over_sub_domains(q, ϕ)
                q[ϕ, k, i] = 0.0
            end
        end

        q[:a, k, gm] = 1.0
        for i in ud
            q[:a, k, i] = bound(0.0, a_bounds)
            k == k_1 && (q[:a, k, i] = bound(SurfaceModel.area, a_bounds))
        end
        q[:a, k, en] = q[:a, k, gm] - sum([q[:a, k, i] for i in ud])

        # Set qt profile
        if z[k] <= 520
            q[:q_tot, k, gm] = (17.0 + (z[k]) * (16.3 - 17.0) / 520.0) / 1000.0
        end
        if z[k] > 520.0 && z[k] <= 1480.0
            q[:q_tot, k, gm] =
                (16.3 + (z[k] - 520.0) * (10.7 - 16.3) / (1480.0 - 520.0)) /
                1000.0
        end
        if z[k] > 1480.0 && z[k] <= 2000.0
            q[:q_tot, k, gm] =
                (10.7 + (z[k] - 1480.0) * (4.2 - 10.7) / (2000.0 - 1480.0)) /
                1000.0
        end
        if z[k] > 2000.0
            q[:q_tot, k, gm] =
                (4.2 + (z[k] - 2000.0) * (3.0 - 4.2) / (3000.0 - 2000.0)) /
                1000.0
        end

        # Set θ_liq profile
        if z[k] <= 520.0
            q[:θ_liq, k, gm] = 298.7
        end
        if z[k] > 520.0 && z[k] <= 1480.0
            q[:θ_liq, k, gm] =
                298.7 + (z[k] - 520) * (302.4 - 298.7) / (1480.0 - 520.0)
        end
        if z[k] > 1480.0 && z[k] <= 2000
            q[:θ_liq, k, gm] =
                302.4 + (z[k] - 1480.0) * (308.2 - 302.4) / (2000.0 - 1480.0)
        end
        if z[k] > 2000.0
            q[:θ_liq, k, gm] =
                308.2 + (z[k] - 2000.0) * (311.85 - 308.2) / (3000.0 - 2000.0)
        end
    end

    @inbounds for k in over_elems_real(grid)
        ts = ActiveThermoState(param_set, q, aux, k, gm)

        aux[:T, k, gm] = air_temperature(ts)

        # Set u profile
        if z[k] <= 700.0
            q[:u, k, gm] = -8.75
        end
        if z[k] > 700.0
            q[:u, k, gm] =
                -8.75 + (z[k] - 700.0) * (-4.61 - -8.75) / (3000.0 - 700.0)
        end
    end # end over_elems

    # Extrapolate to ghost points
    extrap_0th_order!(q, (:θ_liq, :q_tot, :u), grid, gm)
    extrap_0th_order!(aux, :T, grid, gm)
    # Use grid-mean for sub-domain values:

    initialize_updrafts!(q, aux, grid, params, output_dir, case)
    distribute!(q, grid, (:q_tot, :θ_liq))
    distribute!(aux, grid, (:q_liq, :T))
    diagnose_environment!(q, grid, :a, (:q_tot, :θ_liq, :w))
end

function init_forcing!(
    q::StateVec,
    aux::StateVec,
    grid::Grid{FT},
    params,
    output_dir::String,
    case::BOMEX,
) where {FT}
    gm, en, ud, sd, al = allcombinations(q)
    @unpack param_set = params
    z = grid.zc
    for k in over_elems_real(grid)
        # Geostrophic velocity profiles. vg = 0
        aux[:ug, k, gm] = -10.0 + (1.8e-3) * z[k]
        Π = TD.exner_given_pressure(param_set, aux[:p_0, k])
        # Set large-scale cooling
        if z[k] <= 1500.0
            aux[:dTdt, k, gm] = (-2.0 / (3600 * 24.0)) * Π
        else
            aux[:dTdt, k, gm] =
                (
                    -2.0 / (3600 * 24.0) +
                    (z[k] - 1500.0) * (0.0 - -2.0 / (3600 * 24.0)) /
                    (3000.0 - 1500.0)
                ) * Π
        end
        # Set large-scale drying
        if z[k] <= 300.0
            aux[:dqtdt, k, gm] = -1.2e-8   #kg/(kg * s)
        end
        if z[k] > 300.0 && z[k] <= 500.0
            aux[:dqtdt, k, gm] =
                -1.2e-8 + (z[k] - 300.0) * (0.0 - -1.2e-8) / (500.0 - 300.0) #kg/(kg * s)
        end

        #Set large scale subsidence
        if z[k] <= 1500.0
            aux[:subsidence, k, gm] =
                0.0 + z[k] * (-0.65 / 100.0 - 0.0) / (1500.0 - 0.0)
        end
        if z[k] > 1500.0 && z[k] <= 2100.0
            aux[:subsidence, k, gm] =
                -0.65 / 100 +
                (z[k] - 1500.0) * (0.0 - -0.65 / 100.0) / (2100.0 - 1500.0)
        end
    end
end


text:
################################################################################
#                                 DATATYPES
################################################################################
# Extend addchild to take the root of another graph as input
function _LCRST.addchild(parent::_LCRST.Node{T}, newc::_LCRST.Node{T}) where T
    # copy the new node if it is not a root
    # otherwise, we are just merging 2 graphs together
    if !_LCRST.isroot(newc)
        newc = copy(newc)
    end
    # add it on to the tree
    newc.parent = parent
    prevc = parent.child
    if prevc == parent
        parent.child = newc
    else
        prevc = _LCRST.lastsibling(prevc)
        prevc.sibling = newc
    end
    return newc
end

# Extend addchild with convenient nothing dispatch for empty previous child
function _LCRST.addchild(
    parent::_LCRST.Node{T},
    oldc::Nothing,
    newc::_LCRST.Node{T}
    ) where T
    return _LCRST.addchild(parent, newc)
end

# Extend addchild to efficiently add multiple children if the previous is known
function _LCRST.addchild(
    parent::_LCRST.Node{T},
    prevc::_LCRST.Node{T},
    data::T
    ) where T
    # add it on to the tree
    newc = _LCRST.Node(data, parent)
    prevc.sibling = newc
    return newc
end

# Extend addchild to efficiently add multiple children if the previous is known
function _LCRST.addchild(
    parent::_LCRST.Node{T},
    prevc::_LCRST.Node{T},
    newc::_LCRST.Node{T}
    ) where T
    # check if the prev is actually a child of the parent
    @assert prevc.parent === parent "Previous child doesn't belong to parent."
    # copy the new node if it is not a root
    # otherwise, we are just merging 2 graphs together
    if !_LCRST.isroot(newc)
        newc = copy(newc)
    end
    # add it on to the tree
    newc.parent = parent
    prevc.sibling = newc
    return newc
end

# Map a LCRST tree based by operating each node with a function
function _map_tree(map_func::Function, node::_LCRST.Node)
    new_node = map_func(node)
    prev = nothing
    for child in node
        prev = _LCRST.addchild(new_node, prev, _map_tree(map_func, child))
    end
    return new_node
end

# Extend copying for graph nodes
function Base.copy(node::_LCRST.Node)
    return _map_tree(n -> _LCRST.Node(n.data), node)
end

# Replace a node with its only child if it only has 1 child
function _merge_parent_and_child(node::_LCRST.Node)
    if _LCRST.islastsibling(node.child)
        child = node.child
        node.data = child.data
        for n in child
            n.parent = node
        end
        node.child = child.child
        child.child = child
        child.parent = child
    end
    return node
end

# This is ambiguous but faster than the concrete alternatives tested so far
# Even better than using Node{Any}...
"""
    NodeData

A `DataType` for storing values in an expression tree that is used in a
[`NLPExpr`](@ref). Acceptable value types include:
- `Real`: Constants
- `GeneralVariableRef`: Optimization variables
- `JuMP.GenericAffExpr{Float64, GeneralVariableRef}`: Affine expressions
- `JuMP.GenericQuadExpr{Float64, GeneralVariableRef}`: Quadratic expressions
- `Symbol`: Registered NLP function name.

**Fields**
- `value`: The stored value.
"""
struct NodeData
    value
end

# Getter function for the node value (so it is easy to change later on if needed)
function _node_value(data::NodeData)
    return data.value
end

# Recursively determine if node is effectively zero
function _is_zero(node::_LCRST.Node{NodeData})
    raw = _node_value(node.data)
    if isequal(raw, 0)
        return true
    elseif _LCRST.isleaf(node)
        return false
    elseif raw in (:+, :-) && all(_is_zero(n) for n in node)
        return true
    elseif raw == :* && any(_is_zero(n) for n in node)
        return true
    elseif raw in (:/, :^) && _is_zero(node.child)
        return true
    elseif all(_is_zero(n) for n in node) && iszero(get(_NativeNLPFunctions, (raw, length(collect(node))), (i...) -> true)((0.0 for n in node)...))
        return true
    else
        return false
    end
end

# Prone any nodes that are effectively zero
function _drop_zeros!(node::_LCRST.Node{NodeData})
    if _LCRST.isleaf(node)
        return node
    elseif _is_zero(node)
        node.data = NodeData(0.0)
        _LCRST.makeleaf!(node)
        return node
    end
    raw = _node_value(node.data)
    if raw == :+
        for n in node
            if _is_zero(n)
                _LCRST.prunebranch!(n)
            end
        end
        _merge_parent_and_child(node)
    elseif raw == :-
        if _is_zero(node.child)
            _LCRST.prunebranch!(node.child)
        elseif _is_zero(node.child.sibling)
            _LCRST.prunebranch!(node.child.sibling)
            _merge_parent_and_child(node)
        end
    end
    for n in node
        _drop_zeros!(n)
    end
    return node
end

# Extend Base.isequal for our node types
function Base.isequal(n1::_LCRST.Node{NodeData}, n2::_LCRST.Node{NodeData})
    isequal(_node_value(n1.data), _node_value(n2.data)) || return false
    count(i -> true, n1) != count(i -> true, n2) && return false
    for (c1, c2) in zip(n1, n2)
        if !isequal(c1, c2)
            return false
        end
    end
    return true
end

"""
    NLPExpr <: JuMP.AbstractJuMPScalar

A `DataType` for storing scalar nonlinear expressions. It stores the expression
algebraically via an expression tree where each node contains [`NodeData`](@ref)
that can store one of the following:
- a registered function name (stored as a `Symbol`)
- a constant
- a variable
- an affine expression
- a quadratic expression.
Specifically, it employs a left-child right-sibling tree
(from `LeftChildRightSiblingTrees.jl`) to represent the expression tree.

**Fields**
- `tree_root::LeftChildRightSiblingTrees.Node{NodeData}`: The root node of the
  expression tree.
"""
struct NLPExpr <: JuMP.AbstractJuMPScalar
    tree_root::_LCRST.Node{NodeData}

    # Constructor
    function NLPExpr(tree_root::_LCRST.Node{NodeData})
        return new(tree_root)
    end
end

# Extend basic functions
Base.broadcastable(nlp::NLPExpr) = Ref(nlp)
Base.copy(nlp::NLPExpr) = NLPExpr(copy(nlp.tree_root))
Base.zero(::Type{NLPExpr}) = NLPExpr(_LCRST.Node(NodeData(0.0)))
Base.one(::Type{NLPExpr}) = NLPExpr(_LCRST.Node(NodeData(1.0)))
function Base.isequal(nlp1::NLPExpr, nlp2::NLPExpr)
    return isequal(nlp1.tree_root, nlp2.tree_root)
end

"""
    JuMP.drop_zeros!(nlp::NLPExpr)::NLPExpr

Removes the zeros (possibly introduced by deletion) from an nonlinear expression.
Note this only uses a few simple heuristics and will not remove more complex
relationships like `cos(π/2)`.

**Example**
```julia-repl
julia> expr = x^2.3 * max(0, zero(NLPExpr)) - exp(1/x + 0)
x^2.3 * max(0, 0) - exp(1 / x + 0)

julia> drop_zeros!(expr)
-exp(1 / x)
```
"""
function JuMP.drop_zeros!(nlp::NLPExpr)
    _drop_zeros!(nlp.tree_root) # uses a basic simplification scheme
    return nlp
end

# Extend JuMP.isequal_canonical (uses some heuristics but is not perfect)
function JuMP.isequal_canonical(nlp1::NLPExpr, nlp2::NLPExpr)
    n1 = _drop_zeros!(copy(nlp1.tree_root))
    n2 = _drop_zeros!(copy(nlp2.tree_root))
    return isequal(n1, n2)
end

# Print the tree structure of the expression tree
function print_expression_tree(io::IO, nlp::NLPExpr)
    return AbstractTrees.print_tree(io, nlp.tree_root)
end
print_expression_tree(io::IO, expr) = println(io, expr)

"""
    print_expression_tree(nlp::NLPExpr)

Print a tree representation of the nonlinear expression `nlp`.

**Example**
```julia-repl
julia> expr = (x * sin(x)^3) / 2
(x * sin(x)^3) / 2

julia> print_expression_tree(expr)
/
├─ *
│  ├─ x
│  └─ ^
│     ├─ sin
│     │  └─ x
│     └─ 3
└─ 2
```
"""
print_expression_tree(nlp::NLPExpr) = print_expression_tree(stdout::IO, nlp)

# Convenient expression alias
const AbstractInfOptExpr = Union{
    NLPExpr,
    JuMP.GenericQuadExpr{Float64, GeneralVariableRef},
    JuMP.GenericAffExpr{Float64, GeneralVariableRef},
    GeneralVariableRef
}

## Dispatch function for ast mapping
# Constant
function _ast_process_node(map_func::Function, c)
    return c
end

# Variable
function _ast_process_node(map_func::Function, v::GeneralVariableRef)
    return map_func(v)
end

# AffExpr
function _ast_process_node(map_func::Function, aff::JuMP.GenericAffExpr)
    ex = Expr(:call, :+)
    for (v, c) in aff.terms
        if isone(c)
            push!(ex.args, map_func(v))
        else
            push!(ex.args, Expr(:call, :*, c, map_func(v)))
        end
    end
    if !iszero(aff.constant)
        push!(ex.args, aff.constant)
    end
    return ex
end

# QuadExpr
function _ast_process_node(map_func::Function, quad::JuMP.GenericQuadExpr)
    ex = Expr(:call, :+)
    for (xy, c) in quad.terms
        if isone(c)
            push!(ex.args, Expr(:call, :*, map_func(xy.a), map_func(xy.b)))
        else
            push!(ex.args, Expr(:call, :*, c, map_func(xy.a), map_func(xy.b)))
        end
    end
    append!(ex.args, _ast_process_node(map_func, quad.aff).args[2:end])
    return ex
end

# Map an expression tree to a Julia AST tree that is compatible with JuMP
function _tree_map_to_ast(map_func::Function, node::_LCRST.Node)
    if _LCRST.isleaf(node)
        return _ast_process_node(map_func, _node_value(node.data))
    else
        ex = Expr(:call, _node_value(node.data)) # will be function symbol name
        append!(ex.args, (_tree_map_to_ast(map_func, n) for n in node))
        return ex
    end
end

"""
    map_nlp_to_ast(map_func::Function, nlp::NLPExpr)::Expr

Map the nonlinear expression `nlp` to a Julia AST expression where each variable
is mapped via `map_func` and is directly interpolated into the AST expression.
This is intended as an internal method that can be helpful for developers that
wish to map a `NLPExpr` to a Julia AST expression that is compatible with
`JuMP.add_NL_expression`.
"""
function map_nlp_to_ast(map_func::Function, nlp::NLPExpr)
    return _tree_map_to_ast(map_func, nlp.tree_root)
end

################################################################################
#                          EXPRESSION CREATION HELPERS
################################################################################
## Make convenient dispatch methods for raw child input
# NLPExpr
function _process_child_input(nlp::NLPExpr)
    return nlp.tree_root
end

# An InfiniteOpt expression (not general nonlinear)
function _process_child_input(v::AbstractInfOptExpr)
    return NodeData(v)
end

# Function symbol
function _process_child_input(f::Symbol)
    return NodeData(f)
end

# A constant
function _process_child_input(c::Union{Real, Bool})
    return NodeData(c)
end

# Fallback
function _process_child_input(v)
    error("Unrecognized algebraic expression input `$v`.")
end

# Generic graph builder
function _call_graph(func::Symbol, arg1, args...)
    root = _LCRST.Node(NodeData(func))
    prevc = _LCRST.addchild(root, _process_child_input(arg1))
    for a in args
        prevc = _LCRST.addchild(root, prevc, _process_child_input(a))
    end
    return root
end

################################################################################
#                               SUMS AND PRODUCTS
################################################################################
## Define helper functions for sum reductions
# Container of NLPExprs
function _reduce_by_first(::typeof(sum), first_itr::NLPExpr, itr)
    root = _LCRST.Node(NodeData(:+))
    prevc = _LCRST.addchild(root, first_itr.tree_root)
    for ex in itr
        prevc = _LCRST.addchild(root, prevc, _process_child_input(ex))
    end
    return NLPExpr(root)
end

# Container of InfiniteOpt exprs
function _reduce_by_first(
    ::typeof(sum),
    first_itr::JuMP.AbstractJuMPScalar,
    itr
    )
    result = first_itr
    for i in itr
        result = _MA.operate!(_MA.add_mul, result, i)
    end
    return result
end

# Fallback
function _reduce_by_first(::typeof(sum), first_itr, itr; kw...)
    return isempty(itr) ? first_itr : first_itr + sum(identity, itr; kw...)
end

# Hyjack Base.sum for better efficiency on iterators --> this is type piracy...
function Base.sum(itr::Base.Generator; kw...)
    isempty(itr) && return sum(identity, itr; kw...)
    itr1, new_itr = Iterators.peel(itr)
    return _reduce_by_first(sum, itr1, new_itr; kw...)
end

# Extend Base.sum for container of NLPExprs
function Base.sum(arr::AbstractArray{<:NLPExpr}; init = zero(NLPExpr))
    isempty(arr) && return init
    itr1, new_itr = Iterators.peel(arr)
    return _reduce_by_first(sum, itr1, new_itr)
end

# Extend Base.sum for container of InfiniteOpt exprs
function Base.sum(
    arr::AbstractArray{<:AbstractInfOptExpr};
    init = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
    )
    isempty(arr) && return init
    result = _MA.Zero()
    for i in arr
        result = _MA.operate!(_MA.add_mul, result, i)
    end
    return result
end

## Define helper functions for reducing products
# Container of InfiniteOpt exprs
function _reduce_by_first(::typeof(prod), first_itr::AbstractInfOptExpr, itr)
    root = _LCRST.Node(NodeData(:*))
    prevc = _LCRST.addchild(root, _process_child_input(first_itr))
    for ex in itr
        prevc = _LCRST.addchild(root, prevc, _process_child_input(ex))
    end
    return NLPExpr(root)
end

# Fallback
function _reduce_by_first(::typeof(prod), first_itr, itr; kw...)
    return first_itr * prod(identity, itr; kw...)
end

# Hyjack Base.prod for better efficiency on iterators --> this is type piracy...
function Base.prod(itr::Base.Generator; kw...)
    isempty(itr) && return prod(identity, itr; kw...)
    itr1, new_itr = Iterators.peel(itr)
    return _reduce_by_first(prod, itr1, new_itr; kw...)
end

# Extend Base.prod for container of InfiniteOpt exprs
function Base.prod(arr::AbstractArray{<:AbstractInfOptExpr}; init = one(NLPExpr))
    isempty(arr) && return init
    itr1, new_itr = Iterators.peel(arr)
    return _reduce_by_first(prod, itr1, new_itr)
end

################################################################################
#                           MULTIPLICATION OPERATORS
################################################################################
# TODO more intelligently operate with constants

# QuadExpr * expr
function Base.:*(
    quad::JuMP.GenericQuadExpr{Float64, GeneralVariableRef},
    expr::AbstractInfOptExpr
    )
    return NLPExpr(_call_graph(:*, quad, expr))
end

# expr * QuadExpr
function Base.:*(
    expr::AbstractInfOptExpr,
    quad::JuMP.GenericQuadExpr{Float64, GeneralVariableRef}
    )
    return NLPExpr(_call_graph(:*, expr, quad))
end

# QuadExpr * QuadExpr
function Base.:*(
    quad1::JuMP.GenericQuadExpr{Float64, GeneralVariableRef},
    quad2::JuMP.GenericQuadExpr{Float64, GeneralVariableRef}
    )
    return NLPExpr(_call_graph(:*, quad1, quad2))
end

# NLPExpr * QuadExpr
function Base.:*(
    nlp::NLPExpr,
    quad::JuMP.GenericQuadExpr{Float64, GeneralVariableRef}
    )
    return NLPExpr(_call_graph(:*, nlp, quad))
end

# QuadExpr * NLPExpr
function Base.:*(
    quad::JuMP.GenericQuadExpr{Float64, GeneralVariableRef},
    nlp::NLPExpr
    )
    return NLPExpr(_call_graph(:*, quad, nlp))
end

# NLPExpr * expr/constant
function Base.:*(nlp::NLPExpr, expr::Union{AbstractInfOptExpr, Real})
    return NLPExpr(_call_graph(:*, nlp, expr))
end

# expr/constant * NLPExpr
function Base.:*(expr::Union{AbstractInfOptExpr, Real}, nlp::NLPExpr)
    return NLPExpr(_call_graph(:*, expr, nlp))
end

# NLPExpr * NLPExpr
function Base.:*(nlp1::NLPExpr, nlp2::NLPExpr)
    return NLPExpr(_call_graph(:*, nlp1, nlp2))
end

# expr * expr * expr ...
function Base.:*(
    expr1::AbstractInfOptExpr,
    expr2::AbstractInfOptExpr,
    expr3::AbstractInfOptExpr,
    exprs::Vararg{AbstractInfOptExpr}
    )
    return NLPExpr(_call_graph(:*, expr1, expr2, expr3, exprs...))
end

# *NLPExpr
function Base.:*(nlp::NLPExpr)
    return nlp
end

################################################################################
#                              DIVISION OPERATORS
################################################################################
# expr/constant / expr
function Base.:/(
    expr1::Union{AbstractInfOptExpr, Real},
    expr2::AbstractInfOptExpr
    )
    return NLPExpr(_call_graph(:/, expr1, expr2))
end

# NLPExpr / constant
function Base.:/(nlp::NLPExpr, c::Real)
    if iszero(c)
        error("Cannot divide by zero.")
    elseif isone(c)
        return nlp
    else
        return NLPExpr(_call_graph(:/, nlp, c))
    end
end

################################################################################
#                                POWER OPERATORS
################################################################################
# expr ^ Integer
function Base.:^(expr::AbstractInfOptExpr, c::Integer)
    if iszero(c)
        return one(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
    elseif isone(c)
        return expr
    elseif c == 2
        return expr * expr
    else
        return NLPExpr(_call_graph(:^, expr, c))
    end
end

# expr ^ Real
function Base.:^(expr::AbstractInfOptExpr, c::Real)
    if iszero(c)
        return one(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
    elseif isone(c)
        return expr
    elseif c == 2
        return expr * expr
    else
        return NLPExpr(_call_graph(:^, expr, c))
    end
end

# NLPExpr ^ Integer
function Base.:^(expr::NLPExpr, c::Integer)
    if iszero(c)
        return one(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
    elseif isone(c)
        return expr
    else
        return NLPExpr(_call_graph(:^, expr, c))
    end
end

# NLPExpr ^ Real
function Base.:^(expr::NLPExpr, c::Real)
    if iszero(c)
        return one(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
    elseif isone(c)
        return expr
    else
        return NLPExpr(_call_graph(:^, expr, c))
    end
end

# expr/constant ^ expr
function Base.:^(
    expr1::Union{AbstractInfOptExpr, Real},
    expr2::AbstractInfOptExpr
    )
    return NLPExpr(_call_graph(:^, expr1, expr2))
end

################################################################################
#                             SUBTRACTION OPERATORS
################################################################################
# TODO more intelligently operate with constants

# NLPExpr - expr/constant
function Base.:-(nlp::NLPExpr, expr::Union{AbstractInfOptExpr, Real})
    return NLPExpr(_call_graph(:-, nlp, expr))
end

# expr/constant - NLPExpr
function Base.:-(expr::Union{AbstractInfOptExpr, Real}, nlp::NLPExpr)
    return NLPExpr(_call_graph(:-, expr, nlp))
end

# NLPExpr - NLPExpr
function Base.:-(nlp1::NLPExpr, nlp2::NLPExpr)
    return NLPExpr(_call_graph(:-, nlp1, nlp2))
end

# -NLPExpr
function Base.:-(nlp::NLPExpr)
    return NLPExpr(_call_graph(:-, nlp))
end

# Var - Var (to avoid using v == v)
function Base.:-(lhs::V, rhs::V) where {V<:GeneralVariableRef}
    if isequal(lhs, rhs)
        return zero(JuMP.GenericAffExpr{Float64,V})
    else
        return JuMP.GenericAffExpr(0.0,
                DataStructures.OrderedDict(lhs => 1.0, rhs => -1.0))
    end
end

################################################################################
#                              ADDITION OPERATORS
################################################################################
# TODO more intelligently operate with constants

# NLPExpr + expr/constant
function Base.:+(nlp::NLPExpr, expr::Union{AbstractInfOptExpr, Real})
    return NLPExpr(_call_graph(:+, nlp, expr))
end

# expr/constant + NLPExpr
function Base.:+(expr::Union{AbstractInfOptExpr, Real}, nlp::NLPExpr)
    return NLPExpr(_call_graph(:+, expr, nlp))
end

# NLPExpr + NLPExpr
function Base.:+(nlp1::NLPExpr, nlp2::NLPExpr)
    return NLPExpr(_call_graph(:+, nlp1, nlp2))
end

# +NLPExpr
function Base.:+(nlp::NLPExpr)
    return nlp
end

################################################################################
#                               MUTABLE ARITHMETICS
################################################################################
# Define NLPExpr as a mutable type for MA
_MA.mutability(::Type{NLPExpr}) = _MA.IsMutable()

# Extend MA.promote_operation for bettered efficiency
for type in (:Real, :GeneralVariableRef,
             :(JuMP.GenericAffExpr{Float64, GeneralVariableRef}),
             :(JuMP.GenericQuadExpr{Float64, GeneralVariableRef}))
    @eval begin
        function _MA.promote_operation(
            ::Union{typeof(+),typeof(-),typeof(*),typeof(/),typeof(^)},
            ::Type{<:$type},
            ::Type{NLPExpr}
            )
            return NLPExpr
        end
        function _MA.promote_operation(
            ::Union{typeof(+),typeof(-),typeof(*),typeof(/),typeof(^)},
            ::Type{NLPExpr},
            ::Type{<:$type}
            )
            return NLPExpr
        end
    end
end
function _MA.promote_operation(
    ::Union{typeof(+),typeof(-),typeof(*),typeof(/),typeof(^)},
    ::Type{NLPExpr},
    ::Type{NLPExpr}
    )
    return NLPExpr
end
for type in (:GeneralVariableRef,
             :(JuMP.GenericAffExpr{Float64, GeneralVariableRef}))
    @eval begin
        function _MA.promote_operation(
            ::Union{typeof(*),typeof(/),typeof(^)},
            ::Type{<:$type},
            ::Type{JuMP.GenericQuadExpr{Float64, GeneralVariableRef}}
            )
            return NLPExpr
        end
        function _MA.promote_operation(
            ::Union{typeof(*),typeof(/),typeof(^)},
            ::Type{JuMP.GenericQuadExpr{Float64, GeneralVariableRef}},
            ::Type{<:$type}
            )
            return NLPExpr
        end
    end
end
function _MA.promote_operation(
    ::Union{typeof(*),typeof(/),typeof(^)},
    ::Type{<:JuMP.GenericQuadExpr{Float64, GeneralVariableRef}},
    ::Type{<:JuMP.GenericQuadExpr{Float64, GeneralVariableRef}}
    )
    return NLPExpr
end
for type in (:GeneralVariableRef,
             :(JuMP.GenericAffExpr{Float64, GeneralVariableRef}),
             :(JuMP.GenericQuadExpr{Float64, GeneralVariableRef}))
    @eval begin
        function _MA.promote_operation(
            ::Union{typeof(/),typeof(^)},
            ::Type{<:Real},
            ::Type{<:$type}
            )
            return NLPExpr
        end
    end
end
for type in (:GeneralVariableRef,
             :(JuMP.GenericAffExpr{Float64, GeneralVariableRef}))
    @eval begin
        function _MA.promote_operation(
            ::Union{typeof(/),typeof(^)},
            ::Type{GeneralVariableRef},
            ::Type{<:$type}
            )
            return NLPExpr
        end
    end
end
for type in (:GeneralVariableRef,
             :(JuMP.GenericAffExpr{Float64, GeneralVariableRef}))
    @eval begin
        function _MA.promote_operation(
            ::Union{typeof(/),typeof(^)},
            ::Type{JuMP.GenericAffExpr{Float64, GeneralVariableRef}},
            ::Type{<:$type}
            )
            return NLPExpr
        end
    end
end

# Extend MA.scaling in case an NLPExpr needs to be converted to a number
function _MA.scaling(nlp::NLPExpr)
    c = _node_value(nlp.tree_root.data)
    if !(c isa Real)
        error("Cannot convert `$nlp` to `$Float64`.")
    end
    return _MA.scaling(c)
end

# Extend MA.mutable_Copy to avoid unnecessary copying
function _MA.mutable_copy(nlp::NLPExpr)
    return nlp # we don't need to copy since we build from the leaves up
end

# Extend MA.mutable_operate! as required
function _MA.mutable_operate!(
    op::Union{typeof(zero), typeof(one)},
    ::NLPExpr
    )
    return op(NLPExpr) # not actually mutable for safety and efficiency
end
function _MA.mutable_operate!(
    op::Union{typeof(+), typeof(-), typeof(*), typeof(/), typeof(^)},
    nlp::NLPExpr,
    v
    )
    return op(nlp, v)
end
function _MA.mutable_operate!(
    op::Union{typeof(+), typeof(-), typeof(*), typeof(/), typeof(^)},
    v,
    nlp::NLPExpr
    )
    return op(v, nlp)
end
function _MA.mutable_operate!(
    op::typeof(+),
    v::Union{JuMP.GenericAffExpr{Float64, GeneralVariableRef},
             JuMP.GenericQuadExpr{Float64, GeneralVariableRef}},
    nlp::NLPExpr
    )
    return op(v, nlp)
end
function _MA.mutable_operate!(
    op::typeof(-),
    v::Union{JuMP.GenericAffExpr{Float64, GeneralVariableRef},
             JuMP.GenericQuadExpr{Float64, GeneralVariableRef}},
    nlp::NLPExpr
    )
    return op(v, nlp)
end
function _MA.mutable_operate!(
    op::Union{typeof(+), typeof(-), typeof(*), typeof(/), typeof(^)},
    nlp1::NLPExpr,
    nlp2::NLPExpr
    )
    return op(nlp1, nlp2)
end
function _MA.mutable_operate!(op::_MA.AddSubMul, nlp::NLPExpr, args...)
    return _MA.add_sub_op(op)(nlp, *(args...))
end

# TODO maybe extend _MA.add_mul/_MA_.sub_mul as well

################################################################################
#                             NATIVE NLP FUNCTIONS
################################################################################
# Store all of the native registered functions
const _NativeNLPFunctions = Dict{Tuple{Symbol, Int}, Function}(
    (:-, 2) => -,
    (:/, 2) => /,
    (:^, 2) => ^
)

# List of 1 argument base functions to register
const _Base1ArgFuncList = (
    :sqrt => sqrt,
    :cbrt => cbrt,
    :abs => abs,
    :abs2 => abs2,
    :inv => inv,
    :log => log,
    :log10 => log10,
    :log2 => log2,
    :log1p => log1p,
    :exp => exp,
    :exp2 => exp2,
    :expm1 => expm1,
    :sin => sin,
    :cos => cos,
    :tan => tan,
    :sec => sec,
    :csc => csc,
    :cot => cot,
    :sind => sind,
    :cosd => cosd,
    :tand => tand,
    :secd => secd,
    :cscd => cscd,
    :cotd => cotd,
    :asin => asin,
    :acos => acos,
    :atan => atan,
    :asec => asec,
    :acsc => acsc,
    :acot => acot,
    :asind => asind,
    :acosd => acosd,
    :atand => atand,
    :asecd => asecd,
    :acscd => acscd,
    :acotd => acotd,
    :sinh => sinh,
    :cosh => cosh,
    :tanh => tanh,
    :sech => sech,
    :csch => csch,
    :coth => coth,
    :asinh => asinh,
    :acosh => acosh,
    :atanh => atanh,
    :asech => asech,
    :acsch => acsch,
    :acoth => acoth,
    :deg2rad => deg2rad,
    :rad2deg => rad2deg
)

# Setup the base 1 argument functions
for (name, func) in _Base1ArgFuncList
    # add it to the main storage dict
    _NativeNLPFunctions[(name, 1)] = func
    # make an expression constructor
    @eval begin
        function Base.$name(v::AbstractInfOptExpr)
            return NLPExpr(_call_graph($(quot(name)), v))
        end
    end
end

# Setup the Base functions with 2 arguments
for (name, func) in (:min => min, :max => max)
    # add it to the main storage dict
    _NativeNLPFunctions[(name, 2)] = func
    # make an expression constructor
    @eval begin
        function Base.$name(
            v1::Union{AbstractInfOptExpr, Real},
            v2::Union{AbstractInfOptExpr, Real}
            )
            return NLPExpr(_call_graph($(quot(name)), v1, v2))
        end
    end
end

# Setup the ifelse function
_NativeNLPFunctions[(:ifelse, 3)] = Core.ifelse

"""
    InfiniteOpt.ifelse(cond::NLPExpr, v1::Union{AbstractInfOptExpr, Real},
                       v2::Union{AbstractInfOptExpr, Real})::NLPExpr

A symbolic version of `Core.ifelse` that can be used to establish symbolic
expressions with logic conditions. Note that is must be written
`InfiniteOpt.ifelse` since it conflicts with `Core.ifelse`.

**Example**
```julia
julia> InfiniteOpt.ifelse(x >= y, 0, y^3)
ifelse(x >= y, 0, y^3)
```
"""
function ifelse(
    cond::NLPExpr,
    v1::Union{AbstractInfOptExpr, Real},
    v2::Union{AbstractInfOptExpr, Real}
    )
    return NLPExpr(_call_graph(:ifelse, cond, v1, v2))
end
function ifelse(
    cond::Bool,
    v1::Union{AbstractInfOptExpr, Real},
    v2::Union{AbstractInfOptExpr, Real}
    )
    return cond ? v1 : v2
end

# Setup the Base comparison functions
for (name, func) in (:< => Base.:(<), :(==) => Base.:(==), :> => Base.:(>),
                     :<= => Base.:(<=), :>= => Base.:(>=))
    # add it to the main storage dict
    _NativeNLPFunctions[(name, 2)] = func
    # make an expression constructor
    @eval begin
        function Base.$name(v::AbstractInfOptExpr, c::Real)
            return NLPExpr(_call_graph($(quot(name)), v, c))
        end
        function Base.$name(c::Real, v::AbstractInfOptExpr)
            return NLPExpr(_call_graph($(quot(name)), c, v))
        end
        function Base.$name(v1::AbstractInfOptExpr, v2::AbstractInfOptExpr)
            return NLPExpr(_call_graph($(quot(name)), v1, v2))
        end
        if $(quot(name)) in (:<, :>)
            function Base.$name(v1::GeneralVariableRef, v2::GeneralVariableRef)
                if isequal(v1, v2)
                    return false
                else
                    return NLPExpr(_call_graph($(quot(name)), v1, v2))
                end
            end
        else
            function Base.$name(v1::GeneralVariableRef, v2::GeneralVariableRef)
                if isequal(v1, v2)
                    return true
                else
                    return NLPExpr(_call_graph($(quot(name)), v1, v2))
                end
            end
        end
    end
end

# Setup the Base logical functions (we cannot extend && and || directly)
_NativeNLPFunctions[(:&&, 2)] = Base.:&
_NativeNLPFunctions[(:||, 2)] = Base.:|

# Logical And
function Base.:&(v::Union{GeneralVariableRef, NLPExpr}, c::Bool)
    return c ? v : false
end
function Base.:&(c::Bool, v::Union{GeneralVariableRef, NLPExpr})
    return c ? v : false
end
function Base.:&(
    v1::Union{GeneralVariableRef, NLPExpr},
    v2::Union{GeneralVariableRef, NLPExpr}
    )
    return NLPExpr(_call_graph(:&&, v1, v2))
end

# Logical Or
function Base.:|(v::Union{GeneralVariableRef, NLPExpr}, c::Bool)
    return c ? true : v
end
function Base.:|(c::Bool, v::Union{GeneralVariableRef, NLPExpr})
    return c ? true : v
end
function Base.:|(
    v1::Union{GeneralVariableRef, NLPExpr},
    v2::Union{GeneralVariableRef, NLPExpr})
    return NLPExpr(_call_graph(:||, v1, v2)
    )
end

const _Special1ArgFuncList = (
    :erf => SpecialFunctions.erf,
    :erfinv => SpecialFunctions.erfinv,
    :erfc => SpecialFunctions.erfc,
    :erfcinv => SpecialFunctions.erfcinv,
    :erfi => SpecialFunctions.erfi,
    :gamma => SpecialFunctions.gamma,
    :lgamma => SpecialFunctions.lgamma,
    :digamma => SpecialFunctions.digamma,
    :invdigamma => SpecialFunctions.invdigamma,
    :trigamma => SpecialFunctions.trigamma,
    :airyai => SpecialFunctions.airyai,
    :airybi => SpecialFunctions.airybi,
    :airyaiprime => SpecialFunctions.airyaiprime,
    :airybiprime => SpecialFunctions.airybiprime,
    :besselj0 => SpecialFunctions.besselj0,
    :besselj1 => SpecialFunctions.besselj1,
    :bessely0 => SpecialFunctions.bessely0,
    :bessely1 => SpecialFunctions.bessely1,
    :erfcx => SpecialFunctions.erfcx,
    :dawson => SpecialFunctions.dawson
)

# Setup the SpecialFunctions 1 argument functions
for (name, func) in _Special1ArgFuncList
    # add it to the main storage dict
    _NativeNLPFunctions[(name, 1)] = func
    # make an expression constructor
    @eval begin
        function SpecialFunctions.$name(v::AbstractInfOptExpr)
            return NLPExpr(_call_graph($(quot(name)), v))
        end
    end
end

################################################################################
#                               USER FUNCTIONS
################################################################################
"""
    RegisteredFunction{F <: Function, G <: Union{Function, Nothing},
                       H <: Union{Function, Nothing}}

A type for storing used defined registered functions and their information that
is needed by JuMP for build an `NLPEvaluator`. The constructor is of the form:
```julia
    RegisteredFunction(name::Symbol, num_args::Int, func::Function,
                       [gradient::Function, hessian::Function])
```

**Fields**
- `name::Symbol`: The name of the function that is used in `NLPExpr`s.
- `num_args::Int`: The number of function arguments.
- `func::F`: The function itself.
- `gradient::G`: The gradient function if one is given.
- `hessian::H`: The hessian function if one is given.
"""
struct RegisteredFunction{F <: Function, G, H}
    name::Symbol
    num_args::Int
    func::F
    gradient::G
    hessian::H

    # Constructors
    function RegisteredFunction(
        name::Symbol,
        num_args::Int,
        func::F
        ) where {F <: Function}
        return new{F, Nothing, Nothing}(name, num_args, func, nothing, nothing)
    end
    function RegisteredFunction(
        name::Symbol,
        num_args::Int,
        func::F,
        gradient::G
        ) where {F <: Function, G <: Function}
        if isone(num_args) && !hasmethod(gradient, Tuple{Real})
            error("Invalid gradient function form, see the docs for details.")
        elseif !isone(num_args) && !hasmethod(gradient, Tuple{AbstractVector{Real}, ntuple(_->Real, num_args)...})
            error("Invalid multi-variate gradient function form, see the docs for details.")
        end
        return new{F, G, Nothing}(name, num_args, func, gradient, nothing)
    end
    function RegisteredFunction(
        name::Symbol,
        num_args::Int,
        func::F,
        gradient::G,
        hessian::H
        ) where {F <: Function, G <: Function, H <: Function}
        if isone(num_args) && !hasmethod(gradient, Tuple{Real})
            error("Invalid gradient function form, see the docs for details.")
        elseif isone(num_args) && !hasmethod(hessian, Tuple{Real})
            error("Invalid hessian function form, see the docs for details.")
        end
        return new{F, G, H}(name, num_args, func, gradient, hessian)
    end
end

# Helper function for @register
function _register(
    _error::Function,
    call_mod::Module,
    model::InfiniteModel,
    name::Symbol,
    num_args::Int,
    funcs...
    )
    if !all(f -> f isa Function, funcs)
        _error("Gradient and/or hessian must be functions.")
    elseif haskey(_NativeNLPFunctions, (name, num_args)) ||
           haskey(model.func_lookup, (name, num_args))
        _error("A function with name `$name` and $num_args arguments is already " *
               "registered. Please use a function with a different name.")
    elseif !hasmethod(funcs[1], NTuple{num_args, Real})
        _error("The function `$name` is not defined for arguments of type `Real`.")
    elseif length(unique!([m.module for m in methods(funcs[1])])) > 1 ||
           first(methods(funcs[1])).module !== call_mod
        _error("Cannot register function names that are used by packages. Try " *
               "wrapping `$(funcs[1])` in a user-defined function.")
    end
    push!(model.registrations, RegisteredFunction(name, num_args, funcs...))
    model.func_lookup[name, num_args] = funcs[1]
    return
end

# Helper function to check the inputs of created functions
function _check_function_args(model::InfiniteModel, f_name, args...)
    for a in args
        m = _model_from_expr(a)
        if m !== nothing && m !== model
            error("`$f_name` is a registered function in a different model than " *
                  "`$a` belongs to. Try registering `$f_name` to the current " *
                  "model.")
        end
    end
    return
end

"""
    @register(model::InfiniteModel, func_expr, [gradient::Function], [hessian::Function])

Register a user-defined function in accordance with `func_expr` such that it can
be used in `NLPExpr`s that are used with `model` without being traced.

**Argument Information**
Here `func_expr` is of the form: `myfunc(a, b)` where `myfunc` is the function
name and the number of arguments are given symbolically. Note that the choice
of argument symbols is arbitrary. Each function argument must support anything
of type `Real` to specified.

Here we can also specify a gradient function `gradient` which for 1 argument
functions must taken in the 1 argument and return its derivative. For
multi-argument functions the gradient function must be of the form:
```julia
function gradient(g::AbstractVector{T}, args::T...) where {T <: Real}
    # fill g vector with the gradient of the function
end
```

For 1 argument functions we can also specify a hessian function with takes that
argument and return the 2nd derivative. Hessians can ge specified for
multi-argument functions, but `JuMP` backends do not currently support this.

If no gradient and/or hessian is given, the automatic differentation capabilities
of the backend (e.g., `JuMP`) will be used to determine them. Note that the
`JuMP` backend does not use Hessian's for user-defined multi-argument functions.

**Notes**
- When possible, tracing is preferred over registering a function (see
  [Function Tracing](@ref) for more info).
- Only user-defined functions can be specified. If the function is used by a
  package then it can not be used directly. However, we can readily wrap it in a
  new function `newfunc(a) = pkgfunc(a)`.
- We can only register functions in the same scope that they are defined in.
- Registered functions can only be used in or below the scope in which they are
  registered. For instance, if we register some function inside of another
  function then we can only use it inside that function (not outside of it).
- A function with a given name and number of arguments can only be registered
  once in a particular model.

**Examples**
```julia-repl
julia> @variable(model, x)
x

julia> f(a) = a^3;

julia> f(x) # user-function gets traced
x^3

julia> @register(model, f(a)) # register function
f (generic function with 2 methods)

julia> f(x) # function is no longer traced and autodifferentiation will be used
f(x)

julia> f2(a) = a^2; g2(a) = 2 * a; h2(a) = 2;

julia> @register(model, f2(a), g2, h2) # register with explicit gradient and hessian
f2 (generic function with 2 methods)

julia> f2(x)
f2(x)

julia> f3(a, b) = a * b^2;

julia> function g3(v, a, b)
          v[1] = b^2
          v[2] = 2 * a * b
          return
       end;

julia> @register(model, f3(a, b), g3) # register multi-argument function
f3 (generic function with 4 methods)

julia> f3(42, x)
f3(42, x)
```
"""
macro register(model, f, args...)
     # define error message function
     _error(str...) = _macro_error(:register, (f, args...), __source__, str...)

    # parse the arguments and check
    pos_args, extra_kwargs, _, _ = _extract_kwargs(args)
    if !isempty(extra_kwargs)
        _error("Keyword arguments were given, but none are accepted.")
    elseif length(pos_args) > 2
        _error("Too many position arguments given, should be of form " *
               "`@register(myfunc(a), [gradient], [hessian])` where " *
               "`gradient` and `hessian` are optional arguments.")
    end

    # process the function input
    if isexpr(f, :call) && all(a -> a isa Symbol, f.args)
        f_name = f.args[1]
        f_args = f.args[2:end]
        num_args = length(f_args)
    else
        _error("Unexpected function format, should be of form `myfunc(a, b)`.")
    end

    # start creating the register code and register
    code = Expr(:block)
    push!(code.args, quote
        $model isa InfiniteModel || $_error("Expected an `InfiniteModel`.")
    end)
    calling_mod = __module__ # get the module the macro is being called from
    push!(code.args, quote
        InfiniteOpt._register($_error, $calling_mod, $model, $(quot(f_name)),
                              $num_args, $(f_name), $(args...))
    end)

    # define the function overloads needed to create expressions
    Ts = [Real, AbstractInfOptExpr]
    type_combos = vec(collect(Iterators.product(ntuple(_->Ts, num_args)...)))
    filter!(ts -> !all(T -> T == Real, ts), type_combos) # remove combo with only Reals
    annotype(name, T) = :($name :: $T)
    set_args(xs, vs) = (xs = map(annotype, xs, vs); xs)
    for ts in type_combos
        push!(code.args, quote
            function $(f_name)($(set_args(f_args, ts)...))
                InfiniteOpt._check_function_args($model, $(quot(f_name)), $(f_args...))
                return NLPExpr(InfiniteOpt._call_graph($(quot(f_name)), $(f_args...)))
            end
        end)
    end

    # return the code
    return esc(code)
end

"""
    name_to_function(model::InfiniteModel, name::Symbol, num_args::Int)::Union{Function, Nothing}

Return the registered function that corresponds to `name` with `num_args`.
Returns `nothing` if no such registered function exists. This helps retrieve the
functions of function names stored in `NLPExpr`s.
"""
function name_to_function(model::InfiniteModel, name::Symbol, num_args::Int)
    if name == :+
        return +
    elseif name == :*
        return *
    else
        return get(_NativeNLPFunctions, (name, num_args),
                   get(model.func_lookup, (name, num_args), nothing))
    end
end

"""
    all_registered_functions(model::InfiniteModel)::Vector{Function}

Retrieve all the functions that are currently registered to `model`.
"""
function all_registered_functions(model::InfiniteModel)
    funcs = append!(collect(values(_NativeNLPFunctions)), (+, *))
    return append!(funcs, values(model.func_lookup))
end

"""
    user_registered_functions(model::InfiniteModel)::Vector{RegisteredFunction}

Return all the functions (and their associated information) that the user has
registered to `model`. Each is stored as a [`RegisteredFunction`](@ref).
"""
function user_registered_functions(model::InfiniteModel)
    return model.registrations
end

## Define helper function to add registered functions to JuMP
# No gradient or hessian
function _add_func_data_to_jump(
    model::JuMP.Model,
    data::RegisteredFunction{F, Nothing, Nothing}
    ) where {F <: Function}
    JuMP.register(model, data.name, data.num_args, data.func, autodiff = true)
    return
end

# Only gradient information
function _add_func_data_to_jump(
    model::JuMP.Model,
    data::RegisteredFunction{F, G, Nothing}
    ) where {F <: Function, G <: Function}
    JuMP.register(model, data.name, data.num_args, data.func, data.gradient,
                  autodiff = isone(data.num_args))
    return
end

# Gradient and hessian information
function _add_func_data_to_jump(model::JuMP.Model, data::RegisteredFunction)
    if data.num_args > 1
        error("JuMP does not support hessians for multi-argument registered " *
              "functions.")
    end
    JuMP.register(model, data.name, data.num_args, data.func, data.gradient,
                  data.hessian)
    return
end

"""
    add_registered_to_jump(opt_model::JuMP.Model, inf_model::InfiniteModel)::Nothing

Add the user registered functions in `inf_model` to a `JuMP` model `opt_model`.
This is intended as an internal method, but it is provided for developers that
extend `InfiniteOpt` to use other optimizer models.
"""
function add_registered_to_jump(opt_model::JuMP.Model, inf_model::InfiniteModel)
    for data in user_registered_functions(inf_model)
        _add_func_data_to_jump(opt_model, data)
    end
    return
end

################################################################################
#                               LINEAR ALGEBRA
################################################################################
# Extend LinearAlgebra.dot for increased efficiency
LinearAlgebra.dot(lhs::AbstractInfOptExpr, rhs::AbstractInfOptExpr) = lhs * rhs
LinearAlgebra.dot(lhs::AbstractInfOptExpr, rhs::Real) = lhs * rhs
LinearAlgebra.dot(lhs::Real, rhs::AbstractInfOptExpr) = lhs * rhs

# Implement promote_rule to help build better containers
function Base.promote_rule(::Type{NLPExpr}, ::Type{<:Real})
    return NLPExpr
end
function Base.promote_rule(::Type{NLPExpr}, ::Type{GeneralVariableRef})
    return NLPExpr
end
function Base.promote_rule(::Type{NLPExpr}, ::Type{<:JuMP.GenericAffExpr})
    return NLPExpr
end
function Base.promote_rule(::Type{NLPExpr}, ::Type{<:JuMP.GenericQuadExpr})
    return NLPExpr
end

# TODO make proper MA extensions to enable efficient definition

################################################################################
#                                  PRINTING
################################################################################
# Define better printing for NodeData
function Base.show(io::IO, data::NodeData)
    return print(io, string(_node_value(data)))
end

# Map operators to their respective precedence (largest is highest priority)
const _Precedence = (; :^ => 6, Symbol("+u") => 5, Symbol("-u") => 5, :* => 4,
                     :/ => 4, :+ => 3, :- => 3, :(==) => 2, :<= => 2, :>= => 2,
                     :> => 2, :< => 2, :&& => 1, :|| => 1)

## Make functions to determine the precedence of a leaf
# AffExpr
function _leaf_precedence(aff::JuMP.GenericAffExpr)
    has_const = !iszero(JuMP.constant(aff))
    itr = JuMP.linear_terms(aff)
    num_terms = length(itr)
    if iszero(num_terms)
        # we have only a constant
        return 10 # will always have precedence
    elseif has_const || num_terms > 1
        # we have an expr with multiple terms
        return 3
    elseif isone(first(itr)[1])
        # we have a single variable
        return 10 # will always have precedence
    elseif first(itr)[1] == -1
        # we have a single unary negative variable
        return 5
    else
        # we have a single variable multiplied by some coefficient
        return 4
    end
end

# QuadExpr
function _leaf_precedence(quad::JuMP.GenericQuadExpr)
    has_aff = !iszero(quad.aff)
    itr = JuMP.quad_terms(quad)
    num_terms = length(itr)
    if iszero(num_terms)
        # we have an affine expression
        return _leaf_precedence(quad.aff)
    elseif has_aff || num_terms > 1
        # we have a general quadratic expression
        return 3
    else
        # we only have a single quadratic term
        return 4
    end
end

# Other
function _leaf_precedence(v)
    return 10
end

# Recursively build an expression string, starting with a root node
function _expr_string(
    node::_LCRST.Node{NodeData},
    str::String = "";
    prev_prec = 0,
    prev_comm = false
    )
    # prepocess the raw value
    raw_value = _node_value(node.data)
    is_op = raw_value isa Symbol && haskey(_Precedence, raw_value)
    data_str = _string_round(raw_value)
    # make a string according to the node structure
    if _LCRST.isleaf(node) && _leaf_precedence(raw_value) > prev_prec
        # we have a leaf that doesn't require parentheses
        return str * data_str
    elseif _LCRST.isleaf(node)
        # we have a leaf that requires parentheses
        return str * string("(", data_str, ")")
    elseif is_op && !_LCRST.islastsibling(node.child)
        # we have a binary operator
        curr_prec = _Precedence[raw_value]
        has_prec = curr_prec > prev_prec || (prev_comm && curr_prec == prev_prec)
        if !has_prec
            str *= "("
        end
        op_str = data_str == "^" ? data_str : string(" ", data_str, " ")
        is_comm = raw_value == :* || raw_value == :+
        for child in node
            str = string(_expr_string(child, str, prev_prec = curr_prec,
                                      prev_comm = is_comm), op_str)
        end
        str = str[1:prevind(str, end, length(op_str))]
        return has_prec ? str : str * ")"
    elseif is_op
        # we have a unary operator
        curr_prec = _Precedence[Symbol(raw_value, :u)]
        has_prec = curr_prec > prev_prec
        if !has_prec
            str *= "("
        end
        str *= string(data_str, _expr_string(node.child, str,
                                             prev_prec = curr_prec))
        return has_prec ? str : str * ")"
    else
        # we have a function
        str *= string(data_str, "(")
        for child in node
            str = _expr_string(child, str)
            str *= ", "
        end
        return str[1:prevind(str, end, 2)] * ")"
    end
end

# Extend JuMP.function_string for nonlinear expressions
function JuMP.function_string(mode, nlp::NLPExpr)
    return _expr_string(nlp.tree_root)
end


text:
function partialbridgeodeνH!(::R3, t, L, Σ, v, νt, Ht, P, ϵ)
    m, d = size(L)
    # print(typeof(H⁺t))
    # print(typeof(νt))
    # print(typeof(inv(L' * inv(Σ) * L + ϵ * I)   ))
    Ht[end] = (L' * inv(Σ) * L + ϵ * I)
    H⁺ = inv(Ht[end])
    νt[end] = H⁺ * L' * inv(Σ) * v
    ν = νt[end]

    for i in length(t)-1:-1:1
        dt = t[i] - t[i+1]
        ν = kernelr3((t, y, P) -> B(t, P)*y + β(t,P), t[i+1], ν, dt, P)
        H⁺ = kernelr3((t,  y, P) -> B(t, P)*y + y * B(t,P)'-a(t, P), t[i+1], H⁺, dt, P)

        νt[i] = ν
        Ht[i] = inv(H⁺)
    end
    νt, H⁺t
 end

"""
    PartialBridgeνH

Guided proposal process for diffusion bridge using backward recursion.

    PartialBridgeνH(tt, P, Pt,  L, v,ϵ Σ)

Guided proposal process for a partial diffusion bridge of `P` to `v` on
the time grid `tt` using guiding term derived from linear process `Pt`.

Simulate with `bridge!`.
"""
struct PartialBridgeνH{T,TP,TPt,Tv,Tν,TH} <: ContinuousTimeProcess{T}
    Target::TP
    Pt::TPt
    tt::Vector{Float64}
    v::Tv
    ν::Vector{Tν}
    H::Vector{TH}

    function PartialBridgeνH(tt_, P, Pt, L, v::Tv,ϵ, Σ = Bridge.outer(zero(v))) where {Tv}
        tt = collect(tt_)
        N = length(tt)
        m, d = size(L)
        TH = typeof(SMatrix{d,d}(1.0I))
        Ht = zeros(TH, N)
        Tν = typeof(@SVector zeros(d))
        νt = zeros(Tν, N)

        partialbridgeodeνH!(R3(), tt, L, Σ, v, νt, Ht, Pt,ϵ)
        new{Bridge.valtype(P),typeof(P),typeof(Pt),Tv,Tν,TH}(P, Pt, tt, v, νt, Ht)
    end
end

function _b((i,t)::IndexedTime, x, P::PartialBridgeνH)
    b(t, x, P.Target) + a(t, x, P.Target)*(P.H[i]*(P.ν[i] - x))
end

r((i,t)::IndexedTime, x, P::PartialBridgeνH) = P.H[i]*(P.ν[i] - x)
H((i,t)::IndexedTime, x, P::PartialBridgeνH) = inv(P.H⁺[i])

σ(t, x, P::PartialBridgeνH) = σ(t, x, P.Target)
a(t, x, P::PartialBridgeνH) = a(t, x, P.Target)
Γ(t, x, P::PartialBridgeνH) = Γ(t, x, P.Target)
constdiff(P::PartialBridgeνH) = constdiff(P.Target) && constdiff(P.Pt)

function llikelihood(::LeftRule, Xcirc::SamplePath, Po::PartialBridgeνH; skip = 0)
    tt = Xcirc.tt
    xx = Xcirc.yy

    som::Float64 = 0.
    for i in 1:length(tt)-1-skip #skip last value, summing over n-1 elements
        s = tt[i]
        x = xx[i]
        r = Bridge.r((i,s), x, Po)

        som += dot( _b((i,s), x, target(Po)) - _b((i,s), x, auxiliary(Po)), r ) * (tt[i+1]-tt[i])
        if !constdiff(Po)
            H = H(i, x, Po)
            som -= 0.5*tr( (a((i,s), x, target(Po)) - atilde((i,s), x, Po))*(H) ) * (tt[i+1]-tt[i])
            som += 0.5*( r'*(a((i,s), x, target(Po)) - atilde((i,s), x, Po))*r ) * (tt[i+1]-tt[i])
        end
    end
    som
end


text:
using Test
using SimpleTropical

x = Tropical(3.5)
y = Tropical(4)
inf = TropicalInf

@testset "Comparisons" begin
    @test x == x
    @test x != y

    @testset "Infinity" begin
        @test x != inf
        @test inf != y
        @test inf == inf
        @test inf != Tropical(0)
        @test Tropical(0) != inf

        @test !isequal(x, inf)
        @test !isequal(inf, y)
        @test isequal(inf, inf)
        @test !isequal(inf, Tropical(0))
        @test !isequal(Tropical(0), inf)
    end

    @testset "Not-a-number" begin
        nan = Tropical(NaN)
        @test nan != nan
        @test isequal(nan, nan)
    end
end

@testset "Sum" begin
    @test x+y == Tropical(3.5)
    @test x+inf == x
    @test inf+y == y
end

@testset "Product" begin
    z = Tropical(0)
    @test x*y == Tropical(7.5)
    @test x*z == x
    @test y*z == y
    @test inf*x == inf
    @test y*inf == inf
end

@testset "Inv and power" begin
    @test inv(x) == Tropical(-3.5)
    @test y^-1 == inv(y)
    @test x^10 == Tropical(35.0)
    @test y^0 == Tropical(0)
    @test inf^5 == inf
    @test_throws AssertionError inf^-3
    @test_throws AssertionError inv(inf)
end

@testset "Conversions" begin
    @test convert(Tropical, 5) === Tropical(5)
    @test convert(Tropical{Float64}, 5) === Tropical(5.0)
    @test convert(Tropical, x) === x
    @test convert(Tropical{Int}, y) === y
    @test convert(Tropical{Float32}, x) === Tropical(convert(Float32, x.val))
    @test convert(Tropical{Float64}, inf) === Tropical{Float64}(0.0, true)
end


text:
# This file contains shape transfer functions for HLO ops.
# It is lifted into the type domain by a type assertion in execute.jl,
# as well as being used in the compiler.

# HLO ops with specified type and dimensions
function shape_infer(op::Union{GenericHloOp, HloParameter, HloRng, HloConstant},
                     args::Type{<:XRTArray}...)
    return (op.T, op.shape)
end

eltypes_match(_::Type{<:XRTArray}) = true
eltypes_match(a1::Type{<:XRTArray}, a2::Type{<:XRTArray}, args::Type{<:XRTArray}...) =
    eltype(a1) == eltype(a2) && eltypes_match(a2, args...)

ndims_match(_::Type{<:XRTArray}) = true
ndims_match(a1::Type{<:XRTArray}, a2::Type{<:XRTArray}, args::Type{<:XRTArray}...) =
    ndims(a1) == ndims(a2) && eltypes_match(a2, args...)

shapes_match(_::Type{<:XRTArray}) = true
shapes_match(a1::Type{<:XRTArray}, a2::Type{<:XRTArray}, args::Type{<:XRTArray}...) =
    size(a1) == size(a2) && shapes_match(a2, args...)

@Base.pure non_elt_args(T, shape) = shape[1+ndims(T):end]

function shape_infer(op::HloBroadcast,
                     arg::Type{<:XRTArray}) where {T}
    (eltype(arg), op.result_shape)
end

function shape_infer(op::HloReshape,
                     arg::Type{<:XRTArray{T}}, args::Type{<:XRTArray}...) where {T}
    eltypes_match(arg, args...) || error("eltype mismatch in $(typeof(op))")
    (eltype(arg), non_elt_args(T, op.result_shape))
end

@Base.pure drop_reduce_dims(shape, dims) = tuple((shape[i] for i in 1:length(shape) if !(i-1 in dims))...)

function shape_infer(op::HloReduce{fT}, ::Type{<:fT},
                     args::Type{<:XRTArray}...) where {fT}
    # For now, restrict to a single reduction array
    @assert length(args) == 2
    shapes_match(args[1:div(length(args), 2)]...) || error("shape mismatch in $(typeof(op))")
    # For now, assume reductions preserve type
    (eltype(args[1]), drop_reduce_dims(size(args[1]), op.dims))
end

dilated_bound(bound, dilation) = (bound - 1) * dilation + 1
function strided_bound(bound, window_size, stride)
    (bound == 0 || window_size > bound) && return 0
    return div(bound - window_size, stride) + 1
end

@Base.pure function infer_window_output_shape(base_shape, windows)
    ntuple(length(windows)) do i
        window = windows[i]
        dilated_base = dilated_bound(base_shape[i], window.base_dilation)
        padded_dilated_base = window.padding_low + dilated_base + window.padding_high
        dilated_window = dilated_bound(window.size, window.window_dilation)
        return strided_bound(padded_dilated_base, dilated_window, window.stride)
    end
end

function shape_infer(op::HloReduceWindow{fT}, ::Type{<:fT},
                     args::Type{<:XRTArray}...) where {fT}
    # For now, assume reductions preserve type
    (eltype(args[1]), infer_window_output_shape(size(args[1]), op.window))
end

@Base.pure function pad_output_size(op::HloPad, arg::Type{<:XRTArray})
    sz = size(arg)
    ntuple(length(op.padding)) do i
        pad = op.padding[i]
        pad.edge_padding_low + pad.edge_padding_low + (sz[i]-1)*pad.interior_padding + sz[i]
    end
end

function shape_infer(op::HloPad, arg::Type{<:XRTArray}, val::Type{<:XRTArray})
    (eltype(arg), pad_output_size(op, arg))
end

function updims(s)
    dims = collect(map(x->x+1, s))
    ndims(dims) == 0 ? dims[] : dims
end

@Base.pure function shape_infer_conv(op::HloConv, lhs::Type{<:XRTArray}, kernel::Type{<:XRTArray})
    input_spatial_dims = size(lhs)[updims(op.dims.input_spatial_dimensions)]
    window_output_shape = infer_window_output_shape(input_spatial_dims, op.window)
    dimensions = Vector(undef, length(window_output_shape) + 2)
    dimensions[updims(op.dims.output_spatial_dimensions)] .= window_output_shape
    dimensions[updims(op.dims.output_batch_dimension)] =
        size(lhs)[updims(op.dims.input_batch_dimension)]
    dimensions[updims(op.dims.output_feature_dimension)] =
        size(kernel)[updims(op.dims.kernel_output_feature_dimension)]
    tuple(dimensions...)
end

function shape_infer(op::HloConv,
                     lhs::Type{<:XRTArray}, kernel::Type{<:XRTArray})
    shape = shape_infer_conv(op, lhs, kernel)
    # For now, assume convolution preserve type
    (eltype(lhs), shape)
end

function shape_infer(op::HloDynamicUpdateSlice, lhs::Type{<:XRTArray{T, Shp}}, args...) where {T, Shp}
    (T, Shp)
end

function shape_infer(op::HloDynamicSlice, args...) where {T, Shp}
    (eltype(args[1]), non_elt_args(eltype(args[1]), op.sizes))
end

@Base.pure function shape_infer(op::HloGather, A::Type{<:XRTArray}, idxs::Type{<:XRTArray})
    start_idxs = length(size(idxs)) == op.dims.index_vector_dim ?
        (size(idxs)..., 1) : size(idxs)
    result_rank = length(op.dims.offset_dims) + length(start_idxs) - 1
    slice_idx = gather_idx = 1
    shape = ntuple(result_rank) do i
        is_window_index = (i - 1) in op.dims.offset_dims
        if is_window_index
            while (slice_idx-1) in op.dims.collapsed_slice_dims
                slice_idx += 1
            end
            r = op.slice_sizes[slice_idx]; slice_idx += 1
            return r
        else
            if (gather_idx - 1) == op.dims.index_vector_dim
                gather_idx += 1
            end
            r = start_idxs[gather_idx]; gather_idx += 1
            return r
        end
    end
    (eltype(A), shape)
end

@Base.pure slice_dims(dims) = map(dim->div(dim.limit-dim.start, dim.stride), dims)
function shape_infer(op::HloSlice, args...) where {T,Shp}
    (eltype(args[1]), slice_dims(op.dims))
end

@Base.pure function compute_concat_shape(dim::Int64, args...)
    rank = ndims(args[1])
    all(x->ndims(x) == rank, args) || error("All arrays must have the same rank")
    dim <= rank - 1 || error("Concatenation dimension out of bounds")
    ntuple(ndims(args[1])) do i
        if i - 1 == dim
            return sum(j->size(args[j])[i], 1:length(args))
        else
            sz = size(args[1])[i]
            all(j->sz == size(args[j])[i], 1:length(args)) || error("Non-concatenated dimensions must match")
            return sz
        end
    end
end

function shape_infer(op::HloConcatenate, args::Type{<:XRTArray}...)
    eltypes_match(args...) || error("eltype mismatch in $(typeof(op))")
    ndims_match(args...) || error("ndims mismatch in $(typeof(op))")
    shp = compute_concat_shape(op.dim, args...)
    (eltype(args[1]), shp)
end

@Base.pure function compute_dot_dimensions(op::HloDot, lhs::Type{<:XRTArray}, rhs::Type{<:XRTArray})
    lcd = op.dims.lhs_contracting_dimensions[1]
    rcd = op.dims.rhs_contracting_dimensions[1]
    rbatch = op.dims.rhs_batch_dimensions
    dimensions = Int[]
    append!(dimensions, size(lhs)[filter(x->x != lcd+1, 1:ndims(lhs))])
    append!(dimensions, size(rhs)[filter(x->x != rcd+1 && !(x-1 ∈ rbatch), 1:ndims(rhs))])
    tuple(dimensions...)
end

function shape_infer(op::HloDot, lhs::Type{<:XRTArray}, rhs::Type{<:XRTArray})
    (eltype(lhs), compute_dot_dimensions(op, lhs, rhs))
end

function shape_infer(op::HloMap{fT}, ::Type{<:fT}, args::Type{<:XRTArray}...) where fT
    shapes_match(args...) || error("Shape mismatch in HloMap")
    (eltype(args[1]), size(args[1]))
end

@Base.pure transpose_dims(dims, permutation) = map(x->(x + 1) > length(dims) ? 1 : dims[x + 1], permutation)

function shape_infer(op::HloTranspose, args::Type{<:XRTArray}...)
    shapes_match(args...) || error("Shape mismatch in HloMap")
    (eltype(args[1]), transpose_dims(size(args[1]), op.permutation))
end

function shape_infer(op::HloSelectAndScatter{T, S}, ::Type{T}, ::Type{S}, args::Type{<:XRTArray}...) where {T,S}
    (eltype(args[1]), size(args[1]))
end

function shape_infer(op::HloScatter, T::Type, args...)
    (eltype(args[1]), size(args[1]))
end

function shape_infer(op::HloRev, A::Type{<:XRTArray})
    (eltype(A), size(A))
end

shape_infer(op::HloInfeed) = (op.infeed_shape.parameters[1],
                              op.infeed_shape.parameters[2])

shape_infer(op::HloIota) = (op.result_type, op.result_shape)

function infer_rt(op::HloOp, args::Type...)
    T, shape = shape_infer(op, args...)
    XRTArray{T, shape, length(shape)}
end

infer_rt(op::HloInfeed, arg::Type) = Tuple{op.infeed_shape, HloToken}
infer_rt(op::HloAfterAll, args::Type...) = HloToken
infer_rt(op::HloOutfeed, args::Type...) = HloToken

function infer_rt(op::HloCrossReplicaSum, fT::Type, args::Type...)
    Tuple{args...}
end

function infer_rt(op::HloMap, f, args::Type...)
    T, shape = shape_infer(op, typeof(f), args...)
    XRTArray{T, shape, length(shape)}
#=  TODO:
    T = Core.Compiler.return_type(f, Tuple{map(x->XRTArray{eltype(x), (), 0}, args)...})
    if isa(T, UnionAll) || T.abstract
        XRTArray{<:T,
            size(args[1]), length(size(args[1]))}
    else
        XRTArray{T,
            size(args[1]), length(size(args[1]))}
    end
=#
end

infer_rt(op::HloSort, f::Type, args::Type...) = Tuple{args...}


text:
export iframe

mutable struct IFrame
    innerHTML::AbstractString
end

function iframe(content)
    return node(IFrame(stringmime("text/html", content)))
end


text:
using WinReg
if VERSION < v"0.7.0-"
    using Base.Test
else
    using Test
end

@test querykey(WinReg.HKEY_LOCAL_MACHINE,"System\\CurrentControlSet\\Control\\Session Manager\\Environment","OS") == "Windows_NT"