######################## common traits for the plots

max_width = 493.228346

set_theme!(
    font = "Helvetica", #
    fontsize = 10,
    resolution = (max_width, 500.0),
    linewidth = 0.3,
    strokewidth = 1px,
    markersize = 3px,
    rowgap = Fixed(10),
    colgap = Fixed(10),
    LLegend = (markersize = 10px, markerstrokewidth = 1, patchsize = (10, 10), rowgap = Fixed(2), titlegap = Fixed(5), groupgap = Fixed(10), titlehalign = :left, gridshalign = :left, framecolor = :transparent, padding = 0, linewidth = 0.3),
    LAxis = (xticklabelsize = 8, yticklabelsize = 8, xlabel = "X (cm)", ylabel = "Y (cm)", autolimitaspect = 1, xtickalign = 1, xticksize = 3, ytickalign = 1, yticksize = 3, xticklabelpad = 4)
)
markers = Dict("turning_point" => '•', "center_of_search" => '■')
brighten(c, p = 0.5) = weighted_color_mean(p, c, colorant"white")
mydecompose(origin, radii) = [origin + radii .* Iterators.reverse(sincos(t)) for t in range(0, stop = 2π, length = 51)]
mydecompose(x) = mydecompose(x.origin, x.radii)
legendmarkers = OrderedDict(
                            "track" => (linestyle = nothing, linewidth = 0.3, color = :black),
                            "burrow" => (color = :black, marker = '⋆', strokecolor = :black, markerstrokewidth = 0.5, strokewidth = 0.5, markersize = 15px),
                            "fictive burrow" => (color = :white, marker = '⋆', strokecolor = :black, markerstrokewidth = 0.5, strokewidth = 0.5, markersize = 15px),
                            "turning point" => (color = :black, marker = markers["turning_point"], strokecolor = :transparent, markersize = 15px),
                            "center of search" => (color = :black, marker = markers["center_of_search"], strokecolor = :transparent, markersize = 5px),
                            "mean ± FWHM" => [(color = brighten(colorant"black", 0.75), strokecolor = :transparent, polypoints = mydecompose(Point2f0(0.5, 0.5), Vec2f0(0.75, 0.5))),
                                           (color = :white, marker = '+', strokecolor = :transparent, markersize = 10px),
                                          ])
function getellipse(xy)
    n = length(xy)
    X = Array{Float64}(undef, 2, n)
    for i in 1:n
        X[:,i] = xy[i]
    end
    dis = fit(DiagNormal, X)
    radii = sqrt(2log(2))*sqrt.(var(dis)) # half the FWHM
    (origin = Point2f0(mean(dis)), radii = Vec2f0(radii))
end
function distance2nest(track)
    length(searching(track)) < 10 && return Inf
    t = homing(track)
    i = findfirst(>(0) ∘ last, t)
    isnothing(i) ? Inf : abs(first(t[i]))
end
apply_element(xs) = apply_element.(xs)
apply_element(x::NamedTuple) =  :marker ∈ keys(x) ? MarkerElement(; x...) :
                                :linestyle ∈ keys(x) ? LineElement(; x...) :
                                PolyElement(; x...)

label!(scene, ax, letter) = LText(scene, letter, fontsize = 12, padding = (10, 0, 0, 10), halign = :left, valign = :top, bbox = lift(FRect2D, ax.scene.px_area), font ="Noto Sans Bold")
function plottracks!(ax, g::GroupedDataFrame)
    for gg in g
        for r in eachrow(gg)
            lines!(ax, r.track.coords; legendmarkers["track"]..., color = r.color)
        end
    end
end
function plottracks!(ax, g::DataFrame)
    for r in eachrow(g)
        lines!(ax, r.track.coords; legendmarkers["track"]..., color = r.color)
    end
end
function plotpoints!(ax, g, point_type)
    if !ismissing(g[1].nest[1])
        scatter!(ax, [zero(Point2f0)]; legendmarkers["burrow"]...)
    end
    for (k, gg) in pairs(g)
        xy = gg[!, point_type]
        ellipse = getellipse(xy)
        c = gg.groupcolor[1]
        poly!(ax, mydecompose(ellipse), color = RGBA(brighten(c, 0.5), 0.5))
        scatter!(ax, [ellipse.origin]; legendmarkers["mean ± FWHM"][2]...)
        scatter!(ax, xy; legendmarkers[replace(point_type, "_" => " ")]..., color = RGBA(c, 0.75))
        if k.group ≠ "none"
            scatter!(ax, [Point2f0(intended(k.group))]; legendmarkers["fictive burrow"]..., strokecolor = gg.groupcolor[1])
        end
    end
end

######################## Figure 5 ################

function figure5(df)
    gdf = groupby(df, [:group, :nest_coverage])
    g = gdf[[(group = group, nest_coverage = "closed") for group in ("right","left","towards", "away")]]
    polys = OrderedDict(string(k.group) => (color = v.groupcolor[1], strokecolor = :transparent) for (k, v) in pairs(g))
    scene, layout = layoutscene(0, resolution = (max_width, 600.0))
    ax = layout[1,1] = LAxis(scene)
    plottracks!(ax, g)
    hidexdecorations!(ax, grid = false, ticklabels = false, ticks = false)
    hideydecorations!(ax, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax, "a")
    axs = [LAxis(scene) for _ in 1:2]
    for (i, point_type) in enumerate(("turning_point", "center_of_search"))
        plotpoints!(axs[i], g, point_type)
    end
    layout[2,1:2] = axs
    hideydecorations!(axs[1], grid = false, ticklabels = false, ticks = false)
    hideydecorations!(axs[2], grid = false)
    hidexdecorations!.(axs, grid = false, ticklabels = false, ticks = false)
    layout[3, 1:2] = LText(scene, "X (cm)");
    linkaxes!(axs...)
    label!(scene, axs[1], "c")
    label!(scene, axs[2], "d")
    g = DataFrame(gdf[(group = "towards", nest_coverage = "open")], copycols = false)
    sort!(g, :track, by = distance2nest)
    d = g[1:3,:]
    c = d.groupcolor[1]
    ax = layout[1,2] = LAxis(scene)
    scatter!(ax, [zero(Point2f0)]; legendmarkers["burrow"]...)
    plottracks!(ax, d)
    scatter!(ax, turningpoint.(d.track); legendmarkers["turning point"]..., color = RGBA(c, 0.75))
    hidexdecorations!(ax, grid = false)
    hideydecorations!(ax, grid = false, ticklabels = false, ticks = false)
    linkxaxes!(ax, axs[2])
    label!(scene, ax, "b")
    layout[1:2, 0] = LText(scene, "Y (cm)", rotation = π/2);
    layout[4, 2:3] = LLegend(scene, apply_element.(values.([polys, legendmarkers])), collect.(keys.([polys, legendmarkers])), ["Direction of displacements", " "], orientation = :horizontal, nbanks = 2, tellheight = true, height = Auto(), groupgap = 30);
    scene
end
# FileIO.save("a.pdf", scene)

######################## Figure 6 ################
function figure6(df)
    gdf = groupby(df, :group)
    g = gdf[[(group = "zero",)]]
    polys = OrderedDict(string(k.group) => (color = v.groupcolor[1], strokecolor = :transparent) for (k, v) in pairs(g))
    scene, layout = layoutscene(0, resolution = (max_width, 400.0))
    ax = layout[1,1] = LAxis(scene)
    plottracks!(ax, g)
    hidexdecorations!(ax, grid = false, ticklabels = false, ticks = false)
    # hideydecorations!(ax, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax, "a")
    axs = [LAxis(scene) for _ in 1:2]
    for (i, point_type) in enumerate(("turning_point", "center_of_search"))
        plotpoints!(axs[i], g, point_type)
    end
    layout[1,2:3] = axs
    hideydecorations!(axs[1], grid = false, ticklabels = false, ticks = false)
    hideydecorations!(axs[2], grid = false)
    hidexdecorations!.(axs, grid = false, ticklabels = false, ticks = false)
    layout[2, 1:3] = LText(scene, "X (cm)");
    linkaxes!(axs...)
    # linkyaxes!(ax, axs...)
    label!(scene, axs[1], "b")
    label!(scene, axs[2], "c")
    layout[3, 1:3] = LLegend(scene, apply_element(values(legendmarkers)), collect(keys(legendmarkers)), orientation = :horizontal, nbanks = 2, tellheight = true, height = Auto(), groupgap = 30);
    scene
end
# FileIO.save("a.pdf", scene)

######################## Figure 7 ################
function figure7(df)
    gdf = groupby(df, :group)
    g = gdf[[(group = "far",)]]
    polys = OrderedDict(string(k.group) => (color = v.groupcolor[1], strokecolor = :transparent) for (k, v) in pairs(g))
    scene, layout = layoutscene(0, resolution = (max_width, 300.0))
    ax = layout[1,1] = LAxis(scene)
    plottracks!(ax, g)
    hidexdecorations!(ax, grid = false, ticklabels = false, ticks = false)
    # hideydecorations!(ax, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax, "a")
    axs = [LAxis(scene) for _ in 1:2]
    for (i, point_type) in enumerate(("turning_point", "center_of_search"))
        plotpoints!(axs[i], g, point_type)
    end
    layout[1,2:3] = axs
    hideydecorations!(axs[1], grid = false, ticklabels = false, ticks = false)
    hideydecorations!(axs[2], grid = false)
    hidexdecorations!.(axs, grid = false, ticklabels = false, ticks = false)
    layout[2, 1:3] = LText(scene, "X (cm)");
    linkaxes!(axs...)
    # linkaxes!(ax, axs...)
    label!(scene, axs[1], "b")
    label!(scene, axs[2], "c")
    x = copy(legendmarkers);
    delete!(x, "burrow")
    layout[3,1:3] = LLegend(scene, apply_element(values(x)), collect(keys(x)), orientation = :horizontal, nbanks = 2, tellheight = true, height = Auto(), groupgap = 30);
    scene
end
# FileIO.save("a.pdf", scene)

######################## Figure 4 ################

function binit(track, h, nbins, m, M)
    o = Union{Variance, Missing}[Variance() for _ in 1:nbins]
    d = track.rawcoords
    to = findfirst(x -> x.xy[2] > 0, d) - 1
    for (p1, p2) in Iterators.take(zip(d, lag(d, -1, default = d[to])), to)
        y = -(p2.xy[2] + p1.xy[2])/2
        if m < y < M
            i = StatsBase.binindex(h, y)
            v = norm(p2.xy - p1.xy)/(p2.t - p1.t)
            fit!(o[i], v)
        end
    end
    replace!(x -> nobs(x) < 2 ? missing : x, o)
end
function figure4(df)
    gdf = groupby(df, :group)
    g = gdf[[(group = "none",)]]
    polys = OrderedDict(string(k.group) => (color = v.groupcolor[1], strokecolor = :transparent) for (k, v) in pairs(g))
    scene = Scene(resolution = (max_width, 500.0), camera=campixel!);
    sz = round(Int, max_width/4)
    lo = GridLayout(
                    scene, 4, 3,
                    # colsizes = [Auto(), Fixed(sz), Auto(), Auto()],
                    # rowsizes = [Auto(), Auto(), Fixed(sz), Auto()],
                    alignmode = Outside(0, 0, 0, 0)
    )
    lo[2, 1:3] = LText(scene, "X (cm)", tellheight = true);
    ax1 = lo[1,1] = LAxis(scene, autolimitaspect = 1)
    # ax1 = lo[2,2] = LAxis(scene, xaxisposition = :top, xticklabelalign = (:center, :bottom), autolimitaspect = 1)
    # lo[1:2, 1] = LText(scene, "Y (cm)", rotation = π/2);
    plottracks!(ax1, g)
    rect = FRect2D(-10,-10,20,20)
    spinecolor = RGB(only(distinguishable_colors(1, [colorant"white"; g[1].color], dropseed = true))) #:yellow
    lines!(ax1, rect, color = spinecolor);
    hidexdecorations!(ax1, grid = false, ticklabels = false, ticks = false);
    # hideydecorations!(ax1, grid = false, ticklabels = false, ticks = false);
    # hideydecorations!(ax1, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax1, "a")
    axs = [LAxis(scene, autolimitaspect = 1) for _ in 1:2]
    # axs = [LAxis(scene, xaxisposition = :top, xticklabelalign = (:center, :bottom), autolimitaspect = 1) for _ in 1:2]
    for (i, point_type) in enumerate(("turning_point", "center_of_search"))
        plotpoints!(axs[i], g, point_type)
    end
    lo[1,2:3] = axs
    hideydecorations!(axs[1], grid = false, ticklabels = false, ticks = false)
    hideydecorations!(axs[2], grid = false)
    hidexdecorations!.(axs, grid = false, ticklabels = false, ticks = false)
    linkaxes!(axs...)
    label!(scene, axs[1], "b")
    label!(scene, axs[2], "c")
    d = DataFrame(g[1], copycols = false)
    sort!(d, :turning_point, by = norm)
    d.color .= getcolor(d.groupcolor)
    d = d[1:3,:]
    # postTP = 25
    ax2 = lo[3,1] = LAxis(scene,
                          aspect = DataAspect(),
                          autolimitaspect = 1,
                          bottomspinecolor = spinecolor,
                          topspinecolor = spinecolor,
                          leftspinecolor = spinecolor,
                          rightspinecolor = spinecolor,
                          # xaxisposition = :top, xticklabelalign = (:center, :bottom)
                         )
    for r in eachrow(d)
        lines!(ax2, homing(r.track); legendmarkers["track"]..., color = r.color)
        postTP = findfirst(xy -> any(>(10)∘abs, xy), searching(r.track))
        lines!(ax2, searching(r.track)[1:postTP]; legendmarkers["track"]..., color = r.color)
        # lines!(ax2, searching(r.track)[1:postTP]; legendmarkers["track"]..., color = 1:postTP, colormap = [r.color, colorant"white"])
        scatter!(ax2, [r.turning_point]; legendmarkers["turning point"]..., color = RGBA(r.color, 0.75))
    end
    scatter!(ax2, [zero(Point2f0)]; legendmarkers["burrow"]...)
    ax2.targetlimits[] = rect
    hidexdecorations!(ax2, grid = false, ticklabels = false, ticks = false)
    # hideydecorations!(ax2, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax2, "d")
    m, M = (0, 120)
    nbins = 6
    bins = range(m, stop = M, length = nbins + 1)
    mbins = StatsBase.midpoints(bins)
    h = Histogram(bins)
    g = DataFrame(g[1], copycols = false)
    g[!, :id] .= 1:nrow(g)
    DataFrames.transform!(g, :id, :track => ByRow(x -> binit(x, h, nbins, m, M)) => :yv)
    μ = [Variance() for _ in 1:nbins]
    for i in 1:nbins
        reduce(merge!, skipmissing(yv[i] for yv in g.yv), init = μ[i])
    end
    bandcolor = RGB(only(distinguishable_colors(1, [colorant"white"; g.color; spinecolor], dropseed = true))) #:yellow
    ax = lo[3,2:3] = LAxis(scene,
                           aspect = nothing,
                           autolimitaspect = nothing,
                           xlabel = "Distance to burrow (cm)",
    ylabel = "Speed (cm/s)",
    xticks = mbins,
    xreversed = true,
    yaxisposition = :right,
    yticklabelalign = (:left, :center)
   )
    bh = band!(ax, mbins, mean.(μ) .- std.(μ), mean.(μ) .+ std.(μ), color = RGBA(bandcolor, 0.25))
    lh = lines!(ax, mbins, mean.(μ), color = :white, linewidth = 5)
    for r in eachrow(g)
        xy = [Point2f0(x, mean(y)) for (x,y) in zip(mbins, r.yv) if !ismissing(y)]
        lines!(ax, xy; legendmarkers["track"]..., color = r.color)
        scatter!(ax, xy; legendmarkers["turning point"]..., color = r.color)
    end
    ylims!(ax, 0, ax.limits[].origin[2] + ax.limits[].widths[2])
    xlims!(ax, Iterators.reverse(extrema(mbins))...)
    label!(scene, ax, "e")
    x = copy(legendmarkers)
    delete!(x, "fictive burrow")
    x["mean speed ± std"] = [(color = RGBA(bandcolor, 0.25), strokecolor = :transparent), (linestyle = nothing, linewidth = 3, color = :white)]
    x["individual speed"] = [x["track"], x["turning point"]]
    lo[4, 1:3] = LLegend(scene, apply_element(values(x)), collect(keys(x)), orientation = :horizontal, nbanks = 2, tellheight = true, height = Auto(), groupgap = 30);
    tight_ticklabel_spacing!(ax)
    colsize!(lo, 2, Relative(1/3))
    rowsize!(lo, 3, Aspect(2, 1))
    overlay = Scene(scene, camera = campixel!, raw = true)
    for (x, corner) in zip((-10.0, 10.0), (:topleft, :topright))
        point = Node(Point2f0(x, -10.0))
        topleft_ax2 = lift(getfield(MakieLayout, corner), ax2.scene.px_area)
        point_screenspace = lift(ax1.scene.camera.projectionview, ax1.scene.camera.pixel_space, point) do pv, pspace, point
            projected = Point2f0(AbstractPlotting.project(inv(pspace) * pv, point)[1:2]...) .+ AbstractPlotting.origin(ax1.scene.px_area[])
        end
        lines!(overlay, @lift([$point_screenspace, $topleft_ax2]), color = spinecolor)
    end
    scene
end

function save_figures(df)
    # mkpath("figures")
    for figure in (:figure4, :figure5, :figure6, :figure7)
        @eval FileIO.save($("$figure.pdf"), $figure($df))
    end
end


module Soss

import Base.rand
using Random
using Reexport: @reexport

@reexport using Distributions
@reexport using StatsFuns
using NamedTupleTools

import MacroTools: prewalk, postwalk, @q, striplines, replace, flatten, @capture
import MLStyle
@reexport using MonteCarloMeasurements

using LazyArrays
using FillArrays

include("core/statement.jl")
include("core/model.jl")
include("core/jointdistribution.jl")
include("core/canonical.jl")
include("core/dependencies.jl")
include("core/toposort.jl")
include("core/weighted.jl")
include("core/utils.jl")

include("distributions/dist.jl")
include("distributions/for.jl")
include("distributions/iid.jl")
include("distributions/mix.jl")
# include("distributions/flat.jl")

include("primitives/rand.jl")
include("primitives/logpdf.jl")
include("primitives/xform.jl")
include("primitives/likelihood-weighting.jl")
# include("primitives/bijectors.jl")

include("transforms/predictive.jl")
include("transforms/markovblanket.jl")
include("transforms/do.jl")

include("symbolic/symbolic.jl")
include("symbolic/codegen.jl")

include("particles.jl")
include("plots.jl")

include("inference/rejection.jl")
include("inference/dynamicHMC.jl")
include("inference/advancedhmc.jl")

# include("weighted.jl")
#
# # include("graph.jl")
# # # include("optim.jl")
# include("importance.jl")
#
# # # include("sobols.jl")
# # # include("fromcube.jl")
# # # include("tocube.jl")
end # module


export SpeakerListenerEnv

# TODO: review the implementation of this environment

mutable struct SpeakerListenerEnv{T<:Vector{Float64}} <: AbstractEnv
    target::T
    content::T
    player_vel::T
    player_pos::T
    landmarks_pos::Vector{T}
    landmarks_num::Int
    ϵ::Any
    damping::Any
    max_accel::Any
    space_dim::Int
    init_step::Int
    play_step::Int
    max_steps::Int
    continuous::Bool
end

"""
    SpeakerListenerEnv(;kwargs...)

`SpeakerListenerEnv` is a simple cooperative environment of two agents, a `Speaker` and a `Listener`, who are placed in an environment
with `N` landmarks. At each episode, the `Listener` must navigate to a particular landmark(`env.target`) and obtains **reward** based
on its distance to the target. However, while the `Listener` can observe the relative position of the landmarks, it doesn't know which
is the target landmark. Conversely, the `Speaker` can observe the target's landmark, and it can produce a communication output(`env.content`)
at each time step which is observed by the `Listener`.

For more concrete description, you can refer to:
* [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://arxiv.org/abs/1706.02275)
* [multiagent-particle-envs](https://github.com/openai/multiagent-particle-envs)

# Keyword arguments
- `N::Int = 3`, the number of landmarks in the environment.
- `stop = 0.01`, when the distance between the `Listener` and the target is smaller than the `stop`, the game will be terminated.
- `damping = 0.25`, for simulation of the physical space, `Listener`'s action will meet the damping in each step.
- `max_accel = 0.02`, the maximum acceleration of the `Listener` in each step.
- `space_dim::Int = 2`, the dimension of the environment's space.
- `max_steps::Int = 25`, the maximum playing steps in one episode.
- `continuous::Bool = true`, set to `false` if you want the action_space of the players to be discrete. Otherwise, the action_space will be continuous.
"""
function SpeakerListenerEnv(;
    N::Int = 3,
    stop = 0.01,
    damping = 0.25,
    max_accel = 0.5,
    space_dim::Int = 2,
    max_steps::Int = 50,
    continuous::Bool = true,
)
    SpeakerListenerEnv(
        zeros(N),
        zeros(N),
        zeros(space_dim),
        zeros(space_dim),
        [zeros(space_dim) for _ in Base.OneTo(N)],
        N,
        stop,
        damping,
        max_accel,
        space_dim,
        0,
        0,
        max_steps,
        continuous,
    )
end

function RLBase.reset!(env::SpeakerListenerEnv)
    env.init_step = 0
    env.play_step = 0
    env.target = zeros(env.landmarks_num)
    env.content = zeros(env.landmarks_num)
    env.player_pos = zeros(env.space_dim)
    env.landmarks_pos = [zeros(env.space_dim) for _ in Base.OneTo(env.landmarks_num)]
end

RLBase.is_terminated(env::SpeakerListenerEnv) =
    (reward(env) > -env.ϵ) || (env.play_step > env.max_steps)
RLBase.players(::SpeakerListenerEnv) = (:Speaker, :Listener, CHANCE_PLAYER)

RLBase.state(env::SpeakerListenerEnv, ::Observation{Any}, players::Tuple) =
    Dict(p => state(env, p) for p in players)

RLBase.state(env::SpeakerListenerEnv, ::Observation{Any}, player::Symbol) =
# for speaker, it can observe the target and help listener to arrive it.
    if player == :Speaker
        env.target
        # for listener, it can observe current velocity, relative positions of landmarks, and speaker's conveyed information.
    elseif player == :Listener
        vcat(
            env.player_vel...,
            (
                vcat((landmark_pos .- env.player_pos)...) for
                landmark_pos in env.landmarks_pos
            )...,
            env.content...,
        )
    else
        @error "No player $player."
    end

RLBase.state(env::SpeakerListenerEnv, ::Observation{Any}, ::ChancePlayer) =
    vcat(env.landmarks_pos, [env.player_pos])

RLBase.state_space(env::SpeakerListenerEnv, ::Observation{Any}, players::Tuple) =
    Dict(player => state_space(env, player) for player in players)

RLBase.state_space(env::SpeakerListenerEnv, ::Observation{Any}, player::Symbol) =
    if player == :Speaker
        # env.target
        Space(SVector(ntuple(_ -> (0.0, 1.0), env.landmarks_num)))
    elseif player == :Listener
        Space(
            Float64,
            length(env.player_vel) +
            length(env.landmarks_pos) * length(env.player_pos) +
            length(env.content),
        )
    else
        @error "No player $player."
    end

RLBase.state_space(env::SpeakerListenerEnv, ::Observation{Any}, ::ChancePlayer) = [
    # landmarks' positions
    (
        Space((SVector(ntuple(_ -> -Inf .. Inf, env.space_dim)))) for
        _ in Base.OneTo(env.landmarks_num)
    )...,
    # player's position, no bounds.
    Space(SVector(ntuple(_ -> -Inf .. Inf, env.space_dim))),
]

RLBase.action_space(env::SpeakerListenerEnv, players::Tuple) =
    Space(Dict(p => action_space(env, p) for p in players))

RLBase.action_space(env::SpeakerListenerEnv, player::Symbol) =
    if player == :Speaker
        env.continuous ?
        Space([ClosedInterval(0, 1) for _ in Base.OneTo(env.landmarks_num)]) :
        Space([ZeroTo(1) for _ in Base.OneTo(env.landmarks_num)])
    elseif player == :Listener
        # there has two directions in each dimension.
        env.continuous ?
        Space([ClosedInterval(0, 1) for _ in Base.OneTo(2 * env.space_dim)]) :
        Space([ZeroTo(1) for _ in Base.OneTo(2 * env.space_dim)])
    else
        @error "No player $player."
    end

function RLBase.action_space(env::SpeakerListenerEnv, ::ChancePlayer)
    if env.init_step < env.landmarks_num + 1
        Space(SVector(ntuple(_ -> -1.0 .. 1.0, env.space_dim)))
    else
        Space(OneTo(env.landmarks_num))
    end
end

function (env::SpeakerListenerEnv)(action, ::ChancePlayer)
    env.init_step += 1
    if env.init_step <= env.landmarks_num
        env.landmarks_pos[env.init_step] = action
    elseif env.init_step == env.landmarks_num + 1
        env.player_pos = action
    else
        @assert action in Base.OneTo(env.landmarks_num) "The target should be assigned to one of the landmarks."
        env.target[action] = 1.0
    end
end

function (env::SpeakerListenerEnv)(actions::Dict, players::Tuple)
    @assert length(actions) == length(players)
    for p in players
        env(actions[p], p)
    end
    env.play_step += 1
end

function (env::SpeakerListenerEnv)(action::Vector, player::Symbol)
    if player == :Speaker
        # update conveyed information.
        env.content = round.(action)
    elseif player == :Listener
        # update velocity, here env.damping is for simulation physical rule.
        action = round.(action)
        acceleration = [action[2*i] - action[2*i-1] for i in Base.OneTo(env.space_dim)]
        env.player_vel .*= (1 - env.damping)
        env.player_vel .+= (acceleration * env.max_accel)
        # update position
        env.player_pos .+= env.player_vel * 0.1 # velocity * time
    else
        @error "No player $player."
    end
end

RLBase.reward(::SpeakerListenerEnv, ::ChancePlayer) = -Inf

function RLBase.reward(env::SpeakerListenerEnv, p)
    if sum(env.target) == 1
        goal = findfirst(env.target .== 1.0)
        -sum((env.landmarks_pos[goal] .- env.player_pos) .^ 2)
    else
        -Inf
    end
end

RLBase.current_player(env::SpeakerListenerEnv) =
    if env.init_step < env.landmarks_num + 2
        CHANCE_PLAYER
    else
        (:Speaker, :Listener)
    end

RLBase.NumAgentStyle(::SpeakerListenerEnv) = MultiAgent(2)
RLBase.DynamicStyle(::SpeakerListenerEnv) = SIMULTANEOUS
RLBase.ActionStyle(::SpeakerListenerEnv) = MINIMAL_ACTION_SET
RLBase.InformationStyle(::SpeakerListenerEnv) = IMPERFECT_INFORMATION
RLBase.ChanceStyle(::SpeakerListenerEnv) = EXPLICIT_STOCHASTIC


"""
    NeighborLongitudinalResult

A structure to retrieve information about a neihbor in the longitudinal direction i.e. rear and front neighbors on the same lane.
If the neighbor index is equal to `nothing` it means there is no neighbor.

# Fields
- `ind::Union{Nothing, Int64}` index of the neighbor in the scene
- `Δs::Float64` positive distance along the lane between vehicles positions
"""
struct NeighborLongitudinalResult
    ind::Union{Nothing, Int64} # index in scene of the neighbor
    Δs::Float64 # positive distance along lane between vehicles' positions
end

function get_headway(veh_rear::Vehicle1D, veh_fore::Vehicle1D, roadway::StraightRoadway)
    return get_headway(get_front(veh_rear), get_rear(veh_fore), roadway)
end
function get_neighbor_fore(scene::Frame{Entity{State1D, D, I}}, vehicle_index::I, roadway::StraightRoadway) where {D, I}
    ego = scene[vehicle_index]
    best_ind = nothing
    best_gap = Inf
    for (i,veh) in enumerate(scene)
        if i != vehicle_index
            Δs = get_headway(ego, veh, roadway)
            if Δs < best_gap
                best_gap, best_ind = Δs, i
            end
        end
    end
    return NeighborLongitudinalResult(best_ind, best_gap)
end
function get_neighbor_rear(scene::Frame{Entity{State1D, D, I}}, vehicle_index::I, roadway::StraightRoadway) where {D, I}
    ego = scene[vehicle_index]
    best_ind = nothing
    best_gap = Inf
    for (i,veh) in enumerate(scene)
        if i != vehicle_index
            Δs = get_headway(veh, ego, roadway)
            if Δs < best_gap
                best_gap, best_ind = Δs, i
            end
        end
    end
    return NeighborLongitudinalResult(best_ind, best_gap)
end

abstract type VehicleTargetPoint end
struct VehicleTargetPointFront <: VehicleTargetPoint end
get_targetpoint_delta(::VehicleTargetPointFront, veh::Entity{S, D, I}) where {S,D<:AbstractAgentDefinition, I} = length(veh.def)/2*cos(posf(veh.state).ϕ)
struct VehicleTargetPointCenter <: VehicleTargetPoint end
get_targetpoint_delta(::VehicleTargetPointCenter, veh::Entity{S, D, I}) where {S,D<:AbstractAgentDefinition, I} = 0.0
struct VehicleTargetPointRear <: VehicleTargetPoint end
get_targetpoint_delta(::VehicleTargetPointRear, veh::Entity{S, D, I}) where {S,D<:AbstractAgentDefinition, I} = -length(veh.def)/2*cos(posf(veh.state).ϕ)

const VEHICLE_TARGET_POINT_CENTER = VehicleTargetPointCenter()

"""
    get_neighbor_fore_along_lane(scene::EntityFrame{S,D,I}, roadway::Roadway, tag_start::LaneTag, s_base::Float64, targetpoint_primary::VehicleTargetPoint, targetpoint_valid::VehicleTargetPoint;
                                 max_distance_fore::Float64 = 250.0, index_to_ignore::Int=-1) where {S,D<:AbstractAgentDefinition,I}

Return the index and the longitudinal distance of the vehicle that is in the same lane as scene[vehicle_index] and
in front of it with the smallest distance along the lane. The result is returned as a `NeighborLongitudinalResult` object.
The method will search on the current lane first, and if no vehicle is found it will continue to travel along the lane following next_lane(lane, roadway).
If no vehicle is found within `max_distance_fore,` the index takes a value of `nothing`.

# Notes on the optional arguments:
- `targetpoint_primary::VehicleTargetPoint` the reference point whose distance we want to minimize
- `targetpoint_valid::VehicleTargetPoint` the reference point, which if distance to is positive, we include the vehicle

"""
function get_neighbor_fore_along_lane(
    scene::EntityFrame{S,D,I},
    roadway::Roadway,
    tag_start::LaneTag,
    s_base::Float64,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_fore::Float64 = 250.0, # max distance to search forward [m]
    index_to_ignore::Int=-1,
    ) where {S,D<:AbstractAgentDefinition,I}
    best_ind = nothing
    best_dist = max_distance_fore
    tag_target = tag_start

    dist_searched = 0.0
    while dist_searched < max_distance_fore

        lane = roadway[tag_target]

        for (i,veh) in enumerate(scene)
            if i != index_to_ignore

                s_adjust = NaN
                if posf(veh.state).roadind.tag == tag_target
                    s_adjust = 0.0
                elseif is_between_segments_hi(posf(veh.state).roadind.ind, lane.curve) &&
                       is_in_entrances(roadway[tag_target], posf(veh.state).roadind.tag)
                    distance_between_lanes = norm(VecE2(roadway[tag_target].curve[1].pos - roadway[posf(veh.state).roadind.tag].curve[end].pos))
                    s_adjust = -(roadway[posf(veh.state).roadind.tag].curve[end].s + distance_between_lanes)

                elseif is_between_segments_lo(posf(veh.state).roadind.ind) &&
                       is_in_exits(roadway[tag_target], posf(veh.state).roadind.tag)
                    distance_between_lanes = norm(VecE2(roadway[tag_target].curve[end].pos - roadway[posf(veh.state).roadind.tag].curve[1].pos))
                    s_adjust = roadway[tag_target].curve[end].s + distance_between_lanes
                end

                if !isnan(s_adjust)
                    s_valid = posf(veh.state).s + get_targetpoint_delta(targetpoint_valid, veh) + s_adjust
                    dist_valid = s_valid - s_base + dist_searched
                    if dist_valid ≥ 0.0
                        s_primary = posf(veh.state).s + get_targetpoint_delta(targetpoint_primary, veh) + s_adjust
                        dist = s_primary - s_base + dist_searched
                        if dist < best_dist
                            best_dist = dist
                            best_ind = i
                        end
                    end
                end
            end
        end

        if best_ind != nothing
            break
        end

        if !has_next(lane) ||
           (tag_target == tag_start && dist_searched != 0.0) # exit after visiting this lane a 2nd time
            break
        end

        dist_searched += (lane.curve[end].s - s_base)
        s_base = -norm(VecE2(lane.curve[end].pos - next_lane_point(lane, roadway).pos)) # negative distance between lanes
        tag_target = next_lane(lane, roadway).tag
    end

    NeighborLongitudinalResult(best_ind, best_dist)
end
function get_neighbor_fore_along_lane(
    scene::EntityFrame{S,D,I},
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    veh_ego = scene[vehicle_index]
    tag_start = posf(veh_ego.state).roadind.tag
    s_base = posf(veh_ego.state).s + get_targetpoint_delta(targetpoint_ego, veh_ego)

    get_neighbor_fore_along_lane(scene, roadway, tag_start, s_base,
        targetpoint_primary, targetpoint_valid,
        max_distance_fore=max_distance_fore, index_to_ignore=vehicle_index)
end
function get_neighbor_fore_along_left_lane(
    scene::EntityFrame{S,D,I},
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    retval = NeighborLongitudinalResult(nothing, max_distance_fore)

    veh_ego = scene[vehicle_index]
    lane = get_lane(roadway, veh_ego)
    if n_lanes_left(lane, roadway) > 0
        lane_left = roadway[LaneTag(lane.tag.segment, lane.tag.lane + 1)]
        roadproj = proj(posg(veh_ego.state), lane_left, roadway)
        tag_start = roadproj.tag
        s_base = lane_left[roadproj.curveproj.ind, roadway].s + get_targetpoint_delta(targetpoint_ego, veh_ego)

        retval = get_neighbor_fore_along_lane(scene, roadway, tag_start, s_base,
                                              targetpoint_primary, targetpoint_valid,
                                              index_to_ignore=vehicle_index,
                                              max_distance_fore=max_distance_fore)
    end

    retval
end
function get_neighbor_fore_along_right_lane(
    scene::EntityFrame{S,D,I},
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    retval = NeighborLongitudinalResult(nothing, max_distance_fore)

    veh_ego = scene[vehicle_index]
    lane = get_lane(roadway, veh_ego)
    if n_lanes_right(lane, roadway) > 0
        lane_right = roadway[LaneTag(lane.tag.segment, lane.tag.lane - 1)]
        roadproj = proj(posg(veh_ego.state), lane_right, roadway)
        tag_start = roadproj.tag
        s_base = lane_right[roadproj.curveproj.ind, roadway].s + get_targetpoint_delta(targetpoint_ego, veh_ego)

        retval = get_neighbor_fore_along_lane(scene, roadway, tag_start, s_base,
                                              targetpoint_primary, targetpoint_valid,
                                              index_to_ignore=vehicle_index,
                                              max_distance_fore=max_distance_fore)
    end

    retval
end

function get_neighbor_fore_along_lane(
    scene::EntityFrame{S,D,I},
    roadway::Roadway,
    tag_start::LaneTag,
    s_base::Float64;
    max_distance_fore::Float64 = 250.0, # max distance to search forward [m]
    index_to_ignore::Int=-1,
    ) where {S,D<:AbstractAgentDefinition,I}

    get_neighbor_fore_along_lane(scene, roadway, tag_start, s_base,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        max_distance_fore=max_distance_fore,
        index_to_ignore=index_to_ignore)
end
function get_neighbor_fore_along_lane(scene::EntityFrame{S,D,I}, vehicle_index::Int, roadway::Roadway;
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    get_neighbor_fore_along_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_fore=max_distance_fore)
end
function get_neighbor_fore_along_left_lane(scene::EntityFrame{S,D,I}, vehicle_index::Int, roadway::Roadway;
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    get_neighbor_fore_along_left_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_fore=max_distance_fore)
end
function get_neighbor_fore_along_right_lane(scene::EntityFrame{S,D,I}, vehicle_index::Int, roadway::Roadway;
    max_distance_fore::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    get_neighbor_fore_along_right_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_fore=max_distance_fore)
end

function get_neighbor_rear_along_lane(
    scene::EntityFrame{S,D,I},
    roadway::Roadway,
    tag_start::LaneTag,
    s_base::Float64,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_rear::Float64 = 250.0, # max distance to search rearward [m]
    index_to_ignore::Int=-1,
    ) where {S,D<:AbstractAgentDefinition,I}

    best_ind = nothing
    best_dist = max_distance_rear
    tag_target = tag_start

    ignore = Set{I}()

    dist_searched = 0.0
    while dist_searched < max_distance_rear

        lane = roadway[tag_target]

        for (i,veh) in enumerate(scene)
            if i != index_to_ignore && !in(veh.id, ignore)

                s_adjust = NaN

                if posf(veh.state).roadind.tag == tag_target
                    s_adjust = 0.0

                elseif is_between_segments_hi(posf(veh.state).roadind.ind, lane.curve) &&
                       is_in_entrances(roadway[tag_target], posf(veh.state).roadind.tag)

                    distance_between_lanes = norm(VecE2(roadway[tag_target].curve[1].pos - roadway[posf(veh.state).roadind.tag].curve[end].pos))
                    s_adjust = -(roadway[posf(veh.state).roadind.tag].curve[end].s + distance_between_lanes)

                elseif is_between_segments_lo(posf(veh.state).roadind.ind) &&
                       is_in_exits(roadway[tag_target], posf(veh.state).roadind.tag)

                    distance_between_lanes = norm(VecE2(roadway[tag_target].curve[end].pos - roadway[posf(veh.state).roadind.tag].curve[1].pos))
                    s_adjust = roadway[tag_target].curve[end].s + distance_between_lanes
                end

                if !isnan(s_adjust)
                    s_valid = posf(veh.state).s + get_targetpoint_delta(targetpoint_valid, veh) + s_adjust
                    dist_valid = s_base - s_valid + dist_searched
                    if dist_valid ≥ 0.0
                        s_primary = posf(veh.state).s + get_targetpoint_delta(targetpoint_primary, veh) + s_adjust
                        dist = s_base - s_primary + dist_searched
                        if dist < best_dist
                            best_dist = dist
                            best_ind = i
                        end
                    else
                        push!(ignore, veh.id)
                    end
                end
            end
        end

        if best_ind != nothing
            break
        end

        if !has_prev(lane) ||
           (tag_target == tag_start && dist_searched != 0.0) # exit after visiting this lane a 2nd time
            break
        end

        dist_searched += s_base
        s_base = lane.curve[end].s + norm(VecE2(lane.curve[end].pos - prev_lane_point(lane, roadway).pos)) # end of prev lane plus crossover
        tag_target = prev_lane(lane, roadway).tag
    end

    NeighborLongitudinalResult(best_ind, best_dist)
end
function get_neighbor_rear_along_lane(
    scene::EntityFrame{S,D,I},
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_rear::Float64 = 250.0 # max distance to search rearward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    veh_ego = scene[vehicle_index]
    tag_start = posf(veh_ego.state).roadind.tag
    s_base = posf(veh_ego.state).s + get_targetpoint_delta(targetpoint_ego, veh_ego)

    get_neighbor_rear_along_lane(scene, roadway, tag_start, s_base,
        targetpoint_primary, targetpoint_valid,
        max_distance_rear=max_distance_rear, index_to_ignore=vehicle_index)
end
function get_neighbor_rear_along_left_lane(
    scene::EntityFrame{S,D,I},
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    retval = NeighborLongitudinalResult(nothing, max_distance_rear)

    veh_ego = scene[vehicle_index]
    lane = get_lane(roadway, veh_ego)
    if n_lanes_left(lane, roadway) > 0
        lane_left = roadway[LaneTag(lane.tag.segment, lane.tag.lane + 1)]
        roadproj = proj(posg(veh_ego.state), lane_left, roadway)
        tag_start = roadproj.tag
        s_base = lane_left[roadproj.curveproj.ind, roadway].s + get_targetpoint_delta(targetpoint_ego, veh_ego)

        retval = get_neighbor_rear_along_lane(scene, roadway, tag_start, s_base,
                                              targetpoint_primary, targetpoint_valid,
                                              index_to_ignore=vehicle_index,
                                              max_distance_rear=max_distance_rear)
    end

    retval
end
function get_neighbor_rear_along_right_lane(
    scene::EntityFrame{S,D,I},
    vehicle_index::Int,
    roadway::Roadway,
    targetpoint_ego::VehicleTargetPoint,
    targetpoint_primary::VehicleTargetPoint, # the reference point whose distance we want to minimize
    targetpoint_valid::VehicleTargetPoint; # the reference point, which if distance to is positive, we include the vehicle
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    retval = NeighborLongitudinalResult(nothing, max_distance_rear)

    veh_ego = scene[vehicle_index]
    lane = get_lane(roadway, veh_ego)
    if n_lanes_right(lane, roadway) > 0
        lane_right = roadway[LaneTag(lane.tag.segment, lane.tag.lane - 1)]
        roadproj = proj(posg(veh_ego.state), lane_right, roadway)
        tag_start = roadproj.tag
        s_base = lane_right[roadproj.curveproj.ind, roadway].s + get_targetpoint_delta(targetpoint_ego, veh_ego)

        retval = get_neighbor_rear_along_lane(scene, roadway, tag_start, s_base,
                                              targetpoint_primary, targetpoint_valid,
                                              index_to_ignore=vehicle_index,
                                              max_distance_rear=max_distance_rear)
    end

    retval
end

function get_neighbor_rear_along_lane(
    scene::EntityFrame{S,D,I},
    roadway::Roadway,
    tag_start::LaneTag,
    s_base::Float64;
    max_distance_rear::Float64 = 250.0, # max distance to search rearward [m]
    index_to_ignore::Int=-1,
    ) where {S,D<:AbstractAgentDefinition,I}

    get_neighbor_rear_along_lane(scene, roadway, tag_start, s_base,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        max_distance_rear=max_distance_rear,
        index_to_ignore=index_to_ignore)
end
function get_neighbor_rear_along_lane(scene::EntityFrame{S,D,I}, vehicle_index::Int, roadway::Roadway;
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    get_neighbor_rear_along_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_rear=max_distance_rear)
end
function get_neighbor_rear_along_left_lane(
    scene::EntityFrame{S,D,I},
    vehicle_index::Int,
    roadway::Roadway;
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    get_neighbor_rear_along_left_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_rear=max_distance_rear)
end
function get_neighbor_rear_along_right_lane(
    scene::EntityFrame{S,D,I},
    vehicle_index::Int,
    roadway::Roadway;
    max_distance_rear::Float64 = 250.0 # max distance to search forward [m]
    ) where {S,D<:AbstractAgentDefinition,I}

    get_neighbor_rear_along_right_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_CENTER, VEHICLE_TARGET_POINT_CENTER,
        VEHICLE_TARGET_POINT_CENTER, max_distance_rear=max_distance_rear)
end

"""
    Project the given point to the same lane as the given RoadIndex.
This will return the projection of the point, along with the Δs along the lane from the RoadIndex.
"""
struct FrenetRelativePosition
    origin::RoadIndex
    target::RoadIndex
    Δs::Float64
    t::Float64
    ϕ::Float64
end
function get_frenet_relative_position(posG::VecSE2{Float64}, roadind::RoadIndex, roadway::Roadway;
    max_distance_fore::Float64 = 250.0, # max distance to search forward [m]
    max_distance_rear::Float64 = 250.0, # max distance to search backward [m]
    improvement_threshold::Float64 = 1e-4,
    )

    # project to current lane first
    tag_start = roadind.tag
    lane_start = roadway[tag_start]
    curveproj_start = proj(posG, lane_start, roadway, move_along_curves=false).curveproj
    curvept_start = lane_start[curveproj_start.ind, roadway]
    s_base = lane_start[roadind.ind, roadway].s
    s_proj = curvept_start.s
    Δs = s_proj - s_base
    sq_dist_to_curve = normsquared(VecE2(posG - curvept_start.pos))
    retval = FrenetRelativePosition(roadind,
                RoadIndex(curveproj_start.ind, tag_start), Δs, curveproj_start.t, curveproj_start.ϕ)

    # search downstream
    if has_next(lane_start)
        dist_searched = lane_start.curve[end].s - s_base
        s_base = -norm(VecE2(lane_start.curve[end].pos - next_lane_point(lane_start, roadway).pos)) # negative distance between lanes
        tag_target = next_lane(lane_start, roadway).tag
        while dist_searched < max_distance_fore

            lane = roadway[tag_target]
            curveproj = proj(posG, lane, roadway, move_along_curves=false).curveproj
            sq_dist_to_curve2 = normsquared(VecE2(posG - lane[curveproj.ind, roadway].pos))

            if sq_dist_to_curve2 < sq_dist_to_curve - improvement_threshold

                sq_dist_to_curve = sq_dist_to_curve2
                s_proj = lane[curveproj.ind, roadway].s
                Δs = s_proj - s_base + dist_searched
                retval = FrenetRelativePosition(
                            roadind, RoadIndex(curveproj.ind, tag_target),
                            Δs, curveproj.t, curveproj.ϕ)
            end

            if !has_next(lane)
                break
            end

            dist_searched += (lane.curve[end].s - s_base)
            s_base = -norm(VecE2(lane.curve[end].pos - next_lane_point(lane, roadway).pos)) # negative distance between lanes
            tag_target = next_lane(lane, roadway).tag

            if tag_target == tag_start
                break
            end
        end
    end

    # search upstream
    if has_prev(lane_start)
        dist_searched = s_base
        tag_target = prev_lane(lane_start, roadway).tag
        s_base = roadway[tag_target].curve[end].s + norm(VecE2(lane_start.curve[1].pos - prev_lane_point(lane_start, roadway).pos)) # end of the lane

        while dist_searched < max_distance_fore

            lane = roadway[tag_target]
            curveproj = proj(posG, lane, roadway, move_along_curves=false).curveproj
            sq_dist_to_curve2 = normsquared(VecE2(posG - lane[curveproj.ind, roadway].pos))

            if sq_dist_to_curve2 < sq_dist_to_curve - improvement_threshold

                sq_dist_to_curve = sq_dist_to_curve2
                s_proj = lane[curveproj.ind, roadway].s
                dist = s_base - s_proj  + dist_searched
                Δs = -dist
                retval = FrenetRelativePosition(
                            roadind, RoadIndex(curveproj.ind, tag_target),
                            Δs, curveproj.t, curveproj.ϕ)
            end

            if !has_prev(lane)
                break
            end

            dist_searched += s_base
            s_base = lane.curve[end].s + norm(VecE2(lane.curve[1].pos - prev_lane_point(lane, roadway).pos)) # length of this lane plus crossover
            tag_target = prev_lane(lane, roadway).tag

            if tag_target == tag_start
                break
            end
        end
    end

    retval
end
get_frenet_relative_position(veh_fore::Entity{S, D, I}, veh_rear::Entity{S, D, I}, roadway::Roadway) where {S,D, I} = get_frenet_relative_position(posg(veh_fore.state), posf(veh_rear.state).roadind, roadway)


# Testing
# =======
#
# Utilities to assist testing of BioJulia packages.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/BioCore.jl/blob/master/LICENSE.md

module Testing

"""
    random_array(n::Integer, elements, probs)

Create a random array of length `n`, composed of possible `elements`, which are
selected according to their `prob`.
"""
function random_array(n::Integer, elements, probs)
    @assert sum(probs) == 1.0 "the probabilities in probs, don't sum to 1.0!"
    cumprobs = cumsum(probs)
    x = Vector{eltype(elements)}(undef, n)
    for i in 1:n
        x[i] = elements[searchsorted(cumprobs, rand()).start]
    end
    return x
end

"""
    random_seq(n::Integer, nts, probs)

Create a random sequence of length `n`, composed of `nts`, according to their
`prob`.
"""
function random_seq(n::Integer, nts, probs)
    return String(random_array(n, nts, probs))
end

"""
    random_dna(n, probs = [0.24, 0.24, 0.24, 0.24, 0.04])

Create a random DNA sequence of length `n`, by sampling the nucleotides
A, C, G, T, and N, according to their probability in `probs`.
"""
function random_dna(n, probs = [0.24, 0.24, 0.24, 0.24, 0.04])
    return random_seq(n, ['A', 'C', 'G', 'T', 'N'], probs)
end

"""
    random_rna(n, probs = [0.24, 0.24, 0.24, 0.24, 0.04])

Create a random RNA sequence of length `n`, by sampling the nucleotides
A, C, G, U, and N, according to their probability in `probs`.
"""
function random_rna(n, probs = [0.24, 0.24, 0.24, 0.24, 0.04])
    return random_seq(n, ['A', 'C', 'G', 'U', 'N'], probs)
end

"""
    random_aa(n, probs = [0.24, 0.24, 0.24, 0.24, 0.04])

Create a random amino acid sequence of length `n`, by sampling the possible
amino acid characters.
"""
function random_aa(len)
    return random_seq(len,
        ['A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I',
         'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'X' ],
        push!(fill(0.049, 20), 0.02))
end

"""
    intempdir(fn::Function, parent = tempdir())

Execute some function `fn` in a temporary directory.
After the function is executed, the directory is cleaned by doing the
equivalent of `rm -r`.
"""
function intempdir(fn::Function, parent = tempdir())
    dirname = mktempdir(parent)
    try
        cd(fn, dirname)
    finally
        rm(dirname, recursive = true)
    end
end

"""
    random_interval(minstart, maxstop)

Create a random interval between a minimum starting point, and a maximum stop
point.
"""
function random_interval(minstart, maxstop)
    start = rand(minstart:maxstop)
    return start:rand(start:maxstop)
end

end # Module Testing


##########################################################################
# Copyright 2017 Samuel Ridler.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##########################################################################

# For determining transient period duration at start and end of simulation, based on call response durations.
# Useful for deciding duration of simulation warm-up and cool-down periods.
# Loads sim from config, splits calls into batches, runs simulation for each batch,
# then saves and plots the moving average of response durations.

using JEMSS
using Statistics

# parameters
configFilename = "sim_config.xml"
outputPath = "output"
displayPlots = false

println("Loading sim")
sim = initSim(configFilename, doPrint = false)

isdir(outputPath) || mkdir(outputPath)

callArrivalTimes = [c.arrivalTime for c in sim.calls] # backup, as call arrival times will change when calls are batched

# break calls into call batches, each with callBatchNumCalls many calls, or numCallBatches many call batches
# mutates: calls
function createCallBatchesByCount!(calls::Vector{Call}, startTime::Float;
	callBatchNumCalls::Int = nullIndex, numCallBatches::Int = nullIndex)

	@assert(all(c -> c.status == callNullStatus, calls))

	numCalls = length(calls)
	@assert(sum([callBatchNumCalls, numCallBatches] .!= nullIndex) == 1)
	if callBatchNumCalls != nullIndex
		numCallBatches = Int(floor(numCalls / callBatchNumCalls))
	elseif numCallBatches != nullIndex
		callBatchNumCalls = Int(floor(numCalls / numCallBatches))
	end

	callBatches = Vector{Vector{Call}}(undef,numCallBatches)
	callBatchStartTimes = Vector{Float}(undef,numCallBatches)
	callBatchFirstCallIndex = callBatchLastCallIndex = 0 # first and last call indices for current call batch
	for callBatchIndex = 1:numCallBatches
		# get call batch start time
		if callBatchIndex == 1
			callBatchStartTimes[1] = startTime
		else
			callBatchStartTimes[callBatchIndex] = calls[callBatchLastCallIndex].arrivalTime
		end
		# find first and last calls for call batch callBatchIndex
		callBatchFirstCallIndex = callBatchLastCallIndex + 1
		callBatchLastCallIndex = callBatchFirstCallIndex + callBatchNumCalls - 1
		callBatches[callBatchIndex] = calls[callBatchFirstCallIndex:callBatchLastCallIndex]
	end

	# renumber calls from 1 to n within each call batch, set each call batch to start at startTime
	for (callBatchIndex, callBatch) in enumerate(callBatches)
		for (callIndex, call) in enumerate(callBatch)
			call.index = callIndex
			call.arrivalTime -= (callBatchStartTimes[callBatchIndex] - startTime)
		end
	end

	return callBatches
end

# run simulation for each callBatch in callBatches
# mutates: sim, callBatches
function simulateCallBatches!(sim::Simulation, callBatches::Vector{Vector{Call}})
	# run simulation for each call batch
	numCallBatches = length(callBatches)
	for (callBatchIndex, callBatch) in enumerate(callBatches)
		print("\rSimulating call batch ", callBatchIndex, " of ", numCallBatches)
		setSimCalls!(sim, callBatch)
		simulate!(sim)
		sim.calls = Call[] # remove call batch
	end
	println()
end

function calcAvgResponseDurations(callBatches::Vector{Vector{Call}})
	# for callBatches that have already been used in simulation, calculate avgResponseDurations,
	# where avgResponseDurations[i] is the average response duration of call i from all batches
	@assert(all(callBatch -> all(call -> call.responseDuration != nullTime, callBatch), callBatches))
	avgResponseDurations = [mean([callBatches[i][j].responseDuration for i = 1:length(callBatches)]) for j = 1:length(callBatches[1])] * (24*60)
	return avgResponseDurations
end

function calcMovingAvgResponseDurations(avgResponseDurations::Vector{Float}, movingWindowSize::Int)
	callBatchNumCalls = length(avgResponseDurations)

	# # for constant window size:
	# numWindows = max(0, callBatchNumCalls - movingWindowSize + 1)
	# movingAvgResponseDurations = [mean(avgResponseDurations[i + (0:movingWindowSize-1)]) for i = 1:numWindows]

	# # for window size = min(window size, number of calls so far):
	# numWindows = callBatchNumCalls
	# movingAvgResponseDurations = [mean(avgResponseDurations[max(1,i-(movingWindowSize-1)):i]) for i = 1:numWindows]

	# for window size = 1 + 2 * min(movingWindowSize, max allowable window size):
	# (each average value is at centre of window)
	numWindows = callBatchNumCalls
	movingAvgResponseDurations = Vector{Float}(undef,numWindows)
	for windowIndex = 1:numWindows
		currentWindowSize = min(movingWindowSize, windowIndex - 1, numWindows - windowIndex)
		movingAvgResponseDurations[windowIndex] = mean(avgResponseDurations[windowIndex .+ (-currentWindowSize : currentWindowSize)])
	end

	# # for window size = 1 + 2 * min(movingWindowSize, max allowable window size):
	# # (each average value is at centre of window, last n (=movingWindowSize) average values are not calculated)
	# # this follows Welch's procedure in "On the problem of the initial transient in steady-state simulation"
	# numWindows = max(0, callBatchNumCalls - movingWindowSize)
	# movingAvgResponseDurations = Vector{Float}(undef,numWindows)
	# for windowIndex = 1:numWindows
		# currentWindowSize = min(movingWindowSize, max(0, windowIndex-1))
		# movingAvgResponseDurations[windowIndex] = mean(avgResponseDurations[windowIndex .+ (-currentWindowSize : currentWindowSize)])
	# end

	return movingAvgResponseDurations
end

# For calls in sim.calls, split into n batches (n = callSetsNumCallBatches[i], i=1,...), then simulate.
# The response duration of the jth call is averaged over all batches.
function getCallSetsAvgResponseDurations(sim, callSetsNumCallBatches::Vector{Int})::Vector{Vector{Float}}
	numCallSets = length(callSetsNumCallBatches)
	callSetsAvgResponseDurations = []
	calls = sim.calls # shorthand
	for (callSetIndex, numCallBatches) in enumerate(callSetsNumCallBatches)
		println("Call set ", callSetIndex, " of ", numCallSets)

		# create batches (for current set), simulate for each, calculate average response durations
		reset!(sim)
		callBatches = createCallBatchesByCount!(calls, sim.startTime; numCallBatches = numCallBatches)
		simulateCallBatches!(sim, callBatches)
		push!(callSetsAvgResponseDurations, calcAvgResponseDurations(callBatches))

		# undo changes from createCallBatchesByCount!(), and reset calls
		for (i,c) in enumerate(calls)
			c.index = i
			c.arrivalTime = callArrivalTimes[i]
		end
		setSimCalls!(sim, calls)
	end

	return callSetsAvgResponseDurations
end

function writeCallSetsAvgResponseDurationsFile(filename::String, callSetsAvgResponseDurations::Vector{Vector{Float}}, callSetsNumCallBatches::Vector{Int}, movingWindowSizes::Vector{Int})
	tables = Table[]
	for (i, numCallBatches) in enumerate(callSetsNumCallBatches)
		avgResponseDurations = callSetsAvgResponseDurations[i]
		m = length(avgResponseDurations)
		header = ["callIndex"]
		data = Array{Any,2}(undef, m, length(movingWindowSizes) + 1) # to be filled
		data[:] .= "" # empty values
		data[:,1] = 1:m
		for (j, movingWindowSize) in enumerate(movingWindowSizes)
			push!(header, "windowSize_$movingWindowSize")
			movingAvgResponseDurations = calcMovingAvgResponseDurations(avgResponseDurations, movingWindowSize)
			if movingAvgResponseDurations != []
				data[1:length(movingAvgResponseDurations), j+1] = movingAvgResponseDurations
			end
		end
		push!(tables, Table("numCallBatches_$numCallBatches", header, data))
	end
	writeTablesToFile(filename, tables)
end

# Given the response duration of the jth call averaged over all batches, moving averages are calculated using
# each value in movingWindowSizes. One plot is produced for each value in callSetsNumCallBatches.
function plotMovingAvgResponseDurations(callSetsAvgResponseDurations::Vector{Vector{Float}}, callSetsNumCallBatches::Vector{Int}, movingWindowSizes::Vector{Int})
	numCallSets = length(callSetsNumCallBatches)
	@assert(numCallSets == length(callSetsAvgResponseDurations))
	plots = []
	for i = 1:numCallSets
		# create plot of moving average response durations with different moving window sizes
		push!(plots, Plots.plot(title = string("Number of call batches: ", callSetsNumCallBatches[i]),
			xaxis = "Call index", yaxis = "Mean response duration (minutes)"))
		for movingWindowSize in movingWindowSizes
			movingAvgResponseDurations = calcMovingAvgResponseDurations(callSetsAvgResponseDurations[i], movingWindowSize)
			if movingAvgResponseDurations != []
				Plots.plot!(plots[i], movingAvgResponseDurations, label = string(movingWindowSize));
			end
		end
	end
	return plots
end

println("Simulating with different call batches")
callSetsNumCallBatches = [5,10] # for runtests.jl
# callSetsNumCallBatches = [100:100:1000;] # for actual use
movingWindowSizes = [0, 1, 2, 5, 10, 20, 50, 100, 200, 500, 1000];
callSetsAvgResponseDurations = getCallSetsAvgResponseDurations(sim, callSetsNumCallBatches)

println("Saving results to file")
filename = joinpath(outputPath, "transient.csv")
writeCallSetsAvgResponseDurationsFile(filename, callSetsAvgResponseDurations, callSetsNumCallBatches, movingWindowSizes)

if displayPlots
	println("Plotting")
	import Plots
	Plots.plotly()
	plots = plotMovingAvgResponseDurations(callSetsAvgResponseDurations, callSetsNumCallBatches, movingWindowSizes)
	for plot in plots display(plot) end
end


using DFTK
using MPI
using JSON
using JLD2
using Printf
using Plots
using Statistics
using LinearAlgebra
using Dates
using Logging
include("quantum_espresso.jl")

function build_magnetic_moments(atoms::AbstractArray; magmoms...)
    magmoms = Dict(magmoms)
    map(atoms) do (element, positions)
        element => fill(magmoms[element.symbol], length(positions))
    end
end

function attach_pseudos(atoms::AbstractArray; pseudomap...)
    pseudomap = Dict(pseudomap)

    map(atoms) do (element, position)
        pspfile = get(pseudomap, element.symbol, nothing)
        ElementPsp(element.symbol, psp=load_psp(pspfile)) => position
    end
end

function standard_cases()
    cases = Dict{String, Any}(
        "adaptive"    => (damping=DFTK.AdaptiveDamping(), ),  # Parameter-free adaptive damping
        "adaptive_05" => (damping=DFTK.AdaptiveDamping(0.5), ),
    )
    for α in 1:9
        cases["standard_0$α"] = (damping=DFTK.FixedDamping(α / 10), )
    end
    cases["standard_10"] = (damping=DFTK.FixedDamping(1.0), )
    cases
end

function random_logfile(prefix)
    logfile = "$(prefix)_$(@sprintf "%05d" rand(UInt16)).log"
    while isfile(logfile)
        logfile = "$(prefix)_$(@sprintf "%05d" rand(UInt16)).log"
    end
    logfile
end

function run_step_on_cases(payload, cases)
    spay = string(payload)
    header_printed = false

    anymissing = false
    for dir in sort(collect(keys(cases)))
        config = cases[dir]

        MPI.Barrier(MPI.COMM_WORLD)
        if isfile(joinpath(config.directory, string(payload) * ".done"))
            continue
        end

        for i in 1:10  # To avoid infinite loop
            MPI.Barrier(MPI.COMM_WORLD)
            res = nothing
            logfile = nothing
            if mpi_master()
                if !header_printed
                    println("#####" * "#"^length(spay) * "#####")
                    println("#--  $(string(payload))  --#")
                    println("#####" * "#"^length(spay) * "#####")
                    header_printed = true
                end

                logfile = random_logfile(joinpath(config.directory, string(payload)))
                println(basename(dir), "  ->  ", basename(dir) * "/" * basename(logfile))
                res = open(logfile, "w") do fp
                    redirect_stdout(fp) do
                        with_logger(SimpleLogger(stdout, Logging.Debug)) do
                            payload(config)
                        end
                    end
                end
                println("    ... $res\n")
            else
                res = payload(config)
            end

            if res == :exhausted
                if mpi_master()
                    rm(logfile)
                    anymissing || touch(joinpath(config.directory, string(payload) * ".done"))
                end
                break
            elseif res == :missing
                anymissing = true
                mpi_master() && rm(logfile)
            end
        end
    end

    MPI.Barrier(MPI.COMM_WORLD)
    nothing
end

function has_results(fileprefix)
    resultsfile = fileprefix * ".scfres.json"
    isfile(resultsfile) && open(JSON.parse, resultsfile, "r")["done"]
end

function basis_from_config(config)
    model = model_PBE(config.lattice, config.atoms;
                      magnetic_moments=config.magnetic_moments, config.kwargs_model...)
    fft_size = compute_fft_size(model, config.Ecut, supersampling=config.supersampling)
    if !haskey(config, :kgrid)
        kgrid = kgrid_size_from_minimal_spacing(config.lattice, 2π * 0.022)
    else
        kgrid=config.kgrid
    end
    PlaneWaveBasis(model, config.Ecut, kgrid=kgrid, kshift=[0, 0, 0], fft_size=fft_size)
end

function scf(config)
    basis = basis_from_config(config)

    scfkey = nothing
    if mpi_master()
        for scf in sort(collect(keys(config.cases)))
            lockfile = joinpath(config.directory, scf * ".running")
            fileprefix = joinpath(config.directory, scf)
            (isfile(lockfile) || has_results(fileprefix))  && continue
            write(lockfile, read(`hostname`, String))
            scfkey = scf
            break
        end
    end
    scfkey = MPI.bcast(scfkey, 0, MPI.COMM_WORLD)

    scfkey === nothing && return :exhausted
    fileprefix  = joinpath(config.directory, scfkey)
    lockfile    = joinpath(config.directory, scfkey * ".running")
    resultsfile = joinpath(config.directory, scfkey * ".scfres.json")
    guessfile   = joinpath(config.directory, "init.jld2")

    !isfile(guessfile) && return :missing
    ρ, ψ0, V = load_guess(basis, guessfile)

    try
        if mpi_master()
            println("#\n#-- SCF $(basename(config.directory)) -- case=$scfkey\n#")
            println("    keywords: $((; config.cases[scfkey]..., config.kwargs_scf...))")
        end

        MPI.Barrier(MPI.COMM_WORLD)
        mpi_master() && DFTK.reset_timer!(DFTK.timer)
        if haskey(config.cases[scfkey], :type) && config.cases[scfkey].type == :qe
            caseargs = Dict(k => v for (k, v) in pairs(config.cases[scfkey]) if k != :type)
            run_qe(basis, fileprefix, resultsfile;
                   magnetic_moments=config.magnetic_moments, caseargs...)
        else
            run_dftk(basis, resultsfile; ρ=ρ, ψ0=ψ0, V=V,
                     mixing=config.mixing, config.cases[scfkey]..., config.kwargs_scf...)
        end
        MPI.Barrier(MPI.COMM_WORLD)
        mpi_master() && println(DFTK.timer)
    catch
        mpi_master() && rm(joinpath(config.directory, "scf.done"), force=true)
        rethrow()
    finally
        mpi_master() && rm(lockfile)
    end
    mpi_master() && plot_convergence_summary(config)

    :done
end

function dump_guess(config)
    basis     = basis_from_config(config)
    guessfile = joinpath(config.directory, "init.jld2")
    isfile(guessfile) && return :exhausted

    ρ = guess_density(basis, config.magnetic_moments)
    _, ham = energy_hamiltonian(basis, nothing, nothing; ρ=ρ)

    V = nothing
    isnothing(config.potential) || (V = config.potential(basis))
    if !isnothing(V)
        ham = DFTK.hamiltonian_with_total_potential(ham, V)
    end

    n_bands = (haskey(config, :n_bands) ? (n_bands = config.n_bands)
                                        : DFTK.default_n_bands(basis.model))
    res_ψ0 = DFTK.next_density(ham; n_bands=DFTK.default_n_bands(basis.model),
                               n_ep_extra=3, miniter=2, tol=0.0001)
    diagiter = DFTK.mpi_mean(mean(res_ψ0.diagonalization.iterations), basis.comm_kpts)
    mpi_master() && println("guess:    diag = $diagiter  diagtol = 0.0001")
    isnothing(V) && (V = DFTK.total_local_potential(ham))

    # Gather ψ array on master
    ψ = DFTK.gather_kpts(res_ψ0.ψ, basis)
    if mpi_master()
        JLD2.jldopen(guessfile, "w") do jld
            jld["ρ"] = ρ
            jld["V"] = V
            jld["ψ"] = ψ
        end
    end

    :done
end

function load_guess(basis, guessfile)
    JLD2.jldopen(guessfile, "r") do jld
        ψ = jld["ψ"][basis.krange_thisproc]
        jld["ρ"], ψ, jld["V"]
    end
end

function run_dftk(basis, resultsfile;
                  ρ, ψ0, V, mixing=SimpleMixing(),
                  damping, tol=1e-11, type=:potmix, diagtol_max=1e-3, kwargs...)

    n_bands = (haskey(config, :n_bands) ? (n_bands = config.n_bands)
                                        : DFTK.default_n_bands(basis.model))
    if mpi_master()
        println()
        println("hostname        = $(read(`hostname`, String))")
        println("started on      = $(Dates.now())")
        println("julia threads   = $(Threads.nthreads())")
        println("BLAS threads    = $(BLAS.get_num_threads())")
        println("MPI procs       = $(mpi_nprocs())")
        println()
        println("temperature     = $(basis.model.temperature)")
        println("smearing        = $(basis.model.smearing)")
        println("lattice         = $(round.(basis.model.lattice, sigdigits=4))")
        println("Ecut            = $(basis.Ecut)")
        println("fft_size        = $(basis.fft_size)")
        println("kgrid           = $(basis.kgrid)")
        println("kshift          = $(basis.kshift)")
        println("n_irreducible_k = $(sum(length, basis.krange_allprocs))")
        println("n_bands         = $(n_bands)")
        println("n_electrons     = $(basis.model.n_electrons)")
        println("type            = $type")
        println("mixing          = $mixing")
        println("damping         = $damping")
        flush(stdout)
    end

    # Collection of data to dump later
    save_dict = Dict(
        "temperature"  => basis.model.temperature,
        "smearing"     => string(basis.model.smearing),
        "lattice"      => basis.model.lattice,
        "Ecut"         => basis.Ecut,
        "fft_size"     => basis.fft_size,
        "kgrid"        => basis.kgrid,
        "n_kirred"     => sum(length, basis.krange_allprocs),
        "n_bands"      => n_bands,
        "n_electr"     => basis.model.n_electrons,
        "mixing"       => string(mixing),
        "damping"      => string(damping),
        "type"         => string(type),
        "kwargs"       => string(kwargs),
        "done"         => false,
        #
        "energies"     => Float64[],
        "residuals"    => Float64[],
        "ndiag"        => Int[],
        "α"            => Float64[],
    )

    # Empty existing results
    if mpi_master()
        open(fp -> JSON.print(fp, save_dict), resultsfile, "w")
    end

    function ExtractCallback()
        function callback(info)
            if info.stage != :finalize
                push!(save_dict["energies"],  info.energies.total)
                push!(save_dict["ndiag"],     length(info.diagonalization))
                push!(save_dict["α"], isnan(info.α) ? DFTK.trial_damping(damping) : info.α)

                if type == :potmix
                    # Note: This misses a factor sqrt(info.basis.dvol)), which is taken into
                    # account in the analyse.jl script
                    push!(save_dict["residuals"], norm(info.Vout - info.Vin))
                elseif type == :densmix
                    # Note: This misses a factor sqrt(info.basis.dvol)), which is taken into
                    # account in the analyse.jl script
                    push!(save_dict["residuals"], norm(info.ρout - info.ρin))
                end
            end

            if mpi_master()
                open(fp -> JSON.print(fp, save_dict), resultsfile, "w")
            end
        end
    end

    determine_diagtol = DFTK.ScfDiagtol(; ratio_ρdiff=0.01, diagtol_max, diagtol_min=1e-12)
    callback = ExtractCallback() ∘ DFTK.ScfDefaultCallback()

    if type == :potmix
        run_scf = ((damping isa DFTK.AdaptiveDamping) ? DFTK.scf_potential_mixing_adaptive
                                                      : DFTK.scf_potential_mixing)
        scfres = run_scf(basis; tol=tol, ψ=ψ0, determine_diagtol=determine_diagtol,
                         mixing=mixing, callback=callback, damping=damping, ρ=ρ, V=V,
                         diag_miniter=2, kwargs...)
    elseif type == :densmix
        @assert damping isa Number
        scfres = self_consistent_field(basis; tol=tol, ψ=ψ0,
                                       determine_diagtol=determine_diagtol,
                                       mixing=mixing, callback=callback,
                                       damping=damping, ρ=ρ, kwargs...)
    else
        error("Unsupported SCF type: $type")
    end

    if mpi_master()
        println()
        println(scfres.energies)
        flush(stdout)
        save_dict["done"] = true
        open(fp -> JSON.print(fp, save_dict), resultsfile, "w")
    end

    scfres
end

function run_qe(basis, qe_directory, resultsfile; magnetic_moments,
                mixing_beta=0.7, kwargs...)
    if !mpi_master()
        MPI.Barrier(MPI.COMM_WORLD)
        return nothing
    end
    # This only runs on MPI master

    println()
    println("hostname        = $(read(`hostname`, String))")
    println("started on      = $(Dates.now())")
    println()
    flush(stdout)

    model = basis.model
    outputfile = joinpath(qe_directory, "espresso.pwo")
    if !isfile(outputfile)
        qeres = scf_quantum_espresso(model.lattice, model.atoms;
                                     Ecut=basis.Ecut,
                                     fileprefix=joinpath(qe_directory, "espresso"),
                                     temperature=model.temperature,
                                     mixing_beta=mixing_beta,
                                     kgrid=basis.kgrid,
                                     magnetic_moments=magnetic_moments,
                                     kwargs...)
    @assert qeres.outputfile == outputfile
    end

    # Copy output file to stdout
    println.(readlines(outputfile))
    flush(stdout)

    # Extract results
    save_dict = Dict(
        "done"         => true,
        "energies"     => Float64[],
        "residuals"    => Float64[],
        "ndiag"        => Int[],
        "α"            => Float64[],
    )
    for iteration in parse_quantum_espresso_pwo(outputfile)
        push!(save_dict["energies"],  iteration.total_energy)
        push!(save_dict["residuals"], iteration.accuracy)
        push!(save_dict["ndiag"],     1)
        push!(save_dict["α"],         mixing_beta)
    end
    open(fp -> JSON.print(fp, save_dict), resultsfile, "w")

    MPI.Barrier(MPI.COMM_WORLD)
    nothing
end

function collect_minimal_total_energy(config::NamedTuple; kwargs...)
    collect_minimal_total_energy(config.directory, config.cases; kwargs...)
end
function collect_minimal_total_energy(directory::AbstractString, cases::AbstractDict;
                                      kwargs...)
    collect_minimal_total_energy(directory, collect(keys(cases)); kwargs...)
end
function collect_minimal_total_energy(directory::AbstractString, cases::AbstractVector;
                                      tol_resid=0.05, tol_energy=1e-4)
    mintotals = Float64[]
    for scfkey in sort(cases)
        startswith(scfkey, "espresso_") && continue
        resultfile = joinpath(directory, scfkey * ".scfres.json")
        isfile(resultfile) || continue

        residuals = open(JSON.parse, resultfile, "r")["residuals"]
        energies  = open(JSON.parse, resultfile, "r")["energies"]
        energies === nothing && continue
        isempty(energies) && continue

        minidx = argmin(energies[max(end-4, 1):end]) + max(length(energies) - 4, 1) - 1
        residuals[minidx] > tol_resid && continue

        min_matching = findfirst(e -> abs(e - energies[minidx]) < tol_energy, mintotals)
        if min_matching === nothing
            push!(mintotals, energies[minidx])
        else
            mintotals[min_matching] = min(mintotals[min_matching], energies[minidx])
        end
    end

    mintotals
end

function parse_scfkey_label(scfkey)
    if scfkey == "adaptive"
        ("adaptive", 0.2)
    elseif startswith(scfkey, "adaptive_") && all(isnumeric, scfkey[10:end])
        digits = length(scfkey[10:end]) - 1
        alpha  = parse(Int, scfkey[10:end]) / 10^digits
        ("adaptive", alpha)
    elseif startswith(scfkey, "adaptive_")
        ("adaptive", scfkey[10:end])
    elseif startswith(scfkey, "standard_") && all(isnumeric, scfkey[10:end])
        digits = length(scfkey[10:end]) - 1
        alpha  = parse(Int, scfkey[10:end]) / 10^digits
        ("fixed", alpha)
    elseif startswith(scfkey, "espresso_") && all(isnumeric, scfkey[10:end])
        digits = length(scfkey[10:end]) - 1
        alpha  = parse(Int, scfkey[10:end]) / 10^digits
        ("espresso", alpha)
    elseif startswith(scfkey, "highprec_") && all(isnumeric, scfkey[10:end])
        digits = length(scfkey[10:end]) - 1
        alpha  = parse(Int, scfkey[10:end]) / 10^digits
        ("highprec", alpha)
    else
        (scfkey, nothing)
    end
end

function plot_convergence_summary(config; iteration=Inf, tol=1e-10, damping=true,
                                  outputfile=joinpath(config.directory, "summary.svg"),
                                  scfkeys=nothing, mintotals=nothing, kwargs...)
    function deduce_label(scfkey)
        method, alpha = parse_scfkey_label(scfkey)
        method * (isnothing(alpha) ? "" : " (α = $alpha)")
    end

    !mpi_master() && error("Plotting only on MPI master supported.")
    if isnothing(mintotals)
        mintotals = collect_minimal_total_energy(config)
    end
    isempty(mintotals) && return plot()  # Return on no data

    if scfkeys === nothing
        scfkeys = collect(keys(config.cases))
    end

    errors   = Dict{String, Vector{Float64}}()
    cumdiags = Dict{String, Vector{Int}}()
    dampings = Dict{String, Vector{Float64}}()
    whichmin = Dict{String, Int}()
    for scfkey in sort(scfkeys)
        method, _ = parse_scfkey_label(scfkey)
        resultfile = joinpath(config.directory, scfkey * ".scfres.json")
        isfile(resultfile) || continue

        energies  = open(JSON.parse, resultfile, "r")["energies"]
        ndiags    = open(JSON.parse, resultfile, "r")["ndiag"]
        damps     = open(JSON.parse, resultfile, "r")["α"]

        energies === nothing && continue
        if method == "espresso"
            values = open(JSON.parse, resultfile, "r")["residuals"]
            mintrial  = minimum(energies[max(end-4, 1):end])
            bestmatch = sortperm(abs.(mintrial .- mintotals))[1]
        else
            isempty(energies) && continue
            # Find best matching minimum
            mintrial  = minimum(energies[max(end-4, 1):end])
            bestmatch = sortperm(abs.(mintrial .- mintotals))[1]
            values = abs.(energies .- mintotals[bestmatch])
            idx = findfirst(values .< 1e-16)
            if !isnothing(idx)
                values = values[1:idx-1]
                ndiags = ndiags[1:idx-1]
                damps  = damps[1:idx-1]
            end
        end
        isempty(values) && continue

        whichmin[scfkey] = bestmatch
        errors[scfkey]   = values
        dampings[scfkey] = damps
        cumdiags[scfkey] = cumsum(ndiags)
    end

    # If no data return
    isempty(errors) && return plot()

    any_adaptive = false
    ymax = 0
    maxiter = 0
    colors = Dict{Any, Int}()
    p = plot(; yaxis=:log, legend=:topright, title=basename(config.directory), kwargs...)
    for (i, scfkey) in enumerate(sort(collect(keys(errors))))
        label = deduce_label(scfkey)
        maxiter = max(maxiter, length(errors[scfkey]))
        ymax = max(ymax, maximum(errors[scfkey]))

        if length(mintotals) > 1
            label *= " (min $(whichmin[scfkey]))"
        end

        method, alpha = parse_scfkey_label(scfkey)
        if isnothing(alpha)
            color = i
        elseif alpha in keys(colors)
            color = colors[alpha]
        else
            colors[alpha] = i
            color = colors[alpha]
        end

        mark = :x
        if method == "adaptive"
            mark = :o
            any_adaptive = true
        elseif method == "espresso"
            mark = :rtriangle
        elseif method == "highprec"
            mark = :ltriangle
        end
        plot!(p, cumdiags[scfkey], errors[scfkey];
              label=label, color=color, mark=mark, lw=1.3, markersize=4)
    end
    ylims!(p, (tol, 2 * ymax))
    ylabel!(p, "Total energy absolute error")

    if damping && any_adaptive
        q = twinx()
        for scfkey in sort(collect(keys(errors)))
            method, alpha = parse_scfkey_label(scfkey)
            method == "adaptive" || continue
            (alpha in keys(colors)) || continue
            plot!(q, cumdiags[scfkey], abs.(dampings[scfkey]), mark=:+,
                  color=colors[alpha], label="", linestyle=:dot)
        end
        ylims!(q, (0.0, 1.0))
        ylabel!(q, "Damping parameter")
    end
    xlims!(p, (0, min(iteration, maxiter)))
    xlabel!(p, "Iteration / Hamiltonian diagonalization")

    if isnothing(outputfile)
        p
    else
        savefig(p, outputfile)
    end
end


function IterativeLinearSolver(arg0::IterationManager)
    return IterativeLinearSolver((IterationManager,), arg0)
end

function IterativeLinearSolver(arg0::jint)
    return IterativeLinearSolver((jint,), arg0)
end

function solve(obj::IterativeLinearSolver, arg0::RealLinearOperator, arg1::RealVector)
    return jcall(obj, "solve", RealVector, (RealLinearOperator, RealVector), arg0, arg1)
end

function solve(obj::IterativeLinearSolver, arg0::RealLinearOperator, arg1::RealVector, arg2::RealVector)
    return jcall(obj, "solve", RealVector, (RealLinearOperator, RealVector, RealVector), arg0, arg1, arg2)
end

function solve_in_place(obj::IterativeLinearSolver, arg0::RealLinearOperator, arg1::RealVector, arg2::RealVector)
    return jcall(obj, "solveInPlace", RealVector, (RealLinearOperator, RealVector, RealVector), arg0, arg1, arg2)
end


module DynamicsOutputs

using RecursiveArrayTools: ArrayPartition
using DiffEqCallbacks: SavedValues, SavingCallback
using ComponentArrays: ComponentVector

using NQCModels: NQCModels
using NQCDynamics:
    Estimators,
    DynamicsUtils

using ..DynamicsUtils:
    get_positions,
    get_velocities,
    get_quantum_subsystem

using ..RingPolymers: get_centroid

"""
    create_saving_callback(quantities::NTuple{N, Symbol}; saveat=[]) where {N}

Get the `SavingCallback` that will populate `saved_values` with the result obtained
by evaluating the functions provided in `function_names`.
"""
function create_saving_callback(function_names::NTuple{N, Symbol}; saveat=[]) where {N}
    saved_values = SavedValues(Float64, Vector{Any})
    saving_function = OutputSaver(function_names)
    SavingCallback(saving_function, saved_values; saveat=saveat), saved_values
end

struct OutputSaver{N,F}
    function_names::NTuple{N, Symbol}
    output_functions::F
end

"""
    OutputSaver(function_names::NTuple{N, Symbol}) where {N}

Used to obtain the outputs for all functions given in `function_names`.
"""
function OutputSaver(function_names::NTuple{N, Symbol}) where {N}
    output_functions = tuple([getfield(DynamicsOutputs, f) for f in function_names]...)
    OutputSaver(function_names, output_functions)
end

"""
    (output::OutputSaver)(u, t, integrator)

Evaluates every function listed in `output.function_names` and returns all the results.
"""
function (output::OutputSaver)(u, t, integrator)
    out = Any[]
    sizehint!(out, length(output.function_names))
    evaluate_output!(out, u, t, integrator, output.output_functions...)
    return out
end

"""
Used to recursively evaluate every function for the [`OutputSaver`](@ref).

See here for a description of why it is written like this:
https://stackoverflow.com/questions/55840333/type-stability-for-lists-of-closures
"""
function evaluate_output!(out, u, t, integrator, f::F, output_functions...) where {F}
    push!(out, f(u, t, integrator))
    return evaluate_output!(out, u, t, integrator, output_functions...)
end
evaluate_output!(out, u, t, integrator) = out

export force
export velocity
export position
export centroid_position
export potential
export hamiltonian
export kinetic
export u
export quantum_subsystem
export state
export noise
export population
export adiabatic_population
export friction

"Get the force"
force(u, t, integrator) = -copy(integrator.p.calculator.derivative)

"Get the velocity"
velocity(u, t, integrator) = copy(get_velocities(u))

"Get the velocity of the ring polymer centroid"
centroid_velocity(u, t, integrator) = get_centroid(get_velocities(u))

"Get the position"
position(u, t, integrator) = copy(get_positions(u))

"Get the position of the ring polymer centroid"
centroid_position(u, t, integrator) = get_centroid(get_positions(u))

"Evaluate the potential from the model"
potential(u, t, integrator) = NQCModels.potential(integrator.p.calculator.model, get_positions(u))[1]

"Evaluate the classical Hamiltonian"
hamiltonian(u, t, integrator) = DynamicsUtils.classical_hamiltonian(integrator.p, u)

"Evaluate the classical kinetic energy"
kinetic(u, t, integrator) = DynamicsUtils.classical_kinetic_energy(integrator.p, get_velocities(u))

"Get all the dynamics variables. This is the default"
u(u, t, integrator) = copy(u)
u(u::ArrayPartition, t, integrator) = ComponentVector(v=copy(get_velocities(u)), r=copy(get_positions(u)))

"""
Get the quantum subsystem of the dynamics variables.
Requires that `DynamicsUtils.get_quantum_subsystem` is implemented for your chosen method.
"""
quantum_subsystem(u, t, integrator) = copy(get_quantum_subsystem(u))

mapping_position(u, t, integrator) = copy(DynamicsUtils.get_mapping_positions(u))
mapping_momentum(u, t, integrator) = copy(DynamicsUtils.get_mapping_momenta(u))

"""
Get the currently occupied state from the dynamics variables.
Requires that the dynamics variable has a field `state`.
Currently this is for surface hopping methods only.
Use [`population`](@ref) or [`adiabatic_population`](@ref) for most other methods.
"""
state(u, t, integrator) = copy(u.state)

"Get the noise along the stochastic trajectory"

noise(u, t, integrator) = copy(integrator.W.dW) / sqrt(integrator.dt)

"Evaluate the diabatic population"
population(u, t, integrator) = Estimators.diabatic_population(integrator.p, u)

"Evaluate the adiabatic population"
adiabatic_population(u, t, integrator) = Estimators.adiabatic_population(integrator.p, u)

"Evaluate the friction. This is used for MDEF only."
function friction(u, t, integrator)
    integrator.g(integrator.cache.gtmp,get_positions(u),integrator.p,t)
    copy(integrator.cache.gtmp)
end

end # module


using Statistics: mean

struct ReachableFitnessLogger <: Logger
    original_organisms_alive::Vector{UInt64}
    direct_children_alive::Vector{UInt64}
    direct_grandchildren_alive::Vector{UInt64}

    parents::Dict
    children::Vector

    function ReachableFitnessLogger(snapshot)
        new(get_organism_ids(snapshot), [], [], Dict(), [])
    end
end

function get_reachable_fitness(snapshot, rel_tolerance, min_samples, max_samples)
    reachable_fitness = estimate(rel_tolerance, min_samples, max_samples) do
        is_beneficial = []

        while true
            logger = ReachableFitnessLogger(snapshot)
            simulate_snapshot!(should_terminate, snapshot, logger)

            if length(logger.children) == 0 continue end

            is_beneficial = []

            for child in logger.children
                parent = logger.parents[child]

                fitness_child = get_fitness(snapshot, child)
                fitness_parent = get_fitness(snapshot, parent)

                push!(is_beneficial, fitness_parent < fitness_child ? 1 : 0)
            end

            if length(is_beneficial) == 0 continue end

            break
        end

        return rand(is_beneficial, 5)
    end

    return reachable_fitness
end

function should_terminate(logger::ReachableFitnessLogger, snapshot)
    isempty(logger.original_organisms_alive) &&
    isempty(logger.direct_children_alive) &&
    isempty(logger.direct_grandchildren_alive)
end

function log_step(::ReachableFitnessLogger, model) end
function save_log(::ReachableFitnessLogger) end

function log_birth(logger::ReachableFitnessLogger, model, child, parents=nothing)
    if parents === nothing return end

    parent = rand(parents)

    child_genotype = get_genotype(model, child)
    parent_genotype = get_genotype(model, parent)

    if child_genotype == parent_genotype return end

    child_id = get_id(model, child)
    parent_id = get_id(model, parent)

    if parent_id in logger.original_organisms_alive
        push!(logger.direct_children_alive, child_id)
        push!(logger.children, child)
        logger.parents[child] = parent
    elseif parent_id in logger.direct_children_alive
        push!(logger.direct_grandchildren_alive, child_id)
    end
end

function log_death(logger::ReachableFitnessLogger, model, organism)
    id = get_id(model, organism)

    if id in logger.original_organisms_alive
        delete!(logger.original_organisms_alive, id)
    elseif id in logger.direct_children_alive
        delete!(logger.direct_children_alive, id)
    elseif id in logger.direct_grandchildren_alive
        delete!(logger.direct_grandchildren_alive, id)
    end
end


####
#### Stan interface
####

# NOTE:
#   Type fields
#     fieldnames(CmdStan.Sample)
#       num_samples, num_warmup, save_warmup, thin, adapt, algorithm

#     fieldnames(CmdStan.Hmc)
#       engine, metric, stepsize, stepsize_jitter

#     fieldnames(CmdStan.Adapt)
#       engaged, gamma, delta, kappa, t0, init_buffer, term_buffer, window

#   Ref
#     http://goedman.github.io/Stan.jl/latest/index.html#Types-1

function sample(mf::T, ss::CmdStan.Sample) where T
    return sample(mf, ss.num_samples, ss.num_warmup,
                    ss.save_warmup, ss.thin, ss.adapt, ss.alg)
end

function sample(mf::T,
    num_samples::Int,
    num_warmup::Int,
    save_warmup::Bool,
    thin::Int,
    ss::CmdStan.Sample
) where T
    return sample(mf, num_samples, num_warmup, save_warmup, thin, ss.adapt, ss.alg)
end

function sample(mf::T,
    num_samples::Int,
    num_warmup::Int,
    save_warmup::Bool,
    thin::Int,
    adapt::CmdStan.Adapt,
    alg::CmdStan.Hmc
) where T
    if alg.stepsize_jitter != 0.0
        @warn("[Turing.sample] Turing does not support adding noise to stepsize yet.")
    end
    if adapt.engaged == false
        if isa(alg.engine, CmdStan.Static)   # hmc
            stepnum = Int(round(alg.engine.int_time / alg.stepsize))
            sample(mf, HMC(num_samples, alg.stepsize, stepnum); adaptor=NUTSAdaptor(adapt))
        elseif isa(alg.engine, CmdStan.Nuts) # error
            error("[Turing.sample] CmdStan.Nuts cannot be used with adapt.engaged set as false")
        end
    else
        if isa(alg.engine, CmdStan.Static)   # hmcda
            sample(mf, HMCDA(num_samples, num_warmup, adapt.delta, alg.engine.int_time);
                    adaptor=NUTSAdaptor(adapt))
        elseif isa(alg.engine, CmdStan.Nuts) # nuts
            if isa(alg.metric, CmdStan.diag_e)
                sample(mf, NUTS(num_samples, num_warmup, adapt.delta);
                        adaptor=NUTSAdaptor(adapt))
            else # TODO: reove the following since Turing support this feature now.
                @warn("[Turing.sample] Turing does not support full covariance matrix for pre-conditioning yet.")
            end
        end
    end
end

function AHMCAdaptor(adaptor::CmdAdaptorType) where CmdAdaptorType
    if :engaged in fieldnames(typeof(adaptor)) # CmdStan.Adapt
        adaptor.engaged ? spl.alg.n_adapts : 0,
        AHMC.Preconditioner(metric),
        AHMC.NesterovDualAveraging(adaptor.gamma,
            adaptor.t0, adaptor.kappa, adaptor.δ, init_ϵ),
            adaptor.init_buffer,
            adaptor.term_buffer,
            adaptor.window
    else # default adaptor
        @warn "Invalid adaptor type: $(typeof(adaptor)). Default adaptor is used instead."
        adaptor = AHMC.StanHMCAdaptor(
            spl.alg.n_adapts, AHMC.Preconditioner(:DiagEuclideanMetric),
            AHMC.NesterovDualAveraging(spl.alg.δ, init_ϵ)
        )
    end
end