function Plane(arg0::Plane)
    return Plane((Plane,), arg0)
end

function Plane(arg0::Vector3D, arg1::Vector3D, arg2::Vector3D, arg3::jdouble)
    return Plane((Vector3D, Vector3D, Vector3D, jdouble), arg0, arg1, arg2, arg3)
end

function Plane(arg0::Vector3D, arg1::Vector3D, arg2::jdouble)
    return Plane((Vector3D, Vector3D, jdouble), arg0, arg1, arg2)
end

function Plane(arg0::Vector3D, arg1::jdouble)
    return Plane((Vector3D, jdouble), arg0, arg1)
end

function contains(obj::Plane, arg0::Vector3D)
    return jcall(obj, "contains", jboolean, (Vector3D,), arg0)
end

function copy_self(obj::Plane)
    return jcall(obj, "copySelf", Plane, ())
end

function empty_hyperplane(obj::Plane)
    return jcall(obj, "emptyHyperplane", SubPlane, ())
end

function get_normal(obj::Plane)
    return jcall(obj, "getNormal", Vector3D, ())
end

function get_offset(obj::Plane, arg0::JVector)
    return jcall(obj, "getOffset", jdouble, (JVector,), arg0)
end

function get_offset(obj::Plane, arg0::Plane)
    return jcall(obj, "getOffset", jdouble, (Plane,), arg0)
end

function get_offset(obj::Plane, arg0::Point)
    return jcall(obj, "getOffset", jdouble, (Point,), arg0)
end

function get_origin(obj::Plane)
    return jcall(obj, "getOrigin", Vector3D, ())
end

function get_point_at(obj::Plane, arg0::Vector2D, arg1::jdouble)
    return jcall(obj, "getPointAt", Vector3D, (Vector2D, jdouble), arg0, arg1)
end

function get_tolerance(obj::Plane)
    return jcall(obj, "getTolerance", jdouble, ())
end

function get_u(obj::Plane)
    return jcall(obj, "getU", Vector3D, ())
end

function get_v(obj::Plane)
    return jcall(obj, "getV", Vector3D, ())
end

function intersection(::Type{Plane}, arg0::Plane, arg1::Plane, arg2::Plane)
    return jcall(Plane, "intersection", Vector3D, (Plane, Plane, Plane), arg0, arg1, arg2)
end

function intersection(obj::Plane, arg0::Line)
    return jcall(obj, "intersection", Vector3D, (Line,), arg0)
end

function intersection(obj::Plane, arg0::Plane)
    return jcall(obj, "intersection", Line, (Plane,), arg0)
end

function is_similar_to(obj::Plane, arg0::Plane)
    return jcall(obj, "isSimilarTo", jboolean, (Plane,), arg0)
end

function project(obj::Plane, arg0::Point)
    return jcall(obj, "project", Point, (Point,), arg0)
end

function reset(obj::Plane, arg0::Plane)
    return jcall(obj, "reset", void, (Plane,), arg0)
end

function reset(obj::Plane, arg0::Vector3D, arg1::Vector3D)
    return jcall(obj, "reset", void, (Vector3D, Vector3D), arg0, arg1)
end

function revert_self(obj::Plane)
    return jcall(obj, "revertSelf", void, ())
end

function rotate(obj::Plane, arg0::Vector3D, arg1::Rotation)
    return jcall(obj, "rotate", Plane, (Vector3D, Rotation), arg0, arg1)
end

function same_orientation_as(obj::Plane, arg0::Hyperplane)
    return jcall(obj, "sameOrientationAs", jboolean, (Hyperplane,), arg0)
end

function to_space(obj::Plane, arg0::JVector)
    return jcall(obj, "toSpace", Vector3D, (JVector,), arg0)
end

function to_space(obj::Plane, arg0::Point)
    return jcall(obj, "toSpace", Vector3D, (Point,), arg0)
end

function to_sub_space(obj::Plane, arg0::JVector)
    return jcall(obj, "toSubSpace", Vector2D, (JVector,), arg0)
end

function to_sub_space(obj::Plane, arg0::Point)
    return jcall(obj, "toSubSpace", Vector2D, (Point,), arg0)
end

function translate(obj::Plane, arg0::Vector3D)
    return jcall(obj, "translate", Plane, (Vector3D,), arg0)
end

function whole_hyperplane(obj::Plane)
    return jcall(obj, "wholeHyperplane", SubPlane, ())
end

function whole_space(obj::Plane)
    return jcall(obj, "wholeSpace", PolyhedronsSet, ())
end


export K1_1StructuredParams

"""
Type to use the K1.1 formulation with a structured Krylov method, using the package
[`Krylov.jl`](https://github.com/JuliaSmoothOptimizers/Krylov.jl).
This only works for solving Linear Problems.
The outer constructor

    K1_1StructuredParams(; uplo = :L, kmethod = :lsqr, rhs_scale = true,
                         atol0 = 1.0e-4, rtol0 = 1.0e-4,
                         atol_min = 1.0e-10, rtol_min = 1.0e-10,
                         ρ_min = 1e3 * sqrt(eps()), δ_min = 1e4 * sqrt(eps()),
                         mem = 20)

creates a [`RipQP.SolverParams`](@ref) that should be used to create a [`RipQP.InputConfig`](@ref).
The available methods are:
- `:lslq`
- `:lsqr`
- `:lsmr`

"""
mutable struct K1_1StructuredParams <: NormalParams
  uplo::Symbol
  kmethod::Symbol
  rhs_scale::Bool
  atol0::Float64
  rtol0::Float64
  atol_min::Float64
  rtol_min::Float64
  ρ_min::Float64
  δ_min::Float64
  mem::Int
end

function K1_1StructuredParams(;
  uplo::Symbol = :L,
  kmethod::Symbol = :lsqr,
  rhs_scale::Bool = true,
  atol0::T = 1.0e-4,
  rtol0::T = 1.0e-4,
  atol_min::T = 1.0e-10,
  rtol_min::T = 1.0e-10,
  ρ_min::T = 1e3 * sqrt(eps()),
  δ_min::T = 1e4 * sqrt(eps()),
  mem::Int = 20,
) where {T <: Real}
  return K1_1StructuredParams(
    uplo,
    kmethod,
    rhs_scale,
    atol0,
    rtol0,
    atol_min,
    rtol_min,
    ρ_min,
    δ_min,
    mem,
  )
end

mutable struct PreallocatedDataK1_1Structured{T <: Real, S, Ksol <: KrylovSolver} <:
               PreallocatedDataNormalStructured{T, S}
  E::S  # temporary top-left diagonal
  invE::S
  ξ1::S
  ξ2::S # todel
  Δy0::S
  ξ12::S # todel
  rhs_scale::Bool
  regu::Regularization{T}
  KS::Ksol
  atol::T
  rtol::T
  atol_min::T
  rtol_min::T
end

function PreallocatedData(
  sp::K1_1StructuredParams,
  fd::QM_FloatData{T},
  id::QM_IntData,
  itd::IterData{T},
  pt::Point{T},
  iconf::InputConfig{Tconf},
) where {T <: Real, Tconf <: Real}
  @assert iconf.solve_method == :IPF

  # init Regularization values
  E = similar(fd.c, id.nvar)
  if iconf.mode == :mono
    regu =
      Regularization(T(sqrt(eps()) * 1e5), T(sqrt(eps()) * 1e5), T(sp.ρ_min), T(sp.δ_min), :classic)
    E .= T(1.0e0) / 2
  else
    regu = Regularization(
      T(sqrt(eps()) * 1e5),
      T(sqrt(eps()) * 1e5),
      T(sqrt(eps(T)) * 1e0),
      T(sqrt(eps(T)) * 1e0),
      :classic,
    )
    E .= T(1.0e-2)
  end
  if regu.δ_min == zero(T) # gsp for gpmr
    regu.δ = zero(T)
  end
  invE = similar(E)
  invE .= one(T) ./ E

  ξ1 = similar(fd.c, id.nvar)
  ξ2 = similar(fd.c, id.ncon)
  Δy0 = similar(fd.c, id.ncon)
  ξ12 = similar(fd.c, id.nvar)

  KS = init_Ksolver(fd.uplo == :U ? fd.A : fd.A', ξ12, sp)

  return PreallocatedDataK1_1Structured(
    E,
    invE,
    ξ1,
    ξ2,
    Δy0,
    ξ12,
    sp.rhs_scale,
    regu,
    KS,
    T(sp.atol0),
    T(sp.rtol0),
    T(sp.atol_min),
    T(sp.rtol_min),
  )
end

function solver!(
  dd::AbstractVector{T},
  pad::PreallocatedDataK1_1Structured{T},
  dda::DescentDirectionAllocs{T},
  pt::Point{T},
  itd::IterData{T},
  fd::Abstract_QM_FloatData{T},
  id::QM_IntData,
  res::AbstractResiduals{T},
  cnts::Counters,
  T0::DataType,
  step::Symbol,
) where {T <: Real}
  pad.ξ1 .= @views step == :init ? fd.c : dd[1:(id.nvar)]
  pad.ξ2 .= @views (step == :init && all(dd[(id.nvar + 1):end] .== zero(T))) ? one(T) :
         dd[(id.nvar + 1):end]

  if pad.regu.δ == zero(T)
    pad.Δy0 .= zero(T)
  else
    pad.Δy0 .= .-pad.ξ2 ./ pad.regu.δ
  end
  if fd.uplo == :U
    mul!(pad.ξ12, fd.A, pad.Δy0)
  else
    mul!(pad.ξ12, fd.A', pad.Δy0)
  end
  pad.ξ12 .+= pad.ξ1
  if pad.rhs_scale
    ξ12Norm = kscale!(pad.ξ12)
  end
  ksolve!(
    pad.KS,
    fd.uplo == :U ? fd.A : fd.A',
    pad.ξ12,
    Diagonal(pad.invE),
    pad.regu.δ,
    verbose = 0,
    atol = pad.atol,
    rtol = pad.rtol,
  )
  if pad.rhs_scale
    kunscale!(pad.KS.x, ξ12Norm)
  end
  dd[(id.nvar + 1):end] .= pad.KS.x .- pad.Δy0
  if fd.uplo == :U
    @views mul!(pad.ξ1, fd.A, dd[(id.nvar + 1):end], one(T), -one(T))
  else
    @views mul!(pad.ξ1, fd.A', dd[(id.nvar + 1):end], one(T), -one(T))
  end
  dd[1:(id.nvar)] .= pad.ξ1 ./ pad.E
  update_kresiduals_history_K1struct!(
    res,
    fd.uplo == :U ? fd.A' : fd.A,
    pad.E,
    pad.invE, # tmp storage vector
    pad.regu.δ,
    pad.KS.x,
    pad.ξ12,
    :K1_1,
  )

  return 0
end

function update_pad!(
  pad::PreallocatedDataK1_1Structured{T},
  dda::DescentDirectionAllocs{T},
  pt::Point{T},
  itd::IterData{T},
  fd::Abstract_QM_FloatData{T},
  id::QM_IntData,
  res::AbstractResiduals{T},
  cnts::Counters,
  T0::DataType,
) where {T <: Real}
  if cnts.k != 0
    update_regu!(pad.regu)
  end

  if pad.atol > pad.atol_min
    pad.atol /= 10
  end
  if pad.rtol > pad.rtol_min
    pad.rtol /= 10
  end

  pad.E .= pad.regu.ρ
  pad.E[id.ilow] .+= pt.s_l ./ itd.x_m_lvar
  pad.E[id.iupp] .+= pt.s_u ./ itd.uvar_m_x
  pad.invE .= one(T) ./ pad.E

  return 0
end


module TestSndfile

include("testhelpers.jl")

using AudioIO
using FactCheck
import AudioIO: DeviceInfo, render, AudioSample, AudioBuf

facts("WAV file write/read") do
    fname = Pkg.dir("AudioIO", "test", "sinwave.wav")

    srate = 44100
    freq = 440
    t = [0 : 2 * srate - 1] / srate
    phase = 2 * pi * freq * t
    reference = int16((2 ^ 15 - 1) * sin(phase))

    AudioIO.open(fname, "w") do f
        write(f, reference)
    end

    # test basic reading
    AudioIO.open(fname) do f
        @fact f.sfinfo.channels => 1
        @fact f.sfinfo.frames => 2 * srate
        actual = read(f)
        @fact length(reference) => length(actual)
        @fact reference => actual[:, 1]
        @fact samplerate(f) => srate
    end

    # test seeking

    # test rendering as an AudioNode
    AudioIO.open(fname) do f
        # pretend we have a stream at the same rate as the file
        bufsize = 1024
        input = zeros(AudioSample, bufsize)
        test_info = DeviceInfo(srate, bufsize)
        node = FilePlayer(f)
        # convert to floating point because that's what AudioIO uses natively
        expected = convert(AudioBuf, reference ./ (2^15))
        buf = render(node, input, test_info)
        @fact expected[1:bufsize] => buf[1:bufsize]
        buf = render(node, input, test_info)
        @fact expected[bufsize+1:2*bufsize] => buf[1:bufsize]
    end
end

facts("Stereo file reading") do
    fname = Pkg.dir("AudioIO", "test", "440left_880right.wav")
    srate = 44100
    t = [0 : 2 * srate - 1] / srate
    expected = int16((2^15-1) * hcat(sin(2pi*t*440), sin(2pi*t*880)))

    AudioIO.open(fname) do f
        buf = read(f)
        @fact buf => mse(expected, 5)
    end
end

# note - currently AudioIO just mixes down to Mono. soon we'll support this
# new-fangled stereo sound stuff
facts("Stereo file rendering") do
    fname = Pkg.dir("AudioIO", "test", "440left_880right.wav")
    srate = 44100
    bufsize = 1024
    input = zeros(AudioSample, bufsize)
    test_info = DeviceInfo(srate, bufsize)
    t = [0 : 2 * srate - 1] / srate
    expected = convert(AudioBuf, 0.5 * (sin(2pi*t*440) + sin(2pi*t*880)))

    AudioIO.open(fname) do f
        node = FilePlayer(f)
        buf = render(node, input, test_info)
        @fact buf[1:bufsize] => mse(expected[1:bufsize])
        buf = render(node, input, test_info)
        @fact buf[1:bufsize] => mse(expected[bufsize+1:2*bufsize])
    end
end

end # module TestSndfile


__precompile__(true)
module  MyNGSIM1

# Using packages
using DataFrames
using Reexport
@reexport using AutomotiveDrivingModels
@reexport using AutoViz
@reexport using NGSIM

export
    # functions
    get_scene_id,
    get_neighbor_fore_along_lane_NGSIM,
    get_neighbor_fore_along_left_lane_NGSIM,
    get_neighbor_fore_along_right_lane_NGSIM,
    get_neighbor_rear_along_lane_NGSIM,
    get_neighbor_rear_along_left_lane_NGSIM,
    get_neighbor_rear_along_right_lane_NGSIM,
    # Raw trajectory data in Type NGSIMTrajectorydata
    tdraw_my,
    # Processed trajectory data
    trajdata_my,
    # My road way
    ROADWAY_my,
    # Customized data types
    LaneChangeInfo,
    RefinedLaneChangeInfo,
    LaneKeepInfo,
    RefinedLaneKeepInfo

# =============================================================================
# Data Type
# =============================================================================
# Define a structure for storing lane change information of vehicles.
struct LaneChangeInfo
    changto::AbstractString         # Changed to which direction 'left' or 'right'.
    from_to::Array{Int32, 1}       # [laneID_1, laneID_2] From laneID_1 to laneID_2
    cam_frames::Array{Int, 1}       # The camera frames during which the change happened.
    df_frames::Array{Int,1}         # The dataframe frames during which the change happened.
end

struct RefinedLaneChangeInfo
    changto::AbstractString         # Changed to which direction 'left' or 'right'.
    from_to::Array{Int32, 1}       # [laneID_1, laneID_2] From laneID_1 to laneID_2
    cam_frame::Int                 # The exact camera frame where the vehicle crossed the lane.
    df_frame::Int                  # The exact dataframe frame where the vehicle crossed the lane.
end

# Define a structure for storing lane change information of vehicles.
struct LaneKeepInfo
    cam_frames::Array{Int, 1}       # The camera frames during which the change happened.
    df_frames::Array{Int,1}         # The dataframe frames during which the change happened.
end

struct RefinedLaneKeepInfo
    cam_frame::Int                  # The exact camera frame where the vehicle started braking.
    df_frame::Int                   # The exact dataframe frame where the vehicle started braking.
end

# =============================================================================
# Variables
# =============================================================================
# We define the reference points on the vehicle body. For example, when we need
# to measure the distance between two vehicles, the distance can be from the center point
# of vehcile 1 to the center point of vehicle 2, with reference point VEHICLE_TARGET_POINT_CENTER;
# or the distance can be from the front of vehicle 1 to the back of vehicle 2,
# with reference points VEHICLE_TARGET_POINT_FRONT, VEHICLE_TARGET_POINT_REAR, respectively.
const VEHICLE_TARGET_POINT_CENTER = VehicleTargetPointCenter()
const VEHICLE_TARGET_POINT_FRONT = VehicleTargetPointFront()
const VEHICLE_TARGET_POINT_REAR = VehicleTargetPointRear()

# Load the traffic data of us_101, the units are meters, m/s
# 1. "i101_trajectories-0750am-0805am.txt"
# 2. "i101_trajectories-0805am-0820am.txt"
# 3. "i101_trajectories-0820am-0835am.txt"
# 4. "i80_trajectories-0400-0415.txt",
# 5. "i80_trajectories-0500-0515.txt",
# 6. "i80_trajectories-0515-0530.txt"
# Indicate the proper dataset here.
dataset_num = 1
#
if dataset_num == 1
    trajdata_my = load_trajdata(1) # "i101_trajectories-0750am-0805am.txt"
    road_filename = "i101_trajectories-0750am-0805am.txt"
    ROADWAY_my = ROADWAY_101
elseif dataset_num == 2
    trajdata_my = load_trajdata(2) # "i101_trajectories-0805am-0820am.txt"
    road_filename = "i101_trajectories-0805am-0820am.txt"
    ROADWAY_my = ROADWAY_101
elseif dataset_num == 3
    trajdata_my = load_trajdata(3) # "i101_trajectories-0820am-0835am.txt"
    road_filename = "i101_trajectories-0820am-0835am.txt"
    ROADWAY_my = ROADWAY_101
elseif dataset_num == 4
    trajdata_my = load_trajdata(4) # "i80_trajectories-0400-0415.txt",
    road_filename = "i80_trajectories-0400-0415.txt"
    ROADWAY_my = ROADWAY_80
elseif dataset_num == 5
    trajdata_my = load_trajdata(5) # "i80_trajectories-0500-0515.txt"
    road_filename = "i80_trajectories-0500-0515.txt"
    ROADWAY_my = ROADWAY_80
elseif dataset_num == 6
    trajdata_my = load_trajdata(6) # "i80_trajectories-0515-0530.txt"
    road_filename = "i80_trajectories-0515-0530.txt"
    ROADWAY_my = ROADWAY_80
else
    error("not appropriate dataset number")
end
road_filepath = Pkg.dir("NGSIM", "data", road_filename)
# --Load the un-converted data of us_101. This data is in DataFrammes structure,
# The units are feet, foot/s
tdraw_my = NGSIM.load_ngsim_trajdata(road_filepath)

# =============================================================================
# Functions of MyNGSIM
# =============================================================================
# We define a function to find vehicle scene id from its NGSIM id.
function get_scene_id{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int
    )
    # Use NGSIM_vehicle_index to search for the scene vehicel index.
    # --Get the total number of the vehicles in the current scene.
    total_in_scene = scene.n
    # Note: Be careful, in Julia, if you want to return a value from a for-loop,
    # you should define the corresponding variable before the for-loop. For instance,
    # vehicle_index in the following; otherwise, nothing is returned from the for-loop.
    vehicle_index = 0
    for scene_ind = 1:total_in_scene
        # If the NGSIM_vehicle_index is the same in the current vehicle, this vehicle
        # is what we are looking for.
        if scene[scene_ind].id == NGSIM_vehicle_index
            vehicle_index = scene_ind
            break
        # Otherwise, if the current scene_ind is the last one, a.k.a, we have searched
        # all the vehicles in the scene, we return an error saying no finding.
        elseif scene_ind == total_in_scene
            error("This wanted vehicle is not in the current scene.")
        end
        # If neither case happens, this for-loop goes to the next iteration.
    end
    vehicle_index # return the vehicle' index in the scene
end

# Add a new method to function get_neigbor_fore_along_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_fore_along_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_fore::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_fore_along_lane() from the package AutomotiveDrivingModels
    get_neighbor_fore_along_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_FRONT, VEHICLE_TARGET_POINT_REAR,
        VEHICLE_TARGET_POINT_FRONT, max_distance_fore=max_distance_fore)
end

# Add a new method to function get_neigbor_fore_along_left_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_fore_along_left_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_fore::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_fore_along_left_lane() from the package AutomotiveDrivingModels
    get_neighbor_fore_along_left_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_FRONT, VEHICLE_TARGET_POINT_REAR,
        VEHICLE_TARGET_POINT_FRONT, max_distance_fore=max_distance_fore)
end

# Add a new method to function get_neigbor_fore_along_right_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_fore_along_right_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_fore::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_fore_along_right_lane() from the package AutomotiveDrivingModels
    get_neighbor_fore_along_right_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_FRONT, VEHICLE_TARGET_POINT_REAR,
        VEHICLE_TARGET_POINT_FRONT, max_distance_fore=max_distance_fore)
end

# Add a new method to function get_neigbor_rear_along_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_rear_along_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_rear::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_rear_along_lane() from the package AutomotiveDrivingModels
    get_neighbor_rear_along_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_REAR, VEHICLE_TARGET_POINT_FRONT,
        VEHICLE_TARGET_POINT_REAR, max_distance_rear=max_distance_rear)
end

# Add a new method to function get_neigbor_rear_along_left_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_rear_along_left_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_rear::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_rear_along_left_lane() from the package AutomotiveDrivingModels
    get_neighbor_rear_along_left_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_REAR, VEHICLE_TARGET_POINT_FRONT,
        VEHICLE_TARGET_POINT_REAR, max_distance_rear=max_distance_rear)
end

# Add a new method to function get_neigbor_rear_along_right_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_rear_along_right_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_rear::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_rear_along_left_lane() from the package AutomotiveDrivingModels
    get_neighbor_rear_along_right_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_REAR, VEHICLE_TARGET_POINT_FRONT,
        VEHICLE_TARGET_POINT_REAR, max_distance_rear=max_distance_rear)
end

end # end of module


using GeoInterface
using LibGEOS

using LinearAlgebra
using LaTeXStrings
using Formatting
import PyPlot; plt=PyPlot
cbbPalette = ["#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7"]
# using ..NYC: read_distr_reprpoints
using Distributions: MultivariateNormal
using Statistics: mean
using GaussianProcesses: predict_f
using PyCall
# mplot3d = pyimport("mpl_toolkits.mplot3d")

function plot_polygon(poly_coords, facecolor, alpha;
                      edgecolor="none", linestyle="None", linewidth=0.0, kwargs...)
    ax = plt.gca()
    poly_array = Array{Float64}(undef, length(poly_coords), 2)
    for i in 1:length(poly_coords)
        poly_array[i, :] = poly_coords[i]
    end
    polygon = plt.plt.Polygon(poly_array, true; facecolor=facecolor, kwargs...)
    p = ax.add_patch(polygon)
    p.set_edgecolor(edgecolor)
    p.set_linestyle(linestyle)
    p.set_linewidth(linewidth)
    p.set_alpha(alpha)
end

function plot_buffer(shape, border, color, dist, alpha; kwargs...)
    # a buffer zone around the border:
    border_buffer = LibGEOS.buffer(border, dist)
    # prevents a ring self-intersection bug (bad shapefile?):
    shape_rectified = LibGEOS.buffer(shape, 0.0)
    # the part of the buffer on the `shape` side of the border:
    shape_buffer = LibGEOS.intersection(border_buffer, shape_rectified)
    if typeof(shape_buffer) <: LibGEOS.Polygon
        poly_coords = GeoInterface.coordinates(shape_buffer)[1]
        plot_polygon(poly_coords, color, alpha; kwargs...)
    elseif typeof(shape_buffer) <: LibGEOS.MultiPolygon
        mpoly_coords = GeoInterface.coordinates(shape_buffer)
        for poly_coords in mpoly_coords
            plot_polygon(poly_coords[1], color, alpha; kwargs...)
        end
    else
        println(typeof(shape_buffer))
    end
end

function plot_schdistr_labels(;fontsize=12.0, labelcolor="black")
    plt.rc("text", usetex=false)
    # where to place the school district labels:
    schdistr_represent = Dict{SchDistr, Tuple{Float64, Float64}}(25 => (1035426.576356058, 215708.47750849032), 20 => (980957.6301213544, 166113.83169550286), 3 => (992826.0954164147, 226695.6766967253), 4 => (1000248.2836855762, 228298.3092956051), 16 => (1004214.0439845352, 188657.53201288398), 23 => (1008819.7413809987, 183460.7804870081), 18 => (1008787.3316119347, 173443.51080316756), 26 => (1054423.2425350049, 212717.9963988735), 32 => (1006895.1698167967, 192470.89239496225), 1 => (988952.4509786353, 202318.4309997029), 2 => (987095.1379395023, 210542.70410151145), 24 => (1013178.9222047399, 202219.43350214284), 28 => (1036851.1969587662, 195628.9613036562), 10 => (1008845.8110727321, 258295.0202941391), 9 => (1008189.8006258132, 245065.8830871096), 7 => (1006949.3221918349, 235790.44049067114), 30 => (1011393.8551870013, 216959.78259271727), 12 => (1017667.4294655474, 243924.23712153034), 8 => (1022714.7969213917, 237308.7890929639), 14 => (998121.9324939251, 199650.8938903253), 31 => (947923.9119948001, 147916.42111200799), 17 => (997069.2552085526, 180290.71820062993), 13 => (994552.5771594604, 190230.17111200746), 15 => (987418.7441204586, 181218.70071406086), 22 => (1001740.2660473757, 162368.14001459582), 21 => (991081.0983826279, 157637.68869012795), 11 => (1026344.6177574333, 256966.2814940873), 29 => (1055388.4397568335, 188398.91879267507), 19 => (1017942.166879338, 181833.5946349586), 27 => (1040711.4984081909, 178979.62564887386), 5 => (996548.0388532034, 237031.53149408224), 6 => (1001541.6608069859, 248002.3049010716))
    xmin,xmax = plt.xlim()
    ymin,ymax = plt.ylim()
    ax = plt.gca()
    for schdistr in keys(schdistr_represent)
        center_x, center_y = schdistr_represent[schdistr]

        if !(xmin < center_x < xmax)
            continue
        elseif !(ymin < center_y < ymax)
            continue
        end

        ax.annotate(
            @sprintf("%d", schdistr),
            xy=(center_x, center_y),
            xycoords="data",
            xytext=(0,0),
            textcoords="offset points",
            horizontalalignment="center",
            verticalalignment="center",
            fontweight="black",
            color=labelcolor,
            alpha=1.0,
            fontsize = fontsize,
            zorder=3,
            annotation_clip=true,
        )
        ax.annotate(
            @sprintf("%d", schdistr),
            xy=(center_x, center_y),
            xycoords="data",
            xytext=(0,0),
            textcoords="offset points",
            horizontalalignment="center",
            verticalalignment="center",
            fontweight="black",
            color=labelcolor,
            alpha=0.5,
            fontsize = fontsize,
            zorder=20,
            annotation_clip=true,
        )
    end
    plt.rc("text", usetex=true)
end

function plot_all_pairs(τpost_pairs::Dict{Tuple{SchDistr,SchDistr}}, regionData_dict::Dict{SchDistr,GeoRDD.RegionData}; scaleup=200.0, zorder="none")
    for distr_pair in keys(τpost_pairs)
        distrA, distrB = distr_pair

        τpost = τpost_pairs[distr_pair]
        # eff_size = abs(mean(τpost)) / std(τpost) * 2
        # if eff_size < 2.0
            # continue
        # end

        if distrA ∈ (1,2) || distrB ∈ (1,2)
            # don't do Manhattan, it's weird
            continue
        end

        # skip districts that aren't on the map
        if distrA ∉ (13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 32)
            continue
        end
        if distrB ∉ (13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 32)
            continue
        end

        if distrA >= distrB
            continue
        end

        shapeA = regionData_dict[distrA].shape
        shapeB = regionData_dict[distrB].shape
        border = GeoRDD.get_border(shapeA, shapeB, 10.0)

        border_coords = GeoInterface.coordinates(border)
        if typeof(border) <: LibGEOS.MultiLineString
        else
            border_coords = [border_coords]
        end
        ax = plt.gca()

        if mean(τpost) < 0
            cheaper_shape = shapeA
            pricier_shape = shapeB
        else
            cheaper_shape = shapeB
            pricier_shape = shapeA
        end
        plot_buffer(pricier_shape, border, cbbPalette[2], (abs(mean(τpost)) + std(τpost))*scaleup, 1.0, zorder=zorder-1)
        plot_buffer(pricier_shape, border, cbbPalette[1], abs(mean(τpost))*scaleup, 1.0, zorder=zorder)
    end
end

function plot_streets(;zoom=11, brighten=30, rgb=(200, 200, 200), zorder=-9)
    ax = plt.gca()
    ax.set_autoscale_on(false)
    ax.set_autoscalex_on(false)
    ax.set_autoscaley_on(false)
    EPSG=2263 # projection
    topleft =  (40.75, -74.01)
    botright = (40.58, -73.60)
    py"""
    import matplotlib.pyplot as plt

    import geopandas as gpd
    import tilemapbase
    import pyproj
    import PIL.Image as _Image

    def bounds(plotter, epsg):
        epsg_proj = pyproj.Proj('epsg:%d' % epsg, preserve_units=True)
        epsg_3857 = pyproj.Proj('epsg:%d' % tilemapbase.mapping._WEB_MERCATOR)
        self = plotter
        scale = 2 ** self.zoom
        x0, y0 = self.xtilemin / scale, self.ytilemin / scale # in tile space
        x1, y1 = (self.xtilemax + 1) / scale, (self.ytilemax + 1) / scale
        x0_ll, y0_ll = tilemapbase.to_lonlat(x0, y0)
        x1_ll, y1_ll = tilemapbase.to_lonlat(x1, y1)
        y0_ll, y1_ll = y1_ll, y0_ll # for some reason that I don't understand
        print(y0_ll)
        print(y1_ll)
        x0_epsg, y0_epsg = epsg_proj(x0_ll, y0_ll)
        x1_epsg, y1_epsg = epsg_proj(x1_ll, y1_ll)
        return x0_epsg, x1_epsg, y0_epsg, y1_epsg

    def brighten(img, factor):
        # split the image into individual bands
        source = img.split()
        R, G, B, A = 0, 1, 2, 3

        # process each band separately
        out_R = source[R].point(lambda x: min(x+factor, 256))
        out_G = source[G].point(lambda x: min(x+factor, 256))
        out_B = source[B].point(lambda x: min(x+factor, 256))

        # build a new multiband image
        brightened = _Image.merge(img.mode, (out_R, out_G, out_B, source[A]))
        return brightened

    def makecolor(img, rgb):
        # split the image into individual bands
        source = img.split()
        R, G, B, A = 0, 1, 2, 3

        # process each band separately
        out_R = source[R].point(lambda x: rgb[0])
        out_G = source[G].point(lambda x: rgb[1])
        out_B = source[B].point(lambda x: rgb[2])

        # build a new multiband image
        new_im = _Image.merge(img.mode, (out_R, out_G, out_B, source[A]))
        return new_im

    tilemapbase.start_logging()
    tilemapbase.init("/Users/imolk/tmp/tilemapbase.cache", create=True)

    extent = tilemapbase.Extent.from_lonlat($topleft[1], $botright[1],
                      $botright[0], $topleft[0])
    # plotter_ln = tilemapbase.Plotter(extent,tilemapbase.tiles.Carto_Dark_No_Labels, zoom=$zoom)
    plotter_ln = tilemapbase.Plotter(extent,tilemapbase.tiles.Stamen_Terrain_Lines, zoom=$zoom)
    plotter = plotter_ln
    # tile = brighten(plotter.as_one_image(), $brighten)
    tile = makecolor(plotter.as_one_image(), $rgb)
    $ax.imshow(tile, interpolation="lanczos", extent=bounds(plotter, $EPSG), origin="upper", zorder=$zorder)
    """
end

function ticks(xmin, xmax, tickdist)
    return tickdist * (ceil(xmin/tickdist):1.0:floor(xmax/tickdist))
end
function ticks_km(tickdist::Float64)
    _xlim = plt.xlim()
    _ylim = plt.ylim()
    km_in_ft = 3280.84
    _xlim_km = _xlim ./ km_in_ft
    _ylim_km = _ylim ./ km_in_ft
    _xtick_km = ticks(_xlim_km[1], _xlim_km[2], tickdist)
    _ytick_km = ticks(_ylim_km[1], _ylim_km[2], tickdist)
    _xtick_ft = _xtick_km .* km_in_ft
    _ytick_ft = _ytick_km .* km_in_ft
    plt.xticks(_xtick_ft, [@sprintf("%.0f", x) for x in _xtick_km])
    plt.yticks(_ytick_ft, [@sprintf("%.0f", x) for x in _ytick_km])
    plt.xlabel("Eastings (km)")
    plt.ylabel("Northings (km)")
end

function plot_sales(filtered, missing_sqft; edgecolor="black", labelsize=12.0, labelcolor="black", colorbar=true, tickdist=5.0)
    fig=plt.figure()
    background_schdistrs(plt.gca(),
                color="#AAAAAA",
                edgecolor=edgecolor,
                alpha=1.0,
                linestyle="-",
                zorder=-10
                )
    background_schdistrs(plt.gca(),
                color="none",
                edgecolor=edgecolor,
                alpha=1.0,
                linestyle="-",
                linewidth=0.5,
                zorder=11)
    _xlim = (0.98e6, 1.06e6)
    _ylim = (1.5e5, 2.2e5)
    plt.plot(
        missing_sqft[!,:XCoord],
        missing_sqft[!,:YCoord],
        color="white",
        marker="x",
        markersize=1.0,
        linestyle="",
        zorder=8,
        alpha=0.7,
        linewidth=0.1,
        markeredgewidth=0.3,
    )
    plt.scatter(
        filtered[!,:XCoord], filtered[!,:YCoord],
        c=filtered[!,:logSalePricePerSQFT],
        cmap="jet",
        edgecolors="None",
        zorder=10,
        marker="o",
    #     vmin=4.0,
    #     vmax=8.0,
        alpha=1.0,
        s=1,
        )
    # plt.xlim(np.percentile(sales_geocoded_ggl["XCoord"].dropna(), [1,99]))
    # plt.ylim(np.percentile(sales_geocoded_ggl["YCoord"].dropna(), [1,99]))
    # plt.axis("off")
    # plt.title("log(Property Price per SQFT) in New York")
    ax = plt.gca()
    ax.set(aspect="equal", adjustable="box")
    # ax.get_xaxis().set_ticks([])
    # ax.get_yaxis().set_ticks([])
    plt.xlim(_xlim)
    plt.ylim(_ylim)
    ticks_km(tickdist)

    ax.set_facecolor("#B0DAEE")
    if colorbar
        cbar = plt.colorbar(label="Square foot prices (\\\$ per sqft)")
        ax = cbar.ax
        tick_template = [1, 2, 3, 4, 5, 6, 7, 8, 9]
        labeled = [10, 100, 1000, 10^4]
        ax.set_yticks(log.(labeled))
        ax.set_yticks(log.([tick_template; tick_template.*10; tick_template.*100; tick_template.*1000]),
                               minor=true)
        ax.set_yticklabels([format(x, commas=true) for x in labeled])
    end

    ax.set_autoscale_on(false)
    ax.set_autoscalex_on(false)
    ax.set_autoscaley_on(false)
    plot_schdistr_labels(fontsize=labelsize, labelcolor=labelcolor)
end
function plot_cliffface(μ, Σ, color; ndraws::Int=10, label=L"posterior of $\tau(x)$ on residuals")
    plt.plot(μ, color=color, ".", label=label)
    msize=10
    plt.plot(1,μ[1], color=cbbPalette[5],
        markersize=msize, marker="o")
    plt.plot(length(μ),μ[end], color=cbbPalette[4],
        markersize=msize, marker="o")
    posterior_sd = sqrt.(diag(Σ))
    plt.fill_between(1:length(μ), μ.-2*posterior_sd, μ.+2*posterior_sd,
        color=color,
        alpha=0.3,
        linewidth=0,
        zorder=-2,
        )
    posterior_distr = MultivariateNormal(μ, Σ)
    for _ in 1:ndraws
        plt.plot(1:length(μ), rand(posterior_distr), "-", color="white", alpha=0.4, linewidth=1.5, zorder=-1)
    end
    plt.xlim(1,length(μ))
    plt.ylim(minimum(μ.-2*posterior_sd), maximum(μ.+2*posterior_sd))

    # second y-axis for percentage price increase
    ax = plt.gca()
    y1, y2=ax.get_ylim()
    x1, x2=ax.get_xlim()
    ax2=ax.twinx()
    ax2.set_ylim(y1, y2)
    yticks_frac = collect(0.5: 0.1 : 2.0)
    yticks_log = log.(yticks_frac)
    in_window = y1 .< yticks_log .< y2
    ax2.set_yticks(yticks_log[in_window])
    ylabels = [@sprintf("%.0f\\%%", y*100) for y in yticks_frac[in_window]]
    ax2.set_yticklabels(ylabels)

    plt.sca(ax)
end

# plotting convenience functions
function yaxis_right()
    ax = plt.gca()
    ax.yaxis.tick_right()
    ax.yaxis.set_label_position("right")
end
function title_in_axis(s, leftright)
    if leftright == :left
        plt.text(0.05, 0.95, s,
             horizontalalignment="left",
             verticalalignment="top",
             transform = plt.gca().transAxes)
    elseif leftright == :right
        plt.text(0.95, 0.95, s,
             horizontalalignment="right",
             verticalalignment="top",
             transform = plt.gca().transAxes)
    else
        throw("left or right?")
    end
end
function hide_xaxis()
    plt.gca().xaxis.set_ticklabels([])
    plt.xlabel("")
end
function pval_yticks()
    ax = plt.gca()
    plt.ylim(0,1)
    ax.set_yticks(0:0.1:1, minor=true)
    ax.set_yticks(0:0.2:1, minor=false)
end

# NYC_dir is GeoRDD/examples/NYC and is expected to be defined
# in the module that loads this file
NYSD_FILENAME = joinpath(NYC_dir, "NYC_data/nysd_16c/nysd.json")
py"""
import matplotlib.pyplot as plt
import pandas as pd
import geopandas as gpd
plot_dataframe = gpd.plotting.plot_dataframe
import numpy as np
import itertools
plt.interactive(False)

EPSG=2263 # projection
# import os
# dirname = os.path.dirname(__file__)
# nysd_filename = os.path.join(dirname, "NYC_data/nysd_16c/nysd.json")
nycdistrs=gpd.read_file($(NYSD_FILENAME)).to_crs(epsg=EPSG)

def background_schdistrs(ax, **kwargs):
    plot_dataframe(nycdistrs, ax=ax, **kwargs)
    return None
"""
background_schdistrs = py"background_schdistrs"


using Test

# # State debug statistics
#
# Set up a basic environment
using MPI
using StaticArrays
using Random
using ClimateMachine
using ClimateMachine.VariableTemplates
using ClimateMachine.MPIStateArrays
using ClimateMachine.StateCheck
using ClimateMachine.GenericCallbacks

@testset "$(@__FILE__)" begin

    ClimateMachine.init()
    FT = Float64

    # Define some dummy vector and tensor abstract variables with associated types
    # and dimensions
    F1 = @vars begin
        ν∇u::SMatrix{3, 2, FT, 6}
        κ∇θ::SVector{3, FT}
    end
    F2 = @vars begin
        u::SVector{2, FT}
        θ::SVector{1, FT}
    end

    # Create ```MPIStateArray``` variables with arrays to hold elements of the
    # vectors and tensors
    Q1 = MPIStateArray{Float32, F1}(
        MPI.COMM_WORLD,
        ClimateMachine.array_type(),
        4,
        9,
        8,
    )
    Q2 = MPIStateArray{Float64, F2}(
        MPI.COMM_WORLD,
        ClimateMachine.array_type(),
        4,
        3,
        8,
    )

    # ### Create a call-back
    cb = ClimateMachine.StateCheck.sccreate(
        [(Q1, "My gradients"), (Q2, "My fields")],
        1;
        prec = 15,
    )

    # ### Invoke the call-back
    #     Compare on local "realdata", fill via broadcast to keep GPU happy.
    Q1.data = rand(MersenneTwister(0), Float32, size(Q1.data))
    Q2.data = rand(MersenneTwister(0), Float64, size(Q2.data))
    Q1.realdata .= Q1.data
    Q2.realdata .= Q2.data
    GenericCallbacks.init!(cb, nothing, nothing, nothing, nothing)
    GenericCallbacks.call!(cb, nothing, nothing, nothing, nothing)

    # ### Check with reference values
    include("test_statecheck_refvals.jl")

    # This should return true (MPI rank != 0 always returns true)
    test_stat_01 = ClimateMachine.StateCheck.scdocheck(cb, (varr, parr))

    # This should return false (MPI rank != 0 always returns true)
    varr[1][3] = varr[1][3] * 10.0
    test_stat_02 = ClimateMachine.StateCheck.scdocheck(cb, (varr, parr))
    if MPI.Comm_rank(MPI.COMM_WORLD) != 0
        test_stat_02 = false
    end

    test_stat = test_stat_01 && !test_stat_02

    @test test_stat

end


using GLVisualize, GeometryTypes, GLAbstraction, GLWindow
w, r = Screen()
@async r()
a = GLNormalUVMesh(Cube{Float32}(Vec3f0(0), Vec3f0(1)))
view(visualize(a))


import PyCall
import POMDPs: MDP, Policy

using Images, Colors

#
# Atari environment is used for playing OpenAI Atari https://gym.openai.com/envs/#atari
# Games are provided in a form of a 3-channel RGB image.
# P.S. We are taking every second pixel of a frame.
#
# Our environment will consist of Python variable and action space
#
struct AtariMDP <: MDP{Vector{Float64,}, Int64}
    env::PyCall.PyObject
    action_space::Vector{Int64}
end

#
# Classic initializer which is setting action space
#
function AtariMDP(env::PyCall.PyObject)

    actions = Vector{Int64}(0:(env[:action_space][:n] - 1))

    return AtariMDP(env, actions)
end

#
# We define number of function required to play the game
#
states(mdp::AtariMDP) = -1
actions(mdp::AtariMDP) = mdp.action_space
n_states(mdp::AtariMDP) = -1
n_actions(mdp::AtariMDP) = length(mdp.action_space)
initial_state(mdp::AtariMDP, rng::AbstractRNG = MersenneTwister(0)) = state_view(mdp, mdp.env[:reset]())
isterminal(mdp::AtariMDP, s) = iszero(s)
action_index(mdp::AtariMDP, a::Int64) = a

# We take every second pixel, swap the dimensions to be image and
# convert an image to grayscale
state_view(mdp::AtariMDP, s) = channelview(Gray.(colorview(RGB, permutedims(s[1:2:end, 1:2:end, :] * 1.0, (3, 1, 2)) ./ 255.0)))

#
# The only function we can use from OpenAIGym is getting a reward when taking a
# particular action in a specific step
#
function generate_sr(mdp::AtariMDP, s, a::Int64)

    state, reward, done = mdp.env[:step](a)

    if done
        return (zeros(state_view(mdp, state)), reward)
    end

    return (state_view(mdp, state), reward)
end


export OvCF, OvController, OvMsg

import Optim

struct OvMsg{X,Z}
    δx::Timed{X}
    z::Timed{Optional{Z}}
end

to_optional(x::Timed{X}) where X = convert(Timed{Optional{X}}, x)

"""
Matrices are indexed by [time, id].
"""
struct OvLog{X,Z,U}
    ũ::Matrix{U}
    x̃::Matrix{X}
    z̃::Matrix{Z}
    δx::X
    z::Z
end

function to_array(m::MMatrix{n1,n2,T}) where {n1,n2,T}
    r = Matrix{T}(undef, n1, n2)
    copyto!(r, m)
end

"""
The OneVision Controller Framework.
"""
@kwdef struct OvCF{N,X,Z,U,S,H} <: ControllerFramework{X,Z,U,OvMsg{X,Z},OvLog{X,Z,U}}
    central::CentralControl{U,S}
    world_model::WorldDynamics{N}
    delay_model::DelayModel
    x_weights::SVector{N,X}
    u_weights::SVector{N,U}
    optim_options::Optim.Options
    save_log::FuncT{Tuple{ℕ,𝕋,X,Z},Bool} = FuncT(x -> false, Tuple{ℕ,𝕋,X,Z}, Bool)
end

function OvCF(central::CentralControl{U,S}, world_model, delay_model, x_weights, u_weights;
    X, Z, N, H,
    u_tol = 1e-4, loss_tol = 1e-6, loss_rel_tol = 1e-6, max_iters = 100,
    save_log = FuncT(x -> false, Tuple{ℕ,𝕋,X,Z}, Bool)
) where {U,S}
    OvCF{N,X,Z,U,S,H}(;central, world_model, delay_model, x_weights, u_weights, save_log,
        optim_options = Optim.Options(
            x_abstol = u_tol, f_abstol = loss_tol, f_reltol = loss_rel_tol,
            iterations = max_iters))
end

function OvCF(loss_model::RegretLossModel{N,X,U,S}, delay_model;
    Z, H,
    optim_options = Optim.Options(
        x_abstol=1e-4, f_abstol=1e-6, f_reltol=1e-6, iterations=100),
    save_log = FuncT(x -> false, Tuple{ℕ,𝕋,X,Z}, Bool)
) where {N,X,U,S}
    @unpack central, world_model, x_weights, u_weights = loss_model
    OvCF{N,X,Z,U,S,H}(;central, world_model, delay_model,
        x_weights, u_weights, save_log, optim_options)
end

@kwdef mutable struct OvController{N,X,Z,U,S,H,ΔT,Hf,WDy,Ctrl} <: Controller{X,Z,U,OvMsg{X,Z},OvLog{X,Z,U}}
    id::ℕ
    cf::OvCF{N,X,Z,U,S,H}
    τ0::𝕋
    τ::𝕋                                # -- At the start of each controll step --
    u_last::U                           # t = τ+Tu-1
    u_history::TimedQueue{U}            # t ∈ [τ-Tx,τ+Tu-1]
    "predicted x from last control step, used to measure disturbance."
    pred_x::Timed{X}                    # t = τ-Tx
    "self disturbance history"
    self_δx::TimedQueue{X}              # t ∈ [τ-Tx-Tc-ΔT-1, τ-Tx-2]
    self_z::TimedQueue{Optional{Z}}     # t ∈ [τ-Tx-Tc-ΔT, τ-Tx-1]
    ideal_xz::Timed{Each{Tuple{X,Z}}}   # t = τ-Tx-Tc-ΔT
    ideal_s::S                   # t = τ-Tx-Tc-ΔT
    fp_prob::ForwardPredictProblem{N,Hf,X,Z,U,WDy,Ctrl}
    plan_prob::TrajPlanningProblem{H,ΔT,X,U}
    logs::Dict{𝕋,OvLog{X,Z,U}}
end

OneVision.write_logs(ctrl::OvController) = ctrl.logs

"""
# Arguments
- `init_status`: The initial (x,z,u) of the agent. x and z are assumed to be
measured at `t0-Tx`.
"""
function OneVision.make_controllers(
    cf::OvCF{N,X,Z,U,S,H},
    init_status::Each{Tuple{X,Z,U}},
    τ::𝕋,
)::Tuple where {N,X,Z,U,S,H}
    dm = cf.delay_model
    @unpack Tx, Tu, Tc, Ta, ΔT = short_delay_names(dm)
    function mk_controller(id)
        x0, z0, u0 = init_status[id]
        u_history = TimedQueue(u0, τ-Tx, τ+Tu-1)
        x_dy = cf.world_model.dynamics[id]
        pred_x = Timed(τ-Tx, x0)
        self_δx = TimedQueue(zero(X), τ-Tx-Tc-ΔT-1, τ-Tx-2)
        self_z = TimedQueue(missing, τ-Tx-Tc-ΔT, τ-Tx-1, eltype=Optional{Z})
        ideal_xz = Timed(τ-Tx-Tc-ΔT, [(x, z) for (x, z, _) in init_status])
        ideal_s = init_state(cf.central, τ)

        fp_prob = let Hf = H * ΔT + Ta + ΔT
            ForwardPredictProblem(cf.world_model, cf.central; X, Z, Hf)
        end
        x_weights = cf.x_weights[id]
        u_weights = cf.u_weights[id]
        plan_prob = TrajPlanningProblem(
            H, ΔT, x_dy, x_weights, u_weights, optim_options = cf.optim_options,
        )

        OvController(;
            id, cf, τ0 = τ, τ, u_last = u0, u_history, pred_x,
            self_δx, self_z, ideal_xz, ideal_s,
            fp_prob, plan_prob, logs = Dict{𝕋,OvLog{X,Z,U}}(),
        )
    end

    function mk_q(_)
        receives = [[OvMsg(Timed(t-1, zero(x)), Timed{Optional{Z}}(t, missing))
                    for (x, z, u) in init_status
                    ] for t in range(τ-Tc-Tx, length = msg_queue_length(dm))]
        FixedQueue(receives)
    end

    ctrls::NTuple{N,OvController} = ntuple(mk_controller, N)
    msg_qs::NTuple{N,MsgQueue} = ntuple(mk_q, N)
    ctrls, msg_qs
end

function OneVision.control!(
    π::OvController{N,X,Z,U,S,H,ΔT},
    x::X,
    z::Z,
    msgs::Each{OvMsg{X,Z}},
)::Tuple{U,Each{OvMsg{X,Z}}} where {N,X,Z,U,S,H,ΔT}
    @unpack id, τ, τ0 = π
    world = π.cf.world_model
    fp_prob = π.fp_prob
    @unpack Tx, Tu, Tc, Ta = short_delay_names(π.cf.delay_model)
    zt = Timed(τ-Tx, z)
    xt = Timed(τ-Tx, x)

    δx = x - π.pred_x[τ-Tx] |> attime(τ-Tx-1)   # δx: t = τ-Tx-1
    pushpop!(π.self_δx, δx)     # π.self_δx: t ∈ [τ-Tx-1-Tc, τ-Tx-1]
    pushpop!(π.self_z, zt)       # π.self_z: t ∈ [τ-Tx-Tc, τ-Tx]
    new_msg = OvMsg(δx, to_optional(zt))
    should_log = π.cf.save_log((id, π.τ, x, zt.value))

    x_self::Timed{X} = isempty(π.u_history) ? xt : self_estimate(
        world.dynamics[id],
        xt,
        π.u_history.queue
    )[end]  # x_self: t = τ+Tu

    "the time offset within the ΔT period"
    t = mod1(τ - τ0, ΔT)

    # forward predict the ideal fleet trajectory from the state at t-Tx-Tc-1
    for j in 1:N
        if j == id continue end
        @unpack δx, z = msgs[j]
        fp_prob.δx[t, j] = δx[τ-Tx-Tc-1]
        fp_prob.z_obs[t, j] = z[τ-Tx-Tc]
    end

    if t == ΔT
        fp_prob.δx[1:length(π.self_δx),id] = collect(x.value for x in π.self_δx.queue)
        fp_prob.z_obs[1:length(π.self_z),id] = collect(x.value for x in π.self_z.queue)

        # x̃[t in τ-Tx-Tc-ΔT+1: τ+Tu+HΔT], ũ[t in τ-Tx-Tc-ΔT: τ+Tu+HΔT-1]
        ptime = τ-Tx-Tc-ΔT  # the starting time of forward_predict!
        @asserteq π.self_δx[1].time ptime
        ũ, x̃, z̃ = forward_predict!(
            fp_prob, π.ideal_xz[ptime], π.ideal_s, ptime, ΔT)
        π.ideal_xz = [(x̃[ΔT,j], z̃[ΔT,j]) for j in 1:N] |> attime(ptime+ΔT)  # π.ideal_xz: t = τ-Tx-Tc

        # local planning using path following

        # u_plan: [τ+Tu:τ+Tu+HΔT-1], x_plan: [τ+Tu+1:τ+Tu+HΔT]
        u_plan, loss = let
            n_x, n_u = length(X), length(U)
            x_path = SVector{H * ΔT}(x̃[1 + ΔT + Ta : end, id])  # [τ+Tu+1:τ+Tu+HΔT]
            u_path = SVector{H * ΔT}(ũ[1 + ΔT + Ta : end, id])  # [τ+Tu:τ+Tu+HΔT-1]
            plan_trajectory(π.plan_prob, x_self[τ+Tu], x_path, u_path, τ+Tu)
        end

        π.u_last = u_plan[1]
    end
    u = π.u_last

    π.pred_x = let
        x_dy = world.dynamics[id]
        u_old = pushpop!(π.u_history, Timed(τ+Tu, u))[τ-Tx]  # u_old at t = τ-Tx
        x1 = sys_forward(x_dy, x, u_old, τ-Tx)
        Timed(τ-Tx+1, x1)
    end

    if should_log
        π.logs[π.τ] = OvLog(to_array.((ũ, x̃, z̃))..., x, z)
    end

    π.τ += 1
    u, fill(new_msg, N)
end


struct Normal{T} end
struct Exponential{T} end

# sampler => Dict(seed => rand(LehmerRNG(seed), n))
# if key (sampler) isa Integer, it represents 1:key
const STREAMS = Dict{Any,Dict{Int64,Vector}}(
    UInt64 => Dict(
        0 => [0x45a31efc5a35d971, 0xc6106997913e7a3b, 0x949e7d1a64224226,
              0x233c600f92647349, 0x275e265b37c53f51, 0x7eb1431760d5e0d9,
              0x75a91e49d501684d, 0x2eb2f990c5bd12e0, 0xdf459aadbb8c0b08,
              0xcb81434944438ad0, 0x5385088560aef3b0, 0x0702607d43845765],
        1 => [0xd0e95cf50ea18c53, 0x52313cc6b3bb6eb2, 0xbddb774f2c66c674,
              0x69b5202eb72d59db, 0x761a7311a74fbdf5, 0x7c13c9462281a28d,
              0x60fb5add7f0438e8, 0x8c18ecb2513738a0, 0x9dd0d00932a42118,
              0x6283c9dbcccaa072, 0xfa8f1990220cdb11, 0x15072177ca8d0630],
        Int64(2)^32 =>
        [0xfa0ed1dea67579f2, 0xe88d5e0e5b8b40fe, 0x5ce30167dd0c0c22,
         0x480546a1a865a818, 0x96e8a4ff013c8a92, 0x405d04cb572b978f,
         0x1fabeee4affbf0c3, 0xba2d1f50e67a63d0, 0x565db0bdf8d96d6b,
         0x540858ead6088e6d, 0x14e2efe6ac5775fa, 0x8e0b0f48316d364b],
        typemax(Int64) =>
        [0xe07cb1442060b16b, 0x9f15f9ca1910888d, 0x27d667e3a5597f5e,
         0xe7c43a5819b14e87, 0xbd5d7f453bf3bc1b, 0x7c7998438141ee3f,
         0xf7d425f171d5da47, 0xe1abaee77149b9c0, 0x3f6116840ff765f4,
         0xfd613e850d5b6198, 0x524f389f5079bcf4, 0x6ff20ef5be83190b]),
    123 => Dict(
        0 => [0x72, 0x3c, 0x27, 0x4a, 0x52, 0x5a, 0x4e, 0x61,
              0x09, 0x51, 0x31, 0x66, 0x6d, 0x68, 0x65, 0x1b,
              0x5e, 0x07, 0x28, 0x3a, 0x33, 0x70, 0x6e, 0x35],
        1 => [0x54, 0x33, 0x75, 0x5c, 0x76, 0x0e, 0x69, 0x21,
              0x19, 0x73, 0x12, 0x31, 0x46, 0x38, 0x2d, 0x51,
              0x18, 0x15, 0x76, 0x2d, 0x17, 0x50, 0x4a, 0x1d],
        Int64(2)^32 => [0x73, 0x23, 0x19, 0x13, 0x10, 0x44, 0x51, 0x6c,
                        0x6e, 0x7b, 0x4c, 0x12, 0x19, 0x3e, 0x16, 0x6e,
                        0x06, 0x26, 0x05, 0x7a, 0x13, 0x6d, 0x6e, 0x18],
        typemax(Int64) => [0x6c, 0x0e, 0x5f, 0x08, 0x1c, 0x40, 0x48, 0x41,
                           0x75, 0x19, 0x75, 0x0c, 0x21, 0x52, 0x51, 0x43,
                           0x03, 0x2c, 0x3f, 0x6a, 0x74, 0x27, 0x66, 0x74],
    ),
    # can't have 1:UInt128(2)^100 as a key, it's too big, cf. #35603
    UInt128(2)^100 => Dict(
        0 => [0x007913e7a3b45a31efc5a35d972, 0x00f92647349949e7d1a64224227,
              0x00760d5e0d9275e265b37c53f52, 0x000c5bd12e075a91e49d501684e,
              0x00944438ad0df459aadbb8c0b09, 0x00d438457655385088560aef3b1,
              0x002292725e7bf0412ad4e6167ed, 0x0065e6e00fe22801cd76310dc65,
              0x00329916fdd3b9df5bfe8ce461b, 0x00951b9612719a9585e619de307,
              0x009f94fa7326c9f3ed5bdcadeba, 0x0046c73306dba795fe8c75a9770],
        1 => [0x006b3bb6eb2d0e95cf50ea18c54, 0x00eb72d59dbbddb774f2c66c675,
              0x0062281a28d761a7311a74fbdf6, 0x002513738a060fb5add7f0438e9,
              0x00bcccaa0729dd0d00932a42119, 0x007ca8d0630fa8f1990220cdb12,
              0x0067b7571b73d0c3807eb2437c6, 0x0031b4a02fc678056862932952d,
              0x0097cb44f97b2d9e13fba6ad251, 0x00bf52c23754cfc091b24d9a915,
              0x00debeef59645ddbc8139609c2d, 0x00d455991492f6c1fba560fc650],
        Int64(2)^32 =>
        [0x00e5b8b40fefa0ed1dea67579f3, 0x001a865a8185ce30167dd0c0c23,
         0x00b572b978f96e8a4ff013c8a93, 0x000e67a63d01fabeee4affbf0c4,
         0x00ad6088e6d565db0bdf8d96d6c, 0x008316d364b14e2efe6ac5775fb,
         0x002248fb2185bc6e28653152312, 0x00448742715e8a1d59fe134f43e,
         0x00d6a7b16ed0d3a81f5bdfa84fd, 0x0078499dfa5dce51e6be76c4206,
         0x00e6f5d2d79e834fc489240f705, 0x00fffad736c492e8ec852c85213],
        typemax(Int64) =>
        [0x00a1910888de07cb1442060b16c, 0x00819b14e8727d667e3a5597f5f,
         0x0038141ee3fbd5d7f453bf3bc1c, 0x0077149b9c0f7d425f171d5da48,
         0x0050d5b61983f6116840ff765f5, 0x005be83190b524f389f5079bcf5,
         0x0063a7004d1c355cae530dd93a1, 0x0050aaefc421c91ef1545dcbf51,
         0x004e93aea2baef829b0d3d36503, 0x005cc3461e9a93dd7210c9ab0bf,
         0x009de23a326fd9bcd4fa993b0f4, 0x00aaf1de1730b3d7d6857669166],
    ),
    true => Dict(
        0 => [1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1,
              0, 1, 0, 0, 0, 1, 0, 1, 1, 0, 1, 1],
        1 => [1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0,
              1, 1, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1],
        Int64(2)^32 => [0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1,
                        1, 0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0],
        typemax(Int64) => [1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1,
                           0, 1, 0, 0, 0, 1, 0, 1, 1, 0, 1, 1],
    ),
    Float64 => Dict(
        0 =>
        [0.1950648807374728566, 0.0257793115230373626, 0.9055427466847816298,
         0.7734523504688743944, 0.8843643358740858051, 0.0788797171244850492,
         0.5698946304824119568, 0.1859290813601361946, 0.3502633405739192795,
         0.0789272944027068490, 0.3145803238578359639, 0.1485569608098213922,
         0.2545598088968246131, 0.0483571631942767421, 0.0070413465271252917,
         0.0508483589802035673, 0.8724974721884621864, 0.1605102180246753551,
         0.5840743840364566175, 0.4382718269681136202, 0.9528405584114081339,
         0.2329044092861747294, 0.5859153544674560176, 0.8064312206579657971],
        1 =>
        [0.5851946422124185698, 0.0773379345691123099, 0.7166282400543453334,
         0.3203570514066231833, 0.6530930076222578595, 0.2366391513734555918,
         0.7096838914472360926, 0.5577872440804085840, 0.0507900217217578386,
         0.2367818832081209912, 0.9437409715735081139, 0.4456708824294643989,
         0.7636794266904740613, 0.1450714895828306705, 0.0211240395813758752,
         0.1525450769406111462, 0.6174924165653870034, 0.4815306540740260654,
         0.7522231521093702966, 0.3148154809043408608, 0.8585216752342246238,
         0.6987132278585241884, 0.7577460634023684971, 0.4192936619738978354],
        Int64(2)^32 =>
        [0.9262377264309376024, 0.8354629112182574424, 0.1878431925440584571,
         0.3297440126430100094, 0.5402822540569656162, 0.8136704831208338984,
         0.7458235620861592974, 0.8201455119180884878, 0.8556499214817800958,
         0.5217083321266742235, 0.1835695964904275890, 0.6912309580542508325,
         0.4303038831150767241, 0.9418660572220058924, 0.1146544263022086785,
         0.8013346510146608370, 0.6567284985579240342, 0.3401312026411986711,
         0.3199271240192740340, 0.6107096843461274904, 0.3115926468885286127,
         0.1895431852488782543, 0.9098590120021339267, 0.2050781237987520100],
        typemax(Int64) =>
        [0.7932778610948740816, 0.3734837512097641809, 0.4003635844200839066,
         0.2642441753022539696, 0.8435718862417675989, 0.5996737527010649326,
         0.2592634627050232065, 0.7302011895477420466, 0.0679970382392385275,
         0.0777636071926277594, 0.9513238090110549549, 0.1286523286704006618,
         0.3620349797067561326, 0.2187559397428240348, 0.1208698967929713319,
         0.1413011948896074621, 0.5101784014171362180, 0.1411332235862896578,
         0.8650217525819665631, 0.5771692253821305573, 0.7376248001192464709,
         0.1897371938050098627, 0.8431170858467222917, 0.0142466392990996393]
    ),
    Normal{Float64} => Dict(
        0 =>
        [-0.381111861786743722, -0.020452508120591215, +0.610788422210481796,
         -1.385748932245142839, -1.611731486297559135, -0.110517958623616241,
         -0.503001669385000350, +0.265474785013935066, +0.547305897889466042,
         +0.108473227118974030, +0.718753376471714445, -0.282678969833291605,
         +0.736726902332106914, -0.134961030673573013, +0.006949743869390782,
         +0.077725160632003479, +0.510076650626809736, -0.426360569496689379,
         +0.908664368119342147, -0.730524759144574709, -1.232581867613880222,
         +0.398240840524731410, -1.139673658738002970, -0.821968503141728313],
        1 =>
        [-0.532520074864123116, +0.098465514284785005, +0.752886522124523450,
         -0.843537403874254577, -2.013829355631106655, -0.269807641197127135,
         +1.032293472303989157, +0.675566099119190633, +0.082160827543964254,
         +0.246926640359151073, -1.500529147318949708, +0.328632009212124654,
         -1.853198706543207619, -0.337190826546952760, +0.035662516379140363,
         +0.232143558890440915, +0.556345368618102021, -1.027821009268467423,
         +1.206403573149015029, -0.620592199441276415, +0.612308205879786516,
         +1.491394429348988471, -0.675799668973327217, -0.751223710193982930],
        Int64(2)^32 =>
        [+2.804497592250339011, +1.277061645230312958, +0.121477387690437907,
         +0.187076393406817809, +0.627670131300367640, -1.703321215824013413,
         -0.989982760767698200, +2.075029846471020178, -1.649703602248066669,
         -0.531759939161147721, +0.609495986040443971, -0.603631997036122026,
         -0.684174512146213942, +0.534356647282485442, -0.092174031696126188,
         -1.285167816972012433, +0.999411412925283571, -0.499305278581140211,
         -0.091555014119427833, -1.350741201226236088, +0.482633667214938111,
         -0.377021488481488309, +0.463503407362597097, +0.397140231988096948],
        typemax(Int64) =>
        [-1.529449500521834260, -0.425833043072113893, +0.751955071680120102,
         -0.295503514518007215, -0.493166268112920858, -0.488318336729492652,
         -0.221460042280709618, +1.749374475915582039, +0.209670644309441018,
         +0.166809441113434515, +1.422070714069686215, -0.201905699115775966,
         +0.792459370505811722, -0.300645840135806364, +0.220282309568965884,
         +0.117923668733925277, +0.786760475994732755, -0.098886664279309366,
         +1.136538822658988090, -1.629277960475042697, -1.097667207947351953,
         +0.316259702942887966, -1.604310346528319942, -0.027958930823246426],
    ),
    Exponential{Float64} => Dict(
        0 =>
        [+0.290443376756253313, +0.023107801871801649, +0.602445039808046201,
         +0.809080223295992051, +1.000988927033177101, +0.264416567316517359,
         +0.620552542494976045, +0.670911854378504646, +0.088847225724062964,
         +0.242968221532827683, +0.747559288415950429, +0.200756450113355078,
         +1.069994735701830679, +0.189708173540399205, +0.009436464636036371,
         +0.352940169680237725, +0.455343388287083484, +0.560636133492580124,
         +0.124215981573885054, +0.296580109076140352, +2.424871194870999336,
         +0.186103706973448774, +0.858731419298913834, +1.163246417195959559],
        1 =>
        [+0.674990107171162257, +0.186569786854466957, +1.092357077170384105,
         +1.095860857333019078, +3.225480781007138642, +0.435163300108867002,
         +2.820713900366846882, +1.177949909079809343, +0.025211559614314250,
         +0.355336719748536645, +0.378642730871126820, +0.346359637424302036,
         +2.138641751151649562, +0.363542010086324974, +0.015483132528467234,
         +0.937262071676046249, +0.692276643844622619, +0.952454488432565638,
         +0.333481484835768782, +0.483612444263148922, +0.629264342409806110,
         +1.372023120746223457, +0.841369485855095078, +0.438607768727813097],
        Int64(2)^32 =>
        [+4.301903604378438217, +5.798976162863313633, +0.116262999011904108,
         +0.163680985957347869, +1.030481683072774768, +1.518325010610733328,
         +2.054668731968426698, +2.524719717518081552, +1.214564758793303501,
         +0.752544460856171993, +1.040180627988793294, +0.737862246265513999,
         +0.172644234291837201, +0.467531051284808818, +0.105231573151698218,
         +0.366119257230970885, +4.035048035845405501, +1.451422316778351984,
         +0.060772937253374834, +1.337722289946995380, +0.051452552989185357,
         +0.300272614348955136, +0.378059311825467925, +0.293455165778313376],
        typemax(Int64) =>
        [+1.126029828119699872, +0.686811209261761424, +0.509836696167542347,
         +0.464964296843861813, +0.450819900116565198, +0.563226973228172834,
         +0.265664126146503754, +1.963481563365191640, +0.328625588378420153,
         +0.154936598094436634, +4.597690643152912315, +0.039603426460420012,
         +0.764555080732192693, +0.679541058538978282, +0.135508395324274611,
         +0.137243906430860846, +0.053486342196888477, +0.101871384862385714,
         +2.289254085681110773, +2.324958005392129046, +3.492832668884421121,
         +0.126229525545209859, +1.139368981849247842, +0.021547655127562688],
    ),
    Float32 => Dict(
        0 =>
        [0.420698285102844f0, 0.488105177879333f0, 0.267643690109253f0,
         0.784768223762512f0, 0.540994763374329f0, 0.670924305915833f0,
         0.010995507240295f0, 0.477138519287109f0, 0.094086647033691f0,
         0.527673721313477f0, 0.366811752319336f0, 0.033917069435120f0,
         0.760983943939209f0, 0.305844187736511f0, 0.131725788116455f0,
         0.859405279159546f0, 0.611514329910278f0, 0.136226296424866f0,
         0.233490705490112f0, 0.448277354240417f0, 0.584921956062317f0,
         0.622289896011353f0, 0.707746386528015f0, 0.899915337562561f0],
        1 =>
        [0.262094855308533f0, 0.464315652847290f0, 0.802931308746338f0,
         0.354304671287537f0, 0.622984528541565f0, 0.012773156166077f0,
         0.032986640930176f0, 0.431415557861328f0, 0.282259941101074f0,
         0.583021402359009f0, 0.100435376167297f0, 0.101751327514648f0,
         0.282951951026917f0, 0.917532801628113f0, 0.395177364349365f0,
         0.578216075897217f0, 0.834543228149414f0, 0.408678889274597f0,
         0.700472354888916f0, 0.344832062721252f0, 0.754765987396240f0,
         0.866869688034058f0, 0.123239398002625f0, 0.699746251106262f0],
        Int64(2)^32 =>
        [0.917783975601196f0, 0.087920904159546f0, 0.094120264053345f0,
         0.794192314147949f0, 0.472978830337524f0, 0.340562701225281f0,
         0.968284964561462f0, 0.956171035766602f0, 0.698651671409607f0,
         0.066846489906311f0, 0.683287858963013f0, 0.853219389915466f0,
         0.165132641792297f0, 0.122622489929199f0, 0.413703560829163f0,
         0.907442688941956f0, 0.957183361053467f0, 0.961637139320374f0,
         0.845764756202698f0, 0.202137589454651f0, 0.507538318634033f0,
         0.727950215339661f0, 0.565004587173462f0, 0.355084896087646f0],
        typemax(Int64) =>
        [0.755414366722107f0, 0.129167199134827f0, 0.699199438095093f0,
         0.385208964347839f0, 0.904178023338318f0, 0.515083193778992f0,
         0.670723795890808f0, 0.575981140136719f0, 0.932798862457275f0,
         0.713915824890137f0, 0.951078891754150f0, 0.024201750755310f0,
         0.731067657470703f0, 0.875146985054016f0, 0.724588394165039f0,
         0.367073297500610f0, 0.651520013809204f0, 0.460271239280701f0,
         0.208518743515015f0, 0.409237980842590f0, 0.153837561607361f0,
         0.278416395187378f0, 0.801312088966370f0, 0.233442664146423f0],
    ),
    Normal{Float32} => Dict(
        0 =>
        [-0.38111186027527f0, -0.02045250870287f0, +0.61078840494156f0,
         -1.38574898242950f0, -1.61173152923584f0, -0.11051795631647f0,
         -0.50300168991089f0, +0.26547479629517f0, +0.54730588197708f0,
         +0.10847322642803f0, +0.71875339746475f0, -0.28267896175385f0,
         +0.73672688007355f0, -0.13496102392673f0, +0.00694974372163f0,
         +0.07772515714169f0, +0.51007664203644f0, -0.42636057734489f0,
         +0.90866434574127f0, -0.73052477836609f0, -1.23258185386658f0,
         +0.39824083447456f0, -1.13967370986938f0, -0.82196849584579f0],
        1 =>
        [-0.53252005577087f0, +0.09846551716328f0, +0.75288653373718f0,
         -0.84353739023209f0, -2.01382946968079f0, -0.26980763673782f0,
         +1.03229343891144f0, +0.67556607723236f0, +0.08216083049774f0,
         +0.24692663550377f0, -1.50052917003632f0, +0.32863199710846f0,
         -1.85319864749908f0, -0.33719083666801f0, +0.03566251695156f0,
         +0.23214356601238f0, +0.55634534358978f0, -1.02782106399536f0,
         +1.20640361309052f0, -0.62059217691422f0, +0.61230820417404f0,
         +1.49139440059662f0, -0.67579966783524f0, -0.75122368335724f0],
        Int64(2)^32 =>
        [+2.80449748039246f0, +1.27706170082092f0, +0.12147738784552f0,
         +0.18707638978958f0, +0.62767010927200f0, -1.70332121849060f0,
         -0.98998278379440f0, +2.07502985000610f0, -1.64970362186432f0,
         -0.53175991773605f0, +0.60949599742889f0, -0.60363197326660f0,
         -0.68417453765869f0, +0.53435665369034f0, -0.09217403084040f0,
         -1.28516781330109f0, +0.99941140413284f0, -0.49930527806282f0,
         -0.09155501425266f0, -1.35074114799500f0, +0.48263368010521f0,
         -0.37702149152756f0, +0.46350342035294f0, +0.39714023470879f0],
        typemax(Int64) =>
        [-1.52944946289063f0, -0.42583304643631f0, +0.75195509195328f0,
         -0.29550352692604f0, -0.49316626787186f0, -0.48831832408905f0,
         -0.22146004438400f0, +1.74937450885773f0, +0.20967064797878f0,
         +0.16680943965912f0, +1.42207074165344f0, -0.20190569758415f0,
         +0.79245936870575f0, -0.30064582824707f0, +0.22028231620789f0,
         +0.11792366951704f0, +0.78676044940948f0, -0.09888666123152f0,
         +1.13653886318207f0, -1.62927794456482f0, -1.09766721725464f0,
         +0.31625971198082f0, -1.60431039333344f0, -0.02795893140137f0],
    ),
    Exponential{Float32} => Dict(
        0 =>
        [+0.29044339060783f0, +0.02310780249536f0, +0.60244506597519f0,
         +0.80908024311066f0, +1.00098896026611f0, +0.26441657543182f0,
         +0.62055253982544f0, +0.67091184854507f0, +0.08884722739458f0,
         +0.24296821653843f0, +0.74755930900574f0, +0.20075644552708f0,
         +1.06999468803406f0, +0.18970817327499f0, +0.00943646486849f0,
         +0.35294017195702f0, +0.45534339547157f0, +0.56063616275787f0,
         +0.12421598285437f0, +0.29658010601997f0, +2.42487120628357f0,
         +0.18610370159149f0, +0.85873144865036f0, +1.16324639320374f0],
        1 =>
        [+0.67499011754990f0, +0.18656978011131f0, +1.09235703945160f0,
         +1.09586083889008f0, +3.22548079490662f0, +0.43516328930855f0,
         +2.82071399688721f0, +1.17794990539551f0, +0.02521155960858f0,
         +0.35533672571182f0, +0.37864273786545f0, +0.34635964035988f0,
         +2.13864183425903f0, +0.36354202032089f0, +0.01548313256353f0,
         +0.93726205825806f0, +0.69227665662766f0, +0.95245450735092f0,
         +0.33348149061203f0, +0.48361244797707f0, +0.62926435470581f0,
         +1.37202310562134f0, +0.84136950969696f0, +0.43860778212547f0],
        Int64(2)^32 =>
        [+4.30190372467041f0, +5.79897594451904f0, +0.11626300215721f0,
         +0.16368098556995f0, +1.03048169612885f0, +1.51832497119904f0,
         +2.05466866493225f0, +2.52471971511841f0, +1.21456480026245f0,
         +0.75254446268082f0, +1.04018068313599f0, +0.73786222934723f0,
         +0.17264422774315f0, +0.46753105521202f0, +0.10523157566786f0,
         +0.36611926555634f0, +4.03504800796509f0, +1.45142233371735f0,
         +0.06077293679118f0, +1.33772230148315f0, +0.05145255476236f0,
         +0.30027261376381f0, +0.37805929780006f0, +0.29345515370369f0],
        typemax(Int64) =>
        [+1.12602984905243f0, +0.68681120872498f0, +0.50983667373657f0,
         +0.46496430039406f0, +0.45081990957260f0, +0.56322699785233f0,
         +0.26566413044930f0, +1.96348154544830f0, +0.32862558960915f0,
         +0.15493659675121f0, +4.59769058227539f0, +0.03960342705250f0,
         +0.76455509662628f0, +0.67954105138779f0, +0.13550838828087f0,
         +0.13724391162395f0, +0.05348634347320f0, +0.10187138617039f0,
         +2.28925418853760f0, +2.32495808601379f0, +3.49283266067505f0,
         +0.12622952461243f0, +1.13936901092529f0, +0.02154765464365f0],
    ),
    Float16 => Dict(
        0 =>
        Float16[0.360351563, 0.557617188, 0.537109375,
                0.821289063, 0.829101563, 0.211914063,
                0.075195313, 0.718750000, 0.757812500,
                0.703125000, 0.921875000, 0.848632813,
                0.980468750, 0.475585938, 0.097656250,
                0.248046875, 0.525390625, 0.965820313,
                0.755859375, 0.288085938, 0.680664063,
                0.798828125, 0.858398438, 0.106445313],
        1 =>
        Float16[0.081054688, 0.673828125, 0.613281250,
                0.463867188, 0.489257813, 0.637695313,
                0.226562500, 0.156250000, 0.273437500,
                0.111328125, 0.766601563, 0.546875000,
                0.942382813, 0.428710938, 0.292968750,
                0.746093750, 0.578125000, 0.897460938,
                0.269531250, 0.864257813, 0.042968750,
                0.396484375, 0.577148438, 0.321289063],
        Int64(2)^32 =>
        Float16[0.486328125, 0.248046875, 0.033203125,
                0.023437500, 0.642578125, 0.889648438,
                0.190429688, 0.953125000, 0.354492188,
                0.606445313, 0.494140625, 0.573242188,
                0.766601563, 0.523437500, 0.059570313,
                0.770507813, 0.246093750, 0.731445313,
                0.504882813, 0.911132813, 0.753906250,
                0.368164063, 0.517578125, 0.855468750],
        typemax(Int64) =>
        Float16[0.354492188, 0.137695313, 0.841796875,
                0.631835938, 0.026367188, 0.561523438,
                0.569335938, 0.437500000, 0.488281250,
                0.398437500, 0.238281250, 0.260742188,
                0.906250000, 0.204101563, 0.828125000,
                0.064453125, 0.251953125, 0.541992188,
                0.185546875, 0.477539063, 0.237304688,
                0.787109375, 0.348632813, 0.362304688],
    ),
    Normal{Float16} => Dict(
        0 =>
        Float16[-0.381103516, -0.020446777, +0.610839844,
                -1.385742188, -1.611328125, -0.110534668,
                -0.502929688, +0.265380859, +0.547363281,
                +0.108459473, +0.718750000, -0.282714844,
                +0.736816406, -0.135009766, +0.006950378,
                +0.077697754, +0.510253906, -0.426269531,
                +0.908691406, -0.730468750, -1.232421875,
                +0.398193359, -1.139648438, -0.821777344],
        1 =>
        Float16[-0.532714844, +0.098449707, +0.752929688,
                -0.843750000, -2.013671875, -0.269775391,
                +1.032226563, +0.675781250, +0.082153320,
                +0.246948242, -1.500976563, +0.328613281,
                -1.853515625, -0.337158203, +0.035675049,
                +0.232177734, +0.556152344, -1.027343750,
                +1.206054688, -0.620605469, +0.612304688,
                +1.491210938, -0.675781250, -0.751464844],
        Int64(2)^32 =>
        Float16[+2.804687500, +1.277343750, +0.121459961,
                +0.187133789, +0.627441406, -1.703125000,
                -0.989746094, +2.074218750, -1.649414063,
                -0.531738281, +0.609375000, -0.603515625,
                -0.684082031, +0.534179688, -0.092163086,
                -1.285156250, +0.999511719, -0.499267578,
                -0.091552734, -1.350585938, +0.482666016,
                -0.376953125, +0.463623047, +0.397216797],
        typemax(Int64) =>
        Float16[-1.529296875, -0.425781250, +0.751953125,
                -0.295410156, -0.493164063, -0.488281250,
                -0.221435547, +1.749023438, +0.209716797,
                +0.166870117, +1.421875000, -0.201904297,
                +0.792480469, -0.300537109, +0.220336914,
                +0.117919922, +0.786621094, -0.098876953,
                +1.136718750, -1.628906250, -1.097656250,
                +0.316162109, -1.604492188, -0.027954102],
    ),
    Exponential{Float16} => Dict(
        0 =>
        Float16[+0.290527344, +0.023101807, +0.602539063,
                +0.809082031, +1.000976563, +0.264404297,
                +0.620605469, +0.670898438, +0.088867188,
                +0.242919922, +0.747558594, +0.200805664,
                +1.070312500, +0.189697266, +0.009437561,
                +0.353027344, +0.455322266, +0.560546875,
                +0.124206543, +0.296630859, +2.425781250,
                +0.186157227, +0.858886719, +1.163085938],
        1 =>
        Float16[+0.674804688, +0.186523438, +1.092773438,
                +1.095703125, +3.224609375, +0.435058594,
                +2.820312500, +1.177734375, +0.025207520,
                +0.355224609, +0.378662109, +0.346435547,
                +2.138671875, +0.363525391, +0.015480042,
                +0.937500000, +0.692382813, +0.952636719,
                +0.333496094, +0.483642578, +0.629394531,
                +1.372070313, +0.841308594, +0.438720703],
        Int64(2)^32 =>
        Float16[+4.300781250, +5.800781250, +0.116271973,
                +0.163696289, +1.030273438, +1.518554688,
                +2.054687500, +2.525390625, +1.214843750,
                +0.752441406, +1.040039063, +0.737792969,
                +0.172607422, +0.467529297, +0.105224609,
                +0.366210938, +4.035156250, +1.451171875,
                +0.060760498, +1.337890625, +0.051452637,
                +0.300292969, +0.378173828, +0.293457031],
        typemax(Int64) =>
        Float16[+1.125976563, +0.687011719, +0.509765625,
                +0.464843750, +0.450927734, +0.562988281,
                +0.265625000, +1.963867188, +0.328613281,
                +0.154907227, +4.597656250, +0.039611816,
                +0.764648438, +0.679687500, +0.135498047,
                +0.137207031, +0.053497314, +0.101867676,
                +2.289062500, +2.324218750, +3.492187500,
                +0.126220703, +1.139648438, +0.021545410],
    )
)


# # [Rising Thermal Bubble](@id EX-RTB-docs)
#
# In this example, we demonstrate the usage of the `ClimateMachine`
# [AtmosModel](@ref AtmosModel-docs) machinery to solve the fluid
# dynamics of a thermal perturbation in a neutrally stratified background state
# defined by its uniform potential temperature. We solve a flow in a
# [`FlatOrientation`](@ref LESConfig) (Box) configuration - this is
# representative of a large-eddy simulation. Several versions of the problem
# setup may be found in literature, but the general idea is to examine the
# vertical ascent of a thermal _bubble_ (we can interpret these as simple
# representation of convective updrafts).
#
# ## Description of experiment
# 1) Dry Rising Bubble (circular potential temperature perturbation)
# 2) Boundaries
#    - `Impenetrable(FreeSlip())` - no momentum flux, no mass flux through
#      walls.
#    - `Impermeable()` - non-porous walls, i.e. no diffusive fluxes through
#       walls.
#    - Laterally periodic
# 3) Domain - 2500m (horizontal) x 2500m (horizontal) x 2500m (vertical)
# 4) Resolution - 50m effective resolution
# 5) Total simulation time - 1000s
# 6) Mesh Aspect Ratio (Effective resolution) 1:1
# 7) Overrides defaults for
#    - CPU Initialisation
#    - Time integrator
#    - Sources
#    - Smagorinsky Coefficient
# 8) Default settings can be found in `src/Driver/<files>.jl`

#md # !!! note
#md #     This experiment setup assumes that you have installed the
#md #     `ClimateMachine` according to the instructions on the landing page.
#md #     We assume the users' familiarity with the conservative form of the
#md #     equations of motion for a compressible fluid (see the
#md #     [`AtmosModel`](@ref AtmosModel-docs) page).
#md #
#md #     The following topics are covered in this example
#md #     - Package requirements
#md #     - Defining a `model` subtype for the set of conservation equations
#md #     - Defining the initial conditions
#md #     - Applying boundary conditions
#md #     - Applying source terms
#md #     - Choosing a turbulence model
#md #     - Adding tracers to the model
#md #     - Choosing a time-integrator
#md #     - Choosing diagnostics (output) configurations
#md #
#md #     The following topics are not covered in this example
#md #     - Defining new boundary conditions
#md #     - Defining new turbulence models
#md #     - Building new time-integrators
#md #     - Adding diagnostic variables (beyond a standard pre-defined list of
#md #       variables)
#
# ## Boilerplate (Using Modules)
#
# #### [Skip Section](@ref init)
#
# Before setting up our experiment, we recognize that we need to import some
# pre-defined functions from other packages. Julia allows us to use existing
# modules (variable workspaces), or write our own to do so.  Complete
# documentation for the Julia module system can be found
# [here](https://docs.julialang.org/en/v1/manual/modules/#).

# We need to use the `ClimateMachine` module! This imports all functions
# specific to atmospheric and ocean flow modelling.  While we do not cover the
# ins-and-outs of the contents of each of these we provide brief descriptions
# of the utility of each of the loaded packages.
using ClimateMachine
ClimateMachine.init()

using ClimateMachine.Atmos
# - Required so that we inherit the appropriate model types for the large-eddy
#   simulation (LES) and global-circulation-model (GCM) configurations.
using ClimateMachine.ConfigTypes
# - Required so that we may define diagnostics configurations, e.g. choice of
#   file-writer, choice of output variable sets, output-frequency and directory,
using ClimateMachine.Diagnostics
# - Required so that we may define (or utilise existing functions) functions
#   that are `called-back` or executed at frequencies of either timesteps,
#   simulation-time, or wall-clock time.
using ClimateMachine.GenericCallbacks
# - Required so we load the appropriate functions for the time-integration
#   component. Contains ODESolver methods.
using ClimateMachine.ODESolvers
# - Required for utility of spatial filtering functions (e.g. positivity
#   preservation)
using ClimateMachine.Mesh.Filters
# - Required so functions for computation of temperature profiles.
using ClimateMachine.TemperatureProfiles
# - Required so functions for computation of moist thermodynamic quantities is
#   enabled.
using ClimateMachine.MoistThermodynamics
# - Required so we may access our variable arrays by a sensible naming
#   convention rather than by numerical array indices.
using ClimateMachine.VariableTemplates
# - Required so we may access planet parameters
#   ([CLIMAParameters](https://CliMA.github.io/CLIMAParameters.jl/latest/)
#   specific to this problem include the gas constant, specific heats,
#   mean-sea-level pressure, gravity and the Smagorinsky coefficient)

# In ClimateMachine we use `StaticArrays` for our variable arrays.
using StaticArrays
# We also use the `Test` package to help with unit tests and continuous
# integration systems to design sensible tests for our experiment to ensure new
# / modified blocks of code don't damage the fidelity of the physics. The test
# defined within this experiment is not a unit test for a specific
# subcomponent, but ensures time-integration of the defined problem conditions
# within a reasonable tolerance. Immediately useful macros and functions from
# this include `@test` and `@testset` which will allow us to define the testing
# parameter sets.
using Test

using CLIMAParameters
using CLIMAParameters.Atmos.SubgridScale: C_smag
using CLIMAParameters.Planet: R_d, cp_d, cv_d, MSLP, grav
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

# ## [Initial Conditions](@id init)
# This example of a rising thermal bubble can be classified as an initial value
# problem. We must (at the very least) assign values for the initial variables
# in a sensible manner. This example demonstrates the use of functions defined
# in the [`MoistThermodynamics`](@ref MoistThermodynamics-docs) package to
# generate the appropriate initial state for our problem.

#md # !!! note
#md #     The following variables are assigned in the initial condition
#md #     - `state.ρ` = Scalar quantity for initial density profile
#md #     - `state.ρu`= 3-component vector for initial momentum profile
#md #     - `state.ρe`= Scalar quantity for initial total-energy profile
#md #       humidity
#md #     - `state.tracers.ρχ` = Vector of four tracers (here, for demonstration
#md #       only; we can interpret these as dye injections for visualisation
#md #       purposes)
function init_risingbubble!(bl, state, aux, (x, y, z), t)
    # Problem float-type
    FT = eltype(state)

    # Unpack constant parameters
    R_gas::FT = R_d(bl.param_set)
    c_p::FT = cp_d(bl.param_set)
    c_v::FT = cv_d(bl.param_set)
    p0::FT = MSLP(bl.param_set)
    _grav::FT = grav(bl.param_set)
    γ::FT = c_p / c_v

    # Define bubble center and background potential temperature
    xc::FT = 1250
    yc::FT = 1250
    zc::FT = 1000
    r = sqrt((x - xc)^2 + (y - yc)^2 + (z - zc)^2)
    rc::FT = 500
    # TODO: clean this up, or add convenience function:
    # This is configured in the reference hydrostatic state
    θ_ref::FT = bl.ref_state.virtual_temperature_profile.T_surface
    Δθ::FT = 0

    # Compute temperature difference over bubble region
    if r <= rc
        Δθ = FT(5) * cospi(r / rc / 2)
    end

    # Compute perturbed thermodynamic state:
    θ = θ_ref + Δθ                                      # potential temperature
    π_exner = FT(1) - _grav / (c_p * θ) * z             # exner pressure
    ρ = p0 / (R_gas * θ) * (π_exner)^(c_v / R_gas)      # density
    T = θ * π_exner
    e_int = internal_energy(bl.param_set, T)
    ts = PhaseDry(bl.param_set, e_int, ρ)
    ρu = SVector(FT(0), FT(0), FT(0))                   # momentum
    #State (prognostic) variable assignment
    e_kin = FT(0)                                       # kinetic energy
    e_pot = gravitational_potential(bl.orientation, aux)# potential energy
    ρe_tot = ρ * total_energy(e_kin, e_pot, ts)         # total energy

    ρχ = FT(0)                                          # tracer

    # We inject tracers at the initial condition at some specified z coordinates
    if 500 < z <= 550
        ρχ += FT(0.05)
    end

    # We want 4 tracers
    ntracers = 4

    # Define 4 tracers, (arbitrary scaling for this demo problem)
    ρχ = SVector{ntracers, FT}(ρχ, ρχ / 2, ρχ / 3, ρχ / 4)

    # Assign State Variables
    state.ρ = ρ
    state.ρu = ρu
    state.ρe = ρe_tot
    state.tracers.ρχ = ρχ
end

# ## [Model Configuration](@id config-helper)
# We define a configuration function to assist in prescribing the physical
# model. The purpose of this is to populate the
# [`ClimateMachine.AtmosLESConfiguration`](@ref LESConfig) with arguments
# appropriate to the problem being considered.
function config_risingbubble(FT, N, resolution, xmax, ymax, zmax)

    # Choose an Explicit Multi-rate Solver from the existing [ODESolvers](@ref
    # ODESolvers-docs) options Apply the outer constructor to define the
    # `ode_solver`. Here `AtmosAcousticGravityLinearModel` splits the
    # acoustic-gravity wave components from the advection-diffusion dynamics.
    # The 1D-IMEX method is less appropriate for the problem given the current
    # mesh aspect ratio (1:1)
    ode_solver = ClimateMachine.MultirateSolverType(
        linear_model = AtmosAcousticGravityLinearModel,
        slow_method = LSRK144NiegemannDiehlBusch,
        fast_method = LSRK144NiegemannDiehlBusch,
        timestep_ratio = 10,
    )

    # Since we want four tracers, we specify this and include the appropriate
    # diffusivity scaling coefficients (normally these would be physically
    # informed but for this demonstration we use integers corresponding to the
    # tracer index identifier)
    ntracers = 4
    δ_χ = SVector{ntracers, FT}(1, 2, 3, 4)

    # The model coefficient for the turbulence closure is defined via the
    # [CLIMAParameters
    # package](https://CliMA.github.io/CLIMAParameters.jl/latest/) A reference
    # state for the linearisation step is also defined.
    T_surface = FT(300)
    T_min_ref = FT(0)
    T_profile = DryAdiabaticProfile{FT}(param_set, T_surface, T_min_ref)
    ref_state = HydrostaticState(T_profile)

    # The fun part! Here we assemble the `AtmosModel`.
    #md # !!! note
    #md #     Docs on model subcomponent options can be found here:
    #md #     - [`param_set`](https://CliMA.github.io/CLIMAParameters.jl/latest/)
    #md #     - [`turbulence`](@ref Turbulence-Closures-docs)
    #md #     - [`hyperdiffusion`](@ref Hyperdiffusion-docs)
    #md #     - [`source`](@ref atmos-sources)
    #md #     - [`tracers`](@ref Tracers-docs)
    #md #     - [`init_state`](@ref init)

    _C_smag = FT(C_smag(param_set))
    model = AtmosModel{FT}(
        AtmosLESConfigType,                           # Flow in a box, requires the AtmosLESConfigType
        param_set;                                    # Parameter set corresponding to earth parameters
        turbulence = SmagorinskyLilly(_C_smag),       # Turbulence closure model
        moisture = DryModel(),                        # Exclude moisture variables
        hyperdiffusion = StandardHyperDiffusion(60),  # Hyperdiffusion (4th order) model
        source = (Gravity(),),                        # Gravity is the only source term here
        tracers = NTracers{ntracers, FT}(δ_χ),        # Tracer model with diffusivity coefficients
        ref_state = ref_state,                        # Reference state
        init_state_conservative = init_risingbubble!, # Apply the initial condition
    )

    # Finally, we pass a `Problem Name` string, the mesh information, and the
    # model type to  the [`AtmosLESConfiguration`] object.
    config = ClimateMachine.AtmosLESConfiguration(
        "DryRisingBubble",       # Problem title [String]
        N,                       # Polynomial order [Int]
        resolution,              # (Δx, Δy, Δz) effective resolution [m]
        xmax,                    # Domain maximum size [m]
        ymax,                    # Domain maximum size [m]
        zmax,                    # Domain maximum size [m]
        param_set,               # Parameter set.
        init_risingbubble!,      # Function specifying initial condition
        solver_type = ode_solver,# Time-integrator type
        model = model,           # Model type
    )
    return config
end

#md # !!! note
#md #     `Keywords` are used to specify some arguments (see appropriate source
#md #     files).

# ## [Diagnostics](@id config_diagnostics)
# Here we define the diagnostic configuration specific to this problem.
function config_diagnostics(driver_config)
    interval = "10000steps"
    dgngrp = setup_atmos_default_diagnostics(interval, driver_config.name)
    return ClimateMachine.DiagnosticsConfiguration([dgngrp])
end

function main()
    # These are essentially arguments passed to the
    # [`config_risingbubble`](@ref config-helper) function.  For type
    # consistency we explicitly define the problem floating-precision.
    FT = Float64
    # We need to specify the polynomial order for the DG discretization,
    # effective resolution, simulation end-time, the domain bounds, and the
    # courant-number for the time-integrator. Note how the time-integration
    # components `solver_config` are distinct from the spatial / model
    # components in `driver_config`. `init_on_cpu` is a helper keyword argument
    # that forces problem initialisation on CPU (thereby allowing the use of
    # random seeds, spline interpolants and other special functions at the
    # initialisation step.)
    N = 4
    Δh = FT(50)
    Δv = FT(50)
    resolution = (Δh, Δh, Δv)
    xmax = FT(2500)
    ymax = FT(2500)
    zmax = FT(2500)
    t0 = FT(0)
    timeend = FT(1000)
    CFL = FT(20)

    # Assign configurations so they can be passed to the `invoke!` function
    driver_config = config_risingbubble(FT, N, resolution, xmax, ymax, zmax)
    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        init_on_cpu = true,
        Courant_number = CFL,
    )
    dgn_config = config_diagnostics(driver_config)

    # User defined filter (TMAR positivity preserving filter)
    cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do (init = false)
        Filters.apply!(solver_config.Q, 6, solver_config.dg.grid, TMARFilter())
        nothing
    end

    # Invoke solver (calls `solve!` function for time-integrator), pass the driver, solver and diagnostic config
    # information.
    result = ClimateMachine.invoke!(
        solver_config;
        diagnostics_config = dgn_config,
        user_callbacks = (cbtmarfilter,),
        check_euclidean_distance = true,
    )

    # Check that the solution norm is reasonable.
    @test isapprox(result, FT(1); atol = 1.5e-3)
end

# The experiment definition is now complete. Time to run it.

# ## Running the Experiment
# `julia --project /experiments/AtmosLES/risingbubble.jl` will run the
# experiment from the main ClimateMachine.jl directory, with diagnostics output
# at the intervals specified in [`config_diagnostics`](@ref
# config_diagnostics).  You can also prescribe command line arguments (docs
# pending, `Driver.jl`) for simulation update and output specifications.  For
# rapid turnaround, we recommend that you run this experiment on a GPU.

# ## [Output Visualisation](@id output-viz)
#md # See the `ClimateMachine` [command line arguments](@ref ClimateMachine-args)
#md # for generating output.
#md #
#md # Given VTK output,
#md # - [VisIt](https://wci.llnl.gov/simulation/computer-codes/visit/)
#md # - [Paraview](https://wci.llnl.gov/simulation/computer-codes/visit/)
#md # are two commonly used programs for `.vtu` files.
#md #
#md # For NetCDF or JLD2 diagnostics you may use Julia's
#md # [`NCDatasets`](https://github.com/Alexander-Barth/NCDatasets.jl) and
#md # [`JLD2`](https://github.com/JuliaIO/JLD2.jl) packages with a suitable
#md # plotting program.

#

main()


using Makie
import DataStructures: DefaultDict

export Port

@recipe(Steps, t, y) do scene
    Theme(
        color = :black,
        fill = nothing,
    )
end

function Makie.plot!(myplot::Steps)
    t = myplot[1]
    y = myplot[2]

    line = Node(Point2f0[])
    function update!(t, y)
        empty!(line[])
        y_old = y[1]
        for (_t,_y) in zip(t,y)
            push!(line[], Point2f0(_t,y_old))
            push!(line[], Point2f0(_t,_y))
            y_old = _y
        end
    end

    Makie.Observables.onany(update!, t, y)
    update!(t[],y[])

    poly!(myplot, line, color=myplot[:fill])
    lines!(myplot, line, color=myplot[:color])
    myplot
end

@recipe(StateTrace, t, state, color, offset, height) do scene
    Theme(
    )
end

function Makie.plot!(plt::StateTrace)
    t,state,color,offset,height=plt[:t],plt[:state],plt[:color],plt.offset,plt.height
    rects = Node(Rect2D{Float64}[])
    colors = Node(Any[])

    Makie.Observables.onany(t, state, color, offset, height) do t,state,color, offset, height
        t₀,s₀ = first(t),first(state)
        for (t₁,s₁) in zip(t[2:end],state[2:end])
            push!(rects[], Rect(float(t₀), float(offset), float(t₁-t₀), float(height)))
            push!(colors[], color[s₀])
            t₀,s₀ = t₁,s₁
        end
    end

    # Force update
    t[] = t[]

    poly!(plt, rects, color=colors)

    plt
end

struct Port{W,I}
    name::I
end
Port{W}(name::I) where {W,I} = Port{W,I}(name)

function Makie.default_theme(scene::Makie.SceneLike, ::Type{<: Makie.Plot(Neuron)})
    Theme(
        color=RGBAf0(0.9,0.9,0.9,1.0),
        branch_width=0.1,
        branch_length=1.0,
        angle_between=10/180*pi,
        xticks = [],
        root_position = Point2f0(0,0),
        ports = Dict{Symbol,Vector{Symbol}}(),
        fxaa = false,
    )
end

function Makie.plot!(treeplot::Makie.Plot(Neuron))
    # get the actual object to plot
    tree = to_value(treeplot[1])
    # get position of the root = offset of all nodes
    offset = treeplot[:root_position]

    # if obj is an observable dict-type holding named values, get named value ...
    maybe_get(obj::Node, key) = if isa(obj[],Union{DefaultDict,Dict})
        @lift $obj[key]
    else
        obj
    end

    # if obj is a dict-type holding named values, get named value ...
    maybe_get(obj::Union{DefaultDict,Dict}, key) = obj[key]
    # ... otherwise return the obj itself
    maybe_get(obj, key) = obj

    function serialize_tree(tree::NeuronOrSegment{ID,T,WT,IT}, l, w, ω, given_ports) where {ID,T,WT,IT}
        sectors = Float64[]
        depths = Float64[]
        all_points = Vector{Pair{ID,Point2f0}}[]
        all_ports = Vector{Pair{ID,Point2f0}}[]
        all_points_flat = Pair{ID,Point2f0}[]
        all_ports_flat = Pair{ID,Point2f0}[]
        all_parents_flat = Pair{ID,ID}[]

        # go through all branches once to collect information
        for subtree in tree.next_upstream
            # get angle and depth of sector
            α,d,points,ports,parents = serialize_tree(subtree, l, w, ω, given_ports)
            # compute depth of new sector
            ll = maybe_get(l, subtree.id)
            c = √(ll^2+d^2-2*ll*d*cos(π-α/2))
            # compute angle of new sector
            β=2acos((ll^2+c^2-d^2)/(2ll*c))

            push!(sectors, β+ maybe_get(ω, subtree.id))
            push!(depths, c)
            push!(all_points, points)
            push!(all_ports, ports)
            append!(all_parents_flat, parents)
            push!(all_parents_flat, subtree.id=>tree.id)
        end

        # calculate total angle and depth for current (sub-)tree
        sector = sum(sectors)
        depth = isempty(depths) ? 0.0 : maximum(depths)

        # calculate angles for all branches
        branch_angles = cumsum(sectors) .- sectors./2 .- sector/2

        for (branch, branch_angle, points, ports) in zip(tree.next_upstream, branch_angles, all_points, all_ports)
            # shift by ll and rotate by `branch_angle`
            ll = maybe_get(l, branch.id)
            c=cos(branch_angle)
            s=sin(branch_angle)
            transform = ((x,y),) -> Point2f0(c*x-s*(y+ll),s*x+c*(y+ll))

            branch_end = transform(Point2f0(0.0,0.0))
            # add branch for branch
            push!(all_points_flat, branch.id=>branch_end)

            # add ports for branch
            branch_ports = get(given_ports, branch.id, ID[])
            append!(all_ports_flat, Pair.(branch_ports, LinRange(Point2f0(0,0), branch_end, length(branch_ports)+2)[2:end-1]))

            # keep branches from branches' subtree
            if !isempty(points)
                append!(all_points_flat, map(((name,pts),) -> name=>transform(pts), points))
            end

            # keep ports from branches' subtree
            if !isempty(ports)
                append!(all_ports_flat, map(((prt,pos),) -> prt=>transform(pos), ports))
            end
        end

        if isa(tree, Neuron)
            neuron_ports = get(given_ports, tree.id, ID[])

            ww = maybe_get(w, tree.id)
            append!(all_ports_flat, Pair.(neuron_ports, LinRange(Point2f0(0,-ww), Point2f0(0,ww), length(neuron_ports)+2)[2:end-1]))
        end

        return (sector=sector, depth=depth, nodes=all_points_flat, ports=all_ports_flat, parents=all_parents_flat)
    end

    # get all nodes, their ports and parents in flat format
    serialized = lift(treeplot[:branch_length],treeplot[:branch_width],treeplot[:angle_between], treeplot[:ports]) do l,w,ω,prts
        ser=serialize_tree(tree,l,w,ω,prts)
        push!(ser.nodes, tree.id => Point2f0(0.0,0.0))
        ser
    end
    parent = DefaultDict(tree.id, serialized[].parents...)
    # get all the positions of the ports
    ports = lift(serialized,offset) do ser,off
        [k=>v+off for (k,v) in ser.ports]
    end

    # plot branches
    for (name,parent_name) in serialized[].parents
        c  = maybe_get(treeplot[:color], name)
        w1 = maybe_get(treeplot[:branch_width], name)
        w2 = maybe_get(treeplot[:branch_width], parent_name)

        # dynamically recompute polygon
        branch_poly = lift(w1, w2, serialized, offset) do w1,w2,ser,off
            node_dict = Dict(ser.nodes...)
            b1 = node_dict[name]
            b2 = node_dict[parent_name]
            normal = [0 1; -1 0]*(b2-b1)
            normal /= sqrt(normal'*normal)*2
            Ref(off) .+ Point2f0[
                -normal*w1+b1,
                 normal*w1+b1,
                 normal*w2+b2,
                -normal*w2+b2
            ]
        end

        # draw polygon for the branch
        poly!(treeplot, branch_poly, color=c, strokewidth = 1, strokecolor=c)

        # draw circles to cap off branch
        poly!(treeplot, lift(w1,serialized, offset) do w1,ser,off
            node_dict = Dict(ser.nodes...)
            b2 = node_dict[name]
            decompose(Point2f0, Circle(off + b2, w1/2))
        end, color=c, strokewidth = 1, strokecolor=c)

        poly!(treeplot, lift(w1,offset) do w1,off
            decompose(Point2f0, Circle(off, w1/2))
        end, color=c, strokewidth = 1, strokecolor=c)
    end

    # draw polygon for the root
    c = maybe_get(treeplot[:color], tree.id)
    w = maybe_get(treeplot[:branch_width], tree.id)
    root_poly = lift(w,offset) do w,offset
        Point2f0[(-2w/2*√(3), -2w/2), (0,2w), (2w/2*√(3), -2w/2)].+offset
    end
    poly!(treeplot, root_poly, color=c, strokewidth = 1, strokecolor=c)

    treeplot.attributes[:ports] = ports
    treeplot
end


using KernelAbstractions: NoneEvent

using Oceananigans.Grids: Flat, Bounded

import Oceananigans.Utils: launch!

#####
##### Implements a "single column model mode" for HydrostaticFreeSurfaceModel
#####

const SingleColumnGrid = AbstractGrid{<:AbstractFloat, <:Flat, <:Flat, <:Bounded}

PressureField(arch, ::SingleColumnGrid) = (pHY′ = nothing,)
FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, velocities, arch, ::SingleColumnGrid) = nothing
FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, arch, ::SingleColumnGrid) = nothing

validate_momentum_advection(momentum_advection, ::SingleColumnGrid) = nothing
validate_tracer_advection(tracer_advection::AbstractAdvectionScheme, ::SingleColumnGrid) = nothing, NamedTuple()
validate_tracer_advection(tracer_advection::Nothing, ::SingleColumnGrid) = nothing, NamedTuple()

@inline hydrostatic_pressure_y_gradient(i, j, k, grid, ::Nothing) = zero(eltype(grid))

@inline launch!(arch, ::SingleColumnGrid, ::Val{:xy},  args...; kwargs...) = NoneEvent()
@inline launch!(arch, ::SingleColumnGrid, ::Val{dims}, args...; kwargs...) where dims = launch!(arch, grid, dims, args...; kwargs...)


#=
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"] = [
    ((1, 2, 2, 2),),
    ]
insts["fast"] = [
    ((2, 2, 3, 2),),
    ((3, 3, 3, 3),),
    ((3, 3, 5, 4),),
    ((5, 2, 5, 3),),
    ((1, 30, 2, 30),),
    ((10, 1, 3, 1),),
    ]
insts["various"] = [
    ((3, 3, 3, 3),),
    ((3, 3, 5, 4),),
    ((5, 2, 5, 3),),
    ((2, 30, 3, 30),),
    ((2, 15, 6, 15),),
    ((2, 30, 2, 40),),
    ((8, 3, 3, 3),),
    ]
return (PolyEnvelopeJuMP, insts)