# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

# du .= zero(eltype(du)) doesn't scale when using multiple threads.
# See https://github.com/trixi-framework/Trixi.jl/pull/924 for a performance comparison.
function reset_du!(du, dg, cache)
  @threaded for element in eachelement(dg, cache)
      du[.., element] .= zero(eltype(du))
  end

  return du
end

#     pure_and_blended_element_ids!(element_ids_dg, element_ids_dgfv, alpha, dg, cache)
#
# Given blending factors `alpha` and the solver `dg`, fill
# `element_ids_dg` with the IDs of elements using a pure DG scheme and
# `element_ids_dgfv` with the IDs of elements using a blended DG-FV scheme.
function pure_and_blended_element_ids!(element_ids_dg, element_ids_dgfv, alpha, dg::DG, cache)
  empty!(element_ids_dg)
  empty!(element_ids_dgfv)

  for element in eachelement(dg, cache)
    # Clip blending factor for values close to zero (-> pure DG)
    dg_only = isapprox(alpha[element], 0, atol=1e-12)
    if dg_only
      push!(element_ids_dg, element)
    else
      push!(element_ids_dgfv, element)
    end
  end

  return nothing
end

function volume_jacobian(element, mesh::TreeMesh, cache)
  return inv(cache.elements.inverse_jacobian[element])^ndims(mesh)
end

# Indicators used for shock-capturing and AMR
include("indicators.jl")
include("indicators_1d.jl")
include("indicators_2d.jl")
include("indicators_3d.jl")

# Container data structures
include("containers.jl")

# 1D DG implementation
include("dg_1d.jl")

# 2D DG implementation
include("dg_2d.jl")
include("dg_2d_parallel.jl")

# 3D DG implementation
include("dg_3d.jl")

# Auxiliary functions that are specialized on this solver
include("dg_2d_compressible_euler.jl")

end # @muladd


## Common:

### backward compat

function DecisionRule(exo_grid, endo_grid, values::Vector{Matrix{Float64}})
    N,n_x = size(values[1])
    vals = [to_LOP(v) for v in values]
    DecisionRule(exo_grid, endo_grid, vals)
end

function DecisionRule(exo_grid, endo_grid, values::Vector{Matrix{Float64}}, tt)
    N,n_x = size(values[1])
    vals = [reshape(reinterpret(Value{n_x}, vec(v')), (N,)) for v in values]
    DecisionRule(exo_grid, endo_grid, vals, tt)
end

# EmptyGrid x UCGrid

function evaluate(dr::AbstractDecisionRule{EmptyGrid{d1}, UCGrid{d}, n_x}, z::AbstractMatrix{Float64}) where n_x where d where d1
    N = size(z,1)
    @assert size(z,2)==d
    points = to_LOP(Point{d}, z)
    out = evaluate(dr,copy(points))
    return reshape(reinterpret(Float64, vec(out)), (n_x,N))'
end

evaluate(dr::AbstractDecisionRule{EmptyGrid{d1}, <:UCGrid}, z::Vector{Float64})  where d1 = vec(evaluate(dr,z'))
evaluate(dr::AbstractDecisionRule{EmptyGrid{d1}, <:UCGrid}, i::Int, y::Union{Vector,AbstractMatrix})  where d1 = evaluate(dr,y) # kind of nonsensical
evaluate(dr::AbstractDecisionRule{EmptyGrid{d1}, <:UCGrid}, x::Union{Vector,AbstractMatrix}, y::Union{Vector,AbstractMatrix})  where d1 = evaluate(dr,y)
UCGrid
function set_values!(dr::AbstractDecisionRule{EmptyGrid{d1}, <:UCGrid, n_x}, vals::Vector{Matrix{Float64}}) where n_x where d1
    N = size(vals[1],1)
    dims = tuple(dr.grid_endo.n...)
    V = [reshape(reinterpret(Value{n_x}, vec(vals[1]')), dims)]
    set_values!(dr, V)
end

# UCGrid x UCGrid

function evaluate(dr::AbstractDecisionRule{UCGrid{d1}, UCGrid{d2}, n_x}, x::AbstractMatrix, y::AbstractMatrix) where d1 where d2 where n_x
    N = size(x,1)
    @assert size(y,1)==N
    xx = to_LOP(Point{d1}, x)
    yy = to_LOP(Point{d2}, y)
    res = evaluate(dr,copy(xx),copy(yy))
    return reshape(reinterpret(Float64,vec(res)), (n_x, N))'
end
#
evaluate(dr::AbstractDecisionRule{<:UCGrid, <:UCGrid}, x::Vector{Float64}, y::Matrix{Float64}) = evaluate(dr, repeat(x', size(y,1)), y)
evaluate(dr::AbstractDecisionRule{<:UCGrid, <:UCGrid}, x::Vector{Float64}, y::Vector{Float64}) = evaluate(dr, vector_to_matrix(x), vector_to_matrix(y))
evaluate(dr::AbstractDecisionRule{UCGrid{d1}, UCGrid{d2}}, z::AbstractMatrix) where d1 where d2 = evaluate(dr, [z[:,1:d1] z[:,1:(d1+1:end)]])
evaluate(dr::AbstractDecisionRule{<:UCGrid, <:UCGrid}, z::Vector{Float64}) = vec(evaluate(dr,vector_to_matrix(z)))
evaluate(dr::AbstractDecisionRule{UCGrid{d1}, UCGrid{d2}}, i::Int, y::Union{Vector{Float64},AbstractMatrix{Float64}}) where d1 where d2  = evaluate(dr,node(dr.grid_exo,i), y)

function set_values!(dr::AbstractDecisionRule{<:UCGrid, <:UCGrid, n_x}, vals::Vector{Matrix{Float64}}) where n_x
    N = size(vals[1],1)
    dims = tuple(dr.grid_endo.n...)
    V = [reshape(reinterpret(Value{n_x}, vec(v')), dims) for v in vals]
    set_values!(dr, V)
end

# UnstructuredGrid x UCGrid

function evaluate(dr::AbstractDecisionRule{UnstructuredGrid{d1}, UCGrid{d2},n_x}, i::Int, z::AbstractMatrix{Float64}) where n_x where d1 where d2
    N = size(z,1)
    @assert size(z,2)==d2
    points = to_LOP(Point{d2}, z)
    # TODO: remove copy()
    out = evaluate(dr,i,copy(points))
    reshape(reinterpret(Float64, vec(out)), (n_x, N))'
end

evaluate(dr::AbstractDecisionRule{UnstructuredGrid{d1}, UCGrid{d2}, n_x}, i::Int, x::Vector{Float64}) where n_x where d1 where d2 = vec(evaluate(dr,i,vector_to_matrix(x)))
function evaluate(dr::AbstractDecisionRule{UnstructuredGrid{d1}, UCGrid{d2}, n_x}, inds::Vector{Int}, x::AbstractMatrix{Float64}) where n_x where d1 where d2
    l = [evaluate(dr,inds[n],x[n,:])' for n=1:size(x,1)]
    vcat(l...)
end

function evaluate(dr::AbstractDecisionRule{UnstructuredGrid{d1}, UCGrid{d2}, n_x}, inds::Vector{Int}, x::Vector{Float64}) where n_x where d1 where d2
    vcat([evaluate(dr,i,x)' for i in inds]...)
end

function set_values!(dr::AbstractDecisionRule{<:UnstructuredGrid, <:UCGrid, n_x}, vals::Vector{Matrix{Float64}}) where n_x
    N = size(vals[1],1)
    dims = tuple(dr.grid_endo.n...)
    V = [to_LOP(Value{n_x}, v) for v in vals]
    set_values!(dr, V)
end


"""
    DispersalKernel <: AbstractKernelNeighborhood

    DispersalKernel(; kw...)

Dispersal kernel for taking the dot product of the neighborhood and a matching
kernel of weights. May hold any `Neighborhood` object: the kernel
will be built to match the shape, using the `folumation`, `cellsize`
and `distancemethod`.

# Keyword Arguments

- `neighborhood`: Any DynamicGrids.jl `Neighborhood`, or an
    already constructed [`DispersalKernel`](@ref). Using this keyword means `radius` is
    ignored, and for a `DispersalKernel`, all other keywords are ignored.
- `neighborhood`: `Neighborhood` object specifying the range from the origin of the
    discretised dispersal kernal. Defaults to `Window(radius)`.
- `formulation`: kernel formulation object holding the exact form of the kernal.
    Default [`ExponentialKernel`](@ref).
- `cellsize`: the cell size of the discretised kernal (i.e. simulation grid size).
    Default is 1.0.
- `distancemethod`: [`DistanceMethod`](@ref) object for calculating distance between cells.
    The default is [`CentroidToCentroid`](@ref).
"""
struct DispersalKernel{R,L,N<:Neighborhood{R,L},K,F,C,D<:DistanceMethod} <: AbstractKernelNeighborhood{R,L}
    neighborhood::N
    kernel::K
    formulation::F
    cellsize::C
    distancemethod::D
    function DispersalKernel(
        hood::N, kernel, formulation::F, cellsize::C, distancemethod::D
    ) where {N<:Neighborhood{R,L},F,C,D<:DistanceMethod} where {R,L}
        if hood isa DispersalKernel
            hood
        else
            # Build the kernel matrix
            newkernel = scale(buildkernel(hood, formulation, distancemethod, cellsize))
            new{R,L,N,typeof(newkernel),F,C,D}(
                hood, newkernel, formulation, cellsize, distancemethod
            )
        end
    end
    function DispersalKernel{R,L,N,K,F,C,D}(
        hood::N, kernel::K, formulation::F, cellsize::C, distancemethod::D
    ) where {R,L,N,K,F,C,D}
        new{R,L,N,K,F,C,D}(hood, kernel, formulation, cellsize, distancemethod)
    end
end
function DispersalKernel(;
    radius=1,
    neighborhood=Window{radius}(),
    formulation=ExponentialKernel(),
    cellsize=1.0,
    distancemethod=CentroidToCentroid(),
)
    DispersalKernel(neighborhood, nothing, formulation, cellsize, distancemethod)
end
DispersalKernel{R}(; radius=R, kw...) where R = DispersalKernel(; radius=radius, kw...)

ConstructionBase.constructorof(::Type{<:DispersalKernel}) = DispersalKernel

function DG._setbuffer(n::DispersalKernel{R,L,<:Any,K,F,C,D}, buffer) where {R,L,K,F,C,D}
    newhood = DG._setbuffer(neighborhood(n), buffer)
    DispersalKernel{R,L,typeof(newhood),K,F,C,D}(
        newhood, kernel(n), formulation(n), cellsize(n), distancemethod(n)
    )
end

cellsize(hood::DispersalKernel) = hood.cellsize
distancemethod(hood::DispersalKernel) = hood.distancemethod
formulation(hood::DispersalKernel) = hood.formulation

function buildkernel(window::Window{R}, formulation, distancemethod, cellsize) where R
    # The radius doesn't include the center cell, so add it
    S = 2R + 1
    kernel = zeros(typeof(cellsize), S, S)
    r1 = R + one(R)
    # Paper: l. 97
    for x = 0:R, y = 0:R
        # Calculate the distance effect from the center cell to this cell
        prob = dispersalprob(formulation, distancemethod, x, y, cellsize)
        # Update the kernel value based on the formulation and distance
        kernel[ x + r1,  y + r1] = prob
        kernel[ x + r1, -y + r1] = prob
        kernel[-x + r1,  y + r1] = prob
        kernel[-x + r1, -y + r1] = prob
    end
    SMatrix{S,S}(kernel)
end
function buildkernel(window::Neighborhood{<:Any,L}, f, dm, cellsize) where L
    SVector{L}(Tuple(dispersalprob(f, dm, x, y, cellsize) for (x, y) in offsets(window)))
end

scale(x) = x ./ sum(x)


using JuMP; using NCDatasets; using AxisArrays; using CalliopeJuMP.Util;

function load_dispatch_constraints(model_dict)

    sets = model_dict["sets"]
    constraint_dict = Dict()

    if haskey(sets, "loc_tech_carriers_carrier_production_max_constraint")
        constraint_dict["carrier_production_max_constraint"] = (
            build_constraint(
                ["loc_tech_carriers_carrier_production_max_constraint", "timesteps"],
                "carrier_production_max", model_dict
            )
        )
    end

    if haskey(sets, "loc_tech_carriers_carrier_production_min_constraint")
        constraint_dict["carrier_production_min_constraint"] = (
            build_constraint(
                ["loc_tech_carriers_carrier_production_min_constraint", "timesteps"],
                "carrier_production_min", model_dict
            )
        )
    end

    if haskey(sets, "loc_tech_carriers_carrier_consumption_max_constraint")
        constraint_dict["carrier_consumption_max_constraint"] = (
            build_constraint(
                ["loc_tech_carriers_carrier_consumption_max_constraint", "timesteps"],
                "carrier_consumption_max", model_dict
            )
        )
    end

    if haskey(sets, "loc_techs_resource_max_constraint")
        constraint_dict["resource_max_constraint"] = (
            build_constraint(
                ["loc_techs_resource_max_constraint", "timesteps"],
                "resource_max", model_dict
            )
        )
    end

    if haskey(sets, "loc_techs_storage_max_constraint")
        constraint_dict["storage_max_constraint"] = (
            build_constraint(
                ["loc_techs_storage_max_constraint", "timesteps"],
                "storage_max", model_dict
            )
        )
    end

    return constraint_dict
end

function carrier_production_max_constraint_rule(backend_model, set_indices, model_dict)
    """
    Set maximum carrier production. All technologies.
    """
    loc_tech_carrier, timestep = set_indices
    variables = model_dict["variables"]
    loc_tech = get_loc_tech(loc_tech_carrier)
    carrier_prod = variables["carrier_prod"][loc_tech_carrier, timestep]
    timestep_resolution = model_dict["parameters"]["timestep_resolution"][timestep]
    parasitic_eff = get_param(model_dict, "parasitic_eff", [loc_tech], timestep)

    return @constraint(backend_model,
        carrier_prod <= variables["energy_cap"][loc_tech] * timestep_resolution * parasitic_eff
    )
end

function carrier_production_min_constraint_rule(backend_model, set_indices, model_dict)
    """
    Set minimum carrier production. All technologies except conversion_plus
    """
    loc_tech_carrier, timestep = set_indices
    variables = model_dict["variables"]
    loc_tech = get_loc_tech(loc_tech_carrier)
    carrier_prod = variables["carrier_prod"][loc_tech_carrier, timestep]
    timestep_resolution = model_dict["parameters"]["timestep_resolution"][timestep]
    min_use = get_param(model_dict, "energy_cap_min_use", [loc_tech], timestep)

    return @constraint(backend_model, (
        carrier_prod >= variables["energy_cap"][loc_tech] * timestep_resolution * min_use
    ))
end

function carrier_consumption_max_constraint_rule(backend_model, set_indices, model_dict)
    """
    Set maximum carrier consumption for demand, storage, and transmission techs
    """
    loc_tech_carrier, timestep = set_indices
    variables = model_dict["variables"]
    loc_tech = get_loc_tech(loc_tech_carrier)
    carrier_con = variables["carrier_con"][loc_tech_carrier, timestep]
    timestep_resolution = model_dict["parameters"]["timestep_resolution"][timestep]

    return @constraint(backend_model, (
        carrier_con >= -1 * variables["energy_cap"][loc_tech] * timestep_resolution
    ))
end

function resource_max_constraint_rule(backend_model, set_indices, model_dict)
    """
    Set maximum resource consumed by supply_plus techs
    """
    loc_tech, timestep = set_indices
    variables = model_dict["variables"]
    timestep_resolution = model_dict["parameters"]["timestep_resolution"][timestep]

    return @constraint(backend_model, (
        variables["resource_con"][loc_tech, timestep] <=
        timestep_resolution * variables["resource_cap"][loc_tech]
    ))
end

function storage_max_constraint_rule(backend_model, set_indices, model_dict)
    """
    Set maximum stored energy. Supply_plus & storage techs only.
    """
    loc_tech, timestep = set_indices
    variables = model_dict["variables"]

    return @constraint(backend_model, (
        variables["storage"][loc_tech, timestep] <= variables["storage_cap"][loc_tech]
    ))
end


# Compile speed tests
# Tests time to get from nothing to certain tasks
using JuMP  # Ensure compiled before running the tests

# Test 1
# Size 100, aff_str into model_str
test1_time = @elapsed begin
    run(`$(ENV["_"]) -e """
    using JuMP
    m = Model()
    N = 100
    @variable(m, x[1:N])
    @constraint(m, sum(i*x[i] for i=1:N) >= N)
    @time JuMP.aff_str(JuMP.REPLMode, m.linconstr[end].terms)
    @time JuMP.model_str(JuMP.REPLMode, m)"""
    `)
end

# Test 2
# Size 100, model_str
test2_time = @elapsed begin
    run(`$(ENV["_"]) -e """
    using JuMP
    m = Model()
    N = 100
    @variable(m, x[1:N])
    @constraint(m, sum(i*x[i] for i=1:N) >= N)
    @time JuMP.model_str(JuMP.REPLMode, m)"""
    `)
end

# Test 3
# Size 10,000, print model
test3_time = @elapsed begin
    run(`$(ENV["_"]) -e """
    using JuMP
    m = Model()
    N = 10000
    @variable(m, x[1:N])
    @constraint(m, sum(i*x[i] for i=1:N) >= N)
    @time sprint(print, m)"""
    `)
end

# Test 4
# Size 10, 4 variables, print model
test4_time = @elapsed begin
    run(`$(ENV["_"]) -e """
    using JuMP
    m = Model()
    N = 10
    @variable(m, x1[1:N])
    @variable(m, x2[1:N,f=1:N])
    @variable(m, x3[1:N,f=1:2:N])
    @variable(m, x4[[:a,:b,:c]])
    @variable(m, x5[[:a,:b,:c],[:d,"e",4]])
    @constraint(m,
        sum(i*x1[i] for i=1:N) +
        sum(i*f*x2[i,f] for i=1:N,f=1:N) +
        sum(i*f*x3[i,f] for i=1:N,f=1:2:N) +
        sum(x4) >= N)
    @time sprint(print, m)"""
    `)
end

println("Julia ", VERSION)
println(
    "Current branch: $(chomp(read(`git rev-parse --abbrev-ref HEAD`, String)))",
)
println("Current commit: $(chomp(read(`git rev-parse HEAD`, String)))")
@printf("Test 1: %7.2f s\n", test1_time)
@printf("Test 2: %7.2f s\n", test2_time)
@printf("Test 3: %7.2f s\n", test3_time)
@printf("Test 4: %7.2f s\n", test4_time)


# # Dry Atmosphere GCM diagnostics
#
# This file computes selected diagnostics for the GCM and outputs them
# on the spherical interpolated diagnostic grid.
#
# Use it by calling `Diagnostics.setup_atmos_default_diagnostics()`.
#
# TODO:
# - enable zonal means and calculation of covariances using those means
#   - ds.T_zm = mean(.*1., ds.T; dims = 3)
#   - ds.u_zm = mean((ds.u); dims = 3 )
#   - v_zm = mean(ds.v; dims = 3)
#   - w_zm = mean(ds.w; dims = 3)
#   - ds.uvcovariance = (ds.u .- ds.u_zm) * (ds.v .- v_zm)
#   - ds.vTcovariance = (ds.v .- v_zm) * (ds.T .- ds.T_zm)
# - add more variables, including horiz streamfunction from laplacial of vorticity (LN)
# - density weighting
# - maybe change thermo/dyn separation to local/nonlocal vars?

import CUDA
using LinearAlgebra
using Printf
using Statistics

using ..Atmos
using ..Atmos: recover_thermo_state
using ..TurbulenceClosures: turbulence_tensors

"""
    setup_atmos_default_diagnostics(
        ::AtmosGCMConfigType,
        interval::String,
        out_prefix::String;
        writer::AbstractWriter,
        interpol = nothing,
    )

Create and return a `DiagnosticsGroup` containing the "AtmosDefault"
diagnostics for GCM configurations. All the diagnostics in the group will run
at the specified `interval`, be interpolated to the specified boundaries and
resolution, and will be written to files prefixed by `out_prefix` using
`writer`.
"""
function setup_atmos_default_diagnostics(
    ::AtmosGCMConfigType,
    interval::String,
    out_prefix::String;
    writer = NetCDFWriter(),
    interpol = nothing,
)
    # TODO: remove this
    @assert !isnothing(interpol)

    return DiagnosticsGroup(
        "AtmosGCMDefault",
        Diagnostics.atmos_gcm_default_init,
        Diagnostics.atmos_gcm_default_fini,
        Diagnostics.atmos_gcm_default_collect,
        interval,
        out_prefix,
        writer,
        interpol,
    )
end

include("diagnostic_fields.jl")

# 3D variables
function vars_atmos_gcm_default_simple_3d(atmos::AtmosModel, FT)
    @vars begin
        u::FT
        v::FT
        w::FT
        rho::FT
        temp::FT
        pres::FT
        thd::FT                 # θ_dry
        et::FT                  # e_tot
        ei::FT                  # e_int
        ht::FT
        hi::FT
        vort::FT                # Ω₃

        moisture::vars_atmos_gcm_default_simple_3d(atmos.moisture, FT)
    end
end
vars_atmos_gcm_default_simple_3d(::MoistureModel, FT) = @vars()
function vars_atmos_gcm_default_simple_3d(m::EquilMoist, FT)
    @vars begin
        qt::FT                  # q_tot
        ql::FT                  # q_liq
        qv::FT                  # q_vap
        qi::FT                  # q_ice
        thv::FT                 # θ_vir
        thl::FT                 # θ_liq

    end
end
num_atmos_gcm_default_simple_3d_vars(m, FT) =
    varsize(vars_atmos_gcm_default_simple_3d(m, FT))
atmos_gcm_default_simple_3d_vars(m, array) =
    Vars{vars_atmos_gcm_default_simple_3d(m, eltype(array))}(array)

function atmos_gcm_default_simple_3d_vars!(
    atmos::AtmosModel,
    state_prognostic,
    thermo,
    dyni,
    vars,
)
    vars.u = state_prognostic.ρu[1] / state_prognostic.ρ
    vars.v = state_prognostic.ρu[2] / state_prognostic.ρ
    vars.w = state_prognostic.ρu[3] / state_prognostic.ρ
    vars.rho = state_prognostic.ρ
    vars.temp = thermo.temp
    vars.pres = thermo.pres
    vars.thd = thermo.θ_dry
    vars.et = state_prognostic.ρe / state_prognostic.ρ
    vars.ei = thermo.e_int
    vars.ht = thermo.h_tot
    vars.hi = thermo.h_int

    vars.vort = dyni.Ω₃

    atmos_gcm_default_simple_3d_vars!(
        atmos.moisture,
        state_prognostic,
        thermo,
        vars,
    )

    return nothing
end
function atmos_gcm_default_simple_3d_vars!(
    ::MoistureModel,
    state_prognostic,
    thermo,
    vars,
)
    return nothing
end
function atmos_gcm_default_simple_3d_vars!(
    moist::EquilMoist,
    state_prognostic,
    thermo,
    vars,
)
    vars.moisture.qt = state_prognostic.moisture.ρq_tot / state_prognostic.ρ
    vars.moisture.ql = thermo.moisture.q_liq
    vars.moisture.qv = thermo.moisture.q_vap
    vars.moisture.qi = thermo.moisture.q_ice
    vars.moisture.thv = thermo.moisture.θ_vir
    vars.moisture.thl = thermo.moisture.θ_liq_ice

    return nothing
end

# Dynamic variables
function vars_dyn(FT)
    @vars begin
        Ω₁::FT
        Ω₂::FT
        Ω₃::FT
    end
end
dyn_vars(array) = Vars{vars_dyn(eltype(array))}(array)

"""
    atmos_gcm_default_init(dgngrp, currtime)

Initialize the GCM default diagnostics group, establishing the output file's
dimensions and variables.
"""
function atmos_gcm_default_init(dgngrp::DiagnosticsGroup, currtime)
    atmos = Settings.dg.balance_law
    FT = eltype(Settings.Q)
    mpicomm = Settings.mpicomm
    mpirank = MPI.Comm_rank(mpicomm)

    if !(dgngrp.interpol isa InterpolationCubedSphere)
        @warn """
            Diagnostics ($dgngrp.name): currently requires `InterpolationCubedSphere`!
            """
        return nothing
    end

    if mpirank == 0
        # get dimensions for the interpolated grid
        dims = dimensions(dgngrp.interpol)

        # adjust the level dimension for `planet_radius`
        level_val = dims["level"]
        dims["level"] = (
            level_val[1] .- FT(planet_radius(Settings.param_set)),
            level_val[2],
        )

        # set up the variables we're going to be writing
        vars = OrderedDict()
        varnames = map(
            s -> startswith(s, "moisture.") ? s[10:end] : s,
            flattenednames(vars_atmos_gcm_default_simple_3d(atmos, FT)),
        )
        for varname in varnames
            var = Variables[varname]
            vars[varname] = (tuple(collect(keys(dims))...), FT, var.attrib)
        end

        # create the output file
        dprefix = @sprintf(
            "%s_%s_%s",
            dgngrp.out_prefix,
            dgngrp.name,
            Settings.starttime,
        )
        dfilename = joinpath(Settings.output_dir, dprefix)
        init_data(dgngrp.writer, dfilename, dims, vars)
    end

    return nothing
end

"""
    atmos_gcm_default_collect(bl, currtime)

    Master function that performs a global grid traversal to compute various
    diagnostics using the above functions.
"""
function atmos_gcm_default_collect(dgngrp::DiagnosticsGroup, currtime)
    interpol = dgngrp.interpol
    if !(interpol isa InterpolationCubedSphere)
        @warn """
            Diagnostics ($dgngrp.name): currently requires `InterpolationCubedSphere`!
            """
        return nothing
    end

    dg = Settings.dg
    atmos = dg.balance_law
    Q = Settings.Q
    mpicomm = Settings.mpicomm
    mpirank = MPI.Comm_rank(mpicomm)
    grid = dg.grid
    topology = grid.topology
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dimensionality(grid) == 2 ? 1 : Nq
    npoints = Nq * Nq * Nqk
    nrealelem = length(topology.realelems)
    nvertelem = topology.stacksize
    nhorzelem = div(nrealelem, nvertelem)

    # get needed arrays onto the CPU
    device = array_device(Q)
    if device isa CPU
        ArrayType = Array
        state_data = Q.realdata
        aux_data = dg.state_auxiliary.realdata
    else
        ArrayType = CUDA.CuArray
        state_data = Array(Q.realdata)
        aux_data = Array(dg.state_auxiliary.realdata)
    end
    FT = eltype(state_data)

    # TODO: can this be done in one pass?
    #
    # Non-local vars, e.g. relative vorticity
    vgrad = VectorGradients(dg, Q)
    vort = Vorticity(dg, vgrad)

    # Compute thermo variables
    thermo_array = Array{FT}(undef, npoints, num_thermo(atmos, FT), nrealelem)
    @visitQ nhorzelem nvertelem Nqk Nq begin
        state = extract_state(dg, state_data, ijk, e, Prognostic())
        aux = extract_state(dg, aux_data, ijk, e, Auxiliary())

        thermo = thermo_vars(atmos, view(thermo_array, ijk, :, e))
        compute_thermo!(atmos, state, aux, thermo)
    end

    # Interpolate the state, thermo and dyn vars to sphere (u and vorticity
    # need projection to zonal, merid). All this may happen on the GPU.
    istate =
        ArrayType{FT}(undef, interpol.Npl, number_states(atmos, Prognostic()))
    interpolate_local!(interpol, Q.realdata, istate)

    ithermo = ArrayType{FT}(undef, interpol.Npl, num_thermo(atmos, FT))
    interpolate_local!(interpol, ArrayType(thermo_array), ithermo)

    idyn = ArrayType{FT}(undef, interpol.Npl, size(vort.data, 2))
    interpolate_local!(interpol, vort.data, idyn)

    # TODO: get indices here without hard-coding them
    _ρu, _ρv, _ρw = 2, 3, 4
    project_cubed_sphere!(interpol, istate, (_ρu, _ρv, _ρw))
    _Ω₁, _Ω₂, _Ω₃ = 1, 2, 3
    project_cubed_sphere!(interpol, idyn, (_Ω₁, _Ω₂, _Ω₃))

    # FIXME: accumulating to rank 0 is not scalable
    all_state_data = accumulate_interpolated_data(mpicomm, interpol, istate)
    all_thermo_data = accumulate_interpolated_data(mpicomm, interpol, ithermo)
    all_dyn_data = accumulate_interpolated_data(mpicomm, interpol, idyn)

    if mpirank == 0
        # get dimensions for the interpolated grid
        dims = dimensions(dgngrp.interpol)

        # set up the array for the diagnostic variables based on the interpolated grid
        nlong = length(dims["long"][1])
        nlat = length(dims["lat"][1])
        nlevel = length(dims["level"][1])

        simple_3d_vars_array = Array{FT}(
            undef,
            nlong,
            nlat,
            nlevel,
            num_atmos_gcm_default_simple_3d_vars(atmos, FT),
        )

        @visitI nlong nlat nlevel begin
            statei = Vars{vars_state(atmos, Prognostic(), FT)}(view(
                all_state_data,
                lo,
                la,
                le,
                :,
            ))
            thermoi = thermo_vars(atmos, view(all_thermo_data, lo, la, le, :))
            dyni = dyn_vars(view(all_dyn_data, lo, la, le, :))
            simple_3d_vars = atmos_gcm_default_simple_3d_vars(
                atmos,
                view(simple_3d_vars_array, lo, la, le, :),
            )

            atmos_gcm_default_simple_3d_vars!(
                atmos,
                statei,
                thermoi,
                dyni,
                simple_3d_vars,
            )
        end

        # assemble the diagnostics for writing
        varvals = OrderedDict()
        varnames = map(
            s -> startswith(s, "moisture.") ? s[10:end] : s,
            flattenednames(vars_atmos_gcm_default_simple_3d(atmos, FT)),
        )
        for (vari, varname) in enumerate(varnames)
            varvals[varname] = simple_3d_vars_array[:, :, :, vari]
        end

        # write output
        append_data(dgngrp.writer, varvals, currtime)
    end

    MPI.Barrier(mpicomm)
    return nothing
end # function collect

function atmos_gcm_default_fini(dgngrp::DiagnosticsGroup, currtime) end


############# data preparation ###################

function parsetitle(title, r)
    run = r.data
    d = Dict{Symbol, Any}(k => missing for k in (:displace_location, :displace_direction , :nest_coverage, :transfer))
    d[:nest_coverage] = "open"
    d[:nest] = run.originalnest
    d[:feeder] = run.feeder
    d[:fictive_nest] = run.nest
    d[:track] = run.track
    d[:title] = title
    d[:comment] = r.metadata.comment
    for kv in split(title, ' ')
        k, v = split(kv, '#')
        if k ∉ ("person", "pellet")
            if k == "nest"
                d[:nest_coverage] = v
            else
                d[Symbol(k)] = v
            end
        end
    end
    return (; pairs(d)...)
end

function getdf(data)
    data["nest#closed person#therese displace_direction#none displace_location#feeder"] = data["nest#closed person#therese"]
    delete!(data, "nest#closed person#therese")
    df = DataFrame(parsetitle(k, r) for (k, v) in data for r in v.runs)

    @. df[!, :displace_direction] = switchdirections.(df.displace_direction)
    @. df[!, :group] .= getgroup(df.displace_location, df.transfer, df.displace_direction)
    @. df[!, :set] = getset(df.transfer, df.group)

    categorical!(df, [:group, :set, :displace_direction, :displace_location, :nest_coverage, :transfer])
    levels!(df.group, ["none", "left", "right", "away", "towards", "zero", "back", "far"])

    filter!(r -> r.group ≠ "far" || r.title == "transfer#far person#therese", df)

    df[!, :direction_deviation]  = [angle(r.fictive_nest - r.feeder, turningpoint(r.track) - r.feeder) for r in eachrow(df)]
    max_direction_deviation = maximum(r.direction_deviation for r in eachrow(df) if r.group ∉ ("far", "zero"))
    mean_direction_deviation = mean(r.direction_deviation for r in eachrow(df) if r.group ∉ ("far", "zero"))
    filter!(r -> r.group ≠ "far" || r.direction_deviation < 4mean_direction_deviation, df)

    df[!, :turning_point] .= zero.(df.feeder)
    df[!, :center_of_search] .= zero.(df.feeder)
    for r in eachrow(df)
        trans = createtrans(r.nest, r.displace_location, r.fictive_nest, r.feeder)
        @. r.track.coords = trans(r.track.coords)
        @. r.track.rawcoords.xy .= trans(r.track.rawcoords.xy)
        r.feeder = trans(r.feeder)
        r.fictive_nest = trans(r.fictive_nest)
        r.nest = trans(r.nest)
        Δ = intended(r.group) - r.fictive_nest
        r.turning_point = turningpoint(r.track) + Δ
        r.center_of_search = searchcenter(r.track) + Δ
    end

    groups = levels(df.group)
    nc = length(groups)
    colors = OrderedDict(zip(groups, [colorant"black"; distinguishable_colors(nc - 1, [colorant"white", colorant"black"], dropseed = true)]))

    gdf = groupby(df, :group)
    DataFrames.transform!(gdf, :group => (g -> colors[g[1]]) => :groupcolor)
    DataFrames.transform!(gdf, :groupcolor => getcolor => :color)

    df
end

switchdirections(_::Missing) = missing
switchdirections(d) =   d == "left" ? "right" :
                        d == "right" ? "left" :
                        d

getgroup(displace_location::Missing, transfer, displace_direction) = transfer
getgroup(displace_location, transfer, displace_direction) = displace_location == "nest" ? "zero" : displace_direction
getset(_::Missing, d) = d == "none" ? "Closed" : "Displacement"
getset(_, __) = "Transfer"

intended(d::AbstractString) =  d == "none" ? DungBase.Point(0,0) :
                d == "away" ? DungBase.Point(0, -50) :
                d == "towards" ? DungBase.Point(0, 50) :
                d == "right" ? DungBase.Point(50, 0) :
                d == "left" ? DungBase.Point(-50, 0) :
                d == "zero" ? DungBase.Point(0, -130) :
                d == "back" ? DungBase.Point(0,0) :
                d == "far" ? DungBase.Point(0,0) :
                error("unknown displacement")
intended(d) = intended(string(d))

_get_rotation_center(displace_location::Missing, nest, fictive_nest) = fictive_nest
_get_rotation_center(displace_location, nest, fictive_nest) = displace_location == "feeder" ? fictive_nest : nest
_get_zeroing(nest::Missing, fictive_nest) = fictive_nest
_get_zeroing(nest, fictive_nest) = nest
function createtrans(nest, displace_location, fictive_nest, feeder)
    v = feeder - _get_rotation_center(displace_location, nest, fictive_nest)
    α = atan(v[2], v[1])
    rot = LinearMap(Angle2d(-π/2 - α))
    trans = Translation(-_get_zeroing(nest, fictive_nest))
    passmissing(rot ∘ trans)
end

function highlight(c, i, n)
    h = HSL(c)
    HSL(h.h, h.s, i/(n + 1))
end
function getcolor(g)
    n = length(g)
    [highlight(c, i, n) for (i, c) in enumerate(g)]
end


for (f, T) in Base.Iterators.product((:add, :mul, :sub), (Float32, Float64))
    name = Symbol("$(f)_float_contract")
    if T === Float32
        llvmt = "float"
    elseif T === Float64
        llvmt = "double"
    end

    #XXX Use LLVM.jl
    ir = """
        %x = f$f contract $llvmt %0, %1
        ret $llvmt %x
    """
    @eval begin
        # the @pure is necessary so that we can constant propagate.
        Base.@pure function $name(a::$T, b::$T)
            @Base._inline_meta
            Base.llvmcall($ir, $T, Tuple{$T, $T}, a, b)
        end
    end
end


using JSON

#Constants
const e_charge = 1.602176453E-09
const e_mass = 9.1093826E-31
const c_0 = 299792458
const ε_0 = 8.85419E-12
const μ_0 = 4E-07

#Reading the dielectrics files that defines de problem
dielectrics = JSON.Parser.parsefile(".\\src\\Dielectrics.json")

#Reading the config files that defines de problem
scheme_prob = JSON.Parser.parsefile(".\\src\\Scheme.json")

for i in range(1, 20, step=1)
    println(i)
end


type Model
    ws::Vector
end

function (m::Model)(s::State)
    tokens = s.tokens
    s0 = tokens[s.top]
    s1 = isnull(s.left) ? tokens[s.left.top] : nulltoken
    b0 = s.right <= length(tokens) ? tokens[s.right] : nulltoken
    sl = isnull(s.lch) ? tokens[s.lch.top] : nulltoken
    sr = isnull(s.rch) ? tokens[s.rch.top] : nulltoken

end

function train()
    opt = SGD()
    for epoch = 1:nepochs
        println("epoch: $(epoch)")
        opt.rate = 0.0075 / epoch
        loss = fit(train_x, train_y, model, crossentropy, opt)
        println("loss: $(loss)")

        test_z = map(x -> predict(model,x), test_x)
        acc = accuracy(test_y, test_z)
        println("test acc.: $(acc)")
        println("")
    end
end


### Implementation via recursion ###
# recursion approach
function decode(m::Tacotron₂, values::T, keys::T, stepsleft::Integer) where T <: DenseArray{<:Real,3}
   #initialize dimensions
   _, time_in, batchsize = size(values)
   querydim = size(m.attention.dense.W, 2)
   nmelfeatures = length(m.frameproj.b)
   # initialize parameters
   query   = gpu(zeros(Float32, querydim, batchsize))
   weights = gpu(zeros(Float32, time_in, 1, batchsize))
   frame   = gpu(zeros(Float32, nmelfeatures, batchsize))

   context, weights = m.attention(values, keys, query, weights, weights)
   prenetoutput = m.prenet(frame)
   prenetoutput_context = [prenetoutput; context]
   query = m.lstms(prenetoutput_context)
   query_context = [query; context]
   frame = m.frameproj(query_context)
   decoding₁ = (frame, query_context)
   decode(m, values, keys, stepsleft-1, query, weights, weights, frame, decoding₁)
end

function decode(m::Tacotron₂, values::T, keys::T, stepsleft::Integer, query::M, weights::T, Σweights::T, frame::M, decodings) where {M <: DenseMatrix, T <: DenseArray{<:Real,3}}
   (stepsleft == 0) && (return decodings)
   context, weights = m.attention(values, keys, query, weights, Σweights)
   Σweights += weights
   prenetoutput = m.prenet(frame)
   prenetoutput_context = [prenetoutput; context]
   query = m.lstms(prenetoutput_context)
   query_context = [query; context]
   frame = m.frameproj(query_context)
   decoding = (frame, query_context)
   decodings = [decodings; decoding]
   decode(m, values, keys, stepsleft-1, query, weights, Σweights, frame, decodings)
end

###
function (m::Tacotron₂)(textindices::DenseMatrix{<:Integer}, time_out::Integer)
   # dimensions
   batchsize = size(textindices, 2)
   nmelfeatures = length(m.frameproj.b)
   # encoding stage
   chex = m.che(textindices)
   convblock₃x = m.convblock₃(chex)
   values = m.blstm(convblock₃x)
   # @ein keys[a,t,b] := m.attention.V[a,d] * values[d,t,b] # dispatches to batched_contract (vetted)
   keys = einsum(EinCode{((1,2), (2,3,4)), (1,3,4)}(), (m.attention.V, values))
   # #=check=# Vh ≈ reduce(hcat, [reshape(m.V * values[:,t,:], size(m.V,1), 1, :) for t ∈ axes(values,2)])
   decodings = decode(m, values, keys, time_out)
   frames = first.(decodings)
   query_contexts = last.(decodings)
   σ⁻¹stoprobs = reshape(m.stopproj(reduce(hcat, query_contexts)), (batchsize, time_out))
   # #=check=# σ⁻¹stoprobs == reduce(hcat, reshape.(m.stopproj.(query_contexts), Val(1)))
   prediction = reshape(reduce(hcat, frames), (nmelfeatures, batchsize, time_out))
   # #=check=# prediction == cat(frames...; dims=3)
   melprediction = permutedims(prediction, (3,1,2))
   melprediction⁺residual = melprediction + m.postnet(melprediction)
   return melprediction, melprediction⁺residual, σ⁻¹stoprobs
end

### Implementation via iterator ###
# Warning: adjoint of the `collect(x::Decodings)` defined below is likely wrong
struct Decodings{T, M}
   m :: M
   values   :: T
   keys     :: T
   time_out :: Int
end

@adjoint function Decodings(m::Tacotron2, values::T, keys::T, time_out::Int) where T <: DenseArray{<:Real,3}
   return Decodings(m, values, keys, time_out), function (d̄)
      return d̄.m, d̄.values, d̄.keys, nothing
   end
end

# this adjoint is seems wrong
@adjoint function collect(x::Decodings)
   collect(x), function(ȳ)
      return (Decodings(first(ȳ)[3:end]..., length(ȳ)),)
   end
end

Flux.trainable(d::Decodings) = (d.m,)
@functor Decodings

Base.eltype(::Type{<:Decodings{T,M}}) where {T <: DenseArray{<:Real,3}, M <: Tacotron2} = Tuple{matrixof(T), vectorof(T), M, T, T}
Base.length(itr::Decodings) = itr.time_out

function Base.iterate(itr::Decodings)
   (itr.time_out <= 0) && (return nothing)
   #initialize dimensions
   _, time_in, batchsize = size(itr.values)
   querydim = size(m.attention.dense.W, 2)
   nmelfeatures = length(m.frameproj.b)
   # initialize parameters
   query   = gpu(zeros(Float32, querydim, batchsize))
   weights = gpu(zeros(Float32, time_in, 1, batchsize))
   frame   = gpu(zeros(Float32, nmelfeatures, batchsize))

   context, weights = itr.m.attention(itr.values, itr.keys, query, weights, weights)
   prenetoutput = itr.m.prenet(frame)
   prenetoutput_context = [prenetoutput; context]
   query = itr.m.lstms(prenetoutput_context)
   query_context = [query; context]
   frame = itr.m.frameproj(query_context)
   σ⁻¹pstopᵀ = itr.m.stopproj(query_context)
   σ⁻¹pstop = reshape(σ⁻¹pstopᵀ, Val(1))
   i = (frame, σ⁻¹pstop, m, values, keys)
   state = (1, query, weights, weights, frame)
   return i, state
end

function Base.iterate(itr::Decodings, state::Tuple{Int,M,T,T,M}) where {M <: DenseMatrix, T <: DenseArray{<:Real,3}}
   (t, query, weights, Σweights, frame) = state
   (t == itr.time_out) && (return nothing)
   t += 1
   context, weights = itr.m.attention(itr.values, itr.keys, query, weights, Σweights)
   Σweights += weights
   prenetoutput = itr.m.prenet(frame)
   prenetoutput_context = [prenetoutput; context]
   query = itr.m.lstms(prenetoutput_context)
   query_context = [query; context]
   frame = itr.m.frameproj(query_context)
   σ⁻¹pstopᵀ = itr.m.stopproj(query_context)
   σ⁻¹pstop = reshape(σ⁻¹pstopᵀ, Val(1))
   i = (frame, σ⁻¹pstop, m, values, keys)
   state = (t, query, weights, Σweights, frame)
   return i, state
end

function (m::Tacotron2)(textindices::DenseMatrix{<:Integer}, time_out::Integer)
   # dimensions
   batchsize = size(textindices, 2)
   nmelfeatures = length(m.frameproj.b)
   # encoding stage
   chex = m.che(textindices)
   convblock₃x = m.convblock₃(chex)
   values = m.blstm(convblock₃x)
   # @ein keys[a,t,b] := m.attention.V[a,d] * values[d,t,b] # dispatches to batched_contract (vetted)
   keys = einsum(EinCode{((1,2), (2,3,4)), (1,3,4)}(), (m.attention.V, values))
   # #=check=# Vh ≈ reduce(hcat, [reshape(m.V * values[:,t,:], size(m.V,1), 1, :) for t ∈ axes(values,2)])
   # decodings′ = Zygote.bufferfrom(Vector{Tuple{Matrix{Float32},Vector{Float32}}}(undef, time_out))
   # decodings′ .= Decodings(m, values, keys, time_out)
   # decodings = copy(decodings′)
   # decodings = [decoding for decoding ∈ decodings′]
   decodings = collect(Decodings(m, values, keys, time_out))
   frames = first.(decodings)
   σ⁻¹pstops = getindex.(decodings, 2)
   σ⁻¹stoprobs = reduce(hcat, σ⁻¹pstops)
   prediction = reshape(reduce(hcat, frames), nmelfeatures, batchsize, time_out)
   # #=check=# prediction == cat(frames...; dims=3)
   melprediction = permutedims(prediction, (3,1,2))
   melprediction⁺residual = melprediction + m.postnet(melprediction)
   return melprediction, melprediction⁺residual, σ⁻¹stoprobs
end


### indicator data -> string
########################################

vec2string(p::StaticVector{2}) = @sprintf("(%0.3f, %0.3f)", p[1], p[2])
vec2string(p::StaticVector{3}) = @sprintf("(%0.3f, %0.3f, %0.3f)", p[1], p[2], p[3])

position2string(p::StaticVector{2}) = @sprintf("x: %0.6f\ny: %0.6f", p[1], p[2])
position2string(p::StaticVector{3}) = @sprintf("x: %0.6f\ny: %0.6f\nz: %0.6f", p[1], p[2], p[3])

function bbox2string(bbox::Rect3)
    p0 = origin(bbox)
    p1 = p0 .+ widths(bbox)
    @sprintf(
        """
        Bounding Box:
         x: (%0.3f, %0.3f)
         y: (%0.3f, %0.3f)
         z: (%0.3f, %0.3f)
        """,
        p0[1], p1[1], p0[2], p1[2], p0[3], p1[3]
    )
end

function bbox2string(bbox::Rect2)
    p0 = origin(bbox)
    p1 = p0 .+ widths(bbox)
    @sprintf(
        """
        Bounding Box:
         x: (%0.3f, %0.3f)
         y: (%0.3f, %0.3f)
        """,
        p0[1], p1[1], p0[2], p1[2]
    )
end

color2text(c::AbstractFloat) = @sprintf("%0.3f", c)
color2text(c::Symbol) = string(c)
color2text(c) = color2text(to_color(c))
function color2text(c::RGBAf)
    if c.alpha == 1.0
        @sprintf("RGB(%0.2f, %0.2f, %0.2f)", c.r, c.g, c.b)
    else
        @sprintf("RGBA(%0.2f, %0.2f, %0.2f, %0.2f)", c.r, c.g, c.b, c.alpha)
    end
end

color2text(name, i::Integer, j::Integer, c) = "$name[$i, $j] = $(color2text(c))"
function color2text(name, i, j, c)
    idxs = @sprintf("%0.2f, %0.2f", i, j)
    "$name[$idxs] = $(color2text(c))"
end

### dealing with markersize and rotations
########################################

_to_scale(f::AbstractFloat, idx) = Vec3f(f)
_to_scale(v::Vec2f, idx) = Vec3f(v[1], v[2], 1)
_to_scale(v::Vec3f, idx) = v
_to_scale(v::Vector, idx) = _to_scale(v[idx], idx)

_to_rotation(x, idx) = to_rotation(x)
_to_rotation(x::Vector, idx) = to_rotation(x[idx])

### Selecting a point on a nearby line
########################################

function closest_point_on_line(p0::Point2f, p1::Point2f, r::Point2f)
    # This only works in 2D
    AP = P .- A; AB = B .- A
    A .+ AB * dot(AP, AB) / dot(AB, AB)
end

function view_ray(scene)
    inv_projview = inv(camera(scene).projectionview[])
    view_ray(inv_projview, events(scene).mouseposition[], pixelarea(scene)[])
end
function view_ray(inv_view_proj, mpos, area::Rect2)
    # This figures out the camera view direction from the projectionview matrix (?)
    # and computes a ray from a near and a far point.
    # Based on ComputeCameraRay from ImGuizmo
    mp = 2f0 .* (mpos .- minimum(area)) ./ widths(area) .- 1f0
    v = inv_view_proj * Vec4f(0, 0, -10, 1)
    reversed = v[3] < v[4]
    near = reversed ? 1f0 - 1e-6 : 0f0
    far = reversed ? 0f0 : 1f0 - 1e-6

    origin = inv_view_proj * Vec4f(mp[1], mp[2], near, 1f0)
    origin = origin[SOneTo(3)] ./ origin[4]

    p = inv_view_proj * Vec4f(mp[1], mp[2], far, 1f0)
    p = p[SOneTo(3)] ./ p[4]

    dir = normalize(p - origin)
    return origin, dir
end

# These work in 2D and 3D
function closest_point_on_line(A, B, origin, dir)
    closest_point_on_line(
        to_ndim(Point3f, A, 0),
        to_ndim(Point3f, B, 0),
        to_ndim(Point3f, origin, 0),
        to_ndim(Vec3f, dir, 0)
    )
end
function closest_point_on_line(A::Point3f, B::Point3f, origin::Point3f, dir::Vec3f)
    # See:
    # https://en.wikipedia.org/wiki/Line%E2%80%93plane_intersection
    AB_norm = norm(B .- A)
    u_AB = (B .- A) / AB_norm
    u_dir = normalize(dir)
    u_perp = normalize(cross(u_dir, u_AB))
    # e_RD, e_perp defines a plane with normal n
    n = normalize(cross(u_dir, u_perp))
    t = dot(origin .- A, n) / dot(u_AB, n)
    A .+ clamp(t, 0.0, AB_norm) * u_AB
end

function ray_triangle_intersection(A, B, C, origin, dir)
    # See: https://www.iue.tuwien.ac.at/phd/ertl/node114.html
    AO = A .- origin
    BO = B .- origin
    CO = C .- origin
    A1 = 0.5 * dot(cross(BO, CO), dir)
    A2 = 0.5 * dot(cross(CO, AO), dir)
    A3 = 0.5 * dot(cross(AO, BO), dir)

    e = 1e-3
    if (A1 > -e && A2 > -e && A3 > -e) || (A1 < e && A2 < e && A3 < e)
        Point3f((A1 * A .+ A2 * B .+ A3 * C) / (A1 + A2 + A3))
    else
        Point3f(NaN)
    end
end

### Surface positions
########################################

surface_x(xs::ClosedInterval, i, j, N) = minimum(xs) + (maximum(xs) - minimum(xs)) * (i-1) / (N-1)
surface_x(xs, i, j, N) = xs[i]
surface_x(xs::AbstractMatrix, i, j, N) = xs[i, j]

surface_y(ys::ClosedInterval, i, j, N) = minimum(ys) + (maximum(ys) - minimum(ys)) * (j-1) / (N-1)
surface_y(ys, i, j, N) = ys[j]
surface_y(ys::AbstractMatrix, i, j, N) = ys[i, j]

function surface_pos(xs, ys, zs, i, j)
    N, M = size(zs)
    Point3f(surface_x(xs, i, j, N), surface_y(ys, i, j, M), zs[i, j])
end

### Mapping mesh vertex indices to Vector{Polygon} index
########################################

function vertexindex2poly(polys, idx)
    counter = 0
    for i in eachindex(polys)
        step = ncoords(polys[i])
        if idx <= counter + step
            return i
        else
            counter += step
        end
    end
    return length(polys)
end

ncoords(x) = length(coordinates(x))
ncoords(mesh::Mesh) = length(coordinates(mesh))
function ncoords(poly::Polygon)
    N = length(poly.exterior) + 1
    for int in poly.interiors
        N += length(int) + 1
    end
    N
end

### Text bounding box
########################################

function Bbox_from_glyphcollection(text, gc)
    bbox = Rect2f(0, 0, 0, 0)
    bboxes = Rect2f[]
    broadcast_foreach(gc.extents, gc.fonts, gc.scales) do extent, font, scale
        b = FreeTypeAbstraction.height_insensitive_boundingbox(extent, font) * scale
        push!(bboxes, b)
    end
    for (c, o, bb) in zip(text, gc.origins, bboxes)
        c == '\n' && continue
        bbox2 = Rect2f(o[Vec(1,2)] .+ origin(bb), widths(bb))
        if bbox == Rect2f(0, 0, 0, 0)
            bbox = bbox2
        else
            bbox = union(bbox, bbox2)
        end
    end
    bbox
end

## Shifted projection
########################################

function shift_project(scene, plot, pos)
    project(
        camera(scene).projectionview[],
        Vec2f(widths(pixelarea(scene)[])),
        apply_transform(transform_func_obs(plot)[], pos)
    ) .+ Vec2f(origin(pixelarea(scene)[]))
end

################################################################################
### Base pixel-space plot for Indicator
################################################################################

@recipe(_Inspector, x) do scene
    # Attributes starting with _ are modified internally
    Attributes(
        # Text
        text_padding = Vec4f(5, 5, 3, 3), # LRBT
        text_align = (:left, :bottom),
        textcolor = :black,
        textsize = 18,
        font = theme(scene, :font),
        _display_text = " ",
        _text_position = Point2f(0),

        # Background
        background_color = :white,
        outline_color = :grey,
        outline_linestyle = nothing,
        outline_linewidth = 2,

        # pixel BBox/indicator
        indicator_color = :red,
        indicator_linewidth = 2,
        indicator_linestyle = nothing,
        _bbox2D = Rect2f(Vec2f(0), Vec2f(0)),
        _px_bbox_visible = true,

        # general
        tooltip_align = (:center, :top), # default tooltip position relative to cursor
        tooltip_offset = Vec2f(20), # default offset in alignment direction
        depth = 9e3,
        enabled = true,
        range = 10,

        _root_px_projection = Mat4f(1,0,0,0,0,1,0,0,0,0,1,0,0,0,0,1),
        _model = Mat4f(1,0,0,0,0,1,0,0,0,0,1,0,0,0,0,1),
        _visible = true,
        _tooltip_align = (:center, :top),
        _bbox3D = Rect3f(Vec3f(0), Vec3f(0)),
        _bbox_visible = true,
        _position = Point3f(0),
        _color = RGBAf(0,0,0,0)
    )
end

function plot!(plot::_Inspector)
    @extract plot (
        text_padding, text_align, textcolor, textsize, font,
        background_color, outline_color, outline_linestyle, outline_linewidth,
        indicator_linestyle, indicator_linewidth, indicator_color,
        tooltip_offset, depth,
        _display_text, _text_position, _bbox2D, _px_bbox_visible,
        _tooltip_align, _root_px_projection, _visible
    )

    # tooltip text
    _aligned_text_position = Observable(Point2f(0))
    id = Mat4f(1,0,0,0,0,1,0,0,0,0,1,0,0,0,0,1)
    text_plot = text!(plot, _display_text,
        position = _aligned_text_position, visible = _visible, align = text_align,
        color = textcolor, font = font, textsize = textsize,
        inspectable = false,
        # with https://github.com/JuliaPlots/Makie.jl/tree/master/GLMakie/pull/183 this should
        # allow the tooltip to work in any scene.
        space = :data,
        projection = _root_px_projection, view = id, projectionview = _root_px_projection
    )

    # compute text boundingbox and adjust _aligned_text_position
    bbox = map(text_plot.plots[1][1], text_plot.position, text_padding) do gc, pos, pad
        rect = Bbox_from_glyphcollection(_display_text[], gc)
        l, r, b, t = pad
        Rect2f(
            origin(rect) .+ Vec2f(pos[1] - l, pos[2] - b),
            widths(rect) .+ Vec2f(l + r, b + t)
        )
    end
    onany(_text_position, _tooltip_align, tooltip_offset, bbox) do pos, align, offset, bbox
        halign, valign = align
        ox, oy = offset
        wx, wy = widths(bbox)
        if halign == :left
            dx = -wx - ox
        elseif halign == :center
            dx = -0.5wx
        elseif halign == :right
            dx = + ox
        end
        if valign == :top
            dy = oy
        elseif valign == :center
            dy = -0.5wy
        elseif valign == :bottom
            dy = -wy - oy
        end
        new_pos = pos .+ Point2f(dx, dy)
        if new_pos != _aligned_text_position[]
            _aligned_text_position[] = new_pos
        end
    end

    # tooltip background and frame
    background = mesh!(
        plot, bbox, color = background_color, shading = false, #fxaa = false,
        # TODO with fxaa here the text above becomes seethrough on a heatmap
        visible = _visible, inspectable = false,
        projection = _root_px_projection, view = id, projectionview = _root_px_projection
    )
    outline = wireframe!(
        plot, bbox,
        color = outline_color, visible = _visible, inspectable = false,
        linestyle = outline_linestyle, linewidth = outline_linewidth,
        projection = _root_px_projection, view = id, projectionview = _root_px_projection
    )

    # pixel-space marker for selected element (not always used)
    px_bbox = wireframe!(
        plot, _bbox2D,
        color = indicator_color, linewidth = indicator_linewidth,
        linestyle = indicator_linestyle, visible = _px_bbox_visible,
        inspectable = false,
        projection = _root_px_projection, view = id, projectionview = _root_px_projection
    )

    # To make sure inspector plots end up in front
    on(depth) do d
        translate!(background, Vec3f(0,0,d+1))
        translate!(outline,    Vec3f(0,0,d+2))
        translate!(text_plot,  Vec3f(0,0,d+3))
        translate!(px_bbox,    Vec3f(0,0,d))
    end
    depth[] = depth[]
    nothing
end

################################################################################
### Interactive selection via DataInspector
################################################################################

# TODO destructor?
mutable struct DataInspector
    root::Scene
    temp_plots::Vector{AbstractPlot}
    plot::_Inspector
    selection::AbstractPlot
    obsfuncs::Vector{Any}
end

function DataInspector(scene::Scene, plot::AbstractPlot)
    x = DataInspector(scene, AbstractPlot[], plot, plot, Any[])
    # finalizer(cleanup, x) # doesn't get triggered when this is dereferenced
    x
end

function cleanup(inspector::DataInspector)
    off.(inspector.obsfuncs)
    empty!(inspector.obsfuncs)
    delete!(inspector.root, inspector.plot)
    clear_temporary_plots!(inspector, inspector.slection)
    inspector
end

function Base.delete!(::Union{Scene, Figure}, inspector::DataInspector)
    cleanup(inspector)
end

enable!(inspector::DataInspector) = inspector.plot.enabled[] = true
disable!(inspector::DataInspector) = inspector.plot.enabled[] = false

"""
    DataInspector(figure; kwargs...)

Creates a data inspector which will show relevant information in a tooltip
when you hover over a plot. If you wish to exclude a plot you may set
`plot.inspectable[] = false`.

### Keyword Arguments:
- `range = 10`: Controls the snapping range for selecting an element of a plot.
- `enabled = true`: Disables inspection of plots when set to false. Can also be
    adjusted with `enable!(inspector)` and `disable!(inspector)`.
- `text_padding = Vec4f(5, 5, 3, 3)`: Padding for the box drawn around the
    tooltip text. (left, right, bottom, top)
- `text_align = (:left, :bottom)`: Alignment of text within the tooltip. This
    does not affect the alignment of the tooltip relative to the cursor.
- `textcolor = :black`: Tooltip text color.
- `textsize = 20`: Tooltip text size.
- `font = "Dejavu Sans"`: Tooltip font.
- `background_color = :white`: Background color of the tooltip.
- `outline_color = :grey`: Outline color of the tooltip.
- `outline_linestyle = nothing`: Linestyle of the tooltip outline.
- `outline_linewidth = 2`: Linewidth of the tooltip outline.
- `indicator_color = :red`: Color of the selection indicator.
- `indicator_linewidth = 2`: Linewidth of the selection indicator.
- `indicator_linestyle = nothing`: Linestyle of the selection indicator
- `tooltip_align = (:center, :top)`: Default position of the tooltip relative to
    the cursor or current selection. The real align may adjust to keep the
    tooltip in view.
- `tooltip_offset = Vec2f(20)`: Offset from the indicator to the tooltip.
- `depth = 9e3`: Depth value of the tooltip. This should be high so that the
    tooltip is always in front.
- `priority = 100`: The priority of creating a tooltip on a mouse movement or
    scrolling event.
"""
function DataInspector(fig_or_layoutable; kwargs...)
    DataInspector(fig_or_layoutable.scene; kwargs...)
end

function DataInspector(scene::Scene; priority = 100, kwargs...)
    parent = root(scene)
    @assert origin(pixelarea(parent)[]) == Vec2f(0)

    plot = _inspector!(
        parent, 1,
        _root_px_projection = camera(parent).pixel_space;
        kwargs...
    )
    inspector = DataInspector(parent, plot)

    e = events(parent)
    f1 = on(_ -> on_hover(inspector), e.mouseposition, priority = priority)
    f2 = on(_ -> on_hover(inspector), e.scroll, priority = priority)

    push!(inspector.obsfuncs, f1, f2)

    inspector
end

function on_hover(inspector)
    parent = inspector.root
    (inspector.plot.enabled[] && is_mouseinside(parent)) || return Consume(false)

    mp = mouseposition_px(parent)
    should_clear = true
    for (plt, idx) in pick_sorted(parent, mp, inspector.plot.range[])
        if to_value(get(plt.attributes, :inspectable, true))
            # show_data should return true if it created a tooltip
            if show_data_recursion(inspector, plt, idx)
                should_clear = false
                break
            end
        end
    end

    if should_clear
        inspector.plot._visible[] = false
        inspector.plot._bbox_visible[] = false
        inspector.plot._px_bbox_visible[] = false
    end

    return Consume(false)
end

function show_data_recursion(inspector, plot, idx)
    processed = show_data_recursion(inspector, plot.parent, idx, plot)
    if processed
        return true
    else
        return show_data(inspector, plot, idx)
    end
end
show_data_recursion(inspector, plot, idx, source) = false
function show_data_recursion(inspector, plot::AbstractPlot, idx, source)
    processed = show_data_recursion(inspector, plot.parent, idx, source)
    if processed
        return true
    else
        return show_data(inspector, plot, idx, source)
    end
end

# clears temporary plots (i.e. bboxes) and update selection
function clear_temporary_plots!(inspector::DataInspector, plot)
    inspector.selection = plot
    for i in length(inspector.obsfuncs):-1:3
        off(pop!(inspector.obsfuncs))
    end
    for p in inspector.temp_plots
        delete!(parent_scene(p), p)
    end
    empty!(inspector.temp_plots)
end

# update alignment direction
function update_tooltip_alignment!(inspector, proj_pos)
    a = inspector.plot.attributes
    a._text_position[] = proj_pos

    wx, wy = widths(pixelarea(inspector.root)[])
    px, py = proj_pos
    halign, valign = a.tooltip_align[]
    px < wx/4  && (halign = :right)
    px > 3wx/4 && (halign = :left)
    py < wy/4  && (valign = :top)
    py > 3wy/4 && (valign = :bottom)
    a._tooltip_align[] = (halign, valign)
end

################################################################################
### show_data for primitive plots
################################################################################

# TODO: better 3D scaling
function show_data(inspector::DataInspector, plot::Scatter, idx)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    proj_pos = shift_project(scene, plot, to_ndim(Point3f, plot[1][][idx], 0))
    update_tooltip_alignment!(inspector, proj_pos)
    ms = plot.markersize[]

    a._display_text[] = position2string(plot[1][][idx])
    a._bbox2D[] = Rect2f(proj_pos .- 0.5 .* ms .- Vec2f(5), Vec2f(ms) .+ Vec2f(10))
    a._px_bbox_visible[] = true
    a._bbox_visible[] = false
    a._visible[] = true

    return true
end

function show_data(inspector::DataInspector, plot::MeshScatter, idx)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    proj_pos = shift_project(scene, plot, to_ndim(Point3f, plot[1][][idx], 0))
    update_tooltip_alignment!(inspector, proj_pos)
    bbox = Rect{3, Float32}(convert_attribute(
        plot.marker[], Key{:marker}(), Key{Makie.plotkey(plot)}()
    ))

    a._model[] = transformationmatrix(
        plot[1][][idx],
        _to_scale(plot.markersize[], idx),
        _to_rotation(plot.rotations[], idx)
    )

    if inspector.selection != plot
        cc = cameracontrols(scene)
        if cc isa Camera3D
            eyeposition = cc.eyeposition[]
            lookat = cc.lookat[]
            upvector = cc.upvector[]
        end

        # To avoid putting a bbox outside the plots bbox
        a._bbox3D[] = boundingbox(plot)

        clear_temporary_plots!(inspector, plot)
        p = wireframe!(
            scene, a._bbox3D, model = a._model, color = a.indicator_color,
            linewidth = a.indicator_linewidth, linestyle = a.indicator_linestyle,
            visible = a._bbox_visible, inspectable = false
        )
        push!(inspector.temp_plots, p)

        # Restore camera
        cc isa Camera3D && update_cam!(scene, eyeposition, lookat, upvector)
    end

    a._display_text[] = position2string(plot[1][][idx])
    a._bbox3D[] = bbox
    a._px_bbox_visible[] = false
    a._bbox_visible[] = true
    a._visible[] = true

    return true
end

function show_data(inspector::DataInspector, plot::Union{Lines, LineSegments}, idx)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    # cast ray from cursor into screen, find closest point to line
    p0, p1 = plot[1][][idx-1:idx]
    origin, dir = view_ray(scene)
    pos = closest_point_on_line(p0, p1, origin, dir)
    lw = plot.linewidth[] isa Vector ? plot.linewidth[][idx] : plot.linewidth[]

    proj_pos = shift_project(scene, plot, to_ndim(Point3f, pos, 0))
    update_tooltip_alignment!(inspector, proj_pos)

    a._display_text[] = position2string(typeof(p0)(pos))
    a._bbox2D[] = Rect2f(proj_pos .- 0.5 .* lw .- Vec2f(5), Vec2f(lw) .+ Vec2f(10))
    a._px_bbox_visible[] = true
    a._bbox_visible[] = false
    a._visible[] = true

    return true
end

function show_data(inspector::DataInspector, plot::Mesh, idx)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    bbox = boundingbox(plot)
    proj_pos = Point2f(mouseposition_px(inspector.root))
    update_tooltip_alignment!(inspector, proj_pos)

    a._bbox3D[] = bbox

    if inspector.selection != plot
        cc = cameracontrols(scene)
        if cc isa Camera3D
            eyeposition = cc.eyeposition[]
            lookat = cc.lookat[]
            upvector = cc.upvector[]
        end

        clear_temporary_plots!(inspector, plot)
        p = wireframe!(
            scene, a._bbox3D, color = a.indicator_color,
            linewidth = a.indicator_linewidth, linestyle = a.indicator_linestyle,
            visible = a._bbox_visible, inspectable = false
        )
        push!(inspector.temp_plots, p)

        # Restore camera
        cc isa Camera3D && update_cam!(scene, eyeposition, lookat, upvector)
    end

    a._text_position[] = proj_pos
    a._display_text[] = bbox2string(bbox)
    a._px_bbox_visible[] = false
    a._bbox_visible[] = true
    a._visible[] = true

    return true
end

function show_data(inspector::DataInspector, plot::Surface, idx)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    proj_pos = Point2f(mouseposition_px(inspector.root))
    update_tooltip_alignment!(inspector, proj_pos)

    xs = plot[1][]
    ys = plot[2][]
    zs = plot[3][]
    w, h = size(zs)
    _i = mod1(idx, w); _j = div(idx-1, w)

    # This isn't the most accurate so we include some neighboring faces
    origin, dir = view_ray(scene)
    pos = Point3f(NaN)
    for i in _i-1:_i+1, j in _j-1:_j+1
        (1 <= i <= w) && (1 <= j < h) || continue

        if i - 1 > 0
            pos = ray_triangle_intersection(
                surface_pos(xs, ys, zs, i, j),
                surface_pos(xs, ys, zs, i-1, j),
                surface_pos(xs, ys, zs, i, j+1),
                origin, dir
            )
        end

        if i + 1 <= w && isnan(pos)
            pos = ray_triangle_intersection(
                surface_pos(xs, ys, zs, i, j),
                surface_pos(xs, ys, zs, i, j+1),
                surface_pos(xs, ys, zs, i+1, j+1),
                origin, dir
            )
        end

        isnan(pos) || break
    end

    if !isnan(pos)
        a._display_text[] = position2string(pos)
        a._text_position[] = proj_pos
        a._bbox2D[] = Rect2f(proj_pos .- Vec2f(5), Vec2f(10))
        a._bbox_visible[] = false
        a._px_bbox_visible[] = true
        a._visible[] = true
    else
        a._bbox_visible[] = false
        a._px_bbox_visible[] = false
        a._visible[] = false
    end

    return true
end

function show_data(inspector::DataInspector, plot::Heatmap, idx)
    show_imagelike(inspector, plot, "H", true)
end

function show_data(inspector::DataInspector, plot::Image, idx)
    show_imagelike(inspector, plot, "img", false)
end

function show_imagelike(inspector, plot, name, edge_based)
    a = inspector.plot.attributes
    scene = parent_scene(plot)
    mpos = mouseposition(scene)

    i, j, z = if plot.interpolate[]
        _interpolated_getindex(plot[1][], plot[2][], plot[3][], mpos)
    else
        _pixelated_getindex(plot[1][], plot[2][], plot[3][], mpos, edge_based)
    end

    a._color[] = if z isa AbstractFloat
        interpolated_getindex(
            to_colormap(plot.colormap[]), z,
            to_value(get(plot.attributes, :colorrange, (0, 1)))
        )
    else
        z
    end

    a._position[] = to_ndim(Point3f, mpos, 0)
    proj_pos = Point2f(mouseposition_px(inspector.root))
    update_tooltip_alignment!(inspector, proj_pos)

    if plot.interpolate[]
        if inspector.selection != plot || !(inspector.temp_plots[1] isa Scatter)
            clear_temporary_plots!(inspector, plot)
            p = scatter!(
                scene, map(p -> [p], a._position), color = a._color,
                visible = a._bbox_visible,
                inspectable = false,
                marker=Rect, markersize = map(r -> 2r - 4, a.range),
                strokecolor = a.indicator_color,
                strokewidth = a.indicator_linewidth #, linestyle = a.indicator_linestyle no?
            )
            translate!(p, Vec3f(0, 0, a.depth[]))
            push!(inspector.temp_plots, p)
            # Hacky?
            push!(
                inspector.obsfuncs,
                Observables.ObserverFunction(a._position.listeners[end], a._position, false)
            )
        end
        a._display_text[] = color2text(name, mpos[1], mpos[2], z)
    else
        a._bbox2D[] = _pixelated_image_bbox(plot[1][], plot[2][], plot[3][], i, j, edge_based)
        if inspector.selection != plot || !(inspector.temp_plots[1][1][] isa Rect2)
            clear_temporary_plots!(inspector, plot)
            p = wireframe!(
                scene, a._bbox2D, model = a._model, color = a.indicator_color,
                strokewidth = a.indicator_linewidth, linestyle = a.indicator_linestyle,
                visible = a._bbox_visible, inspectable = false
            )
            translate!(p, Vec3f(0, 0, a.depth[]))
            push!(inspector.temp_plots, p)
        end
        a._display_text[] = color2text(name, i, j, z)
    end

    a._bbox_visible[] = true
    a._px_bbox_visible[] = false
    a._visible[] = true
    return true
end

function _interpolated_getindex(xs, ys, img, mpos)
    x0, x1 = extrema(xs)
    y0, y1 = extrema(ys)
    x, y = clamp.(mpos, (x0, y0), (x1, y1))

    i = clamp((x - x0) / (x1 - x0) * size(img, 1) + 0.5, 1, size(img, 1))
    j = clamp((y - y0) / (y1 - y0) * size(img, 2) + 0.5, 1, size(img, 2))
    l = clamp(floor(Int, i), 1, size(img, 1)-1);
    r = clamp(l+1, 2, size(img, 1))
    b = clamp(floor(Int, j), 1, size(img, 2)-1);
    t = clamp(b+1, 2, size(img, 2))
    z = ((r-i) * img[l, b] + (i-l) * img[r, b]) * (t-j) +
        ((r-i) * img[l, t] + (i-l) * img[r, t]) * (j-b)

    # float, float, value (i, j are no longer used)
    return i, j, z
end
function _pixelated_getindex(xs, ys, img, mpos, edge_based)
    x0, x1 = extrema(xs)
    y0, y1 = extrema(ys)
    x, y = clamp.(mpos, (x0, y0), (x1, y1))

    i = clamp(round(Int, (x - x0) / (x1 - x0) * size(img, 1) + 0.5), 1, size(img, 1))
    j = clamp(round(Int, (y - y0) / (y1 - y0) * size(img, 2) + 0.5), 1, size(img, 2))

    # int, int, value
    return i, j, img[i,j]
end

function _interpolated_getindex(xs::Vector, ys::Vector, img, mpos)
    # x, y = mpos
    # i, j, _ = _pixelated_getindex(xs, ys, img, mpos, false)
    # w = (xs[i+1] - xs[i]); h = (ys[j+1] - ys[j])
    # z = ((xs[i+1] - x) / w * img[i, j]   + (x - xs[i]) / w * img[i+1, j])   * (ys[j+1] - y) / h +
    #     ((xs[i+1] - x) / w * img[i, j+1] + (x - xs[i]) / w * img[i+1, j+1]) * (y - ys[j]) / h
    # return i, j, z
    _interpolated_getindex(minimum(xs)..maximum(xs), minimum(ys)..maximum(ys), img, mpos)
end
function _pixelated_getindex(xs::Vector, ys::Vector, img, mpos, edge_based)
    if edge_based
        x, y = mpos
        i = max(1, something(findfirst(v -> v >= x, xs), length(xs))-1)
        j = max(1, something(findfirst(v -> v >= y, ys), length(ys))-1)
        return i, j, img[i, j]
    else
        _pixelated_getindex(minimum(xs)..maximum(xs), minimum(ys)..maximum(ys), img, mpos, edge_based)
    end
end

function _pixelated_image_bbox(xs, ys, img, i::Integer, j::Integer, edge_based)
    x0, x1 = extrema(xs)
    y0, y1 = extrema(ys)
    nw, nh = ((x1 - x0), (y1 - y0)) ./ size(img)
    Rect2f(x0 + nw * (i-1), y0 + nh * (j-1), nw, nh)
end
function _pixelated_image_bbox(xs::Vector, ys::Vector, img, i::Integer, j::Integer, edge_based)
    if edge_based
        Rect2f(xs[i], ys[j], xs[i+1] - xs[i], ys[j+1] - ys[j])
    else
        _pixelated_image_bbox(
            minimum(xs)..maximum(xs), minimum(ys)..maximum(ys),
            img, i, j, edge_based
        )
    end
end

function show_data(inspector::DataInspector, plot, idx, source=nothing)
    return false
end

################################################################################
### show_data for Combined/recipe plots
################################################################################

function show_data(inspector::DataInspector, plot::BarPlot, idx, ::Lines)
    return show_data(inspector, plot, div(idx-1, 6)+1)
end

function show_data(inspector::DataInspector, plot::BarPlot, idx, ::Mesh)
    return show_data(inspector, plot, div(idx-1, 4)+1)
end

function show_data(inspector::DataInspector, plot::BarPlot, idx)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    pos = plot[1][][idx]
    proj_pos = shift_project(scene, plot, to_ndim(Point3f, pos, 0))
    update_tooltip_alignment!(inspector, proj_pos)
    a._model[] = plot.model[]
    a._bbox2D[] = plot.plots[1][1][][idx]

    if inspector.selection != plot
        clear_temporary_plots!(inspector, plot)
        p = wireframe!(
            scene, a._bbox2D, model = a._model, color = a.indicator_color,
            strokewidth = a.indicator_linewidth, linestyle = a.indicator_linestyle,
            visible = a._bbox_visible, inspectable = false
        )
        translate!(p, Vec3f(0, 0, a.depth[]))
        push!(inspector.temp_plots, p)
    end

    a._display_text[] = position2string(pos)
    a._bbox_visible[] = true
    a._px_bbox_visible[] = false
    a._visible[] = true

    return true
end

function show_data(inspector::DataInspector, plot::Arrows, idx, ::LineSegments)
    return show_data(inspector, plot, div(idx+1, 2), nothing)
end
function show_data(inspector::DataInspector, plot::Arrows, idx, source)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    pos = plot[1][][idx]
    proj_pos = shift_project(scene, plot, to_ndim(Point3f, pos, 0))

    mpos = Point2f(mouseposition_px(inspector.root))
    update_tooltip_alignment!(inspector, mpos)

    p = vec2string(pos)
    v = vec2string(plot[2][][idx])

    a._text_position[] = mpos
    a._display_text[] = "Position:\n  $p\nDirection\n  $v"
    a._bbox2D[] = Rect2f(proj_pos .- Vec2f(5), Vec2f(10))
    a._bbox_visible[] = false
    a._px_bbox_visible[] = true
    a._visible[] = true

    return true
end

# This should work if contourf would place computed levels in colors and let the
# backend handle picking colors from a colormap
function show_data(inspector::DataInspector, plot::Contourf, idx, source::Mesh)
    a = inspector.plot.attributes
    scene = parent_scene(plot)
    idx, ext = show_poly(inspector, plot.plots[1], idx, source)
    level = plot.plots[1].color[][idx]

    a._text_position[] = Point2f(mouseposition_px(inspector.root))
    a._display_text[] = @sprintf("level = %0.3f", level)
    return true
end

# What should this display?
# function show_data(
#         inspector::DataInspector, plot::Poly{<: Tuple{<: AbstractVector}},
#         idx, source::Mesh
#     )
#     @info "PolyMesh"
#     idx, ext = show_poly(inspectable, plot, idx, source)
#     return true
# end

function show_poly(inspector, plot, idx, source)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    idx = vertexindex2poly(plot[1][], idx)
    m = GeometryBasics.mesh(plot[1][][idx])

    clear_temporary_plots!(inspector, plot)
    ext = plot[1][][idx].exterior
    p = lines!(
        scene, ext, color = a.indicator_color,
        strokewidth = a.indicator_linewidth, linestyle = a.indicator_linestyle,
        visible = a._visible, inspectable = false
    )
    translate!(p, Vec3f(0,0,a.depth[]))
    push!(inspector.temp_plots, p)

    for int in plot[1][][idx].interiors
        p = lines!(
            scene, int, color = a.indicator_color,
            strokewidth = a.indicator_linewidth, linestyle = a.indicator_linestyle,
            visible = a._visible, inspectable = false
        )
        translate!(p, Vec3f(0,0,a.depth[]))
        push!(inspector.temp_plots, p)
    end

    a._px_bbox_visible[] = false
    a._bbox_visible[] = true
    a._visible[] = true

    return idx, ext
end

function show_data(inspector::DataInspector, plot::VolumeSlices, idx, child::Heatmap)
    a = inspector.plot.attributes
    scene = parent_scene(plot)

    proj_pos = Point2f(mouseposition_px(inspector.root))
    update_tooltip_alignment!(inspector, proj_pos)

    qs = extrema(child[1][])
    ps = extrema(child[2][])
    data = child[3][]
    trans = child.transformation.translation[]

    child_idx = findfirst(isequal(child), plot.plots)
    if child_idx == 2
        vs = [ # clockwise
            Point3f(trans[1], qs[1], ps[1]),
            Point3f(trans[1], qs[1], ps[2]),
            Point3f(trans[1], qs[2], ps[2]),
            Point3f(trans[1], qs[2], ps[1])
        ]
    elseif child_idx == 3
        vs = [ # clockwise
            Point3f(qs[1], trans[2], ps[1]),
            Point3f(qs[1], trans[2], ps[2]),
            Point3f(qs[2], trans[2], ps[2]),
            Point3f(qs[2], trans[2], ps[1])
        ]
    else
        vs = [ # clockwise
            Point3f(qs[1], ps[1], trans[3]),
            Point3f(qs[1], ps[2], trans[3]),
            Point3f(qs[2], ps[2], trans[3]),
            Point3f(qs[2], ps[1], trans[3])
        ]
    end

    origin, dir = view_ray(scene)
    pos = Point3f(NaN)
    pos = ray_triangle_intersection(vs[1], vs[2], vs[3], origin, dir)
    if isnan(pos)
        pos = ray_triangle_intersection(vs[3], vs[4], vs[1], origin, dir)
    end

    if !isnan(pos)
        if child_idx == 2
            x = pos[2]; y = pos[3]
        elseif child_idx == 3
            x = pos[1]; y = pos[3]
        else
            x = pos[1]; y = pos[2]
        end
        i = clamp(round(Int, (x - qs[1]) / (qs[2] - qs[1]) * size(data, 1) + 0.5), 1, size(data, 1))
        j = clamp(round(Int, (y - ps[1]) / (ps[2] - ps[1]) * size(data, 2) + 0.5), 1, size(data, 2))
        val = data[i, j]

        a._display_text[] = @sprintf(
            "x: %0.6f\ny: %0.6f\nz: %0.6f\n%0.6f0",
            pos[1], pos[2], pos[3], val
        )
        a._text_position[] = proj_pos
        a._bbox2D[] = Rect2f(proj_pos .- Vec2f(5), Vec2f(10))
        a._bbox_visible[] = false
        a._px_bbox_visible[] = true
        a._visible[] = true
    else
        a._bbox_visible[] = false
        a._px_bbox_visible[] = false
        a._visible[] = false
    end

    return true
end


using AutoViz
using Colors

#############################################################
# Radar Sensor

type RadarSensor
    angles::Vector{Float64}
    ranges::Vector{Float64}
    range_rates::Vector{Float64}
    max_range::Float64
    poly::ConvexPolygon
end
function RadarSensor(nbeams::Int;
    max_range::Float64=100.0,
    angle_offset::Float64=0.0,
    )
    angles = collect(linspace(angle_offset,2*pi+angle_offset,nbeams+1))[2:end]
    ranges = Array(Float64, nbeams)
    range_rates = Array(Float64, nbeams)
    RadarSensor(angles, ranges, range_rates, max_range, ConvexPolygon(4))
end
function observe!(radar::RadarSensor, scene::Scene, roadway::Roadway, vehicle_index::Int)
    state_ego = scene[vehicle_index].state
    egoid = scene[vehicle_index].def.id
    ego_vel = polar(state_ego.v, state_ego.posG.θ)
    for (i,angle) in enumerate(radar.angles)
        ray_angle = state_ego.posG.θ + angle
        ray_vec = polar(1.0, ray_angle)
        ray = VecSE2(state_ego.posG.x, state_ego.posG.y, ray_angle)

        range = radar.max_range
        range_rate = 0.0
        for veh in scene
            if veh.def.id != egoid
                to_oriented_bounding_box!(radar.poly, veh)
                range2 = AutomotiveDrivingModels.get_collision_time(ray, radar.poly, 1.0)
                if !isnan(range2) && range2 < range
                    range = range2
                    relative_speed = polar(veh.state.v, veh.state.posG.θ) - ego_vel
                    range_rate = proj(relative_speed, ray_vec, Float64)
                end
            end
        end
        radar.ranges[i] = range
        radar.range_rates[i] = range_rate
    end

    radar
end
function render_radar!(rendermodel::RenderModel, radar::RadarSensor, posG::VecSE2;
    color::Colorant=colorant"white",
    line_width::Float64 = 0.05,
    )

    for (angle, range, range_rate) in zip(radar.angles, radar.ranges, radar.range_rates)
        ray = VecSE2(posG.x, posG.y, posG.θ + angle)
        t = 2 / (1 + exp(-range_rate))
        ray_color = t > 1.0 ? lerp(RGBA(1.0,1.0,1.0,0.5), RGBA(0.2,0.2,1.0,0.5), t/2) :
                              lerp(RGBA(1.0,0.2,0.2,0.5), RGBA(1.0,1.0,1.0,0.5), t)
        render_ray!(rendermodel::RenderModel, ray, color=ray_color, length=range, line_width=line_width)
    end
    rendermodel
end

#############################################################
# Road Line Radar Sensor

type RoadlineRadarSensor
    angles::Vector{Float64}
    ranges::Matrix{Float64} # [n_lanes by nbeams]
    max_range::Float64
    poly::ConvexPolygon
end
function RoadlineRadarSensor(nbeams::Int;
    max_range::Float64=100.0,
    max_depth::Int=2,
    angle_offset::Float64=0.0,
    )

    angles = collect(linspace(angle_offset,2*pi+angle_offset,nbeams+1))[2:end]
    ranges = Array(Float64, max_depth, nbeams)
    RoadlineRadarSensor(angles, ranges, max_range, ConvexPolygon(4))
end
function _update_radar!(radar::RoadlineRadarSensor, ray, p_lo, p_hi, beam_index)
    test_range = get_collision_time(ray, AutomotiveDrivingModels.LineSegment(p_lo, p_hi), 1.0)

    n_ranges = size(radar.ranges, 1)
    if !isnan(test_range)
        for k in 1 : n_ranges
            if test_range < radar.ranges[k,beam_index]
                for l in k+1 : n_ranges
                    radar.ranges[l,beam_index] = radar.ranges[l-1,beam_index]
                end
                radar.ranges[k,beam_index] = test_range
                break
            end
        end
    end
    radar
end
function observe!(radar::RoadlineRadarSensor, scene::Scene, roadway::Roadway, vehicle_index::Int)
    state_ego = scene[vehicle_index].state
    egoid = scene[vehicle_index].def.id
    ego_vel = polar(state_ego.v, state_ego.posG.θ)

    fill!(radar.ranges, radar.max_range)
    for (i,angle) in enumerate(radar.angles)
        ray_angle = state_ego.posG.θ + angle
        ray_vec = polar(1.0, ray_angle)
        ray = VecSE2(state_ego.posG.x, state_ego.posG.y, ray_angle)

        for seg in roadway.segments
            for lane in seg.lanes

                N = length(lane.curve)
                halfwidth = lane.width/2

                # always render the left lane marking
                p_lo = convert(VecE2, lane.curve[1].pos + polar(halfwidth, lane.curve[1].pos.θ + π/2))
                for j in 2:length(lane.curve)
                    p_hi = convert(VecE2, lane.curve[j].pos + polar(halfwidth, lane.curve[j].pos.θ + π/2))
                    _update_radar!(radar, ray, p_lo, p_hi, i)
                    p_lo = p_hi
                end
                if has_next(lane)
                    lane2 = next_lane(lane, roadway)
                    pt = lane2.curve[1]
                    p_hi = convert(VecE2, pt.pos + polar(lane2.width/2, pt.pos.θ + π/2))
                    _update_radar!(radar, ray, p_lo, p_hi, i)
                end

                # only check the right lane marking if this is the first lane
                if lane.tag.lane == 1
                    p_lo = convert(VecE2, lane.curve[1].pos + polar(halfwidth, lane.curve[1].pos.θ - π/2))
                    for j in 2:length(lane.curve)
                        p_hi = convert(VecE2, lane.curve[j].pos + polar(halfwidth, lane.curve[j].pos.θ - π/2))
                        _update_radar!(radar, ray, p_lo, p_hi, i)
                        p_lo = p_hi
                    end
                    if has_next(lane)
                        lane2 = next_lane(lane, roadway)
                        pt = lane2.curve[1]
                        p_hi = convert(VecE2, pt.pos + polar(lane2.width/2, pt.pos.θ - π/2))
                        _update_radar!(radar, ray, p_lo, p_hi, i)
                    end
                end
            end
        end
    end

    radar
end
function render_radar!(rendermodel::RenderModel, radar::RoadlineRadarSensor, posG::VecSE2;
    color::Colorant=RGBA(1.0,1.0,1.0,0.5),
    line_width::Float64 = 0.05,
    depth_level::Int = 1, # if 1, render first lanes struck by beams, if 2 render 2nd ...
    )

    for (angle, range) in zip(radar.angles, radar.ranges[depth_level,:])
        ray = VecSE2(posG.x, posG.y, posG.θ + angle)
        render_ray!(rendermodel::RenderModel, ray, color=color, length=range, line_width=line_width)
    end
    rendermodel
end

#############################################################
# Feature Extraction

const RADAR = RadarSensor(200, max_range=100.0, angle_offset=-π) # car distances and range rates
const LINERADAR = RoadlineRadarSensor(200, max_range=100.0, angle_offset=-π, max_depth=3) # lane edges
const N_FEATURES = 2*length(RADAR.ranges) + length(LINERADAR.ranges)
function AutomotiveDrivingModels.pull_features!{Fl<:AbstractFloat}(features::Array{Fl}, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0;
    radar::RadarSensor = RADAR,
    lineradar::RoadlineRadarSensor = LINERADAR,
    )

    #=
    let N = length(radar.ranges)
    let M = size(lineradar.ranges, 2)
    let max_depth = size(lineradar.ranges, 1)

    1 : N is the radar ranges; distance to the struck object, maxes out
    N+1 : 2N is the radar range rates
    2N+1 : 2N + M is line radar ranges at depth 1
    2N+M+1 : 2N + 2M is line radar ranges at depth 2
    ...
    2N+(max_depth-1)×M+1 : 2N + max_depth×M is line radar ranges at depth max_depth
    =#

    scene = get_scene(rec, pastframe)

    observe!(radar, scene, roadway, vehicle_index)
    observe!(lineradar, scene, roadway, vehicle_index)

    N = length(radar.ranges)
    M = size(lineradar.ranges, 2)
    max_depth = size(lineradar.ranges, 1)

    copy!(features, 1, radar.ranges, 1)
    copy!(features, N+1, radar.range_rates, 1)
    copy!(features, 2N+1, lineradar.ranges, 1)

    if findfirst(v->isnan(v), features) != 0
        error("feature $(findfirst(v->isnan(v), features)) is nan")
    end

    features
end


# Test that the land model still runs, even with the lowest/simplest
# version of soil (prescribed heat and prescribed water - no state
# variables)
using MPI
using OrderedCollections
using StaticArrays
using Test

using CLIMAParameters
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

using ClimateMachine
using ClimateMachine.Land
using ClimateMachine.Land.SoilWaterParameterizations
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.DGMethods: BalanceLaw, LocalGeometry
using ClimateMachine.MPIStateArrays
using ClimateMachine.GenericCallbacks
using ClimateMachine.ODESolvers
using ClimateMachine.VariableTemplates
using ClimateMachine.SingleStackUtils
using ClimateMachine.BalanceLaws:
    BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state

@testset "Prescribed Models" begin
    ClimateMachine.init()
    FT = Float64

    function init_soil_water!(land, state, aux, localgeo, time) end

    soil_water_model = PrescribedWaterModel()
    soil_heat_model = PrescribedTemperatureModel()
    soil_param_functions = nothing
    m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)
    sources = ()
    m = LandModel(
        param_set,
        m_soil;
        source = sources,
        init_state_prognostic = init_soil_water!,
    )

    N_poly = 5
    nelem_vert = 10

    # Specify the domain boundaries
    zmax = FT(0)
    zmin = FT(-1)

    driver_config = ClimateMachine.SingleStackConfiguration(
        "LandModel",
        N_poly,
        nelem_vert,
        zmax,
        param_set,
        m;
        zmin = zmin,
        numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
    )

    t0 = FT(0)
    timeend = FT(60)
    dt = FT(1)

    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        ode_dt = dt,
    )
    mygrid = solver_config.dg.grid
    Q = solver_config.Q
    aux = solver_config.dg.state_auxiliary

    ClimateMachine.invoke!(solver_config;)

    t = ODESolvers.gettime(solver_config.solver)
    state_vars = SingleStackUtils.get_vars_from_nodal_stack(
        mygrid,
        Q,
        vars_state(m, Prognostic(), FT),
    )
    aux_vars = SingleStackUtils.get_vars_from_nodal_stack(
        mygrid,
        aux,
        vars_state(m, Auxiliary(), FT),
    )
    #Make sure it runs, and that there are no state variables, and only "x,y,z" as aux.
    @test t == timeend
    @test size(Base.collect(keys(aux_vars)))[1] == 3
    @test size(Base.collect(keys(state_vars)))[1] == 0
end


"""
Return leadfield for target location
"""
function calculate_specific_leadfield2(ldf::Leadfield, target::Coordinate; kwargs...)

    calculate_specific_leadfield2(ldf, find_location(ldf, target); kwargs...)
end

"""
Return leadfield for target location index
"""
function calculate_specific_leadfield2(ldf::Leadfield, target::Int; kwargs...)

    calculate_specific_leadfield2(ldf.L, target; kwargs...)
end

"""
Return leadfield for target location with null placed at supress location
"""
function calculate_specific_leadfield2(ldf::Leadfield, target::Union{Coordinate, Int}, supress::Union{Coordinate, Int}; bilateral::Bool=false, kwargs...)

    if bilateral & isa(supress, Coordinate)
        supress.x = -1 * supress.x
    elseif bilateral & isa(supress, Int)
        supress = find_location(ldf, Talairach(-1 * ldf.x[supress], ldf.y[supress], ldf.z[supress]))
    end

    target_ldf = calculate_specific_leadfield2(ldf, target; kwargs...)
    supress_ldf = calculate_specific_leadfield2(ldf, supress; kwargs...)

    hcat(target_ldf, supress_ldf)
end

"""
Return leadfield for target location with nulls placed in supress region
"""
function calculate_specific_leadfield2(ldf::Leadfield, target::Union{Coordinate, Int}, supress::Union{Coordinate, Int}, region_size::Real; bilateral::Bool=false, kwargs...)

    if bilateral & isa(supress, Coordinate)
        supress.x = -1 * supress.x
    elseif bilateral & isa(supress, Int)
        supress = find_location(ldf, Talairach(-1 * ldf.x[supress], ldf.y[supress], ldf.z[supress]))
    end

    target_ldf = calculate_specific_leadfield2(ldf, target; kwargs...)
    supress_ldf = calculate_specific_leadfield_region(ldf, supress, region_size; kwargs...)

    hcat(target_ldf, supress_ldf )
end

function calculate_specific_leadfield2(ldf::Leadfield, target::Union{Coordinate, Int}, supress::Array{EEG.Talairach,1}, region_size::Real;
    bilateral::Bool=false, keep_vecs::Int=3, kwargs...)

    target_ldf = calculate_specific_leadfield2(ldf, target; kwargs...)

    for s in supress

        supress_ldf = calculate_specific_leadfield_region(ldf, s, region_size; kwargs...)

        supress_ldf = reduce_leadfield(supress_ldf, keep_vecs, keep_dims = 0; kwargs...)

        target_ldf = hcat(target_ldf, supress_ldf)
    end

    return target_ldf
end

"""
Return leadfield for location at given index
"""
function calculate_specific_leadfield2{A <: AbstractFloat}(ldf::Array{A, 3}, target::Int; kwargs...)

    squeeze(ldf[target,:,:], 1)'
end

"""
Return leadfield for given region given index
"""
function calculate_specific_leadfield_region(ldf::Leadfield, target::Int, region::Real; bilateral::Bool=false, kwargs...)

    supress_ldf = calculate_specific_leadfield_region(ldf, Talairach(ldf.x[target], ldf.y[target], ldf.z[target]), region; kwargs...)
end

"""
Return leadfield for given region given coordinaates
"""
function calculate_specific_leadfield_region(ldf::Leadfield, target::Coordinate, region::Real; bilateral::Bool=false, kwargs...)

    to_supress = falses(size(ldf.x))

    if bilateral
        target.x = -1 * target.x
    end

    for loc in 1:length(ldf.x)

        if euclidean(target, [ldf.x[loc], ldf.y[loc], ldf.z[loc]]) < region
            to_supress[loc] = true
        end
    end
    Ls = ldf.L[to_supress, :, :]
    Ls = reshape(permutedims(Ls, [3, 2, 1]), size(ldf.L, 3), size(Ls, 1) * 3)
end

"""
Reduce leadfield dimensionality
"""
function reduce_leadfield{F <: AbstractFloat}(ldf::Array{F, 2}, keep_vecs; keep_dims = 3, kwargs...)
    hcat(ldf[:, 1:keep_dims], svdfact(ldf[:, keep_dims+1:end]).U[:, 1:keep_vecs])

end


using LinearAlgebra

import Tulip

include("LPBenchmark.jl")

"""
    run_tulip

Run benchmark for Tulip.
"""
function run_tulip(fname::String, fsol::String="")

    BLAS.set_num_threads(1)

    tulip = Tulip.Optimizer()

    # Read file
    load_problem!(tulip, fname)
    MOI.set(tulip, MOI.Silent(), false)
    MOI.set(tulip, MOI.NumberOfThreads(), 1)
    MOI.set(tulip, MOI.TimeLimitSec(), 10_000)
    MOI.set(tulip, MOI.RawParameter("BarrierIterationsLimit"), 500)

    # Solve instance
    run(tulip)
    return nothing
end

if abspath(PROGRAM_FILE) == @__FILE__
    run_tulip(ARGS[1])
end


module computer

export Context, io, run

using DelimitedFiles, Match

mutable struct Context
    program::Array{Int64,1}
    instruction::Int64
    relative_base::Int64
    debug::Bool
    buffer::String
end

function io()
    return (Channel{Int64}(Inf), Channel{Int64}(Inf))
end

function run(program::Array{Int64,1}, debug=false)
    stdin, stdout = io()
    run(program, stdin, stdout, debug)
end

function run(program::Array{Int64,1}, input::Int64, debug=false)
    stdin, stdout = io()
    put!(stdin, input)
    run(program, stdin, stdout, debug)
end

function run(program::Array{Int64,1}, stdin::Channel{Int64}, stdout::Channel{Int64}, debug=false)
    c = Context(program, 1, 0, debug, "")
    while execute(c, stdin, stdout); end
    return program
end

function execute(context::Context, stdin::Channel{Int64}, stdout::Channel{Int64})
    dbg(context, "@@")
    dbg(context, lpad(context.program[context.instruction], 8))
    op = rem(context.program[context.instruction], 100)
    @match op begin
        # ADD
        1 => begin
            dbg(context, rpad("  ADD", 8))
            x = rval(context, 1)
            y = rval(context, 2)
            wval(context, 3, x + y)
            context.instruction += 4
        end
        # MULTIPLY
        2 => begin
            dbg(context, rpad("  MULT", 8))
            x = rval(context, 1)
            y = rval(context, 2)
            wval(context, 3, x * y)
            context.instruction += 4
        end
        # INPUT
        3 => begin
            dbg(context, rpad("  STDIN", 8))
            val = take!(stdin)
            wval(context, 1, val)
            context.instruction += 2
        end
        # OUTPUT
        4 => begin
            dbg(context, rpad("  STDOUT", 8))
            val = rval(context, 1)
            put!(stdout, val)
            context.instruction += 2
        end
        # JUMP IF TRUE
        5 => begin
            dbg(context, rpad("  JMPT", 8))
            val = rval(context, 1)
            dest = rval(context, 2)
            val != 0 ?
                context.instruction = dest + 1 :
                context.instruction += 3
        end
        # JUMP IF FALSE
        6 => begin
            dbg(context, rpad("  JMPF", 8))
            val = rval(context, 1)
            dest = rval(context, 2)
            val == 0 ?
                context.instruction = dest + 1 :
                context.instruction += 3
        end
        # LESS THAN
        7 => begin
            dbg(context, rpad("  CMPLT", 8))
            x = rval(context, 1)
            y = rval(context, 2)
            x < y ? wval(context, 3, 1) : wval(context, 3, 0)
            context.instruction += 4
        end
        # EQUALS
        8 => begin
            dbg(context, rpad("  CMPEQ", 8))
            x = rval(context, 1)
            y = rval(context, 2)
            x == y ? wval(context, 3, 1) : wval(context, 3, 0)
            context.instruction += 4
        end
        # ADJUST RELATIVE BASE
        9 => begin
            dbg(context, rpad("  REL", 8))
            x = rval(context, 1)
            context.relative_base += x
            context.instruction += 2
        end
        # HALT
        99 => begin
            dbg(context, rpad("  HALT", 8))
            context.debug && println(context.buffer)
            context.buffer = ""
            return false
        end
        _ => throw(ErrorException("unknown op code: $op"))
    end
    context.debug && println(context.buffer)
    context.buffer = ""
    return true
end

function rval(context::Context, parameter::Int)
    program = context.program
    instruction = context.instruction
    rel = context.relative_base
    op = program[instruction]
    address = instruction + parameter
    mode = rem(div(op, ^(10, parameter + 1)), 10)
    dbg(context, lpad(address, 8))
    dbg(context, rpad(glyph(mode), 2))
    @match mode begin
        0 => begin
            idx = read(context, instruction + parameter)
            return read(context, idx + 1)
        end
        1 => return read(context, instruction + parameter)
        2 => begin
            offset = read(context, instruction + parameter)
            return read(context, rel + offset + 1)
        end
        _ => throw(ErrorException("unknown parameter mode: $mode"))
    end
end

function wval(context::Context, parameter::Int, value::Int64)
    program = context.program
    instruction = context.instruction
    rel = context.relative_base
    op = program[instruction]
    address = instruction + parameter
    mode = rem(div(op, ^(10, parameter + 1)), 10)
    dbg(context, lpad(address, 8))
    dbg(context, rpad(glyph(mode), 2))
    @match mode begin
        0 => begin
            idx = read(context, instruction + parameter)
            write(context, idx + 1, value)
        end
        1 => throw(ErrorException("invalid parameter mode: $mode"))
        2 => begin
            offset = read(context, instruction + parameter)
            write(context, rel + offset + 1, value)
        end
        _ => throw(ErrorException("unknown parameter mode: $mode"))
    end
end

function read(context::Context, index::Int64)
    return index > length(context.program) ? 0 : context.program[index]
end

function write(context::Context, index::Int64, value::Int64)
    if index > length(context.program)
        mem = zeros(Int64, index - length(context.program))
        context.program = cat(context.program, mem, dims=(1))
    end
    context.program[index] = value
end

function dbg(context, string)
    if context.debug
        context.buffer *= string
    end
end

function glyph(mode::Int)
    @match mode begin
        0 => return "x"
        1 => return "i"
        2 => return "r"
    end
end

function load_program(path::String)
    return vec(readdlm(path, ',', Int64))
end

end


info("Compare JuLIP EAM Implementation against ASE EAM Implementation")

info("Loading libraries...")
using PyCall, JuLIP, ASE
@pyimport ase.calculators.eam as eam

pot_file = Pkg.dir("JuLIP") * "/data/w_eam4.fs"

info("Generate the ASE potential")
eam4_ase = eam.EAM(potential=pot_file) |> ASECalculator

info("Generate low-accuracy JuLIP potential")
eam4_jl1 = EAM(pot_file)
info("Generate medium-accuracy JuLIP potential")
eam4_jl2 = EAM(pot_file; s = 1e-4)
info("Generate high-accuracy JuLIP potential")
eam4_jl3 = EAM(pot_file; s = 1e-6)

info("Generate test configurations")
at1 = rattle!(bulk(:W, cubic=true) * 3, 0.1)
at2 = deleteat!(bulk(:W, cubic=true) * 3, 1)
at1_ase = ASEAtoms(at1)
at2_ase = ASEAtoms(at2)

info("Test 1")
println("   Low Accuracy energy error:",
      (energy(eam4_ase, at1_ase) - energy(eam4_jl1, at1)) / length(at1))
println("Medium Accuracy energy error:",
      (energy(eam4_ase, at1_ase) - energy(eam4_jl2, at1)) / length(at1))
println("  High Accuracy energy error:",
      (energy(eam4_ase, at1_ase) - energy(eam4_jl3, at1)) / length(at1))

info("Test 2")
println("   Low Accuracy energy error:",
      (energy(eam4_ase, at2_ase) - energy(eam4_jl1, at2)) / length(at2))
println("Medium Accuracy energy error:",
      (energy(eam4_ase, at2_ase) - energy(eam4_jl2, at2)) / length(at2))
println("  High Accuracy energy error:",
      (energy(eam4_ase, at2_ase) - energy(eam4_jl3, at2)) / length(at2))