################################################# FILE DESCRIPTION #########################################################

# Graft allows users to refer to vertices externally, through arbitrary Julia types. The LabelMap is responsible
# for the resolution of these arbitrary objects into the internally used Integer indices.

# The label map is the datastructure that is used to realize two separate labelling schemes:
# 1. By default, vertices don't have labels, and are referred to by their internally used integer indices.
#    The IdentityLM type is used in this case. This labelling scheme is similar to the one used in LightGraphs.jl
#    and is confusing to use when vertices are removed etc. Sometimes its very inconvenient to think of vertices
#    as numbers! However this scheme incurs no performance overhead in label resolution.
#
# 2. Once the user assigns labels to all vertices, the DictLM datastructure is used to accomplish forward and
#    reverse translations. While both translations are done in constant time, a user defined vertex labelling
#    scheme can incur significant overheads in label translations, especially in large graphs.
#
# Since the type of the map can change depending on the nature of the operation, the modified label map is returned.
# Most label operations should therefore be called as:
# g.lmap = mutate!(g.lmap)

################################################# IMPORT/EXPORT #######################################################

export
# Types
setlabel!, relabel!, haslabel, decode, encode

"""
A type that forward maps labels into internally used vertex identifiers,
and reverse maps vertex identifiers into labels.
"""
abstract LabelMap{T}

################################################# CONSTRUCTORS ########################################################

LabelMap(nv::Int) = IdentityLM(nv)

LabelMap{T}(ls::Vector{T}) = DictLM(ls)

################################################# IDENTITYLM ##########################################################

"""
The default label map, that indicates the absence of meaningful vertex labels.
The usage of this type incurs zero overhead in label resolution.

Since vertices are referred to by their internally used indices, the usage
of this labelling scheme can be problematic when vertices are deleted.
"""
type IdentityLM <: LabelMap{Int}
   nv::Int
end

nv(x::IdentityLM) = x.nv

################################################# DICTLM ##############################################################

"""
This label map is used when vertices are assigned meaningful labels. This type uses
a dictionary to map labels onto vertex identifies, and a vector to map vertex
identifiers onto labels.

Labels can be of any user defined type.
"""
type DictLM{T} <: LabelMap{T}
   nv::Int
   fmap::Dict{T,VertexID}
   rmap::Vector{T}
end

""" Construct a label map from a list of labels """
function DictLM{T}(labels::AbstractVector{T})
   nv = length(labels)
   fmap = Dict(l => i for (i,l) in enumerate(labels))
   return DictLM{T}(nv, fmap, copy(labels))
end

"""
Construct a label map from the internally used vertex identifiers.
"""
DictLM(x::IdentityLM) = DictLM([i for i in 1 : nv(x)])

# Accessors
nv(x::DictLM) = x.nv
fmap(x::DictLM) = x.fmap
rmap(x::DictLM) = x.rmap

################################################# MSC ################################################################

###
# ELTYPE
###
Base.eltype(x::IdentityLM) = Int
Base.eltype{T}(x::DictLM{T}) = T

###
# ==
###
(==)(x::LabelMap, y::LabelMap) = encode(x) == encode(y)

###
# COPY
###
Base.copy(x::IdentityLM) = IdentityLM(nv(x))
Base.copy(x::DictLM) = DictLM(nv(x), copy(fmap(x)), copy(rmap(x)))

###
# DEEPCOPY
###
Base.deepcopy(x::IdentityLM) = IdentityLM(nv(x))
Base.deepcopy(x::DictLM) = DictLM(nv(x), deepcopy(fmap(x)), deepcopy(rmap(x)))

################################################# SETLABEL ############################################################

""" Remove vertex labels """
setlabel!(x::LabelMap) = IdentityLM(x.nv)

""" Set new labels for all vertices """
setlabel!(x::IdentityLM, ls::AbstractVector) = DictLM(ls)
setlabel!(x::DictLM, ls::AbstractVector) = DictLM(ls)

################################################# RELABEL SINGLE ######################################################

""" Change the label of a single vertex """
relabel!(x::IdentityLM, v::VertexID, l) = relabel!(DictLM(x), v, l)

function relabel!{T}(x::DictLM{T}, v::VertexID, l)
   l = convert(T, l)
   fmap(x)[l] = v
   rmap(x)[v] = l
   return x
end

################################################# RELABEL MUTLI #######################################################

""" Change labels for a list of vertices """
relabel!(x::IdentityLM, vs::VertexList, ls::AbstractVector) = relabel!(DictLM(x), vs, ls)

function relabel!{T}(x::DictLM{T}, vs::VertexList, ls::AbstractVector)
   ls = collect(T, ls)
   D = fmap(x)
   for i in eachindex(vs, ls)
      D[ls[i]] = vs[i]
   end
   rmap(x)[vs] = ls
   return x
end

################################################# HASLABEL ############################################################

""" Check if the input vertex label is valid """
haslabel(x::IdentityLM, v::VertexID) = 1 <= v <= nv(x)
haslabel(x::IdentityLM, l) = false

haslabel(x::DictLM, l) = haskey(fmap(x), l)

################################################# DECODE VERTEX #######################################################
""" Translate a vertex label into the internally used vertex identifier """
decode(x::IdentityLM, v::VertexID) = 1 <= v <= nv(x) ? v : error("Couldn't decode vertex label $v")

decode(x::IdentityLM, l) = error("Couldn't decode vertex label $l")

function decode(x::DictLM, l)
   get(fmap(x), l) do
      error("Couldn't decode vertex label $l")
   end
end

################################################# DECODE VERTICES ##########################################################

""" Translate a list of labels into the internally used vertex identifiers """
decode(x::IdentityLM, vs::VertexList) = vs
decode(x::DictLM, ls::AbstractVector) = [decode(x, l) for l in ls]

################################################# DECODE EDGE ##############################################################

""" Decode both vertices in a labelled edge """
decode(x::IdentityLM, e::EdgeID) = e
decode(x::DictLM, el::Pair) = EdgeID(decode(x, el.first), decode(x, el.second))

################################################# ENCODE VERTEX ############################################################

""" Map the input vertex identifier into the externally used label """
encode(x::IdentityLM, v::VertexID) = v
encode(x::DictLM, v::VertexID) = getindex(rmap(x), v)

################################################# ENCODE VERTICES ##########################################################

""" Map the input vertex identifier into its externally used label """
encode(x::IdentityLM, vs::VertexList) = collect(vs)
encode(x::DictLM, vs::VertexList) = getindex(rmap(x), vs)

################################################# FETCH ALL LABELS #########################################################

""" Map the input list of vertex identifiers into their externally used labels """
encode(x::IdentityLM) = collect(1 : x.nv)
encode(x::DictLM) = copy(rmap(x))

################################################# ENCODE EDGE ##############################################################

""" Encode both vertices in an edge identifier """
encode(x::IdentityLM, e::EdgeID) = e
encode{T}(x::DictLM{T}, e::EdgeID) = Pair{T,T}(encode(x, e.first), encode(x, e.second))

################################################# ENCODE EDGES #############################################################

""" Encode all edges in an input edge list """
encode(x::IdentityLM, es::EdgeList) = es

encode(x::DictLM, es::EdgeList) = [encode(x, e) for e in es]

################################################# ADDVERTEX ################################################################

"""
Add an unlabelled vertex to the graph. This method throws
an error, when used on a labelled graph.
"""
function addvertex!(x::IdentityLM)
   x.nv += 1
   return x
end

addvertex!(x::DictLM) = error("Please supply a label")

"""
Add a labelled vertex to the graph. If used on an
unlabelled graph, this method will assign the internally
used vertex identifiers as labels to existing vertices.
"""
function addvertex!(x::IdentityLM, l)
   if l == nv(x) + 1
      addvertex!(x)
   else
      addvertex!(DictLM(x), l)
   end
end

function addvertex!{T}(x::DictLM{T}, l)
   l = convert(T, l)
   x.nv += 1
   push!(rmap(x), l)
   fmap(x)[l] = nv(x)
   return x
end

################################################# RMVERTEX ###################################################################

""" Remove a list of vertices from the graph. """
function rmvertex!(x::IdentityLM, vs)
   x.nv -= length(vs)
   return x
end

function rmvertex!(x::DictLM, vs)
   F = fmap(x)
   R = rmap(x)

   ls = encode(x, vs)

   # Remove labels from Forward map
   for l in ls
      delete!(F, l)
   end

   # Remove labels from Reverse map
   deleteat!(R, vs)

   # Redirect labels in Forward map to new indexes
   for v in minimum(vs) : length(R)
      F[R[v]] = v
   end

   # Decrement Label Map size
   x.nv -= length(vs)

   return x
end

################################################# SUBGRAPH ####################################################################

""" Retrieve a subset of vertex labels """
subgraph(x::IdentityLM, vlist::VertexList) = IdentityLM(length(vlist))

function subgraph(x::DictLM, vlist::VertexList)
   DictLM(rmap(x)[vlist])
end


module TestExamplesTree1D

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "tree_1d_dgsem")

# Start with a clean environment: remove Trixi output directory if it exists
outdir = "out"
isdir(outdir) && rm(outdir, recursive=true)

@testset "TreeMesh1D" begin

# Run basic tests
@testset "Examples 1D" begin
  # Linear scalar advection
  include("test_tree_1d_advection.jl")

  # Burgers
  include("test_tree_1d_burgers.jl")

  # Hyperbolic diffusion
  include("test_tree_1d_hypdiff.jl")

  # Compressible Euler
  include("test_tree_1d_euler.jl")

  # Compressible Euler Multicomponent
  include("test_tree_1d_eulermulti.jl")

  # MHD
  include("test_tree_1d_mhd.jl")

  # MHD Multicomponent
  include("test_tree_1d_mhdmulti.jl")

  # Compressible Euler with self-gravity
  include("test_tree_1d_eulergravity.jl")
end

# Coverage test for all initial conditions
@testset "Tests for initial conditions" begin
  # Linear scalar advection
  @trixi_testset "elixir_advection_extended.jl with initial_condition_sin" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_extended.jl"),
      l2   = [0.00017373554109980247],
      linf = [0.0006021275678165239],
      maxiters = 1,
      initial_condition = Trixi.initial_condition_sin)
  end

  @trixi_testset "elixir_advection_extended.jl with initial_condition_constant" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_extended.jl"),
      l2   = [2.441369287653687e-16],
      linf = [4.440892098500626e-16],
      maxiters = 1,
      initial_condition = initial_condition_constant)
  end

  @trixi_testset "elixir_advection_extended.jl with initial_condition_linear_x" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_extended.jl"),
      l2   = [1.9882464973192864e-16],
      linf = [1.4432899320127035e-15],
      maxiters = 1,
      initial_condition = Trixi.initial_condition_linear_x,
      boundary_conditions = Trixi.boundary_condition_linear_x,
      periodicity=false)
  end

  @trixi_testset "elixir_advection_extended.jl with initial_condition_convergence_test" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_extended.jl"),
      l2   = [6.1803596620800215e-6],
      linf = [2.4858560899509996e-5],
      maxiters = 1,
      initial_condition = initial_condition_convergence_test,
      boundary_conditions = BoundaryConditionDirichlet(initial_condition_convergence_test),
      periodicity=false)
  end
end

@testset "Displaying components 1D" begin
  @test_nowarn include(joinpath(EXAMPLES_DIR, "elixir_advection_amr.jl"))

  # test both short and long printing formats
  @test_nowarn show(mesh); println()
  @test_nowarn println(mesh)
  @test_nowarn display(mesh)

  @test_nowarn show(equations); println()
  @test_nowarn println(equations)
  @test_nowarn display(equations)

  @test_nowarn show(solver); println()
  @test_nowarn println(solver)
  @test_nowarn display(solver)

  @test_nowarn show(solver.basis); println()
  @test_nowarn println(solver.basis)
  @test_nowarn display(solver.basis)

  @test_nowarn show(solver.mortar); println()
  @test_nowarn println(solver.mortar)
  @test_nowarn display(solver.mortar)

  @test_nowarn show(solver.volume_integral); println()
  @test_nowarn println(solver.volume_integral)
  @test_nowarn display(solver.volume_integral)

  @test_nowarn show(semi); println()
  @test_nowarn println(semi)
  @test_nowarn display(semi)

  @test_nowarn show(summary_callback); println()
  @test_nowarn println(summary_callback)
  @test_nowarn display(summary_callback)

  @test_nowarn show(amr_controller); println()
  @test_nowarn println(amr_controller)
  @test_nowarn display(amr_controller)

  @test_nowarn show(amr_callback); println()
  @test_nowarn println(amr_callback)
  @test_nowarn display(amr_callback)

  @test_nowarn show(stepsize_callback); println()
  @test_nowarn println(stepsize_callback)
  @test_nowarn display(stepsize_callback)

  @test_nowarn show(save_solution); println()
  @test_nowarn println(save_solution)
  @test_nowarn display(save_solution)

  @test_nowarn show(analysis_callback); println()
  @test_nowarn println(analysis_callback)
  @test_nowarn display(analysis_callback)

  @test_nowarn show(alive_callback); println()
  @test_nowarn println(alive_callback)
  @test_nowarn display(alive_callback)

  @test_nowarn println(callbacks)

  # Check whether all output is suppressed if the summary, analysis and alive
  # callbacks are set to the TrivialCallback(). Modelled using `@test_nowarn`
  # as basis.
  let fname = tempname()
    try
      open(fname, "w") do f
        redirect_stderr(f) do
          trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_extended.jl"),
                        summary_callback=TrivialCallback(),
                        analysis_callback=TrivialCallback(),
                        alive_callback=TrivialCallback())
        end
      end
      @test isempty(read(fname, String))
    finally
      rm(fname, force=true)
    end
  end
end

@testset "Additional tests in 1D" begin
  @testset "compressible Euler" begin
    eqn = CompressibleEulerEquations1D(1.4)

    @test isapprox(Trixi.entropy_thermodynamic([1.0, 2.0, 20.0], eqn), 1.9740810260220094)
    @test isapprox(Trixi.entropy_math([1.0, 2.0, 20.0], eqn), -4.935202565055024)
    @test isapprox(Trixi.entropy([1.0, 2.0, 20.0], eqn), -4.935202565055024)

    @test isapprox(energy_total([1.0, 2.0, 20.0], eqn), 20.0)
    @test isapprox(energy_kinetic([1.0, 2.0, 20.0], eqn), 2.0)
    @test isapprox(energy_internal([1.0, 2.0, 20.0], eqn), 18.0)
  end
end

@trixi_testset "Nonconservative terms in 1D (linear advection)" begin
  # Same setup as docs/src/adding_new_equations/nonconservative_advection.md

  # Define new physics
  using Trixi
  using Trixi: AbstractEquations, get_node_vars

  # Since there is no native support for variable coefficients, we use two
  # variables: one for the basic unknown `u` and another one for the coefficient `a`
  struct NonconservativeLinearAdvectionEquation <: AbstractEquations{1 #= spatial dimension =#,
                                                                     2 #= two variables (u,a) =#}
  end

  Trixi.varnames(::typeof(cons2cons), ::NonconservativeLinearAdvectionEquation) = ("scalar", "advectionvelocity")

  Trixi.default_analysis_integrals(::NonconservativeLinearAdvectionEquation) = ()

  # The conservative part of the flux is zero
  Trixi.flux(u, orientation, equation::NonconservativeLinearAdvectionEquation) = zero(u)

  # Calculate maximum wave speed for local Lax-Friedrichs-type dissipation
  function Trixi.max_abs_speed_naive(u_ll, u_rr, orientation::Integer, ::NonconservativeLinearAdvectionEquation)
    _, advectionvelocity_ll = u_ll
    _, advectionvelocity_rr = u_rr

    return max(abs(advectionvelocity_ll), abs(advectionvelocity_rr))
  end

  # We use nonconservative terms
  Trixi.have_nonconservative_terms(::NonconservativeLinearAdvectionEquation) = Val(true)

  function flux_nonconservative(u_mine, u_other, orientation,
                                equations::NonconservativeLinearAdvectionEquation)
    _, advectionvelocity = u_mine
    scalar, _            = u_other

    return SVector(advectionvelocity * scalar, zero(scalar))
  end

  # Create a simulation setup
  using Trixi
  using OrdinaryDiffEq

  equation = NonconservativeLinearAdvectionEquation()

  # You can derive the exact solution for this setup using the method of
  # characteristics
  function initial_condition_sine(x, t, equation::NonconservativeLinearAdvectionEquation)
    x0 = -2 * atan(sqrt(3) * tan(sqrt(3) / 2 * t - atan(tan(x[1] / 2) / sqrt(3))))
    scalar = sin(x0)
    advectionvelocity = 2 + cos(x[1])
    SVector(scalar, advectionvelocity)
  end

  # Create a uniform mesh in 1D in the interval [-π, π] with periodic boundaries
  mesh = TreeMesh(-Float64(π), Float64(π), # min/max coordinates
                  initial_refinement_level=4, n_cells_max=10^4)

  # Create a DGSEM solver with polynomials of degree `polydeg`
  volume_flux  = (flux_central, flux_nonconservative)
  surface_flux = (flux_lax_friedrichs, flux_nonconservative)
  solver = DGSEM(polydeg=3, surface_flux=surface_flux,
                volume_integral=VolumeIntegralFluxDifferencing(volume_flux))

  # Setup the spatial semidiscretization containing all ingredients
  semi = SemidiscretizationHyperbolic(mesh, equation, initial_condition_sine, solver)

  # Create an ODE problem with given time span
  tspan = (0.0, 1.0)
  ode = semidiscretize(semi, tspan);

  summary_callback = SummaryCallback()
  analysis_callback = AnalysisCallback(semi, interval=50)
  callbacks = CallbackSet(summary_callback, analysis_callback);

  # OrdinaryDiffEq's `solve` method evolves the solution in time and executes
  # the passed callbacks
  sol = solve(ode, Tsit5(), abstol=1.0e-6, reltol=1.0e-6,
              save_everystep=false, callback=callbacks);

  @test analysis_callback(sol).l2 ≈ [0.00029610274971929974, 5.573684084938363e-6]
end

# Clean up afterwards: delete Trixi output directory
@test_nowarn rm(outdir, recursive=true)

end # TreeMesh1D

end # module


function SQ._collect(tbl, g::SQ.GroupbyNode)
    src = copy(tbl)
    predicate_aliases = Dict{Expr, Symbol}()
    pre_groupby!(src, g, predicate_aliases)
    return groupby(src, g.args, predicate_aliases)
end

"""
"""
function pre_groupby!(src, q, predicate_aliases)
    i = 1 # count number of predicate aliases
    for (h, arg) in zip(q.helpers, q.args)
        is_predicate, f, arg_fields = SQ.parts(h)
        if is_predicate
            predicate_alias = Symbol("pred_$i")
            predicate_aliases[arg] = predicate_alias
            # include the result of applying the predicate as a new column
            apply!(src, SQ.SelectHelper(predicate_alias, f, arg_fields), src)
            i += 1
        end
    end
end

"""
"""
function groupby{T<:AbstractTable}(
    src::T, groupbys, predicate_aliases
)::Grouped{T}
    # obtain the field names of the groupbys (either the names of the selected
    # columns or the predicate aliases given to groupby predicates)
    # TODO: store this information in the `groupbys` field of the result
    groupby_fields = map(x->isa(x, Symbol) ? x : predicate_aliases[x], groupbys)
    group_indices = build_group_indices(src, groupby_fields)
    group_levels = build_group_levels(group_indices, length(groupby_fields))
    return Grouped(
        src,
        group_indices,
        GroupLevels(group_levels),
        groupbys,
        predicate_aliases
    )
end

function build_group_levels(group_indices, ngroupbys)
    joint_group_levels = collect(keys(group_indices))
    group_levels = Vector{Vector}(ngroupbys)
    for j in 1:ngroupbys
        group_levels[j] = [ level[j] for level in joint_group_levels ]
    end
    map!(unique, group_levels)
    return group_levels
end

function build_group_indices(tbl, groupbys)
    row_itr = eachrow(tbl, groupbys...)
    # TODO: type this more strongly
    group_indices = Dict{Any, Vector{Int}}()
    grow_indices!(group_indices, row_itr)
    return group_indices
end

@noinline function grow_indices!(group_indices, row_itr)
    for (i, row) in enumerate(row_itr)
        group_level = row
        if haskey(group_indices, group_level)
            push!(group_indices[group_level], i)
        else
            group_indices[group_level] = [i]
        end
    end
    return group_indices
end


using Random
using POMDPSimulators
using PMDPs.LP

@testset "HistoryReplayer.jl" begin
    mg, me = simple_mdps()
    pp = PMDPs.pp(mg)

    # History replayer takes either full SimulationHistory structure
    simhistory = PMDPs.simulate_trace(mg, MersenneTwister(1))
    # Or more limited structure containing Abstract Arrays of NamedTuples that contain state and info fields.
    trace = simple_short_trace(pp)

    @test isa(PMDPs.HistoryReplayer(mg, simhistory), PMDPs.PMDP)
    @test isa(PMDPs.HistoryReplayer(mg, trace), PMDPs.PMDP)

    hrpl = PMDPs.HistoryReplayer(mg, trace)

    # Test some basic properties
    @test POMDPs.actions(hrpl) == POMDPs.actions(mg)

    # Test gen method
    s = trace[1].s
    (sp, r, info) = POMDPs.gen(hrpl, s, 1000., MersenneTwister(99))
    @test sp == trace[2].s
    @test r == 0.
    @test info == trace[1].info

    # s = PMDPs.State(SA[2,3], 5, SA[true, true])

    # test that the new trace from replayer matches the input trace
    hrpl_sim = PMDPs.HistoryReplayer(mg, simhistory)
    t = PMDPs.simulate_trace(hrpl_sim, MersenneTwister(321))
    @test t == simhistory

    # TODO: The following should likely be tested in tests of the specific policies

    hrec = HistoryRecorder(max_steps = PMDPs.selling_period_end(mg), rng = MersenneTwister(4321))

    # test VI with HistoryReplayer
    policy = PMDPs.get_VI_policy(me)
    hᵥ = simulate(hrec, hrpl, policy)
    @test length(hᵥ) == length(trace)
    @test sum(collect(hᵥ[:r])) > 0.

    # test MCTS with historyReplayer
    planner = PMDPs.get_MCTS_planner(mg; params_mcts = Dict(:rng=>MersenneTwister(1)))
    hₘ = simulate(hrec, hrpl, planner)
    @test length(hₘ) == length(trace)
    @test sum(collect(hₘ[:r])) > 0.

    # test hindsight with historyReplayer
    hindsight = PMDPs.LP.get_MILP_hindsight_policy(mg, trace)
    hₕ = simulate(hrec, hrpl, hindsight)
    @test length(hₕ) == length(trace)
    @test sum(collect(hₕ[:r])) > 0.

    # test flatrate with historyReplayer
    # R, U = PMDPs.optimize_flatrate_policy(mg, [trace, trace])
    flatrate = PMDPs.get_flatrate_policy(mg, [trace, trace])
    hᵣ = simulate(hrec, hrpl, flatrate)
    @test length(hᵣ) == length(trace)
    @test sum(collect(hᵣ[:r])) > 0.

end


mutable struct GridsContinuumParams
    epsilon::Float64
    num_generative_samples::Int64
    num_bridge_samples::Int64
    horizon_limit::Int64
    vals_along_axis::Int64 # Number of points along axis for interpolation
    move_coefficient::Float64
    time_coefficient::Float64
    beta::Float64
    move_amt::Float64
    move_std::Float64
    cworlds::Dict{Int64,CWorld} # The exemplar continuum worlds
end

function continuum_parse_params(filename::AbstractString)

    params_key = TOML.parsefile(filename)

    return GridsContinuumParams(params_key["EPSILON"],
                                params_key["NUM_GEN_SAMPLES"],
                                params_key["NUM_BRIDGE_SAMPLES"],
                                params_key["HORIZON_LIMIT"],
                                params_key["AXIS_VALS"],
                                params_key["MOVE_COEFFICIENT"],
                                params_key["TIME_COEFFICIENT"],
                                params_key["BETA"],
                                params_key["MOVE_AMT"],
                                params_key["MOVE_STD"],
                                Dict{Int64, CWorld}())

end

const GridsContinuumStateType = CMSSPState{Int64, Vec2}
const GridsContinuumActionType = CMSSPAction{Int64, Vec2}
const GridsContinuumContextType = Dict{Tuple{Int64, Int64}, Vec2}

mutable struct GridsContinuumContextSet
    curr_timestep::Int64
    curr_context::GridsContinuumContextType
    future_contexts::Vector{GridsContinuumContextType}
end

struct ContinuumDummyBookkeeping
    flag::Bool
end
ContinuumDummyBookkeeping() = ContinuumDummyBookkeeping(false)

struct ContinuumDummyMetadata
    flag::Bool
end
Base.zero(::Type{ContinuumDummyMetadata}) = ContinuumDummyMetadata(false)

GridsContinuumContextSet() = GridsContinuumContextSet(0, GridsContinuumContextType(), Vector{GridsContinuumContextType}(undef, 0))

const GridsContinuumCMSSPType = CMSSP{Int64, Vec2, Int64, Vec2, GridsContinuumParams, GridsContinuumContextSet}
const GridsContinuumMDPType = ModalFinHorMDP{Int64, Vec2, Vec2, GridsContinuumParams}
const GridsContinuumInfhorMDPType = ModalInfHorMDP{Int64, Vec2, Vec2, GridsContinuumParams}
const GridsContinuumBridgeSample = BridgeSample{Vec2, ContinuumDummyMetadata}
const GridsContinuumOLV = OpenLoopVertex{Int64, Vec2, Int64, ContinuumDummyMetadata}
const GridsContinuumGraphTracker = GraphTracker{Int64, Vec2, Int64, ContinuumDummyMetadata, ContinuumDummyBookkeeping, RNG} where {RNG <: AbstractRNG}
const GridsContinuumSolverType = HHPCSolver{Int64, Vec2, Int64, ContinuumDummyMetadata, ContinuumDummyBookkeeping, RNG} where {RNG <: AbstractRNG}

# Const mode values
const CONTINUUM_MODES = [1, 2, 3, 4]
const CONTINUUM_GOAL_MODE = 4

# Discrete MDP - 4 states, 2 actions
# 1 -- 2 -- 4
# |         |
# 3 --------|
function get_grids_continuum_mdp()

    T = zeros(4, 2, 4)
    T[1, 1, 2] = 1.0
    T[1, 2, 3] = 1.0
    T[2, 1, 4] = 1.0
    T[3, 1, 4] = 1.0

    # No cost for mode switches
    R = zeros(4, 2)

    return TabularMDP(T, R, 1.0)
end

function get_grids_continuum_actions(params::GridsContinuumParams)

    actions = Vector{GridsContinuumActionType}(undef, 0)

    # First enter the mode actions
    push!(actions, GridsContinuumActionType(1, 1))
    push!(actions, GridsContinuumActionType(2, 2))

    # Then enter the control actions
    # Use actions of CWorld 1
    cworld_actions = params.cworlds[1].actions

    for (i, cwa) in enumerate(cworld_actions)
        idx = i+2
        push!(actions, GridsContinuumActionType(cwa, idx))
    end

    return actions
end

function create_continuum_cmssp(params::GridsContinuumParams, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    actions = get_grids_continuum_actions(params)
    switch_mdp = get_grids_continuum_mdp()
    temp_goal_state = GridsContinuumStateType(CONTINUUM_GOAL_MODE, sample_continuum_state(rng))

    return GridsContinuumCMSSPType(actions, CONTINUUM_MODES, switch_mdp, temp_goal_state, params, GridsContinuumContextSet())
""
end

function POMDPs.isterminal(mdp::GridsContinuumMDPType, relative_state::Vec2)
    return norm(relative_state) < mdp.params.epsilon
end

function POMDPs.isterminal(cmssp::GridsContinuumCMSSPType, state::GridsContinuumStateType)
    return state.mode == CONTINUUM_GOAL_MODE
end

function HHPC.get_relative_state(mdp::Union{GridsContinuumMDPType, GridsContinuumInfhorMDPType},
                                 source::Vec2, target::Vec2)
    return (source - target)
end

function sample_continuum_state(rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}
    d = Distributions.Uniform(0.0, 1.0)
    return Vec2(rand(rng, d), rand(rng, d))
end

function generate_start_state(cmssp::GridsContinuumCMSSPType, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}
    return GridsContinuumStateType(1, sample_continuum_state(rng))
end

function POMDPs.reward(mdp::GridsContinuumMDPType, state::Vec2, action::Vec2, statep::Vec2)

    # Copy over reward of underlying Continuum World
    delta_state = norm(statep - state)
    rew = -1.0*(mdp.params.move_coefficient*delta_state + mdp.params.time_coefficient)
    return rew

end

function POMDPs.convert_s(::Type{Vector{Float64}}, s::Vec2,
                          mdp::GridsContinuumMDPType) where V <: AbstractVector{Float64}
    return [s[1], s[2]]
end

function POMDPs.generate_s(mdp::GridsContinuumMDPType, state::Vec2, action::Vec2, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    params = mdp.params
    cworld = params.cworlds[mdp.mode]
    return generate_s(cworld, state, action, rng)
end

function POMDPs.generate_sr(mdp::Union{GridsContinuumMDPType, GridsContinuumInfhorMDPType},
                            state::Vec2, action::Vec2, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    statep = generate_s(mdp, state, action, rng)
    rew = reward(mdp, state, action, statep)
    return (statep, rew)
end

function HHPC.expected_reward(mdp::Union{GridsContinuumMDPType, GridsContinuumInfhorMDPType}, state::Vec2,
                              action::Vec2, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    params = mdp.params
    cworld = params.cworlds[mdp.mode]
    avg_reward = 0.0

    for i = 1:params.num_generative_samples
        statep = generate_s(cworld, state, action, rng)
        avg_reward += POMDPs.reward(mdp, state, action, statep)
    end

    avg_reward /= params.num_generative_samples
    return avg_reward
end

# Will never actually be called because any vertices in grid 4 satisfy the goal condition
# Just return the same vertex and the current goal idx
function HHPC.generate_goal_vertex_set!(cmssp::GridsContinuumCMSSPType, vertex::GridsContinuumOLV,
                                        graph_tracker::GridsContinuumGraphTracker, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}
    return [], [graph_tracker.curr_goal_idx]
end

# Just generate a new set of neighbours at each timestep
# So vertices_to_add has vertices and nbrs_to_add is empty
function HHPC.generate_bridge_vertex_set!(cmssp::GridsContinuumCMSSPType, vertex::GridsContinuumOLV,
                                          mode_pair::Tuple{Int64, Int64}, graph_tracker::GridsContinuumGraphTracker,
                                          action::Int64, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    context_set = cmssp.curr_context_set
    future_contexts = context_set.future_contexts # Vector of contexts
    curr_timestep = context_set.curr_timestep
    params = cmssp.params

    vertices_to_add = Vector{GridsContinuumOLV}(undef, 0)

    @assert !isempty(future_contexts)
    avg_samples = convert(Int64, ceil(cmssp.params.num_bridge_samples/length(future_contexts)))

    # @show avg_samples
    # @show length(future_contexts)

    for (hor, context) in enumerate(future_contexts)

        # Iterate over (M1,M2) -> point
        for (mpair, switch_point) in context

            # Generate bridge point
            if mpair == mode_pair

                d1 = Distributions.Uniform(switch_point[1] - params.epsilon/2.0, switch_point[1] + params.epsilon/2.0)
                d2 = Distributions.Uniform(switch_point[2] - params.epsilon/2.0, switch_point[2] + params.epsilon/2.0)

                for i = 1:avg_samples

                    bp1 = clamp(rand(rng, d1), 0.0, 1.0)
                    bp2 = clamp(rand(rng, d2), 0.0, 1.0)
                    cont_state = Vec2(bp1, bp2)

                    bridge_sample = GridsContinuumBridgeSample(cont_state, cont_state,
                                        TPDistribution([curr_timestep + hor], [1.0]), zero(ContinuumDummyMetadata))
                    new_bridge_vtx = GridsContinuumOLV(mode_pair[1], mode_pair[2], bridge_sample, action)

                    push!(vertices_to_add, new_bridge_vtx)
                end
            end
        end
    end

    # Return all new vertices and no existing neighbours
    return vertices_to_add, []
end

function HHPC.generate_next_valid_modes(cmssp::GridsContinuumCMSSPType, mode::Int64)

    context_set = cmssp.curr_context_set
    future_contexts = context_set.future_contexts

    next_actions_modes = Set{Tuple{Int64, Int64}}()

    for (hor, context) in enumerate(future_contexts) # Vector of dictionaries

        for (mpair, switch_point) in context # Iterate over current dictionary

            if mpair[1] == mode # Switch from current mode

                next_mode = mpair[2]

                for action = 1:2
                    if cmssp.modeswitch_mdp.T[mode, action, next_mode] > 0.0
                        push!(next_actions_modes, (action, next_mode))
                    end
                end
            end
        end
    end

    return next_actions_modes
end

# Dummy function as vertices will always be regenerated
function HHPC.update_vertices_with_context!(cmssp::GridsContinuumCMSSPType, graph_tracker::GridsContinuumGraphTracker,
                                            curr_timestep::Int64)
    return
end

function HHPC.update_next_target!(cmssp::GridsContinuumCMSSPType, solver::GridsContinuumSolverType, curr_timestep::Int64)
    return
end

function set_start_context_set!(cmssp::GridsContinuumCMSSPType,  rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    params = cmssp.params
    init_context = GridsContinuumContextType()

    # Generate initial switch point regions
    for m1 in cmssp.modes
        for m2 in cmssp.modes
            for a in cmssp.mode_actions
                if cmssp.modeswitch_mdp.T[m1, a ,m2] > 0.0
                    switch_point = sample_continuum_state(rng)
                    init_context[(m1, m2)] = switch_point
                end
            end
        end
    end

    future_contexts = Vector{GridsContinuumContextType}(undef, 0)
    temp_context = init_context

    # Sample within a radius of twice the epsilon value
    r = 2.0*params.epsilon

    for idx = 1:params.horizon_limit

        next_context = GridsContinuumContextType()

        # Iterate over current switch points and sample around them
        for (mpair, switch_point) in temp_context

            theta = rand(rng, Distributions.Uniform(0.0, 2.0*pi))
            new_pt1 = clamp(switch_point[1] + r*cos(theta), 0.0, 1.0)
            new_pt2 = clamp(switch_point[2] + r*sin(theta), 0.0, 1.0)
            new_point = Vec2(new_pt1, new_pt2)

            next_context[mpair] = new_point
        end

        push!(future_contexts, next_context)
        temp_context = next_context

    end

    cmssp.curr_context_set = GridsContinuumContextSet(0, init_context, future_contexts)
end

function HHPC.get_bridging_action(vertex::GridsContinuumOLV)
    return vertex.bridging_action
end

function HHPC.display_context_future(context_set::GridsContinuumContextSet, future_timestep::Int64)
    @assert future_timestep > context_set.curr_timestep
    println(context_set.future_contexts[future_timestep - context_set.curr_timestep])
end

function HHPC.update_context_set!(cmssp::GridsContinuumCMSSPType, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    curr_context_set = cmssp.curr_context_set
    params = cmssp.params

    # Update timestep
    curr_context_set.curr_timestep += 1

    # First update the current context based on the first future context
    curr_context = curr_context_set.curr_context
    next_context = curr_context_set.future_contexts[1]
    new_curr_context = GridsContinuumContextType()

    for (mpair, curr_switch_point) in curr_context

        next_switch_point = next_context[mpair]

        # Choose an interpolated point between curr and next chosen switch point
        interp_frac = rand(rng, Distributions.Uniform(0.5, 1.0))
        new_switch_point = curr_switch_point + interp_frac*(next_switch_point - curr_switch_point)
        new_curr_context[mpair] = new_switch_point
    end

    # Update new current context
    curr_context_set.curr_context = new_curr_context

    # Pop first future context
    popfirst!(curr_context_set.future_contexts)

    # Add a new future context
    pre_final_context = curr_context_set.future_contexts[end]
    new_final_context = GridsContinuumContextType()

    r = params.epsilon/2.0

    for (mpair, switch_point) in pre_final_context

        theta = rand(rng, Distributions.Uniform(0.0, 2.0*pi))
        new_pt1 = clamp(switch_point[1] + r*cos(theta), 0.0, 1.0)
        new_pt2 = clamp(switch_point[2] + r*sin(theta), 0.0, 1.0)
        new_point = Vec2(new_pt1, new_pt2)

        new_final_context[mpair] = new_point
    end

    push!(curr_context_set.future_contexts, new_final_context)

    @assert length(curr_context_set.future_contexts) == params.horizon_limit

end

function HHPC.simulate_cmssp!(cmssp::GridsContinuumCMSSPType, state::GridsContinuumStateType, a::Union{Nothing, GridsContinuumActionType},
                              curr_timestep::Int64, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    # Default - will delete when timeout removed
    timeout = false

    curr_context_set = cmssp.curr_context_set
    curr_context = curr_context_set.curr_context

    curr_mode = state.mode
    curr_cont_state = state.continuous
    params = cmssp.params

    if typeof(a.action) <: Int64

        # Action is mode switch
        next_mode = findfirst(x -> x > 0.0, cmssp.modeswitch_mdp.T[curr_mode, a.action, :])
        reward = cmssp.modeswitch_mdp.R[curr_mode, a.action]

        # Check if it can be done based on current context
        for (mpair, switch_point) in curr_context

            if mpair == (curr_mode, next_mode)

                if norm(switch_point - curr_cont_state) < 2.0*params.epsilon
                    @info "Successful mode switch to grid ",next_mode
                    next_state = GridsContinuumStateType(next_mode, switch_point)
                    return (next_state, reward, false, timeout)
                else
                    @info "Mode switch to grid ",next_mode, "failed "
                    next_state = state
                    return (next_state, reward, true, timeout)
                end
            end
        end

        @warn "Attempted mode switch to ", next_mode," unavailable from context"
        next_state = curr_state
        return (next_state, reward, true, timeout)
    else
        # control action within mode
        cworld = params.cworlds[curr_mode]
        new_cont_state = generate_s(cworld, state.continuous, a.action, rng)
        reward = POMDPs.reward(cworld, state.continuous, a.action, new_cont_state)
        next_state = GridsContinuumStateType(curr_mode, new_cont_state)
        return (next_state, reward, false, timeout)
    end
end


export predict

# similar functions as in apply_tree_fn.jl
# however, for the functions in this file we do not require the featurepools to match
# i.e. DecisionTrees.assert_consistent_features could evaluate to false

function predict(x::Tree, f::DataFrame)
    asrt = DecisionTrees.assert_consistent_features(x.featurepools, f)
    if asrt
        return predictConsistentFeatures(x, f)
    else
        @warn("DTM: Applying model even though model features do not match data features. Results might be inconsistent, you may want to review the predictions.")
        fp = x.featurepools
        res = apply_tree_by_leaf_inconsistent_features(x.rootnode, f, fp)
        return res
    end
    # return apply_tree_by_leaf(x.rootnode,f)
end

function predict(x::BoostedTree, f::DataFrame;prroduceEstAndLeafMatrices::Bool=false)
    asrt = DecisionTrees.assert_consistent_features(x.featurepools, f)
    if asrt
        return predictConsistentFeatures(x::BoostedTree, f::DataFrame;prroduceEstAndLeafMatrices=prroduceEstAndLeafMatrices)
    else
        @show DecisionTrees.assert_consistent_features(x.featurepools, f)
        @error("in the works....this is not yet done for boosting")
    end
end

function apply_tree_by_leaf_inconsistent_features(t::Union{Leaf,Node{UInt8},Node{UInt16}}, features::DataFrame, fp)
    nobs = size(features, 1)
    fit = zeros(Float64, nobs)
    leafnrs = zeros(Int, nobs)
    idx = collect(1:nobs)
    apply_tree_by_leaf_inconsistent_features_iteration!(idx, t, features, fit, leafnrs, fp)
    return fit, leafnrs
end

function apply_tree_by_leaf_inconsistent_features_iteration!(idx::Vector{Int}, t::Node{T}, features::DataFrame, fit::Vector{Float64}, leaf::Vector{Int}, fp) where T <: Unsigned
    if length(idx) == 0
        return nothing
    end
    subset = t.subset
    subsetWhichIsStringOrFloatVector = fp[t.featid_new_positive][subset]
    idxlInteger, idxrInteger = lrIndicesInconsitentFeatures(idx, features[t.featid_new_positive], subsetWhichIsStringOrFloatVector)
    apply_tree_by_leaf_inconsistent_features_iteration!(idxlInteger, t.left, features, fit, leaf, fp)
    apply_tree_by_leaf_inconsistent_features_iteration!(idxrInteger, t.right, features, fit, leaf, fp)
    return nothing
end

function apply_tree_by_leaf_inconsistent_features_iteration!(idx, t::Leaf, features::DataFrame, fit, leaf, fp)
    @inbounds for i in idx
        leaf[i] = t.id
        fit[i] = t.fitted
    end
    return nothing
end

"""
function lrIndicesInconsitentFeatures(trnidx::Vector{Int},f,subset::Array)
modifies trnidx to 'become' the left child index, returns r
l and r are defined by subset. l corresponds to the observations 'in' subset
this function is used for new/unseen features which have a different pool definition than the model data
"""
function lrIndicesInconsitentFeatures(trnidx::Vector{Int}, f, subset::Array)
    if eltype(f.pool) <: Number
        return lrIndicesForNumericalVariableInconsitentFeatures(trnidx, f, subset)
    else
        return lrIndicesDefaultInconsitentFeatures(trnidx, f, subset)
    end

end

"""
function lrIndicesDefaultInconsitentFeatures(trnidx::Vector{Int},f,subset::Array)
modifies trnidx to 'become' the left child index, returns r
l and r are defined by subset. l corresponds to the observations 'in' subset
this function is used for new/unseen features which have a different pool definition than the model data
"""
function lrIndicesDefaultInconsitentFeatures(trnidx::Vector{Int}, f, subset::Array)
	l = Vector{Int}(undef, 0)
	r = Vector{Int}(undef, 0)
	sizehint!(l, length(trnidx))
	sizehint!(r, length(trnidx))
	for i in trnidx
		@inbounds thisref = f[i]
		if in(thisref, subset)
			push!(l, i)
		else
			push!(r, i)
		end
	end
	return l, r
end

"""
function lrIndicesForNumericalVariableInconsitentFeatures(trnidx::Vector{Int},f,subset::Array)
modifies trnidx to 'become' the left child index, returns r
l and r are defined by subset. l corresponds to the observations 'in' subset
this function is used for new/unseen features which have a different pool definition than the model data
"""
function lrIndicesForNumericalVariableInconsitentFeatures(trnidx::Vector{Int}, f, subset::Array)
	l = Vector{Int}(undef, 0)
	r = Vector{Int}(undef, 0)
	sizehint!(l, length(trnidx))
    sizehint!(r, length(trnidx))
    subsetEnd = subset[end]
	for i in trnidx
		@inbounds thisref = f[i]
        if thisref <= subsetEnd
            push!(l, i)
        else
            push!(r, i)
        end
	end
	return l, r
end


using GeoData, RasterDataSources, Test, Dates, NCDatasets, ArchGDAL

# Too big to test on CI
# if !haskey(ENV, "CI")
#     @testset "load WorldClim Weather" begin
#         # Weather time-series
#         dates = (Date(2001), Date(2002))
#         ser = series(WorldClim{Weather}, (:prec,); date=dates)
#         ser[Date(2001, 1)][:prec]
#         A = geoarray(WorldClim{Weather}, :prec; date=DateTime(2001, 05), mappedcrs=EPSG(4326))
#     end
# end

@testset "load WorldClim Climate" begin
    # Weather time-series
    ser = GeoSeries(WorldClim{Climate}, :prec; res="10m", month=Jan:March, mappedcrs=EPSG(4326))
    # Select Australia, using regular lat/lon selectors
    A = ser[month=Jan]
    @test A isa GeoArray
    A[Y(Between(-10, -45)), X(Between(110, 160))]
    st = GeoStack(WorldClim{Climate}, (:prec, :tmax); month=1)
    st[:prec]
    @test st isa GeoStack
end

@testset "load WorldClim BioClim" begin
    A = GeoArray(WorldClim{BioClim}, :Bio_1; mappedcrs=EPSG(4326))
    A[Y(Between(-10, -45)), X(Between(110, 160))]
    @test A isa GeoArray
    st = stack(WorldClim{BioClim}, (1, 2))
    st[:bio1]
    @test st isa GeoStack
    @test A isa GeoArray
end

@testset "load CHELSA BioClim" begin
    A = GeoArray(CHELSA{BioClim}, 1; mappedcrs=EPSG(4326))
    @test GeoData.name(A) == :bio1
    st = GeoStack(CHELSA{BioClim}, (:bio1, :BIO2))
    @test keys(st) == (:bio1, :bio2)
    @test A isa GeoArray
    @test st isa GeoStack
    @test st[:bio2] isa GeoArray
end

@testset "load EarthEnv HabitatHeterogeneity" begin
    A = GeoArray(EarthEnv{HabitatHeterogeneity}, :cv; mappedcrs=EPSG(4326))
    A[Y(Between(-10, -45)), X(Between(110, 160))]
    st = GeoStack(EarthEnv{HabitatHeterogeneity}, (:cv, :evenness))
    @test A isa GeoArray
    @test st isa GeoStack
    @test st[:evenness] isa GeoArray
end

@testset "load EarthEnv LandCover" begin
    A = GeoArray(EarthEnv{LandCover}, 2; mappedcrs=EPSG(4326))
    @test GeoData.name(A) == :evergreen_broadleaf_trees
    A[Y(Between(-10, -45)), X(Between(110, 160))]
    @test A isa GeoArray
    st = GeoStack(EarthEnv{LandCover}, (:evergreen_broadleaf_trees, :deciduous_broadleaf_trees); mappedcrs=EPSG(4326))
    keys(st) = (:evergreen_broadleaf_trees, :deciduous_broadleaf_trees)
    @test st isa GeoStack
end

@testset "load ALWB" begin
    A = GeoArray(ALWB{Deciles,Day}, :rain_day; date=DateTime(2019, 10, 19))
    @test crs(A) == EPSG(4326)
    A = GeoArray(ALWB{Values,Day}, :ss_pct; date=DateTime(2019, 10, 19))
    @test crs(A) == EPSG(4326)
    st = GeoStack(ALWB{Values,Day}, (:s0_pct, :ss_pct); date=DateTime(2019, 10, 19))
    @test crs(st) == EPSG(4326)
    @test crs(st[:s0_pct]) == EPSG(4326)
    dates = DateTime(2019, 10, 19), DateTime(2021, 11, 20)
    s = GeoSeries(ALWB{Values,Day}, (:s0_pct, :ss_pct); date=dates)
    s[1]
    @test A isa GeoArray
    @test st isa GeoStack
    @test s isa GeoSeries
end

# Obscure .Z format may not work on windows
if Sys.islinux()
    @testset "load AWAP" begin
        A = GeoArray(AWAP, :rainfall; date=DateTime(2019, 10, 19))
        @test crs(A) == EPSG(4326)
        st = GeoStack(AWAP; date=DateTime(2019, 10, 19), resize=crop)
        @test crs(st) == EPSG(4326)
        dates = DateTime(2019, 09, 19), DateTime(2019, 11, 19)
        s = GeoSeries(AWAP; date=dates, resize=crop)
        # s = GeoSeries(AWAP; date=dates, resize=resample, crs=EPSG(4326)) TODO: all the same
        # s = GeoSeries(AWAP; date=dates, resize=extend) TODO: this is slow !!!
        @test crs(s[1][:solar]) == EPSG(4326)
        @test A isa GeoArray
        @test st isa GeoStack
        @test s isa GeoSeries
    end
end


"""
$(TYPEDEF)
$(TYPEDFIELDS)

"""
struct Float <: Types.AbstractDelayed
    arg::Types.AbstractDelayed
end

Base.float(arg::Types.AbstractDelayed) = Float(arg)


using HexEdit

# construct type
hexedit = HexEd("test.bin")
# dump whole file
dump!(hexedit, 0x00)
# write foobar beginning at offset 0x04 using hex string as arg
edit!(hexedit, "0x666f6f626172", 4)
# dump 16 bytes beginning at offset 0x04
dump!(hexedit, 4, 16)
# write foobar beginning at offset 0x04 using ascii as arg
edit!(hexedit, "foobar", 0x04)
# dump 16 bytes beginning at offset 0x04
dump!(hexedit, 0x04, 16)
# find foobar
fbl = find!(hexedit, "foobar", 0x00)
if fbl != nothing
    println("found foobar at $(uppercase(string(fbl, base=16, pad=8)))")
end
# find hex signature
fbl = find!(hexedit, "0x07ae0971")
if fbl != nothing
    println("found hex signature at $(uppercase(string(fbl, base=16, pad=8)))")
end


struct Chain_Report
    id::Chain_ID
    final_hmm::BHMM
    test_lh::AbstractFloat
    naive_lh::AbstractFloat
    final_delta::AbstractFloat
    state_run_lengths::AbstractVector{AbstractFloat}
    convergence_values::Chains
    convergence_diagnostic::ChainDataFrame
    converged::Bool
end

function Base.show(io::IO, report::Chain_Report)
    nominal_dict=Dict(0=>"th",1=>"st",2=>"nd",3=>"rd",4=>"th",5=>"th")
    haskey(nominal_dict,report.id.order) ? (nom_str=nominal_dict[report.id.order]) : (nom_str="th")
    printstyled(io, stdout, "BHMM EM Chain Results\n"; bold=true)
    println(io, "$(report.id.K)-state, $(report.id.order)$nom_str order BHMM")
    println(io, "Trained on observation set \"$(report.id.obs_id)\"")
    report.test_lh > report.naive_lh ? (lh_str="greater"; lh_color=:green) : (lh_str="less"; lh_color=:red)
    printstyled(io, "Test logP(O|θ): $(report.test_lh), $lh_str than the naive model's $(report.naive_lh)\n"; color=lh_color)
    println(io, "Replicate $(report.id.replicate)")
    report.converged ? println(io, "Converged with final step δ $(report.final_delta)") : println(io, "Failed to converge!")
    println(io, " ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶")
    println(io, "Last θ:")
    show(report.final_hmm)
    println(io, "State Mean Feature Length (bp)")
    for i in 1:length(report.state_run_lengths)
        println(io, "K$i: $(report.state_run_lengths[i].*float(report.id.order+1))")
    end
    println(io, " ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶ ̶̶ ̶ ̶ ̶")
    zidx=findfirst(!iszero,report.convergence_values["logP(O|θ)"].data)[1]
    lh_vec=Vector([report.convergence_values["logP(O|θ)"].data...][zidx:end])

    lh_plot=lineplot([zidx:length(lh_vec)+zidx-1 ...],lh_vec;title="Chain likelihood evolution", xlabel="Training iterate", xlim=(0,length(lh_vec)+zidx+1), ylim= (floor(minimum(lh_vec),sigdigits=2),ceil(maximum(lh_vec),sigdigits=2)), name="logP(O|θ)")
    lineplot!(lh_plot,[report.naive_lh for i in 1:length(lh_vec)], color=:magenta,name="naive")
    show(io, lh_plot)
    println(io, "\n")

    k1vec=Vector([report.convergence_values["K1"].data...])
    k_plot=lineplot(k1vec, title="State p(Auto) evolution", xlabel="Training iterate", ylabel="prob", name="K1",ylim=(0,1),xlim=(0,length(k1vec)))
    for k in 2:report.id.K
        kvec=Vector([report.convergence_values["K$k"].data...])
        lineplot!(k_plot,kvec,name="K$k")
    end
    show(io,k_plot)
    println(io)

    printstyled(io,"\nConvergence Diagnostics\n",bold=true)
    if report.convergence_diagnostic.name == "short"
        printstyled(io,"Convergence diagnostics unavailable for chains <10 steps!\n", color=:yellow)
    elseif report.convergence_diagnostic.name == "error"
        printstyled(io,"Convergence diagnostics errored! Zeros in autotransition matrix?\n", color=:red)
    else
        all(Bool.(report.convergence_diagnostic.nt.stationarity)) && all(Bool.(report.convergence_diagnostic.nt.test)) ? printstyled(io, "Likelihood and autotransition probabilites converged and passing tests.\n", color=:green) : printstyled(io, "Not all parameters converged or passing tests!\n",color=:red)
        display(report.convergence_diagnostic)
    end
end

function latex_report(report::Chain_Report)

end

function report_chains(chains::Dict{Chain_ID,Vector{EM_step}}, test_sets::Dict{String,Vector{LongSequence{DNAAlphabet{2}}}}; naive_hmm=BHMM(ones(1,1),[Categorical(4)]))
    job_ids=[id for (id,chain) in chains] # make list of chain_ids
    partitions=unique([id.obs_id for id in job_ids]) #get list of unique partitions represented among chain_ids
    naive_order = Integer(log(4,length(naive_hmm.B[1].support))-1)
    for partition in partitions #check that all partitions are available for testing and make sure codes for naive hmm are generated by adding job_ids
        !(partition in keys(test_sets)) && throw(ArgumentError("Observation set $partition required by chains for testing not present in test sets!"))
        push!(job_ids, Chain_ID(partition,length(naive_hmm.a),naive_order, 1))
    end

    code_dict=code_job_obs(job_ids, test_sets) #code all the obs sets that are necessary ot perform our tests
    naive_lhs=Dict{String,Float64}() #construct dict of naive likelihoods for all partitions to be tested prior to individual tests
    @showprogress 1 "Testing naive hmm..." for ((partition,order),obs_set) in code_dict
        order==naive_order && (naive_lhs[partition]=obs_lh_given_hmm(code_dict[(partition,naive_order)],naive_hmm))
    end

    reports=Dict{Chain_ID,Chain_Report}()
    @showprogress 1 "Testing chains..." for (id, chain) in chains
        chains_array=zeros(length(chain),1+id.K)
        for (n,step) in enumerate(chain)
            chains_array[n,1]=step.log_p
            chains_array[n,2:end]=get_diagonal_array(step.hmm)
        end
        convergence_values=Chains(chains_array,["logP(O|θ)", ["K"*string(i) for i in 1:id.K]...])
        if length(chain)>=10
            try
                convergence_diagnostic=heideldiag(convergence_values)[1]
            catch
                convergence_diagnostic=ChainDataFrame("error",(;:a=>[1.]))
            end
        else
            convergence_diagnostic=ChainDataFrame("short",(;:a=>[1.]))
        end

        id.K > 1 ? (mean_run_lengths=sim_run_lengths(get_diagonal_array(chain[end].hmm),1000)) : (mean_run_lengths=[Inf])
        test_lh=obs_lh_given_hmm(code_dict[(id.obs_id,id.order)], chain[end].hmm)
        reports[id]=Chain_Report(id, chain[end].hmm, test_lh, naive_lhs[id.obs_id], chain[end].delta, mean_run_lengths, convergence_values, convergence_diagnostic, chain[end].converged)
    end

    return reports
end

function get_diagonal_array(hmm::BHMM)
    k = length(hmm.a)
    diagonal = zeros(k)
    @inbounds for i in 1:k
        diagonal[i] = hmm.A[i,i]
    end
    return diagonal
end

#function to simulate run lengths for vector of diagonal values
function sim_run_lengths(diagonal_value::AbstractArray, samples::Integer)
    mean_run_lengths = zeros(length(diagonal_value))
    for (i, value) in enumerate(diagonal_value)
        if value == 0.
            mean_run_lengths[i]=0.
        elseif value ==1.
            mean_run_lengths[i]=Inf
        else
            runlengths = zeros(Integer, samples)
            for s in 1:samples
                run = true
                runlength = 0
                while run
                    runlength += 1
                    if rand(1)[1] > value
                        run = false
                    end
                end
                runlengths[s] = runlength
            end
            mean_run_lengths[i] = mean(runlengths)
        end
    end
    return mean_run_lengths
end


export OffPolicy

"""
    OffPolicy(π_target::P, π_behavior::B) -> OffPolicy{P,B}
"""
Base.@kwdef struct OffPolicy{P,B} <: AbstractPolicy
    π_target::P
    π_behavior::B
end

(π::OffPolicy)(obs) = π.π_behavior(obs)

function RLBase.update!(π::OffPolicy, t::AbstractTrajectory)
    experience = extract_experience(t, π)
    isnothing(experience) || update!(π, experience)
end

function RLBase.update!(π::OffPolicy{<:VBasedPolicy}, transitions::NamedTuple)
    # ??? define a `get_batch_prob` function for efficiency
    weights = [
        get_prob(π.π_target, (state = s,), a) / get_prob(π.π_behavior, (state = s,), a)
        for (s, a) in zip(transitions.states, transitions.actions)
    ]  # TODO: implement iterate interface for (SubArray of) CircularArrayBuffer
    experience = merge(transitions, (weights = weights,))
    update!(π.π_target, experience)
end

RLBase.extract_experience(t::AbstractTrajectory, π::OffPolicy) =
    extract_experience(t, π.π_target)


export plot

if Pkg.installed("MATLAB") == Nothing()

	function missingPackage()

		warn("Install Julia package MATLAB to enable TensorMesh3D mesh plotting.")
		return

	end

	plot(mesh::TensorMesh3D)                      = missingPackage()
	plot(mesh::TensorMesh3D, u::Array{Int64,1})   = missingPackage()
	plot(mesh::TensorMesh3D, u::Array{Int64,3})   = missingPackage()
	plot(mesh::TensorMesh3D, u::Array{Float64,1}) = missingPackage()
	plot(mesh::TensorMesh3D, u::Array{Float64,3}) = missingPackage()

else

	using MATLAB

	function plot(mesh::TensorMesh3D; f = 1)

		v = diag(getVolume(mesh))
		u = log10(v ./ minimum(v)) .+ 1.0

		plot(mesh, u; f = f)

		return

	end

	function plot(mesh::TensorMesh3D, u::Array{Int64,1}; f = 1)

		u = reshape(u, (mesh.n[1], mesh.n[2], mesh.n[3]))

		plot(mesh, u; f = f)

	end

	function plot(mesh::TensorMesh3D, u::Array{Float64,1}; f = 1)

		u = reshape(u, (mesh.n[1], mesh.n[2], mesh.n[3]))

		plot(mesh, u; f = f)

	end

	function plot(mesh::TensorMesh3D, u::Array{Int64,3}; f = 1)

		if (mesh.n[1],mesh.n[2],mesh.n[3]) != size(u)
			error("Incompatible array sizes")
		end

		x,y,z = getNodalAxes(mesh)

		# set path to Matlab function plotTensorMesh3D.m which resides in the same directory like this Julia function
		(dname,fname) = splitdir(@__FILE__())
		mxcall(:addpath, 0, dname)

		# plot
		mxcall(:plotTensorMesh3D, 0, f, x, y, z, u)

		return

	end

	function plot(mesh::TensorMesh3D, u::Array{Float64,3}; f = 1)

		if (mesh.n[1],mesh.n[2],mesh.n[3]) != size(u)
			error("Incompatible array sizes")
		end

		x,y,z = getNodalAxes(mesh)

		# set path to Matlab function plotTensorMesh3D.m which resides in the same directory like this Julia function
		(dname,fname) = splitdir(@__FILE__())
		mxcall(:addpath, 0, dname)

		# plot
		mxcall(:plotTensorMesh3D, 0, f, x, y, z, u)

		return

	end

end


using Mux, JSON

function todict(s::AbstractString)
    d = Dict{String,String}()
    for pair in split(s,"&")
        a = split(pair,"=")
        d[a[1]] = a[2]
    end
    d
end

function readresults(annualEmissions, results, firstyear, lastyear)
    years = iyear(firstyear):iyear(lastyear)
    res = results[years]
    temp = getfield.(res, :Temp_global)
    GAS3 = [:CO2, :CH4, :N2O]
    conc = getfield.(res, :Concentration)
    concentrations = Dict(g => get.(conc, g, 0.0) for g in GAS3)
    emissions = Dict(g => get(annualEmissions, g, 0.0)[years] for g in GAS3)
    return Dict(
            "temperature" => temp,
            "concentrations" => concentrations,
            "emissions" => emissions
    )
end

function runclimatemodel(req)
    cccdata = JSON.parse(transcode(String, req[:data]))
    firstyear = cccdata["firstyear"]
    lastyear = cccdata["lastyear"]
    # if cccdata["firstyear"] != 2010 || cccdata["lastyear"] != 2100
    #     throw("Received emission data for $(cccdata["firstyear"])-$(cccdata["lastyear"]), " *
    #         "but currently only 2010-2100 is supported.")
    # end
    rcp = "RCP3PD"
    annualEmissions = getscenario(rcp)
    annualEmissions[:CO2][iyear(firstyear):iyear(lastyear)] = 12/44*(cccdata["emissions"]["FossilCO2"] + cccdata["emissions"]["OtherCO2"])
    annualEmissions[:CH4][iyear(firstyear):iyear(lastyear)] = cccdata["emissions"]["CH4"] .+ 270    # natural background emissions: 270 MtCH4
    annualEmissions[:N2O][iyear(firstyear):iyear(lastyear)] = cccdata["emissions"]["N2O"] .+ 10.7   # natural background emissions: 10.7 MtN (TAR p253)
    results, p = solveclimate(annualEmissions, usecache=true, lambda=cccdata["climatesensitivity"]/3.7, rcp=rcp)
    printresults(firstyear:10:lastyear, results, p, annualEmissions, rcp)
    return Dict(
        :body => JSON.json(readresults(annualEmissions, results, firstyear, lastyear)),
        :headers => Dict("Content-Type" => "application/json")
    )
end

function printPOSTinfo(req)
    println("\n$(req[:method])\n")
    display(req[:headers])
    println("\n")
    d = JSON.parse(transcode(String, req[:data]))
    display(d)
    JSON.print(d, 4)
    "hello kitty! :)"
    #"
method: $(req[:method])
data: $(transcode(String, req[:data]))
headers: $(req[:headers])"
end

global webserver

function startserver()
    staticroute = route("static", files(@__DIR__), Mux.notfound())
    #POSTroute = page("/user/:user", req -> "
Hello, $(req[:params][:user])!
")
    #POSTroute = method("POST", respond("POST request received."))
    #POSTroute = branch(req -> req[:method] == "POST", req -> printPOSTinfo(req))
    POSTroute = branch(req -> req[:method] == "POST", req -> runclimatemodel(req))

    serverstarted = isdefined(ClimateCalculator, :webserver)
    defaults = Mux.stack(Mux.todict, Mux.splitquery, Mux.toresponse)
    # @app webserver = (defaults, Mux.basiccatch, POSTroute, staticroute, Mux.notfound())
    @app webserver = (defaults, POSTroute, staticroute, Mux.notfound())
    !serverstarted && serve(webserver)    # only run this once, modify the line with @app test = ... instead
    println("Open http://localhost:8000/static/UI_layout.html in your web browser.")
    # login to NAS as admin, do "sudo -i" to become root
    # cd /volume1/Sync/julia-1.3.1
    # nohup bin/julia -e 'using ClimateCalculator; @sync startserver(); wait(Condition())' > serverlog.txt 2>&1 &
    # the @sync prefix might not be needed anymore
    # nohup bin/julia -e 'using ClimateCalculator; @sync startserver(); wait(Condition())' > /dev/null 2>&1 &
    # ps ax | grep julia
    # kill procnumber
end


# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# EAGO
# A development environment for robust and global optimization
# See https://github.com/PSORLab/EAGO.jl
#############################################################################
# src/eago_optimizer/types/incremental.jl
# A type-stable wrapper with for optimizers used by EAGO to enable bridging and
# incremental loading. This is taylored to the internal routines used by EAGO.jl
# so methods may be specialized by optimizer types and error checking is often
# avoided.
#############################################################################

#=
mutable struct IncrementalCache{S <: MOI.AbstractOptimizer} <: MOI.AbstractOptimizer}
end
=#

"""
    Incremental{Q,S}

A type-stable cache used to wrapper for an optimizer that enables incremental
modification of solvers that don't inherently suppport this. Explicitly checks
support of MOI functionality used in EAGO. For `Q = Val{true}`, the subsolver
supports incremental loading. For `Q = Val{false}`, the subsolver does not.
"""
mutable struct Incremental{S <: MOI.AbstractOptimizer} <: MOI.AbstractOptimizer
    optimizer::MOIB.LazyBridgeOptimizer{S}
    cache::MOIB.LazyBridgeOptimizer{MOIU.CachingOptimizer{S,MOIU.Model{Float64}}}
end
function Incremental(m::S) where {S <: MOI.AbstractOptimizer}
    b = MOIB.full_bridge_optimizer(m, Float64)
    cache = MOIB.full_bridge_optimizer(MOIU.CachingOptimizer(MOIU.Model{Float64}(), m), Float64)
    return Incremental{S}(b, cache)
end

function MOI.copy_to(model::Incremental{S}, src::MOI.ModelLike) where S <: MOI.AbstractOptimizer
    return MOI.Utilities.default_copy_to(model, src, copy_names)
end

_is_incremental(x) = false
_get_storage(d::Incremental{S}) where S = _is_incremental(S) ? d.optimizer : d.cache

# Set attributes
MOI.set(d::Incremental, ::MOI.ObjectiveFunction{VI}, f::VI)   = (MOI.set(_get_storage(d), MOI.ObjectiveFunction{VI}(), f); nothing)
MOI.set(d::Incremental, ::MOI.ObjectiveFunction{SAF}, f::SAF) = (MOI.set(_get_storage(d), MOI.ObjectiveFunction{SAF}(), f); nothing)
MOI.set(d::Incremental, ::MOI.ObjectiveFunction{SQF}, f::SQF) = (MOI.set(_get_storage(d), MOI.ObjectiveFunction{SQF}(), f); nothing)

function MOI.set(d::Incremental, ::MOI.ObjectiveSense, s)
    MOI.set(_get_storage(d), MOI.ObjectiveSense(), s)
    return
end
function MOI.set(d::Incremental{S}, ::MOI.Silent, s) where S <: MOI.AbstractOptimizer
    if _is_incremental(S)
        MOI.set(d.optimizer, MOI.Silent(), s)
    else
        MOI.set(d.cache.model.optimizer, MOI.Silent(), s)
    end
    return
end

function MOI.set(d::Incremental, ::MOI.VariablePrimalStart, v::VI, x)
    MOI.set(_get_storage(d), MOI.VariablePrimalStart(), v, x)
    return
end

MOI.set(d::Incremental, ::MOI.NLPBlock, ::Nothing) = nothing
function MOI.set(d::Incremental, ::MOI.NLPBlock, s)
    MOI.set(_get_storage(d), MOI.NLPBlock(), s)
    return
end

function MOI.set(d::Incremental, p::MOI.RawOptimizerAttribute, s)
    MOI.set(_get_storage(d), p, s)
    return
end

# Add variable/constraint
function MOI.add_variable(d::Incremental)
    MOI.add_variable(_get_storage(d))::VI
end

MOI.add_constraint(d::Incremental, f::VI, s::LT) = MOI.add_constraint(_get_storage(d), f, s)::CI{VI,LT}
MOI.add_constraint(d::Incremental, f::VI, s::GT) = MOI.add_constraint(_get_storage(d), f, s)::CI{VI,GT}
MOI.add_constraint(d::Incremental, f::VI, s::ET) = MOI.add_constraint(_get_storage(d), f, s)::CI{VI,ET}
MOI.add_constraint(d::Incremental, f::VI, s::IT) = MOI.add_constraint(_get_storage(d), f, s)::CI{VI,IT}
MOI.add_constraint(d::Incremental, f::VI, s::ZO) = MOI.add_constraint(_get_storage(d), f, s)::CI{VI,ZO}
MOI.add_constraint(d::Incremental, f::VI, s::MOI.Integer) = MOI.add_constraint(_get_storage(d), f, s)::CI{VI,MOI.Integer}

MOI.add_constraint(d::Incremental, f::SAF, s::LT) = MOI.add_constraint(_get_storage(d), f, s)::CI{SAF,LT}
MOI.add_constraint(d::Incremental, f::SAF, s::GT) = MOI.add_constraint(_get_storage(d), f, s)::CI{SAF,GT}
MOI.add_constraint(d::Incremental, f::SAF, s::ET) = MOI.add_constraint(_get_storage(d), f, s)::CI{SAF,ET}
MOI.add_constraint(d::Incremental, f::SAF, s::IT) = MOI.add_constraint(_get_storage(d), f, s)::CI{SAF,IT}

MOI.add_constraint(d::Incremental, f::SQF, s::LT) = MOI.add_constraint(_get_storage(d), f, s)::CI{SQF,LT}
MOI.add_constraint(d::Incremental, f::SQF, s::GT) = MOI.add_constraint(_get_storage(d), f, s)::CI{SQF,GT}
MOI.add_constraint(d::Incremental, f::SQF, s::ET) = MOI.add_constraint(_get_storage(d), f, s)::CI{SQF,ET}
MOI.add_constraint(d::Incremental, f::SQF, s::IT) = MOI.add_constraint(_get_storage(d), f, s)::CI{SQF,IT}

# Delete
function MOI.delete(d::Incremental, ci::CI{VI,T}) where T <: Union{LT,GT,ET,IT,MOI.Integer}
    MOI.delete(_get_storage(d), ci)
    return
end
function MOI.delete(d::Incremental, ci::CI{SAF,LT})
     MOI.delete(_get_storage(d), ci)
     return
end

# Set modifications
function MOI.set(d::Incremental, ::MOI.ConstraintSet, ci::CI{VI,T}, s::T) where T <: Union{LT,GT,ET,IT}
     MOI.set(_get_storage(d), MOI.ConstraintSet(), ci, s)
     return
end

# Get attributes
function MOI.get(d::Incremental{S}, ::MOI.TerminationStatus) where S
    MOI.get(d.optimizer, MOI.TerminationStatus())::MOI.TerminationStatusCode
end
function MOI.get(d::Incremental{S}, ::MOI.PrimalStatus) where S
    MOI.get(d.optimizer, MOI.PrimalStatus())::MOI.ResultStatusCode
end
function MOI.get(d::Incremental{S}, ::MOI.DualStatus) where S
    MOI.get(d.optimizer, MOI.DualStatus())::MOI.ResultStatusCode
end
function MOI.get(d::Incremental{S}, ::MOI.RawStatusString) where S
    MOI.get(d.optimizer, MOI.RawStatusString())::MOI.String
end

function MOI.get(d::Incremental{S}, ::MOI.ObjectiveBound) where S
    MOI.get(d.optimizer, MOI.ObjectiveBound())::Float64
end
function MOI.get(d::Incremental{S}, ::MOI.ObjectiveValue) where S
    MOI.get(d.optimizer, MOI.ObjectiveValue())::Float64
end
function MOI.get(d::Incremental{S}, ::MOI.DualObjectiveValue) where S
    MOI.get(d.optimizer, MOI.DualObjectiveValue())::Float64
end

function MOI.get(d::Incremental{S}, ::MOI.VariablePrimal, vi::VI) where S
    MOI.get(d.optimizer, MOI.VariablePrimal(), vi)::Float64
end

const SAF_CI_TYPES = Union{CI{SAF,LT},CI{SAF,GT},CI{SAF,ET}}
function MOI.get(d::Incremental{S}, ::MOI.ConstraintPrimal, ci::SAF_CI_TYPES) where S
    MOI.get(d.optimizer, MOI.ConstraintPrimal(), ci)::Float64
end

const SQF_CI_TYPES = Union{CI{SQF,LT},CI{SQF,GT},CI{SQF,ET}}
function MOI.get(d::Incremental{S}, ::MOI.ConstraintPrimal, ci::SQF_CI_TYPES) where S
    MOI.get(d.optimizer, MOI.ConstraintPrimal(), ci)::Float64
end

function MOI.get(d::Incremental{S}, ::MOI.ConstraintDual, ci::Union{CI{VI,LT},CI{VI,GT}}) where S
    MOI.get(d.optimizer, MOI.ConstraintDual(), ci)::Float64
end
function MOI.get(d::Incremental{S}, ::MOI.ResultCount) where S
    MOI.get(d.optimizer, MOI.ResultCount())::Int
end
function MOI.get(d::Incremental{S}, n::MOI.SolverName) where S
    MOI.get(d.optimizer, n)::String
end

# define optimize!
function MOI.optimize!(d::Incremental{S}) where S
    MOI.optimize!(_get_storage(d))
    return
end

function MOI.empty!(d::Incremental{S}) where S
    MOI.empty!(_get_storage(d))
    return
end

MOI.is_empty(d::Incremental{S}) where S = MOI.is_empty(_get_storage(d))


module Parser

using CombinedParsers
using TextParse

# @syntax subterm = Either{Rational{Int}}(Any[TextParse.Numeric(Int)]);
# @syntax for parenthesis in subterm
#     mult         = evaluate |> join(subterm, CharIn("*/"), infix=:prefix )
#     @syntax term = evaluate |> join(mult,    CharIn("+-"), infix=:prefix )
#     Sequence(2,'(',term,')')
# end;

end


using MPIMeasurements
using PyPlot

include("config.jl")

scanner = MPIScanner(conf)
daq = getDAQ(scanner)

params = toDict(daq.params)

params["studyName"] = "Test"
params["studyDescription"] = "A very cool measurement"
params["scannerOperator"] = "Tester"
params["dfStrength"] = [1e-3]
params["dfWaveform"] = "SINE"
params["acqNumFrames"] = 10
params["acqNumAverages"] = 10

u, uSlowADC = measurement(daq, params, controlPhase=true)
#filename = measurement(daq, params, MDFStore, controlPhase=true)

figure(1)
clf()
plot(vec(u))
savefig("images/im1.png")

scanner = MPIScanner(conf) #what???
daq = getDAQ(scanner)
params["dfWaveform"] = "TRIANGLE"
u, uSlowADC = measurement(daq, params, controlPhase=false)
figure(1)
clf()
plot(vec(u))
savefig("images/im2.png")


@testset "event" begin
    seed!(42)
    y = rand(1:10, 10)
    l = PoissonLikelihood(3.0)
    @test AGP.treat_labels!(y, l) == y
    @test_throws ErrorException AGP.treat_labels!(rand(10), l)
end


function tdms_header!(TDMS::TDMSbuf, io::IO, v::Integer)
  fastskip(io, 4)
  TDMS.flags  = fastread(io, UInt32)
  fastskip(io, 4)
  TDMS.nsos   = fastread(io, UInt64) + LEAD_IN_LENGTH
  TDMS.rdos   = fastread(io, UInt64) + LEAD_IN_LENGTH
  TDMS.n_ch   = fastread(io, UInt32) - 0x00000002
  l_path      = fastread(io, UInt32)
  fastskip(io, 4 + l_path)
  nprops      = fastread(io, UInt32)

  for i = 1:nprops
    # Get variable name and code
    L         = fastread(io, UInt32)
    name      = String(fastread(io, L))

    # Add dictionaries as needed
    D = TDMS.hdr
    if occursin(".", name)
      k = String.(split(name, "."))
      L = length(k)
      for j in 1:L-1
        if haskey(D, k[j])
          D = D[k[j]]
        else
          D[k[j]] = Dict{String, Any}()
          D = D[k[j]]
        end
      end
      name = k[L]
    end
    pcode     = fastread(io, UInt32)
    (pcode == 0x00000000) && continue

    # Undefined/bad codes
    T         = get(tdms_codes, pcode, nothing)
    (T == nothing) && error("Undefined property type!")

    if pcode == 0x00000044
      # TDMS documentation has the order of these two variables reversed
      nf = fastread(io, UInt64)
      ns = fastread(io, Int64)

      dt = tdms_dtos + ns + nf*2^(-64)
      if name == "GPSTimeStamp"
        TDMS.ts = round(Int64, sμ*dt)
      else
        D[name] = u2d(dt)
        (v > 1) && println(name, " = ", D[name])
      end
      # Will be corrected to UTC by adding SystemInfomation.GPS.UTCOffset

    elseif pcode == 0x00000020
      # No documentation here either
      nn = fastread(io, UInt32)
      val = fastread(io, nn)
      if name == "name"
        TDMS.name = String(val)
      else
        D[name] = String(val)
        (v > 1) && println(name, " = ", D[name])
      end

    else
      if name == "SamplingFrequency[Hz]"
        TDMS.fs = fastread(io, T)
      elseif name == "Latitude"
        TDMS.oy = fastread(io, T)
      elseif name == "Longitude"
        TDMS.ox = fastread(io, T)
      elseif name == "Altitude"
        TDMS.oz = fastread(io, T)
      else
        D[name] = fastread(io, T)
        (v > 1) && println(name, " = ", D[name])
      end

      #=
       In which direction is StartPosition[m]?
                  what about Start Distance (m)?

      =#
    end
  end
  return nothing
end

function read_silixa_tdms(file::String, nn::String, s::TimeSpec, t::TimeSpec, chans::ChanSpec, memmap::Bool, v::Integer)

  io = memmap ? IOBuffer(Mmap.mmap(file)) : open(file, "r")
  tdms_header!(TDMS, io, v)

  # not documented; inferred from manufacturer's Matlab script
  # uses general formula: mask = 2^(n-1); N & mask != mask
  #                       N = UInt8(data_flags & 0b00100000)
  #                       n = 6 (6th bit = 1 if data are decimated)
  decimated = Bool(UInt8(TDMS.flags & DECIMATE_MASK) != 2^5)
  # I don't think this is actually part of the standard TDMS data spec

  # data header info
  fastskip(io, fastread(io, UInt32)+8) # Group information
  fastskip(io, fastread(io, UInt32)+4) # first channel path and length
  DataType = fastread(io, UInt32)
  fastskip(io, 4)
  chunk_size = Int64(fastread(io, UInt32))
  T = get(tdms_codes, DataType, nothing)
  (T == nothing) && error("Unsupported data type!")

  seg1_length   = Int64(div(div(TDMS.nsos - TDMS.rdos, TDMS.n_ch), sizeof(T)))
  if v > 0
    println("data type = ", T,
            ", chunk size = ", chunk_size,
            ", channel length = ", seg1_length,
            ", decimated = ", decimated)
  end

  # parse start time to get zero-indexed si, ei in each channel
  if (typeof(s) <: Real) && (typeof(t) <: Real)
    si = max(round(Int64, TDMS.fs*s), 0)
    ei = min(round(Int64, TDMS.fs*t), seg1_length) - 1
  else
    Δ = round(Int64, SeisIO.sμ / TDMS.fs)
    if s == "0001-01-01T00:00:00"
      s = u2d(TDMS.ts*μs)
    elseif isa(s, Real)
      s = u2d(TDMS.ts*μs + s)
    end
    if t == "9999-12-31T12:59:59"
      t = u2d(TDMS.ts*μs + seg1_length/TDMS.fs)
    elseif isa(t, Real)
      t = u2d(TDMS.ts*μs + t)
    end
    (d0, d1) = parsetimewin(s, t)
    t0 = DateTime(d0).instant.periods.value*1000 - dtconst
    t1 = DateTime(d1).instant.periods.value*1000 - dtconst
    (v > 2) && println("t0 = ", t0, ", t1 = ", t1)
    si = max(div(t0-TDMS.ts, Δ), 0)
    ei = min(div(t1-TDMS.ts, Δ), seg1_length) - 1
  end
  si = Int64(si)
  ei = Int64(ei)

  # this prints one-indexed forms
  (v > 1) && println("reading from si = ", si+1, " to ei = ", ei+1)

  # Chunk bounds
  first_chunk = div(si, chunk_size) + 1
  last_chunk = div(ei, chunk_size) + 1
  n_chunks = last_chunk - first_chunk + 1

  # Skip to start of data
  nskip = TDMS.rdos - position(io)
  if first_chunk > 1
    chunk_os = (first_chunk-1)*chunk_size*sizeof(T)*TDMS.n_ch
    nskip += chunk_os
    (v > 2) && println("To skip: ", chunk_os, " bytes of chunks.")
  end
  if nskip != 0
    (v > 0) && println("Skipping ", nskip, " bytes total.")
    fastskip(io, nskip)
  end

  # Read Data
  buf = Array{T, 2}(undef, TDMS.n_ch, chunk_size)
  data = Array{Float32, 2}(undef, ei-si+1, TDMS.n_ch)

  # ===================================================================
  # Read chunks
  # sj = starting index within chunk
  # ej = ending index within chunk
  j = 1
  jmax = div(seg1_length, chunk_size)*chunk_size - si
  for i in first_chunk:last_chunk
    if j > jmax
      if VERSION < v"1.4"
        buf = Array{T, 2}(undef, TDMS.n_ch, rem(seg1_length, chunk_size))
        read!(io, buf)
      else
        vbuf = view(buf, :, 1:rem(seg1_length, chunk_size))
        read!(io, vbuf)
      end
    else
      read!(io, buf)
    end
    sj = (i == first_chunk ? rem(si+1, chunk_size) : 1)
    ej = (i == last_chunk ? rem(ei, chunk_size) + 1 : chunk_size)
    nj = ej-sj+1
    vx = view(data, j:j+nj-1, :)
    vb = view(buf, :, sj:ej)
    (v > 2) && println("chunk #", i, "/", n_chunks, ", sj = ", sj, ", ej = ", ej, ", j = ", j, ":", j+nj-1)

    transpose!(vx, vb) # converts to Float32 if needed
    j += nj
  end

  # Done with file
  close(io)

  # ----------------------------------------------------------------
  # String values for :name, :id
  name = (try
    string(TDMS.hdr["SystemInfomation"]["OS"]["HostName"], "_",
           TDMS.hdr["SystemInfomation"]["Devices0"]["Model"], "_",
           TDMS.hdr["SystemInfomation"]["Devices1"]["Model"])
  catch
    ""
  end)
  net = nn * "."
  cha = string("..O", getbandcode(TDMS.fs), "0")

  # -----------------------------------------------------------------
  # Time values
  utc_os = get(TDMS.hdr, "SystemInfomation.GPS.UTCOffset", zero(Float64))
  ts = TDMS.ts + round(Int64, sμ*utc_os)

  # =================================================================
  # Parse to NodalData
  if isempty(chans)
    chans = 1:TDMS.n_ch
  end
  S = NodalData(data, TDMS.hdr, chans, ts)
  fill!(S.fs, TDMS.fs)
  fill!(S.src, realpath(file))
  fill!(S.units, "m/m")
  S.ox = TDMS.ox
  S.oy = TDMS.oy
  S.oz = TDMS.oz
  for (i,j) in enumerate(chans)
    S.id[i] = string(net, lpad(j, 5, '0'), cha)
    S.name[i] = string(name, "_", j)
  end
  # =================================================================

  return S
end