@pyimport wntr.network.controls as controls
@pyimport wntr.sim.epanet as sims
function UpdateExtrema(wn::PyCall.PyObject, wn_dict::Dict{Any,Any}, Qmax_dict::Dict{String, Float64},
                        Qmin_dict::Dict{String,Float64}, Hmax_dict::Dict{String, Float64},
                        Hmin_dict::Dict{String,Float64}, num_sim::Int64, sim::Int64,
                        tight_coef::Float64, Q_lb::Float64)
    link_results, node_results = run_sims(wn)
    UpdateFlowRange(wn, wn_dict, link_results, node_results, Qmax_dict, Qmin_dict, num_sim, sim, tight_coef, Q_lb)
    UpdateHeadRange(wn_dict, Hmax_dict, Hmin_dict, link_results, node_results)
end

function run_sims(wn::PyCall.PyObject)
    wns = sims.EpanetSimulator(wn)
    results = wns[:run_sim]()
    link_results = results[:link]
    node_results = results[:node]
    return link_results, node_results
end
function UpdateFlowRange(wn::PyCall.PyObject, wn_dict::Dict{Any,Any},
                    link_results::Dict{Any,Any}, node_results::Dict{Any,Any},
                    Qmax_dict::Dict{String, Float64}, Qmin_dict::Dict{String,Float64},
                    num_sim::Int64, sim::Int64, tight_coef::Float64, Q_lb::Float64)

    pipes = wn_dict["pipes"]
    valves = wn_dict["valves"]
    links = merge(pipes, valves)
    for (key, link) in links
        name = link["name"]
        flows = link_results["flowrate"][name][:values]
        convert(Array{Float64}, flows)
        if sim == 1
            Qmin_dict[name] = convert(Float64,minimum(flows))
            Qmax_dict[name] = convert(Float64, maximum(flows))
        elseif sim < num_sim
            Qmin = Qmin_dict[name]
            Qmax = Qmax_dict[name]

            Qmin_dict[name] = minimum([Qmin, minimum(flows)])
            Qmax_dict[name] = maximum([Qmax, maximum(flows)])
        else
            Qmax_val = Qmax_dict[name]
            Qmin_val = Qmin_dict[name]
            gap = Qmax_val - Qmin_val
            if Qmin_val <= 0
                Qmin_dict[name] = Qmin_val - tight_coef*gap
            else
                Qmin_dict[name] = minimum([Q_lb, Qmin_val])
            end
            if Qmax_val >= 0
                Qmax_dict[name] = Qmax_val + tight_coef*gap
            else
                Qmax_dict[name] = maximum([Qmax_val, -Q_lb])
            end
        end
    end

    for (key, pump) in wn_dict["pumps"]
        name = pump["name"]
        flow = link_results["flowrate"][name][:values]
        convert(Array{Float64}, flow)
        if sim != 1
            flow = flow[flow .> 0]
        end

        if length(flow) !=0 || sim == num_sim
            if sim == 1
                Qmin_dict[name] = minimum(flow)
                Qmax_dict[name]= maximum(flow)
            elseif sim < num_sim
                Qmax = Qmax_dict[name]
                Qmin = Qmin_dict[name]
                Qmin_dict[name] = minimum([Qmin, minimum(flow)])
                Qmax_dict[name] = maximum([Qmax, maximum(flow)])
            else
                Qmin_val = Qmin_dict[name]
                Qmax_val = Qmax_dict[name]

                gap = Qmax_val - Qmin_val
                if pump["pump_type"] == "HEAD"
                    pump_curve_name = pump["pump_curve_name"]
                    pump_curve = wn_dict["curves"][pump_curve_name]["points"]
                else
                    pump_curve = [(pump["power"],0.0)] #power pump types gives fixed power value,
                end

                Qmin_dict[name] = maximum([pump_curve[1][1], Qmin_val - tight_coef*gap])
                Qmax_dict[name] = minimum([pump_curve[length(pump_curve)][1], Qmax_val + tight_coef*gap])
            end
        end
    end
end

function UpdateHeadRange(wn_dict::Dict{Any,Any}, Hmax_dict::Dict{String, Float64},
                        Hmin_dict::Dict{String,Float64}, link_results::Dict{Any,Any},
                        node_results::Dict{Any,Any})
    pumps = wn_dict["pumps"]
    nodes = wn_dict["nodes"]
    for (key, pump) in pumps
        name1 = pump["start_node_name"]
        elev1 = nodes[name1]["elevation"]

        name2 = pump["end_node_name"]
        elev2 = nodes[name2]["elevation"]

        pump_head = (node_results["pressure"][name2][:values] + elev2) - (node_results["pressure"][name1][:values] + elev1)
        pump_status = link_results["status"][key][:values]

        pump_head = pump_head .* pump_status #only care about head when pump is on
        pump_head = pump_head[pump_head .> 0]

        if length(pump_head) > 0
            try
                Hmin_dict[key] = min(minimum(pump_head), Hmin_dict[key])
                Hmax_dict[key] = max(maximum(pump_head), Hmax_dict[key])
            catch
                Hmin_dict[key] = minimum(pump_head)
                Hmax_dict[key] = maximum(pump_head)
            end
        end
    end
    return Hmin_dict, Hmax_dict
end


# ============================================================
# Structs of Control Volume
# Solver stores data in arrays of struct (AoS)
# ============================================================

# ------------------------------------------------------------
# 1D
# ------------------------------------------------------------

"""
    mutable struct ControlVolume1D{A,B} <: AbstractControlVolume1D
        w::A
        prim::B
        sw::A
    end

1D control volume with no distribution function

"""
mutable struct ControlVolume1D{A,B} <: AbstractControlVolume1D
    w::A
    prim::B
    sw::A
end

function ControlVolume1D(W, PRIM)
    w = deepcopy(W)
    prim = deepcopy(PRIM)
    sw = zero(W)

    return ControlVolume1D{typeof(w),typeof(prim)}(w, prim, sw)
end

function Base.show(io::IO, ctr::ControlVolume1D{A,B}) where {A,B}
    print(
        io,
        "ControlVolume1D{$A,$B}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
    )
end

"""
    mutable struct ControlVolume1D1F{A,B} <: AbstractControlVolume1D
        w::A
        prim::A
        sw::A

        f::B
        sf::B
    end

1D control volume with 1 distribution function

"""
mutable struct ControlVolume1D1F{A,B} <: AbstractControlVolume1D
    w::A
    prim::A
    sw::A

    f::B
    sf::B
end

function ControlVolume1D1F(W::T, PRIM::T, F) where T
    w = deepcopy(W)
    prim = deepcopy(PRIM)
    sw = zero(W)

    f = deepcopy(F)
    sf = zero(f)

    return ControlVolume1D1F{typeof(w),typeof(f)}(w, prim, sw, f, sf)
end

function Base.show(io::IO, ctr::ControlVolume1D1F{A,B}) where {A,B}
    print(
        io,
        "ControlVolume1D1F{$A,$B}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
        "pdf vars: f\n",
        "pdf slopes: sf\n",
    )
end

"""
    mutable struct ControlVolume1D2F{A,B} <: AbstractControlVolume1D
        w::A
        prim::A
        sw::A

        h::B
        b::B
        sh::B
        sb::B
    end

1D control volume with 2 distribution functions

"""
mutable struct ControlVolume1D2F{A,B} <: AbstractControlVolume1D
    w::A
    prim::A
    sw::A

    h::B
    b::B
    sh::B
    sb::B
end

function ControlVolume1D2F(
    W::T1,
    PRIM::T1,
    H::T2,
    B::T2,
) where {T1<:AbstractArray,T2<:AbstractArray}

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    sw = zero(W)

    h = deepcopy(H)
    b = deepcopy(B)
    sh = zero(h)
    sb = zero(b)

    return ControlVolume1D2F{typeof(w),typeof(h)}(w, prim, sw, h, b, sh, sb)

end

function Base.show(io::IO, ctr::ControlVolume1D2F{A,B}) where {A,B}
    print(
        io,
        "ControlVolume1D2F{$A,$B}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
        "pdf vars: h, b\n",
        "pdf slopes: sh, sb\n",
    )
end

"""
    mutable struct ControlVolume1D3F{A,B,C,D,E} <: AbstractControlVolume1D
        w::A
        prim::A
        sw::A

        h0::B
        h1::B
        h2::B
        sh0::B
        sh1::B
        sh2::B

        E::C
        B::C
        ϕ::D
        ψ::D
        lorenz::E
    end

1D control volume with 3 distribution functions

"""
mutable struct ControlVolume1D3F{A,B,C,D,E} <: AbstractControlVolume1D
    w::A
    prim::A
    sw::A

    h0::B
    h1::B
    h2::B
    sh0::B
    sh1::B
    sh2::B

    E::C
    B::C
    ϕ::D
    ψ::D
    lorenz::E
end

# deterministic
function ControlVolume1D3F(
    W::AbstractArray{<:Real,2},
    PRIM::AbstractArray{<:Real,2},
    H0::AbstractArray{<:AbstractFloat,3},
    H1::AbstractArray{<:AbstractFloat,3},
    H2::AbstractArray{<:AbstractFloat,3},
    E0::AbstractArray{<:AbstractFloat,1},
    B0::AbstractArray{<:AbstractFloat,1},
    L::AbstractArray{<:AbstractFloat,2},
)
    w = deepcopy(W)
    prim = deepcopy(PRIM)
    sw = zero(W)

    h0 = deepcopy(H0)
    h1 = deepcopy(H1)
    h2 = deepcopy(H2)
    sh0 = zero(H0)
    sh1 = zero(H1)
    sh2 = zero(H2)

    E = deepcopy(E0)
    B = deepcopy(B0)
    ϕ = 0.0
    ψ = 0.0
    lorenz = deepcopy(L)

    return ControlVolume1D3F{typeof(w),typeof(h0),typeof(E),typeof(ϕ),typeof(lorenz)}(
        w,
        prim,
        sw,
        h0,
        h1,
        h2,
        sh0,
        sh1,
        sh2,
        E,
        B,
        ϕ,
        ψ,
        lorenz,
    )
end

# Rykov
function ControlVolume1D3F(
    W::AbstractArray{<:Real,1},
    PRIM::AbstractArray{<:Real,1},
    H0::AbstractArray{<:AbstractFloat,1},
    H1::AbstractArray{<:AbstractFloat,1},
    H2::AbstractArray{<:AbstractFloat,1},
)
    w = deepcopy(W)
    prim = deepcopy(PRIM)
    sw = zero(W)

    h0 = deepcopy(H0)
    h1 = deepcopy(H1)
    h2 = deepcopy(H2)
    sh0 = zero(H0)
    sh1 = zero(H1)
    sh2 = zero(H2)

    E = nothing
    B = nothing
    ϕ = nothing
    ψ = nothing
    lorenz = nothing

    return ControlVolume1D3F{typeof(w),typeof(h0),typeof(E),typeof(ϕ),typeof(lorenz)}(
        w,
        prim,
        sw,
        h0,
        h1,
        h2,
        sh0,
        sh1,
        sh2,
        E,
        B,
        ϕ,
        ψ,
        lorenz,
    )
end

# stochastic
function ControlVolume1D3F(
    W::AbstractArray{<:Real,3},
    PRIM::AbstractArray{<:Real,3},
    H0::AbstractArray{<:AbstractFloat,4},
    H1::AbstractArray{<:AbstractFloat,4},
    H2::AbstractArray{<:AbstractFloat,4},
    E0::AbstractArray{<:AbstractFloat,2},
    B0::AbstractArray{<:AbstractFloat,2},
    L::AbstractArray{<:AbstractFloat,3},
)
    w = deepcopy(W)
    prim = deepcopy(PRIM)
    sw = zero(W)

    h0 = deepcopy(H0)
    h1 = deepcopy(H1)
    h2 = deepcopy(H2)
    sh0 = zero(H0)
    sh1 = zero(H1)
    sh2 = zero(H2)

    E = deepcopy(E0)
    B = deepcopy(B0)
    ϕ = zero(E[1, :]) # here is difference
    ψ = zero(B[1, :])
    lorenz = deepcopy(L)

    return ControlVolume1D3F{typeof(w),typeof(h0),typeof(E),typeof(ϕ),typeof(lorenz)}(
        w,
        prim,
        sw,
        h0,
        h1,
        h2,
        sh0,
        sh1,
        sh2,
        E,
        B,
        ϕ,
        ψ,
        lorenz,
    )
end

function Base.show(io::IO, ctr::ControlVolume1D3F{A,B,C,D,E}) where {A,B,C,D,E}
    print(
        io,
        "ControlVolume1D3F{$A,$B,$C,$D,$E}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
        "pdf vars: h0, h1, h2\n",
        "pdf slopes: sh0, sh1, sh2\n",
        "electric field: $(ctr.E)\n",
        "magnetic field: $(ctr.B)\n",
        "Lorenz force: $(ctr.lorenz)\n",
    )
end

"""
    mutable struct ControlVolume1D4F{A,B,C,D,E} <: AbstractControlVolume1D
        w::A
        prim::A
        sw::A

        h0::B
        h1::B
        h2::B
        h3::B
        sh0::B
        sh1::B
        sh2::B
        sh3::B

        E::C
        B::C
        ϕ::D
        ψ::D
        lorenz::E
    end

1D control volume with 4 distribution functions

"""
mutable struct ControlVolume1D4F{A,B,C,D,E} <: AbstractControlVolume1D
    w::A
    prim::A
    sw::A

    h0::B
    h1::B
    h2::B
    h3::B
    sh0::B
    sh1::B
    sh2::B
    sh3::B

    E::C
    B::C
    ϕ::D
    ψ::D
    lorenz::E
end

#--- deterministic ---#
function ControlVolume1D4F(
    W::AbstractArray{<:Real,2},
    PRIM::AbstractArray{<:Real,2},
    H0::AbstractArray{<:AbstractFloat,2},
    H1::AbstractArray{Float64,2},
    H2::AbstractArray{Float64,2},
    H3::AbstractArray{Float64,2},
    E0::AbstractArray{Float64,1},
    B0::AbstractArray{Float64,1},
    L::AbstractArray{Float64,2},
)

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    sw = zero(W)

    h0 = deepcopy(H0)
    h1 = deepcopy(H1)
    h2 = deepcopy(H2)
    h3 = deepcopy(H3)
    sh0 = deepcopy(H0)
    sh1 = deepcopy(H1)
    sh2 = deepcopy(H2)
    sh3 = deepcopy(H3)

    E = deepcopy(E0)
    B = deepcopy(B0)
    ϕ = 0.0
    ψ = 0.0
    lorenz = deepcopy(L)

    return ControlVolume1D4F{typeof(w),typeof(h0),typeof(E),typeof(ϕ),typeof(lorenz)}(
        w,
        prim,
        sw,
        h0,
        h1,
        h2,
        h3,
        sh0,
        sh1,
        sh2,
        sh3,
        E,
        B,
        ϕ,
        ψ,
        lorenz,
    )

end

#--- uncertainty quantification ---#
function ControlVolume1D4F(
    W::AbstractArray{<:Real,3},
    PRIM::AbstractArray{<:Real,3},
    H0::AbstractArray{<:AbstractFloat,3},
    H1::AbstractArray{Float64,3},
    H2::AbstractArray{Float64,3},
    H3::AbstractArray{Float64,3},
    E0::AbstractArray{Float64,2},
    B0::AbstractArray{Float64,2},
    L::AbstractArray{Float64,3},
)

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    sw = zero(W)

    h0 = deepcopy(H0)
    h1 = deepcopy(H1)
    h2 = deepcopy(H2)
    h3 = deepcopy(H3)
    sh0 = zero(H0)
    sh1 = zero(H1)
    sh2 = zero(H2)
    sh3 = zero(H3)

    E = deepcopy(E0)
    B = deepcopy(B0)
    ϕ = zero(E[1, :])
    ψ = zero(B[1, :])
    lorenz = deepcopy(L)

    return ControlVolume1D4F{typeof(w),typeof(h0),typeof(E),typeof(ϕ),typeof(lorenz)}(
        w,
        prim,
        sw,
        h0,
        h1,
        h2,
        h3,
        sh0,
        sh1,
        sh2,
        sh3,
        E,
        B,
        ϕ,
        ψ,
        lorenz,
    )

end

function Base.show(io::IO, ctr::ControlVolume1D4F{A,B,C,D,E}) where {A,B,C,D,E}
    print(
        io,
        "ControlVolume1D4F{$A,$B,$C,$D,$E}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
        "pdf vars: h0, h1, h2, h3\n",
        "pdf slopes: sh0, sh1, sh2, sh3\n",
        "electric field: $(ctr.E)\n",
        "magnetic field: $(ctr.B)\n",
        "Lorenz force: $(ctr.lorenz)\n",
    )
end

# ------------------------------------------------------------
# 2D
# ------------------------------------------------------------

"""
    mutable struct ControlVolume2D{A,B} <: AbstractControlVolume2D
        w::A
        prim::A
        sw::B
    end

2D control volume with no distribution function

"""
mutable struct ControlVolume2D{A,B} <: AbstractControlVolume2D
    w::A
    prim::A
    sw::B
end

function ControlVolume2D(W, PRIM)
    w = deepcopy(W)
    prim = deepcopy(PRIM)
    #sw = zeros(eltype(W), (axes(W)..., Base.OneTo(2)))
    sw = slope_array(W)

    return ControlVolume2D{typeof(w),typeof(sw)}(w, prim, sw)
end

function Base.show(io::IO, ctr::ControlVolume2D{A,B}) where {A,B}
    print(
        io,
        "ControlVolume2D{$A,$B}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
    )
end

"""
    mutable struct ControlVolume2D1F{A,B,C,D} <: AbstractControlVolume2D
        w::A
        prim::A
        sw::B

        f::C
        sf::D
    end

2D control volume with 1 distribution function

"""
mutable struct ControlVolume2D1F{A,B,C,D} <: AbstractControlVolume2D
    w::A
    prim::A
    sw::B

    f::C
    sf::D
end

function ControlVolume2D1F(W, PRIM, F::AbstractArray)
    w = deepcopy(W)
    prim = deepcopy(PRIM)
    #sw = zeros(eltype(W), (axes(W)..., Base.OneTo(2)))
    sw = slope_array(W)

    f = deepcopy(F)
    #sf = zeros(eltype(F), (axes(F)..., Base.OneTo(2)))
    sf = slope_array(F)

    return ControlVolume2D1F{typeof(w),typeof(sw),typeof(f),typeof(sf)}(w, prim, sw, f, sf)
end

function Base.show(io::IO, ctr::ControlVolume2D1F{A,B,C,D}) where {A,B,C,D}
    print(
        io,
        "ControlVolume2D1F{$A,$B,$C,$D}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
        "pdf vars: f\n",
        "pdf slopes: sf\n",
    )
end

"""
    mutable struct ControlVolume2D2F{A,B,C,D} <: AbstractControlVolume2D
        w::A
        prim::A
        sw::B

        h::C
        b::C
        sh::D
        sb::D
    end

2D control volume with 2 distribution functions

"""
mutable struct ControlVolume2D2F{A,B,C,D} <: AbstractControlVolume2D
    w::A
    prim::A
    sw::B

    h::C
    b::C
    sh::D
    sb::D
end

function ControlVolume2D2F(
    W::AbstractArray,
    PRIM::AbstractArray,
    H::AbstractArray,
    B::AbstractArray,
)

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    #sw = zeros(eltype(W), (axes(W)..., Base.OneTo(2)))
    sw = slope_array(W)

    h = deepcopy(H)
    b = deepcopy(B)
    #sh = zeros(eltype(H), (axes(H)..., Base.OneTo(2)))
    #sb = zeros(eltype(B), (axes(B)..., Base.OneTo(2)))
    sh = slope_array(H)
    sb = slope_array(B)

    return ControlVolume2D2F{typeof(w),typeof(sw),typeof(h),typeof(sh)}(
        w,
        prim,
        sw,
        h,
        b,
        sh,
        sb,
    )

end

function Base.show(io::IO, ctr::ControlVolume2D2F{A,B,C,D}) where {A,B,C,D}
    print(
        io,
        "ControlVolume2D2F{$A,$B,$C,$D}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
        "pdf vars: h, b\n",
        "pdf slopes: sh, sb\n",
    )
end

"""
    mutable struct ControlVolume2D3F{T2,T3,T4,T5,T6,T7,T8} <: AbstractControlVolume2D
        w::T2
        prim::T2
        sw::T3

        h0::T4
        h1::T4
        h2::T4
        sh0::T5
        sh1::T5
        sh2::T5

        E::T6
        B::T6
        ϕ::T7
        ψ::T7
        lorenz::T8
    end

2D control volume with 3 distribution functions

"""
mutable struct ControlVolume2D3F{T2,T3,T4,T5,T6,T7,T8} <: AbstractControlVolume2D
    w::T2
    prim::T2
    sw::T3

    h0::T4
    h1::T4
    h2::T4
    sh0::T5
    sh1::T5
    sh2::T5

    E::T6
    B::T6
    ϕ::T7
    ψ::T7
    lorenz::T8
end

#--- deterministic & stochastic ---#
function ControlVolume2D3F(
    W::AbstractArray,
    PRIM::AbstractArray,
    H0::AbstractArray,
    H1::AbstractArray,
    H2::AbstractArray,
    E0::AbstractArray,
    B0::AbstractArray,
    L::AbstractArray,
)

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    #sw = zeros(eltype(W), (axes(W)..., Base.OneTo(2))) # 2D
    sw = slope_array(W)

    h0 = deepcopy(H0)
    h1 = deepcopy(H1)
    h2 = deepcopy(H2)
    #sh0 = zeros(eltype(H0), (axes(H0)..., Base.OneTo(2)))
    #sh1 = zeros(eltype(H1), (axes(H1)..., Base.OneTo(2)))
    #sh2 = zeros(eltype(H2), (axes(H2)..., Base.OneTo(2)))
    sh0 = slope_array(H0)
    sh1 = slope_array(H1)
    sh2 = slope_array(H2)

    E = deepcopy(E0)
    B = deepcopy(B0)
    ϕ = 0.0
    ψ = 0.0
    lorenz = deepcopy(L)

    return ControlVolume2D3F{
        typeof(w),
        typeof(sw),
        typeof(h0),
        typeof(sh2),
        typeof(E),
        typeof(ϕ),
        typeof(lorenz),
    }(
        w,
        prim,
        sw,
        h0,
        h1,
        h2,
        sh0,
        sh1,
        sh2,
        E,
        B,
        ϕ,
        ψ,
        lorenz,
    )

end

#--- Rykov ---#
function ControlVolume2D3F(
    W::AbstractArray{<:Real,1},
    PRIM::AbstractArray{<:Real,1},
    H0::AbstractArray{<:AbstractFloat,2},
    H1::AbstractArray{<:AbstractFloat,2},
    H2::AbstractArray{<:AbstractFloat,2},
)

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    #sw = zeros(eltype(W), (axes(W)..., Base.OneTo(2))) # 2D
    sw = slope_array(W)

    h0 = deepcopy(H0)
    h1 = deepcopy(H1)
    h2 = deepcopy(H2)
    #sh0 = zeros(eltype(H0), (axes(H0)..., Base.OneTo(2)))
    #sh1 = zeros(eltype(H1), (axes(H1)..., Base.OneTo(2)))
    #sh2 = zeros(eltype(H2), (axes(H2)..., Base.OneTo(2)))
    sh0 = slope_array(H0)
    sh1 = slope_array(H1)
    sh2 = slope_array(H2)

    E = nothing
    B = nothing
    ϕ = nothing
    ψ = nothing
    lorenz = nothing

    return ControlVolume2D3F{
        typeof(w),
        typeof(sw),
        typeof(h0),
        typeof(sh2),
        typeof(E),
        typeof(ϕ),
        typeof(lorenz),
    }(
        w,
        prim,
        sw,
        h0,
        h1,
        h2,
        sh0,
        sh1,
        sh2,
        E,
        B,
        ϕ,
        ψ,
        lorenz,
    )

end

function Base.show(
    io::IO,
    ctr::ControlVolume2D3F{T1,T2,T3,T4,T5,T6,T7},
) where {T1,T2,T3,T4,T5,T6,T7}
    print(
        io,
        "ControlVolume2D3F{$T1,$T2,$T3,$T4,$T5,$T6,$T7}\n",
        "conservative vars: $(ctr.w)\n",
        "primitive vars: $(ctr.prim)\n",
        "conservative slopes: $(ctr.sw)\n",
        "pdf vars: h0, h1, h2\n",
        "pdf slopes: sh0, sh1, sh2\n",
        "electric field: $(ctr.E)\n",
        "magnetic field: $(ctr.B)\n",
        "Lorenz force: $(ctr.lorenz)\n",
    )
end

# ------------------------------------------------------------
# Unstructured cell
# ------------------------------------------------------------

"""
    mutable struct ControlVolumeUS{E,F,A,B,C} <: AbstractUnstructControlVolume
        n::E
        x::F
        dx::F

        w::A
        prim::B
        sw::C
    end

Unstructured control volume with no distribution function

"""
mutable struct ControlVolumeUS{E,F,A,B,C} <: AbstractUnstructControlVolume
    n::E
    x::F
    dx::F

    w::A
    prim::B
    sw::C
end

function ControlVolumeUS(N, X::T, DX::T, W, PRIM) where {T<:Union{Real,AbstractVector}}
    n = deepcopy(N)
    x = deepcopy(X)
    dx = deepcopy(DX)

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    #sw = zeros(eltype(W), axes(W)..., length(N[1]))
    sw = ifelse(
        length(N[1]) == 2,
        slope_array(W; reduction = true),
        slope_array(W; reduction = false),
    )

    return ControlVolumeUS{typeof(n),typeof(x),typeof(w),typeof(prim),typeof(sw)}(
        n,
        x,
        dx,
        w,
        prim,
        sw,
    )
end

"""
    mutable struct ControlVolumeUS1F{E,F,A,B,C,D} <: AbstractUnstructControlVolume
        n::E
        x::F
        dx::F

        w::A
        prim::A
        sw::B

        f::C
        sf::D
    end

Unstructured control volume with 1 distribution function

"""
mutable struct ControlVolumeUS1F{E,F,A,B,C,D} <: AbstractUnstructControlVolume
    n::E
    x::F
    dx::F

    w::A
    prim::A
    sw::B

    f::C
    sf::D
end

function ControlVolumeUS1F(N, X, DX, W, PRIM, F::AbstractArray{T}) where T
    n = deepcopy(N)
    x = deepcopy(X)
    dx = deepcopy(DX)

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    #sw = zeros(eltype(W), axes(W)..., length(N[1]))
    sw = ifelse(
        length(N[1]) == 2,
        slope_array(W; reduction = true),
        slope_array(W; reduction = false),
    )

    f = deepcopy(F)
    #sf = zeros(eltype(F), axes(F)..., length(N[1]))
    sf = ifelse(
        length(N[1]) == 2,
        slope_array(F; reduction = true),
        slope_array(F; reduction = false),
    )

    return ControlVolumeUS1F{typeof(n),typeof(x),typeof(w),typeof(sw),typeof(f),typeof(sf)}(
        n,
        x,
        dx,
        w,
        prim,
        sw,
        f,
        sf,
    )
end

"""
    mutable struct ControlVolumeUS2F{E,F,A,B,C,D} <: AbstractUnstructControlVolume
        n::E
        x::F
        dx::F

        w::A
        prim::A
        sw::B

        h::C
        b::C
        sh::D
        sb::D
    end

Unstructured control volume with 2 distribution functions

"""
mutable struct ControlVolumeUS2F{E,F,A,B,C,D} <: AbstractUnstructControlVolume
    n::E
    x::F
    dx::F

    w::A
    prim::A
    sw::B

    h::C
    b::C
    sh::D
    sb::D
end

function ControlVolumeUS2F(
    N,
    X,
    DX,
    W::T1,
    PRIM::T1,
    H::T2,
    B::T2,
) where {T1<:AbstractArray,T2<:AbstractArray}

    n = deepcopy(N)
    x = deepcopy(X)
    dx = deepcopy(DX)

    w = deepcopy(W)
    prim = deepcopy(PRIM)
    #sw = zeros(eltype(W), axes(W)..., length(N[1]))
    sw = ifelse(
        length(N[1]) == 2,
        slope_array(W; reduction = true),
        slope_array(W; reduction = false),
    )

    h = deepcopy(H)
    b = deepcopy(B)
    #sh = zeros(eltype(H), axes(H)..., length(N[1]))
    #sb = zeros(eltype(B), axes(B)..., length(N[1]))
    sh = ifelse(
        length(N[1]) == 2,
        slope_array(H; reduction = true),
        slope_array(H; reduction = false),
    )
    sb = ifelse(
        length(N[1]) == 2,
        slope_array(B; reduction = true),
        slope_array(B; reduction = false),
    )

    return ControlVolumeUS2F{typeof(n),typeof(x),typeof(w),typeof(sw),typeof(h),typeof(sh)}(
        n,
        x,
        dx,
        w,
        prim,
        sw,
        h,
        b,
        sh,
        sb,
    )

end


################################################################################
#                             BASIC REFERENCE METHODS
################################################################################
"""
    JuMP.owner_model(cref::InfOptConstraintRef)::InfiniteModel

Extend [`JuMP.owner_model`](@ref JuMP.owner_model(::JuMP.ConstraintRef)) to
return the infinite model associated with `cref`.

**Example**
```julia-repl
julia> model = owner_model(cref)
An InfiniteOpt Model
Minimization problem with:
Finite Parameters: 0
Infinite Parameters: 3
Variables: 3
Derivatives: 0
Measures: 0
Objective function type: GeneralVariableRef
`GenericAffExpr{Float64,GeneralVariableRef}`-in-`MathOptInterface.EqualTo{Float64}`: 1 constraint
Names registered in the model: g, t, h, x
Optimizer model backend information:
Model mode: AUTOMATIC
CachingOptimizer state: NO_OPTIMIZER
Solver name: No optimizer attached.
```
"""
JuMP.owner_model(cref::InfOptConstraintRef)::InfiniteModel = cref.model

"""
    JuMP.index(cref::InfOptConstraintRef)::ConstraintIndex

Extend [`JuMP.index`](@ref JuMP.index(::JuMP.ConstraintRef)) to return the index
of an `InfiniteOpt` constraint `cref`.

**Example**
```julia-repl
julia> index(cref)
ConstrainIndex(2)
```
"""
JuMP.index(cref::InfOptConstraintRef)::ConstraintIndex = cref.index

# Extend Base and JuMP functions
function Base.:(==)(v::InfOptConstraintRef, w::InfOptConstraintRef)::Bool
    return v.model === w.model && v.index == w.index
end
Base.broadcastable(cref::InfOptConstraintRef) = Ref(cref)
JuMP.constraint_type(m::InfiniteModel) = InfOptConstraintRef

################################################################################
#                             CORE OBJECT METHODS
################################################################################
# Extend _add_data_object
function _add_data_object(model::InfiniteModel,
                          object::ConstraintData)::ConstraintIndex
    return MOIUC.add_item(model.constraints, object)
end

# Extend _data_dictionary
function _data_dictionary(cref::InfOptConstraintRef
    )::MOIUC.CleverDict{ConstraintIndex, ConstraintData}
    return JuMP.owner_model(cref).constraints
end

# Extend _data_object
function _data_object(cref::InfOptConstraintRef)::ConstraintData
  object = get(_data_dictionary(cref), JuMP.index(cref), nothing)
  object === nothing && error("Invalid constraint reference, cannot find " *
                        "corresponding constraint in the model. This is likely " *
                        "caused by using the reference of a deleted constraint.")
  return object
end

# Return the core constraint object
function _core_constraint_object(cref::InfOptConstraintRef)::JuMP.AbstractConstraint
    return _data_object(cref).constraint
end

# Update the core constraint object
function _set_core_constraint_object(cref::InfOptConstraintRef,
                                     constr::JuMP.AbstractConstraint)::Nothing
    _data_object(cref).constraint = constr
    return
end

# Extend _object_numbers
function _object_numbers(cref::InfOptConstraintRef)::Vector{Int}
    return _data_object(cref).object_nums
end

# Extend _measure_dependencies
function _measure_dependencies(cref::InfOptConstraintRef)::Vector{MeasureIndex}
    return _data_object(cref).measure_indices
end

# Return if this constraint is an info constraint
function _is_info_constraint(cref::InfOptConstraintRef)::Bool
    return _data_object(cref).is_info_constraint
end

# Extend _delete_data_object
function _delete_data_object(cref::InfOptConstraintRef)::Nothing
    delete!(_data_dictionary(cref), JuMP.index(cref))
    return
end

################################################################################
#                             DEFINITION METHODS
################################################################################
# This might not be necessary...
function JuMP.build_constraint(_error::Function,
    v::GeneralVariableRef,
    set::MOI.AbstractScalarSet;
    parameter_bounds::ParameterBounds{GeneralVariableRef} = ParameterBounds()
    )::JuMP.AbstractConstraint
    # make the constraint
    if length(parameter_bounds) != 0
        _check_bounds(parameter_bounds)
        return BoundedScalarConstraint(v, set, parameter_bounds,
                                       copy(parameter_bounds))
    else
        return JuMP.ScalarConstraint(v, set)
    end
end

"""
    JuMP.build_constraint(_error::Function, expr::JuMP.AbstractJuMPScalar,
                          set::MOI.AbstractScalarSet;
                          [parameter_bounds::ParameterBounds = ParameterBounds()]
                          )::JuMP.AbstractConstraint

Extend `JuMP.build_constraint` to accept the ```parameter_bounds``` argument
and return a [`BoundedScalarConstraint`](@ref) if the ```parameter_bounds``` keyword
argument is specifed or return a [`JuMP.ScalarConstraint`](@ref) otherwise. This is
primarily intended to work as an internal function for constraint macros.

**Example**
```jldoctest; setup = :(using JuMP, InfiniteOpt; model = InfiniteModel())
julia> @infinite_parameter(model, t in [0, 10]);

julia> @infinite_variable(model, g(t));

julia> @hold_variable(model, x);

julia> constr = build_constraint(error, g + x, MOI.EqualTo(42.0),
              parameter_bounds = ParameterBounds((t => IntervalSet(0, 1),)));

julia> isa(constr, BoundedScalarConstraint)
true
```
"""
function JuMP.build_constraint(_error::Function,
    expr::Union{JuMP.GenericAffExpr{C, V}, JuMP.GenericQuadExpr{C, V}},
    set::MOI.AbstractScalarSet;
    parameter_bounds::ParameterBounds{GeneralVariableRef} = ParameterBounds()
    )::JuMP.AbstractConstraint where {C, V <: GeneralVariableRef}
    # make the constraint
    offset = JuMP.constant(expr)
    JuMP.add_to_expression!(expr, -offset)
    if !isempty(parameter_bounds)
        _check_bounds(parameter_bounds)
        return BoundedScalarConstraint(expr, MOIU.shift_constant(set, -offset),
                                       parameter_bounds, copy(parameter_bounds))
    else
        return JuMP.ScalarConstraint(expr, MOIU.shift_constant(set, -offset))
    end
end

## Perfrom bound checks and update them if needed, then return the updated constraint
# BoundedScalarConstraint
function _check_and_update_bounds(model::InfiniteModel,
                                  c::BoundedScalarConstraint,
                                  vrefs::Vector{GeneralVariableRef}
                                  )::JuMP.AbstractConstraint
    # look for bounded hold variables and update bounds
    for vref in vrefs
        _update_var_bounds(vref, c.bounds)
    end
    # now validate and return
    _validate_bounds(model, c.bounds)
    # TODO should we check that bounds don't violate point variables?
    return c
end

# ScalarConstraint
function _check_and_update_bounds(model::InfiniteModel,
                                  c::JuMP.ScalarConstraint,
                                  vrefs::Vector{GeneralVariableRef}
                                  )::JuMP.AbstractConstraint
    bounds = ParameterBounds()
    # check for bounded hold variables and build the intersection of the bounds
    for vref in vrefs
        _update_var_bounds(vref, bounds)
    end
    # if we added bounds, change to a bounded constraint and validate
    if !isempty(bounds)
        new_c = BoundedScalarConstraint(c.func, c.set, bounds, ParameterBounds())
        _validate_bounds(model, new_c.bounds)
        return new_c
    end
    # TODO should we check that bounds don't violate point variables?
    return c
end

# Used to update the model.var_to_constrs field
function _update_var_constr_mapping(vrefs::Vector{GeneralVariableRef},
                                    cref::InfOptConstraintRef)::Nothing
    for vref in vrefs
        dvref = dispatch_variable_ref(vref)
        push!(_constraint_dependencies(dvref), JuMP.index(cref))
        if dvref isa MeasureRef
            push!(_measure_dependencies(cref), JuMP.index(dvref))
        end
        # TODO maybe make mapping for derivatives...
    end
    return
end

# Extend functions for bounded constraints
JuMP.shape(c::BoundedScalarConstraint) = JuMP.shape(JuMP.ScalarConstraint(c.func, c.set))
JuMP.jump_function(c::BoundedScalarConstraint) = c.func
JuMP.moi_set(c::BoundedScalarConstraint) = c.set
parameter_bounds(c::BoundedScalarConstraint) = c.bounds
original_parameter_bounds(c::BoundedScalarConstraint) = c.orig_bounds

"""
    JuMP.add_constraint(model::InfiniteModel, c::JuMP.AbstractConstraint,
                        [name::String = ""])::InfOptConstraintRef

Extend [`JuMP.add_constraint`](@ref JuMP.add_constraint(::JuMP.Model, ::JuMP.AbstractConstraint, ::String))
to add a constraint `c` to an infinite model
`model` with name `name`. Returns an appropriate constraint reference whose type
depends on what variables are used to define the constraint. Errors if a vector
constraint is used, the constraint only constains parameters, or if any
variables do not belong to `model`. This is primarily used as an internal
method for the cosntraint macros.

**Example**
```jldoctest; setup = :(using JuMP, InfiniteOpt; model = InfiniteModel())
julia> @infinite_parameter(model, t in [0, 10]);

julia> @infinite_variable(model, g(t));

julia> @hold_variable(model, x);

julia> constr = build_constraint(error, g + x, MOI.EqualTo(42));

julia> cref = add_constraint(model, constr, "name")
name : g(t) + x = 42.0, ∀ t ∈ [0, 10]
```
"""
function JuMP.add_constraint(model::InfiniteModel,
                             c::JuMP.AbstractConstraint,
                             name::String = "";
                             is_info_constr::Bool = false
                             )::InfOptConstraintRef
    # gather the unique list of variable references for testing and mapping
    vrefs = _all_function_variables(JuMP.jump_function(c))
    # test in the model
    for vref in vrefs
        JuMP.check_belongs_to_model(vref, model)
    end
    # test parameter bounds and update with hold variable bounds
    if model.has_hold_bounds
        c = _check_and_update_bounds(model, c, vrefs)
    elseif c isa BoundedScalarConstraint
        _validate_bounds(model, parameter_bounds(c))
    end
    # get the parameter object numbers
    object_nums = sort!(_object_numbers(vrefs))
    # add the constaint to the model
    constr_object = ConstraintData(c, object_nums, name, MeasureIndex[],
                                   is_info_constr)
    cindex = _add_data_object(model, constr_object)
    cref = InfOptConstraintRef(model, cindex)
    # update the variable mappings and model status
    _update_var_constr_mapping(vrefs, cref)
    set_optimizer_model_ready(model, false)
    return cref
end

# Fallback for vector constraints
function JuMP.add_constraint(model::InfiniteModel, c::JuMP.VectorConstraint,
                             name::String = "";
                             is_info_constr::Bool = false)
    error("Vector constraints not supported in InfiniteOpt.")
end

################################################################################
#                                 JuMP METHODS
################################################################################
"""
    JuMP.is_valid(model::InfiniteModel, cref::InfOptConstraintRef)::Bool

Extend [`JuMP.is_valid`](@ref JuMP.is_valid(::JuMP.Model, ::JuMP.ConstraintRef{JuMP.Model}))
to return `Bool` whether an `InfiniteOpt` constraint reference is valid.

**Example**
```jldoctest; setup = :(using JuMP, InfiniteOpt; model = InfiniteModel(); cref = @constraint(model, 2 == 2))
julia> is_valid(model, cref)
true
```
"""
function JuMP.is_valid(model::InfiniteModel, cref::InfOptConstraintRef)::Bool
    return (model === JuMP.owner_model(cref) &&
            JuMP.index(cref) in keys(_data_dictionary(cref)))
end

"""
    JuMP.constraint_object(cref::InfOptConstraintRef)::JuMP.AbstractConstraint

Extend [`JuMP.constraint_object`](@ref JuMP.constraint_object)
to return the constraint object associated with `cref`.

**Example**
```jldoctest; setup = :(using JuMP, InfiniteOpt; model = InfiniteModel())
julia> @infinite_parameter(model, t in [0, 10]);

julia> @hold_variable(model, x <= 1);

julia> cref = UpperBoundRef(x);

julia> obj = constraint_object(cref)
ScalarConstraint{GeneralVariableRef,MathOptInterface.LessThan{Float64}}(x,
MathOptInterface.LessThan{Float64}(1.0))
```
"""
function JuMP.constraint_object(cref::InfOptConstraintRef)::JuMP.AbstractConstraint
    return _core_constraint_object(cref)
end

"""
    JuMP.name(cref::InfOptConstraintRef)::String

Extend [`JuMP.name`](@ref JuMP.name(::JuMP.ConstraintRef{JuMP.Model,<:JuMP._MOICON})
to return the name of an `InfiniteOpt` constraint.

**Example**
```jldoctest; setup = :(using JuMP, InfiniteOpt; model = InfiniteModel(); cref = @constraint(model, constr_name, 2 == 2))
julia> name(cref)
"constr_name"
```
"""
function JuMP.name(cref::InfOptConstraintRef)::String
    object = get(_data_dictionary(cref), JuMP.index(cref), nothing)
    return object === nothing ? "" : object.name
end

"""
    JuMP.set_name(cref::InfOptConstraintRef, name::String)::Nothing

Extend [`JuMP.set_name`](@ref JuMP.set_name(::JuMP.ConstraintRef{JuMP.Model,<:JuMP._MOICON}, ::String))
to specify the name of a constraint `cref`.

**Example**
```jldoctest; setup = :(using JuMP, InfiniteOpt; model = InfiniteModel(); cref = @constraint(model, 2 == 2))
julia> set_name(cref, "new_name")

julia> name(cref)
"new_name"
```
"""
function JuMP.set_name(cref::InfOptConstraintRef, name::String)::Nothing
    _data_object(cref).name = name
    JuMP.owner_model(cref).name_to_constr = nothing
    return
end

# Return a constraint set with an updated value
function _set_set_value(set::S, value::Real) where {T, S <: Union{MOI.LessThan{T},
                                            MOI.GreaterThan{T}, MOI.EqualTo{T}}}
    return S(convert(T, value))
end

"""
    JuMP.set_normalized_rhs(cref::InfOptConstraintRef, value::Real)::Nothing

Set the right-hand side term of `constraint` to `value`.
Note that prior to this step, JuMP will aggregate all constant terms onto the
right-hand side of the constraint. For example, given a constraint `2x + 1 <=
2`, `set_normalized_rhs(con, 4)` will create the constraint `2x <= 4`, not `2x +
1 <= 4`.

```julia-repl
julia> @constraint(model, con, 2x + 1 <= 2)
con : 2 x ≤ 1.0

julia> set_normalized_rhs(con, 4)

julia> con
con : 2 x ≤ 4.0
```
"""
function JuMP.set_normalized_rhs(cref::InfOptConstraintRef, value::Real)::Nothing
    old_constr = JuMP.constraint_object(cref)
    new_set = _set_set_value(JuMP.moi_set(old_constr), value)
    if old_constr isa BoundedScalarConstraint
        new_constr = BoundedScalarConstraint(JuMP.jump_function(old_constr),
                        new_set, parameter_bounds(old_constr),
                        original_parameter_bounds(old_constr))
    else
        new_constr = JuMP.ScalarConstraint(JuMP.jump_function(old_constr),
                                           new_set)
    end
    _set_core_constraint_object(cref, new_constr)
    return
end

"""
    JuMP.normalized_rhs(cref::InfOptConstraintRef)::Float64

Return the right-hand side term of `cref` after JuMP has converted the
constraint into its normalized form.
"""
function JuMP.normalized_rhs(cref::InfOptConstraintRef)::Float64
    constr = JuMP.constraint_object(cref)
    return MOI.constant(JuMP.moi_set(constr))
end

"""
    JuMP.add_to_function_constant(cref::InfOptConstraintRef, value::Real)::Nothing

Add `value` to the function constant term.
Note that for scalar constraints, JuMP will aggregate all constant terms onto the
right-hand side of the constraint so instead of modifying the function, the set
will be translated by `-value`. For example, given a constraint `2x <=
3`, `add_to_function_constant(c, 4)` will modify it to `2x <= -1`.
```
"""
function JuMP.add_to_function_constant(cref::InfOptConstraintRef,
                                       value::Real)::Nothing
    current_value = JuMP.normalized_rhs(cref)
    JuMP.set_normalized_rhs(cref, current_value - value)
    return
end

"""
    JuMP.set_normalized_coefficient(cref::InfOptConstraintRef,
                                    variable::GeneralVariableRef,
                                    value::Real)::Nothing

Set the coefficient of `variable` in the constraint `constraint` to `value`.
Note that prior to this step, JuMP will aggregate multiple terms containing the
same variable. For example, given a constraint `2x + 3x <= 2`,
`set_normalized_coefficient(con, x, 4)` will create the constraint `4x <= 2`.

```julia-repl
julia> con
con : 5 x ≤ 2.0

julia> set_normalized_coefficient(con, x, 4)

julia> con
con : 4 x ≤ 2.0
```
"""
function JuMP.set_normalized_coefficient(cref::InfOptConstraintRef,
                                         variable::GeneralVariableRef,
                                         value::Real)::Nothing
    # update the constraint expression and update the constraint
    old_constr = JuMP.constraint_object(cref)
    new_expr = _set_variable_coefficient!(JuMP.jump_function(old_constr),
                                          variable, value)
    if old_constr isa BoundedScalarConstraint
        new_constr = BoundedScalarConstraint(new_expr, JuMP.moi_set(old_constr),
                        parameter_bounds(old_constr),
                        original_parameter_bounds(old_constr))
    else
        new_constr = JuMP.ScalarConstraint(new_expr, JuMP.moi_set(old_constr))
    end
    _set_core_constraint_object(cref, new_constr)
    return
end

"""
    JuMP.normalized_coefficient(cref::InfOptConstraintRef,
                                variable::GeneralVariableRef)::Float64

Return the coefficient associated with `variable` in `constraint` after JuMP has
normalized the constraint into its standard form.
"""
function JuMP.normalized_coefficient(cref::InfOptConstraintRef,
                                     variable::GeneralVariableRef)::Float64
    constr = JuMP.constraint_object(cref)
    expr = JuMP.jump_function(constr)
    if expr isa GeneralVariableRef && expr == variable
        return 1.0
    elseif expr isa GeneralVariableRef
        return 0.0
    else
        return JuMP._affine_coefficient(expr, variable)
    end
end

# Return the appropriate constraint reference given the index and model
function _make_constraint_ref(model::InfiniteModel,
                              index::ConstraintIndex)::InfOptConstraintRef
    constr = model.constraints[index].constraint
    return InfOptConstraintRef(model, index)
end

"""
    JuMP.constraint_by_name(model::InfiniteModel,
                            name::String)::Union{InfOptConstraintRef, Nothing}

Extend [`JuMP.constraint_by_name`](@ref JuMP.constraint_by_name)
to return the constraint reference
associated with `name` if one exists or returns nothing. Errors if more than
one constraint uses the same name.

**Example**
```julia-repl
julia> constraint_by_name(model, "constr_name")
constr_name : x + pt = 3.0
```
"""
function JuMP.constraint_by_name(model::InfiniteModel,
    name::String)::Union{InfOptConstraintRef, Nothing}
    if model.name_to_constr === nothing
        # Inspired from MOI/src/Utilities/model.jl
        model.name_to_constr = Dict{String, Int}()
        for (index, data_object) in model.constraints
            constr_name = data_object.name
            if haskey(model.name_to_constr, constr_name)
                model.name_to_constr[constr_name] = ConstraintIndex(-1)
            else
                model.name_to_constr[constr_name] = index
            end
        end
    end
    index = get(model.name_to_constr, name, nothing)
    if index === nothing
        return nothing
    elseif index == ConstraintIndex(-1)
        error("Multiple constraints have the name $name.")
    else
        return _make_constraint_ref(model, index)
    end
end

"""
    JuMP.num_constraints(model::InfiniteModel,
                         function_type::Type{<:JuMP.AbstractJuMPScalar},
                         set_type::Type{<:MOI.AbstractSet})::Int

Extend [`JuMP.num_constraints`](@ref JuMP.num_constraints(::JuMP.Model, ::Type{<:Union{JuMP.AbstractJuMPScalar, Vector{<:JuMP.AbstractJuMPScalar}}}, ::Type{<:MOI.AbstractSet}))
to return the number of constraints
with a partiuclar function type and set type.

**Example**
```julia-repl
julia> num_constraints(model, HoldVariableRef, MOI.LessThan)
1
```
"""
function JuMP.num_constraints(model::InfiniteModel,
                              function_type::Type{<:JuMP.AbstractJuMPScalar},
                              set_type::Type{<:MOI.AbstractSet})::Int
    counter = 0
    for (index, data_object) in model.constraints
        if isa(JuMP.jump_function(data_object.constraint), function_type) &&
           isa(JuMP.moi_set(data_object.constraint), set_type)
            counter += 1
        end
    end
    return counter
end

"""
    JuMP.num_constraints(model::InfiniteModel,
                         function_type::Type{<:JuMP.AbstractJuMPScalar})::Int

Extend [`JuMP.num_constraints`](@ref JuMP.num_constraints(::JuMP.Model, ::Type{<:Union{JuMP.AbstractJuMPScalar, Vector{<:JuMP.AbstractJuMPScalar}}}, ::Type{<:MOI.AbstractSet}))
to search by function types for all MOI
sets and return the total number of constraints with a particular function type.

```julia-repl
julia> num_constraints(model, HoldVariableRef)
3
```
"""
function JuMP.num_constraints(model::InfiniteModel,
                            function_type::Type{<:JuMP.AbstractJuMPScalar})::Int
    return JuMP.num_constraints(model, function_type, MOI.AbstractSet)
end

"""
    JuMP.num_constraints(model::InfiniteModel,
                         function_type::Type{<:MOI.AbstractSet})::Int

Extend [`JuMP.num_constraints`](@ref JuMP.num_constraints(::JuMP.Model, ::Type{<:Union{JuMP.AbstractJuMPScalar, Vector{<:JuMP.AbstractJuMPScalar}}}, ::Type{<:MOI.AbstractSet}))
to search by MOI set type for all function
types and return the total number of constraints that use a particular MOI set
type.

```julia-repl
julia> num_constraints(model, MOI.LessThan)
2
```
"""
function JuMP.num_constraints(model::InfiniteModel,
                              set_type::Type{<:MOI.AbstractSet})::Int
    return JuMP.num_constraints(model, JuMP.AbstractJuMPScalar, set_type)
end

"""
    JuMP.num_constraints(model::InfiniteModel)::Int

Extend [`JuMP.num_constraints`](@ref JuMP.num_constraints(::JuMP.Model, ::Type{<:Union{JuMP.AbstractJuMPScalar, Vector{<:JuMP.AbstractJuMPScalar}}}, ::Type{<:MOI.AbstractSet}))
to return the total number of constraints
in an infinite model `model`.

```julia-repl
julia> num_constraints(model)
4
```
"""
function JuMP.num_constraints(model::InfiniteModel)::Int
    return length(model.constraints)
end

"""
    JuMP.all_constraints(model::InfiniteModel,
                         function_type::Type{<:JuMP.AbstractJuMPScalar},
                         set_type::Type{<:MOI.AbstractSet}
                         )::Vector{InfOptConstraintRef}

Extend [`JuMP.all_constraints`](@ref JuMP.all_constraints(::JuMP.Model, ::Type{<:Union{JuMP.AbstractJuMPScalar, Vector{<:JuMP.AbstractJuMPScalar}}}, ::Type{<:MOI.AbstractSet}))
to return a list of all the constraints with a particular function type and set type.

```julia-repl
julia> all_constraints(model, GeneralVariableRef, MOI.LessThan)
1-element Array{InfOptConstraintRef,1}:
 x ≤ 1.0
```
"""
function JuMP.all_constraints(model::InfiniteModel,
                              function_type::Type{<:JuMP.AbstractJuMPScalar},
                              set_type::Type{<:MOI.AbstractSet}
                              )::Vector{InfOptConstraintRef}
    constr_list = Vector{InfOptConstraintRef}(undef,
                           JuMP.num_constraints(model, function_type, set_type))
    counter = 1
    for (index, object) in model.constraints
        if isa(JuMP.jump_function(object.constraint), function_type) &&
           isa(JuMP.moi_set(object.constraint), set_type)
            constr_list[counter] = _make_constraint_ref(model, index)
            counter += 1
        end
    end
    return constr_list
end

"""
    JuMP.all_constraints(model::InfiniteModel,
                         function_type::Type{<:JuMP.AbstractJuMPScalar}
                         )::Vector{InfOptConstraintRef}

Extend [`JuMP.all_constraints`](@ref JuMP.all_constraints(::JuMP.Model, ::Type{<:Union{JuMP.AbstractJuMPScalar, Vector{<:JuMP.AbstractJuMPScalar}}}, ::Type{<:MOI.AbstractSet}))
to search by function types for all MOI
sets and return a list of all constraints use a particular function type.

```julia-repl
julia> all_constraints(model, GeneralVariableRef)
3-element Array{InfOptConstraintRef,1}:
 x ≥ 0.0
 x ≤ 3.0
 x integer
```
"""
function JuMP.all_constraints(model::InfiniteModel,
                              function_type::Type{<:JuMP.AbstractJuMPScalar}
                              )::Vector{InfOptConstraintRef}
    return JuMP.all_constraints(model, function_type, MOI.AbstractSet)
end

"""
    JuMP.all_constraints(model::InfiniteModel,
                         set_type::Type{<:MOI.AbstractSet}
                         )::Vector{InfOptConstraintRef}

Extend [`JuMP.all_constraints`](@ref JuMP.all_constraints(::JuMP.Model, ::Type{<:Union{JuMP.AbstractJuMPScalar, Vector{<:JuMP.AbstractJuMPScalar}}}, ::Type{<:MOI.AbstractSet}))
to search by MOI set type for all function
types and return a list of all constraints that use a particular set type.

```julia-repl
julia> all_constraints(model, MOI.GreaterThan)
3-element Array{InfOptConstraintRef,1}:
 x ≥ 0.0
 g(t) ≥ 0.0, ∀ t ∈ [0, 6]
 g(0.5) ≥ 0.0
```
"""
function JuMP.all_constraints(model::InfiniteModel,
                              set_type::Type{<:MOI.AbstractSet}
                              )::Vector{InfOptConstraintRef}
    return JuMP.all_constraints(model, JuMP.AbstractJuMPScalar, set_type)
end

"""
    JuMP.all_constraints(model::InfiniteModel)::Vector{InfOptConstraintRef}

Extend [`JuMP.all_constraints`](@ref JuMP.all_constraints(::JuMP.Model, ::Type{<:Union{JuMP.AbstractJuMPScalar, Vector{<:JuMP.AbstractJuMPScalar}}}, ::Type{<:MOI.AbstractSet}))
to return all a list of all the constraints
in an infinite model `model`.

```julia-repl
julia> all_constraints(model)
5-element Array{InfOptConstraintRef,1}:
 x ≥ 0.0
 x ≤ 3.0
 x integer
 g(t) ≥ 0.0, ∀ t ∈ [0, 6]
 g(0.5) ≥ 0.0
```
"""
function JuMP.all_constraints(model::InfiniteModel)::Vector{InfOptConstraintRef}
    return [_make_constraint_ref(model, index)
            for (index, object) in model.constraints]
end

"""
    JuMP.list_of_constraint_types(model::InfiniteModel)::Vector{Tuple{DataType, DataType}}

Extend [`JuMP.list_of_constraint_types`](@ref JuMP.list_of_constraint_types(::JuMP.Model))
to return a list of tuples that
contain all the used combinations of function types and set types in the model.

```julia-repl
julia> all_constraints(model)
3-element Array{Tuple{DataType,DataType},1}:
 (GeneralVariableRef, MathOptInterface.LessThan{Float64})
 (GeneralVariableRef, MathOptInterface.GreaterThan{Float64})
 (GeneralVariableRef, MathOptInterface.Integer)
```
"""
function JuMP.list_of_constraint_types(model::InfiniteModel
    )::Vector{Tuple{DataType, DataType}}
    type_set = Set{Tuple{DataType, DataType}}()
    for (index, object) in model.constraints
        push!(type_set, (typeof(JuMP.jump_function(object.constraint)),
                          typeof(JuMP.moi_set(object.constraint))))
    end
    return collect(type_set)
end

################################################################################
#                         PARAMETER REFERENCE METHODS
################################################################################
"""
    parameter_refs(cref::InfOptConstraintRef)::Tuple

Return the tuple of infinite parameter references that determine the infinite
dependencies of `cref`.

**Example**
```julia-repl
julia> parameter_refs(cref)
(t,)
```
"""
function parameter_refs(cref::InfOptConstraintRef)::Tuple
    model = JuMP.owner_model(cref)
    obj_indices = _param_object_indices(model)[_object_numbers(cref)]
    return Tuple(_make_param_tuple_element(model, idx) for idx in obj_indices)
end

################################################################################
#                           PARAMETER BOUND METHODS
################################################################################
"""
    has_parameter_bounds(cref::InfOptConstraintRef)::Bool

Return a `Bool` indicating if `cref` is limited to a sub-domain as defined
by a [`ParameterBounds`](@ref) object.

**Example**
```julia-repl
julia> has_parameter_bounds(cref)
true
```
"""
function has_parameter_bounds(cref::InfOptConstraintRef)::Bool
    if JuMP.constraint_object(cref) isa BoundedScalarConstraint
        return !isempty(JuMP.constraint_object(cref).bounds)
    else
        return false
    end
end

"""
    parameter_bounds(cref::InfOptConstraintRef)::ParameterBounds{GeneralVariableRef}

Return the [`ParameterBounds`](@ref) object associated with the constraint
`cref`. Errors if `cref` does not have parameter bounds.

**Example**
```julia-repl
julia> parameter_bounds(cref)
Subdomain bounds (1): t ∈ [0, 2]
```
"""
function parameter_bounds(cref::InfOptConstraintRef)::ParameterBounds{GeneralVariableRef}
    !has_parameter_bounds(cref) && error("$cref does not have parameter bounds.")
    return JuMP.constraint_object(cref).bounds
end

# Internal function used to change the parameter bounds of a constraint
function _update_constr_param_bounds(cref::InfOptConstraintRef,
                                     bounds::ParameterBounds{GeneralVariableRef},
                                     orig_bounds::ParameterBounds{GeneralVariableRef}
                                     )::Nothing
    old_constr = JuMP.constraint_object(cref)
    func = JuMP.jump_function(old_constr)
    set = JuMP.moi_set(old_constr)
    if !isempty(bounds)
        new_constr = BoundedScalarConstraint(func, set, bounds, orig_bounds)
    else
        new_constr = JuMP.ScalarConstraint(func, set)
    end
    _set_core_constraint_object(cref, new_constr)
    return
end

"""
    set_parameter_bounds(cref::InfOptConstraintRef,
                         bounds:ParameterBounds{GeneralVariableRef};
                         [force::Bool = false])::Nothing

Specify a new [`ParameterBounds`](@ref) object `bounds` for the constraint `cref`.
This is meant to be primarily used by [`@set_parameter_bounds`](@ref) which
provides a more intuitive syntax.

**Example**
```julia-repl
julia> set_parameter_bounds(cref, ParameterBounds((t => IntervalSet(0, 2),)))

julia> parameter_bounds(cref)
Subdomain bounds (1): t ∈ [0, 2]
```
"""
function set_parameter_bounds(cref::InfOptConstraintRef,
                              bounds::ParameterBounds{GeneralVariableRef};
                              force::Bool = false, _error::Function = error
                              )::Nothing
    if has_parameter_bounds(cref) && !force
        _error("$cref already has parameter bounds. Consider adding more using " *
               "`add_parameter_bounds` or overwriting them by setting " *
               "the keyword argument `force = true`")
    else
        # check that bounds are valid and add support(s) if necessary
        _check_bounds(bounds, _error = _error)
        _validate_bounds(JuMP.owner_model(cref), bounds, _error = _error)
        # consider hold variables
        orig_bounds = copy(bounds)
        vrefs = _all_function_variables(JuMP.jump_function(JuMP.constraint_object(cref)))
        for vref in vrefs
            _update_var_bounds(vref, bounds)
        end
        # set the new bounds
        _update_constr_param_bounds(cref, bounds, orig_bounds)
        # update status
        set_optimizer_model_ready(JuMP.owner_model(cref), false)
    end
    return
end

"""
    add_parameter_bounds(cref::InfOptConstraintRef,
                         new_bounds::ParameterBounds{GeneralVariableRef}
                         )::Nothing

Add an additional parameter bound to `cref` such that it is defined over the
sub-domain based on `pref` from `lower` to `upper`. This is primarily meant to be
used by [`@add_parameter_bounds`](@ref).

```julia-repl
julia> add_parameter_bounds(cref, ParameterBounds((t => IntervalSet(0, 2),))

julia> parameter_bounds(cref)
Subdomain bounds (1): t ∈ [0, 2]
```
"""
function add_parameter_bounds(cref::InfOptConstraintRef,
                              new_bounds::ParameterBounds{GeneralVariableRef};
                              _error::Function = error)::Nothing
    # check the new bounds
    _check_bounds(new_bounds, _error = _error)
    _validate_bounds(JuMP.owner_model(cref), new_bounds, _error = _error)
    # add the bounds
    if JuMP.constraint_object(cref) isa BoundedScalarConstraint
        _update_bounds(parameter_bounds(cref), new_bounds, _error = _error)
        _update_bounds(original_parameter_bounds(JuMP.constraint_object(cref)),
                       new_bounds, _error = _error)
    else
        _update_constr_param_bounds(cref, new_bounds, copy(new_bounds))
    end
    # update the optimizer model status
    set_optimizer_model_ready(JuMP.owner_model(cref), false)
    return
end

"""
    delete_parameter_bounds(cref::InfOptConstraintRef)::Nothing

Delete all the parameter bounds of the constraint `cref`. Note any bounds that
are needed for hold variables inside in `cref` will be unaffected.

**Example**
```julia-repl
julia> @BDconstraint(model, c1(x == 0), y <= 42)
c1 : y(x) ≤ 42, ∀ x[1] = 0, x[2] = 0

julia> delete_parameter_bounds(c1)

julia> c1
c1 : y(x) ≤ 42, ∀ x[1] ∈ [-1, 1], x[2] ∈ [-1, 1]
```
"""
function delete_parameter_bounds(cref::InfOptConstraintRef)::Nothing
    # check if has bounds on pref from the constraint
    constr = JuMP.constraint_object(cref)
    if has_parameter_bounds(cref) && !isempty(original_parameter_bounds(constr))
        # consider hold variables
        new_bounds = ParameterBounds()
        vrefs = _all_function_variables(JuMP.jump_function(constr))
        for vref in vrefs
            _update_var_bounds(vref, new_bounds)
        end
        # set the new bounds
        _update_constr_param_bounds(cref, new_bounds, ParameterBounds())
    end
    return
end

################################################################################
#                                 DELETION
################################################################################
"""
    JuMP.delete(model::InfiniteModel, cref::InfOptConstraintRef)::Nothing

Extend [`JuMP.delete`](@ref JuMP.delete(::JuMP.Model, ::JuMP.ConstraintRef{JuMP.Model}))
to delete an `InfiniteOpt` constraint and all associated information. Errors
if `cref` is invalid.

**Example**
```julia-repl
julia> print(model)
Min measure(g(t)*t) + z
Subject to
 z ≥ 0.0
 g(t) + z ≥ 42.0, ∀ t ∈ [0, 6]

julia> delete(model, cref)

julia> print(model)
Min measure(g(t)*t) + z
Subject to
 z ≥ 0.0
```
"""
function JuMP.delete(model::InfiniteModel, cref::InfOptConstraintRef)::Nothing
    # check valid reference
    @assert JuMP.is_valid(model, cref) "Invalid constraint reference."
    # update variable dependencies
    constr = _core_constraint_object(cref)
    all_vrefs = _all_function_variables(JuMP.jump_function(constr))
    for vref in all_vrefs
        filter!(e -> e != JuMP.index(cref), _constraint_dependencies(vref))
    end
    # delete constraint information
    _delete_data_object(cref)
    # reset optimizer model status
    set_optimizer_model_ready(model, false)
    return
end


#!/usr/bin/env julia

#=
 Licensing Information: You are free to use or extend these projects for
 education or reserach purposes provided that (1) you retain this notice
 and (2) you provide clear attribution to UC Berkeley, including a link
 to http://barc-project.com

 Attibution Information: The barc project ROS code-base was developed
 at UC Berkeley in the Model Predictive Control (MPC) lab by Jon Gonzales
 (jon.gonzales@berkeley.edu). The cloud services integation with ROS was developed
 by Kiet Lam  (kiet.lam@berkeley.edu). The web-server app Dator was
 based on an open source project by Bruce Wootton
=#

using RobotOS
@rosimport barc.msg: ECU, Vel_est, pos_info
@rosimport geometry_msgs.msg: Vector3
@rosimport sensor_msgs.msg: Imu, NavSatFix
# @rosimport marvelmind_nav.msg: hedge_pos
@rosimport std_msgs.msg: Header
rostypegen()
using barc.msg
using geometry_msgs.msg
using sensor_msgs.msg
using std_msgs.msg
using marvelmind_nav.msg
using JLD

include("barc_lib/classes.jl")
include("barc_lib/simModel.jl")

# This type contains measurement data (time, values and a counter)
type Measurements{T}
    i::Int64          # measurement counter
    t::Array{Float64}       # time data
    t_msg::Array{Float64}
    z::Array{T}       # measurement values
end

# This function cleans the zeros from the type above once the simulation is finished
function clean_up(m::Measurements)
    m.t     = m.t[1:m.i-1]
    m.t_msg = m.t_msg[1:m.i-1]
    m.z     = m.z[1:m.i-1,:]
end

function ECU_callback(msg::ECU,u_current::Array{Float64},cmd_log::Measurements)
    u_current[:] = convert(Array{Float64,1},[msg.motor, msg.servo])
    cmd_log.t_msg[cmd_log.i]    = to_sec(get_rostime())
    cmd_log.t[cmd_log.i]        = to_sec(get_rostime())
    cmd_log.z[cmd_log.i,:]      = u_current
    cmd_log.i += 1
end

function main()
    u_current = zeros(Float64,2)      # msg ECU is Float32 !

    buffersize = 60000
    gps_meas = Measurements{Float64}(1,zeros(buffersize),zeros(buffersize),zeros(buffersize,2))
    imu_meas = Measurements{Float64}(1,zeros(buffersize),zeros(buffersize),zeros(buffersize,2))
    cmd_log  = Measurements{Float64}(1,ones(buffersize)*Inf,ones(buffersize)*Inf,zeros(buffersize,2))
    z_real   = Measurements{Float64}(1,zeros(buffersize),zeros(buffersize),zeros(buffersize,8))
    slip_a   = Measurements{Float64}(1,zeros(buffersize),zeros(buffersize),zeros(buffersize,2))

    # initiate node, set up publisher / subscriber topics
    init_node("barc_sim")
    pub_gps = Publisher("fix", hedge_pos, queue_size=1)::RobotOS.Publisher{marvelmind_nav.msg.hedge_pos}
    pub_imu = Publisher("imu/data", Imu, queue_size=1)::RobotOS.Publisher{sensor_msgs.msg.Imu}
    pub_vel = Publisher("vel_est", Vel_est, queue_size=1)::RobotOS.Publisher{barc.msg.Vel_est}
    real_val = Publisher("real_val", pos_info, queue_size=1)::RobotOS.Publisher{barc.msg.pos_info}

    s1  = Subscriber("ecu", ECU, ECU_callback, (u_current,cmd_log,), queue_size=1)::RobotOS.Subscriber{barc.msg.ECU}

    z_current = zeros(buffersize,8)
    z_current[1,:] = [0.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0]
    slip_ang = zeros(buffersize,2)

    dt = 0.01
    loop_rate = Rate(1/dt)

    i = 2

    dist_traveled = randn(3)        # encoder positions of three wheels
    last_updated  = 0.0

    r_tire      = 0.05             # radius from tire center to perimeter along magnets [m]

    imu_drift = 0.0       # simulates yaw-sensor drift over time (slow sine)

    modelParams     = ModelParams()
    run_id          = get_param("run_id")
    # modelParams.l_A = copy(get_param("L_a"))      # always throws segmentation faults *after* execution!!! ??
    # modelParams.l_B = copy(get_param("L_a"))
    # modelParams.m   = copy(get_param("m"))
    # modelParams.I_z = copy(get_param("I_z"))

    modelParams.l_A = 0.18
    modelParams.l_B = 0.14
    modelParams.m   = 1.98
    modelParams.I_z = 0.03#0.24             # using homogenous distributed mass over a cuboid

    println("Publishing sensor information. Simulator running.")
    imu_data    = Imu()
    vel_est     = Vel_est()
    t0          = to_sec(get_rostime())
    gps_data    = hedge_pos()
    real_data   = pos_info()

    z_real.t_msg[1] = t0
    slip_a.t_msg[1] = t0
    z_real.t[1]     = t0
    slip_a.t[1]     = t0
    t               = 0.0

    sim_gps_interrupt   = 0                 # counter if gps interruption is simulated
    vel_dist_update     = 2*pi*r_tire/8      # distance to travel until velocity is updated (1/8 wheel rotation)

    gps_header = Header()
    while ! is_shutdown()
        t_ros   = get_rostime()
        t       = to_sec(t_ros)
        # if sizeof(cmd_log.z[t.>cmd_log.t+0.2,2]) >= 1
        #    u_current[2] = cmd_log.z[t.>=cmd_log.t+0.2,2][end]       # artificial steering input delay
        # end
        # update current state with a new row vector
        z_current[i,:],slip_ang[i,:]  = simDynModel_exact_xy(z_current[i-1,:], u_current', dt, modelParams)

        z_real.t_msg[i] = t
        z_real.t[i]     = t
        slip_a.t_msg[i] = t
        slip_a.t[i]     = t

        # IMU measurements
        if i%2 == 0                 # 50 Hz
            imu_drift   = 1+(t-t0)/100#sin(t/100*pi/2)     # drifts to 1 in 100 seconds (and add random start value 1)
            rand_yaw = 0.05*randn()
            if rand_yaw > 0.1
                rand_yaw = 0.1
            elseif rand_yaw<-0.1
                rand_yaw=-0.1
            end

            yaw         = z_current[i,5] + imu_drift + rand_yaw#+ 0.002*randn()

            rand_psiDot = 0.01*randn()
            if rand_psiDot > 0.1
                rand_psiDot = 0.1
            elseif rand_psiDot<-0.1
                rand_psiDot=-0.1
            end
            psiDot      = z_current[i,6] +rand_psiDot#+ 0.001*randn()
            imu_meas.t_msg[imu_meas.i] = t
            imu_meas.t[imu_meas.i] = t
            imu_meas.z[imu_meas.i,:] = [yaw psiDot]
            imu_meas.i += 1
            imu_data.orientation = geometry_msgs.msg.Quaternion(cos(yaw/2), sin(yaw/2), 0, 0)
            imu_data.angular_velocity = Vector3(0,0,psiDot)
            imu_data.header.stamp = t_ros

            rand_accX = 0.01*randn()
            if rand_accX > 0.1
                rand_accX = 0.1
            elseif rand_accX<-0.1
                rand_accX=-0.1
            end

            imu_data.linear_acceleration.x = diff(z_current[i-1:i,3])[1]/dt - z_current[i,6]*z_current[i,4] #+rand_accX#+ randn()*0.3*1.0

            rand_accY = 0.01*randn()
            if rand_accY > 0.1
                rand_accY = 0.1
            elseif rand_accY<-0.1
                rand_accY=-0.1
            end

            imu_data.linear_acceleration.y = diff(z_current[i-1:i,4])[1]/dt + z_current[i,6]*z_current[i,3] #+rand_accY#+ randn()*0.3*1.0
            publish(pub_imu, imu_data)      # Imu format is defined by ROS, you can look it up by google "rosmsg Imu"
                                            # It's sufficient to only fill the orientation part of the Imu-type (with one quaternion)
        end

        # real values
        if i%2 == 0
            real_data.psiDot = z_current[i,6]
            real_data.psi    = z_current[i,5]
            real_data.v_x    = z_current[i,3]
            real_data.v_y    = z_current[i,4]
            real_data.x      = z_current[i,1]
            real_data.y      = z_current[i,2]
            publish(real_val,real_data)
        end

        # Velocity measurements
        dist_traveled += norm(diff(z_current[i-1:i,1:2]))
        if i%5 == 0                 # 20 Hz
            if sum(dist_traveled .>= vel_dist_update)>=1 && z_current[i,3] > 0.1     # only update if at least one of the magnets has passed the sensor
                dist_traveled[dist_traveled.>=vel_dist_update] = 0
                vel_est.vel_est = convert(Float32,norm(z_current[i,3:4]))#+0.00*randn())
                vel_est.vel_fl = convert(Float32,0)
                vel_est.vel_fr = convert(Float32,0)
                vel_est.vel_bl = convert(Float32,0)
                vel_est.vel_br = convert(Float32,0)
            end
            vel_est.header.stamp = t_ros
            publish(pub_vel, vel_est)
        end

        # GPS measurements
        if i%6 == 0               # 16 Hz

            rand_x = 0.01*randn()
            if rand_x > 0.1
                rand_x = 0.1
            elseif rand_x<-0.1
                rand_x=-0.1
            end

            x = round(z_current[i,1] +  rand_x,2)#0.002*randn(),2)       # Indoor gps measures, rounded on cm

            rand_y = 0.01*randn()
            if rand_y > 0.1
                rand_y = 0.1
            elseif rand_y<-0.1
                rand_y=-0.1
            end

            y = round(z_current[i,2] + rand_y,2)#0.002*randn(),2)

            if randn()>10            # simulate gps-outlier (probability about 0.13% for randn()>3, 0.62% for randn()>2.5, 2.3% for randn()>2.0 )
                x += 1#randn()        # add random value to x and y
                y -= 1#randn()
                #sim_gps_interrupt = 6       # also do a little interruption
            elseif randn()>10 && sim_gps_interrupt < 0
                sim_gps_interrupt = 10      # simulate gps-interrupt (10 steps at 25 Hz is 0.4 seconds)
            end
            if sim_gps_interrupt < 0
                gps_meas.t_msg[gps_meas.i] = t
                gps_meas.t[gps_meas.i] = t
                gps_meas.z[gps_meas.i,:] = [x y]
                gps_meas.i += 1
                # gps_data.header.stamp = get_rostime()
                gps_data.x_m = x
                gps_data.y_m = y
                publish(pub_gps, gps_data)
            end
            sim_gps_interrupt -= 1
        end
        i += 1
        rossleep(loop_rate)
    end

    # Clean up buffers
    clean_up(gps_meas)
    clean_up(imu_meas)
    clean_up(cmd_log)
    z_real.z[1:i-1,:] = z_current[1:i-1,:]
    slip_a.z[1:i-1,:] = slip_ang[1:i-1,:]
    z_real.i = i
    slip_a.i = i
    clean_up(z_real)
    clean_up(slip_a)

    # Save simulation data to file
    # log_path = "$(homedir())/simulations/output-SIM-$(run_id[1:4]).jld"
    # save(log_path,"gps_meas",gps_meas,"z",z_real,"imu_meas",imu_meas,"cmd_log",cmd_log,"slip_a",slip_a)
    println("Exiting node... Saving data to $log_path. Simulated $((i-1)*dt) seconds.")

end

if ! isinteractive()
    main()
end


# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# EAGO
# A development environment for robust and global optimization
# See https://github.com/PSORLab/EAGO.jl
#############################################################################
# src/eago_optimizer/optimizer.jl
# Defines optimizer structure used by EAGO. Namely, ObjectiveType, ProblemType
# EAGOParameters, InputProblem, ParsedProblem, and Optimizer.
#############################################################################

#=
Holds outputs in the Optimizer object, passing parameters EAGO specific parameters
to _options, the GlobalOptimizer is remade to ensure type stability of subroutines
prior to initializing a global solve.
=#
export Optimizer
mutable struct Optimizer{T<:AbstractFloat} <: MOI.AbstractOptimizer
    _model::InputModel{T}
    _options::GlobalOptimizerOptions{T}
    _solver::GlobalOptimizer
    _problem_type::ProblemType

    _termination_status_code::MOI.TerminationStatusCode
    _primal_status_code::MOI.ResultStatusCode

    _objective_value::T
    _objective_bound::T
    _relative_gap::T
    _iteration_count::Int
    _node_count::Int

    _is_silent::Bool

    _start_time::T
    _time_limit::T
    _time_left::T
    _run_time::T
    _parse_time::T

    _solution::Vector{T}
    _constraint_primal::Dict{CI,Union{Vector{T},T}}
    _input_to_solution_map
end
function Optimizer{T}(; kwargs...) where T<:Real
    model   = InputModel{T}()
    options = GlobalOptimizerOptions{T}(;kwargs...)
    solver  = GlobalOptimizer(T)
    problem_type = UNCLASSIFIED

    termination_status_code = MOI.OPTIMIZE_NOT_CALLED
    primal_status_code      = MOI.OTHER_RESULT_STATUS

    objective_value = Inf
    objective_bound = -Inf
    relative_gap = Inf
    iteration_count = 0
    node_count      = 0

    is_silent       = false

    start_time   = 0.0
    time_limit   = Inf
    time_left    = 1000.0
    run_time     = 0.0
    parse_time   = 0.0

    solution          = T[]
    constraint_primal = Dict{CI,T}()
    input_to_solution_map = nothing

    return Optimizer{T}(model,
                        options,
                        solver,
                        problem_type,
                        termination_status_code,
                        primal_status_code,
                        objective_value,
                        objective_bound,
                        relative_gap,
                        iteration_count,
                        node_count,
                        is_silent,
                        start_time,
                        time_limit,
                        time_left,
                        run_time,
                        parse_time,
                        solution,
                        constraint_primal,
                        input_to_solution_map)
end
Optimizer(; kwargs...) = Optimizer{Float64}(; kwargs...)

for f in fieldnames(GlobalOptimizerOptions)
    @eval @inline $(Symbol("_"*String(f)))(x::Optimizer) = x._options.$f
end

_input_model(m::Optimizer{T}) where T = m._model._input_model
_input_nlp_data(m::Optimizer) = m._model._nlp_data

function _set_cons_primal!(m::Optimizer{T}, ci::CI{F,S}, v::T) where {F <: MOI.AbstractFunction,
                                                                     S <: MOI.AbstractScalarSet,
                                                                     T <: AbstractFloat}
    m._constraint_primal[ci] = v
    return
end

function _set_cons_primal!(m::Optimizer{T}, ci::CI{F,S}, v::Vector{T}) where {F <: MOI.AbstractFunction,
                                                                              S <: MOI.AbstractVectorSet,
                                                                              T <: AbstractFloat}
    m._constraint_primal[ci] = v
    return
end


# TernaryLayer and TernaryMERA types, and methods thereof.
# To be `included` in MERAKit.jl.

"""
    TernaryLayer{ST, ET, Tan} <: SimpleLayer

The type for layers of a ternary MERA.

Each layer consists of two tensors, a 2-to-2 disentangler, often called `u`, and a 3-to-1
isometry, often called `w`.

The type parameters are `ST` for space type, e.g. `ComplexSpace` or `SU2Space`; `ET` for
element type, e.g. `Complex{Float64}`; and `Tan` for whether this layer is a tangent layer
or not.  If `Tan = false`, the layer in question is an actual MERA layer. If `Tan = true` it
consists, instead of the actual tensors, of Stiefel/Grassmann tangent vectors of these
tensors.

`TernaryLayer` implements the iteration interface, returning first the disentangler, then
the isometry.

Index numbering convention is as follows, where the physical indices are at the bottom:
Disentangler:
```
 3│    4│
 ┌┴─────┴┐
 │   u   │
 └┬─────┬┘
 1│    2│
```

Isometry:
```
    4│
 ┌───┴───┐
 │   w   │
 └┬──┬──┬┘
 1│ 2│ 3│
```
"""
struct TernaryLayer{ST, ET, Tan} <: SimpleLayer
    # Even though the types are here marked as any, they are restricted to be specific,
    # concrete types dependent on ST, ET and Tan, both in the constructor and in
    # getproperty.
    disentangler::Any
    isometry::Any

    function TernaryLayer{ST, ET, Tan}(disentangler, isometry) where {ST, ET, Tan}
        DisType = disentangler_type(ST, ET, Tan)
        IsoType = ternaryisometry_type(ST, ET, Tan)
        disconv = convert(DisType, disentangler)::DisType
        isoconv = convert(IsoType, isometry)::IsoType
        return new{ST, ET, Tan}(disconv, isoconv)
    end
end

function TernaryLayer(disentangler::DisType, isometry::IsoType) where {
    ST,
    DisType <: AbstractTensorMap{ST, 2, 2},
    IsoType <: AbstractTensorMap{ST, 3, 1}
}
    ET = eltype(DisType)
    @assert eltype(IsoType) === ET
    return TernaryLayer{ST, ET, false}(disentangler, isometry)
end

function TernaryLayer(disentangler::DisTanType, isometry::IsoTanType) where {
    ST,
    DisType <: AbstractTensorMap{ST, 2, 2},
    IsoType <: AbstractTensorMap{ST, 3, 1},
    DisTanType <: Stiefel.StiefelTangent{DisType},
    IsoTanType <: Grassmann.GrassmannTangent{IsoType},
}
    ET = eltype(DisType)
    @assert eltype(IsoType) === ET
    return TernaryLayer{ST, ET, true}(disentangler, isometry)
end

function Base.getproperty(l::TernaryLayer{ST, ET, Tan}, sym::Symbol) where {ST, ET, Tan}
    if sym === :disentangler
        T = disentangler_type(ST, ET, Tan)
    elseif sym === :isometry
        T = ternaryisometry_type(ST, ET, Tan)
    else
        T = Any
    end
    return getfield(l, sym)::T
end

"""
    TernaryMERA{N}

A ternary MERA is a MERA consisting of `TernaryLayer`s.
"""
TernaryMERA{N} = GenericMERA{N, T, O} where {T <: TernaryLayer, O}
layertype(::Type{TernaryMERA}) = TernaryLayer
#Base.show(io::IO, ::Type{TernaryMERA}) = print(io, "TernaryMERA")
#Base.show(io::IO, ::Type{TernaryMERA{N}}) where {N} = print(io, "TernaryMERA{($N)}")

# Implement the iteration and indexing interfaces.
# See simplelayer.jl for details.
_tuple(layer::TernaryLayer) = (layer.disentangler, layer.isometry)

function operatortype(::Type{TernaryLayer{ST, ET, false}}) where {ST, ET}
    return tensormaptype(ST, 2, 2, ET)
end
operatortype(::Type{TernaryLayer{ST, ET, true}}) where {ST, ET} = Nothing

TensorKit.spacetype(::Type{TernaryLayer{ST, ET, Tan}}) where {ST, ET, Tan} = ST

scalefactor(::Type{<:TernaryLayer}) = 3

causal_cone_width(::Type{<:TernaryLayer}) = 2

Base.eltype(::Type{TernaryLayer{ST, ET, Tan}}) where {ST, ET, Tan} = ET

outputspace(layer::TernaryLayer) = space(layer.disentangler, 1)
inputspace(layer::TernaryLayer) = space(layer.isometry, 4)'
internalspace(layer::TernaryLayer) = space(layer.isometry, 1)
internalspace(m::TernaryMERA, depth) = internalspace(getlayer(m, depth))

function expand_inputspace(layer::TernaryLayer, V_new)
    u, w = layer
    w = pad_with_zeros_to(w, 4 => V_new')
    return TernaryLayer(u, w)
end

function expand_outputspace(layer::TernaryLayer, V_new)
    u, w = layer
    u = pad_with_zeros_to(u, 1 => V_new, 2 => V_new)
    w = pad_with_zeros_to(w, 2 => V_new)
    return TernaryLayer(u, w)
end

function expand_internalspace(layer::TernaryLayer, V_new)
    u, w = layer
    u = pad_with_zeros_to(u, 3 => V_new', 4 => V_new')
    w = pad_with_zeros_to(w, 1 => V_new, 3 => V_new)
    return TernaryLayer(u, w)
end

function randomlayer(
    ::Type{TernaryLayer}, ::Type{T}, Vin, Vout, Vint = Vout; random_disentangler = false
) where {T}
    w = randomisometry(T, Vint ⊗ Vout ⊗ Vint, Vin)
    u = initialize_disentangler(T, Vout, Vint, random_disentangler)
    return TernaryLayer(u, w)
end

function ascending_fixedpoint(layer::TernaryLayer)
    width = causal_cone_width(layer)
    Vtotal = ⊗(ntuple(n->inputspace(layer), Val(width))...)
    return id(storagetype(operatortype(layer)), Vtotal)
end

function scalingoperator_initialguess(l::TernaryLayer, irreps...)
    width = causal_cone_width(l)
    V = inputspace(l)
    interlayer_space = ⊗(ntuple(n->V, Val(width))...)
    outspace = interlayer_space
    inspace = interlayer_space
    for irrep in irreps
        inspace = inspace ⊗ spacetype(V)(irrep => 1)
    end
    typ = eltype(l)
    t = TensorMap(randn, typ, outspace ← inspace)
    return t
end

function gradient(layer::TernaryLayer, env::TernaryLayer; metric = :euclidean)
    u, w = layer
    uenv, wenv = env
    # The environment is the partial derivative. We need to turn that into a tangent vector
    # of the Stiefel manifold point u or w.
    # The factor of two is from the partial_x + i partial_y derivative of the cost function,
    # and how it depends on both v and v^dagger.
    ugrad = Stiefel.project!(2*uenv, u; metric = metric)
    wgrad = Grassmann.project!(2*wenv, w)
    return TernaryLayer(ugrad, wgrad)
end

function precondition_tangent(layer::TernaryLayer, tan::TernaryLayer, rho)
    u, w = layer
    utan, wtan = tan
    @planar rho_wl[-1; -11] := rho[-1 1; -11 1]
    @planar rho_wr[-1; -11] := rho[1 -1; 1 -11]
    rho_w = (rho_wl + rho_wr) / 2.0
    @planar(
        rho_u[-1 -2; -11 -12] :=
        w'[12; 1 2 -11] * w'[22; -12 3 4] *
        rho[11 21; 12 22] *
        w[1 2 -1; 11] * w[-2 3 4; 21]
    )
    utan_prec = precondition_tangent(utan, rho_u)
    wtan_prec = precondition_tangent(wtan, rho_w)
    return TernaryLayer(utan_prec, wtan_prec)
end

# # # Invariants

function space_invar_intralayer(layer::TernaryLayer)
    u, w = layer
    matching_bonds = (
        (space(u, 3)', space(w, 3)),
        (space(u, 4)', space(w, 1)),
    )
    allmatch = all(pair -> ==(pair...), matching_bonds)
    # Check that the dimensions are such that isometricity can hold.
    allmatch &= all((u, w)) do v
        codom, dom = fuse(codomain(v)), fuse(domain(v))
        infimum(dom, codom) == dom
    end
    return allmatch
end

function space_invar_interlayer(layer::TernaryLayer, next_layer::TernaryLayer)
    u, w = layer.disentangler, layer.isometry
    unext, wnext = next_layer.disentangler, next_layer.isometry
    matching_bonds = (
        (space(w, 4)', space(unext, 1)),
        (space(w, 4)', space(unext, 2)),
        (space(w, 4)', space(wnext, 2)),
    )
    allmatch = all(pair -> ==(pair...), matching_bonds)
    return allmatch
end

# # # Ascending and descending superoperators

# TODO Fix this to work with anyons.
#"""
#Return the ascending superoperator of the one site in the middle of the isometries in a
#TernaryMERA, as a TensorMap. Unlike most ascending superoperators, this one is actually
#affordable to construct as a full tensor.
#"""
#ascending_superop_onesite(m::TernaryMERA) = ascending_superop_onesite(getlayer(m, Inf))
#
#function ascending_superop_onesite(layer::TernaryLayer)
#    w = layer.isometry
#    @planar superop[-1 -2; -11 -12] := w'[-11; 1 -1 2] * w[1 -2 2; -12]
#    return superop
#end

const TernaryOperator{S} = AbstractTensorMap{S, 2, 2}
const ChargedTernaryOperator{S} = AbstractTensorMap{S, 2, 3}
const DoubleChargedTernaryOperator{S} = AbstractTensorMap{S, 2, 4}

function ascend(
    op::Union{TernaryOperator, ChargedTernaryOperator, DoubleChargedTernaryOperator},
    layer::TernaryLayer
) where {S1}
    u, w = layer
    l = ascend_left(op, layer)
    r = ascend_right(op, layer)
    m = ascend_mid(op, layer)
    scaled_op = (l+r+m)/3
    return scaled_op
end

function ascend(op::SquareTensorMap{1}, layer::TernaryLayer)
    return ascend(expand_support(op, causal_cone_width(TernaryLayer)), layer)
end

# TODO Think about how to best remove the code duplication of having the separate methods
# for ordinary, charged, and double charged operators. Note also that there's a mix of using
# @tensor and @planar.
function ascend_left(op::ChargedTernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @tensor(
        scaled_op[-100 -200; -300 -400 -1000] :=
        w[51 52 53; -300] * w[54 11 12; -400] *
        u[41 42; 53 54] *
        op[31 32; 52 41 -1000] *
        u'[21 55; 32 42] *
        w'[-100; 51 31 21] * w'[-200; 55 11 12]
    )
    return scaled_op
end

function ascend_right(op::ChargedTernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @tensor(
        scaled_op[-100 -200; -300 -400 -1000] :=
        w[11 12 65; -300] * w[63 61 62; -400] *
        u[51 52; 65 63] *
        op[31 41; 52 61 -1000] *
        u'[64 21; 51 31] *
        w'[-100; 11 12 64] * w'[-200; 21 41 62]
    )
    return scaled_op
end

function ascend_mid(op::ChargedTernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 6X^6
    @tensor(
        scaled_op[-100 -200; -300 -400 -1000] :=
        w[31 32 41; -300] * w[51 21 22; -400] *
        u[1 2; 41 51] *
        op[11 12; 1 2 -1000] *
        u'[42 52; 11 12] *
        w'[-100; 31 32 42] * w'[-200; 52 21 22]
    )
    return scaled_op
end

function ascend_left(op::TernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        scaled_op[-100 -200; -300 -400] :=
        w[51 52 53; -300] * w[54 11 12; -400] *
        u[41 42; 53 54] *
        op[31 32; 52 41] *
        u'[21 55; 32 42] *
        w'[-100; 51 31 21] * w'[-200; 55 11 12]
    )
    return scaled_op
end

function ascend_right(op::TernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        scaled_op[-100 -200; -300 -400] :=
        w[11 12 65; -300] * w[63 61 62; -400] *
        u[51 52; 65 63] *
        op[31 41; 52 61] *
        u'[64 21; 51 31] *
        w'[-100; 11 12 64] * w'[-200; 21 41 62]
    )
    return scaled_op
end

function ascend_mid(op::TernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 6X^6
    @planar(
        scaled_op[-100 -200; -300 -400] :=
        w[31 32 41; -300] * w[51 21 22; -400] *
        u[1 2; 41 51] *
        op[11 12; 1 2] *
        u'[42 52; 11 12] *
        w'[-100; 31 32 42] * w'[-200; 52 21 22]
    )
    return scaled_op
end

function ascend_left(op::DoubleChargedTernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        temp1[-1 -2 -3 -4; -11 -12 -13 -14] :=
        op[31 32; -1 -11 -12 -13] *
        u'[21 -4; 32 -14] *
        w'[-3; -2 31 21]
    )
    temp2 = braid(temp1, (11, 12, 13, 14, 15, 1, 100, 16), (1, 2, 3, 7, 6, 4), (5, 8))
    @planar(
        temp3[-100 -1000 -2000 -200; -300 -400] :=
        w[51 52 53; -300] * w[54 11 12; -400] *
        u[41 42; 53 54] *
        temp2[52 51 -100 -1000 -2000 55; 41 42] *
        w'[-200; 55 11 12]
    )
    scaled_op = braid(temp3, (11, 100, 1, 12, 13, 14), (1, 4), (5, 6, 3, 2))
    return scaled_op
end

function ascend_right(op::DoubleChargedTernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        temp1[-1 -2 -3 -4 -5 -6; -11 -12] :=
        op[31 41; -12 -6 -5 -4] *
        u'[-1 21; -11 31] *
        w'[-2; 21 41 -3]
    )
    temp2 = braid(temp1, (11, 12, 13, 100, 1, 14, 15, 16), (1, 2, 4, 5, 3, 6), (7, 8))
    @planar(
        temp3[-100 -200 -1000 -2000; -300 -400] :=
        w[11 12 65; -300] * w[63 61 62; -400] *
        u[51 52; 65 63] *
        temp2[64 -200 -1000 -2000 62 61; 51 52] *
        w'[-100; 11 12 64]
    )
    scaled_op = braid(temp3, (11, 12, 100, 1, 13, 14), (1, 2), (5, 6, 4, 3))
    return scaled_op
end

function ascend_mid(op::DoubleChargedTernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 6X^6
    @planar(
        temp1[-1 -2; -3 -4 -5 -6] :=
        u[1 2; -3 -4] *
        op[11 12; 1 2 -5 -6] *
        u'[-1 -2; 11 12]
    )
    temp2 = braid(temp1, (11, 12, 13, 14, 1, 100), (1, 6, 5, 2), (3, 4))
    @planar(
        temp3[-100 -1000 -2000 -200; -300 -400] :=
        w[31 32 41; -300] * w[51 21 22; -400] *
        temp2[42 -1000 -2000 52; 41 51] *
        w'[-100; 31 32 42] * w'[-200; 52 21 22]
    )
    scaled_op = braid(temp3, (11, 100, 1, 12, 13, 14), (1, 4), (5, 6, 3, 2))
    return scaled_op
end

function descend_left(rho::TernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        scaled_rho[-100 -200; -300 -400] :=
        u'[61 62; -400 63] *
        w'[51; 52 -300 61] * w'[21; 62 11 12] *
        rho[42 22; 51 21] *
        w[52 -100 41; 42] * w[31 11 12; 22] *
        u[-200 63; 41 31]
    )
    return scaled_rho
end

function descend_right(rho::TernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        scaled_rho[-100 -200; -300 -400] :=
        u'[62 61; 63 -300] *
        w'[21; 11 12 62] * w'[51; 61 -400 52] *
        rho[22 42; 21 51] *
        w[11 12 41; 22] * w[31 -200 52; 42] *
        u[63 -100; 41 31]
    )
    return scaled_rho
end

function descend_mid(rho::TernaryOperator, layer::TernaryLayer)
    u, w = layer
    # Cost: 6X^6
    @planar(
        scaled_rho[-100 -200; -300 -400] :=
        u'[61 62; -300 -400] *
        w'[21; 11 12 61] * w'[41; 62 31 32] *
        rho[22 42; 21 41] *
        w[11 12 51; 22] * w[52 31 32; 42] *
        u[-100 -200; 51 52]
    )
    return scaled_rho
end

function descend(rho::TernaryOperator, layer::TernaryLayer)
    u, w = layer
    l = descend_left(rho, layer)
    r = descend_right(rho, layer)
    m = descend_mid(rho, layer)
    scaled_rho = (l+r+m)/3
    return scaled_rho
end

# # # Optimization

function environment(op, layer::TernaryLayer, rho; vary_disentanglers = true)
    if vary_disentanglers
        env_u = environment_disentangler(op, layer, rho)
    else
        # The adjoint is just for type stability.
        env_u = zero(layer.disentangler')'
    end
    env_w = environment_isometry(op, layer, rho)
    return TernaryLayer(env_u, env_w)
end

function minimize_expectation_ev(
    layer::TernaryLayer, env::TernaryLayer; vary_disentanglers = true
)
    u = if vary_disentanglers
        projectisometric(env.disentangler; alg = Polar())
    else
        layer.disentangler
    end
    w = projectisometric(env.isometry; alg = Polar())
    return TernaryLayer(u, w)
end

function environment_disentangler(h::TernaryOperator, layer, rho)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        env1[-1 -2; -3 -4] :=
        rho[31 21; 63 22] *
        w[61 51 41; 31] * w[42 11 12; 21] *
        u[52 -2; 41 42] *
        h[62 -1; 51 52] *
        w'[63; 61 62 -3] * w'[22; -4 11 12]
    )

    # Cost: 6X^6
    @planar(
        env2[-1 -2; -3 -4] :=
        rho[41 51; 42 52] *
        w[21 22 61; 41] * w[62 31 32; 51] *
        u[11 12; 61 62] *
        h[-1 -2; 11 12] *
        w'[42; 21 22 -3] * w'[52; -4 31 32]
    )

    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        env3[-1 -2; -3 -4] :=
        rho[21 31; 22 63] *
        w[12 11 42; 21] * w[41 51 61; 31] *
        u[-1 52; 42 41] *
        h[-2 62; 52 51] *
        w'[22; 12 11 -3] * w'[63; -4 62 61]
    )

    env = (env1 + env2 + env3)/3
    return env
end

function environment_isometry(h::TernaryOperator, layer, rho)
    u, w = layer
    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        env1[-1 -2 -3; -4] :=
        rho[81 84; 82 -4] *
        w[62 41 31; 81] * w[83 -2 -3; 84] *
        u[42 52; 31 83] *
        h[61 51; 41 42] *
        u'[63 -1; 51 52] *
        w'[82; 62 61 63]
    )

    # Cost: 6X^6
    @planar(
        env2[-1 -2 -3; -4] :=
        rho[41 62; 42 -4] *
        w[11 12 52; 41] * w[61 -2 -3; 62] *
        u[31 32; 52 61] *
        h[21 22; 31 32] *
        u'[51 -1; 21 22] *
        w'[42; 11 12 51]
    )

    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        env3[-1 -2 -3; -4] :=
        rho[31 33; 32 -4] *
        w[21 11 51; 31] * w[41 61 -3; 33] *
        u[72 62; 51 41] *
        h[71 -2; 62 61] *
        u'[73 -1; 72 71] *
        w'[32; 21 11 73]
    )

    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        env4[-1 -2 -3; -4] :=
        rho[33 31; -4 32] *
        w[-1 61 41; 33] * w[51 11 21; 31] *
        u[62 72; 41 51] *
        h[-2 71; 61 62] *
        u'[-3 73; 71 72] *
        w'[32; 73 11 21]
    )

    # Cost: 6X^6
    @planar(
        env5[-1 -2 -3; -4] :=
        rho[62 41; -4 42] *
        w[-1 -2 61; 62] * w[52 12 11; 41] *
        u[32 31; 61 52] *
        h[22 21; 32 31] *
        u'[-3 51; 22 21] *
        w'[42; 51 12 11]
    )

    # Cost: 2X^8 + 2X^7 + 2X^6
    @planar(
        env6[-1 -2 -3; -4] :=
        rho[84 81; -4 82] *
        w[-1 -2 83; 84] * w[31 41 62; 81] *
        u[52 42; 83 31] *
        h[51 61; 42 41] *
        u'[-3 63; 52 51] *
        w'[82; 63 61 62]
    )

    env = (env1 + env2 + env3 + env4 + env5 + env6)/3
    return env
end


@testset "File Reading" begin
    # Test reading a hex edited Nimrod file. The contents are an 8x8 grid
    # with "5" on the diagonals and "-1" on the off diagonals. Data is
    # stored as 2 Byte integers. The coordinate system is set up with (0,0)
    # in the top left corner and each resolution cell being 10x10m.
    @testset "16-Bit Integer Data" begin
        expectedData = ones(Int16, (8,8)) * -1;
        expectedData[[i for i in 1:9:64]] = 5
        expectedData[1,2] = 5
        testFile = readnimrod("test_hexed_composite_2BI.dat")

        @test size(testFile.data) == (8, 8)
        @test testFile.data == expectedData
        @test testFile.gridtype == NimrodFiles.nationalgrid
        @test isapprox(testFile.xpix_size, 10f0)
        @test isapprox(testFile.ypix_size, 10f0)
        @test testFile.llcorner == (0f0, -80f0)
        @test testFile.llcenter == (5f0, -75f0)
        @test testFile.trcenter == (75f0, -5f0)

        @test testFile.hdr_dump[105] == "mm/h*32"
        @test testFile.hdr_dump[106] == "Plr single site radars"
        @test testFile.hdr_dump[107] == "Rainfall rate Composite"
    end

    # Same thing but for 32-Bit integer data
    @testset "32-Bit Integer data" begin
        expectedData = ones(Int32, (8,8)) * -1;
        expectedData[[i for i in 1:9:64]] = 5
        expectedData[1,2] = 5
        testFile = readnimrod("test_hexed_composite_4BI.dat")

        @test size(testFile.data) == (8, 8)
        @test testFile.data == expectedData
        @test testFile.gridtype == NimrodFiles.nationalgrid
        @test isapprox(testFile.xpix_size, 10f0)
        @test isapprox(testFile.ypix_size, 10f0)
        @test testFile.llcorner == (0f0, -80f0)
        @test testFile.llcenter == (5f0, -75f0)
        @test testFile.trcenter == (75f0, -5f0)

        @test testFile.hdr_dump[105] == "mm/h*32"
        @test testFile.hdr_dump[106] == "Plr single site radars"
        @test testFile.hdr_dump[107] == "Rainfall rate Composite"
    end

    # And now 32-Bit Floating Point Data
    @testset "32-Bit Floating Point data" begin
        expectedData = ones(Float32, (8,8)) * -1;
        expectedData[[i for i in 1:9:64]] = 5
        expectedData[1,2] = 5.0
        testFile = readnimrod("test_hexed_composite_4BF.dat")

        @test size(testFile.data) == (8, 8)
        @test testFile.data == expectedData
        @test testFile.gridtype == NimrodFiles.nationalgrid
        @test isapprox(testFile.xpix_size, 10f0)
        @test isapprox(testFile.ypix_size, 10f0)
        @test testFile.llcorner == (0f0, -80f0)
        @test testFile.llcenter == (5f0, -75f0)
        @test testFile.trcenter == (75f0, -5f0)

        @test testFile.hdr_dump[105] == "mm/h*32"
        @test testFile.hdr_dump[106] == "Plr single site radars"
        @test testFile.hdr_dump[107] == "Rainfall rate Composite"
    end
end


include("textureatlas.jl")

isnewline(x) = x == '\n'

type Text
    data
    text
    atlas
    cursors
    # default values
    font
    scale
    offset
    color
    startposition
    lineheight
end

function Sprite{N, T}(char, position::Point{N, T}, text)
    Sprite(
        char, position, text.scale, text.offset,
        text.color, text.font, text.atlas
    )
end
function Sprite{N, T}(
        char, position::Point{N, T}, scale, offset, color,
        font = defaultfont(),  atlas = get_texture_atlas()
    )
    Sprite{N, T}(
        position,
        glyph_bearing!(atlas, char, font, scale) + offset,
        glyph_scale!(atlas, char, font, scale),
        glyph_uv_width!(atlas, char, font),
        color
    )
end

function nextposition(sprite::Sprite, char, text)
    advance_x, advance_y = glyph_advance!(text.atlas, char, text.font, text.scale)
    position = sprite.position
    if isnewline(char)
        return Point2f0(text.startposition[1], position[2] - advance_y * text.lineheight) #reset to startx
    else
        return position + Point2f0(advance_x, 0)
    end
end

function printat(text::Text, idx::Integer, char::Char)
    position = if checkbounds(Bool, text.data, idx)
        sprite = text.data[idx]
        nextposition(sprite, text.text[idx], text)
    else
        text.startposition
    end
    nextsprite = Sprite(char, position, text)
    idx += 1
    insert!(text.data, idx, nextsprite)
    insert!(text.text, idx, char)
    idx
end

function printat(text::Text, idx::Int, str::String)
    sprite = text.data[idx]
    position = sprite.position
    for char in str
        char == '\r' && continue # stupid windows!
        idx = printat(text, idx, char)
    end
    idx
end

function Base.print(text::Text, char::Union{Char, String})
    map!(text.cursors, text.cursors) do idx
        idx = printat(text, idx, char)
        return idx
    end
    nothing
end
Base.String(text::Text) = join(text.text)