@testset "DataCollection" begin
    mutable struct Nested
        data::Vector{Float64}
    end

    function initialize()
        model = ABM(
            Agent3,
            GridSpace((10, 10));
            properties = Dict(
                :year => 0,
                :tick => 0,
                :flag => false,
                :container => Float64[],
                :deep => Nested([20.0, 52.1])
            ),
        )
        add_agent!((4, 3), model, 0.1)
        add_agent!((7, 5), model, 0.35)
        add_agent!((2, 9), model, 0.67)
        return model
    end

    function agent_step!(agent, model)
        agent.weight += 0.05
        if model.tick % 365 == 0
            agent.weight *= 2
        end
    end
    function model_step!(model)
        model.tick += 1
        model.flag = !model.flag
        if model.tick % 365 == 0
            model.year += 1
            model.deep.data[1] += 0.5
        end
    end

    x_position(agent) = first(agent.pos)
    model = initialize()

    @testset "DataFrame init" begin
        @test init_agent_dataframe(model, nothing) == DataFrame()
        @test collect_agent_data!(DataFrame(), model, nothing, 1) == DataFrame()

        @test init_model_dataframe(model, nothing) == DataFrame()
        @test collect_model_data!(DataFrame(), model, nothing, 1) == DataFrame()

        props = [:weight]
        @test sprint(
            show,
            "text/csv",
            describe(init_agent_dataframe(model, props), :eltype),
        ) ==
              "\"variable\",\"eltype\"\n\"step\",\"Int64\"\n\"id\",\"Int64\"\n\"weight\",\"Float64\"\n"
        props = [:year]
        @test sprint(
            show,
            "text/csv",
            describe(init_model_dataframe(model, props), :eltype),
        ) == "\"variable\",\"eltype\"\n\"step\",\"Int64\"\n\"year\",\"Int64\"\n"

        @test_throws ErrorException init_agent_dataframe(model, [:UNKNOWN])
    end

    @testset "aggname" begin
        @test aggname(:weight) == "weight"
        @test aggname(:weight, mean) == "mean_weight"
        @test aggname(x_position, length) == "length_x_position"
        @test aggname((x_position, length)) == "length_x_position"
        ypos(a) = a.pos[2] > 5
        @test aggname((x_position, length, ypos)) == "length_x_position_ypos"
    end

    @testset "Aggregate Collections" begin
        props = [:weight]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props, 1)
        # Expecting weight values of all three agents. ID and step included.
        @test size(df) == (3, 3)
        @test propertynames(df) == [:step, :id, :weight]
        @test mean(df[!, :weight]) ≈ 0.37333333333

        props = [(:weight, mean)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props, 1)
        # Activate aggregation. Weight column is expected to be one value for this step,
        # renamed mean(weight). ID is meaningless and will therefore be dropped.
        @test size(df) == (1, 2)
        @test propertynames(df) == [:step, :mean_weight]
        @test df[1, aggname(:weight, mean)] ≈ 0.37333333333

        # Add a function as a property
        props = [:weight, x_position]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props, 1)
        @test size(df) == (3, 4)
        @test propertynames(df) == [:step, :id, :weight, :x_position]
        @test mean(df[!, :x_position]) ≈ 4.3333333

        props = [(:weight, mean), (x_position, mean)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props, 1)
        @test size(df) == (1, 3)
        @test propertynames(df) == [:step, :mean_weight, :mean_x_position]
        @test df[1, aggname(x_position, mean)] ≈ 4.3333333

        xtest(agent) = agent.pos[1] > 5
        ytest(agent) = agent.pos[2] > 5

        props = [(:weight, mean, ytest)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props, 1)
        @test size(df) == (1, 2)
        @test propertynames(df) == [:step, :mean_weight_ytest]
        @test df[1, aggname((:weight, mean, ytest))] ≈ 0.67

        props = [(:weight, mean), (:weight, mean, ytest)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props, 1)
        @test size(df) == (1, 3)
        @test propertynames(df) == [:step, :mean_weight, :mean_weight_ytest]
        @test df[1, aggname(:weight, mean)] ≈ 0.37333333333
        @test df[1, aggname(:weight, mean, ytest)] ≈ 0.67

        props = [(:weight, mean, xtest), (:weight, mean, ytest)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props, 1)
        @test size(df) == (1, 3)
        @test propertynames(df) == [:step, :mean_weight_xtest, :mean_weight_ytest]
        @test df[1, aggname(:weight, mean, xtest)] ≈ 0.35
        @test df[1, aggname(:weight, mean, ytest)] ≈ 0.67
    end

    @testset "High-level API for Collections" begin
        # Extract data from the model every year for five years,
        # with the average `weight` of all agents every six months.
        each_year(model, step) = step % 365 == 0
        six_months(model, step) = step % 182 == 0
        agent_data, model_data = run!(
            model,
            agent_step!,
            model_step!,
            365 * 5;
            when_model = each_year,
            when = six_months,
            mdata = [:flag, :year],
            adata = [(:weight, mean)],
        )

        @test size(agent_data) == (11, 2)
        @test propertynames(agent_data) == [:step, :mean_weight]
        @test maximum(agent_data[!, :step]) == 1820

        @test size(model_data) == (6, 3)
        @test propertynames(model_data) == [:step, :flag, :year]
        @test maximum(model_data[!, :step]) == 1825

        agent_data, model_data = run!(
            model,
            agent_step!,
            model_step!,
            365 * 5;
            when_model = [1, 365 * 5],
            when = false,
            mdata = [:flag, :year],
            adata = [(:weight, mean)],
        )
        @test size(agent_data) == (0, 2)
        @test size(model_data) == (2, 3)

        _, model_data = run!(
            model,
            agent_step!,
            model_step!,
            365 * 5;
            when_model = [1, 365 * 5],
            when = false,
            mdata = [:deep],
            obtainer = deepcopy,
        )
        @test model_data[1, :deep].data[1] < model_data[end, :deep].data[1]
    end

    @testset "Low-level API for Collections" begin
        # Generate three separate dataframes using the low level API.
        # User controls the evolution, as well as the identity of each `step`.
        # Note in this example daily data are labelled with a daily `step`,
        # and yearly data with a yearly `step`.
        model = initialize()
        model_props = [:flag, :year]
        agent_agg = [(:weight, mean)]
        agent_props = [:weight]
        daily_model_data = init_model_dataframe(model, model_props)
        daily_agent_aggregate = init_agent_dataframe(model, agent_agg)
        yearly_agent_data = init_agent_dataframe(model, agent_props)

        for year in 1:5
            for day in 1:365
                step!(model, agent_step!, model_step!, 1)
                collect_model_data!(daily_model_data, model, model_props, day * year)
                collect_agent_data!(daily_agent_aggregate, model, agent_agg, day * year)
            end
            collect_agent_data!(yearly_agent_data, model, agent_props, year)
        end

        @test size(daily_model_data) == (1825, 3)
        @test propertynames(daily_model_data) == [:step, :flag, :year]
        @test maximum(daily_model_data[!, :step]) == 1825

        @test size(daily_agent_aggregate) == (1825, 2)
        @test propertynames(daily_agent_aggregate) == [:step, :mean_weight]
        @test maximum(daily_agent_aggregate[!, :step]) == 1825

        @test size(yearly_agent_data) == (15, 3)
        @test propertynames(yearly_agent_data) == [:step, :id, :weight]
        @test maximum(yearly_agent_data[!, :step]) == 5

        @test dummystep(model) == nothing
        @test dummystep(model[1], model) == nothing
        tick = model.tick
        step!(model, agent_step!, 1)
        @test tick == model.tick
        stop(m, s) = m.year == 6
        step!(model, agent_step!, model_step!, stop)
        @test model.tick == 365 * 6
        step!(model, agent_step!, model_step!, 1, false)
    end

    @testset "Mixed-ABM collections" begin
        model = ABM(Union{Agent3, Agent4}, GridSpace((10,10)); warn = false)
        add_agent_pos!(Agent3(1, (6,8), 54.65),model)
        add_agent_pos!(Agent4(2, (10,7), 5),model)

        # Expect position type (both agents have it)
        props = [:pos]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (2, 4)
        @test propertynames(df) == [:step, :id, :agent_type, :pos]
        @test df[1, :pos] == model[1].pos
        @test df[2, :pos] == model[2].pos

        # Expect weight for Agent3, but missing for Agent4
        props = [:weight]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (2, 4)
        @test df[1, :weight] == model[1].weight
        @test ismissing(df[2, :weight])

        # Expect similar output, but using functions
        wpos(a) = a.pos[1] + a.weight
        props = [wpos]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (2, 4)
        @test df[1, :wpos] == model[1].pos[1] + model[1].weight
        @test ismissing(df[2, :wpos])

        add_agent_pos!(Agent3(3, (2,4), 19.81),model)
        add_agent_pos!(Agent4(4, (4,1), 3),model)

        props = [:pos, :weight, :p, wpos]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (4, 7)
        @test typeof(df.pos) <: Vector{Tuple{Int, Int}}
        @test typeof(df.weight) <: Vector{Union{Missing, Float64}}
        @test typeof(df.p) <: Vector{Union{Missing, Int}}
        @test typeof(df.wpos) <: Vector{Union{Missing, Float64}}

        # Expect something completely unknown to fail
        props = [:UNKNOWN]
        @test_throws ErrorException init_agent_dataframe(model, props)

        # Aggregates should behave in a similiar fashion
        pos1(a) = a.pos[1]
        props = [(:pos, length), (pos1, sum)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (1, 3)
        @test typeof(df.length_pos) <: Vector{Int}
        @test df[1, :length_pos] == 4
        @test typeof(df.sum_pos1) <: Vector{Int}
        @test df[1, :sum_pos1] == 22

        # Expect aggregate to filter out Agent4's
        a3(a) = a isa Agent3
        props = [(pos1, sum, a3)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (1, 2)
        @test typeof(df.sum_pos1_a3) <: Vector{Int}
        @test df[1, :sum_pos1_a3] == 8

        # When aggregating, missing data must be handled explicitly.
        props = [(:weight, sum)]
        df = init_agent_dataframe(model, props)
        @test_throws ErrorException collect_agent_data!(df, model, props)

        # Filtering out missings makes this work.
        props = [(:weight, sum, a3)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (1, 2)
        @test typeof(df.sum_weight_a3) <: Vector{Float64}
        @test df[1, :sum_weight_a3] ≈ 74.46

        # Handle missmatches
        # In this example, weight exists in both agents, but they have different types
        mutable struct Agent3Int <: AbstractAgent
            id::Int
            pos::Dims{2}
            weight::Int
        end
        model = ABM(Union{Agent3, Agent3Int}, GridSpace((10,10)); warn = false)
        add_agent_pos!(Agent3(1, (6,8), 54.65),model)
        add_agent_pos!(Agent3Int(2, (10,7), 5),model)
        add_agent_pos!(Agent3(3, (2,4), 19.81),model)
        add_agent_pos!(Agent3Int(4, (4,1), 3),model)

        props = [:weight]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (4, 4)
        @test typeof(df.weight) <: Vector{Union{Float64, Int}}

        # Expect a1.weight <: Float64, a2.weight <: Int64 to fail in aggregate
        props = [(:weight, sum)]
        @test_throws ErrorException init_agent_dataframe(model, props)

        # Promotion is the fix in this case
        fweight(a) = Float64(a.weight)
        props = [(fweight, sum)]
        df = init_agent_dataframe(model, props)
        collect_agent_data!(df, model, props)
        @test size(df) == (1, 2)
        @test df[1, :sum_fweight] ≈ 82.46
    end

    @testset "Observers" begin
        model = initialize()
        model_props = [:container]
        model_data = init_model_dataframe(model, model_props)
        push!(model.container, 50.0)
        collect_model_data!(model_data, model, model_props, 0; obtainer = copy)
        push!(model.container, 37.2)
        collect_model_data!(model_data, model, model_props, 1; obtainer = copy)
        model.container[1] += 21.9
        collect_model_data!(model_data, model, model_props, 2; obtainer = copy)
        @test model_data.container[1][1] ≈ 50.0
        @test model_data.container[3][1] ≈ 71.9
        @test length.(model_data.container) == [1, 2, 2]

        model = initialize()
        model_props = [:deep]
        model_data = init_model_dataframe(model, model_props)
        push!(model.deep.data, 17.5)
        collect_model_data!(model_data, model, model_props, 0; obtainer = deepcopy)
        push!(model.deep.data, 1.2)
        collect_model_data!(model_data, model, model_props, 1; obtainer = deepcopy)
        model.deep.data[1] += 0.9
        collect_model_data!(model_data, model, model_props, 2; obtainer = deepcopy)
        @test model_data[1,:deep].data[1] ≈ 20.0
        @test model_data[3,:deep].data[1] ≈ 20.9
        @test [length(d.data) for d in model_data[!,:deep]] == [3, 4, 4]

        model = initialize()
        agent_data, model_data = run!(
            model,
            agent_step!,
            model_step!,
            365 * 5;
            when_model = [1, 365 * 5],
            when = false,
            mdata = [:flag, :year, :container, :deep],
            adata = [(:weight, mean)],
            obtainer = deepcopy,
        )
        @test size(agent_data) == (0, 2)
        @test size(model_data) == (2, 5)
    end
end

@testset "Parameter scan" begin
    n = 10
    parameters = Dict(
        :density => [0.6, 0.7, 0.8],
        :griddims => (20, 20),
    )

    forest, agent_step!, forest_step! = Models.forest_fire()
    forest_initiation(; kwargs...) = Models.forest_fire(; kwargs...)[1]

    burnt(a) = a.status == :burnt
    unburnt(a) = a.status == :green
    @testset "Standard Scan" begin
        adata = [(unburnt, count), (burnt, count)]
        data, _ = paramscan(
            parameters,
            forest_initiation;
            n = n,
            agent_step! = agent_step!,
            model_step! = forest_step!,
            adata = adata,
            progress = false,
        )
        # 3 is the number of combinations of changing params
        @test size(data) == ((n + 1) * 3, 4)
        data, _ = paramscan(
            parameters,
            forest_initiation;
            n = n,
            agent_step! = agent_step!,
            model_step! = forest_step!,
            include_constants = true,
            adata = adata,
            progress = false,
        )
        # 3 is the number of combinations of changing params,
        # 5 is 3+2, where 2 is the number of constant parameters
        @test size(data) == ((n + 1) * 3, 5)

        adata = [:status]
        data, _ = paramscan(
            parameters,
            forest_initiation;
            n = n,
            agent_step! = agent_step!,
            model_step! = forest_step!,
            adata = adata,
            progress = false,
        )
        @test unique(data.step) == 0:10
        @test unique(data.density) == [0.6, 0.7, 0.8]
    end

    @testset "Scan with replicates" begin
        adata = [(unburnt, count), (burnt, count)]
        data, _ = paramscan(
            parameters,
            forest_initiation;
            n = n,
            agent_step! = dummystep,
            model_step! = forest_step!,
            replicates = 3,
            adata = adata,
            progress = false,
        )
        # the first 3 is the number of combinations of changing params
        @test size(data) == (((n + 1) * 3) * 3, 5)
    end
end

@testset "Issue #179 fix" begin
    # only ids sorted, not properties
    model = ABM(Agent2)
    for i in 1:5
        add_agent!(model, i * 0.2)
    end
    data, _ = run!(model, dummystep, 2; adata = [:weight])
    @test data[1, :id] == 1 && data[1, :weight] ≈ 0.2
    @test data[3, :id] == 3 && data[3, :weight] ≈ 0.6
    @test data[6, :id] == 1 && data[6, :weight] ≈ 0.2
end


function >>(left::Expr, right::Expr)
    quote
        $left
        $right
    end
end

function mapcatexpr(f, xs)
    body = quote end
    for x in xs
        next = f(x)
        body = body >> next
    end
    prettify(body)
end

function brackets(i)
    var -> :($var[$i])
end

function condname(var)
    Symbol("cond__" * string(var))
end

function footer(observe)
    quote
        loglikelihood += current.loglikelihood
        if loglikelihood == -Inf
            $observe
            return traceroot
        end
        previous = current
        current = current.next
    end
end

function varbody(model, dist, var, access)
    invar = access(var)
    random = !haskey(model.conditionings, var)
    value = random ? :(rand($dist)) : access(condname(var))

    quote
        if current isa Leaf
            $invar = $value
            current = extend!(previous, $(quot(var)), $dist, $invar, $random)

            # TODO: we could move this if out of the runtime to avoid checking this
            # on each node
            if traceroot isa Leaf
                traceroot = current
            end
        else
            $invar = current.value
        end
    end
end


abstract type AbstractSyntheticSatellite{T} <: AbstractEmitter{T} end

struct SyntheticSatellite{
    T <: AbstractFloat,
    CS <: ConstantDopplerSatellite{<: AbstractGNSS, T},
    S <: Union{Nothing, <:AbstractVector}
} <: AbstractSyntheticSatellite{T}
    sat::CS
    steer_vec::S
end

function SyntheticSatellite(sat::ConstantDopplerSatellite{<:AbstractGNSS, T}) where T <: AbstractFloat
    SyntheticSatellite(sat, nothing)
end

function gen_signal!(
    synthetic_sat::SyntheticSatellite,
    sample_frequency,
    intermediate_frequency,
    n0,
    num_samples::Integer,
    rng
)
    gen_signal!(
        synthetic_sat.sat,
        sample_frequency,
        intermediate_frequency,
        n0,
        num_samples,
        rng
    )
end

function propagate(
    synthetic_sat::SyntheticSatellite,
    num_samples,
    intermediate_frequency,
    sample_frequency,
    n0,
    rng
)
    SyntheticSatellite(
        propagate(
            synthetic_sat.sat,
            num_samples,
            intermediate_frequency,
            sample_frequency,
            n0,
            rng
        ),
        synthetic_sat.steer_vec
    )
end

function get_steer_vec(
    manifold::AbstractManifold{N},
    emitter::SyntheticSatellite{T, CS, Nothing},
    attitude
) where {N, T <: AbstractFloat, CS <: ConstantDopplerSatellite}
    ones(SVector{N, T})
end

function get_steer_vec(
    manifold::IdealManifold{1},
    emitter::SyntheticSatellite{T, CS, Nothing},
    attitude
) where {T <: AbstractFloat, CS <: ConstantDopplerSatellite}
    one(T)
end

function get_steer_vec(
    manifold::AbstractManifold,
    emitter::SyntheticSatellite,
    attitude
)
    emitter.steer_vec
end

@inline get_doa(ss::SyntheticSatellite) = get_doa(ss.sat)
@inline get_existence(ss::SyntheticSatellite) = get_existence(ss.sat)
@inline get_carrier_doppler(ss::SyntheticSatellite) = get_carrier_doppler(ss.sat)
@inline get_code_doppler(ss::SyntheticSatellite) = get_code_doppler(ss.sat)
@inline get_carrier_phase(ss::SyntheticSatellite) = get_carrier_phase(ss.sat)
@inline get_code_phase(ss::SyntheticSatellite) = get_code_phase(ss.sat)
@inline get_prn(ss::SyntheticSatellite) = get_prn(ss.sat)
@inline get_amplitude(ss::SyntheticSatellite, n0) = get_amplitude(ss.sat, n0)
@inline get_gnss_system(ss::SyntheticSatellite) = get_gnss_system(ss.sat)
@inline get_carrier_to_noise_density_ratio(ss::SyntheticSatellite) =
    get_carrier_to_noise_density_ratio(ss.sat)


using HypothesisTests, Test

@testset "White tests" begin

X = [ 4.18      0.77      1.00;
     -3.41      0.73      1.00;
      5.75      0.81      1.00;
      0.05      0.80      1.00;
     -2.18      0.82      1.00;
      3.88      0.81      1.00;
      0.73      0.77      1.00;
      5.70      0.77      1.00;
     -0.69      0.83      1.00;
     -8.14      0.87      1.00;
      5.37      0.99      1.00;
      1.87      0.95      1.00;
      5.76      0.80      1.00;
     -0.79      0.89      1.00;
    -13.23      1.21      1.00;
      3.97      1.26      1.00;
      5.20      0.81      1.00;
      3.16      0.61      1.00;
      6.41      0.53      1.00;
      1.72      0.64      1.00;
      2.20      0.75      1.00;
      1.06      0.95      1.00;
      9.54      0.96      1.00;
     -4.75      1.31      1.00;
     -5.05      1.04      1.00;
      0.48      1.07      1.00;
      3.41      1.21      1.00;
     -2.21      1.08      1.00;
      0.21      1.15      1.00;
     -2.37      1.35      1.00;
     -1.55      1.24      1.00;
     -6.91      1.28      1.00;
     -7.62      1.24      1.00;
      4.81      1.21      1.00;
      3.51      1.07      1.00;
     -3.68      0.87      1.00;
     -3.42      0.80      1.00;
     -6.03      0.92      1.00;
     -1.99      0.98      1.00;
      3.20      1.13      1.00;
     -3.88      1.06      1.00;
     -3.35      0.96      1.00;
     -3.10      1.05      1.00;
     11.14      0.76      1.00;
      1.17      0.51      1.00;
     11.27      0.59      1.00;
      4.56      0.63      1.00;
      0.78      0.67      1.00;
      3.50      0.69      1.00;
      2.40      0.62      1.00;
      2.84      0.63      1.00;
      6.71      0.71      1.00;
      0.63      0.69      1.00;
      3.11      0.67      1.00;
     -3.90      0.74      1.00;
     -0.41      0.76      1.00;
      0.85      0.76      1.00;
     -3.56      0.76      1.00;
      2.26      0.70      1.00;
     -1.78      0.73      1.00;
     -2.06      0.76      1.00;
     -4.62      0.71      1.00;
      0.61      0.73      1.00;
     -0.56      0.81      1.00;
     -6.01      0.78      1.00;
      1.59      0.75      1.00;
     -2.88      0.82      1.00;
     10.44      0.83      1.00;
     -0.82      0.86      1.00;
     -1.01      1.00      1.00;
     -1.80      0.73      1.00;
      1.73      0.64      1.00;
      7.92      0.65      1.00;
      1.11      0.58      1.00;
     -0.79      0.62      1.00;
     -0.94      0.72      1.00;
      4.92      0.66      1.00;
      1.16      0.55      1.00;
     -0.65      0.62      1.00;
     -1.03      0.55      1.00;
     -4.58      0.60      1.00;
      3.79      0.65      1.00;
      6.31      0.61      1.00;
      3.66      0.65      1.00;
      0.42      0.56      1.00;
      6.72      0.53      1.00;
      4.79      0.60      1.00;
     -1.31      0.52      1.00;
      4.59      0.49      1.00;
      0.90      0.52      1.00;
     -6.49      0.52      1.00;
      6.16      0.46      1.00;
     -8.35      0.45      1.00;
      4.47      0.46      1.00;
      1.12      0.39      1.00;
     -3.13      0.49      1.00;
     12.43      0.42      1.00;
      4.36      0.43      1.00;
      1.90      0.47      1.00;
     -2.14      0.44      1.00 ]

e = [ 5.19346;
      2.22296;
      3.82622;
      2.93541;
      0.70453;
      1.98955;
      0.56655;
      1.54306;
      0.57363;
     -0.13346;
      1.68565;
      9.24558;
      7.39317;
      2.45484;
     -4.83724;
     -0.74742;
     -0.22633;
      0.51537;
      2.82822;
      9.35943;
      5.38263;
      6.59691;
      1.79080;
      0.54506;
      4.06661;
     -4.18808;
      1.69414;
      5.40963;
      2.92762;
     -2.67351;
     -2.74613;
     -3.52752;
     -4.27775;
      2.88245;
     -8.12566;
      0.08760;
      2.20495;
     -0.45509;
     -0.06959;
      1.51330;
      1.37878;
     -2.33046;
      0.35730;
    -11.41189;
     -0.48083;
      1.08917;
      2.52333;
      1.73218;
     15.60426;
     -0.15961;
     -0.22174;
     -1.19174;
     13.92022;
      1.16044;
      0.08087;
     -5.81547;
     -5.14199;
     -5.69918;
     -0.36585;
     -3.70294;
      2.83687;
     -1.37631;
     -2.55403;
     -4.10796;
      1.32092;
     -1.64056;
     -4.09963;
     -4.85728;
     -1.45496;
     -3.12612;
     -4.36648;
     -3.23379;
      5.22893;
      5.06730;
     -0.63041;
     -1.70377;
     -5.71335;
     -2.76831;
      2.20443;
     -0.10174;
     -2.36727;
     -5.49859;
     -4.78093;
     -3.36665;
      5.12027;
     -3.03180;
     -2.52035;
      3.10913;
     -3.51388;
     -0.70389;
     -1.01963;
     -8.78989;
      2.95170;
     -3.95464;
     -4.01258;
     -1.92315;
     -4.94209;
      5.29138;
      1.28435;
      0.28832 ]

    atol = 1e-4
    w_test = WhiteTest(X, e, type = :linear)
    @test pvalue(w_test) ≈ 0.1287 atol=atol

    w_test = WhiteTest(X, e, type = :linear_and_squares)
    @test pvalue(w_test) ≈ 0.2774 atol=atol

    w_test = WhiteTest(X, e)
    w_pval = pvalue(w_test)
    @test w_pval ≈ 0.3458 atol=atol

    refstr = """
    White's (or Breusch-Pagan's) test for heteroskedasticity
    --------------------------------------------------------
    Population details:
        parameter of interest:   T*R2
        value under h_0:         0
        point estimate:          $(w_test.lm)

    Test summary:
        outcome with 95% confidence: fail to reject h_0
        one-sided p-value:           $(round(w_pval,digits=4))

    Details:
        T*R^2 statistic:        $(w_test.lm)
        degrees of freedom:     $(w_test.dof)
        type:                   $(w_test.type)
    """

    buffer = IOBuffer()
    show(buffer, w_test)
    str = String(take!(buffer))
    @test str == refstr

    @test_throws DimensionMismatch WhiteTest([rand(3) ones(3)], rand(4))

    @test_throws ArgumentError WhiteTest(ones(4,1), rand(4))

    @test_throws ArgumentError WhiteTest([zeros(4) rand(4)], rand(4))

    @test_throws ArgumentError WhiteTest(rand(4,2), rand(4))

    bp_test = BreuschPaganTest(X, e)
    @test pvalue(bp_test) ≈ 0.1287 atol=atol
    show(IOBuffer(), bp_test)
end


### Abstract MALA state

abstract MALAState{F<:VariateForm} <: LMCSamplerState{F}

### MALA state subtypes

## UnvMALAState holds the internal state ("local variables") of the MALA sampler for univariate parameters

type UnvMALAState <: MALAState{Univariate}
  pstate::ParameterState{Continuous, Univariate} # Parameter state used internally by MALA
  tune::MCTunerState
  ratio::Real
  μ::Real

  function UnvMALAState(pstate::ParameterState{Continuous, Univariate}, tune::MCTunerState, ratio::Real, μ::Real)
    if !isnan(ratio)
      @assert ratio > 0 "Acceptance ratio should be positive"
    end
    new(pstate, tune, ratio, μ)
  end
end

UnvMALAState(pstate::ParameterState{Continuous, Univariate}, tune::MCTunerState=BasicMCTune()) =
  UnvMALAState(pstate, tune, NaN, NaN)

## MuvMALAState holds the internal state ("local variables") of the MALA sampler for multivariate parameters

type MuvMALAState <: MALAState{Multivariate}
  pstate::ParameterState{Continuous, Multivariate} # Parameter state used internally by MALA
  tune::MCTunerState
  ratio::Real
  μ::RealVector

  function MuvMALAState(pstate::ParameterState{Continuous, Multivariate}, tune::MCTunerState, ratio::Real, μ::RealVector)
    if !isnan(ratio)
      @assert ratio > 0 "Acceptance ratio should be positive"
    end
    new(pstate, tune, ratio, μ)
  end
end

MuvMALAState(pstate::ParameterState{Continuous, Multivariate}, tune::MCTunerState=BasicMCTune()) =
  MuvMALAState(pstate, tune, NaN, Array(eltype(pstate), pstate.size))

### Metropolis-adjusted Langevin Algorithm (MALA)

immutable MALA <: LMCSampler
  driftstep::Real

  function MALA(driftstep::Real)
    @assert driftstep > 0 "Drift step is not positive"
    new(driftstep)
  end
end

MALA() = MALA(1.)

### Initialize MALA sampler

## Initialize parameter state

function initialize!{F<:VariateForm}(
  pstate::ParameterState{Continuous, F},
  parameter::Parameter{Continuous, F},
  sampler::MALA
)
  parameter.uptogradlogtarget!(pstate)
  @assert isfinite(pstate.logtarget) "Log-target not finite: initial value out of support"
  @assert all(isfinite(pstate.gradlogtarget)) "Gradient of log-target not finite: initial values out of support"
end

## Initialize MALA state

sampler_state(
  parameter::Parameter{Continuous, Univariate},
  sampler::MALA,
  tuner::MCTuner,
  pstate::ParameterState{Continuous, Univariate},
  vstate::VariableStateVector
) =
  UnvMALAState(generate_empty(pstate), tuner_state(parameter, sampler, tuner))

sampler_state(
  parameter::Parameter{Continuous, Multivariate},
  sampler::MALA,
  tuner::MCTuner,
  pstate::ParameterState{Continuous, Multivariate},
  vstate::VariableStateVector
) =
  MuvMALAState(generate_empty(pstate), tuner_state(parameter, sampler, tuner))

## Reset parameter state

function reset!(
  pstate::ParameterState{Continuous, Univariate},
  x::Real,
  parameter::Parameter{Continuous, Univariate},
  sampler::MALA
)
  pstate.value = x
  parameter.uptogradlogtarget!(pstate)
end

function reset!(
  pstate::ParameterState{Continuous, Multivariate},
  x::RealVector,
  parameter::Parameter{Continuous, Multivariate},
  sampler::MALA
)
  pstate.value = copy(x)
  parameter.uptogradlogtarget!(pstate)
end

Base.show(io::IO, sampler::MALA) = print(io, "MALA sampler: drift step = $(sampler.driftstep)")

Base.writemime(io::IO, ::MIME"text/plain", sampler::MALA) = show(io, sampler)


## ExoplanetsSysSim/examples/hsu_etal_2018/misc_utils/convert_koi_cat_csv_to_jld.jl
## (c) 2018 Danley C. Hsu
# Converts KOI catalog data file from csv to jld format

using ExoplanetsSysSim
using HDF5, JLD

include(joinpath(Pkg.dir(),"ExoplanetsSysSim","examples","hsu_etal_2018", "christiansen_func.jl"))

koi_catalog_file_in = joinpath(Pkg.dir("ExoplanetsSysSim"), "data", "q1_q16_koi_cand.csv")
koi_catalog_file_out = joinpath(Pkg.dir("ExoplanetsSysSim"), "data", "q1_q16_koi_cand.jld")

koi_catalog, koi_usable = read_koi_catalog(ascii(koi_catalog_file_in))

save(koi_catalog_file_out,"koi_catalog", koi_catalog, "koi_catalog_usable", koi_usable)


module PhaseNet

# Julia translation from:
# Automatic picking of seismic waves using Generalized Phase Detection
# See http://scedc.caltech.edu/research-tools/deeplearning.html for more info
#
# Ross et al. (2018), Generalized Seismic Phase Detection with Deep Learning,
#                     Bull. Seismol. Soc. Am., doi:10.1785/0120180080
#
# Original Author: Zachary E. Ross (2018)
# Contact: zross@gps.caltech.edu
# Website: http://www.seismolab.caltech.edu/ross_z.html

using Dates
using PyCall, SeisIO

# load torch modules
export torch,nn, UNet

const torch = PyNULL()
const nn = PyNULL()

function __init__()
    copy!(torch, pyimport_conda("torch","pytorch","pytorch"))
    copy!(nn, pyimport("torch.nn"))
    include(joinpath(@__DIR__, "UNet.jl"))
end

# include files
include("detection.jl")

end


import CUDA.CUDNN: cudnnNormalizationForwardAD
import AutoGrad: forw
using CUDA.CUDNN: handle, CU_NULL, cudnnNormalizationBackward, cudnnNormalizationForwardTraining, cudnnNormalizationForwardInference, cudnnGetNormalizationForwardTrainingWorkspaceSize, cudnnGetNormalizationTrainingReserveSpaceSize, cudnnTempSpace, scalingParameter
using AutoGrad: AutoGrad, @primitive1, forwargs, recording, Result, value

@primitive1((cudnnNormalizationForwardAD(x, scale, bias, z; training, mean, variance, y, mode, normOps, algo, alpha, beta, epsilon, groupCnt, exponentialAverageFactor, savedMean, savedInvVariance, activationDesc, xDesc, yDesc, zDesc, normScaleBiasDesc, normMeanVarDesc, workspace, reserveSpace, dx, dscale, dbias, dz, dready), dy, _y),
            (dready[] || cudnnNormalizationBack(dy, x, scale, bias, z; training, mean, variance, y, mode, normOps, algo, alpha, beta, epsilon, groupCnt, exponentialAverageFactor, savedMean, savedInvVariance, activationDesc, xDesc, yDesc, zDesc, normScaleBiasDesc, normMeanVarDesc, workspace, reserveSpace, dx, dscale, dbias, dz, dready); dx[]),
            (dready[] || cudnnNormalizationBack(dy, x, scale, bias, z; training, mean, variance, y, mode, normOps, algo, alpha, beta, epsilon, groupCnt, exponentialAverageFactor, savedMean, savedInvVariance, activationDesc, xDesc, yDesc, zDesc, normScaleBiasDesc, normMeanVarDesc, workspace, reserveSpace, dx, dscale, dbias, dz, dready); dscale[]),
            (dready[] || cudnnNormalizationBack(dy, x, scale, bias, z; training, mean, variance, y, mode, normOps, algo, alpha, beta, epsilon, groupCnt, exponentialAverageFactor, savedMean, savedInvVariance, activationDesc, xDesc, yDesc, zDesc, normScaleBiasDesc, normMeanVarDesc, workspace, reserveSpace, dx, dscale, dbias, dz, dready); dbias[]),
            (dready[] || cudnnNormalizationBack(dy, x, scale, bias, z; training, mean, variance, y, mode, normOps, algo, alpha, beta, epsilon, groupCnt, exponentialAverageFactor, savedMean, savedInvVariance, activationDesc, xDesc, yDesc, zDesc, normScaleBiasDesc, normMeanVarDesc, workspace, reserveSpace, dx, dscale, dbias, dz, dready); dz[]))

function cudnnNormalizationBack(dy, x, scale, bias, z; training, mean, variance, y, mode, normOps, algo, alpha, beta, epsilon, groupCnt, exponentialAverageFactor, savedMean, savedInvVariance, activationDesc, xDesc, yDesc, zDesc, normScaleBiasDesc, normMeanVarDesc, workspace, reserveSpace, dx, dscale, dbias, dz, dready)
    # Make sure backward gets called only once
    dready[] && return
    dready[] = true
    (x, scale, bias, z) = value.((x, scale, bias, z))
    # Allocate gradient buffers if necessary
    !isassigned(dx) || dx[]===nothing ? dx[] = similar(x) : @assert issimilar(x, dx[])
    !isassigned(dscale) || dscale[]===nothing ? dscale[] = similar(scale) : @assert issimilar(scale, dscale[])
    !isassigned(dbias) || dbias[]===nothing ? dbias[]  = similar(bias) : @assert issimilar(bias, dbias[])
    z === nothing ? dz[] = nothing : !isassigned(dz) || dz[]===nothing ? dz[] = zero(z) : (@assert issimilar(z, dz[]); dz[] .= 0) # z may not be used, dz may not be modified, should be zeroed
    s0,s1 = scalingParameter(eltype(x),0),scalingParameter(eltype(x),1)
    cudnnNormalizationBackward(handle(), mode, normOps, algo, alpha, s0, alpha, s0, xDesc, x, yDesc, y, yDesc, dy, something(zDesc,C_NULL), something(dz[],CU_NULL), xDesc, dx[], normScaleBiasDesc, scale, bias, dscale[], dbias[], epsilon, something(normMeanVarDesc,C_NULL), something(savedMean,CU_NULL), something(savedInvVariance,CU_NULL), something(activationDesc,C_NULL), something(workspace,CU_NULL), sizeof(workspace), something(reserveSpace,CU_NULL), sizeof(reserveSpace), groupCnt)
end

# CUDNN API controls training vs inference mode with a simple keyword argument.
# In Knet things are a bit more complicated: whether or not we are in @diff mode and we have Value args enter the picture.
# With Value args the AutoGrad.forw method is called regardless of the kwarg / @diff mode.
# The training keyword argument defaults to Knet.training() in Knet API, so by default @diff=training, but the user can override this.
# We have 8 cases to consider for Value, training, @diff:
# When there are no Value arguments, i.e. scale/bias not Params and x is not a Param or Result:
# cudnnNormalizationForwardAD is called, @diff=false, training=false: AD calls Inference: ok
# cudnnNormalizationForwardAD is called, @diff=false, training=true: AD calls Training: user is forcing training mode
# cudnnNormalizationForwardAD is called, @diff=true, training=false: AD calls Inference: user is forcing inference mode
# cudnnNormalizationForwardAD is called, @diff=true, training=true: AD calls Training: ok (but no Value args no back pass)
# When there are some Value arguments:
# forw is called, @diff=false, training=false: Standard test mode. forw->AD->Inference: ok
# forw is called, @diff=false, training=true: forw->AD->Training: We are not taking gradients but user is forcing training mode.
# forw is called, @diff=true, training=false: forw->AD->Inference: User wants to force inference mode?
# forw is called, @diff=true, training=true: Standard training mode. forw->AD->Training: ok


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    CookieCutter(master, others)

A cookie-cutter simulation solver.

## Parameters

* `master` - Master simulation solver (a.k.a. facies solver)
* `others` - Solvers for each value of the master realization

## Examples

Simulate lithofacies with image quilting and fill porosity
values with LU Gaussian simulation:

```julia
julia> f  = IQ(:facies => (TI=IMG, template=(30,30,1)))
julia> p₀ = LUGS(:poro => (variogram=SphericalVariogram(range=10.),))
julia> p₁ = LUGS(:poro => (variogram=SphericalVariogram(range=20.),))
julia> CookieCutter(f, Dict(0=>p₀, 1=>p₁))
```
"""
struct CookieCutter <: SimulationSolver
  master::SimulationSolver
  others::Dict{Number,SimulationSolver}
end

function solve(problem::SimulationProblem, solver::CookieCutter)
  # retrieve problem info
  pdata   = data(problem)
  pdomain = domain(problem)
  preals  = nreals(problem)

  # master variable
  pvars = Dict(name(v) => mactype(v) for v in variables(problem))
  mvars = targets(solver.master)
  @assert length(mvars) == 1 "one single variable must be specified in master solver"
  mvar = mvars[1]
  @assert mvar ∈ keys(pvars) "invalid variable in master solver"
  mtype = pvars[mvar]

  # other variables
  ovars = Tuple(var => V for (var, V) in pvars if var ≠ mvar)
  @assert length(ovars) > 0 "cookie-cutter requires problem with more than one target variable"

  # define master and others problem
  if hasdata(problem)
    mproblem = SimulationProblem(pdata, pdomain, mvar, preals)
    oproblem = SimulationProblem(pdata, pdomain, first.(ovars), preals)
  else
    mproblem = SimulationProblem(pdomain, mvar => mtype, preals)
    oproblem = SimulationProblem(pdomain, ovars, preals)
  end

  # realizations of master variable
  msol = solve(mproblem, solver.master)
  mreals = msol[mvar]

  # pre-allocate memory for realizations
  reals = Dict(var => [Vector{V}(undef, nelements(pdomain)) for i in 1:preals] for (var,V) in ovars)

  # solve other problems
  for (mval, osolver) in solver.others
    osol = solve(oproblem, osolver)

    # cookie-cutter step
    for (var, V) in ovars
      vreals = osol[var]
      for i in 1:preals
        mask = mreals[i] .== mval
        reals[var][i][mask] .= vreals[i][mask]
      end
    end
  end

  realizations = merge(reals, Dict(mvar => mreals))

  Ensemble(pdomain, realizations)
end

# -----------
# IO METHODS
# -----------

function Base.show(io::IO, solver::CookieCutter)
  print(io, "CookieCutter")
end

function Base.show(io::IO, ::MIME"text/plain", solver::CookieCutter)
  mvar = targets(solver.master)[1]
  println(io, solver)
  println(io, "  └─", mvar, " ⇨ ", solver.master)
  lines = ["    └─$val ⇨ $osolver" for (val,osolver) in solver.others]
  print(io, join(lines, "\n"))
end


const GQE{C,V} = GenericQuadExpr{C,V}
const _VariableQuadExpr{C} = GenericQuadExpr{C, VariableRef}
const _DecisionQuadExpr{C} = GenericQuadExpr{C, DecisionRef}
const _KnownQuadExpr{C} = GenericQuadExpr{C, KnownRef}
const _VQE = _VariableQuadExpr{Float64}
const _DQE = _DecisionQuadExpr{Float64}
const _KQE = _KnownQuadExpr{Float64}

struct DecisionCrossTerm
    decision::DecisionRef
    variable::VariableRef
end
Base.first(t::DecisionCrossTerm) = t.decision
Base.last(t::DecisionCrossTerm) = t.variable
Base.hash(t::DecisionCrossTerm, h::UInt) = hash(hash(t.decision) + hash(t.variable), h)
function Base.isequal(t1::DecisionCrossTerm, t2::DecisionCrossTerm)
    return (t1.decision == t2.decision && t1.variable == t2.variable)
end

struct KnownCrossTerm{V}
    known::KnownRef
    variable::V
end
Base.first(t::KnownCrossTerm) = t.known
Base.last(t::KnownCrossTerm) = t.variable
KnownVariableCrossTerm = KnownCrossTerm{VariableRef}
KnownDecisionCrossTerm = KnownCrossTerm{DecisionRef}

Base.hash(t::KnownCrossTerm, h::UInt) = hash(hash(t.known) + hash(t.variable), h)
function Base.isequal(t1::KnownCrossTerm, t2::KnownCrossTerm)
    return (t1.known == t2.known && t1.variable == t2.variable)
end

mutable struct DecisionQuadExpr{C} <: JuMP.AbstractJuMPScalar
    variables::GenericQuadExpr{C, VariableRef}
    decisions::GenericQuadExpr{C, DecisionRef}
    knowns::GenericQuadExpr{C, KnownRef}
    cross_terms::OrderedDict{DecisionCrossTerm, C}
    known_variable_terms::OrderedDict{KnownVariableCrossTerm, C}
    known_decision_terms::OrderedDict{KnownDecisionCrossTerm, C}
end
const DQE = DecisionQuadExpr{Float64}

DQE() = zero(DQE{Float64})

mutable struct BilinearDecisionExpr{C} <: JuMP.AbstractJuMPScalar
    variables::GenericAffExpr{C, VariableRef}
    decisions::GenericAffExpr{C, DecisionRef}
    knowns::GenericQuadExpr{C, KnownRef}
    known_variable_terms::OrderedDict{KnownVariableCrossTerm, C}
    known_decision_terms::OrderedDict{KnownDecisionCrossTerm, C}
end
const BDE = BilinearDecisionExpr{Float64}

# Base overrides #
# ========================== #
function Base.iszero(quad::DecisionQuadExpr)
    return iszero(quad.variables) &&
        iszero(quad.decisions) &&
        iszero(quad.knowns)
end
function Base.zero(::Type{DecisionQuadExpr{C}}) where C
    return DecisionQuadExpr{C}(
        zero(_VariableQuadExpr{C}),
        zero(_DecisionQuadExpr{C}),
        zero(_KnownQuadExpr{C}),
        OrderedDict{DecisionCrossTerm, C}(),
        OrderedDict{KnownVariableCrossTerm, C}(),
        OrderedDict{KnownDecisionCrossTerm, C}())
end
function Base.one(::Type{DecisionQuadExpr{C}}) where C
    return DecisionQuadExpr{C}(
        one(_VariableQuadExpr{C}),
        zero(_DecisionQuadExpr{C}),
        zero(_KnownQuadExpr{C}),
        OrderedDict{DecisionCrossTerm, C}(),
        OrderedDict{KnownVariableCrossTerm, C}(),
        OrderedDict{KnownDecisionCrossTerm, C}())
end
Base.zero(quad::DecisionQuadExpr) = zero(typeof(quad))
Base.one(quad::DecisionQuadExpr) =  one(typeof(quad))
function Base.copy(quad::DecisionQuadExpr{C}) where C
    return DecisionQuadExpr{C}(
        copy(quad.variables),
        copy(quad.decisions),
        copy(quad.knowns),
        copy(quad.cross_terms),
        copy(quad.known_variable_terms),
        copy(quad.known_decision_terms))
end
Base.broadcastable(quad::DecisionQuadExpr) = Ref(quad)

function Base.convert(::Type{DecisionQuadExpr{C}}, c::Number) where C
    return convert(DecisionQuadExpr{C}, convert(DecisionAffExpr{C}, c))
end
function Base.convert(::Type{DecisionQuadExpr{C}}, expr::Union{VariableRef, _VariableAffExpr}) where C
    return convert(DecisionQuadExpr{C}, convert(DecisionAffExpr{C}, expr))
end
function Base.convert(::Type{DecisionQuadExpr{C}}, expr::_VariableQuadExpr) where C
    return DecisionQuadExpr{C}(
        expr,
        zero(_DecisionQuadExpr{C}),
        zero(_KnownQuadExpr{C}),
        OrderedDict{DecisionCrossTerm, C}(),
        OrderedDict{KnownVariableCrossTerm, C}(),
        OrderedDict{KnownDecisionCrossTerm, C}())
end
function Base.convert(::Type{DecisionQuadExpr{C}}, expr::Union{DecisionRef, _DecisionAffExpr}) where C
    return convert(DecisionQuadExpr{C}, convert(DecisionAffExpr{C}, expr))
end
function Base.convert(::Type{DecisionQuadExpr{C}}, expr::_DecisionQuadExpr) where C
    return DecisionQuadExpr{C}(
        zero(_VariableQuadExpr{C}),
        expr,
        zero(_KnownQuadExpr{C}),
        OrderedDict{DecisionCrossTerm, C}(),
        OrderedDict{KnownVariableCrossTerm, C}(),
        OrderedDict{KnownDecisionCrossTerm, C}())
end
function Base.convert(::Type{DecisionQuadExpr{C}}, expr::Union{KnownRef, _KnownAffExpr}) where C
    return convert(DecisionQuadExpr{C}, convert(DecisionAffExpr{C}, expr))
end
function Base.convert(::Type{DecisionQuadExpr{C}}, expr::_KnownQuadExpr) where C
    return DecisionQuadExpr{C}(
        zero(_VariableQuadExpr{C}),
        zero(_DecisionQuadExpr{C}),
        expr,
        OrderedDict{DecisionCrossTerm, C}(),
        OrderedDict{KnownVariableCrossTerm, C}(),
        OrderedDict{KnownDecisionCrossTerm, C}())
end
function Base.convert(::Type{DecisionQuadExpr{C}}, aff::DecisionAffExpr{C}) where C
    return DecisionQuadExpr{C}(
        convert(_VariableQuadExpr{C}, aff.variables),
        convert(_DecisionQuadExpr{C}, aff.decisions),
        convert(_KnownQuadExpr{C}, aff.knowns),
        OrderedDict{DecisionCrossTerm, C}(),
        OrderedDict{KnownVariableCrossTerm, C}(),
        OrderedDict{KnownDecisionCrossTerm, C}())
end

function Base.isequal(quad::DecisionQuadExpr{C}, other::DecisionQuadExpr{C}) where C
    return isequal(quad.variables, other.variables) &&
        isequal(quad.decisions, other.decisions) &&
        isequal(quad.knowns, other.knowns) &&
        isequal(quad.cross_terms, other.cross_terms) &&
        isequal(quad.known_variable_terms, other.known_variable_terms) &&
        isequal(quad.known_decision_terms, other.known_decision_terms)
end

Base.hash(quad::DecisionQuadExpr, h::UInt) =
    hash(hash(quad.variables, h),
         hash(quad.decisions, h),
         hash(quad.knowns, h),
         hash(quad.cross_terms, h),
         hash(quad.known_variable_terms, h),
         hash(quad.known_decision_terms, h))

function sizehint!(quad::DecisionQuadExpr, n::Int, ::Type{_VAE})
    Base.sizehint!(quad.variables.aff, n)
end
function sizehint!(quad::DecisionQuadExpr, n::Int, ::Type{_DAE})
    Base.sizehint!(quad.decisions.aff, n)
end
function sizehint!(quad::DecisionQuadExpr, n::Int, ::Type{_KAE})
    Base.sizehint!(quad.knowns.aff, n)
end

# JuMP overrides #
# ========================== #
function JuMP.drop_zeros!(quad::DecisionQuadExpr)
    JuMP.drop_zeros!(quad.variables)
    JuMP.drop_zeros!(quad.decisions)
    JuMP.drop_zeros!(quad.knowns)
    dropzeros!(terms) = begin
        for (key, coef) in terms
            if iszero(coef)
                delete!(terms, key)
            end
        end
    end
    dropzeros!(quad.cross_terms)
    dropzeros!(quad.known_variable_terms)
    dropzeros!(quad.known_decision_terms)
    return nothing
end

function JuMP._affine_coefficient(f::DecisionQuadExpr, variable::VariableRef)
    return JuMP._affine_coefficient(f.variables, variable)
end
function JuMP._affine_coefficient(f::DecisionQuadExpr, decision::DecisionRef)
    return JuMP._affine_coefficient(f.decisions, decision)
end
function JuMP._affine_coefficient(f::DecisionQuadExpr, known::KnownRef)
    return JuMP._affine_coefficient(f.knowns, known)
end

function JuMP.map_coefficients_inplace!(f::Function, quad::DecisionQuadExpr)
    JuMP.map_coefficients_inplace!(f, quad.variables)
    JuMP.map_coefficients_inplace!(f, quad.decisions)
    JuMP.map_coefficients_inplace!(f, quad.knowns)
    _map_cross_terms!(quad.cross_terms)
    _map_cross_terms!(quad.known_variable_terms)
    _map_cross_terms!(quad.known_decision_terms)
    return quad
end

function _map_cross_terms!(f::Function, terms)
    for (key, value) in terms
        terms[key] = f(value)
    end
    return nothing
end

function _map_cross_terms(f::Function, terms)
    res = copy(terms)
    _map_cross_terms!(f, res)
    return res
end

function JuMP.map_coefficients(f::Function, quad::DecisionQuadExpr)
    return JuMP.map_coefficients_inplace!(f, copy(quad))
end

function JuMP.value(quad::DecisionQuadExpr, value_func::Function)
    varvalue = JuMP.value(quad.variables, value_func)
    dvarvalue = JuMP.value(quad.decisions, value_func)
    kvarvalue = JuMP.value(quad.knowns, value_func)
    T = promote_type(typeof(varvalue), typeof(dvarvalue), typeof(kvarvalue))
    ret = convert(T, varvalue) + convert(T, dvarvalue) + convert(T, kvarvalue)
    val(terms, T) = begin
        ret = zero(T)
        for (vars, coef) in quad.cross_terms
            ret += coef * value_func(vars.a) * value_func(vars.b)
        end
        ret
    end
    ret += val(quad.cross_terms, T)
    ret += val(quad.known_variable_terms, T)
    ret += val(quad.known_decision_terms, T)
    return ret
end

function JuMP.constant(quad::DecisionQuadExpr)
    return quad.variables.constant
end

function JuMP.linear_terms(quad::DecisionQuadExpr, ::Type{_VAE})
    JuMP.linear_terms(quad.variables)
end
function JuMP.linear_terms(quad::DecisionQuadExpr, ::Type{_DAE})
    JuMP.linear_terms(quad.decisions)
end
function JuMP.linear_terms(quad::DecisionQuadExpr, ::Type{_KAE})
    JuMP.linear_terms(quad.knowns)
end

function SparseArrays.dropzeros(quad::DecisionQuadExpr{C}) where C
    variables = SparseArrays.dropzeros(quad.variables)
    decisions = SparseArrays.dropzeros(quad.decisions)
    knowns = SparseArrays.dropzeros(quad.knowns)
    dropzeros!(terms) = begin
        for (key, value) in terms
            if iszero(value)
                delete!(terms, key)
            end
        end
    end
    cross_terms = copy(quad.cross_terms)
    dropzeros!(cross_terms)
    known_variable_terms = copy(quad.known_variable_terms)
    dropzeros!(known_variable_terms)
    known_decision_terms = copy(quad.known_decision_terms)
    dropzeros!(known_decision_terms)
    return DecisionQuadExpr(variables, decisions, knowns, cross_terms, known_variable_terms, known_decision_terms)
end

function JuMP._assert_isfinite(quad::Union{_DecisionQuadExpr, _KnownQuadExpr})
    JuMP._assert_isfinite(quad.aff)
    for (coef, var1, var2) in quad_terms(quad)
        isfinite(coef) || error("Invalid coefficient $coef on quadratic term $var1*$var2.")
    end
    return nothing
end

function JuMP._assert_isfinite(quad::DecisionQuadExpr)
    JuMP._assert_isfinite(quad.variables)
    JuMP._assert_isfinite(quad.decisions)
    JuMP._assert_isfinite(quad.knowns)
    _assert_isfinite(terms) = begin
        for (vars, coef) in terms
            isfinite(coef) || error("Invalid coefficient $coef on quadratic cross term $first(vars)*$last(vars).")
        end
    end
    _assert_isfinite(quad.cross_terms)
    _assert_isfinite(quad.known_variable_terms)
    _assert_isfinite(quad.known_decision_terms)
    return nothing
end

function JuMP.check_belongs_to_model(quad::DecisionQuadExpr, model::AbstractModel)
    JuMP.check_belongs_to_model(quad.variables, model)
    JuMP.check_belongs_to_model(quad.decisions, model)
    JuMP.check_belongs_to_model(quad.knowns, model)
    check_belongs_to_model(terms) = begin
        for variable_pair in keys(terms)
            JuMP.check_belongs_to_model(first(variable_pair), model)
            JuMP.check_belongs_to_model(last(variable_pair), model)
        end
    end
    check_belongs_to_model(quad.cross_terms)
    check_belongs_to_model(quad.known_variable_terms)
    check_belongs_to_model(quad.known_decision_terms)
    return nothing
end

function JuMP.function_string(mode, quad::DecisionQuadExpr, show_constant=true)
    # Decisions
    ret = ""
    decision_terms = JuMP.function_string(mode, quad.decisions)
    if decision_terms != "0"
        ret = decision_terms
    end
    known_decision_terms = _cross_terms_function_string(mode, quad.known_decision_terms)
    ret = _add_terms(ret, known_decision_terms)
    cross_terms = _cross_terms_function_string(mode, quad.cross_terms)
    ret = _add_terms(ret, cross_terms)
    # Knowns
    known_terms = JuMP.function_string(mode, quad.knowns)
    ret = _add_terms(ret, known_terms)
    known_variable_terms = _cross_terms_function_string(mode, quad.known_variable_terms)
    ret = _add_terms(ret, known_variable_terms)
    # Variables
    variable_terms = JuMP.function_string(mode, quad.variables)
    if variable_terms != "0"
        if ret == ""
            ret = variable_terms
        else
            if variable_terms[1] == '-'
                ret = ret * " - " * variable_terms[2:end]
            else
                ret = ret * " + " * variable_terms
            end
        end
    else
        if ret == ""
            return variable_terms
        end
    end
    return ret
end

function _cross_terms_function_string(mode, terms)
    term_str = Array{String}(undef, 2 * length(terms))
    elm = 1
    if length(term_str) > 0
        for (vars, coef) in terms
            pre = JuMP._is_one_for_printing(coef) ? "" : JuMP._string_round(abs(coef)) * " "

            x = function_string(mode, first(vars))
            y = function_string(mode, last(vars))

            term_str[2 * elm - 1] = JuMP._sign_string(coef)
            term_str[2 * elm] = "$pre$x"
            term_str[2 * elm] *= string(JuMP._math_symbol(mode, :times), y)
            if elm == 1
                # Correction for first term as there is no space
                # between - and variable coefficient/name
                term_str[1] = coef < zero(coef) ? "-" : ""
            end
            elm += 1
        end
    end
    ret = join(term_str[1 : 2 * (elm - 1)])
    return ret
end

# With one factor.
function JuMP.add_to_expression!(quad::DQE, other::Number)
    JuMP.add_to_expression!(quad.variables, other)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_var::VariableRef)
    JuMP.add_to_expression!(quad.variables, new_var)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_dvar::DecisionRef)
    JuMP.add_to_expression!(quad.decisions, new_dvar)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_kvar::KnownRef)
    JuMP.add_to_expression!(quad.knowns, new_kvar)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_vae::_VAE)
    JuMP.add_to_expression!(quad.variables, new_vae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_dae::_DAE)
    JuMP.add_to_expression!(quad.decisions, new_dae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_kae::_KAE)
    JuMP.add_to_expression!(quad.knowns, new_kae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_vqe::_VQE)
    JuMP.add_to_expression!(quad.variables, new_vqe)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_dqe::_DQE)
    JuMP.add_to_expression!(quad.decisions, new_dqe)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_kqe::_KQE)
    JuMP.add_to_expression!(quad.knowns, new_kqe)
    return quad
end

function JuMP.add_to_expression!(lhs_quad::DQE, rhs_aff::DAE)
    JuMP.add_to_expression!(lhs_quad.variables, rhs_aff.variables)
    JuMP.add_to_expression!(lhs_quad.decisions, rhs_aff.decisions)
    JuMP.add_to_expression!(lhs_quad.knowns, rhs_aff.knowns)
    return lhs_quad
end

function JuMP.add_to_expression!(lhs_quad::DQE, rhs_quad::DQE)
    JuMP.add_to_expression!(lhs_quad.variables, rhs_quad.variables)
    JuMP.add_to_expression!(lhs_quad.decisions, rhs_quad.decisions)
    JuMP.add_to_expression!(lhs_quad.knowns, rhs_quad.knowns)
    add_cross_terms!(terms, other) = begin
        for (key, term_coef) in other
            JuMP._add_or_set!(terms, key, term_coef)
        end
    end
    add_cross_terms!(lhs_quad.cross_terms, rhs_quad.cross_terms)
    add_cross_terms!(lhs_quad.known_variable_terms, rhs_quad.known_variable_terms)
    add_cross_terms!(lhs_quad.known_decision_terms, rhs_quad.known_decision_terms)
    return lhs_quad
end

# With two factors.
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_var::VariableRef)
    JuMP.add_to_expression!(quad.variables, new_coef, new_var)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_var::VariableRef, new_coef::_Constant)
    JuMP.add_to_expression!(quad.variables, new_coef, new_var)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_dvar::DecisionRef)
    JuMP.add_to_expression!(quad.decisions, new_coef, new_dvar)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_dvar::DecisionRef, new_coef::_Constant)
    JuMP.add_to_expression!(quad.decisions, new_coef, new_dvar)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_kvar::KnownRef)
    JuMP.add_to_expression!(quad.knowns, new_coef, new_kvar)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_kvar::KnownRef, new_coef::_Constant)
    JuMP.add_to_expression!(quad.knowns, new_coef, new_kvar)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_vae::_VAE)
    JuMP.add_to_expression!(quad.variables, new_coef, new_vae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_vae::_VAE, new_coef::_Constant)
    JuMP.add_to_expression!(quad.variables, new_coef, new_vae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_dae::_DAE)
    JuMP.add_to_expression!(quad.decisions, new_coef, new_dae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_dae::_DAE, new_coef::_Constant)
    JuMP.add_to_expression!(quad.decisions, new_coef, new_dae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_kae::_KAE)
    JuMP.add_to_expression!(quad.knowns, new_coef, new_kae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_kae::_KAE, new_coef::_Constant)
    JuMP.add_to_expression!(quad.knowns, new_coef, new_kae)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_vqe::_VQE)
    JuMP.add_to_expression!(quad.variables, new_coef, new_vqe)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_vqe::_VQE, new_coef::_Constant)
    JuMP.add_to_expression!(quad.variables, new_coef, new_vqe)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_dqe::_DQE)
    JuMP.add_to_expression!(quad.decisions, new_coef, new_dqe)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_dqe::_DQE, new_coef::_Constant)
    JuMP.add_to_expression!(quad.decisions, new_coef, new_dqe)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant, new_kqe::_KQE)
    JuMP.add_to_expression!(quad.knowns, new_coef, new_kqe)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, new_kqe::_KQE, new_coef::_Constant)
    JuMP.add_to_expression!(quad.knowns, new_coef, new_kqe)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, coef::_Constant, other::DQE)
    JuMP.add_to_expression!(quad.variables, coef, other.variables)
    JuMP.add_to_expression!(quad.decisions, coef, other.decisions)
    JuMP.add_to_expression!(quad.knowns, coef, other.knowns)
    add_cross_terms!(terms, other) = begin
        for (key, term_coef) in other
            JuMP._add_or_set!(terms, key, coef * term_coef)
        end
    end
    add_cross_terms!(quad.cross_terms, other.cross_terms)
    add_cross_terms!(quad.known_variable_terms, other.known_variable_terms)
    add_cross_terms!(quad.known_decision_terms, other.known_decision_terms)
    return quad
end

function JuMP.add_to_expression!(quad::DQE, other::DQE, coef::_Constant)
    JuMP.add_to_expression!(quad, coef, other)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 var_1::VariableRef,
                                 var_2::Union{DecisionRef, KnownRef})
    return JuMP.add_to_expression!(quad, 1.0, var_1, var_2)
end

function JuMP.add_to_expression!(quad::DQE,
                                 var_1::VariableRef,
                                 var_2::KnownRef)
    return JuMP.add_to_expression!(quad, 1.0, var_1, var_2)
end

function JuMP.add_to_expression!(quad::DQE,
                                 var_1::DecisionRef,
                                 var_2::Union{VariableRef,DecisionRef,KnownRef})
    return JuMP.add_to_expression!(quad, 1.0, var_1, var_2)
end

function JuMP.add_to_expression!(quad::DQE,
                                 var_1::KnownRef,
                                 var_2::Union{VariableRef,DecisionRef,KnownRef})
    return JuMP.add_to_expression!(quad, 1.0, var_1, var_2)
end

function JuMP.add_to_expression!(quad::DQE,
                                 var::VariableRef,
                                 aff::_VAE)
    JuMP.add_to_expression!(quad.variables, var, aff)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 dvar::DecisionRef,
                                 aff::_VAE)
    for (coef, term_var) in linear_terms(aff)
        key = DecisionCrossTerm(dvar, term_var)
        JuMP._add_or_set!(quad.cross_terms, key, coef)
    end
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 kvar::KnownRef,
                                 aff::_VAE)
    for (coef, term_var) in linear_terms(aff)
        key = KnownVariableCrossTerm(kvar, term_var)
        JuMP._add_or_set!(quad.known_variable_terms, key, coef)
    end
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 var::VariableRef,
                                 aff::_DAE)
    for (coef, term_dvar) in linear_terms(aff)
        key = DecisionCrossTerm(term_dvar, var)
        JuMP._add_or_set!(quad.cross_terms, key, coef)
    end
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 dvar::DecisionRef,
                                 aff::_DAE)
    JuMP.add_to_expression!(quad.decisions, dvar, aff)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 kvar::KnownRef,
                                 aff::_DAE)
    for (coef, term_dvar) in linear_terms(aff)
        key = KnownDecisionCrossTerm(kvar, term_dvar)
        JuMP._add_or_set!(quad.known_decision_terms, key, coef)
    end
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 var::VariableRef,
                                 aff::_KAE)
    for (coef, term_kvar) in linear_terms(aff)
        key = KnownVariableCrossTerm(term_kvar, var)
        JuMP._add_or_set!(quad.known_variable_terms, key, coef)
    end
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 dvar::DecisionRef,
                                 aff::_KAE)
    for (coef, term_kvar) in linear_terms(aff)
        key = KnownDecisionCrossTerm(term_kvar, dvar)
        JuMP._add_or_set!(quad.known_decision_terms, key, coef)
    end
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 kvar::KnownRef,
                                 aff::_KAE)
    JuMP.add_to_expression!(quad.knowns, kvar, aff)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 var::Union{VariableRef, DecisionRef, KnownRef},
                                 aff::_DAE)
    JuMP.add_to_expression!(quad, var, aff.variables)
    JuMP.add_to_expression!(quad, var, aff.decisions)
    JuMP.add_to_expression!(quad, var, aff.knowns)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 aff::Union{_VAE, _DAE, _KAE, DAE},
                                 var::Union{VariableRef, DecisionRef, KnownRef})
    return JuMP.add_to_expression!(quad, var, aff)
end

function JuMP.add_to_expression!(quad::DQE,
                                 lhs::_VAE,
                                 rhs::_VAE)
    JuMP.add_to_expression!(quad.variables, lhs, rhs)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 lhs::_DAE,
                                 rhs::_DAE)
    JuMP.add_to_expression!(quad.decisions, lhs, rhs)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 lhs::_KAE,
                                 rhs::_KAE)
    JuMP.add_to_expression!(quad.knowns, lhs, rhs)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 lhs::Union{_VAE,_DAE,_KAE},
                                 rhs::Union{_VAE,_DAE,_KAE})
    lhs_length = length(linear_terms(lhs))
    rhs_length = length(linear_terms(rhs))

    # Quadratic terms
    for (lhscoef, lhsvar) in linear_terms(lhs)
        for (rhscoef, rhsvar) in linear_terms(rhs)
            add_to_expression!(quad, lhscoef*rhscoef, lhsvar, rhsvar)
        end
    end

    # Try to preallocate space for aff
    cur = length(linear_terms(quad, typeof(rhs)))
    if !iszero(lhs.constant)
        sizehint!(quad, cur + lhs_length + rhs_length, typeof(rhs))
    end
    if !iszero(rhs.constant)
        sizehint!(quad, cur + lhs_length, typeof(lhs))
    end

    # [LHS constant] * [RHS linear terms]
    if !iszero(lhs.constant)
        c = lhs.constant
        for (rhscoef, rhsvar) in linear_terms(rhs)
            if rhs isa _VAE
                add_to_expression!(quad.variables.aff, c*rhscoef, rhsvar)
            elseif rhs isa _DAE
                add_to_expression!(quad.decisions.aff, c*rhscoef, rhsvar)
            elseif rhs isa _KAE
                add_to_expression!(quad.knowns.aff, c*rhscoef, rhsvar)
            end
        end
    end

    # [RHS constant] * [LHS linear terms]
    if !iszero(rhs.constant)
        c = rhs.constant
        for (lhscoef, lhsvar) in linear_terms(lhs)
            if lhs isa _VAE
                add_to_expression!(quad.variables.aff, c*lhscoef, lhsvar)
            elseif lhs isa _DAE
                add_to_expression!(quad.decisions.aff, c*lhscoef, lhsvar)
            elseif lhs isa _KAE
                add_to_expression!(quad.knowns.aff, c*lhscoef, lhsvar)
            end
        end
    end

    quad.variables.aff.constant += lhs.constant * rhs.constant

    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 lhs::Union{_VAE, _DAE, _KAE},
                                 rhs::DAE)
    JuMP.add_to_expression!(quad, lhs, rhs.variables)
    JuMP.add_to_expression!(quad, lhs, rhs.decisions)
    JuMP.add_to_expression!(quad, lhs, rhs.knowns)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 lhs::DAE,
                                 rhs::Union{_VAE, _DAE, _KAE})
    JuMP.add_to_expression!(quad, rhs, lhs)
    return quad
end

function JuMP.add_to_expression!(quad::DQE,
                                 lhs::DAE,
                                 rhs::DAE)
    # Variables
    JuMP.add_to_expression!(quad, lhs.variables, rhs.variables)
    JuMP.add_to_expression!(quad, lhs.variables, rhs.decisions)
    JuMP.add_to_expression!(quad, lhs.variables, rhs.knowns)
    # Decisions
    JuMP.add_to_expression!(quad, lhs.decisions, rhs.variables)
    JuMP.add_to_expression!(quad, lhs.decisions, rhs.decisions)
    JuMP.add_to_expression!(quad, lhs.decisions, rhs.knowns)
    # Knowns
    JuMP.add_to_expression!(quad, lhs.knowns, rhs.variables)
    JuMP.add_to_expression!(quad, lhs.knowns, rhs.decisions)
    JuMP.add_to_expression!(quad, lhs.knowns, rhs.knowns)
    return quad
end

# With three factors.
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::VariableRef,
                                 new_var2::VariableRef)
    JuMP.add_to_expression!(quad.variables, new_coef, new_var1, new_var2)
    return quad
end
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::VariableRef,
                                 new_var2::DecisionRef)
    key = DecisionCrossTerm(new_var2, new_var1)
    JuMP._add_or_set!(quad.cross_terms, key, new_coef)
    return quad
end
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::VariableRef,
                                 new_var2::KnownRef)
    key = KnownVariableCrossTerm(new_var2, new_var1)
    JuMP._add_or_set!(quad.known_variable_terms, key, new_coef)
    return quad
end
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::DecisionRef,
                                 new_var2::VariableRef)
    key = DecisionCrossTerm(new_var1, new_var2)
    JuMP._add_or_set!(quad.cross_terms, key, new_coef)
    return quad
end
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::DecisionRef,
                                 new_var2::DecisionRef)
    JuMP.add_to_expression!(quad.decisions, new_coef, new_var1, new_var2)
    return quad
end
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::DecisionRef,
                                 new_var2::KnownRef)
    key = KnownDecisionCrossTerm(new_var2, new_var1)
    JuMP._add_or_set!(quad.known_decision_terms, key, new_coef)
    return quad
end
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::KnownRef,
                                 new_var2::VariableRef)
    key = KnownVariableCrossTerm(new_var1, new_var2)
    JuMP._add_or_set!(quad.known_variable_terms, key, new_coef)
    return quad
end
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::KnownRef,
                                 new_var2::DecisionRef)
    key = KnownDecisionCrossTerm(new_var1, new_var2)
    JuMP._add_or_set!(quad.known_decision_terms, key, new_coef)
    return quad
end
function JuMP.add_to_expression!(quad::DQE, new_coef::_Constant,
                                 new_var1::KnownRef,
                                 new_var2::KnownRef)
    JuMP.add_to_expression!(quad.knowns, new_coef, new_var1, new_var2)
    return quad
end


#Contain the Zone_target structure to load hold and manipulate the targets in a zone
"""loads the targets in the zone and keep track of whether a newzone in needed to be loaded
"""
struct Zone_Target
    #Loads the targets relvant for a particular zone
    #Any thing stored in two element array will always asume body first and head second
    max_target::Integer #maximum number of target
    max_usp::Integer #total maximum number of unique priority for the full tile

    zones::Array{Integer,1}
    #first entry for actual number of target and second entry for number of unique priority
    ntarget::Array{Integer,1}

    index::Array{Integer,1}
    X_FP::Array{Float64,1} #To store the x-position on focal plane of target
    Y_FP::Array{Float64,1} #To store y-position on focal plane of target
    PRIORITY::Array{Integer,1}

    #boolean array
    pot_ass::Array{Bool,1} #To set for the potential assignment
    pot_coll::Array{Bool,1} #To set for potential collision

    #To store theta,phi for a particular positioner
    theta_pos::Array{Float64,1}
    phi_pos::Array{Float64,1}

    #unique sorted priority for this tile
    usp::Array{Integer,1}
    usp_bool_sel::Array{Integer,2} #To store the boolean selection of each priority

    #To store the possible assignments
    possible_ass::Array{Int32,1}

    function  Zone_Target(maximum_target::Integer,nunique_priority::Integer)
        new(maximum_target,nunique_priority,zeros(3),zeros(2),zeros(maximum_target),zeros(maximum_target),zeros(maximum_target),
        zeros(maximum_target),zeros(Bool,maximum_target),zeros(Bool,maximum_target),
        zeros(maximum_target),zeros(maximum_target),
        zeros(nunique_priority),zeros(Bool,maximum_target,nunique_priority),
        zeros(Int32,maximum_target))
    end
end

function load_target_zone(zone_target::Zone_Target,zone_x,zone_y,fp_bound,targets_dic)
    #This zone is already loaded so nothing to be done
    if((zone_target.zones[1]==zone_x) & (zone_target.zones[2]==zone_y))
        return
    else
        zone_target.zones[1]=zone_x
        zone_target.zones[2]=zone_y
        zone_target.zones[3]=get_zonexy(zone_x,zone_y,fp_bound)
    end

    #indsel=zeros(Int32,1000)

    count_targ=1
    for zx in maximum([zone_x-1,1]):minimum([zone_x+1,fp_bound.nx])
        for zy in maximum([zone_y-1,1]):minimum([zone_y+1,fp_bound.ny])
            zone_xy=get_zonexy(zx,zy,fp_bound)

            if(!(zone_xy in targets_dic["zone_map_keys"]))
                continue
            end
            imin=targets_dic["zone_map"][zone_xy][1]
            imax=targets_dic["zone_map"][zone_xy][2]
            #@show zx,zy,zone_xy,targets_dic["zone_map"][zone_xy]
            #@show targets_dic["zone_x"][imin:imax]
            count_this=imax-imin
            zone_target.index[count_targ:count_targ+count_this] .= collect(imin:imax)
            count_targ=count_targ+count_this+1
        end
    end

    #set everything else to 0
    if(count_targ<=zone_target.max_target)
        zone_target.index[count_targ:end] .= 0
    end
    #Due to Julia indexing convetion 1 need to be subtracte
    count_targ=count_targ-1

    tmp_ind=view(zone_target.index,1:count_targ)
    #sort the indices
    sort!(tmp_ind)

    zone_target.ntarget[1]=count_targ

    zone_target.X_FP[1:count_targ]=targets_dic["X_FP"][tmp_ind]
    zone_target.Y_FP[1:count_targ]=targets_dic["Y_FP"][tmp_ind]
    zone_target.PRIORITY[1:count_targ]=targets_dic["PRIORITY"][tmp_ind]

    us_priority=sort(unique(zone_target.PRIORITY[1:zone_target.ntarget[1]]),rev=true)
    nup=size(us_priority,1)
    zone_target.ntarget[2]=nup
    zone_target.usp[1:nup] .= us_priority

    imin=1;imax=zone_target.ntarget[1]
    for uu in 1:zone_target.ntarget[2]
        zone_target.usp_bool_sel[imin:imax,uu] .= (zone_target.PRIORITY[imin:imax].== us_priority[uu])
    end

end


include("../../validation/loocv.jl")

"""
Computes logdet hat(ΣθinvX) given logdet(ΣθinvX), where the hat indicates that
a new point has been incorporated.
"""
function extendlogdetΣθ(old_ΣθinvX, new_choleskyΣθ, new_X)
    lin_sys_loocv()
end


using DataFrames
using DataFramesMeta
using LinearAlgebra
using NearestNeighbors
using Statistics
using StatsBase: countmap

mutable struct BmmData{L}

    # Static data
    ## Segmentation data
    x::DataFrame;
    position_data::Matrix{Float64};
    composition_data::Vector{<:Union{Int, Missing}};
    confidence::Vector{Float64}

    cluster_per_molecule::Vector{Int}
    segment_per_molecule::Vector{Int}
    nuclei_prob_per_molecule::Vector{Float64}

    ## MRF
    adjacent_points::Array{Vector{Int}, 1};
    adjacent_weights::Array{Vector{Float64}, 1};
    real_edge_weight::Float64;
    position_knn_tree::KDTree;
    knn_neighbors::Array{Vector{Int}, 1};

    # Distribution-related
    components::Array{Component{L}, 1};
    distribution_sampler::Component{L};
    assignment::Vector{Int};
    max_component_guid::Int;

    noise_density::Float64;

    cluster_per_cell::Vector{Int};

    # Prior segmentation

    n_molecules_per_segment::Vector{Int}
    main_segment_per_cell::Vector{Int}

    # Utils
    tracer::Dict{Symbol, Any}
    misc::Dict{Symbol, Any}
    center_sample_cache::Vector{Int}

    # Parameters
    prior_seg_confidence::Float64
    cluster_penalty_mult::Float64
    use_gene_smoothing::Bool
    min_nuclei_frac::Float64
    mrf_strength::Float64

    """
    ...
    # Arguments
    - `components::Array{Component, 1}`:
    - `x::DataFrame`:
    - `adjacent_points::Array{Array{Int, 1}, 1}`:
    - `adjacent_weights::Array{Array{Float64, 1}, 1}`: edge weights, used for smoothness penalty
    - `real_edge_weight::Float64`: weight of an edge for "average" real point
    - `distribution_sampler::Component`:
    - `assignment::Array{Int, 1}`:
    """
    function BmmData(components::Array{Component{N}, 1}, x::DataFrame, adjacent_points::Array{Vector{Int}, 1}, adjacent_weights::Array{Vector{Float64}, 1},
                     assignment::Vector{Int}, distribution_sampler::Component{N}; real_edge_weight::Float64=1.0, k_neighbors::Int=20,
                     cluster_penalty_mult::Float64=0.25, use_gene_smoothing::Bool=true, prior_seg_confidence::Float64=0.5,
                     min_nuclei_frac::Float64=0.1, mrf_strength::Float64=0.1, na_genes::Vector{Int}=Int[]) where N
        @assert maximum(assignment) <= length(components)
        @assert minimum(assignment) >= 0
        @assert length(assignment) == size(x, 1)

        if !all(s in propertynames(x) for s in [:x, :y, :gene])
            error("`x` data frame must have columns 'x', 'y' and 'gene'")
        end

        p_data = position_data(x)
        @assert size(p_data, 1) == N

        position_knn_tree = KDTree(p_data)
        knn_neighbors = knn(position_knn_tree, p_data, k_neighbors)[1]

        comp_data = [ifelse(g in na_genes, missing, g) for g in composition_data(x)]
        n_genes = maximum(skipmissing(composition_data(x)))

        x = deepcopy(x)
        if !(:confidence in propertynames(x))
            x[!, :confidence] .= 0.95
        end

        nuclei_probs = :nuclei_probs in propertynames(x) ? x[!, :nuclei_probs] : Float64[]
        cluster_per_molecule = :cluster in propertynames(x) ? x.cluster : Int[]
        self = new{N}(
            x, p_data, comp_data, confidence(x), cluster_per_molecule, Int[], nuclei_probs,
            adjacent_points, adjacent_weights, real_edge_weight, position_knn_tree, knn_neighbors,
            components, deepcopy(distribution_sampler), assignment, length(components), 0.0,
            Int[],
            Int[], Int[], # prior segmentation info
            Dict{Symbol, Any}(), Dict{Symbol, Any}(), Int[],
            prior_seg_confidence, cluster_penalty_mult, use_gene_smoothing, min_nuclei_frac, mrf_strength
        )

        for c in self.components
            c.n_samples = 0
        end

        for c_id in assignment[assignment .> 0]
            self.components[c_id].n_samples += 1
        end

        # Resolve component guids
        guids = [c.guid for c in self.components]

        if maximum(guids) <= 0
            self.max_component_guid = length(self.components)
            for (i,c) in enumerate(self.components)
                c.guid = i
            end
        else
            if minimum(guids) <= 0
                error("Either all or no guids can be <= 0")
            end

            self.max_component_guid = maximum(guids)
        end

        if :prior_segmentation in propertynames(x)
            self.segment_per_molecule = deepcopy(x.prior_segmentation);
            self.n_molecules_per_segment = count_array(self.segment_per_molecule, drop_zero=true);
            update_n_mols_per_segment!(self);
        end

        return self
    end
end

function position_data(df::AbstractDataFrame)::Matrix{Float64}
    if (:z in propertynames(df))
        return copy(Matrix{Float64}(df[:, [:x, :y, :z]])')
    end

    return copy(Matrix{Float64}(df[:, [:x, :y]])')
end
position_data(data::BmmData)::Matrix{Float64} = data.position_data
composition_data(df::AbstractDataFrame)::Union{Vector{Int}, Vector{Union{Missing, Int}}} = df.gene
composition_data(data::BmmData) = data.composition_data
confidence(df::AbstractDataFrame)::Vector{Float64} = df.confidence
confidence(data::BmmData)::Vector{Float64} = data.confidence

num_of_molecules_per_cell(data::BmmData) = count_array(data.assignment, max_value=length(data.components), drop_zero=true)

function assign!(data::BmmData, point_ind::Int, component_id::Int)
    old_id = data.assignment[point_ind]
    if old_id == component_id
        return
    end

    @assert component_id <= length(data.components) "Too large component id: $component_id, maximum available: $(length(data.components))"

    segment_id = isempty(data.segment_per_molecule) ? 0 : data.segment_per_molecule[point_ind]
    if segment_id > 0
        if component_id > 0
            data.components[component_id].n_molecules_per_segment[segment_id] = get(data.components[component_id].n_molecules_per_segment, segment_id, 0) + 1
        end

        if old_id > 0
            data.components[old_id].n_molecules_per_segment[segment_id] -= 1
        end
    end

    data.assignment[point_ind] = component_id
end

function estimate_assignment_by_history(data::BmmData)
    # TODO: it doesn't guarantee connectivity. Can try to run deterministic EM, or use some better estimate here
    if !(:assignment_history in keys(data.tracer)) || (length(data.tracer[:assignment_history]) == 0)
        @warn "Data has no saved history of assignments. Fall back to the basic assignment"
        return data.assignment, ones(length(data.assignment)) / 2
    end

    guid_map = Dict(c.guid => i for (i,c) in enumerate(data.components))
    current_guids = Set(vcat(collect(keys(guid_map)), [0]));
    assignment_mat = hcat(data.tracer[:assignment_history]...);

    reassignment = mapslices(assignment_mat, dims=2) do row
        c_row = row[in.(row, Ref(current_guids))]
        if length(c_row) == 0
            return 0
        end

        c_counts = countmap(c_row);
        count_vals = collect(values(c_counts))
        return maximum(collect(keys(c_counts))[count_vals .== maximum(count_vals)])
    end

    return get.(Ref(guid_map), vec(reassignment), 0), vec(mean(assignment_mat .== reassignment, dims=2))
end

function get_cell_qc_df(segmented_df::DataFrame, cell_assignment::Vector{Int}=segmented_df.cell; sigdigits::Int=4, max_cell::Int=maximum(cell_assignment), dapi_arr::Union{Matrix{<:Real}, Nothing}=nothing)
    seg_df_per_cell = split(segmented_df, cell_assignment .+ 1; max_factor=max_cell+1)[2:end];
    pos_data_per_cell = [position_data(df)[1:2,:] for df in seg_df_per_cell];

    df = DataFrame(:n_transcripts => size.(pos_data_per_cell, 2));
    large_cell_mask = (df.n_transcripts .> 2)

    df[!,:density] = fill(NaN, size(df, 1))
    df[!,:elongation] = fill(NaN, size(df, 1))
    df[!,:area] = fill(NaN, size(df, 1))

    df.area[large_cell_mask] = round.(area.(convex_hull.(pos_data_per_cell[large_cell_mask])), sigdigits=sigdigits);
    df.density[large_cell_mask] = round.(df.n_transcripts[large_cell_mask] ./ df.area[large_cell_mask], sigdigits=sigdigits);
    df.elongation[large_cell_mask] = [round(x[2] / x[1], sigdigits=sigdigits) for x in eigvals.(cov.(transpose.(pos_data_per_cell[large_cell_mask])))];
    if :confidence in propertynames(segmented_df)
        df[!,:avg_confidence] = round.([mean(df.confidence) for df in seg_df_per_cell], sigdigits=sigdigits)
    end

    if dapi_arr !== nothing # TODO: forward DAPI from CLI
        brightness_per_mol = staining_value_per_transcript(segmented_df, dapi_arr);
        df[!, :mean_dapi_brightness] = mean.(trim.(split(brightness_per_mol, cell_assignment .+ 1)[2:end]; prop=0.2))
    end

    return df
end

function get_cell_stat_df(data::BmmData, segmented_df::Union{DataFrame, Nothing}=nothing; add_qc::Bool=true, sigdigits::Int=4)
    df = DataFrame(:cell => 1:length(data.components))

    centers = hcat([Vector(c.position_params.μ) for c in data.components]...)

    for s in [:x, :y]
        df[!,s] = mean.(split(data.x[!,s], data.assignment .+ 1, max_factor=length(data.components) + 1)[2:end])
    end

    if !isempty(data.cluster_per_cell)
        df[!, :cluster] = data.cluster_per_cell
    end

    if add_qc
        if segmented_df === nothing
            segmented_df = get_segmentation_df(data);
        end

        df = hcat(df, get_cell_qc_df(segmented_df; sigdigits=sigdigits, max_cell=length(data.components)))
    end

    return df[num_of_molecules_per_cell(data) .> 0,:]
end

function get_segmentation_df(data::BmmData, gene_names::Union{Nothing, Array{String, 1}}=nothing; use_assignment_history::Bool=true)
    df = deepcopy(data.x)
    df[!,:cell] = deepcopy(data.assignment);

    if use_assignment_history && (:assignment_history in keys(data.tracer)) && (length(data.tracer[:assignment_history]) > 1)
        df[!,:cell], df[!,:assignment_confidence] = estimate_assignment_by_history(data)
        df.assignment_confidence .= round.(df.assignment_confidence, digits=5)
    end

    if :confidence in propertynames(df)
        df.confidence = round.(df.confidence, digits=5)
    end

    df[!,:is_noise] = (df.cell .== 0);

    if gene_names !== nothing
        df[!,:gene] = gene_names[df[!,:gene]]
    end

    if !isempty(data.cluster_per_molecule)
        df[!, :cluster] = data.cluster_per_molecule
    end

    return df
end

function global_assignment_ids(data::BmmData)::Vector{Int}
    cur_guids = [c.guid for c in data.components]
    res = deepcopy(data.assignment)
    non_noise_mask = (res .> 0)
    res[non_noise_mask] .= cur_guids[res[non_noise_mask]]

    return res
end

function update_n_mols_per_segment!(bm_data::BmmData)
    if isempty(bm_data.segment_per_molecule)
        return
    end

    # Estimate number of molecules per segment per component
    for comp in bm_data.components
        empty!(comp.n_molecules_per_segment)
    end

    for i in 1:length(bm_data.assignment)
        c_cell = bm_data.assignment[i]
        c_seg = bm_data.segment_per_molecule[i]
        if (c_cell == 0) || (c_seg == 0)
            continue
        end

        bm_data.components[c_cell].n_molecules_per_segment[c_seg] = get(bm_data.components[c_cell].n_molecules_per_segment, c_seg, 0) + 1
    end

    # Estimate the main segment per cell
    resize!(bm_data.main_segment_per_cell, length(bm_data.components))
    bm_data.main_segment_per_cell .= 0

    for ci in 1:length(bm_data.components)
        if isempty(bm_data.components[ci].n_molecules_per_segment)
            continue
        end

        i_max = 0
        f_max = 0.0
        for (si,nms) in bm_data.components[ci].n_molecules_per_segment
            seg_size = bm_data.n_molecules_per_segment[si]
            if ((nms / seg_size) > (f_max + 1e-10)) || ((nms == seg_size) && (seg_size > bm_data.n_molecules_per_segment[i_max]))
                f_max = nms / seg_size
                i_max = si
            end
        end
        bm_data.main_segment_per_cell[ci] = i_max
    end
end


#=
  本文件配置个性化参数，具体参见 initAlarmSegparas! 中的注释
=#

function initParaTuple( ;segID::Int,StaticThreshold::Int , DynamicThreshold_CalTimePeriod::UInt,DynamicThresHold_calCFactor::Int,
                        microMergeTimePeriod::Int,
                        microFilterTimePeriod::Int,
                        microMergeSpaceLength::UInt,
                        microFilterSpaceLength::UInt,
                        macroMaxSpaceLength::UInt,
                        macroMinSpaceLength::UInt,
                        macroMaxTimeInterval::UInt,
                        macroMinTimeInterval::UInt,
            )

end

function initAlarmSegparas!(SegAlarmParasDictObj::Dict )   #####################
    SegParaNameTuple1 =  ( segID = 0 , microtime=Seg_MicroTimeParaStruct() ,
    microSpace=Seg_MicroSpaceParaStruct(),
    macroSpace= Seg_MacroSpaceParaStruct(),
    macroTime= Seg_MacroTimeParaStruct(),
    )
    # 静态阈值
    SegParaNameTuple1.microtime.StaticThreshold = Int(2200) # if equal to the typeMAxInt, the Static TH  does not be in  effect
    #动态阈值
    SegParaNameTuple1.microtime.DynamicThreshold_CalTimePeriod = UInt( 60 ) ## in 0.426 seconds  # calculated from Seconds.seconds ,may be any time period ,between 10seconds and the cached data period(about 10 minutes above )
    #= the  TH is: the Calculated Avg * factor / 100
        如果动态阈值不适用该点，则DynamicThresHold_calCFactor = typemax(Int)
    =#
    SegParaNameTuple1.microtime.DynamicThresHold_calCFactor=Int( typemax(Int) )

    #微观过滤和合并 时间
    SegParaNameTuple1.microtime.microMergeTimePeriod  = Int(2) # in 0.426 seconds # calculated from :seconds * 0.426
    SegParaNameTuple1.microtime.microFilterTimePeriod = Int(2) # in 0.426 seconds # same as above

    #微观过滤和合并 空间
    SegParaNameTuple1.microSpace.microMergeSpaceLength = UInt(1) # in Segment or  grain. Calculated from the space length and the segPermeters
    SegParaNameTuple1.microSpace.microFilterSpaceLength = UInt(3) # in Segment   or  grain

    #宏观 空间 限制范围
    SegParaNameTuple1.macroSpace.macroMaxSpaceLength = UInt(100)# in Segment or  grain. Calculated from the space length and the segPermeters
    SegParaNameTuple1.macroSpace.macroMinSpaceLength = UInt(4)# in Segment   or  grain
    #宏观 时间  限制范围
    SegParaNameTuple1.macroTime.macroMaxTimeInterval = UInt(30) # in 0.426 seconds
    SegParaNameTuple1.macroTime.macroMinTimeInterval = UInt(15) # in  0.426 seconds

    SegAlarmParasDictObj[SegParaNameTuple1.segID ] = SegParaNameTuple1
end


#include("./computer.jl")
import AoC2019

mutable struct BoardState
  dims
  my_display::Array{Int8, 2}
  current_score
end

function update_board_state!(state, x, y, arg)
  if x == -1 && y == 0
      state.current_score = arg
  else
    state.my_display[y+1, x+1] = arg
  end
end

function move_controller(state)
  the_ball = findfirst(x -> (x == 4), state.my_display)
  the_pad = findfirst(x -> (x == 3), state.my_display)
  #println(the_ball)
  #println(the_pad)
  if the_ball === nothing || the_pad === nothing
    return 0
  end
  direction = AoC2019.pseudoNormalize(the_ball[2] - the_pad[2])
  #println(stderr, direction)
  return direction
end

function draw_state(state)
  color_map = Dict(0 => ' ', 1 => 'w', 2 => 'b', 3 => '-', 4 => 'o')
  println(AoC2019.visualizeArray(state.my_display, color_map))
end

function arcade_cabinet(input_data)

  game_input = Channel{Int64}(Inf)
  game_output = Channel{Int64}(Inf)
  #dims = (24, 42)
  dims = (60, 60)
  my_display = zeros(Int8, dims...)
  current_score = 0

  state = BoardState((50,50), my_display, current_score)
  input_data[1] = 2
  condition=Condition()

  @sync begin
    my_program = @async AoC2019.run_program_async(input_data, game_input, game_output)

    @async while true
      #println(stderr, "Reading X")

      x = take!(game_output)
      #println(stderr, "Reading Y")
      y = take!(game_output)
      #println(stderr, "Reading Arg")
      arg = take!(game_output)
      #println(stderr, "X: $x, Y: $y, arg: $arg")
      #println(stderr, "X == -1: $(x == -1)")
      update_board_state!(state, x, y, arg)
      #      #display(plot!(state.my_display))
    end
    @async while true
      try
        #wait(condition)
        #println(stderr, "woken up")
        draw_state(state)
        sleep(0.033)
        #println(stderr, "Broken")
        direction = move_controller(state)
        #println(stderr, "Inputting $direction")
        put!(game_input, direction)
        #println(stderr, "Input $direction")
        yield()
      catch
        break
      end
    end
  end

  println("Game finished")
  draw_state(state)

  return state.current_score
end

input_data = AoC2019.read_numbers("input13.txt")
println("Good luck")
out = arcade_cabinet(input_data)
println(out)


# RELOG: Reverse Logistics Optimization
# Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

using JuMP

function resolve(model_old, filename::AbstractString; kwargs...)::OrderedDict
    @info "Reading $filename..."
    instance = RELOG.parsefile(filename)
    return resolve(model_old, instance; kwargs...)
end

function resolve(model_old, instance::Instance; optimizer = nothing)::OrderedDict
    milp_optimizer = lp_optimizer = optimizer
    if optimizer === nothing
        milp_optimizer = _get_default_milp_optimizer()
        lp_optimizer = _get_default_lp_optimizer()
    end

    @info "Building new graph..."
    graph = build_graph(instance)
    _print_graph_stats(instance, graph)

    @info "Building new optimization model..."
    model_new = RELOG.build_model(instance, graph, milp_optimizer)

    @info "Fixing decision variables..."
    _fix_plants!(model_old, model_new)
    JuMP.set_optimizer(model_new, lp_optimizer)

    @info "Optimizing MILP..."
    JuMP.optimize!(model_new)

    if !has_values(model_new)
        @warn("No solution available")
        return OrderedDict()
    end

    @info "Extracting solution..."
    solution = get_solution(model_new, marginal_costs = true)

    return solution
end

function _fix_plants!(model_old, model_new)::Nothing
    T = model_new[:instance].time

    # Fix open_plant variables
    for ((node_old, t), var_old) in model_old[:open_plant]
        value_old = JuMP.value(var_old)
        node_new = model_new[:graph].name_to_process_node_map[(
            node_old.location.plant_name,
            node_old.location.location_name,
        )]
        var_new = model_new[:open_plant][node_new, t]
        JuMP.unset_binary(var_new)
        JuMP.fix(var_new, value_old)
    end

    # Fix is_open variables
    for ((node_old, t), var_old) in model_old[:is_open]
        value_old = JuMP.value(var_old)
        node_new = model_new[:graph].name_to_process_node_map[(
            node_old.location.plant_name,
            node_old.location.location_name,
        )]
        var_new = model_new[:is_open][node_new, t]
        JuMP.unset_binary(var_new)
        JuMP.fix(var_new, value_old)
    end

    # Fix plant capacities
    for ((node_old, t), var_old) in model_old[:capacity]
        value_old = JuMP.value(var_old)
        node_new = model_new[:graph].name_to_process_node_map[(
            node_old.location.plant_name,
            node_old.location.location_name,
        )]
        var_new = model_new[:capacity][node_new, t]
        JuMP.delete_lower_bound(var_new)
        JuMP.delete_upper_bound(var_new)
        JuMP.fix(var_new, value_old)
    end

    # Fix plant expansion
    for ((node_old, t), var_old) in model_old[:expansion]
        value_old = JuMP.value(var_old)
        node_new = model_new[:graph].name_to_process_node_map[(
            node_old.location.plant_name,
            node_old.location.location_name,
        )]
        var_new = model_new[:expansion][node_new, t]
        JuMP.delete_lower_bound(var_new)
        JuMP.delete_upper_bound(var_new)
        JuMP.fix(var_new, value_old)
    end
end


export CreateLayoutTree, getInheritWidth
function getInheritWidth(document::Page, p)
        while p.shape.width < 1
            p == p.parent && (return document.width)
            p = p.parent
        end
        return p.shape.width
end
# width, h = getSize(p.shape)
# return width
# ======================================================================================
## CreateLayoutTree(document::Page, node::Element)
# Generate a layout tree starting at `node` in `document`. This should be callable
# starting at any point in the document tree.
# ======================================================================================, shape::Draw
function CreateLayoutTree(document::Page, parent)

    isa(parent.shape, TextLine) && return # TODO: may be able to make this part of multiple dispatch later

    parent.rows  = []  # Clean out old data (?)
    l,t,w,h = getContentBox(parent.shape)
    w = getInheritWidth(document::Page, parent) # because width may not be set
    children = parent.children

    for child in children
                child.parent = parent # attach to parent
                if isdefined(child, :children) && length(child.children) > 0
                    child.rows = []
                end
                #AtributesToLayout(document, child) # Create Child's DOM (DomToLayout.jl)
            if !child.shape.flags[DisplayNone] # only draw if not "display: none"
                PushToRow(document, child, l,t,w) # Put child into row
                CreateLayoutTree(document, child) # Create child's children
            end
    end

    # Clean-up! Generally the last row of each child is not yet finalized.
    if length(parent.rows) > 0
        lastRow = parent.rows[end]
        FinalizeRow(lastRow)
        # Set content height and width for scroller
        parent.scroll.contentHeight = lastRow.top + lastRow.height - parent.shape.top
        parent.scroll.contentWidth  = lastRow.left + lastRow.width  - parent.shape.left

        if !parent.shape.flags[FixedHeight]
            parent.shape.height = parent.scroll.contentHeight
        end
        # This is to be done after the parent node's size is finalised!
        if parent.shape.flags[HasAbsolute]
          # get node metrics again since the height etc. might have changed.
          l,t,w,h = getContentBox(parent.shape)
          for child in children
                shape = child.shape
                width, height = getSize(shape)
                if shape.flags[Absolute]
                  top,left = shape.top, shape.left # temp variables
                  setDistanceFromBounds(l,t,w,h,  shape )
                  rows = child.rows
                  for row in rows         # row.nodes[i]
                    for n in row.nodes
                      MoveAll(n, shape.left - left, shape.top - top)
                    end
                  end
                end
          end
        end
    end
end


module FluxModels

using DocStringExtensions
using Flux
using Flux: @functor
using Parameters
using ModelUtils

abstract type ModuleSpec end

include("./OutputSizes.jl")
using .OutputSizes

include("./activations.jl")
include("./layers.jl")
include("./blocks.jl")
include("./heads.jl")

include("./efficientnet.jl")
include("./mobilenetv3.jl")
include("./xresnet.jl")

export
    outputsize,
    SqueezeExcitation,
    PixelShuffle,

    efficientnetb0,
    mobilenetv3_small,
    mobilenetv3_large,
    xresnet18,
    xresnet50

end # module