## delete this on NCI
# using Distributed
# addprocs(2)
## end delete
## MPI Init
using MPIClusterManagers, Distributed
import MPI
MPI.Init()
rank = MPI.Comm_rank(MPI.COMM_WORLD)
size = MPI.Comm_size(MPI.COMM_WORLD)
if rank == 0
    @info "size is $size"
end
manager = MPIClusterManagers.start_main_loop(MPI_TRANSPORT_ALL)
@info "there are $(nworkers()) workers"
@everywhere @info gethostname()
include("01_read_data_aseggdf.jl")
include("02_set_options.jl")
## MPI checks
# split into sequential iterations of parallel soundings
nsoundings = length(soundings)
ncores = nworkers()
nsequentialiters = ceil(Int, nsoundings/ncores)
@info "will require $nsequentialiters iterations of $ncores soundings"
## set up transD_GP
@everywhere using Distributed
@everywhere using HiQGA.transD_GP
## do the parallel soundings
@info "starting"
transD_GP.SkyTEM1DInversion.loopacrosssoundings(soundings, σstart, σ0,
                            nsequentialiters   = nsequentialiters,
                            zfixed             = zfixed          ,
                            ρfixed             = ρfixed          ,
                            zstart             = zstart          ,
                            extendfrac         = extendfrac      ,
                            dz                 = dz              ,
                            ρbg                = ρbg             ,
                            nlayers            = nlayers         ,
                            ntimesperdecade    = ntimesperdecade ,
                            nfreqsperdecade    = nfreqsperdecade ,
                            modelprimary       = modelprimary    ,
                            regtype            = regtype         ,
                            nstepsmax          = nstepsmax       ,
                            ntries             = ntries          ,
                            target             = target          ,
                            lo                 = lo              ,
                            hi                 = hi              ,
                            λ²min              = λ²min           ,
                            λ²max              = λ²max           ,
                            λ²frac             = λ²frac          ,
                            β²                 = β²              ,
                            ntestdivsλ²        = ntestdivsλ²     ,
                            αmin               = αmin            ,
                            αmax               = αmax            ,
                            αfrac              = αfrac           ,
                            ntestdivsα         = ntestdivsα      ,
                            regularizeupdate   = regularizeupdate,
                            knownvalue         = knownvalue      ,
                            firstvalue         = firstvalue      ,
                            κ                  = κ               ,
                            breakonknown       = breakonknown    ,
                            dobo               = dobo,
                            zipsaveprefix      = zipsaveprefix)

MPIClusterManagers.stop_main_loop(manager)
rmprocs(workers())
exit()


using CGP
using Logging
using PyCall
using ArgParse
using Printf
using Random

@pyimport gym
@pyimport pybullet
@pyimport pybullet_envs
CGP.Config.init("cfg/base.yaml")
CGP.Config.init("cfg/classic.yaml")

function play_env(c::Chromosome, env, render::Bool=false)
    env[:seed](0)
    if render
        env[:render](mode="human")
    end
    ob = env[:reset]()
    total_reward = 0.0
    done = false
    reward = 0.0

    while ~done
        action = process(c, ob./5.0)
        try
            ob, reward, done, _ = env[:step](action)
            if render
                env[:render](mode="human")
            end
        catch
            done = true
        end
        total_reward += reward
    end

    total_reward
end

function get_args()
    s = ArgParseSettings()

    @add_arg_table(
        s,
        "--seed", arg_type = Int, default = 0,
        "--log", arg_type = String, default = "gym.log",
        "--id", arg_type = String, default = "MountainCarContinuous-v0",
        "--ea", arg_type = String, default = "oneplus",
        "--chromosome", arg_type = String, default = "CGPChromo",
        "--cfg", arg_type = String,
    )

    parse_args(CGP.Config.add_arg_settings!(s))
end

if ~isinteractive()
    args = get_args()

    if args["cfg"] != nothing
        CGP.Config.init(args["cfg"])
    end

    CGP.Config.init(Dict([k=>args[k] for k in setdiff(
        keys(args), ["seed", "log", "id", "ea", "chromosome", "cfg", "graph"])]...))

    Random.seed!(args["seed"])
    global_logger(SimpleLogger(open(args["log"], "a+")))
    ea = eval(Meta.parse(args["ea"]))
    ctype = eval(Meta.parse(args["chromosome"]))
    env = gym.make(args["id"])
    nin = length(env[:observation_space][:low])
    nout = length(env[:action_space][:low])

    fit = x->play_env(x, env, false)
    maxfit, best = ea(nin, nout, fit; seed=args["seed"], ctype=ctype, id=args["id"])

    @info(@sprintf("E%0.6f", -maxfit))

    best_ind = ctype(best, nin, nout)
    play_env(best_ind, env, true)
end


export PolarForm, bounds, powerflow
export State, Control, Parameters, NumberOfState, NumberOfControl

"""
    AbstractStructure

The user may specify a mapping to the single input vector `x` for AD.

"""
abstract type AbstractStructure end

"""
    AbstractFormulation

Second layer of the package, implementing the interface between
the first layer (the topology of the network) and the
third layer (implementing the callbacks for the optimization solver).

"""
abstract type AbstractFormulation end

"""
    AbstractFormAttribute

Attributes attached to an `AbstractFormulation`.
"""
abstract type AbstractFormAttribute end

"Number of states attached to a particular formulation."
struct NumberOfState <: AbstractFormAttribute end

"Number of controls attached to a particular formulation."
struct NumberOfControl <: AbstractFormAttribute end

"""
    AbstractVariable

Variables corresponding to a particular formulation.
"""
abstract type AbstractVariable end

"""
    State <: AbstractVariable

All variables ``x`` depending on the variables `Control` ``u`` through
a non-linear equation ``g(x, u) = 0``.

"""
struct State <: AbstractVariable end

"""
    Control <: AbstractVariable

Independent variables ``u`` used in the reduced-space
formulation.

"""
struct Control <: AbstractVariable end

"""
    PhysicalState <: AbstractVariable

All physical variables describing the current physical state
of the underlying network.

`PhysicalState` variables are encoded in a `AbstractNetworkBuffer`,
storing all the physical values needed to describe the current
state of the network.

"""
struct PhysicalState <: AbstractVariable end

# Templates
"""
    get(form::AbstractFormulation, attr::AbstractFormAttribute)

Return value of attribute `attr` attached to the particular
formulation `form`.

## Examples

```julia
get(form, NumberOfState())
get(form, NumberOfControl())

```
"""
function get end

"""
    setvalues!(form::AbstractFormulation, attr::PS.AbstractNetworkAttribute, values)

Update inplace the attribute's values specified by `attr`.

## Examples

```julia
setvalues!(form, ActiveLoad(), new_ploads)
setvalues!(form, ReactiveLoad(), new_qloads)

```
"""
function setvalues! end

"""
    bounds(form::AbstractFormulation, var::AbstractVariable)

Return the bounds attached to the variable `var`.

    bounds(form::AbstractFormulation, func::Function)

Return a tuple of vectors `(lb, ub)` specifying the admissible range
of the constraints specified by the function `cons_func`.

## Examples

```julia
u_min, u_max = bounds(form, Control())
h_min, h_max = bounds(form, reactive_power_constraints)

```
"""
function bounds end

"""
    initial(form::AbstractFormulation, var::AbstractVariable)

Return an initial position for the variable `var`.

## Examples

```julia
u₀ = initial(form, Control())
x₀ = initial(form, State())

```
"""
function initial end

"""
    powerflow(form::AbstractFormulation,
              algo::AbstractNonLinearSolver;
              kwargs...)

    powerflow(form::AbstractFormulation,
              jacobian::AutoDiff.Jacobian,
              buffer::AbstractNetworkBuffer,
              algo::AbstractNonLinearSolver;
              kwargs...) where VT <: AbstractVector

Solve the power flow equations ``g(x, u) = 0`` w.r.t. the state ``x``,
using the algorithm specified in `algo` ([`NewtonRaphson`](@ref) by default).
The initial state ``x`` is specified inside
`buffer`. The object `buffer` is modified inplace in the function.

The algorithm stops when a tolerance `tol` or a maximum number of
iterations `maxiter` is reached (these parameters being specified
in the object `algo`).

## Notes
If only the arguments `form` and `algo` are specified to the function,
then the Jacobian `jacobian` and the cache `buffer` are inferred
from the object `form`.

## Arguments

* `form::AbstractFormulation`: formulation of the power flow equation
* `jacobian::AutoDiff.Jacobian`: Jacobian
* `buffer::AbstractNetworkBuffer`: buffer storing current state `x` and control `u`
* `algo::AbstractNonLinearSolver`: non-linear solver. Currently only `NewtonRaphson` is being implemented.

## Optional arguments

* `linear_solver::AbstractLinearSolver` (default `DirectSolver()`): solver to solve the linear systems ``J x = y`` arising at each iteration of the Newton-Raphson algorithm.

"""
function powerflow end

# Cost function
"""
    cost_production(form::AbstractFormulation, buffer::AbstractNetworkBuffer)::Float64

Get operational cost.
"""
function cost_production end

# Generic constraints

"""
    voltage_magnitude_constraints(form::AbstractFormulation, cons::AbstractVector, buffer::AbstractNetworkBuffer)

Bounds the voltage magnitudes at PQ nodes:
```math
v_{pq}^♭ ≤ v_{pq} ≤ v_{pq}^♯ .
```
The result is stored inplace, inside `cons`.

## Note
The constraints on the voltage magnitudes at PV nodes ``v_{pv}``
are taken into account when bounding the control ``u``.
"""
function voltage_magnitude_constraints end

"""
    active_power_constraints(form::AbstractFormulation, cons::AbstractVector, buffer::AbstractNetworkBuffer)

Evaluate the constraints on the **active power production** at the generators
that are not already taken into account in the bound constraints.
```math
p_g^♭ ≤ p_g ≤ p_g^♯  .
```

The result is stored inplace, inside the vector `cons`.
"""
function active_power_constraints end

"""
    active_power_generation(form::AbstractFormulation, pg::AbstractVector, buffer::AbstractNetworkBuffer)

Evaluate the **active power production** of all generators.
Used to evaluate the operational cost.
The result is stored inplace, inside the vector `pg`.
"""
function active_power_generation end

"""
    reactive_power_constraints(form::AbstractFormulation, cons::AbstractVector, buffer::AbstractNetworkBuffer)

Evaluate the constraints on the **reactive power production** at the generators:
```math
q_g^♭ ≤ q_g ≤ q_g^♯  .
```
The result is stored inplace, inside the vector `cons`.
"""
function reactive_power_constraints end

"""
    flow_constraints(form::AbstractFormulation, cons::AbstractVector, buffer::AbstractNetworkBuffer)

Evaluate the thermal limit constraints porting on the lines of the network.
The result is stored inplace, inside the vector `cons`.
"""
function flow_constraints end

@doc raw"""
    power_balance(form::AbstractFormulation, cons::AbstractVector, buffer::AbstractNetworkBuffer)

Evaluate the power balance in the network:
```math
g(x, u) = 0 ,
```
corresponding to the balance equations
```math
\begin{aligned}
    p_i &= v_i \sum_{j}^{n} v_j (g_{ij}\cos{(\theta_i - \theta_j)} + b_{ij}\sin{(\theta_i - \theta_j})) \,, &
    ∀ i ∈ \{PV, PQ\} \\
    q_i &= v_i \sum_{j}^{n} v_j (g_{ij}\sin{(\theta_i - \theta_j)} - b_{ij}\cos{(\theta_i - \theta_j})) \,. &
    ∀ i ∈ \{PQ\} \\
\end{aligned}
```

The result is stored inplace, inside the vector `cons`.
"""
function power_balance end

# Interface for the constraints
"""
    size_constraint(cons_func::Function)::Bool
Return whether the function `cons_func` is a supported constraint
in the powerflow model.
"""
function is_constraint end

"""
    size_constraint(form::AbstractFormulation, cons_func::Function)::Int

Get number of constraints specified by the function `cons_func`
in the formulation `form`.
"""
function size_constraint end

"""
    adjoint!(form::AbstractFormulation, pbm::AutoDiff.TapeMemory, adj_h, h, buffer)

Return the adjoint w.r.t. the variables of the network (voltage magnitudes
and angles, power injection) for the constraint stored inside the [`AutoDiff.TapeMemory`](@ref)
object `pbm`. The results are stored directly inside the stack stored
inside `pbm`.
"""
function adjoint! end

"""
    jacobian_transpose_product!(form::AbstractFormulation, pbm::AutoDiff.TapeMemory, buffer, v)

Return the two transpose-Jacobian vector product ``(Jᵤ^⊤ v, Jₓ^⊤ v)``  w.r.t. the
control ``u`` and the state ``x``. Store the two resulting vectors directly inside
the [`AutoDiff.TapeMemory`](@ref) `pbm`.

"""
function jacobian_transpose_product! end

"""
    matpower_jacobian(form::AbstractFormulation, X::Union{State,Control}, cons_func::Function, V::Vector{Complex})
    matpower_jacobian(form::AbstractFormulation, X::Union{State,Control}, cons_func::Function, buffer::AbstractNetworkBuffer)

For the constraint `cons_func`, return the expression of the Jacobian ``J``
w.r.t. the state or the control (depending on the argument `X`),
as given by MATPOWER.
"""
function matpower_jacobian end

@doc raw"""
    matpower_hessian(form::AbstractFormulation, cons_func::Function, buffer::AbstractNetworkBuffer, λ::AbstractVector)

For constraint `cons_func`, return the three matrices ``(λ^⊤ H_{xx},
λ^⊤ H_{xu},λ^⊤ H_{uu})`` storing the product of the Hessian tensor ``H`` with the vector ``\lambda``.
The expressions of the Hessian matrices are given by MATPOWER.

"""
function matpower_hessian end

"""
    jacobian_sparsity(form::AbstractFormulation, cons_func::Function, X::Union{State,Control})

For the constraint `cons_func`, return the sparsity pattern of the Jacobian ``J``
w.r.t. the state or the control (depending on the argument `X`).

"""
function jacobian_sparsity end


using AutomotiveDrivingModels
using NGSIM
using HDF5, JLD

ΔT = NGSIM_TIMESTEP # [s]

immutable Interval
    id::Int
    frame_lo::Int
    frame_hi::Int
end
type IntervalInProgress
    id::Int
    frame_lo::Int
    updated::Bool
end

# include("tim2dfeatures.jl")
# include("radar_features.jl")
include("multifeatureset.jl")

extract_core = true
extract_temporal = true
extract_well_behaved = false
extract_neighbor_features = false
extract_carlidar_rangerate = true
carlidar_nbeams = 20
roadlidar_nbeams = 0
roadlidar_nlanes = 0
carlidar_max_range = 100.0 # [m]
roadlidar_max_range = 50.0 # [m]

extractor = MultiFeatureExtractor(
    extract_core,
    extract_temporal,
    extract_well_behaved,
    extract_neighbor_features,
    extract_carlidar_rangerate,
    carlidar_nbeams,
    roadlidar_nbeams,
    roadlidar_nlanes,
    carlidar_max_range=carlidar_max_range,
    roadlidar_max_range=roadlidar_max_range,
    )

N_FEATURES = length(extractor)

# const ROADWAY = open(io->read(io, Roadway), "roadway_passive_aggressive.txt", "r")
# for trajdata_path in ["trajdata_passive_aggressive1.txt"]
# println("FILE: ", trajdata_path); tic()
#     trajdata = open(io->read(io, Trajdata, ROADWAY), trajdata_path, "r")

culling_fidelity = 5.0 # [m]
lane_portion_max_range = roadlidar_max_range + culling_fidelity/2 + 5.0 # [m] note extra margin of error

err = 0.0

for trajdata_path in TRAJDATA_PATHS
    println("FILE: ", trajdata_path); tic()
    trajdata = load_trajdata(trajdata_path)

    extractor.road_lidar_culling = RoadwayLidarCulling(trajdata.roadway, lane_portion_max_range, culling_fidelity)
    @assert extractor.road_lidar_culling.is_leaf == false

    segments = pull_continuous_segments(trajdata, 1)

    # extract the dataset for each segment
    scene = Scene()
    rec = SceneRecord(2, ΔT)
    n_states = sum(i->length(i.frame_lo:i.frame_hi)-1, segments)
    features = Array(Float64, N_FEATURES)
    featureset = Array(Float64, length(features), n_states)
    targetset = Array(Float64, 2, n_states)
    interval_starts = Array(Int, length(segments))

    i = 0
    for (segind, seg) in enumerate(segments)
        empty!(rec)
        interval_starts[segind] = i+1
        tic()
        for frame in seg.frame_lo : seg.frame_hi
            get!(scene, trajdata, frame)
            update!(rec, scene)
            vehicle_index = get_index_of_first_vehicle_with_id(rec, seg.egoid)
            @assert(vehicle_index != 0)
            pull_features!(extractor, features, rec, trajdata.roadway, vehicle_index)

            if frame > seg.frame_lo
                targetset[1, i] = get(ACC, rec, trajdata.roadway, vehicle_index)
                targetset[2, i] = get(TURNRATEG, rec, trajdata.roadway, vehicle_index)
            end

            if frame < seg.frame_hi
                i += 1
                featureset[:,i] = features
            end
        end
        toc()
    end

    # Save it all to a file
    println("saving...")
    outfile = ("./../2d_drive_data/data_" * splitdir(splitext(trajdata_path)[1])[2]
    * "_clb" * string(carlidar_nbeams)
    * "_rlb" * string(roadlidar_nbeams)
    * "_rll" * string(roadlidar_nlanes)
    * "_clmr" * string(Int(carlidar_max_range))
    * "_rlmr" * string(Int(roadlidar_max_range))
    * ".jld")

    #outfile = "./data_" * splitdir(splitext(trajdata_path)[1])[2] * ".jld"
    JLD.save(outfile, "features", featureset, "targets", targetset, "intervals", interval_starts, "timestep", ΔT)
    toc()
end


export modelimage

"""
    $(TYPEDEF)
Container for non-analytic model that contains a image cache which will hold the image,
a Fourier transform cache C, which usually an instance of a <: FourierCache which usually
holds a interpolator that allows you to compute visibilities.

# Notes
This is an internal implementation detail that shouldn't usually be called directly.
Instead the user should use the exported function `modelimage`, for example

```julia
using Comrade
m = ExtendedRing(20.0, 5.0)

# This creates an version where the image is dynamically specified according to the
# radial extent of the image
mimg = modelimage(m) # you can also optionally pass the number of pixels nx and ny

# Or you can create an IntensityMap
img = intensitymap(m, 100.0, 100.0, 512, 512)
mimg = modelimage(m, img)
```
"""
struct ModelImage{M,I,C} <: AbstractModelImage{M}
    model::M
    image::I
    cache::C
end
@inline visanalytic(::Type{<:ModelImage{M}}) where {M} = visanalytic(M)
@inline imanalytic(::Type{<:ModelImage{M}}) where {M} = imanalytic(M)
@inline isprimitive(::Type{<:ModelImage{M}}) where {M} = isprimitive(M)

function Base.show(io::IO, mi::ModelImage)
   #io = IOContext(io, :compact=>true)
   #s = summary(mi)
   ci = first(split(summary(mi.cache), "{"))
   println(io, "ModelImage")
   println(io, "\tmodel: ", summary(mi.model))
   println(io, "\timage: ", summary(mi.image))
   println(io, "\tcache: ", ci)
end

model(m::AbstractModelImage) = m.model
flux(mimg::ModelImage) = flux(intensitymap!(mimg.image, mimg.model))

# function intensitymap(mimg::ModelImage)
#     intensitymap!(mimg.image, mimg.model)
#     mimg.image
# end

radialextent(m::ModelImage) = fov(m.image)[1]/2

#@inline visibility_point(m::AbstractModelImage, u, v) = visibility_point(model(m), u, v)

@inline intensity_point(m::AbstractModelImage, x, y) = intensity_point(model(m), x, y)

"""
    $(SIGNATURES)
Construct a `ModelImage` from a `model`, `image` and the optionally
specified visibility algorithm `alg`

# Notes
For analytic models this is a no-op and just return the model.
For non-analytic models this wraps the model in a object with an image
and precomputes the fourier transform using `alg`.
"""
@inline function modelimage(model::M, image::ComradeBase.AbstractIntensityMap; alg=FFT()) where {M}
    return modelimage(visanalytic(M), model, image; alg)
end

@inline function modelimage(::IsAnalytic, model, image; alg=FFT())
    return model
end

@inline function modelimage(::NotAnalytic, model, image; alg=FFT())
    intensitymap!(image, model)
    cache = create_cache(alg, image)
    return ModelImage(model, image, cache)
end

"""
    modelimage(m)
Construct a `ModelImage` where just the model `m` is specified

# Notes
If m `IsAnalytic()` is the visibility domain this is a no-op and just returns the model itself.
Otherwise `modelimage` will *guess* a reasonable field of view based on the `radialextent`
function. One can optionally pass the number of pixels nx and ny in each direction.
"""
function modelimage(m::M;
                    fovx=2*radialextent(m),
                    fovy=2*radialextent(m),
                    nx=512,
                    ny=512,
                    alg=FFT()) where {M}
    if visanalytic(M) == IsAnalytic()
        return m
    else
        T = typeof(intensity_point(m, 0.0, 0.0))
        img = IntensityMap(zeros(T,ny,nx), fovx, fovy)
        modelimage(m, img; alg)
        #throw(ArgumentError("$m is not an analytic model a image must be specified"))
    end
end


using Test
using Pkg
using Statistics

using Flux

using LyceumBase
using LyceumMuJoCo
using LinearAlgebra
using LyceumAI

function testrollout(getaction!, env::AbstractEnvironment, T)
    s = Array(undef, statespace(env))
    a = Array(undef, actionspace(env))
    o = Array(undef, observationspace(env))
    for t = 1:T
        getstate!(s, env)
        getobservation!(o, env)
        getaction!(a, s, o)
        setaction!(env, a)
        step!(env)
    end
    env
end

@testset "LyceumAI.jl" begin
    tseed!(1)

    @testset "algorithms" begin
        include("algorithms/MPPI.jl")
        include("algorithms/naturalpolicygradient.jl")
    end
    @testset "util" begin
        include("util/misc.jl")
    end
    @testset "vectorproducts" begin
        include("vectorproducts.jl")
    end

end


# Use SymbolicMDPs wrapper and ReinforcementLearning library
using SymbolicMDPs, ReinforcementLearning
using Test

# Import Flux optimizers
import ReinforcementLearningZoo: Flux.Descent

# Import various MDP tools
import POMDPs
import POMDPModelTools: MDPCommonRLEnv

# Load PDDL functions
import PDDL: PDDL, load_domain, load_problem, @pddl
using PlanningDomains

# Load and compile blocksworld domain and problem with 2 blocks
domain = load_domain(JuliaPlannersRepo, "blocksworld")
problem = load_problem(JuliaPlannersRepo, "blocksworld", "problem-1")
domain, _ = PDDL.compiled(domain, problem)

# Construct SymbolicMDP
mdp = SymbolicMDP(domain, problem)

# Test successful plan to goal state using POMDPs.jl interface
s = rand(POMDPs.initialstate(mdp))
for act in PDDL.@pddl("(pick-up a)", "(stack a b)")
    s = rand(POMDPs.transition(mdp, s, act))
end
@test POMDPs.isterminal(mdp, s)

# Convert MDP to RL environment
STATE_TYPE = Int # Convert states to their indices for tabular RL
env = SymbolicRLEnv(STATE_TYPE, mdp)
env = convert(RLBase.AbstractEnv, env);
env = RL.discrete2standard_discrete(env) # Convert actions to their indices

# Test successful plan to goal state
reset!(env)
for act in PDDL.@pddl("(pick-up a)", "(stack a b)")
    a_idx = POMDPs.actionindex(mdp, act) # Look up action index
    env(a_idx) # Execute action
end
@assert is_terminated(env)

# Construct tabular Q-learning agent
agent = Agent(
    policy=QBasedPolicy(
        learner=TDLearner(
            approximator=TabularQApproximator( # Linear Q-value approximator
                n_state=length(state_space(env)), # Number of states
                n_action=length(action_space(env)), # Number of actions
                opt=Descent(0.5)
                ),
            method=:SARS, # SARS is equivalent to Q-learning,
            n=2, # Update Q-values after 2 steps (necessary for convergence)
        ),
        explorer=GreedyExplorer()
    ),
    trajectory=VectorSARTTrajectory(;
        state=typeof(state(env)),
        action=Int,
        reward=Float64,
        terminal=Bool
    )
)

# Run agent on environment for 10 episodes of 100 steps
reset!(env)
for i in 1:10
    run(agent, env, StopAfterStep(100))
end
q_idxs = findall(!=(0), agent.policy.learner.approximator.table)
q_vals = agent.policy.learner.approximator.table[q_idxs]

# Run learned policy and extract plan
plan = PDDL.Term[]
reset!(env)
for i in 1:10
    a_idx = agent.policy(env)
    env(a_idx)
    push!(plan, mdp.actions[a_idx])
    is_terminated(env) && break
end

display(plan) # Display plan
@test plan == PDDL.@pddl("(pick-up a)", "(stack a b)") # Test for correct plan
@test is_terminated(env) # Check that terminal state is reached

# Check that goal is satisfied
goal_state = collect(mdp.states)[state(env)]
PDDL.satisfy(domain, goal_state, problem.goal)


export
    Invertible

"""
    struct Invertible{I,O} <: SFunc{Tuple{I},O,Nothing}

An invertible sfunc, with both a `forward` and a `inverse` function provided.

# Additional supported operators
- `support`
- `support_quality`
- `sample`
- `logcpdf`
- `bounded_probs`
- `make_factors`
- `compute_pi`
- `send_lambda`
"""
struct Invertible{I,O} <: SFunc{Tuple{I},O}
    forward :: Function
    inverse :: Function
end

@impl begin
    struct InvertibleSupport end
    function support(sf::Invertible{I,O}, parranges, size, curr) where {I,O}
        invsup = map(sf.forward, parranges[1])
        return invsup
    end
end

@impl begin
    struct InvertibleSupportQuality end

    function support_quality(::Invertible{I,O}, parranges) where {I,O}
        return :CompleteSupport
    end
end

@impl begin
    struct InvertibleSample end
    function sample(sf::Invertible{I,O}, i::Tuple{I}) where {I,O}
        return sf.forward(i[1])
    end
end

@impl begin
    struct InvertibleLogcpdf end
    function logcpdf(sf::Invertible{I,O}, i::Tuple{I}, o::O) where {I,O}
        return sf.forward(i[1]) == o ? 0.0 : -Inf
    end
end

@impl begin
    struct InvertibleBoundedProbs end

    function bounded_probs(sf::Invertible{I,O}, range::Vector{<:O},
            parranges) where {I,O}
        result = Float64[]
        for i in parranges[1]
            for o in range
                p = sf.forward(i) == o ? 1.0 : 0.0
                push!(result, p)
            end
        end
        return (result, result)
    end
end

@impl begin
    struct InvertibleMakeFactors end

    function make_factors(sf::Invertible{I,O}, range::Vector{<:O}, parranges,
            id, parids) where {I,O}
        dims = (length(parranges[1]), length(range))
        keys = (parids[1], id)
        entries = bounded_probs(sf, range, parranges)[1]
        facts = [Factor(dims, keys, entries)]
        return (facts, facts)
    end
end

@impl begin
    struct InvertibleComputePi end

    function compute_pi(sf::Invertible{I,O}, range::Vector{<:O}, parranges,
            incoming_pis) where {I,O}
        ps = [cpdf(incoming_pis[1], (), sf.inverse(o)) for o in range]
        return Cat(range, ps)
    end
end

@impl begin
    struct InvertibleSendLambda end

    function send_lambda(sf::Invertible{I,O}, lambda, range, parranges, incoming_pis,
            parent_idx) where {I,O}
        @assert parent_idx == 1
        ls = [get_log_score(lambda, sf.forward(x)) for x in parranges[1]]
        return LogScore(parranges[1], ls)
    end
end
#=
@impl begin
    struct InvertibleInitialStats end
    initial_stats(::Invertible) = nothing
end

@impl begin
    struct InvertibleAccumulateStats end
    accumulate_stats(::Invertible, existing_stats, new_stats) = nothing
end

@impl begin
    struct InvertibleExpectedStats end
    expected_stats(::Invertible, range, parranges, incoming_pis, child_lambda) =
        nothing
end

@impl begin
    struct InvertibleMaximizeStats end
    maximize_stats(::Invertible, stats) = nothing
end
=#


"""
Module to define an interactive player. (Intended for testing purposes.)
"""
module InteractivePlay

import PlayingCards: Card, createdeck

import Cribbage: Player
import Cribbage.RandomPlay: discardtocrib, pickcardtoplay
import Cribbage.SmartPlay: averagescore, DECK
import Cribbage.Rules: cardvalue, validhands, validplays
import Cribbage.Scoring: scorehand, scoreplay

struct InteractivePlayer <: Player
    name :: String
end

function InteractivePlayer()
    return InteractivePlayer(UUIDs.uuid4())
end

function discardtocrib(player::InteractivePlayer, cards::Tuple{Vararg{Card}})
    println("You were dealt: $(string(cards))")
    subsets = validhands(cards)
    others = setdiff(DECK, cards)
    sort!(subsets, by=x -> averagescore(x, others), rev=true)
    for (i, h) in enumerate(subsets)
        if i ≥ 10
            break
        end
        println("[$i] $(string(h))")
    end
    println("Which hand would you like to keep?"); userinput = readline()
    handint = parse(Int, userinput)
    hand = subsets[handint]
    crib = setdiff(cards, hand)
    @debug "keeping $(string(hand)), discarding $(string(crib))"
    return hand, crib
end

function pickcardtoplay(player::InteractivePlayer, prior::Tuple{Vararg{Card}}, cards::Tuple{Vararg{Card}})::Union{Missing,Card}
    if length(prior) == 0
        println("No prior cards")
    else
        println("Prior cards $(sum(cardvalue(c) for c in prior)): $(string(prior))")
    end
    println("You have: $(string(cards))")
    possible_plays = validplays(prior, cards)
    if length(possible_plays) == 0
        println("You have no possible plays.")
        return missing
    end
    if length(prior) == 0
        priorarray = Card[]
    else
        priorarray = collect(prior)
    end
    sort!(possible_plays, by=x -> scoreplay(priorarray, x), rev=true)
    for (i, c) in enumerate(possible_plays)
        if i ≥ 10
            break
        end
        println("[$i] $(string(c))")
    end
    println("Which card would you like to play?"); userinput = readline()
    cardint = parse(Int, userinput)
    return possible_plays[cardint]
end

end  # module InteractivePlay


@testset "leadtime.jl" begin
    lt = LeadTime(2, 3, 10)
    @test state(lt) == [10,10,10]
    @test state_size(lt) == 3
    dest = []
    InventoryModels.push!(lt, 5, dest)
    @test state(lt) == [10,10,10,5]
    @test only(lt.pull_orders) == (dest => 5)
    InventoryModels.dispatch!(lt)
    @test dest == [10, lt]
    @test state(lt) == [10,10,5]
    @test InventoryModels.reward!(lt) == -25*2
    @test isempty(lt.pull_orders)
    @test InventoryModels.inventory_position(lt) == 25
    reset!(lt)
    @test state(lt) == [10,10,10]
    @test isempty(lt.pull_orders)

    lt = LeadTime(3, Uniform(5,15))
    @test all(5 .< state(lt) .< 15)
    @test state_size(lt) == 3
    next = first(state(lt))
    dest = []
    InventoryModels.push!(lt, 5, dest)
    @test state(lt)[4] == 5
    @test only(lt.pull_orders) == (dest => 5)
    InventoryModels.dispatch!(lt)
    @test dest == [next, lt]
    @test state(lt)[3] == 5
    InventoryModels.reward!(lt) == 0
    @test isempty(lt.pull_orders)
    reset!(lt)
    @test all(5 .< state(lt) .< 15)
    @test isempty(lt.pull_orders)
end


# # Basic usage of `StandardModel`

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/@__NAME__.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/@__NAME__.ipynb)

# In this tutorial we will use `COBREXA`'s `StandardModel` and functions that
# specifically operate on it. As usual we will use the toy model of *E. coli*
# for demonstration.

!isfile("e_coli_core.json") &&
    download("http://bigg.ucsd.edu/static/models/e_coli_core.json", "e_coli_core.json")

using COBREXA

# ## Loading a model

model = load_model(StandardModel, "e_coli_core.json") # we specifically want to load a StandardModel from the model file

#md # !!! note "Note: Loading `StandardModel`s implicitly uses `convert`"
#md #       When using `load_model(StandardModel, file_location)` the model at
#md #       `file_location` is first loaded into its inferred format and is then
#md #       converted to a `StandardModel` using the generic accessor interface.
#md #       Thus, data loss may occur. Always check your model to ensure that
#md #       nothing important has been lost.

#nb # When using `load_model(StandardModel, file_location)` the model at
#nb # `file_location` is first loaded into its inferred format and is then
#nb # converted to a `StandardModel` using the generic accessor interface.
#nb # Thus, data loss may occur. Always check your model to ensure that
#nb # nothing important has been lost.

# ## Basic analysis

# As before, for optimization based analysis we need to load an optimizer. Here we
# will use [`Tulip.jl`](https://github.com/ds4dm/Tulip.jl) to solve the linear
# programs of this tutorial. Refer to the basic constraint-based analysis
# tutorial for more informaiton.

# All the normal analysis functions work on `StandardModel`, due to it also
# having the same generic accessor interface as all the other model types.

using Tulip

fluxes = flux_balance_analysis_dict(
    model,
    Tulip.Optimizer;
    modifications = [
        change_objective("BIOMASS_Ecoli_core_w_GAM"),
        change_constraint("EX_glc__D_e"; lb = -12, ub = -12),
        change_constraint("EX_o2_e"; lb = 0, ub = 0),
    ],
)

# This is not very exciting yet, since every other model type can also do this.
# However, deeper inspection of flux results is possible when using
# `StandardModel`.

# ## Inspecting the flux solution: `atom_fluxes`

# It is sometimes interesting to keep track of the atoms entering and leaving
# the system through boundary reactions. This can be inspected by calling
# [`atom_fluxes`](@ref). That gives you the flux of individual atoms getting
# consumed and produced by all reactions, based on `fluxes`. We erase the
# reaction that consumes the atoms for creating biomass, to see how much mass
# the "rest" of the reaction produces for it:

fluxes_without_biomass = copy(fluxes);
delete!(fluxes_without_biomass, "BIOMASS_Ecoli_core_w_GAM");
atom_fluxes(model, fluxes_without_biomass)

# ## Inspecting the flux solution: `metabolite_fluxes`

# Another useful flux result analysis function is [`metabolite_fluxes`](@ref).
# This function gives an overview of reactions consuming and producing each
# metabolite.

consuming, producing = metabolite_fluxes(model, fluxes)

consuming["atp_c"] # reactions consuming `atp_c`

# ## Internals of `StandardModel`

# Another benefit of `StandardModel` is that it supports a richer internal
# infrastructure that can be used to manipulate internal model attributes in a
# systematic way. Specifically, the genes, reactions, and metabolites with of a
# model each have a type. This is particularly useful when modifying or even
# constructing a model from scratch.

# ## `Gene`s, `Reaction`s, and `Metabolite`s

# `StandardModel` is composed of ordered dictionaries of `Gene`s, `Metabolite`s
# and `Reaction`s. Ordered dictionaries are used because the order of the
# reactions and metabolites are important for constructing a stoichiometric
# matrix since the rows and columns should correspond to the order of the metabolites
# and reactions returned by calling the accessors `metabolites` and `reactions`.

# Each `StandardModel` is composed of the following fields:

fieldnames(StandardModel) # fields of a StandardModel

# The `:genes` field of a `StandardModel` contains an ordered dictionary of gene ids mapped to `Gene`s.

model.genes # the keys of this dictionary are the same as genes(model)

# The `Gene` type is a struct that can be used to store information about genes
# in a `StandardModel`. Each `Gene` is composed of the following fields:

fieldnames(Gene)

#md # !!! tip "Tip: Use <tab> complete to explore the structure of types"
#md #       Use <tab> to quickly explore the fields of a struct. For example,
#md #       Gene.<tab> will list all the fields shown above.

#nb # Use <tab> to quickly explore the fields of a struct. For example,
#nb # Gene.<tab> will list all the fields shown above.

# The keys used in the ordered dictionaries in
# `model.genes` are the ids returned using the generic accessor `genes`. `Gene`s
# have pretty printing, as demonstrated below for a random gene drawn from the
# model:

random_gene_id = genes(model)[rand(1:n_genes(model))]
model.genes[random_gene_id]

# The same idea holds for both metabolites (stored as `Metabolite`s) and
# reactions (stored as `Reaction`s). This is demonstrated below.

random_metabolite_id = metabolites(model)[rand(1:n_metabolites(model))]
model.metabolites[random_metabolite_id]
#
random_reaction_id = reactions(model)[rand(1:n_reactions(model))]
model.reactions[random_reaction_id]

# `StandardModel` can be used to build your own metabolic model or modify an
# existing one. One of the main use cases for `StandardModel` is that it can be
# used to merge multiple models or parts of multiple models together. Since the
# internals are uniform inside each `StandardModel`, attributes of other model
# types are squashed into the required format (using the generic accessors).
# This ensures that the internals of all `StandardModel`s are the same -
# allowing easy systematic evaluation.

#md # !!! warning "Warning: Combining models with different namespaces is tricky"
#md #       Combining models that use different namespaces requires care.
#md #       For example, in some models the water exchange reaction is called
#md #       `EX_h2o_e`, while in others it is called `R_EX_h2o_s`. This needs to
#md #       manually addressed (for now) to prevent duplicate, e.g. reactions,
#md #       from being added.

# ## Checking the internals of `StandardModel`s: `annotation_index`

# Often when models are automatically reconstructed duplicate genes, reactions
# or metabolites end up in a model. `COBREXA` exports `annotation_index` to
# check for cases where the id of a struct may be different, but the annotations
# the same (possibly suggesting a duplication). `annotation_index` builds a
# dictionary mapping annotation features to the ids of whatever struct you are
# inspecting. This makes it easy to find structs that share certain annotation features.

rxn_annotations = annotation_index(model.reactions)
#
rxn_annotations["ec-code"]

# The `annotation_index` function can also be used on `Reaction`s and
# `Gene`s in the same way.

# ## Checking the internals of `StandardModel`s: `check_duplicate_reaction`

# Another useful function is `check_duplicate_reaction`, which checks for
# reactions that have duplicate (or similar) reaction equations.

pgm_duplicate = Reaction()
pgm_duplicate.id = "pgm2" # Phosphoglycerate mutase
pgm_duplicate.metabolites = Dict{String,Float64}("3pg_c" => 1, "2pg_c" => -1)
pgm_duplicate
#
check_duplicate_reaction(pgm_duplicate, model.reactions; only_metabolites = false) # can also just check if only the metabolites are the same but different stoichiometry is used

# ## Checking the internals of `StandardModel`s: `reaction_mass_balanced`

# Finally, [`reaction_mass_balanced`](@ref) can be used to check if a reaction is mass
# balanced based on the formulas of the reaction equation.

rxn_dict = Dict{String,Float64}("3pg_c" => 1, "2pg_c" => -1, "h2o_c" => 1)
reaction_mass_balanced(model, rxn_dict)

# Now to determine which atoms are unbalanced, you can use `reaction_atom_balance`
reaction_atom_balance(model, rxn_dict)

# Note, since `pgm_duplicate` is not in the model, we cannot use the other variants of this
# function because they find the reaction equation stored inside the `model`.


using BlockArrays
using LinearAlgebra
using MixedModels
using PooledArrays
using Random
using SparseArrays
using Statistics
using StatsModels
using Tables
using Test

using MixedModels: dataset, likelihoodratiotest

const io = IOBuffer()

include("modelcache.jl")

@testset "Dyestuff" begin
    fm1 = only(models(:dyestuff))

    @test length(fm1.allterms) == 3
    @test size(fm1.reterms) == (1, )
    @test lowerbd(fm1) == zeros(1)
    @test fm1.lowerbd == zeros(1)
    @test fm1.optsum.initial == ones(1)
    fm1.θ = ones(1)
    @test fm1.θ == ones(1)

    @test_throws ArgumentError fit!(fm1)

    fm1.optsum.feval = -1
    @test_logs (:warn, "Model has not been fit") show(fm1)

    @test objective(updateL!(setθ!(fm1, [0.713]))) ≈ 327.34216280955366
    @test_deprecated MixedModels.describeblocks(IOBuffer(), fm1)

    show(io, BlockDescription(fm1))
    @test countlines(seekstart(io)) == 3
    output = String(take!(io))
    @test startswith(output, "rows:")

    refit!(fm1)

    @test :θ in propertynames(fm1)
    @test objective(fm1) ≈ 327.3270598811428 atol=0.001
    @test fm1.θ ≈ [0.752580] atol=1.e-5
    @test fm1.λ ≈ [LowerTriangular(reshape(fm1.θ, (1,1)))] atol=1.e-5
    @test deviance(fm1) ≈ 327.32705988 atol=0.001
    @test aic(fm1) ≈ 333.3270598811394 atol=0.001
    @test bic(fm1) ≈ 337.5306520261259 atol=0.001
    @test fixef(fm1) ≈ [1527.5]
    @test dispersion_parameter(fm1)
    @test first(first(fm1.σs)) ≈ 37.26034462135931 atol=0.0001
    @test fm1.β ≈ [1527.5]
    @test dof(fm1) == 3
    @test nobs(fm1) == 30
    @test MixedModels.fixef!(zeros(1),fm1) ≈ [1527.5]
    @test coef(fm1) ≈ [1527.5]
    fm1β = fm1.βs
    @test fm1β isa NamedTuple
    @test isone(length(fm1β))
    @test first(values(fm1β)) ≈ 1527.5
    fm1σρ = fm1.σρs
    @test fm1σρ isa NamedTuple
    @test isone(length(fm1σρ))
    @test isone(length(getproperty(first(fm1σρ), :σ)))
    @test isempty(getproperty(first(fm1σρ), :ρ))
    @test fm1.σ == sdest(fm1)
    @test fm1.b == ranef(fm1)
    @test fm1.u == ranef(fm1, uscale=true)
    @test fm1.stderror == stderror(fm1)
    @test isone(length(fm1.pvalues))
    @test fm1.objective == objective(fm1)
    @test fm1.σ ≈ 49.510099986291145 atol=1.e-5
    @test fm1.X == ones(30,1)
    ds = MixedModels.dataset(:dyestuff)
    @test fm1.y == ds[:yield]
    @test cond(fm1) == ones(1)
    @test first(leverage(fm1)) ≈ 0.15650534392640486 rtol=1.e-5
    @test sum(leverage(fm1)) ≈ 4.695160317792145 rtol=1.e-5
    cm = coeftable(fm1)
    @test length(cm.rownms) == 1
    @test length(cm.colnms) == 4
    @test fnames(fm1) == (:batch,)
    @test response(fm1) == ds[:yield]
    rfu = ranef(fm1, uscale = true)
    rfb = ranef(fm1)
    @test abs(sum(rfu[1])) < 1.e-5
    cv = condVar(fm1)
    @test length(cv) == 1
    @test size(first(cv)) == (1, 1, 6)
    show(IOBuffer(), fm1.optsum)

    @test logdet(fm1) ≈ 8.06014522999825 atol=0.001
    @test varest(fm1) ≈ 2451.2501089607676 atol=0.001
    @test pwrss(fm1) ≈ 73537.49947885796 atol=0.001
    @test stderror(fm1) ≈ [17.69455188898009] atol=0.0001

    vc = VarCorr(fm1)
    show(io, vc)
    str = String(take!(io))
    @test startswith(str, "Variance components:")
    @test vc.s == sdest(fm1)

    refit!(fm1, REML=true)
    @test objective(fm1) ≈ 319.65427684225216 atol=0.0001
    @test_throws ArgumentError loglikelihood(fm1)
    @test dof_residual(fm1) ≥ 0

    print(io, fm1)
    @test startswith(String(take!(io)), "Linear mixed model fit by REML")

    fm1.optsum.maxfeval = 5
    fm1.optsum.feval = -1
    @test_logs (:warn, "NLopt optimization failure: MAXEVAL_REACHED") fit!(fm1)
    fm1.optsum.maxfeval = -1

    vc = fm1.vcov
    @test isa(vc, Matrix{Float64})
    @test only(vc) ≈ 409.79495436473167 rtol=1.e-6
    # since we're caching the fits, we should get it back to being correctly fitted
    refit!(fm1)
end

@testset "Dyestuff2" begin
    ds2 = MixedModels.dataset(:dyestuff2)
    fm = only(models(:dyestuff2))
    @test lowerbd(fm) == zeros(1)
    show(IOBuffer(), fm)
    @test fm.θ ≈ zeros(1)
    @test objective(fm) ≈ 162.87303665382575
    @test abs(std(fm)[1][1]) < 1.0e-9
    @test std(fm)[2] ≈ [3.653231351374652]
    @test stderror(fm) ≈ [0.6669857396443261]
    @test coef(fm) ≈ [5.6656]
    @test logdet(fm) ≈ 0.0
    @test issingular(fm)
    #### modifies the model
    refit!(fm, float(MixedModels.dataset(:dyestuff)[:yield]))
    @test objective(fm) ≈ 327.3270598811428 atol=0.001
    refit!(fm, float(MixedModels.dataset(:dyestuff2)[:yield])) # restore the model in the cache
end

@testset "penicillin" begin
    fm = only(models(:penicillin))
    @test size(fm) == (144, 1, 30, 2)
    @test fm.optsum.initial == ones(2)
    @test lowerbd(fm) == zeros(2)

    @test objective(fm) ≈ 332.18834867227616 atol=0.001
    @test coef(fm) ≈ [22.97222222222222] atol=0.001
    @test fixef(fm) ≈ [22.97222222222222] atol=0.001
    @test coef(fm)[1] ≈ mean(dataset(:penicillin).diameter)
    @test stderror(fm) ≈ [0.7445960346851368] atol=0.0001
    @test fm.θ ≈ [1.5375772376554968, 3.219751321180035] atol=0.001
    @test first(std(fm)) ≈ [0.8455645948223015] atol=0.0001
    @test std(fm)[2] ≈ [1.770647779277388] atol=0.0001
    @test varest(fm) ≈ 0.3024263987592062 atol=0.0001
    @test logdet(fm) ≈ 95.74614821367786 atol=0.001

    @test_throws ArgumentError condVar(fm)

    rfu = ranef(fm, uscale=true)
    @test length(rfu) == 2
    @test first(first(rfu)) ≈ 0.523162392717432 rtol=1.e-4

    rfb = ranef(fm)
    @test length(rfb) == 2
    @test last(last(rfb)) ≈ -3.001823834230942 rtol=1.e-4

    show(io, BlockDescription(fm))
    @test countlines(seekstart(io)) == 4
    tokens = Set(split(String(take!(io)), r"\s+"))
    @test "Diag/Dense" in tokens
    @test "Diagonal" in tokens
end

@testset "pastes" begin
    fm = last(models(:pastes))
    @test size(fm) == (60, 1, 40, 2)
    @test fm.optsum.initial == ones(2)
    @test lowerbd(fm) == zeros(2)

    @test objective(fm) ≈ 247.99446586289676 atol=0.001
    @test coef(fm) ≈ [60.05333333333329] atol=0.001
    @test fixef(fm) ≈ [60.05333333333329] atol=0.001
    @test stderror(fm) ≈ [0.6421359883527029] atol=0.0001
    @test fm.θ ≈ [3.5268858714382905, 1.3299230213750168] atol=0.001
    @test first(std(fm)) ≈ [2.904069002535747] atol=0.001
    @test std(fm)[2] ≈ [1.095070371687089] atol=0.0001
    @test std(fm)[3] ≈ [0.8234088395243269] atol=0.0001
    @test varest(fm) ≈ 0.6780020742644107 atol=0.0001
    @test logdet(fm) ≈ 101.0381339953986 atol=0.001

    show(io, BlockDescription(fm))
    @test countlines(seekstart(io)) == 4
    tokens = Set(split(String(take!(io)), r"\s+"))
    @test "Sparse" in tokens
    @test "Diagonal" in tokens

    lrt = likelihoodratiotest(models(:pastes)...)
    @test length(lrt.deviance) == length(lrt.formulas) == length(lrt.models )== 2
    @test first(lrt.tests.pvalues) ≈ 0.5233767966395597 atol=0.0001
end

@testset "InstEval" begin
    fm1 = first(models(:insteval))
    @test size(fm1) == (73421, 2, 4114, 3)
    @test fm1.optsum.initial == ones(3)
    @test lowerbd(fm1) == zeros(3)

    @test objective(fm1) ≈ 237721.7687745563 atol=0.001
    ftd1 = fitted(fm1);
    @test size(ftd1) == (73421, )
    @test ftd1 == predict(fm1)
    @test first(ftd1) ≈ 3.17876 atol=0.0001
    resid1 = residuals(fm1);
    @test size(resid1) == (73421, )
    @test first(resid1) ≈ 1.82124 atol=0.00001

    @testset "PCA" begin
        @test length(fm1.rePCA) == 3
        pca = MixedModels.PCA(fm1)
        @test length(pca) == 3
        @test :covcor in propertynames(first(pca))
        str = String(take!(io))
        show(io, first(pca), stddevs=true, variances=true)
        str = String(take!(io))
        @test !isempty(findall("Standard deviations:", str))
        @test !isempty(findall("Variances:", str))
    end

    show(io, BlockDescription(fm1))
    @test countlines(seekstart(io)) == 5
    tokens = Set(split(String(take!(io)), r"\s+"))
    @test "Sparse" in tokens
    @test "Sparse/Dense" in tokens
    @test "Diag/Dense" in tokens

    fm2 = last(models(:insteval))
    @test objective(fm2) ≈ 237585.5534151694 atol=0.001
    @test size(fm2) == (73421, 28, 4100, 2)
end

@testset "sleep" begin
    fm = last(models(:sleepstudy))
    @test lowerbd(fm) == [0.0, -Inf, 0.0]
    A11 = getblock(fm.A, 1,1)
    @test isa(A11, UniformBlockDiagonal{Float64})
    @test isa(getblock(fm.L, 1, 1), UniformBlockDiagonal{Float64})
    @test size(A11) == (36, 36)
    a11 = view(A11.data, :, :, 1)
    @test a11 == [10. 45.; 45. 285.]
    @test size(A11.data, 3) == 18
    λ = first(fm.λ)
    b11 = LowerTriangular(view(getblock(fm.L, 1, 1).data, :, :, 1))
    @test b11 * b11' ≈ λ'a11*λ + I rtol=1e-5
    @test count(!iszero, Matrix(getblock(fm.L, 1, 1))) == 18 * 4
    @test rank(fm) == 2

    @test objective(fm) ≈ 1751.9393444647046
    @test fm.θ ≈ [0.929221307, 0.01816838, 0.22264487096] atol=1.e-6
    @test pwrss(fm) ≈ 117889.46144025437
    @test logdet(fm) ≈ 73.90322021999222 atol=0.001
    @test stderror(fm) ≈ [6.632257721914501, 1.5022354739749826] atol=0.0001
    @test coef(fm) ≈ [251.40510484848477,10.4672859595959]
    @test fixef(fm) ≈  [251.40510484848477,10.4672859595959]
    @test std(fm)[1] ≈ [23.780468100188497, 5.716827903196682] atol=0.01
    @test logdet(fm) ≈ 73.90337187545992 atol=0.001
    @test cond(fm) ≈ [4.175251] atol=0.0001
    @test loglikelihood(fm) ≈ -875.9696722323523
    @test sum(leverage(fm)) ≈ 28.611525700136877 rtol=1.e-5
    σs = fm.σs
    @test length(σs) == 1
    @test keys(σs) == (:subj,)
    @test length(σs.subj) == 2
    @test first(values(σs.subj)) ≈ 23.7804686 atol=0.0001
    @test last(values(first(σs))) ≈ 5.7168278 atol=0.0001
    @test fm.corr ≈ [1.0 -0.1375451787621904; -0.1375451787621904 1.0] atol=0.0001

    u3 = ranef(fm, uscale=true)
    @test length(u3) == 1
    @test size(first(u3)) == (2, 18)
    @test first(u3)[1, 1] ≈ 3.030300122575336 atol=0.001

    b3 = ranef(fm)
    @test length(b3) == 1
    @test size(first(b3)) == (2, 18)
    @test first(first(b3)) ≈ 2.815819441982976 atol=0.001

    b3tbl = raneftables(fm)
    @test length(b3tbl) == 1
    @test keys(b3tbl) == (:subj,)
    @test isa(b3tbl, NamedTuple)
    @test Tables.istable(only(b3tbl))

    simulate!(fm)  # to test one of the unscaledre methods
    # must restore state of fm as it is cached in the global fittedmodels
    slp = MixedModels.dataset(:sleepstudy)
    copyto!(fm.y, slp.reaction)
    updateL!(MixedModels.reevaluateAend!(fm))
    @test objective(fm) ≈ 1751.9393444647046 # check the model is properly restored

    fmnc = models(:sleepstudy)[2]
    @test size(fmnc) == (180,2,36,1)
    @test fmnc.optsum.initial == ones(2)
    @test lowerbd(fmnc) == zeros(2)

    @testset "zerocorr PCA" begin
        @test length(fmnc.rePCA) == 1
        @test fmnc.rePCA.subj ≈ [0.5, 1.0]
        @test any(Ref(fmnc.PCA.subj.loadings) .≈ (I(2), I(2)[:, [2,1]]))
        @test show(IOBuffer(), MixedModels.PCA(fmnc)) === nothing
    end

    @test deviance(fmnc) ≈ 1752.0032551398835 atol=0.001
    @test objective(fmnc) ≈ 1752.0032551398835 atol=0.001
    @test coef(fmnc) ≈ [251.40510484848585, 10.467285959595715]
    @test fixef(fmnc) ≈ [251.40510484848477, 10.467285959595715]
    @test stderror(fmnc) ≈ [6.707710260366577, 1.5193083237479683] atol=0.001
    @test fmnc.θ ≈ [0.9458106880922268, 0.22692826607677266] atol=0.0001
    @test first(std(fmnc)) ≈ [24.171449463289047, 5.799379721123582]
    @test last(std(fmnc)) ≈ [25.556130034081047]
    @test logdet(fmnc) ≈ 74.46952585564611 atol=0.001
    ρ = first(fmnc.σρs.subj.ρ)
    @test ρ === -0.0   # test that systematic zero correlations are returned as -0.0

    MixedModels.likelihoodratiotest(fm, fmnc)
    fmrs = fit(MixedModel, @formula(reaction ~ 1+days + (0+days|subj)), slp);
    @test objective(fmrs) ≈ 1774.080315280528 rtol=0.00001
    @test fmrs.θ ≈ [0.24353985679033105] rtol=0.00001

    fm_ind = models(:sleepstudy)[3]
    @test objective(fm_ind) ≈ objective(fmnc)
    @test coef(fm_ind) ≈ coef(fmnc)
    @test fixef(fm_ind) ≈ fixef(fmnc)
    @test stderror(fm_ind) ≈ stderror(fmnc)
    @test fm_ind.θ ≈ fmnc.θ
    @test std(fm_ind) ≈ std(fmnc)
    @test logdet(fm_ind) ≈ logdet(fmnc)

    # combining [ReMat{T,S1}, ReMat{T,S2}] for S1 ≠ S2
    slpcat = (subj = slp.subj, days = PooledArray(string.(slp.days)), reaction = slp.reaction)
    fm_cat = fit(MixedModel, @formula(reaction ~ 1+days+(1|subj)+(0+days|subj)),slpcat)
    @test fm_cat isa LinearMixedModel
    σρ = fm_cat.σρs
    @test σρ isa NamedTuple
    @test isone(length(σρ))
    @test first(keys(σρ)) == :subj
    @test keys(σρ.subj) == (:σ, :ρ)
    @test length(σρ.subj) == 2
    @test length(first(σρ.subj)) == 10
    @test length(σρ.subj.ρ) == 45
    # test that there's no correlation between the intercept and days columns
    ρs_intercept = σρ.subj.ρ[1 .+ cumsum(0:8)]
    @test all(iszero.(ρs_intercept))
    # amalgamate should set these to -0.0 to indicate structural zeros
    @test all(ρs_intercept .=== -0.0)

    # also works without explicitly dropped intercept
    fm_cat2 = fit(MixedModel, @formula(reaction ~ 1+days+(1|subj)+(days|subj)),slpcat)
    @test fm_cat2 isa LinearMixedModel
    σρ = fm_cat2.σρs
    @test σρ isa NamedTuple
    @test isone(length(σρ))
    @test first(keys(σρ)) == :subj
    @test keys(σρ.subj) == (:σ, :ρ)
    @test length(σρ.subj) == 2
    @test length(first(σρ.subj)) == 10
    @test length(σρ.subj.ρ) == 45
    # test that there's no correlation between the intercept and days columns
    ρs_intercept = σρ.subj.ρ[1 .+ cumsum(0:8)]
    @test all(iszero.(ρs_intercept))
    # amalgamate should set these to -0.0 to indicate structural zeros
    @test all(ρs_intercept .=== -0.0)

    show(io, BlockDescription(first(models(:sleepstudy))))
    @test countlines(seekstart(io)) == 3
    @test "Diagonal" in Set(split(String(take!(io)), r"\s+"))

    show(io, BlockDescription(last(models(:sleepstudy))))
    @test countlines(seekstart(io)) == 3
    @test "BlkDiag" in Set(split(String(take!(io)), r"\s+"))

end

@testset "d3" begin
    fm = only(models(:d3))
    @test pwrss(fm) ≈ 5.30480294295329e6 rtol=1.e-4
    @test objective(fm) ≈ 884957.5540213 rtol = 1e-4
    @test coef(fm) ≈ [0.4991229873, 0.31130780953] atol = 1.e-4
    @test length(ranef(fm)) == 3

    show(io, BlockDescription(fm))
    tokens = Set(split(String(take!(io)), r"\s+"))
    @test "BlkDiag/Dense" in tokens
end

@testset "kb07" begin
    global io
    pca = last(models(:kb07)).PCA
    @test keys(pca) == (:subj, :item)
    show(io, models(:kb07)[2])
    tokens = Set(split(String(take!(io)), r"\s+"))
    @test "Corr." in tokens
    @test "-0.89" in tokens
end

@testset "oxide" begin
    # this model has an interesting structure with two diagonal blocks
    m = first(models(:oxide))
    @test all(isapprox.(m.θ, [1.689182746, 2.98504262]; atol=1e-4))
    m = last(models(:oxide))
    # NB: this is a poorly defined fit
    # lme4 gives all sorts of convergence warnings for the different
    # optimizers and even quite different values
    # the overall estimates of the standard deviations are similar-ish
    # but the correlation structure seems particular unstable
    θneldermead = [1.6454, 8.6373e-02, 8.2128e-05, 8.9552e-01, 1.2014, 2.9286]
    # two different BOBYQA implementations
    θnlopt = [1.645, -0.221, 0.986, 0.895, 2.511, 1.169]
    θminqa = [1.6455, -0.2430, 1.0160, 0.8955, 2.7054, 0.0898]
    # very loose tolerance for unstable fit
    # but this is a convenient test of rankUpdate!(::UniformBlockDiagonal)
    @test all(isapprox.(m.θ, θnlopt; atol=1e-2))
end

@testset "Rank deficient" begin
    rng = MersenneTwister(0);
    x = rand(rng, 100);
    data = (x = x, x2 = 1.5 .* x, y = rand(rng, 100), z = repeat('A':'T', 5))
    model = fit(MixedModel, @formula(y ~ x + x2 + (1|z)), data)
    @test length(fixef(model)) == 2
    @test rank(model) == 2
    @test length(coef(model)) == 3
    ct = coeftable(model)
    @test ct.rownms ==  ["(Intercept)", "x", "x2"]
    @test length(fixefnames(model)) == 2
    @test coefnames(model) == ["(Intercept)", "x", "x2"]
    piv = first(model.feterms).piv
    r = first(model.feterms).rank
    @test coefnames(model)[piv][1:r] == fixefnames(model)
end

@testset "coeftable" begin
    ct = coeftable(only(models(:dyestuff)));
    @test [3,4] == [ct.teststatcol, ct.pvalcol]
end

@testset "wts" begin
    # example from https://github.com/JuliaStats/MixedModels.jl/issues/194
    data = (
        a = [1.55945122,0.004391538,0.005554163,-0.173029772,4.586284429,0.259493671,-0.091735715,5.546487603,0.457734831,-0.030169602],
        b = [0.24520519,0.080624178,0.228083467,0.2471453,0.398994279,0.037213859,0.102144973,0.241380251,0.206570975,0.15980803],
        c = PooledArray(["H","F","K","P","P","P","D","M","I","D"]),
        w1 = [20,40,35,12,29,25,65,105,30,75],
        w2 = [0.04587156,0.091743119,0.080275229,0.027522936,0.066513761,0.05733945,0.149082569,0.240825688,0.068807339,0.172018349],
    )

    #= no need to fit yet another model without weights, but here are the reference values from lme4
    m1 = fit(MixedModel, @formula(a ~ 1 + b + (1|c)), data)
    @test m1.θ ≈ [0.0]
    @test stderror(m1) ≈  [1.084912, 4.966336] atol = 1.e-4
    @test vcov(m1) ≈ [1.177035 -4.802598; -4.802598 24.664497] atol = 1.e-4
    =#

    m2 = fit(MixedModel, @formula(a ~ 1 + b + (1|c)), data, wts = data.w1)
    @test m2.θ ≈ [0.295181729258352]  atol = 1.e-4
    @test stderror(m2) ≈  [0.9640167, 3.6309696] atol = 1.e-4
    @test vcov(m2) ≈ [0.9293282 -2.557527; -2.5575267 13.183940] atol = 1.e-4
end

@testset "unifying ReMat eltypes" begin
    sleepstudy = dataset(:sleepstudy)

    re = LinearMixedModel(@formula(reaction ~ 1 + days + (1|subj) + (days|subj)), sleepstudy).reterms
    # make sure that the eltypes are still correct
    # otherwise this test isn't checking what it should be
    @test eltype(sleepstudy.days) == Int8
    @test eltype(sleepstudy.reaction) == Float64

    # use explicit typeof() and == is to remind us that things may break
    # if we change things and don't check their type implications now
    # that we're starting to support a non trivial type hierarchy
    @test typeof(re) == Vector{AbstractReMat{Float64}}
end


using DataStructures

"""
HuffmanTree{M<:Integer, N<:Integer}
    A Huffman encoding of words comprising an array of nodes parents, and an
    array of branching direction for each node.  Array location corresponds to
    the hash value of corresponding word.
"""
struct HuffmanTree{M<:Integer, N<:Integer}
    nodeparent::AbstractArray{M}
    branch::AbstractArray{N}
end

"""
    HuffmanTree(pq, vocab_hash)

Take in a PriorityQueue(pq) and a vocabulary hash, and return a Huffman encoding
of the words in pq.
"""
function HuffmanTree(pq::PriorityQueue, vocab_hash::Dict)
    tmp_hash = copy(vocab_hash)

    # initialize return values
    asize = 2*length(tmp_hash)-1
    nodeparent = zeros(Int, asize)
    branch = zeros(Int, asize)

    # create nodes
    while length(pq) > 2
         _newnode!(pq, tmp_hash, nodeparent, branch)
    end
    _lastnode!(pq, tmp_hash, nodeparent, branch)

    return HuffmanTree(nodeparent, branch)
end

function _newnode!(pq::PriorityQueue, tmp_hash::Dict, nodeparent::Array, branch::Array)
    #create node ids
    (nodea, prioritya) = dequeue_pair!(pq)
    (nodeb, priorityb) = dequeue_pair!(pq)
    idxa = tmp_hash[nodea]
    idxb = tmp_hash[nodeb]

    #add a node to the graph
    idxnode = length(tmp_hash)+1
    prioritynode = prioritya + priorityb
    node = "node-$idxnode" # for debugging
    tmp_hash[node] = idxnode

    #connect the child nodes to the new node
    nodeparent[idxa] = idxnode
    nodeparent[idxb] = idxnode

    #set the binary value for the child nodes
    branch[idxa] = 1
    branch[idxb] = 0

    #enqueue the new node
    enqueue!(pq, node, prioritynode)
end

function _lastnode!(pq::PriorityQueue, tmp_hash::Dict, nodeparent::Array, branch::Array)
    #create node ids
    (nodea, prioritya) = dequeue_pair!(pq)
    (nodeb, priorityb) = dequeue_pair!(pq)
    idxa = tmp_hash[nodea]
    idxb = tmp_hash[nodeb]

    #add a node to the graph
    idxnode = length(tmp_hash)+1
    node = "root"
    tmp_hash[node] = idxnode

    #connect the child nodes to the new node
    nodeparent[idxa] = idxnode
    nodeparent[idxb] = idxnode

    #set the binary value for the child nodes
    branch[idxa] = 1
    branch[idxb] = 0
end


using Revise
using HybridFunnelGraphs

# graph = create_funnel_graph("test/hpd/domain.hpd", "test/hpd/problem.hpd"; max_levels=170)
graph = create_funnel_graph("test/hpd/dom2.hpd", "test/hpd/prob2.hpd"; max_levels=300, has_placement_constraint=false)
graph.num_levels


module JQL
# JQL is not a Query Language
#
# Wrap lazy-eager hash joins, of various flavours, in an iterator

import Base.∩, Base.∪

export AbstractJoinType, Inner, Outer, Left, Right, LeftSemi, RightSemi, AbstractJoin, Join, ∪, ∩

abstract AbstractJoinType <: Enum
abstract AbstractJoin{J<:AbstractJoinType, I, T}
immutable Inner <: AbstractJoinType end
immutable Outer <: AbstractJoinType end
immutable Left <: AbstractJoinType end
immutable Right <: AbstractJoinType end
immutable LeftSemi <: AbstractJoinType end
immutable RightSemi <: AbstractJoinType end

# Light weight pointers and metadata of the sources to join
immutable Join{J<:AbstractJoinType, I, T} <: AbstractJoin{J, I, T}
	large::I
	small::I
	parity::Bool
	ratio::Int64
	remainder::Int64
	function call{J<:AbstractJoinType, I}(::Type{Join}, ::Type{J}, left::I, right::I)
		if eltype(left) != eltype(right)     # this would be weird if it happened; nonetheless
			error("Column types do not match") # it is possible to write type unstable iterators
		end
		if length(left) <= length(right)
			return new{J, I, eltype(left)}(right, left, true, div(right, left), right % left )
		else
			return new{J, I, eltype(left)}(left, right, false, div(left, right), left % right)
		end
	end
end

# Nicer than subscripting tuples
immutable _JoinValue
	large::Vector{Int64}
	small::Vector{Int64}
end

# Would this be more efficient as object id dictionaries?
immutable _JoinState{T}
	reserve::Dict{T, _JoinValue}
	consign::Dict{T, _JoinValue}
	largemarchin::Int64
	smallmarchin::Int64
	largemarchout::Int64
	smallmarchout::Int64
end

# ToDo: drop J and fix on Inner, as this algorithm is just the inner join
# ToDo: use ratio/remainder to skip march in the small compared to the large
# with an initial preallocation of the small of remainder size.

# Alternating march in until there is a hit, then send through done to next for march out
function Base.start{J<:AbstractJoinType, I, T}(iterator::Join{J, I, T})
	reserve = Dict{T, _JoinValue}()
	consign = Dict{T, _JoinValue}()
	sizehint!(reserve, length(iterator.small))
	for i = 1:iterator.smalllength

		# March in from the small source
		smallvalue = get!(reserve, iterator.small(i), _JoinValue(Vector{Int64}(0), Vector{Int64}(0)))
		push!(smallvalue.small, i)

		# March in from the large source
		largevalue = get!(reserve, iterator.large(i), _JoinValue(Vector{Int64}(0), Vector{Int64}(0)))
		push!(largevalue.large, i)

		# Return the hit
		if length(smallvalue.large) > 0 || length(largevalue.small) > 0
			return _JoinState(reserve, consign, i, i, length(largevalue.small), length(smallvalue.large))
		end
	end

	# Both sources are the same size, so we are done
	if length(iterator.small) == length(iterator.large)
		return _JoinState(reserve, consign, length(iterator.large) + 1, length(iterator.small) + 1, 0, 0)
	end

	# Still data left in the large source, we only need roll call, without marching in
	for i = (length(iterator.small) + 1):length(iterator.large)
		largevalue = get(reserve, iterator.large(i), _JoinValue(Vector{Int64}(0), Vector{Int64}(0)))
		if length(largevalue.small) > 0
			return _JoinState(reserve, consign, i, length(iterator.small) + 1, length(largevalue.small), 0)
		end
	end

	# No matches found
	_JoinState(reserve, consign, length(iterator.large) + 1, length(iterator.small) + 1, 0, 0)
end

# March out the finds, small then big, until there are none, and then alternatingly march
# in until there is another hit
function Base.next{J<:AbstractJoinType, I, T}(iterator::Join{J, I}, state::_JoinState{T})

	# We are in next so there must be something to return, start by marching out the small
	# finds, on the left
	if state.smallmarchout > 0 && iterator.parity
		item = (iterator.small[state.smallmarchin], state.smallmarchin, state.reserve[iterator.small[state.smallmarchin]].large[state.smallmarchout])
		largemarchout = state.largemarchout
		smallmarchout = state.smallmarchout - 1

	# March out the small finds, on the right
	elseif state.smallmarchout > 0
		item = (iterator.small[state.smallmarchin], state.reserve[iterator.small[state.smallmarchin]].large[state.smallmarchout], state.smallmarchin)
		largemarchout = state.largemarchout
		smallmarchout = state.smallmarchout - 1

	# March out the large finds, on the right
	elseif state.largemarchout > 0 && iterator.parity
		ritem = (iterator.large[state.largemarchin], state.reserve[iterator.large[state.largemarchin]].small[state.largemarchout], state.largemarchin)
		largemarchout = state.largemarchout - 1
		smallmarchout = state.smallmarchout

	# March out the large finds, on the left
	elseif state.largemarchout > 0
		item = (iterator.large[state.largemarchin], state.largemarchin, state.reserve[iterator.large[state.largemarchin]].small[state.largemarchout])
		largemarchout = state.largemarchout
		smallmarchout = state.smallmarchout - 1

	# Opps something went wrong, maybe next was called out of loop, or sequence
	else
		error("next() was called on an invalid state")
	end

	# More left to march out after current, safe to loop through done to next
	if largemarchout > 0 || smallmarchout > 0
		return (item, _JoinState(state.reserve, state.consign, state.largemarchin, state.smallmarchin, largemarchout, smallmarchout))
	end

	# Nothing left to march out after the current item, so proceed with alternating march in
	if state.smallmarchin < length(iterator.small)
		for i = (state.smallmarchin + 1):length(iterator.small)

			# March in from the small source
			smallvalue = get!(state.reserve, iterator.small(i), _JoinValue(Vector{Int64}(0), Vector{Int64}(0)))
			push!(smallvalue.small, i)

			# March in from the large source
			largevalue = get!(state.reserve, iterator.large(i), _JoinValue(Vector{Int64}(0), Vector{Int64}(0)))
			push!(largevalue.large, i)

			# Return the hit
			if length(smallvalue.large) > 0 || length(largevalue.small) > 0
				return (item, _JoinState(state.reserve, state.consign, i, i, length(largevalue.small), length(smallvalue.large)))
			end
		end
	end

	# Both sources are the same size, so we are done
	if length(iterator.small) == length(iterator.large)
		return (item, _JoinState(state.reserve, state.consign, length(iterator.small) + 1, length(iterator.large) + 1, 0, 0))
	end

	# Still data left in the large source, we only need roll call, without marching in
	for i = (length(iterator.small) + 1):length(iterator.large)
		largevalue = get(state.reserve, iterator.large(i), _JoinValue(Vector{Int64}(0), Vector{Int64}(0)))
		if length(largevalue.small) > 0
			return (item, _JoinState(state.reserve, state.consign, i, length(iterator.small) + 1, length(largevalue.small), 0))
		end
	end

	# Whew! made it all the way through and found no more matches than the current one.
	(item, _JoinState(state.reserve, state.consign, length(iterator.large) + 1, length(iterator.small) + 1, 0, 0))
end
Base.done{J<:AbstractJoinType, I, T}(iterator::Join{J, I, T}, state::_JoinState{T}) =
	state.largemarchin > length(iterator.large) &&
	state.smallmarchin > length(iterator.small) &&
	state.largemarchout < 1 &&
	state.smallmarchout < 1
Base.eltype{J<:AbstractJoinType, I, T}(::Type{Join{J, I, T}}) = Tuple{T, Int64, Int64}

# These are just examples. Eventually they will be replaced with relational algebra symbols
∩(left::Tuple{Vararg{AbstractVector}}, right::Tuple{Vararg{AbstractVector}}) = Join(Inner, zip(left...), zip(right...))
∪(left::Tuple{Vararg{AbstractVector}}, right::Tuple{Vararg{AbstractVector}}) = Join(Outer, zip(left...), zip(right...))

end


import Oceananigans.Grids: required_halo_size
using Oceananigans.Utils: prettysummary

struct ScalarDiffusivity{TD, F, N, K} <: AbstractScalarDiffusivity{TD, F}
    ν :: N
    κ :: K

    function ScalarDiffusivity{TD, F}(ν::N, κ::K) where {TD, F, N, K}
        return new{TD, F, N, K}(ν, κ)
    end
end

"""
    ScalarDiffusivity(time_disc=ExplicitTimeDiscretization,
                      formulation=ThreeDimensionalFormulation, FT=Float64;
                      ν=0, κ=0,
                      discrete_form = false)

Return `ScalarDiffusivity` with viscosity `ν` and tracer diffusivities `κ`
for each tracer field in `tracers`. If a single `κ` is provided, it is
applied to all tracers. Otherwise `κ` must be a `NamedTuple` with values
for every tracer individually.

`ν` and the fields of `κ` may be constants (converted to `FT`), arrays, fields or
    - functions of `(x, y, z, t)` if `discrete_form = false`
    - functions of `(i, j, k, grid, LX, LY, LZ)` with `LX`, `LY` and `LZ` are either `Face()` or `Center()` if
      `discrete_form = true`.
"""
function ScalarDiffusivity(time_disc=ExplicitTimeDiscretization(),
                           formulation=ThreeDimensionalFormulation(), FT=Float64;
                           ν=0, κ=0,
                           discrete_form = false)

    if formulation == HorizontalFormulation() && time_discretization == VerticallyImplicitTimeDiscretization()
        throw(ArgumentError("VerticallyImplicitTimeDiscretization is only supported for `HorizontalFormulation` or `ThreeDimensionalFormulation`"))
    end

    κ = convert_diffusivity(FT, κ; discrete_form)
    ν = convert_diffusivity(FT, ν; discrete_form)

    return ScalarDiffusivity{typeof(time_disc), typeof(formulation)}(ν, κ)
end

const VerticalScalarDiffusivity{TD} = ScalarDiffusivity{TD, VerticalFormulation} where TD
const HorizontalScalarDiffusivity{TD} = ScalarDiffusivity{TD, HorizontalFormulation} where TD

  VerticalScalarDiffusivity(time_disc=ExplicitTimeDiscretization(), FT::DataType=Float64; kwargs...) = ScalarDiffusivity(time_disc, VerticalFormulation(), FT; kwargs...)
HorizontalScalarDiffusivity(time_disc=ExplicitTimeDiscretization(), FT::DataType=Float64; kwargs...) = ScalarDiffusivity(time_disc, HorizontalFormulation(), FT; kwargs...)

# Aliases that allow specify the floating type, assuming that the discretization is Explicit in time
          ScalarDiffusivity(FT::DataType; kwargs...) = ScalarDiffusivity(ExplicitTimeDiscretization(), ThreeDimensionalFormulation(), FT; kwargs...)
  VerticalScalarDiffusivity(FT::DataType; kwargs...) = ScalarDiffusivity(ExplicitTimeDiscretization(), VerticalFormulation(), FT; kwargs...)
HorizontalScalarDiffusivity(FT::DataType; kwargs...) = ScalarDiffusivity(ExplicitTimeDiscretization(), HorizontalFormulation(), FT; kwargs...)

required_halo_size(closure::ScalarDiffusivity) = 1

function with_tracers(tracers, closure::ScalarDiffusivity{TD, F}) where {TD, F}
    κ = tracer_diffusivities(tracers, closure.κ)
    return ScalarDiffusivity{TD, F}(closure.ν, κ)
end

@inline viscosity(closure::ScalarDiffusivity, args...) = closure.ν
@inline diffusivity(closure::ScalarDiffusivity, ::Val{tracer_index}, args...) where tracer_index = closure.κ[tracer_index]

# z_viscosity and z_diffusivity fallback to viscosity and diffusivity
# for both ThreeDimensional and VerticalScalarDiffusivity
@inline z_viscosity(closure::HorizontalScalarDiffusivity, args...) = 0
@inline z_diffusivity(closure::HorizontalScalarDiffusivity, args...) = 0

calculate_diffusivities!(diffusivities, ::ScalarDiffusivity, args...) = nothing

# Note: we could compute ν and κ (if they are Field):
# function calculate_diffusivities!(diffusivities, closure::ScalarDiffusivity, args...)
#     compute!(viscosity(closure, diffusivities))
#     !isnothing(closure.κ) && Tuple(compute!(diffusivity(closure, Val(c), diffusivities) for c=1:length(closure.κ)))
#     return nothing
# end

function Base.summary(closure::ScalarDiffusivity)
    TD = summary(time_discretization(closure))
    prefix = replace(summary(formulation(closure)), "Formulation" => "")
    prefix === "ThreeDimensional" && (prefix = "")
    return string(prefix, "ScalarDiffusivity{$TD}(ν=", prettysummary(closure.ν), ", κ=", prettysummary(closure.κ), ")")
end

Base.show(io::IO, closure::ScalarDiffusivity) = print(io, summary(closure))


abstract type AbstractContainerView{T} end

struct ContainerView{T} <: AbstractContainerView{T}
    model::EnergyModel
    devices::Dict{Symbol,T}
end

struct SubContainerView{T} <: AbstractContainerView{T}
    model::EnergyModel
    devices::Dict{Symbol,T}
    buses::Axis
end

tupleaxisvals(c::Bus, A) = A.val
tupleaxisvals(c) = tupleaxisvals(c, axis(c))
tupleaxisvals(c, A) = AxisArrays.CategoricalVector(tuple.(c.class, A.val))

_maybesubset(cv::ContainerView, c, A) = A
_maybesubset(cv::SubContainerView, c::Component, A) = A[findall(in(cv.buses), c), ntuple(i->:,ndims(A)-1)...]

_wrapaxis(cv::AbstractContainerView{T}, a) where T = Axis{Symbol(naming(T))}(a)
function _wrapaxisarray(cv::AbstractContainerView, c::Component, a::AxisArray)
    a = _maybesubset(cv, c, a)
    ax = first(AxisArrays.axes(a))
    AxisArray(a, _wrapaxis(cv, tupleaxisvals(c, ax)))
end
function _wrapaxisarray(cv::AbstractContainerView, c::Component, a::Union{Float64,Symbol})
    ax = _maybesubset(cv, c, tupleaxisvals(c))
    AxisArray(fill(a, length(ax)), _wrapaxis(cv, ax))
end

devices(cv::AbstractContainerView) = values(cv.devices)
devices(cv::AbstractContainerView, T::Type{<:Device}) = (c for c = devices(cv) if isa(c, T))

axis(cv::AbstractContainerView) = _wrapaxis(cv, vcat((_maybesubset(cv, c, tupleaxisvals(c)) for c = devices(cv))...))
axis(cv::AbstractContainerView, T::Type{<:Component}) = axis(m[T])

# Base.cat(cv::AbstractContainerView, cs::Vector{<:Component}, as::Vector) = @consense(as, "Single values do not agree")
"""
    cat(containerview, components, attrvalues)

Combine the `attrvalues` for the `components` in a `containerview` to a single AxisArray.
"""
Base.cat(cv::AbstractContainerView, components::Vector{<:Component}, attrvalues::Vector) =
    cat((_wrapaxisarray(cv, c, a) for (c, a) = zip(components, attrvalues))...; dims=1)

mapcat(f::Function, cv::AbstractContainerView) = cat(cv, collect(devices(cv)), map(f, devices(cv)))

Base.get(cv::AbstractContainerView, i) = mapcat(c->get(c, i), cv)
Base.get(cv::AbstractContainerView, i, axes...) = WrappedArray(get(cv, i), axes...)

Base.getindex(cv::AbstractContainerView, i) = get(cv, i)


"""
controls changing plane for example from transverse to saggital ...
"""
module ChangePlane
using ModernGL, ..DisplayWords, ..StructsManag, Setfield, ..PrepareWindow,   ..DataStructs , Rocket, GLFW,Dictionaries,  ..ForDisplayStructs, ..TextureManag,  ..OpenGLDisplayUtils,  ..Uniforms, Match, Parameters,DataTypesBasic

"""
In case we want to change the dimansion of scrolling so for example from transverse to coronal ...
    toBeSavedForBack - just marks weather we wat to save the info how to undo latest action
    - false if we invoke it from undoing
"""

function processKeysInfo(toScrollDatPrim::Identity{DataToScrollDims}
                    ,actor::SyncActor{Any, ActorWithOpenGlObjects}
                    ,keyInfo::KeyboardStruct
                    ,toBeSavedForBack::Bool = true ) where T
    toScrollDat= toScrollDatPrim.value

    old = actor.actor.onScrollData.dimensionToScroll

    newCalcDim= getHeightToWidthRatio(actor.actor.calcDimsStruct,toScrollDat )|>
                    getMainVerticies
     actor.actor.calcDimsStruct = newCalcDim
#In order to make the  background black  before we will render quad of possibly diffrent dimensions we will set all to invisible - and obtain black background
textSpecs = actor.actor.mainForDisplayObjects.listOfTextSpecifications

for textSpec in textSpecs
    setTextureVisibility(false,textSpec.uniforms )
end#for
basicRender(actor.actor.mainForDisplayObjects.window)

    #we need to change textures only if dimensions do not match
  #  if(actor.actor.calcDimsStruct.imageTextureWidth!=newCalcDim.imageTextureWidth  || actor.actor.calcDimsStruct.imageTextureHeight!=newCalcDim.imageTextureHeight )
        # first we need to update information about dimensions etc

        #next we need to delete all textures and create new ones

        arr = map(it->it.ID[],textSpecs)
glFinish()# make open gl ready for work

glDeleteTextures(length(arr), arr)# deleting

        #getting new
        initializeTextures(textSpecs,newCalcDim)

   # end#if

actor.actor.onScrollData.dimensionToScroll = toScrollDat.dimensionToScroll
actor.actor.onScrollData.dataToScrollDims = toScrollDat

actor.actor.onScrollData.slicesNumber = getSlicesNumber(actor.actor.onScrollData)
#getting  the slice of intrest based on last recorded mouse position

current=actor.actor.lastRecordedMousePosition[toScrollDat.dimensionToScroll]

#displaying all

singleSlDat= actor.actor.onScrollData.dataToScroll|>
(scrDat)-> map(threeDimDat->threeToTwoDimm(threeDimDat.type,Int64(current),toScrollDat.dimensionToScroll,threeDimDat ),scrDat) |>
(twoDimList)-> SingleSliceDat(listOfDataAndImageNames=twoDimList
                            ,sliceNumber=current
                            ,textToDisp = getTextForCurrentSlice(actor.actor.onScrollData, Int32(current))  )

                            # glFinish()
                            # glFlush()
                            # glClearColor(0.0, 0.0, 0.0 , 1.0)
                            # GLFW.SwapBuffers(actor.actor.mainForDisplayObjects.window)

dispObj = actor.actor.mainForDisplayObjects
#for displaying new quad - to accomodate new proportions
reactivateMainObj(dispObj.shader_program, dispObj.vbo,newCalcDim  )

glClear(GL_COLOR_BUFFER_BIT)
actor.actor.currentlyDispDat=singleSlDat = singleSlDat
updateImagesDisplayed(singleSlDat
                    ,actor.actor.mainForDisplayObjects
                    ,actor.actor.textDispObj
                    ,newCalcDim
                    ,actor.actor.valueForMasToSet      )

     #saving information about current slice for future reference
actor.actor.currentDisplayedSlice = current
# to enbling getting back
if(toBeSavedForBack)
    addToforUndoVector(actor, ()-> processKeysInfo( Option(old),actor, keyInfo,false ))
end#if

end#processKeysInfo
end#ChangePlane


export ProductionModel, ObservationProduce, ObservationGroupProduce, DataProduce

"Production"
struct ProductionModel <: AbstractProductionModel end

# Abstract data strucutres
#---------------------------

struct ObservationProduce{PM<:ProductionModel,T<:Real, V1<:AbstractVector{T}, V2<:AbstractVector{T}, M <:AbstractMatrix{T}} <: AbstractObservation
    model::PM
    y::V1
    x::M
    xsum::V2
    nu::V1
    xpnu::V2
    nusum::T
    nusumsq::T
    function ObservationProduce(model::PM, y::V1, x::M, xsum::V2, nu::V1, xpnu::V2,nusum::T,nusumsq::T) where {PM<:ProductionModel,T<:Real, V1<:AbstractVector{T}, V2<:AbstractVector{T}, M <:AbstractMatrix{T}}
        k,n = size(x)
        length(nu) == length(y) == n  || throw(DimensionMismatch())
        size(xpnu,1) == size(xsum,1) == k || throw(DimensionMismatch())
        return new{PM,T,V1,V2,M}(model,y,x,xsum,nu,xpnu,nusum,nusumsq)
    end
end

struct DataProduce{M<:ProductionModel,T<:Real} <: AbstractDataSet
    model::M
    y::Vector{T}
    x::Matrix{T}
    xsum::Matrix{T}
    nu::Vector{T}
    xpnu::Matrix{T}

    nusum::Vector{T}
    nusumsq::Vector{T}

    obs_ptr::Vector{Int}
    group_ptr::Vector{Int}
    function DataProduce(
        model::M, y::Vector{T}, x::Matrix{T}, xsum::Matrix{T}, nu::Vector{T}, xpnu::Matrix{T},
        nusum::Vector{T}, nusumsq::Vector{T}, obs_ptr::Vector{Int}, group_ptr::Vector{Int}
    ) where {M<:ProductionModel,T<:Real}
        k,n = size(x)
        length(nu) == length(y) == n  || throw(DimensionMismatch())
        size(xsum,1) == k || throw(DimensionMismatch())
        size(xpnu) == size(xsum) || throw(DimensionMismatch())
        issorted(obs_ptr) || throw(error("obs_ptr not sorted"))
        issorted(group_ptr) || throw(error("group_ptr not sorted"))
        length(nusum)+1 == length(nusumsq)+1 == last(group_ptr) == length(obs_ptr) || throw(DimensionMismatch("last(group_ptr)-1 != length(obs_ptr)"))
        last(obs_ptr)-1 == n || throw(DimensionMismatch("last(obs_ptr)-1 != length(y)"))
        return new{M,T}(model,y,x,xsum,nu,xpnu, nusum, nusumsq, obs_ptr,group_ptr)
    end
end

const ObservationGroupProduce = ObservationGroup{<:DataProduce}
const DataOrObsProduction = Union{ObservationProduce,DataProduce}
const AbstractDataStructureProduction = Union{ObservationProduce,DataProduce,ObservationGroupProduce}

# Common interfaces
#---------------------------

_xpnu( d::DataOrObsProduction) = d.xpnu
_nu(   d::DataOrObsProduction) = d.nu
_xsum( d::DataOrObsProduction) = d.xsum
_num_x(d::DataOrObsProduction) = size(_xsum(d),1)
_nusum(d::DataOrObsProduction) = d.nusum
_nusumsq(d::DataOrObsProduction) = d.nusumsq

# Data-specific interfaces
#---------------------------

obs_ptr(  d::DataProduce) = d.obs_ptr
group_ptr(d::DataProduce) = d.group_ptr

# Observation Group interfaces
#---------------------------

function Observation(d::DataProduce, j::Integer)
    rng  = obsrange(d,j)
    y    = view(_y(d), rng)
    nu   = view(_nu(d), rng)
    x    = view(_x(d), :, rng)
    xsum = view(_xsum(d), :, j)
    xpnu = view(_xpnu(d), :, j)
    nusum = getindex(_nusum(d), j)
    nusumsq = getindex(_nusumsq(d),j)
    return ObservationProduce(_model(d),y, x, xsum, nu, xpnu, nusum, nusumsq)
end

getindex(g::ObservationGroupProduce, k) = Observation(_data(g), getindex(grouprange(g), k))
Observation(g::ObservationGroupProduce, k) = getindex(g,k)

# Observation-specific interfaces
#---------------------------

function ==(o1::ObservationProduce, o2::ObservationProduce)
    _y(o1)       == _y(o2) &&
    _x(o1)       == _x(o2) &&
    _xsum(o1)    == _xsum(o2) &&
    _nu(o1)      == _nu(o2) &&
    _xpnu(o1)    == _xpnu(o2) &&
    _nusum(o1)   == _nusum(o2) &&
    _nusumsq(o1) == _nusumsq(o2)
end

function update_nu!(d::DataOrObsProduction, theta)
    length(theta) == _nparm(d) || throw(DimensionMismatch())
    _nu(d) .= _y(d) - _x(d)'*theta_produce_β(d,theta)
    for j in OneTo(_num_obs(d))
        obs = Observation(d, j)
        nu = _nu(obs)
        setindex!(_nusum(d),   sum(nu),    j)
        setindex!(_nusumsq(d), dot(nu,nu), j)
    end
    return nothing
end

function update_xsum!(obs::ObservationProduce)
    xsum = reshape(_xsum(obs), :, 1)
    sum!(xsum, _x(obs))
    return nothing
end

function update_xpnu!(obs::ObservationProduce)
    mul!(_xpnu(obs), _x(obs), _nu(obs))
    return nothing
end

update_xsum!(data::DataProduce) = update_over_obs(update_xsum!, data)
update_xpnu!(data::DataProduce) = update_over_obs(update_xpnu!, data)

function update!(d::DataProduce, theta)
    update_nu!(d, theta)
    update_xpnu!(d)
end

# Abstract data strucutres
#---------------------------

_nparm(d::AbstractDataStructureProduction) = _num_x(d)+3
idx_produce(d::AbstractDataStructureProduction) = OneTo(_nparm(d))
theta_produce(d, theta) = view(theta, idx_produce(d))

idx_produce_ψ() = 1
idx_produce_g() = 2
idx_produce_t() = 3
idx_produce_D() = 4
idx_produce_ψ(  d::Union{AbstractDataStructureProduction,ProductionModel}) =
    idx_produce_ψ()
idx_produce_g(  d::Union{AbstractDataStructureProduction,ProductionModel}) =
    idx_produce_g()
idx_produce_t(  d::Union{AbstractDataStructureProduction,ProductionModel}) =
    idx_produce_t()
idx_produce_D(  d::Union{AbstractDataStructureProduction,ProductionModel}) =
    idx_produce_D()
idx_produce_β(  d::AbstractDataStructureProduction) = 1 .+ (1:_num_x(d))
idx_produce_σ2η(d::AbstractDataStructureProduction) = 2 + _num_x(d)
idx_produce_σ2u(d::AbstractDataStructureProduction) = 3 + _num_x(d)

theta_produce_ψ(  d, theta) = theta[idx_produce_ψ(d)]
theta_produce_β(  d, theta) = view(theta, idx_produce_β(d))
theta_produce_σ2η(d, theta) = theta[idx_produce_σ2η(d)]
theta_produce_σ2u(d, theta) = theta[idx_produce_σ2u(d)]

theta_produce_β(  d::ProductionModel, theta) = view(theta, 2:length(theta)-2)
theta_produce_σ2η(d::ProductionModel, theta) = theta[end-1] # ACTUALLY ση
theta_produce_σ2u(d::ProductionModel, theta) = theta[end]   # ACTUALLY σu

function coefnames(d::AbstractDataStructureProduction)
    nms = Vector{String}(undef, _nparm(d))
    nms[idx_produce_ψ(d)]   = "\\alpha_\\psi"
    beta_idx = idx_produce_β(d) .- 1
    nms[idx_produce_β(d)]  .= ["\\gamma_{$i}" for i in beta_idx]
    nms[idx_produce_g(d)]   = "\\alpha_g"
    nms[idx_produce_σ2η(d)] = "\\sigma_\\eta"
    nms[idx_produce_σ2u(d)] = "\\sigma_u"
    return nms
end

# Dataset generator
#---------------------------

"""
    DataProduce(y, x, xsum, obsptr, grpptr)

Makes dataset and creats xsum, nu automatically
"""
function DataProduce(y::Vector{T}, x::Matrix{T}, obs_ptr::Vector, group_ptr::Vector) where {T}

    k,n = size(x)
    nwells = length(obs_ptr)-1

    xsum = Matrix{T}(undef, k, nwells)
    xpnu = similar(xsum)
    nu   = Vector{T}(undef, n)
    nusum = Vector{T}(undef, nwells)
    nusumsq = similar(nusum)

    data = DataProduce(ProductionModel(),y,x,xsum,nu,xpnu,nusum,nusumsq,obs_ptr,group_ptr)

    update_xsum!(data)

    return data
end

"""
    DataProduce(ngroups, maxwells, ntrange, theta)

Use `theta` to make random dataset with `ngroups` of wells that have
a random number of wells in `0:maxwells` that each have
a random draw of `ntrange` observations.
"""
function DataProduce(psi::Vector, maxwells::Int, ntrange::UnitRange, theta::Vector, args...)
    ngroups = length(psi)
    grouplens = vcat(collect(0:maxwells)..., sample(0:maxwells, ngroups-maxwells-1))
    ncoef = length(theta)-3

    return DataProduce(psi,grouplens,ntrange,theta, args...)
end

function DataProduce(psi::Vector, grouplens::Vector, ntrange::UnitRange, theta::Vector, ivars::Union{Number,Vector}=0)

    length(psi) == length(grouplens) || throw(DimensionMismatch())
    length(theta) > 3 || throw(DimensionMismatch())
    alphapsi = theta[1]
    beta = theta[2:end-2]
    sigeta, sigu = theta[end-1:end]

    ncoef = length(beta)

    # groups
    groupptr = 1 .+ cumsum(vcat(0, grouplens))
    nwells = sum(grouplens)

    us = randn(nwells)
    obslens = vcat(0, sample(ntrange, nwells))
    obsptr = 1 .+ cumsum(obslens)

    # observations
    nobs = last(obsptr)-1

    eta  = randn(nobs)
    x    = randn(ncoef,nobs)
    y    = similar(eta)

    data = DataProduce(y,x,obsptr,groupptr)

    if length(ivars) == length(psi)
        for (i,g) in enumerate(data)
            for (k,o) in enumerate(g)
                j = getindex(grouprange(g), k)
                xj = _x(o)
                xj[1,:] .= ivars[i]
            end
        end
    end

    y .= x'*beta .+ sigeta .* eta

    for (i,g) in enumerate(data)
        for (k,o) in enumerate(g)
            j = getindex(grouprange(g), k)
            y = _y(o) # if we don't do this, errors on julia v1.1.1
            y .+= sigu .* us[j] .+ alphapsi .* psi[i]
        end
    end

    update_xsum!(data)
    update_nu!(data, theta)

    return data
end

DataProduce(ngroups::Int, args...) = DataProduce(randn(ngroups), args...)


# +
import Revise
import GLRenderer
import PoseComposition
import Rotations
import Images
I = Images

R = Rotations
P = PoseComposition
GL = GLRenderer
# -

obj_path = joinpath(@__DIR__, "035_power_drill/textured_simple.obj")
texture_path = joinpath(@__DIR__, "035_power_drill/texture_map.png")

# +
camera_intrinsics = GL.CameraIntrinsics(
    640, 480,
    1000.0, 1000.0,
    320.0, 240.0,
    0.01, 5.0
)

renderer = GL.setup_renderer(camera_intrinsics, GL.TextureMixedMode())

mesh_data = GL.get_mesh_data_from_obj_file(obj_path; tex_path=texture_path)
GL.load_object!(renderer, mesh_data)

rgb_image, depth_image = GL.gl_render(
    renderer, [1], [P.Pose([0.0, 0.0, 1.0], R.RotXYZ(0.1, 0.4, 0.9))],
    P.IDENTITY_POSE
)
img = I.colorview(I.RGBA, permutedims(rgb_image,(3,1,2)))

# +
camera_intrinsics = GL.CameraIntrinsics(
    640, 480,
    1000.0, 1000.0,
    320.0, 240.0,
    0.01, 30.0
)

renderer = GL.setup_renderer(camera_intrinsics, GL.TextureMixedMode())

mesh_data = GL.get_mesh_data_from_obj_file(obj_path; tex_path=texture_path)
GL.load_object!(renderer, mesh_data)

box_mesh = GL.box_mesh_from_dims([10.0, 0.5, 5.0])
GL.load_object!(renderer, box_mesh)

rgb_image, depth_image = GL.gl_render(
    renderer, [1,2], [P.Pose([0.0, 0.0, 1.0], R.RotXYZ(0.1, 0.4, 0.9)), P.Pose([0.0, 2.0, 10.0],P.IDENTITY_ORN)],
    P.IDENTITY_POSE; colors = [I.colorant"red" for _ in 1:2]
)

img = I.colorview(I.RGBA, permutedims(rgb_image,(3,1,2)))
I.save(joinpath(@__DIR__, "imgs/texture_mixed_rendering_rgb_image.png"), img)

img = I.colorview(I.Gray, depth_image ./ maximum(depth_image))
I.save(joinpath(@__DIR__, "imgs/texture_mixed_rendering_depth_image.png"), img)


"""
    EncodedTableRow{M, N} <: Block

Block for processed rows having a tuple of M categorical and
N continuous value collections.
"""

struct EncodedTableRow{M, N, T} <: Block
    catcols::NTuple{M}
    contcols::NTuple{N}
    categorydict::T
end

function EncodedTableRow(catcols, contcols, categorydict)
    EncodedTableRow{length(catcols), length(contcols)}(catcols, contcols, categorydict)
end

function checkblock(::EncodedTableRow{M, N}, x::Tuple{Vector, Vector}) where {M, N}
    length(x[1]) == M && length(x[2]) == N
end

"""
    TabularPreprocessing <: Encoding

Encodes a `TableRow` by applying the following preprocessing steps:
- [`DataAugmentation.NormalizeRow`](#) (for normalizing a row of data for continuous columns)
- [`DataAugmentation.FillMissing`](#) (for filling missing values)
- [`DataAugmentation.Categorify`](#) (for label encoding categorical columns,
    which can be later used for indexing into embedding matrices)
or a sequence of these transformations.
"""
struct TabularPreprocessing{T} <: Encoding
	tfms::T
end

TabularPreprocessing(td::Datasets.TableDataset) = TabularPreprocessing(gettransforms(td))

function encodedblock(::TabularPreprocessing, block::TableRow)
    EncodedTableRow(block.catcols, block.contcols, block.categorydict)
end

function encode(tt::TabularPreprocessing, _, block::TableRow, row)
    columns = Tables.columnnames(row)
    usedrow = NamedTuple(filter(
            x -> x[1] ∈ block.catcols || x[1] ∈ block.contcols,
            collect(zip(columns, row))
        ))
    tfmrow = DataAugmentation.apply(
        tt.tfms,
        DataAugmentation.TabularItem(usedrow, keys(usedrow))
    ).data
    catvals = collect(map(col -> tfmrow[col], block.catcols))
    contvals = collect(map(col -> tfmrow[col], block.contcols))
    (catvals, contvals)
end

function setup(::Type{TabularPreprocessing}, block::TableRow, data::TableDataset)
    return TabularPreprocessing(gettransforms(data, block.catcols, block.contcols))
end

# Utils

"""
The helper functions defined below can be used for quickly constructing a dictionary,
which will be required for creating various tabular transformations available in DataAugmentation.jl.

These functions assume that the table in the TableDataset object td has Tables.jl columnaccess interface defined.
"""
function gettransformdict(td, ::Type{DataAugmentation.NormalizeRow}, cols)
    dict = Dict()
    map(cols) do col
        vals = skipmissing(Tables.getcolumn(td.table, col))
        dict[col] = (Statistics.mean(vals), Statistics.std(vals))
    end
    dict
end

function gettransformdict(td, ::Type{DataAugmentation.FillMissing}, cols)
    dict = Dict()
    map(cols) do col
        vals = skipmissing(Tables.getcolumn(td.table, col))
        dict[col] = Statistics.median(vals)
    end
    dict
end

function gettransformdict(td, ::Type{DataAugmentation.Categorify}, cols)
    dict = Dict()
    map(cols) do col
        vals = Tables.getcolumn(td.table, col)
        dict[col] = unique(vals)
    end
    dict
end

"""
    getcoltypes(td::Datasets.TableDataset)

Returns the categorical and continuous columns present in a `TableDataset`.
"""
function getcoltypes(td::Datasets.TableDataset)
    schema = Tables.schema(td.table)
    catcols = Tuple(name for (name, T) in zip(schema.names, schema.types)
        if T <: Union{Union{Missing, String}, Union{Missing, Symbol}, String, Symbol})
    contcols = Tuple(name for (name, T) in zip(schema.names, schema.types)
        if T <: Union{<:Number, <:Union{Missing, <:Number}})
    catcols, contcols
end

"""
    gettransforms(td::Datasets.TableDataset)

Returns a composition of basic tabular transformations constructed
for the given TableDataset.
"""
function gettransforms(td::TableDataset, catcols, contcols)
    normstats = FastAI.gettransformdict(td, DataAugmentation.NormalizeRow, contcols)
    fmvals = FastAI.gettransformdict(td, DataAugmentation.FillMissing, contcols)
    catdict = FastAI.gettransformdict(td, DataAugmentation.Categorify, catcols)
    normalize = DataAugmentation.NormalizeRow(normstats, contcols)
    categorify = DataAugmentation.Categorify(catdict, catcols)
    fm = DataAugmentation.FillMissing(fmvals, contcols)

    return fm |> normalize |> categorify
end

gettransforms(td::TableDataset) = gettransforms(td, getcoltypes(td)...)


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct CurrentModeControl <: InnerControl
        kpc::Float64
        kic::Float64
        kffv::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
    end

Parameters of an inner loop PI current control using based on Purba, Dhople, Jafarpour, Bullo and Johnson.
"Reduced-order Structure-preserving Model for Parallel-connected Three-phase Grid-tied Inverters."
2017 IEEE 18th Workshop on Control and Modeling for Power Electronics (COMPEL): 1-7.

# Arguments
- `kpc::Float64`: Current controller proportional gain, validation range: `(0, nothing)`
- `kic::Float64`: Current controller integral gain, validation range: `(0, nothing)`
- `kffv::Float64`: Gain to enable feed-forward gain of voltage., validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: The states of the CurrentModeControl model are:
	γd_ic: d-axis integrator state of the PI current controller,
	γq_ic: q-axis integrator state of the PI current controller
- `n_states::Int`: CurrentControl has 2 states
"""
mutable struct CurrentModeControl <: InnerControl
    "Current controller proportional gain"
    kpc::Float64
    "Current controller integral gain"
    kic::Float64
    "Gain to enable feed-forward gain of voltage."
    kffv::Float64
    ext::Dict{String, Any}
    "The states of the CurrentModeControl model are:
	γd_ic: d-axis integrator state of the PI current controller,
	γq_ic: q-axis integrator state of the PI current controller"
    states::Vector{Symbol}
    "CurrentControl has 2 states"
    n_states::Int
end

function CurrentModeControl(kpc, kic, kffv, ext=Dict{String, Any}(), )
    CurrentModeControl(kpc, kic, kffv, ext, [:γd_ic, :γq_ic], 2, )
end

function CurrentModeControl(; kpc, kic, kffv, ext=Dict{String, Any}(), states=[:γd_ic, :γq_ic], n_states=2, )
    CurrentModeControl(kpc, kic, kffv, ext, states, n_states, )
end

# Constructor for demo purposes; non-functional.
function CurrentModeControl(::Nothing)
    CurrentModeControl(;
        kpc=0,
        kic=0,
        kffv=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`CurrentModeControl`](@ref) `kpc`."""
get_kpc(value::CurrentModeControl) = value.kpc
"""Get [`CurrentModeControl`](@ref) `kic`."""
get_kic(value::CurrentModeControl) = value.kic
"""Get [`CurrentModeControl`](@ref) `kffv`."""
get_kffv(value::CurrentModeControl) = value.kffv
"""Get [`CurrentModeControl`](@ref) `ext`."""
get_ext(value::CurrentModeControl) = value.ext
"""Get [`CurrentModeControl`](@ref) `states`."""
get_states(value::CurrentModeControl) = value.states
"""Get [`CurrentModeControl`](@ref) `n_states`."""
get_n_states(value::CurrentModeControl) = value.n_states

"""Set [`CurrentModeControl`](@ref) `kpc`."""
set_kpc!(value::CurrentModeControl, val) = value.kpc = val
"""Set [`CurrentModeControl`](@ref) `kic`."""
set_kic!(value::CurrentModeControl, val) = value.kic = val
"""Set [`CurrentModeControl`](@ref) `kffv`."""
set_kffv!(value::CurrentModeControl, val) = value.kffv = val
"""Set [`CurrentModeControl`](@ref) `ext`."""
set_ext!(value::CurrentModeControl, val) = value.ext = val