export MOVE_FORWARD, TURN_LEFT, TURN_RIGHT, LRUD, WALL

using Crayons
using Colors

const COLORS = (:red, :green, :blue, :magenta, :yellow, :white)

#####
# Actions
#####

struct MoveForward end
const MOVE_FORWARD = MoveForward()

struct Up end
const UP = Up()
(x::Up)(p::CartesianIndex{2}) = p + CartesianIndex(-1, 0)

struct Down end
const DOWN = Down()
(x::Down)(p::CartesianIndex{2}) = p + CartesianIndex(1, 0)

struct Left end
const LEFT = Left()
(x::Left)(p::CartesianIndex{2}) = p + CartesianIndex(0, -1)

struct Right end
const RIGHT = Right()
(x::Right)(p::CartesianIndex{2}) = p + CartesianIndex(0, 1)

const LRUD = Union{Left, Right, Up, Down}

struct TurnRight end
const TURN_RIGHT = TurnRight()
struct TurnLeft end
const TURN_LEFT = TurnLeft()

(x::TurnRight)(::Left) = UP
(x::TurnRight)(::Up) = RIGHT
(x::TurnRight)(::Right) = DOWN
(x::TurnRight)(::Down) = LEFT
(x::TurnLeft)(::Left) = DOWN
(x::TurnLeft)(::Up) = LEFT
(x::TurnLeft)(::Right) = UP
(x::TurnLeft)(::Down) = RIGHT

#####
# Objects
#####

abstract type AbstractObject end

Base.show(io::IO, x::AbstractObject) = print(io, Crayon(foreground=get_color(x), reset=true), convert(Char, x))

struct Empty <: AbstractObject end
const EMPTY = Empty()
Base.convert(::Type{Char}, ::Empty) = '⋅'
get_color(::Empty) = :white

struct Wall <: AbstractObject end
const WALL = Wall()
Base.convert(::Type{Char}, ::Wall) = '█'
get_color(::Wall) = :white

struct Goal <: AbstractObject end
const GOAL = Goal()
Base.convert(::Type{Char}, ::Goal) = '♥'
get_color(::Goal) = :red

struct Door{C} <: AbstractObject end
Door(c) = Door{c}()
Base.convert(::Type{Char}, ::Door) = '🚪'
get_color(::Door{C}) where C = C

Base.@kwdef mutable struct Agent <: AbstractObject
    color::Symbol=:red
    dir::LRUD
end
function Base.convert(::Type{Char}, a::Agent)
    if        a.dir === UP
        '↑'
    elseif  a.dir === DOWN
        '↓'
    elseif  a.dir === LEFT
        '←'
    elseif a.dir === RIGHT
        '→'
    end
end
get_color(a::Agent) = a.color
get_dir(a::Agent) = a.dir
set_dir!(a::Agent, d) = a.dir = d


### State type

# just an ADM scene
const OCState = Records.Frame{Records.Entity{AutomotiveDrivingModels.VehicleState,AutomotiveDrivingModels.VehicleDef,Int64}}

#### Observaton type

# # for LidarFeatures
# const OCObs = LidarSensor

# for vector features
const OCObs = Array{Float64, 1}
#### Action type

struct OCAction
    acc::Float64
end

function Base.copyto!(a::OCAction, b::OCAction)
    a.acc = b.acc
end

function Base.hash(a::OCAction, h::UInt64 = zero(UInt64))
    return hash(a.acc, h)
end

function Base.:(==)(a::OCAction, b::OCAction)
    return a.acc == b.acc
end

#### POMDP type

@with_kw mutable struct OCPOMDP <: POMDP{OCState, OCAction, OCObs}
    env::CrosswalkEnv = CrosswalkEnv(CrosswalkParams(obstacles_visible=true))
    sensor::AbstractSensor = GaussianSensor()
    models::Dict{Int64, DriverModel} = Dict{Int64, DriverModel}(1=>EgoDriver(OCAction(0.)))
    ego_type::VehicleDef =  VehicleDef()
    ped_type::VehicleDef = VehicleDef(AgentClass.PEDESTRIAN, 1.0, 1.0)
    max_acc::Float64 = 2.0
    ego_start::Float64 = 5.0
    ped_start::Float64 = -env.params.crosswalk_length/4
    ego_goal::Float64 = env.params.roadway_length/2 + 2*env.params.crosswalk_width
    ped_goal::Float64 = env.params.crosswalk_length/4
    ΔT::Float64 = 0.5 # decision frequency
    p_birth::Float64 = 0.3 # probability of appearance
    max_peds::Int64 = 10
    no_ped_prob::Float64 = 0.3 # probability that there is no pedestrian (for training)
    no_ped::Bool = false
    v_noise::Float64 = 1.0
    pos_obs_noise::Float64 = 0.5
    vel_obs_noise::Float64 = 0.5
    collision_cost::Float64 = -1.0
    action_cost::Float64 = 0.0
    goal_reward::Float64 = 1.0
    γ::Float64 = 0.95 # discount factor
end

# function OCPOMDP(; env::CrosswalkEnv = CrosswalkEnv(),
#                    sensor::LidarSensor = LidarSensor(50, max_range=30., angle_spread=float(pi)),
#                    models::Dict{Int64, DriverModel} = Dict{Int64, DriverModel}(1=>EgoDriver(OCAction(0.))),
#                    ego_type::VehicleDef = VehicleDef(),
#                    ped_type::VehicleDef = VehicleDef(AgentClass.PEDESTRIAN, 1.0, 1.0),#,env.params.crosswalk_width),
#                    max_acc::Float64 = 2.0,
#                    ego_start::Float64 = 5.,
#                    ped_start::Float64 = -env.params.crosswalk_length/4,
#                    ego_goal::Float64 = env.params.roadway_length/2 + 2*env.params.crosswalk_width,
#                    ped_goal::Float64 = env.params.crosswalk_length/4,
#                    ΔT::Float64  = 0.5,
#                    p_birth::Float64 = 0.3,
#                    max_peds::Int64=10,
#                    no_ped_prob::Float64 = 0.3,
#                    no_ped::Bool = false,
#                    v_noise::Float64 = 1.,
#                    pos_obs_noise::Float64 = 0.5,
#                    vel_obs_noise::Float64 = 0.5,
#                    collision_cost::Float64 = -1.,
#                    action_cost::Float64 = 0.0,
#                    goal_reward::Float64 = 1.,
#                    γ::Float64  = 0.95)
#     return OCPOMDP(env,
#                    sensor,
#                    models,
#                    ego_type,
#                    ped_type,
#                    max_acc,
#                    ego_start,
#                    ped_start,
#                    ego_goal,
#                    ped_goal,
#                    ΔT,
#                    p_birth,
#                    max_peds,
#                    no_ped_prob,
#                    no_ped,
#                    v_noise,
#                    pos_obs_noise,
#                    vel_obs_noise,
#                    collision_cost,
#                    action_cost,
#                    goal_reward,
#                    γ)
# end

### HELPERS

function POMDPs.discount(pomdp::OCPOMDP)
    return pomdp.γ
end

POMDPs.actions(pomdp::OCPOMDP) = [OCAction(-4.0), OCAction(-2.0), OCAction(0.0), OCAction(2.0)]

function POMDPs.actionindex(pomdp::OCPOMDP, action::OCAction)
    if action.acc == -4.0
        return 1
    elseif action.acc == -2.0
        return 2
    elseif action.acc == 0.
        return 3
    else
        return 4
    end
end


# Note this is asynchronous wrt the host, but you can wait on the stream
function cudasleep(secs; dev::Integer=device(), stream=null_stream)
    device(dev)
    rate = attribute(dev, rt.cudaDevAttrClockRate)
    func = ptxdict[dev].fns["clock_block"]
    twatchdog = 1.95  # watchdog timer kicks in after 2 secs
    while secs > 0
        tics = round(Int64, 1000*rate*min(twatchdog, secs))  # rate is in kHz
        secs -= twatchdog
        CUDAdrv.cudacall(func, 1, 1, (Int64,), tics; stream=convert(CUDAdrv.CuStream, stream))
    end
end


function testniter_3stages(niter, K, maxncuts, cutgen, detectlb)
    if isa(cutgen, StructProg.MultiCutGenerator)
        if K == -1
            if maxncuts == -1
                if detectlb
                    # == 14 with Gurobi
                    # == 15 on Mac OS, == 17 on Linux
                    @test 14 <= niter <= 17
                else
                    # 14 on Gurobi, 15 on CPLEX and 19 otherwise
                    @test 14 <= niter <= 19
                end
            else
                if detectlb
                    # down to 14 with Gurobi
                    # down to 15 and up to 19 on Mac OS, 17 <= ... <= 18 on Linux
                    @test 14 <= niter <= 19
                else
                    # between 16, 17 on Gurobi
                    # between 15, 16 on CPLEX
                    # up to 22 on 32 bits, <= 22 on 64 bits
                    @test 15 <= niter <= 22
                end
            end
        else
            if detectlb
                # down to 11 with CPLEX
                # up to 16 on Windows, up to 15 on Mac OS, <= 14 on Linux
                @test 11 <= niter <= 16
            else
                # 12 on CPLEX
                # 13, 14, 21, 24 on Gurobi
                # up to 20 on Windows, up to 19 on Mac OS, <= 18 on Linux
                @test 12 <= niter <= 24
            end
        end
    else
        if maxncuts == -1
            if K == -1
                # 22 on Gurobi
                # == 19 on Mac OS, == 21 on Linux
                @test 19 <= niter <= 22
            else
                if detectlb
                    # == 14 on Mac OS, == 16 on Linux
                    @test 14 <= niter <= 16
                else
                    # == 14 on Mac OS, == 17 on Linux
                    @test 14 <= niter <= 17
                end
            end
        else
            if K == -1
                if detectlb
                    # 27 on Gurobi
                    # down to 15 Mac OS, >= 19 on Linux
                    @test 15 <= niter <= 27
                else
                    # 27 on Gurobi
                    @test 20 <= niter <= 27
                end
            else
                if detectlb
                    # down to 13 Mac OS, >= 15 on Linux
                    @test 13 <= niter <= 19
                else
                    # down to 14 Mac OS, >= 17 on Linux
                    @test 14 <= niter <= 19
                end
            end
        end
    end
end


using Makie
using MakieLayout
using Makie: px
using KernelDensity
using StatsBase
using GLFW; GLFW.WindowHint(GLFW.FLOATING, 1)
using StatsMakie
using DataFrames
using Colors

struct Infrastructure{T<:NamedTuple}
    content::T
end

function Base.getproperty(infra::Infrastructure, sym::Symbol)
    if sym in fieldnames(Infrastructure)
        getfield(infra, sym)
    else
        infra.content[sym]
    end
end

function create_infrastructure(T::Type, scene::Scene, args...; kwargs...)
    error("create_infrastructure not defined for type $T.")
end

function myplot(T::Type, args...; kwargs...)
    dkwargs = Dict(kwargs)
    scenekw = pop!(dkwargs, :scene, NamedTuple())
    layoutkw = pop!(dkwargs, :layout, NamedTuple())

    scene, layout = layoutscene(; scenekw..., layoutkw...)
    infrastructure, plots = myplot(T, layout[1, 1], args...; dkwargs...)
    (scene = scene, layout = layout, infra = infrastructure, plots = plots)
end

function myplot(T::Type, scene::Scene, args...; kwargs...)
    dkwargs = Dict(kwargs)
    scenekw = pop!(dkwargs, :scene, NamedTuple())
    layoutkw = pop!(dkwargs, :layout, NamedTuple())

    layout = GridLayout(scene; alignmode = Outside(30), scenekw..., layoutkw...)
    infrastructure, plots = myplot(T, layout[1, 1], args...; dkwargs...)
    (layout = layout, infra = infrastructure, plots = plots)
end

function get_layout_top_scene(layout::GridLayout)
    if isnothing(layout.layoutobservables.gridcontent)
        return layout.parentscene
    else
        parentgrid = layout.layoutobservables.gridcontent.parent
        if isnothing(parentgrid)
            return layout.parentscene
        else
            return get_layout_top_scene(parentgrid)
        end
    end
end

abstract type InfrastructureType end
struct SingleAxis <: InfrastructureType end
struct FacetAxes <: InfrastructureType end

infrastructure_type(any::Type) = any # if nothing is set, just use the type directly
infrastructure_type(::Type{Scatter}) = SingleAxis
infrastructure_type(::Type{Lines}) = SingleAxis

function myplot(T::Type, gp::GridPosition, args...; kwargs...)
    scene = get_layout_top_scene(gp.layout)
    if isnothing(scene)
        error("Could not retrieve a parent scene from the layout's tree. Can't plot using a GridLayout without top level parent scene.")
    end

    dkwargs = Dict(kwargs)
    infrakw = pop!(dkwargs, :infra, NamedTuple())

    infrastructure = create_infrastructure(
        infrastructure_type(T),
        scene, args...; infrakw...)

    gp[] = infrastructure.layout
    plots = myplot(T, infrastructure, args...; dkwargs...)
    infrastructure, plots
end

function myplot(T::Type, infra::Infrastructure, args...; kwargs...)
    error("myplot not defined for type $T")
end

function create_infrastructure(::Type{SingleAxis}, scene::Scene, args...; kwargs...)
    dkwargs = Dict(kwargs)
    axiskw = pop!(dkwargs, :axis, NamedTuple())
    layoutkw = pop!(dkwargs, :layout, NamedTuple())

    layout = GridLayout(;layoutkw...)
    axis = layout[1, 1] = Axis(scene; axiskw...)
    Infrastructure((axis = axis, layout = layout))
end

function create_infrastructure(::Union{Type{Heatmap}, Type{Image}}, scene::Scene, args...; kwargs...)
    infra = create_infrastructure(SingleAxis, scene, args...; kwargs...)
    tightlimits!(infra.axis)
    infra
end

function Base.display(nt::NamedTuple{(:scene, :layout, :infra, :plots), Tuple{S,G,I,P}}) where {S<:SceneLike, G, I, P}
    display(nt.scene)
end

function myplot(::Type{Scatter}, infra::Infrastructure, args...; kwargs...)
    scat = scatter!(infra.axis, args...; kwargs...)
end

function myplot(::Type{Lines}, infra::Infrastructure, args...; kwargs...)
    lin = lines!(infra.axis, args...; kwargs...)
end

function myplot(::Type{Heatmap}, infra::Infrastructure, args...; kwargs...)
    hm = heatmap!(infra.axis, args...; kwargs...)
end

function myplot(::Type{Image}, infra::Infrastructure, args...; kwargs...)
    hm = image!(infra.axis, args...; kwargs...)
end

##
scene, layout, _ = myplot(Heatmap, rand(100, 100), infra = (;axis = (;aspect = DataAspect(), xlabel = "helloo"))); scene
myplot(Lines, layout[1, 2], 1:10, sin.(1:10))
scene
##

myplot(Scatter, Group(color = rand(1:3, 100)), rand(100), rand(100), markersize = 20px)

struct MarginDensityScatter end

function create_infrastructure(::Type{MarginDensityScatter}, scene::Scene, args...; kwargs...)

    dkwargs = Dict(kwargs)

    mainax = Axis(scene, xlabel = get(dkwargs, :xlabel, " "), ylabel = get(dkwargs, :ylabel, " "))
    topax = Axis(scene; xlabelvisible = false, xticklabelsvisible = false, xticksvisible = false)
    rightax = Axis(scene; ylabelvisible = false, yticklabelsvisible = false, yticksvisible = false)

    linkxaxes!(mainax, topax)
    linkyaxes!(mainax, rightax)

    tightlimits!(topax, Bottom())
    tightlimits!(rightax, Left())

    marginfraction = get(dkwargs, :marginfraction, 1/4)
    gap = get(dkwargs, :gap, Fixed(20))

    layout = GridLayout(2, 2; rowsizes = [Relative(marginfraction), Auto()],
        colsizes = [Auto(), Relative(marginfraction)], addedcolgaps = gap, addedrowgaps = gap)

    layout[2, 1] = mainax
    layout[1, 1] = topax
    layout[2, 2] = rightax

    Infrastructure((axes = (main = mainax, top = topax, right = rightax), layout = layout))
end

function myplot(::Type{MarginDensityScatter}, infra::Infrastructure, x, y; kwargs...)

    dkwargs = Dict(kwargs)

    xkde = kde(x; npoints = 300)
    ykde = kde(y; npoints = 300)

    color = get(dkwargs, :color, :red)

    scat = scatter!(infra.axes.main, x, y; markersize = get(dkwargs, :markersize, 10px), color = color, strokecolor = get(dkwargs, :strokecolor, :black), strokewidth = get(dkwargs, :strokewidth, 1))
    kde1 = poly!(infra.axes.top, Point2f0.(xkde.x, xkde.density), color = color)
    kde2 = poly!(infra.axes.right, Point2f0.(ykde.density, ykde.x), color = color)

    autolimits!(infra.axes.main)

    (scatter = scat, xkde = kde1, ykde = kde2)
end

##
scene, layout = layoutscene()

for i in 1:3, j in 1:3
    myplot(MarginDensityScatter, layout[i, j], randn(100), randn(100);
        color = rand(RGBAf0),
        infra = (
            marginfraction = 0.2,
            xlabel = "$i & $j",
            gap = Fixed(10)
        ))
end

scene
##

##
scene, layout = layoutscene()
myplot(Scatter, layout[1, 1], rand(100), rand(100), markersize = 20px, color = :green, strokecolor = :black, strokewidth = 1)
myplot(Scatter, layout[2, 1], rand(100), rand(100), markersize = 20px, color = :red, strokecolor = :black, strokewidth = 1)
myplot(Scatter, layout[1:2, 2], rand(100), rand(100), markersize = 20px, color = :blue, strokecolor = :black, strokewidth = 1)
myplot(MarginDensityScatter, layout[1:2, 3:4], randn(100), randn(100))
scene
##

struct FacetPlot end

function create_infrastructure(::Type{FacetPlot}, scene::Scene, plottype, rows, cols, args...; kwargs...)

    dkwargs = Dict(kwargs)

    axiskw = pop!(dkwargs, :axes, NamedTuple())
    layoutkw = pop!(dkwargs, :layout, NamedTuple())

    nrows = length(unique(rows))
    ncols = length(unique(cols))

    layout = GridLayout(; layoutkw...)

    axs = [Axis(scene; axiskw...) for x in CartesianIndices((nrows, ncols))]

    layout[] = axs

    for ax in axs[:, 2:end]
        ax.ylabelvisible = false
        ax.yticklabelsvisible = false
        ax.yticksvisible[] = false
    end

    for ax in axs[1:end-1, :]
        ax.xlabelvisible[] = false
        ax.xticklabelsvisible = false
        ax.xticksvisible[] = false
    end

    linkxaxes!(axs...)
    linkyaxes!(axs...)

    Infrastructure((axes = axs, layout = layout))
end

function myplot(::Type{FacetPlot}, infra::Infrastructure, plottype, rows, cols, x, y; kwargs...)

    dkwargs = Dict(kwargs)

    df = DataFrame(rows = rows, cols = cols, x = x, y = y)

    for (ax, sdf) in zip(infra.axes, groupby(df, [:rows, :cols]))
        myplot(plottype, Infrastructure((axis = ax,)), sdf.x, sdf.y; kwargs...)
    end

    autolimits!(infra.axes[1])
end

##

msize = Node(10px)
scene, layout, infra = myplot(FacetPlot, Scatter,
    rand(1:3, 1000), rand(1:4, 1000), randn(1000), randn(1000), markersize = msize, color = :orange, strokewidth = 1, strokecolor = :black, infra = (; axes = (;backgroundcolor = Gray(0.9))))
myplot(FacetPlot, infra, Lines,
    rand(1:3, 1000), rand(1:4, 1000), randn(1000), randn(1000); color = :red)

for c in 1:infra.layout.ncols
    infra.layout[1, c, Top()] = Box(scene, color = Gray(0.8))
    infra.layout[1, c, Top()] = Label(scene, "Column $c", padding = (0, 0, 10, 10))
end

for r in 1:infra.layout.nrows
    infra.layout[r, end, Right()] = Box(scene, color = Gray(0.8))
    infra.layout[r, end, Right()] = Label(scene, "Row $r", rotation = -pi/2, padding = (10, 10, 0, 0))
end

layout[0, :] = Box(scene, color = Gray(0.7))
layout[1, :] = Label(scene, "Columns", padding = (0, 0, 10, 10))
layout.content[end].content.tellwidth = false

foreach(Axis, layout) do ax; tight_ticklabel_spacing!(ax); end


"""
    setoptions!

This takes `options` arguments for setting DSP options.

# Arguments
- `options`: possible keys are `:param` (path to parameter file),
    `:is_stochastic` (`true` if stochastic model; otherwise, `false`),
    and `:solve_type` (algorithm type; see `Methods`).
"""
function setoptions!(options)
    check_dsp()
    for (optname, optval) in options
        if optname == :param
            readParamFile(dspenv, optval)
        elseif optname == :is_stochastic
            dspenv.is_stochastic = optval
        elseif optname == :solve_type
            if optval in instances(Methods)
                dspenv.solve_type = optval
            else
                @warn("solve_type $optval is not available.")
            end
        else
            @warn("Options $optname is not available.")
        end
    end
end

"""
    load_problem!

Load problem from StructJuMP

# Arguments
- `m`: StructJuMP model
"""
function load_problem!(m::SJ.StructuredModel)

    # set number of blocks (scenarios)
    dspenv.nblocks = SJ.num_scenarios(m)

    if dspenv.is_stochastic
        loadStochasticProblem!(m)
    else
        loadStructuredProblem!(m)
    end

    setBlocks()
end

"""
    loadStochasticProblem!

Load stochastic programming problem from StructJuMP

# Arguments
- `model`: StructJuMP model
"""
function loadStochasticProblem!(model::SJ.StructuredModel)

    nscen = dspenv.nblocks
    ncols1 = length(model.variables)
    nrows1 = length(dspenv.linConstrs[0])
    ncols2 = 0
    nrows2 = 0
    for (id, subm) in SJ.getchildren(model)
        ncols2 = length(subm.variables)
        nrows2 = length(dspenv.linConstrs[id])
        break
    end

    # set scenario indices for each MPI processor
    if dspenv.comm_size > 1
        ncols2 = MPI.Allreduce([ncols2], MPI.MAX, dspenv.comm)[1]
        nrows2 = MPI.Allreduce([nrows2], MPI.MAX, dspenv.comm)[1]
    end

    # set DSPProblem data
    dspenv.numCols[0] = ncols1
    dspenv.numRows[0] = nrows1
    dspenv.colVal[0] = Vector{Float64}(undef, ncols1)
    for (id, blk) in SJ.getchildren(model)
        dspenv.numCols[id] = ncols2
        dspenv.numRows[id] = nrows2
        dspenv.colVal[id] = Vector{Float64}(undef, ncols2)
    end

    setNumberOfScenarios(dspenv, nscen)
    setDimensions(dspenv, ncols1, nrows1, ncols2, nrows2)

    qc_supported = true

    # set problem data
    start, index, value, rlbd, rubd, obj, clbd, cubd, ctype, cname, nqrows, linnzcnt, quadnzcnt, rhs, sense, linstart, linind, linval, quadstart, quadrow, quadcol, quadval = get_model_data(model)
    if nqrows == 0
        loadFirstStage(dspenv, start, index, value, clbd, cubd, ctype, obj, rlbd, rubd)
    else
        if getVersionMajor(dspenv) >= 2
            loadQCQPFirstStage(dspenv, start, index, value, clbd, cubd, ctype, obj, C_NULL, C_NULL, C_NULL, 0, rlbd, rubd, nqrows, linnzcnt, quadnzcnt, rhs, sense, linstart, linind, linval, quadstart, quadrow, quadcol, quadval)
        else
            loadFirstStage(dspenv, start, index, value, clbd, cubd, ctype, obj, rlbd, rubd)
            qc_supported = false
        end
    end

    for (id, blk) in SJ.getchildren(model)
        probability = SJ.getprobability(model)[id]
        start, index, value, rlbd, rubd, obj, clbd, cubd, ctype, cname, nqrows, linnzcnt, quadnzcnt, rhs, sense, linstart, linind, linval, quadstart, quadrow, quadcol, quadval = get_model_data(blk, id)
        if nqrows == 0
            loadSecondStage(dspenv, id-1, probability, start, index, value, clbd, cubd, ctype, obj, rlbd, rubd)
        else
            if getVersionMajor(dspenv) >= 2
                loadQCQPSecondStage(dspenv, id-1, probability, start, index, value, clbd, cubd, ctype, obj, C_NULL, C_NULL, C_NULL, 0, rlbd, rubd, nqrows, linnzcnt, quadnzcnt, rhs, sense, linstart, linind, linval, quadstart, quadrow, quadcol, quadval)
            else
                loadSecondStage(dspenv, id-1, probability, start, index, value, clbd, cubd, ctype, obj, rlbd, rubd)
                qc_supported = false
            end
        end
    end

    if qc_supported == false
        @warn "QCQP is not supported with DSP version $(getVersion(dspenv)). The unsupported objective/constraints are ignored."
    end

    # Set DRO data
    loadDroData(dspenv, dspenv.dro)
end

"""
    loadStructuredProblem!

Load generic block-structured problem from StructJuMP

# Arguments
- `model`: StructJuMP model
"""
function loadStructuredProblem!(model::SJ.StructuredModel)

    ncols1 = length(model.variables)
    nrows1 = length(model.constraints)

    # set DSPProblem data
    dspenv.numCols[0] = ncols1
    dspenv.numRows[0] = nrows1
    dspenv.colVal[0] = Vector{Float64}(undef, ncols1)
    for (id, blk) in SJ.getchildren(model)
        ncols2 = length(blk.variables)
        nrows2 = length(dspenv.linConstrs[id])
        dspenv.numCols[id] = ncols2
        dspenv.numRows[id] = nrows2
        dspenv.colVal[id] = Vector{Float64}(undef, ncols2)
    end

    # load master
    start, index, value, rlbd, rubd, obj1, clbd1, cubd1, ctype1, cname, nqrows, linnzcnt, quadnzcnt, rhs, sense, linstart, linind, linval, quadstart, quadrow, quadcol, quadval = get_model_data(model)
    loadBlockProblem(dspenv, 0, ncols1, nrows1, start[end],
        start, index, value, clbd1, cubd1, ctype1, obj1, rlbd, rubd)

    # Check if the master block has quadratic constraints.
    has_qc = nqrows == 0 ? false : true

    # going over blocks
    for (id, blk) in SJ.getchildren(model)
        ncols2 = length(blk.variables)
        nrows2 = length(dspenv.linConstrs[id])
        start, index, value, rlbd, rubd, obj, clbd, cubd, ctype, cname, nqrows, linnzcnt, quadnzcnt, rhs, sense, linstart, linind, linval, quadstart, quadrow, quadcol, quadval = get_model_data(blk, id)
        loadBlockProblem(dspenv, id, ncols1 + ncols2, nrows2, start[end],
            start, index, value, [clbd1; clbd], [cubd1; cubd], [ctype1; ctype], [obj1; obj], rlbd, rubd)

        # Check if the sub-block has quadratic constraints.
        if !has_qc && nqrows > 0
            has_qc = true
        end
    end

    if has_qc
        @warn "Quadratic constraints are not supported for generic structured programs and thus will be ignored."
    end

    # Finalize loading blocks
    updateBlocks(dspenv)
end

function solve!()
    if dspenv.comm_size == 1
        if dspenv.solve_type == Dual
            if dspenv.is_stochastic
                solveDd(dspenv);
            else
                @warn("Dual decomposition is available for stochastic programming only.")
                return
            end
        elseif dspenv.solve_type == Benders
            if dspenv.is_stochastic
                solveBd(dspenv);
            else
                @warn("Benders decomposition is available for stochastic programming only.")
                return
            end
        elseif dspenv.solve_type == ExtensiveForm
            solveDe(dspenv);
        elseif dspenv.solve_type == DW
            solveDw(dspenv);
        else
            @error("Unexpected error")
            return
        end
    elseif dspenv.comm_size > 1
        if dspenv.solve_type == Dual
            if dspenv.is_stochastic
                solveDdMpi(dspenv);
            else
                @warn("Dual decomposition is available for stochastic programming only.")
                return
            end
        elseif dspenv.solve_type == Benders
            if dspenv.is_stochastic
                solveBdMpi(dspenv);
            else
                @warn("Benders decomposition is available for stochastic programming only.")
                return
            end
        elseif dspenv.solve_type == DW
            solveDwMpi(dspenv);
        elseif dspenv.solve_type == ExtensiveForm
            solveDe(dspenv);
        else
            @error("Unexpected error")
            return
        end
    end

    # solution status
    dspenv.status = getStatus(dspenv)
end

function post_solve!()
    if dspenv.status == 3998
        return
    end

    # get solution time
    dspenv.solve_time = getWallTime(dspenv)

    # primal and dual objective values
    dspenv.primVal = getPrimalBound(dspenv) * dspenv.objective_sense
    dspenv.dualVal = getDualBound(dspenv) * dspenv.objective_sense

    if dspenv.solve_type == Dual
        dspenv.rowVal = getDualSolution(dspenv)
    end

    primVal = dspenv.primVal
    if mysize() > 1
        primVal = MPI.bcast(dspenv.primVal, 0, dspenv.comm)
    end

    if abs(primVal) < 1.0e+20
        primsol = getSolution(dspenv)

        # parse solution to each block
        n_start = 1
        n_end = dspenv.numCols[0]
        dspenv.colVal[0] = primsol[n_start:n_end]
        n_start += dspenv.numCols[0]

        numBlockCols = getNumBlockCols()
        for i in 1:dspenv.nblocks
            n_end += numBlockCols[i]
            dspenv.colVal[i] = primsol[n_start:n_end]
            n_start += numBlockCols[i]
        end
    end
end


using Reactive

function standard_push_test(non_input::Signal)
    m = map(x->2x, non_input)
    pval = number()
    push!(non_input, pval)
    step()

    @fact value(non_input) --> pval
    @fact value(m) --> 2pval
end

facts("Push to non-input nodes") do

    a = Signal(number(); name="a")
    b = map(x -> x*x, a; name="b")

    context("push to map") do
        m = map(x->2x, b)

        push!(b, 10.0)
        step()
        @fact value(b) --> 10.0
        @fact value(m) --> 2value(b)

        push!(a, 2.0)
        step()
        @fact value(b) --> 4.0
        @fact value(m) --> 2value(b)

        push!(b, 3.0)
        step()
        @fact value(b) --> 3.0
        @fact value(m) --> 2value(b)
    end

    context("push to merge") do
        ## Merge
        d = Signal(number(); name="d")
        e = merge(b, d, a; name="e")
        m = map(x->2x, e)
        # precedence to d
        @fact value(e) --> value(d)
        @fact value(m) --> 2value(e)

        standard_push_test(e)
    end

    context("push to foldp") do
        gc()
        x = Signal(number())
        f = foldp(+, 0, x)

        standard_push_test(f)

        # Check for correct continuation
        pval = value(f)
        sval = number()
        push!(x, sval)
        step()

        @fact value(f) --> pval + sval
    end

    context("push to filter") do
        g = Signal(0)
        pred = x -> x % 2 != 0
        h = filter(pred, 1, g)
        standard_push_test(h)
    end

    context("push to sampleon") do
        # sampleon
        g = Signal(0)
        nv = number()
        push!(g, nv)
        step()
        i = Signal(true)
        j = sampleon(i, g)
        standard_push_test(j)
    end

    context("push to droprepeats") do
        # droprepeats
        k = Signal(1)
        l = droprepeats(k)

        standard_push_test(l)
    end

    context("push to filterwhen") do
        # filterwhen
        b = Signal(false)
        n = Signal(1)
        dw = filterwhen(b, 0, n)
        standard_push_test(dw)
    end

    context("push to previous") do
        x = Signal(0)
        y = previous(x)
        standard_push_test(y)
    end

    context("push to delay") do
        x = Signal(0)
        y = delay(x)
        standard_push_test(y)
    end

    context("bind non-input") do
        s = Signal(1; name="sig 1")
        m = map(x->2x, s; name="m")
        s2 = Signal(3; name="sig 2")
        push!(m, 10)
        step()
        @fact value(m) --> 10

        bind!(m, s2) #two-way bind
        @fact value(m) --> 3
        @fact value(s2) --> 3

        push!(m, 6)
        step()
        @fact value(m) --> 6
        @fact value(s2) --> 6

        push!(s2, 10)
        step()
        @fact value(m) --> 10
        @fact value(s2) --> 10
    end

end


# Problem generation #
# ========================== #
"""
    stage_one_model(stochasticprogram::StochasticProgram)

Return a generated copy of the first stage model in `stochasticprogram`.
"""
function stage_one_model(stochasticprogram::StochasticProgram; solver = JuMP.UnsetSolver())
    has_generator(stochasticprogram, :stage_1) || error("First-stage problem not defined in stochastic program. Consider @stage 1.")
    stage_one_model = Model(solver=solver)
    generator(stochasticprogram, :stage_1)(stage_one_model, stage_parameters(stochasticprogram, 1))
    return stage_one_model
end
function _stage_model(generator::Function, stage_params::Any, scenario::AbstractScenario, parent::JuMP.Model)
    stage_model = Model(solver=JuMP.UnsetSolver())
    generator(stage_model, stage_params, scenario, parent)
    return stage_model
end
"""
    stage_two_model(stochasticprogram::StochasticProgram)

Return a generated second stage model corresponding to `scenario`, in `stochasticprogram`.
"""
function stage_model(stochasticprogram::StochasticProgram{N}, stage::Integer, scenario::AbstractScenario; solver = JuMP.UnsetSolver()) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N")
    stage == 1 && return stage_one_model(stochasticprogram; solver = solver)
    stage_key = Symbol(:stage_, stage)
    has_generator(stochasticprogram, stage_key) || error("Stage problem $stage not defined in stochastic program. Consider @stage $stage")
    return _stage_model(generator(stochasticprogram, stage_key),
                        stage_parameters(stochasticprogram, stage),
                        scenario,
                        parentmodel(stochasticprogram, stage))
end
function generate_parent!(scenarioproblems::ScenarioProblems, generator::Function, parent_params::Any)
    generator(parentmodel(scenarioproblems), parent_params)
    return nothing
end
function generate_parent!(scenarioproblems::DScenarioProblems, generator::Function, parent_params::Any)
    generator(parentmodel(scenarioproblems), parent_params)
    @sync begin
        for w in workers()
            @async remotecall_fetch((sp,generator,params)->generate_parent!(fetch(sp),generator,params),
                                    w,
                                    scenarioproblems[w-1],
                                    generator,
                                    parent_params)
        end
    end
    return nothing
end
function generate_parent!(stochasticprogram::StochasticProgram, stage::Integer)
    if stage == 1
        @warn "The first stage has no predecessors."
        return nothing
    end
    parent_key = Symbol(:stage_, stage - 1, :_vars)
    has_generator(stochasticprogram, parent_key) || error("Stage problem $(stage - 1) not defined in stochastic program. Consider @stage $(stage - 1).")
    generate_parent!(scenarioproblems(stochasticprogram, stage), generator(stochasticprogram, parent_key), stage_parameters(stochasticprogram, stage - 1))
    return nothing
end
function generate_stage_one!(stochasticprogram::StochasticProgram)
    haskey(stochasticprogram.problemcache, :stage_1) && return nothing
    has_generator(stochasticprogram, :stage_1) && has_generator(stochasticprogram, :stage_1_vars) || error("First-stage problem not defined in stochastic program. Consider @stage 1.")
    stochasticprogram.problemcache[:stage_1] = JuMP.Model()
    generator(stochasticprogram, :stage_1)(stochasticprogram.problemcache[:stage_1], stage_parameters(stochasticprogram, 1))
    return nothing
end
function generate!(scenarioproblems::ScenarioProblems{S}, stage_params::Any, generator::Function) where S <: AbstractScenario
    for i in nsubproblems(scenarioproblems)+1:nscenarios(scenarioproblems)
        push!(scenarioproblems.problems, _stage_model(generator, stage_params, scenario(scenarioproblems,i), parentmodel(scenarioproblems)))
    end
    return nothing
end
function generate!(scenarioproblems::DScenarioProblems{S}, stage_params::Any, generator::Function) where S <: AbstractScenario
    @sync begin
        for w in workers()
            @async remotecall_fetch((sp,params,generator)->generate!(fetch(sp), params, generator), w, scenarioproblems[w-1], stage_params, generator)
        end
    end
    return nothing
end
function generate_stage!(stochasticprogram::StochasticProgram{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $s not in range 1 to $N")
    if stage == 1
        if haskey(stochasticprogram.problemcache, :stage_1)
            remove_stages!(stochasticprogram, 1)
            invalidate_cache!(stochasticprogram)
        end
        generate_stage_one!(stochasticprogram)
        generate_parent!(stochasticprogram, 2)
    else
        if nsubproblems(stochasticprogram, stage) > 0
            remove_stages!(stochasticprogram, stage)
            invalidate_cache!(stochasticprogram)
        end
        stage_key = Symbol(:stage_, stage)
        has_generator(stochasticprogram, stage_key) || error("Stage problem $stage not defined in stochastic program. Consider @stage $stage.")
        if nscenarios(stochasticprogram, stage) > 0
            p = stage_probability(stochasticprogram, stage)
            abs(p - 1.0) <= 1e-6 || @warn "Scenario probabilities do not add up to one. The probability sum is given by $p"
        end
        stage < N && generate_parent!(stochasticprogram, stage + 1)
        generate!(scenarioproblems(stochasticprogram, stage),
                  stage_parameters(stochasticprogram, stage),
                  generator(stochasticprogram, stage_key))
    end
    return nothing
end
"""
    generate!(stochasticprogram::StochasticProgram)

Generate the `stochasticprogram` using the model definitions from @stage and available data.
"""
function generate!(stochasticprogram::StochasticProgram{N}) where N
    for stage in 1:N
        generate_stage!(stochasticprogram, stage)
    end
    return stochasticprogram
end

function _outcome_model!(outcome_model::JuMP.Model,
                         stage_one_generator::Function,
                         stage_two_generator::Function,
                         stage_one_params::Any,
                         stage_two_params::Any,
                         decision::AbstractVector,
                         scenario::AbstractScenario)
    stage_one_generator(outcome_model, stage_one_params)
    for obj in values(outcome_model.objDict)
        if isa(obj,JuMP.Variable)
            val = decision[obj.col]
            outcome_model.colCat[obj.col] = :Fixed
            outcome_model.colVal[obj.col] = val
            outcome_model.colLower[obj.col] = val
            outcome_model.colUpper[obj.col] = val
        elseif isa(obj,JuMP.JuMPArray{JuMP.Variable})
            for var in obj.innerArray
                val = decision[var.col]
                outcome_model.colCat[var.col] = :Fixed
                outcome_model.colVal[var.col] = val
                outcome_model.colLower[var.col] = val
                outcome_model.colUpper[var.col] = val
            end
        elseif isa(obj,Array{JuMP.Variable})
            for var in obj
                val = decision[var.col]
                outcome_model.colCat[var.col] = :Fixed
                outcome_model.colVal[var.col] = val
                outcome_model.colLower[var.col] = val
                outcome_model.colUpper[var.col] = val
            end
        else
            continue
        end
    end
    stage_two_generator(outcome_model, stage_two_params, scenario, outcome_model)
    return outcome_model
end
"""
    outcome_model(stochasticprogram::TwoStageStochasticProgram,
                  decision::AbstractVector,
                  scenario::AbstractScenario;
                  solver::MathProgBase.AbstractMathProgSolver = JuMP.UnsetSolver())

Return the resulting second stage model if `decision` is the first-stage decision in scenario `i`, in `stochasticprogram`. Optionally, supply a capable `solver` to the outcome model.
"""
function outcome_model(stochasticprogram::StochasticProgram{2}, decision::AbstractVector, scenario::AbstractScenario; solver::MPB.AbstractMathProgSolver = JuMP.UnsetSolver())
    has_generator(stochasticprogram,:stage_1_vars) || error("First-stage not defined in stochastic program. Consider @first_stage or @stage 1.")
    has_generator(stochasticprogram,:stage_2) || error("Second-stage problem not defined in stochastic program. Consider @second_stage.")
    outcome_model = Model(solver = solver)
    _outcome_model!(outcome_model,
                    generator(stochasticprogram,:stage_1_vars),
                    generator(stochasticprogram,:stage_2),
                    stage_parameters(stochasticprogram, 1),
                    stage_parameters(stochasticprogram, 2),
                    decision,
                    scenario)
    return outcome_model
end
"""
    outcome_model(stochasticprogram::StochasticProgram{N},
                  decisions::NTuple{N-1,AbstractVector}
                  scenario_path::NTuple{N-1,AbstractScenario},
                  solver::MathProgBase.AbstractMathProgSolver = JuMP.UnsetSolver())

Return the resulting `N`:th stage model if `decisions` are the decisions taken in the previous stages and `scenario_path` are the realized scenarios up to stage `N` in `stochasticprogram`. Optionally, supply a capable `solver` to the outcome model.
"""
function outcome_model(stochasticprogram::StochasticProgram{N},
                       decisions::NTuple{M,AbstractVector},
                       scenario_path::NTuple{M,AbstractScenario};
                       solver::MPB.AbstractMathProgSolver = JuMP.UnsetSolver()) where {N,M}
    N == M - 1 || error("Inconsistent number of stages $N and number of decisions and scenarios $M")
    # TODO
end
# ========================== #


struct FeatureSubset{T, R} <: FeatureCreator
    fc::T
    subset::R
end

Base.size(fc::FeatureSubset) = Base.size(fc.fc)
Base.length(fc::FeatureSubset) = Base.size(fc.fc)

(fc::FeatureSubset)(s; kwargs...) =
    fc.fc(s[fc.subset]; kwargs...)

(fc::FeatureSubset)(s,a,s_tp1) = fc.fc(s[fc.subset],a,s_tp1[fc.subset])

struct Sparsify{FC}
    fc::FC
end

function (fc::Sparsify)(s; kwargs...)
    ints = fc.fc(s; kwargs...)
    s = spzeros(size(fc))
    if ints isa AbstractArray
        s[ints] .= 1
    else
        s[ints] = 1
    end
    return s
end

function get_active_action_state_vector(state::SparseVector, action, feature_size, num_actions)
    vec_length = feature_size * num_actions
    new_ind = (state.nzind .- 1) * num_actions .+ action
    active_state_action = sparsevec(new_ind, state.nzval, vec_length)
    return active_state_action
end

mutable struct ValueFeatureProjector <: FeatureCreator
    func::Function
    pf_length::Int
end

function project_features(FP::ValueFeatureProjector, state)
    return FP.func(state, action)
end
(FP::ValueFeatureProjector)(state,action) = project_features(FP, state, action)

Base.size(VFP::ValueFeatureProjector) = VFP.pf_length

mutable struct ActionValueFeatureProjector{T} <: FeatureCreator
    func::T
    num_actions::Int
end

function project_features(FP::ActionValueFeatureProjector,state,action,state_tp1)
    state_action = spzeros(size(FP.func) * FP.num_actions)
    state = FP.func(state,action,state_tp1)
    state_length = size(FP.func)
    state_action[(action-1)*state_length + 1: action*state_length] = state
    return state_action
end

(FP::ActionValueFeatureProjector)(state,action,next_state) = project_features(FP, state, action, next_state)

Base.length(AVFP::ActionValueFeatureProjector) = size(AVFP.func) * AVFP.num_actions
Base.size(AVFP::ActionValueFeatureProjector) = size(AVFP.func) * AVFP.num_actions

struct FeatureProjector{FC} <: FeatureCreator
    fc::FC
    next_obs::Bool
end

Base.size(fc::FeatureProjector) = Base.size(fc.fc)
Base.length(fc::FeatureProjector) = Base.size(fc.fc)

(fc::FeatureProjector)(s, a, s_prime; kwargs...) = if fc.next_obs
    fc.fc(s_prime; kwargs...)
else
    fc.fc(s; kwargs...)
end


# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

# This file extends some JuMP functions so that decision variables can be safely
# replaced by (constant) floating point numbers.

import JuMP: value, fix, set_name

function value(x::Float64)
    return x
end

function fix(x::Float64, v::Float64; force)
    return abs(x - v) < 1e-6 || error("Value mismatch: $x != $v")
end

function set_name(x::Float64, n::String)
    # nop
end

function _init(model::JuMP.Model, key::Symbol)::OrderedDict
    if !(key in keys(object_dictionary(model)))
        model[key] = OrderedDict()
    end
    return model[key]
end

function _set_names!(model::JuMP.Model)
    @info "Setting variable and constraint names..."
    time_varnames = @elapsed begin
        _set_names!(object_dictionary(model))
    end
    @info @sprintf("Set names in %.2f seconds", time_varnames)
end

function _set_names!(dict::Dict)
    for name in keys(dict)
        dict[name] isa AbstractDict || continue
        for idx in keys(dict[name])
            if dict[name][idx] isa AffExpr
                continue
            end
            idx_str = join(map(string, idx), ",")
            set_name(dict[name][idx], "$name[$idx_str]")
        end
    end
end


dir_path = @__DIR__
push!(LOAD_PATH, dir_path * "/sim") # simulation
push!(LOAD_PATH, dir_path * "/vis") # visualization
push!(LOAD_PATH, dir_path * "/adv") # adversary policies
push!(LOAD_PATH, dir_path * "/dis") # discretization
push!(LOAD_PATH, dir_path * "/olm")
push!(LOAD_PATH, dir_path * "/pomdp")
push!(LOAD_PATH, dir_path * "/tap")
using sim
using vis
using adv
using dis
using olm
using pomdp
using tap

using Serialization
using Printf
using DelimitedFiles
using Statistics

include("update_info.jl")

function discretize_adv(world::World)
  pts_per_acc = 3
  pts_per_ste = 3
  min_acc = -15
  max_acc = 5
  min_ste = -1e-1
  max_ste = 1e-1
  ctrl_d = Discretization([pts_per_acc, pts_per_ste],
                          [min_acc, min_ste],
                          [max_acc, max_ste])
  return ctrl_d
end

# Scenario: Racer scenario
function main(use_tape::Bool=false)
  # load the graphical context (OpenGL handle, the graphical window, etc.)
  context = vis.setup()
  vis_scaling = 0.01

  # make the road
  path = make_figure8_path()
  road_width = 20.0
  road = Road(path, road_width,
              vis.make_road(context, path.X, path.Y, road_width))
  road.road.T = vis.scale_mat(vis_scaling, vis_scaling)

  # make the world
  world = World(road, [], vis_scaling)

	len_rd = path.S[end]

  # forward search agent ---------------------------------------------------- #
  x01 = [len_rd - 50.0; 25.0; 0.0]
  agent1 = Agent(1, copy(x01), vis.make_car(context))
	agent1.controller! = olm.controller_pofs!
  agent1.custom = [0.0, 0.0]
  b0 = fill(1 / length(pomdp.DRIVERS), length(pomdp.DRIVERS))
  b = b0
	ctrl_d = discretize_adv(world)
  # ------------------------------------------------------------------------- #

  v_track = 25.5
	p_track = 0.0
  x02 = [len_rd / 2 - 50.0; v_track; 0]
  agent2 = Agent(2, copy(x02), vis.make_car(context, [1.0, 0.0, 0.0]))
	agent2.controller! = pomdp.adv_controller!
	agent2.custom = [v_track, p_track, pomdp.NOTHING, ctrl_d, pomdp.WEAK]

  push!(world.agents, agent1)
  push!(world.agents, agent2)

  # make diagnostics render objects
  info1 = vis.InfoBox(context, 0.75, 0.75, vis_scaling)
  info2 = vis.InfoBox(context, -0.75, -0.25, vis_scaling)

  # main loop for rendering and simulation
  window = true
  h = 1e-2
  dt = 1 / 60

  tf = 4.0
  N = round(Int, tf / dt)
  D = 4
  global T = Array{Float64, 2}(undef, (N, D))
  global Cpofs = Array{Float64, 2}(undef, (N, D))
  global Cpomc = Array{Float64, 2}(undef, (N, D))

  do_vis = false
  for j in 1:D
    t = 0.0
    agent1.x = copy(x01)
    agent2.x = copy(x02)
    b = b0
    agent2.custom[5] = pomdp.WEAK
    for i in 1:N
      if do_vis == true
        tape = tap.Tape(world)
        tape.info_objs[1] = info1
        tape.info_objs[2] = info2
        tap.add_frame!(tape)
      end

      t1 = time_ns() / 1e9
      (u, ret) = olm.plan_pomc(agent1.x, b, agent1, agent2, world,
                               pomdp.reward, ctrl_d, j)
      t2 = time_ns() / 1e9
      Cpomc[i, j] = t2 - t1

      t1 = time_ns() / 1e9
      (u, ret) = olm.plan_pofs(agent1.x, b, agent1, agent2, world,
                               pomdp.reward, ctrl_d, j)
      t2 = time_ns() / 1e9
      Cpofs[i, j] = t2 - t1
      agent1.custom = u

      (o, _) = adv.replan_adv(agent2, world)
      agent2.custom[3] = o
      (c, cv) = predict_collision(agent2.x, agent1.x, world)
      b = pomdp.update_belief(b, o, c)

      ## advance one frame in time
      advance!(sim.default_dynamics!, agent1.x, Pair(agent1, world), 0.0, dt, h)
      advance!(agent2.dynamics!, agent2.x, Pair(agent2, world), 0.0, dt, h)

      ## visualize
      if do_vis == true
        diag = diagnostic(agent1, world, 0.0)
        (x, y, sx, sy, dx, u) = diag
        (_, cv) = predict_collision(agent1.x, agent2.x, world)
        tap.add_to_frame!(tape, agent1.x, diag, cv)

        diag = diagnostic(agent2, world, 0.0)
        (x, y, sx, sy, dx, u) = diag
        (_, cv) = predict_collision(agent2.x, agent1.x, world)
        tap.add_to_frame!(tape, agent2.x, diag, cv)
      end

      T[i, j] = t
      t += dt

      if mod(i, div(N, 10)) == 0
        println("$(j) -> $(t)")
      end

      if do_vis == true
        to_vis = tap.visualize_frame(tape, 1)
        window = vis.visualize(context, to_vis)
      end

    end
    println(agent1.x)
    println("Finished depth = $(j)")
  end

  DATApofs = fill(0.0, (N, 2 * D))
  DATApomc = fill(0.0, (N, 2 * D))
  for i in 1:D
    DATApofs[:, 2 * (i - 1) + 1] = T[:, i]
    DATApofs[:, 2 * (i - 1) + 2] = Cpofs[:, i]

    DATApomc[:, 2 * (i - 1) + 1] = T[:, i]
    DATApomc[:, 2 * (i - 1) + 2] = Cpomc[:, i]
  end
  writedlm("../data/pofs_data.txt", DATApofs)
  writedlm("../data/pomc_data.txt", DATApomc)

  AVGCpofs = map(i -> mean(DATApofs[:, 2 * (i - 1) + 2]), 1:D)
  AVGCpomc = map(i -> mean(DATApomc[:, 2 * (i - 1) + 2]), 1:D)

  writedlm("../data/pofs_avgc.txt", AVGCpofs)
  writedlm("../data/pomc_avgc.txt", AVGCpomc)
end


# http://web.mit.edu/puzzle/www/2012/puzzles/watson_2_0/r_e_s_p_e_c_t/
@testset "R.E.S.P.E.C.T" begin
    f = best_feature(String.(shuffle.(collect.(wordlist("""
        ABMNOT
        AENORTY
        BCEKLORSTU
        BFLU
        CDEILNOTU
        CIK
        GIOPS
        ABCEKNO
        ABLMNOR
        ACEILMPR
        ACEILTUV
        HINOPSY
        ACDEIKLNST
        ACDEMY
        ACEMZ
        ACHIMNST
        DEIO
        GLNOTU
        ABCEGILORTY
        ABEINORTX
        ACRTUY
        AGHTUY
        CDEIMNOPU
        ACDILMOPT
        ACHILP
        BCEKUY
        MNOTU
        """)))), 5)
    @test f.description == "has a 1-letter transaddition"
end


mutable struct Atoms
    Natoms::Int64
    Nspecies::Int64
    positions::Array{Float64,2}
    atm2species::Array{Int64,1}
    atsymbs::Array{String,1}
    SpeciesSymbols::Array{String,1}  # unique symbols
end

# Overload println for Atoms
import Base.println
function println( a::Atoms )
    @printf("\n")
    @printf("                                     -----\n")
    @printf("                                     Atoms\n")
    @printf("                                     -----\n")
    @printf("\n")
    @printf("Natoms   = %5d\n", a.Natoms)
    @printf("Nspecies = %5d\n", a.Nspecies)
    @printf("\n")
    for isp = 1:a.Nspecies
        @printf("Species %3d: %s\n", isp, a.SpeciesSymbols[isp])
    end
    @printf("\n")
    @printf("Coordinates in bohr:\n")
    @printf("\n")
    for ia = 1:a.Natoms
        @printf("%5s %18.10f %18.10f %18.10f\n", a.atsymbs[ia],
                a.positions[1,ia], a.positions[2,ia], a.positions[3,ia])
    end
end

# dummy atoms, contains only one atom
function Atoms()
    Natoms = 1
    Nspecies = 1
    positions = zeros(3,Natoms)
    atm2species = [1]
    atsymbs = ["X"]
    SpeciesSymbols = ["X"]  # unique symbols
    return Atoms( Natoms, Nspecies, positions, atm2species, atsymbs, SpeciesSymbols )
end

"""
Initialize from xyz file
"""
function init_atoms_xyz(filexyz; in_bohr=false, verbose=false)
    f = open(filexyz, "r")
    l = readline(f)
    Natoms = parse(Int64, l)
    positions = zeros(3,Natoms)
    atsymbs = Array{String}(undef,Natoms)
    l = readline(f)
    for ia = 1:Natoms
        ll = split( readline(f) )
        atsymbs[ia] = ll[1]
        positions[1,ia] = parse( Float64, ll[2] )
        positions[2,ia] = parse( Float64, ll[3] )
        positions[3,ia] = parse( Float64, ll[4] )
    end
    close(f)

    # convert from angstrom to bohr
    if !in_bohr
        if verbose
            println("Coordinate in xyz file is assumed to be given in angstrom")
            println("It will be converted to bohr")
        end
        positions[:,:] = positions[:,:] * ANG2BOHR
    else
        if verbose
            println("Coordinate in xyz file is assumed to be given in bohr")
        end
    end

    # Determine number of species
    Nspecies = 0
    for ia = 1:Natoms
        k2 = 0
        for k1 = 1:ia-1
            if atsymbs[k1] == atsymbs[ia]
                k2 = 1
            end
        end
        # find different
        if k2 == 0
            Nspecies = Nspecies + 1
        end
    end

    SpeciesSymbols = Array{String}(undef,Nspecies)

    idx1 = 0
    for ia = 1:Natoms
        k2 = 0
        for k1 = 1:ia-1
            if atsymbs[k1] == atsymbs[ia]
                k2 = 1
            end
        end
        # Found different species
        if k2==0
            idx1 = idx1 + 1
            SpeciesSymbols[idx1] = atsymbs[ia]
        end
    end

    # Mapping of atoms to species index
    atm2species = Array{Int64}(undef,Natoms)
    for ia = 1:Natoms
        for isp = 1:Nspecies
            if atsymbs[ia] == SpeciesSymbols[isp]
                atm2species[ia] = isp
            end
        end
    end

    return Atoms(Natoms, Nspecies, positions, atm2species, atsymbs, SpeciesSymbols )

    #ALLOCATE( AtomicMasses(Nspecies) )
    #AtomicMasses(:) = 0.d0 ! FIXME: Use internal database to set this

end

const ZATOMS = Dict(
    "H"  => 1,
    "He" => 2,
    "Li" => 3,
    "Be" => 4,
    "B"  => 5,
    "C"  => 6,
    "N"  => 7,
    "O"  => 8,
    "F"  => 9,
    "Ne" => 10,
    "Na" => 11,
    "Mg" => 12,
    "Al" => 13,
    "Si" => 14,
    "P"  => 15,
    "S"  => 16,
    "Cl" => 17,
    "Ar" => 18,
    "K"  => 19,
    "Ca" => 20,
    "Sc" => 21,
    "Ti" => 22,
    "V"  => 23,
    "Cr" => 24,
    "Mn" => 25,
    "Fe" => 26,
    "Co" => 27,
    "Ni" => 28,
    "Cu" => 29,
    "Zn" => 30,
    "Ga" => 31,
    "Ge" => 32,
    "As" => 33,
    "Se" => 34,
    "Br" => 34,
    "Kr" => 36,
    "Rb" => 37,
    "Sr" => 38,
    "Y"  => 39,
    "Zr" => 40,
    "Nb" => 41,
    "Mo" => 42,
    "Tc" => 43,
    "Ru" => 44,
    "Rh" => 45,
    "Pd" => 46,
    "Ag" => 47,
    "Cd" => 48,
    "In" => 49,
    "Sn" => 50,
    "Sb" => 51,
    "Te" => 52,
    "I"  => 53,
    "Xe" => 54, )

function get_Zatoms( atoms::Atoms )
    Nspecies = atoms.Nspecies
    SpeciesSymbols = atoms.SpeciesSymbols
    Zatoms = zeros(Float64,Nspecies)
    for isp = 1:Nspecies
        Zatoms[isp] = ZATOMS[SpeciesSymbols[isp]]
    end
    return Zatoms
end

function get_Nelectrons( atoms:: Atoms )
    Zatoms = get_Zatoms(atoms)
    Natoms = atoms.Natoms
    atm2species = atoms.atm2species
    Nelectrons = 0.0
    for ia = 1:Natoms
        isp = atm2species[ia]
        Nelectrons = Nelectrons + Zatoms[isp]
    end
    return Nelectrons
end


function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    ::DeviceModel{T, FixedOutput},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    nodal_expression!(optimization_container, devices, S)

    return
end

"""
This function creates the model for a full thermal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, D},
    ::Type{S},
) where {
    T <: PSY.ThermalGen,
    D <: AbstractStandardUnitCommitment,
    S <: PM.AbstractPowerModel,
}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(optimization_container, ActivePowerVariable, devices, D())
    add_variables!(optimization_container, ReactivePowerVariable, devices, D())
    add_variables!(optimization_container, OnVariable, devices, D())
    add_variables!(optimization_container, StartVariable, devices, D())
    add_variables!(optimization_container, StopVariable, devices, D())

    # Initial Conditions
    initial_conditions!(optimization_container, devices, D())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    commitment_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    time_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

"""
This function creates the model for a full thermal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, D},
    ::Type{S},
) where {
    T <: PSY.ThermalGen,
    D <: AbstractStandardUnitCommitment,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(optimization_container, ActivePowerVariable, devices, D())
    add_variables!(optimization_container, OnVariable, devices, D())
    add_variables!(optimization_container, StartVariable, devices, D())
    add_variables!(optimization_container, StopVariable, devices, D())

    # Initial Conditions
    initial_conditions!(optimization_container, devices, D())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    commitment_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    time_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

"""
This function creates the model for a full thermal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, ThermalBasicUnitCommitment},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractPowerModel}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        ReactivePowerVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        OnVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StartVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StopVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalBasicUnitCommitment())

    # Constraints
    # TODO: active_power_constraints
    # TODO: refactor constraints such that ALL variables for all devices are added first, and then the constraint creation is trigged
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    commitment_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

"""
This function creates the model for a full thermal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, ThermalBasicUnitCommitment},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        OnVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StartVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StopVariable,
        devices,
        ThermalBasicUnitCommitment(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalBasicUnitCommitment())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    commitment_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

"""
This function creates the model for a full thermal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, ThermalRampLimited},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractPowerModel}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalRampLimited(),
    )
    add_variables!(
        optimization_container,
        ReactivePowerVariable,
        devices,
        ThermalRampLimited(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalRampLimited())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

"""
This function creates the model for a full thermal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, ThermalRampLimited},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalRampLimited(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalRampLimited())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, D},
    ::Type{S},
) where {
    T <: PSY.ThermalGen,
    D <: AbstractThermalDispatchFormulation,
    S <: PM.AbstractPowerModel,
}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(optimization_container, ActivePowerVariable, devices, D())
    add_variables!(optimization_container, ReactivePowerVariable, devices, D())

    # Initial Conditions

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, D},
    ::Type{S},
) where {
    T <: PSY.ThermalGen,
    D <: AbstractThermalDispatchFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(optimization_container, ActivePowerVariable, devices, D())

    # Initial Conditions

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))

    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{PSY.ThermalMultiStart, ThermalMultiStartUnitCommitment},
    ::Type{S};
    kwargs...,
) where {S <: PM.AbstractPowerModel}
    devices = PSY.get_components(PSY.ThermalMultiStart, sys)

    if !validate_available_devices(PSY.ThermalMultiStart, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        ReactivePowerVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        OnVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StopVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StartVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        ColdStartVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        WarmStartVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        HotStartVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalMultiStartUnitCommitment())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    commitment_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    time_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    startup_time_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    startup_type_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    startup_initial_condition_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    must_run_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    initial_range_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))
    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{PSY.ThermalMultiStart, ThermalMultiStartUnitCommitment},
    ::Type{S},
) where {S <: PM.AbstractActivePowerModel}
    devices = PSY.get_components(PSY.ThermalMultiStart, sys)

    if !validate_available_devices(PSY.ThermalMultiStart, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        OnVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StopVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StartVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        ColdStartVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        WarmStartVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        HotStartVariable,
        devices,
        ThermalMultiStartUnitCommitment(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalMultiStartUnitCommitment())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    commitment_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    time_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    startup_time_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    startup_type_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    startup_initial_condition_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    must_run_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    initial_range_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    feedforward!(optimization_container, devices, model, get_feedforward(model))
    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, ThermalCompactUnitCommitment},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractPowerModel}
    devices = PSY.get_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        ReactivePowerVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        OnVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StartVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StopVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalCompactUnitCommitment())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    commitment_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    time_constraints!(optimization_container, devices, model, S, get_feedforward(model))

    feedforward!(optimization_container, devices, model, get_feedforward(model))
    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, ThermalCompactUnitCommitment},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractActivePowerModel}
    devices = PSY.get_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        OnVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StartVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )
    add_variables!(
        optimization_container,
        StopVariable,
        devices,
        ThermalCompactUnitCommitment(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalCompactUnitCommitment())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    commitment_constraints!(
        optimization_container,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    time_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    feedforward!(optimization_container, devices, model, get_feedforward(model))
    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, ThermalCompactDispatch},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractPowerModel}
    devices = PSY.get_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalCompactDispatch(),
    )
    add_variables!(
        optimization_container,
        ReactivePowerVariable,
        devices,
        ThermalCompactDispatch(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalCompactDispatch())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    feedforward!(optimization_container, devices, model, get_feedforward(model))
    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end

function construct_device!(
    optimization_container::OptimizationContainer,
    sys::PSY.System,
    model::DeviceModel{T, ThermalCompactDispatch},
    ::Type{S},
) where {T <: PSY.ThermalGen, S <: PM.AbstractActivePowerModel}
    devices = PSY.get_components(T, sys)

    if !validate_available_devices(T, devices)
        return
    end

    # Variables
    add_variables!(
        optimization_container,
        ActivePowerVariable,
        devices,
        ThermalCompactDispatch(),
    )

    # Initial Conditions
    initial_conditions!(optimization_container, devices, ThermalCompactDispatch())

    # Constraints
    add_constraints!(
        optimization_container,
        RangeConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
        get_feedforward(model),
    )
    ramp_constraints!(optimization_container, devices, model, S, get_feedforward(model))
    feedforward!(optimization_container, devices, model, get_feedforward(model))
    # Cost Function
    cost_function!(optimization_container, devices, model, S, get_feedforward(model))

    return
end


include("./src/ShellModel.jl")
using .ShellModel

##########################
### solve EC for approx. eigenpairs
##########################
"""
calc_moment=true   # calc. mu&Q-moment
write_appwav0false # to write .appwav you need to make sample w.f.
"""
function solve(write_appwav=false)
    sntpath = "snts/random_input/validsnts_sigma1/sdshl/"
    Hs = [  sntpath*"tmp_$i"*".snt" for i = 0:99 ]
    #tpath = "/Users/Sota/Desktop/develop_ShellModel.jl/"
    tpath = "/Users/sotauu/Desktop/develop_ShellModel.jl/"
    elogf = "runlog_sigma1_valid.dat"

    num_ev_target = 1
    verbose = false # to show all Energy(EC)
    tJNs = [ [tJ,num_ev_target] for tJ in 0:2:8]
    for target_nuc in ["Mg24","Al26","Si28"]
        solveEC(Hs,target_nuc,tJNs;wpath=tpath,verbose=verbose,exact_logf=elogf)
    end

    tJNs = [ [tJ,num_ev_target] for tJ in 1:2:9]
    for target_nuc in ["Mg25"]
        solveEC(Hs,target_nuc,tJNs;wpath=tpath,verbose=verbose,exact_logf=elogf)
    end
end

solve()


"""
at this point we already have res list filled with all fp and fn's in some cases it is enough in others we want futher actions to be done
  - for example in case of lack of isometry in voxels we can now apply some corrections for it
"""
module ResultProcess

"""
it will scan through result list - ignorea all that are zeroces in res indicies, add only the iteration numbers and in the end we will just dividy it by (fp+fn)
entriesPerBlock - amount of entries per block of threads
totalLength - total length of result list
iterLoopResList - how many times given block need to iterate to aalyze all entries given to it
globalSum - global variable accessed atomically holding sum of iter numbs of results
"""
function getAverage(resList,resListIndicies,entriesPerBlock,totalLength,iterLoopResList ,globalSum )
 locSum = UInt32(0)
 shmemSum = @cuStaticSharedMem(UInt32, (33,1))
 offsetIter= UInt8(1)

  @iterateLinearlyForResList(iterLoopResList,entriesPerBlock,totalLength,begin
      if(resListIndicies[i]>0)# we do not want 0 entries or duplicated ones
        locSum+=resList[i,6]#adding iteration number
      end
  end )

  #time for reduction
  @redWitAct(offsetIter,shmemSum, locSum,+   )
  sync_threads()
  @addAtomic(shmemSum, globalSum)
end #getAverage

"""
iterate linearly result list
entriesPerBlock - amount of entries per block of threads
totalLength - total length of result list
"""
macro iterateLinearlyForResList(iterLoop,entriesPerBlock,totalLength, ex)
  return  esc(quote
  i = UInt32(0)
  @unroll for j in 0:($iterLoop)
    offset = ((blockIdxX()-1) *$entriesPerBlock)
    i= threadIdxX()+(threadIdxY()-1)*blockDimX()+ j* blockDimX()*blockDimY()
    if((i+offset)<=$totalLength && i<$entriesPerBlock)
      i+=offset
      $ex
    end
  end
   end)
  end

"""
We will iterate over result list take all entries that are non zeo in res  indicies list now in a while loop we will lok for the source voxel that was responsible for the  result
so we can start fom naive approach where each lane will be responsible for finding source voxel that is the closest voxel from other mask
In order to achieve it we will go in the direction indicated by the dir value in the distance indicated by the dir in the result matrix  and the number of voxels indicated by the iteration number
It is probable that we will not find there he source voxel so we will need to iteratively increase number of voxels that go in orthogonal direction so for example
  when direction is down and iteration number 10 we will ook 10 voxels blow in straight line at first
  then if wew ould not find we will look into voxel 9 down and in 1 left,right,anterior,posterior - so over orthogonal axes
  then we will  try going 8 pixels down and 1 pixel in 4 directions and futher pixel also in 4 directions  so 16 permutations
  basically we will want to scan the area of the sphere with radius of iteration number and getting circle (or sth more like a square/rhombus as we analze 1 norm ...) crossections of it  starting from the point in the end of the axis in given
  direction and progress up circle crossesction by circle crossesction
so the amount of voxels that we have to use at each iteration will increase in each step hence w have no free voxels to move at first step - then we have 4 possible position then I suppose 16...
so we can do it like that as we are fro example on step 3 and we are circling around with sth like circle around with radius 3 we can first analyze one axis for example left right so we go
3 voxls laeft and right we have a part of line now we get shorter line (by 1 on each end) anteriorly and posteriorly to the first axis then next row even shorter anteriorly and posteriorly ...
until we will reduce row length to 1 - this should lead to check the voxels at the end of anterior posterior ends of anterior posterior axis

We can xperiment with getting corrections one by single line at the begining with some maximum number of ccles allowed - so if single lane wuld not be able to find result we would cease looking for it
and then scan one more time whole warps  for scanning those points that were not yet resolved

additionally if we want to create intresting visualization we can add information about corrected distance to all points between source and target voxel..

  entriesPerBlock - amount of entries per block of threads
  totalLength - total length of result list
  iterLoopResList - how many times given block need to iterate to aalyze all entries given to it
  globalSum - global variable accessed atomically holding sum of iter numbs of results - so we will accumulate corrected values - to later return average
  maxSingleThrIterNumb - maximum iteration number that should be analyzed using just a single lane; all points bigger than that will be analyzed used whole warps or maybe blocks
  """
function applyCorrection(resList,resListIndicies,entriesPerBlock,totalLength,iterLoopResList,globalSum,maxSingleThrIterNumb, referenceArrs )
  locSum = UInt32(0)
  shmemSum = @cuStaticSharedMem(UInt32, (33,1))
  offsetIter= Int16(1)
  locIterNumb=Int16(0)
  #indicates how coordinates should change relatively to the covered point from result list
  xChange= Int16(0)
  yChange= Int16(0)
  zChange= Int16(0)

@iterateLinearlyForResList(iterLoopResList,entriesPerBlock,totalLength,begin
  if(resListIndicies[i]>0)# we do not want to analyze 0 entries or duplicated ones

  end
end )

end

### looking for the true entry in given array  we iteratively scan area  where the original point from which dilatation started and led to this result that we are currently analyzing

"""
using single thread we iteratively scan area  where the original point from which dilatation started and led to this result that we are currently analyzing
referenceArrs= (reducedGoldB,reducedSegmB)
                top 6
                bottom 5
                left 2
                right 1
                anterior 3
                posterior 4
"""
macro singleThreadScan(resList, i, referenceArrs )

  return  esc(quote

  offsetIter=$resList[$i,6] #storing iteration number
  locIterNumb=0 #will get incemented every time we move centrally in main axis
  while offsetIter>0
    #first we need to  modify x,y or z depending on the direction given in dir variable of result list
    dir = $resList[$i,5]
    xChange-= (dir==2)*offsetIter
    xChange+= (dir==1)*offsetIter

    yChange-= (dir==4)*offsetIter
    yChange+= (dir==3)*offsetIter

    zChange-= (dir==6)*offsetIter
    zChange+= (dir==5)*offsetIter

    #now we have established position over main axis - what is left is to scan by modifying orthogonal axes

    @unroll for orthoAxAmove in -locIterNumb:locIterNumb#so we need to iterate over one of the
      @unroll for orthoAxBmove in [-locIterNumb+ CUDA.abs(orthoAxAmove),locIterNumb-CUDA.abs(orthoAxAmove)]
        #isXMain,isYMain - true if x or in second case y is main axis - we define main axis on the basis of dir variable
        isXMain = (dir==2 || dir==1)
        isYMain = (dir==3 || dir==4)

        #(dir!=2 && dir!=1) so x axis is not main axis
        xChange+= (!isXMain)*orthoAxAmove
        #(dir==2 || dir==1)*orthoAxBmove - so if x is main axis - it means that y is not  and it will be axis A
        yChange+=(!isXMain && !isYMain)*orthoAxBmove+ (isXMain)*orthoAxAmove
        #zChange can never affect axis A  will affect axis B if it is not main axis - so when either x or y is main axis
        zChange+=(!isXMain && !isYMain)*orthoAxBmove

        #here if we will find that there is true in this position we will assume  that those coordinates are coordinates of
        if(referenceArrs[2-$resList[$i,4]][ $resList[$i,1]+xChange,$resList[$i,2]+yChange,$resList[$i,3]+zChange])
            #now on the basis of new coordinates we need to calculate correction for non isometric voxels
            corrected =getCorrectedDistance(yDimSize,zDimSize, $resList[$i,1]+xChange,$resList[$i,2]+yChange,$resList[$i,3]+zChange)
            locSum+=corrected
            $resList[$i,6]=corrected
        end

        #resetting values
    xChange= Int16(0)
    yChange= Int16(0)
    zChange= Int16(0)
      end#
    end#orthoAxAmove
    offsetIter-=1
    locIterNumb+=1

  end#while
end)#quote
end #singleThreadScan

"""
using single warp we iteratively scan area  where the original point from which dilatation started and led to this result that we are currently analyzing
referenceArrs= (reducedGoldB,reducedSegmB)
                top 6
                bottom 5
                left 2
                right 1
                anterior 3
                posterior 4
"""
macro singleWarpScan(resList, i, referenceArrs , is)

  return  esc(quote
  #this will enable simple way of parallelization as steps will be just executing only those that modulus is compatible with thread id ...

  operationIndex = UInt16(0)

  offsetIter=$resList[$i,6] #storing iteration number
  locIterNumb=0 #will get incemented every time we move centrally in main axis
  while offsetIter>0
    #first we need to  modify x,y or z depending on the direction given in dir variable of result list
    dir = $resList[$i,5]
    xChange-= (dir==2)*offsetIter
    xChange+= (dir==1)*offsetIter

    yChange-= (dir==4)*offsetIter
    yChange+= (dir==3)*offsetIter

    zChange-= (dir==6)*offsetIter
    zChange+= (dir==5)*offsetIter

    #now we have established position over main axis - what is left is to scan by modifying orthogonal axes

    @unroll for orthoAxAmove in -locIterNumb:locIterNumb#so we need to iterate over one of the
      @unroll for orthoAxBmove in [-locIterNumb+ CUDA.abs(orthoAxAmove),locIterNumb-CUDA.abs(orthoAxAmove)]
        #isXMain,isYMain - true if x or in second case y is main axis - we define main axis on the basis of dir variable
        isXMain = (dir==2 || dir==1)
        isYMain = (dir==3 || dir==4)

        #(dir!=2 && dir!=1) so x axis is not main axis
        xChange+= (!isXMain)*orthoAxAmove
        #(dir==2 || dir==1)*orthoAxBmove - so if x is main axis - it means that y is not  and it will be axis A
        yChange+=(!isXMain && !isYMain)*orthoAxBmove+ (isXMain)*orthoAxAmove
        #zChange can never affect axis A  will affect axis B if it is not main axis - so when either x or y is main axis
        zChange+=(!isXMain && !isYMain)*orthoAxBmove

        #here if we will find that there is true in this position we will assume  that those coordinates are coordinates of
        if(referenceArrs[2-$resList[$i,4]][ $resList[$i,1]+xChange,$resList[$i,2]+yChange,$resList[$i,3]+zChange])
            #now on the basis of new coordinates we need to calculate correction for non isometric voxels
            corrected =getCorrectedDistance(yDimSize,zDimSize, $resList[$i,1]+xChange,$resList[$i,2]+yChange,$resList[$i,3]+zChange)
            locSum+=corrected
            $resList[$i,6]=corrected
        end

        #resetting values
    xChange= Int16(0)
    yChange= Int16(0)
    zChange= Int16(0)
      end#
    end#orthoAxAmove
    offsetIter-=1
    locIterNumb+=1

  end#while
end)#quote
end #singleThreadScan

"""
function invoked in case image has non isometric voxels - this ussually is related to the
  fact that voxel is smaller in z dimension than in other dimensions
  IMPORTANT !!! the correction is not perfect as in dilatation step the lack of isometry is not taken into account source voxel may be diffrent than one we would find below, yet in most cases approximation should be good

  x dimension is generally set as a reference dimension and set to be 1 ; y dimension ussually is the same but in theory do not have to be
  z dimension will be given to this function as units where 1 unit is x dimension of voxel in mm
    yDimSize,zDimSize - given in scale where 1 is x dimension of the voxel and presumed to be uniform in whole image
    xDisp,yDisp,zDisp - displacement in x,y,z axis from the covered voxel to source voxel
"""
function getCorrectedDistance(yDimSize,zDimSize, xDisp,yDisp,zDisp)
  return xDisp+ yDimSize*yDisp+ zDimSize*zDisp
end

end#ResultProcess


include(joinpath(@__DIR__, "..", "src/dynamics", "flamingo", "visuals.jl"))
T = Float64
vis = Visualizer()
render(vis)
open(vis)

# get model
model_sim = deepcopy(flamingo_piecewise)

plot(x, model_sim.env.surf.(x))
dir = joinpath(module_dir(), "src/dynamics/flamingo")

path_base = joinpath(dir, "dynamics/base.jld2")
path_dyn = joinpath(dir, "dynamics/dynamics.jld2")
path_res = joinpath(dir, "piecewise/residual.jld2")
path_jac = joinpath(dir, "piecewise/sparse_jacobians.jld2")
path_linearized = joinpath(dir, "piecewise/linearized.jld2")

instantiate_base!(model_sim, path_base)

expr_dyn = generate_dynamics_expressions(model_sim, derivs = true)
save_expressions(expr_dyn, path_dyn, overwrite=true)
instantiate_dynamics!(model_sim, path_dyn, derivs = true)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(model_sim, jacobians = :approx)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(model_sim, path_res, jacobians = :approx)
model_sim.spa.rz_sp = copy(rz_sp)
model_sim.spa.rθ_sp = copy(rθ_sp)

model = get_model("flamingo", surf="flat")
nq = model.dim.q
nu = model.dim.u
nc = model.dim.c
nb = model.dim.b
nd = nq + nc + nb
nr = nq + nu + nc + nb + nd
nz = num_var(model)
nθ = num_data(model)

# get trajectory
# ref_traj = get_trajectory("flamingo", "gait1", load_type=:split_traj_alt, model=model)
# ref_traj = get_trajectory("flamingo", "gait_forward_36_3", load_type=:split_traj_alt, model=model)
ref_traj = get_trajectory("flamingo", "gait_forward_36_4", load_type=:split_traj_alt, model=model)

H = ref_traj.H
h = ref_traj.h
N_sample = 5
H_mpc = 15
h_sim = h / N_sample
H_sim = 2500#15000

# barrier parameter
κ_mpc = 1.0e-4

obj = TrackingVelocityObjective(H_mpc, model.dim,
    v = [Diagonal(1e-3 * [1e0,1,1e4,1,1,1,1,1e4,1e4]) for t = 1:H_mpc],
    q = [Diagonal(1e-1 * [3e2, 1e-6, 3e2, 1, 1, 1, 1, 0.1, 0.1]) for t = 1:H_mpc],
    u = [Diagonal(3e-1 * [0.1; 0.1; 0.3; 0.3; ones(nu-6); 2; 2]) for t = 1:H_mpc],
    γ = [Diagonal(1.0e-100 * ones(model.dim.c)) for t = 1:H_mpc],
    b = [Diagonal(1.0e-100 * ones(model.dim.b)) for t = 1:H_mpc])

p = linearized_mpc_policy(ref_traj, model, obj,
    H_mpc = H_mpc,
    N_sample = N_sample,
    κ_mpc = κ_mpc,
    n_opts = NewtonOptions(
        r_tol = 3e-4,
        max_iter = 5),
    mpc_opts = LinearizedMPCOptions(
        # live_plotting=true,
        altitude_update = true,
        altitude_impact_threshold = 0.02,
        altitude_verbose = true,
        )
    )

q1_ref = copy(ref_traj.q[2])
q0_ref = copy(ref_traj.q[1])
q1_sim = SVector{model.dim.q}(q1_ref)
q0_sim = SVector{model.dim.q}(copy(q1_sim - (q1_ref - q0_ref) / N_sample))
@assert norm((q1_sim - q0_sim) / h_sim - (q1_ref - q0_ref) / h) < 1.0e-8

sim = simulator(model_sim, q0_sim, q1_sim, h_sim, H_sim,
    p = p,
    ip_opts = InteriorPointOptions(
        r_tol = 1.0e-8,
        κ_init = 1.0e-8,
        κ_tol = 2.0e-8),
    sim_opts = SimulatorOptions(warmstart = true)
    )

@time status = simulate!(sim)

# save trajectory
@save joinpath(module_dir(), "src/dynamics/flamingo/simulations/piecewise.jld2") sim
@load joinpath(module_dir(), "src/dynamics/flamingo/simulations/piecewise.jld2") sim

l = 9
lu = 1
plt = plot(layout=(3,1), legend=false)
plot!(plt[1,1], hcat(Vector.(vcat([fill(ref_traj.q[i], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[1,1], hcat(Vector.([q[l:l] for q in sim.traj.q])...)', color=:blue, linewidth=1.0)
plot!(plt[2,1], hcat(Vector.(vcat([fill(ref_traj.u[i][lu:lu], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[3,1], hcat(Vector.(vcat([fill(ref_traj.γ[i][1:nc], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[2,1], hcat(Vector.([u[lu:lu] for u in sim.traj.u]*N_sample)...)', color=:blue, linewidth=1.0)
# plot!(plt[3,1], hcat(Vector.([γ[1:nc] for γ in sim.traj.γ]*N_sample)...)', color=:blue, linewidth=1.0)
# plot!(plt[3,1], hcat(Vector.([b[1:nb] for b in sim.traj.b]*N_sample)...)', color=:red, linewidth=1.0)

plot_lines!(vis, model, sim.traj.q[1:N_sample:end], offset=-0.01)
plot_surface!(vis, model_sim.env, xlims=[-1, 9])
anim = visualize_meshrobot!(vis, model_sim, sim.traj, sample=10)
anim = visualize_force!(vis, model_sim, sim.traj, anim=anim, h=h_sim, sample=10)
open(vis)

filename = "flamingo_40_steps_sine"
MeshCat.convert_frames_to_video(
    "/home/simon/Downloads/$filename.tar",
    "/home/simon/Documents/$filename.mp4", overwrite=true)

convert_video_to_gif(
    "/home/simon/Documents/$filename.mp4",
    "/home/simon/Documents/$filename.gif", overwrite=true)

# Ghost
flamingo_ghost!(vis, sim, piecewise_smoothed)

# Animation
anim, shift_traj = flamingo_animation!(vis, sim, piecewise_smoothed)
anim = visualize_force!(vis, sim.model, shift_traj, anim=anim, h=h_sim, sample=10)


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#

insts = OrderedDict()
insts["minimal"] = [
    ((2, 4, 1.0),),
    ((2, 4, 0.5),),
    ]
insts["fast"] = [
    ((15, 20, 1.0),),
    ((15, 20, 0.25),),
    ((100, 100, 1.0),),
    ((100, 100, 0.15),),
    ((500, 100, 1.0),),
    ((500, 100, 0.15),),
    ]
insts["various"] = [
    ((250, 1000, 0.1),),
    ((500, 5000, 0.05),),
    ((500, 2000, 1.0),),
    ]
return (LinearOptNative, insts)