module SokobanDirectedModule

import ..GridWorlds as GW
import Random
import ReinforcementLearningBase as RLBase
import ..SokobanUndirectedModule as SUM

#####
##### game logic
#####

const NUM_OBJECTS = SUM.NUM_OBJECTS
const AGENT = SUM.AGENT
const WALL = SUM.WALL
const BOX = SUM.BOX
const TARGET = SUM.TARGET
const NUM_ACTIONS = 4

mutable struct SokobanDirected{R, RNG} <: GW.AbstractGridWorld
    env::SUM.SokobanUndirected{R, RNG}
    agent_direction::Int
end

function SokobanDirected(; R = Float32, file = joinpath(dirname(functionloc(GW.eval)[1]), "envs/sokoban/boxoban-levels/medium/train/000.txt"), rng = Random.GLOBAL_RNG)
    env = SUM.SokobanUndirected(R = R, file = file, rng = rng)
    agent_direction = rand(rng, 0:GW.NUM_DIRECTIONS-1)
    env = SokobanDirected(env, agent_direction)
    GW.reset!(env)
    return env
end

function GW.reset!(env::SokobanDirected)
    GW.reset!(env.env)
    env.agent_direction = rand(env.env.rng, 0:GW.NUM_DIRECTIONS-1)
    return nothing
end

function GW.act!(env::SokobanDirected, action)
    @assert action in Base.OneTo(NUM_ACTIONS) "Invalid action $(action). Action must be in Base.OneTo($(NUM_ACTIONS))"

    inner_env = env.env
    tile_map = inner_env.tile_map
    box_positions = inner_env.box_positions
    box_reward = inner_env.box_reward

    r1 = sum(pos -> tile_map[TARGET, pos], box_positions) * box_reward
    agent_position = inner_env.agent_position
    agent_direction = env.agent_direction

    if action in Base.OneTo(2)
        if action == 1
            dest = GW.move_forward(agent_position, agent_direction)
            beyond_dest = GW.move_forward(dest, agent_direction)
        else
            dest = GW.move_backward(agent_position, agent_direction)
            beyond_dest = GW.move_backward(dest, agent_direction)
        end

        if !tile_map[WALL, dest]
            if !tile_map[BOX, dest]
                tile_map[AGENT, agent_position] = false
                inner_env.agent_position = dest
                tile_map[AGENT, dest] = true
            else
                if !tile_map[BOX, beyond_dest] && !tile_map[WALL, beyond_dest]
                    tile_map[BOX, dest] = false
                    tile_map[BOX, beyond_dest] = true

                    box_idx = findfirst(==(dest), box_positions)
                    box_positions[box_idx] = beyond_dest
                    tile_map[AGENT, agent_position] = false
                    inner_env.agent_position = dest
                    tile_map[AGENT, dest] = true
                end
            end
        end
    elseif action == 3
        env.agent_direction = GW.turn_left(agent_direction)
    else
        env.agent_direction = GW.turn_right(agent_direction)
    end

    inner_env.done = all(pos -> tile_map[TARGET, pos], box_positions)
    r2 = sum(pos -> tile_map[TARGET, pos], box_positions) * box_reward
    inner_env.reward = r2 - r1

    return nothing
end

#####
##### miscellaneous
#####

GW.get_height(env::SokobanDirected) = GW.get_height(env.env)
GW.get_width(env::SokobanDirected) = GW.get_width(env.env)

GW.get_action_names(env::SokobanDirected) = (:MOVE_FORWARD, :MOVE_BACKWARD, :TURN_LEFT, :TURN_RIGHT)
GW.get_object_names(env::SokobanDirected) = GW.get_object_names(env.env)

function GW.get_pretty_tile_map(env::SokobanDirected, position::CartesianIndex{2})
    characters = ('☻', '█', '▒', '✖', '→', '↑', '←', '↓', '⋅')

    object = findfirst(@view env.env.tile_map[:, position])
    if isnothing(object)
        return characters[end]
    elseif object == AGENT
        return characters[NUM_OBJECTS + 1 + env.agent_direction]
    else
        return characters[object]
    end
end

function GW.get_pretty_sub_tile_map(env::SokobanDirected, window_size, position::CartesianIndex{2})
    tile_map = env.env.tile_map
    agent_position = env.env.agent_position
    agent_direction = env.agent_direction

    characters = ('☻', '█', '▒', '✖', '→', '↑', '←', '↓', '⋅')

    sub_tile_map = GW.get_sub_tile_map(tile_map, agent_position, window_size, agent_direction)

    object = findfirst(@view sub_tile_map[:, position])
    if isnothing(object)
        return characters[end]
    elseif object == AGENT
        return '↓'
    else
        return characters[object]
    end
end

function Base.show(io::IO, ::MIME"text/plain", env::SokobanDirected)
    str = "tile_map:\n"
    str = str * GW.get_pretty_tile_map(env)
    str = str * "\nsub_tile_map:\n"
    str = str * GW.get_pretty_sub_tile_map(env, GW.get_window_size(env))
    str = str * "\nreward: $(env.env.reward)"
    str = str * "\ndone: $(env.env.done)"
    str = str * "\nlevel_number: $(env.env.level_number)"
    str = str * "\naction_names: $(GW.get_action_names(env))"
    str = str * "\nobject_names: $(GW.get_object_names(env))"
    print(io, str)
    return nothing
end

GW.get_action_keys(env::SokobanDirected) = ('w', 's', 'a', 'd')

#####
##### RLBase API
#####

RLBase.StateStyle(env::GW.RLBaseEnv{E}) where {E <: SokobanDirected} = RLBase.InternalState{Any}()
RLBase.state_space(env::GW.RLBaseEnv{E}, ::RLBase.InternalState) where {E <: SokobanDirected} = nothing
RLBase.state(env::GW.RLBaseEnv{E}, ::RLBase.InternalState) where {E <: SokobanDirected} = (env.env.env.tile_map, env.env.agent_direction)

RLBase.reset!(env::GW.RLBaseEnv{E}) where {E <: SokobanDirected} = GW.reset!(env.env)

RLBase.action_space(env::GW.RLBaseEnv{E}) where {E <: SokobanDirected} = Base.OneTo(NUM_ACTIONS)
(env::GW.RLBaseEnv{E})(action) where {E <: SokobanDirected} = GW.act!(env.env, action)

RLBase.reward(env::GW.RLBaseEnv{E}) where {E <: SokobanDirected} = env.env.env.reward
RLBase.is_terminated(env::GW.RLBaseEnv{E}) where {E <: SokobanDirected} = env.env.env.done

end # module


#=
This file is auto-generated. Do not edit.
=#

#! format: off

"""
    mutable struct HydroTurbineGov <: TurbineGov
        R::Float64
        r::Float64
        Tr::Float64
        Tf::Float64
        Tg::Float64
        VELM::Float64
        gate_position_limits::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
        Tw::Float64
        At::Float64
        D_T::Float64
        q_nl::Float64
        P_ref::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
        states_types::Vector{StateTypes}
        internal::InfrastructureSystemsInternal
    end

Hydro Turbine-Governor.

# Arguments
- `R::Float64`: Permanent droop parameter, validation range: `(0, 0.1)`, action if invalid: `warn`
- `r::Float64`: Temporary Droop, validation range: `(0, 2)`, action if invalid: `warn`
- `Tr::Float64`: Governor time constant, validation range: `(eps(), 30)`, action if invalid: `error`
- `Tf::Float64`: Filter Time constant, validation range: `(eps(), 0.1)`
- `Tg::Float64`: Servo time constant, validation range: `(eps(), 1)`, action if invalid: `error`
- `VELM::Float64`: gate velocity limit, validation range: `(eps(), 0.3)`, action if invalid: `error`
- `gate_position_limits::NamedTuple{(:min, :max), Tuple{Float64, Float64}}`: Gate position limits
- `Tw::Float64`: water time constant, validation range: `(eps(), 3)`, action if invalid: `error`
- `At::Float64`: Turbine gain, validation range: `(0.8, 1.5)`, action if invalid: `warn`
- `D_T::Float64`: Turbine Damping, validation range: `(0, 0.5)`, action if invalid: `warn`
- `q_nl::Float64`: No-power flow, validation range: `(0, nothing)`, action if invalid: `warn`
- `P_ref::Float64`: Reference Power Set-point, validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: The states of the HydroTurbineGov model are:
	x_g1: filter_output,
	x_g2: desired gate,
	x_g3: gate opening,
	x_g4: turbine flow
- `n_states::Int`: HYGOV has 4 states
- `states_types::Vector{StateTypes}`: HYGOV has 4 differential states
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct HydroTurbineGov <: TurbineGov
    "Permanent droop parameter"
    R::Float64
    "Temporary Droop"
    r::Float64
    "Governor time constant"
    Tr::Float64
    "Filter Time constant"
    Tf::Float64
    "Servo time constant"
    Tg::Float64
    "gate velocity limit"
    VELM::Float64
    "Gate position limits"
    gate_position_limits::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
    "water time constant"
    Tw::Float64
    "Turbine gain"
    At::Float64
    "Turbine Damping"
    D_T::Float64
    "No-power flow"
    q_nl::Float64
    "Reference Power Set-point"
    P_ref::Float64
    ext::Dict{String, Any}
    "The states of the HydroTurbineGov model are:
	x_g1: filter_output,
	x_g2: desired gate,
	x_g3: gate opening,
	x_g4: turbine flow"
    states::Vector{Symbol}
    "HYGOV has 4 states"
    n_states::Int
    "HYGOV has 4 differential states"
    states_types::Vector{StateTypes}
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function HydroTurbineGov(R, r, Tr, Tf, Tg, VELM, gate_position_limits, Tw, At, D_T, q_nl, P_ref=1.0, ext=Dict{String, Any}(), )
    HydroTurbineGov(R, r, Tr, Tf, Tg, VELM, gate_position_limits, Tw, At, D_T, q_nl, P_ref, ext, [:x_g1, :x_g2, :x_g3, :x_g4], 4, [StateTypes.Differential, StateTypes.Differential, StateTypes.Differential, StateTypes.Differential], InfrastructureSystemsInternal(), )
end

function HydroTurbineGov(; R, r, Tr, Tf, Tg, VELM, gate_position_limits, Tw, At, D_T, q_nl, P_ref=1.0, ext=Dict{String, Any}(), states=[:x_g1, :x_g2, :x_g3, :x_g4], n_states=4, states_types=[StateTypes.Differential, StateTypes.Differential, StateTypes.Differential, StateTypes.Differential], internal=InfrastructureSystemsInternal(), )
    HydroTurbineGov(R, r, Tr, Tf, Tg, VELM, gate_position_limits, Tw, At, D_T, q_nl, P_ref, ext, states, n_states, states_types, internal, )
end

# Constructor for demo purposes; non-functional.
function HydroTurbineGov(::Nothing)
    HydroTurbineGov(;
        R=0,
        r=0,
        Tr=0,
        Tf=0,
        Tg=0,
        VELM=0,
        gate_position_limits=(min=0.0, max=0.0),
        Tw=0,
        At=0,
        D_T=0,
        q_nl=0,
        P_ref=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`HydroTurbineGov`](@ref) `R`."""
get_R(value::HydroTurbineGov) = value.R
"""Get [`HydroTurbineGov`](@ref) `r`."""
get_r(value::HydroTurbineGov) = value.r
"""Get [`HydroTurbineGov`](@ref) `Tr`."""
get_Tr(value::HydroTurbineGov) = value.Tr
"""Get [`HydroTurbineGov`](@ref) `Tf`."""
get_Tf(value::HydroTurbineGov) = value.Tf
"""Get [`HydroTurbineGov`](@ref) `Tg`."""
get_Tg(value::HydroTurbineGov) = value.Tg
"""Get [`HydroTurbineGov`](@ref) `VELM`."""
get_VELM(value::HydroTurbineGov) = value.VELM
"""Get [`HydroTurbineGov`](@ref) `gate_position_limits`."""
get_gate_position_limits(value::HydroTurbineGov) = value.gate_position_limits
"""Get [`HydroTurbineGov`](@ref) `Tw`."""
get_Tw(value::HydroTurbineGov) = value.Tw
"""Get [`HydroTurbineGov`](@ref) `At`."""
get_At(value::HydroTurbineGov) = value.At
"""Get [`HydroTurbineGov`](@ref) `D_T`."""
get_D_T(value::HydroTurbineGov) = value.D_T
"""Get [`HydroTurbineGov`](@ref) `q_nl`."""
get_q_nl(value::HydroTurbineGov) = value.q_nl
"""Get [`HydroTurbineGov`](@ref) `P_ref`."""
get_P_ref(value::HydroTurbineGov) = value.P_ref
"""Get [`HydroTurbineGov`](@ref) `ext`."""
get_ext(value::HydroTurbineGov) = value.ext
"""Get [`HydroTurbineGov`](@ref) `states`."""
get_states(value::HydroTurbineGov) = value.states
"""Get [`HydroTurbineGov`](@ref) `n_states`."""
get_n_states(value::HydroTurbineGov) = value.n_states
"""Get [`HydroTurbineGov`](@ref) `states_types`."""
get_states_types(value::HydroTurbineGov) = value.states_types
"""Get [`HydroTurbineGov`](@ref) `internal`."""
get_internal(value::HydroTurbineGov) = value.internal

"""Set [`HydroTurbineGov`](@ref) `R`."""
set_R!(value::HydroTurbineGov, val) = value.R = val
"""Set [`HydroTurbineGov`](@ref) `r`."""
set_r!(value::HydroTurbineGov, val) = value.r = val
"""Set [`HydroTurbineGov`](@ref) `Tr`."""
set_Tr!(value::HydroTurbineGov, val) = value.Tr = val
"""Set [`HydroTurbineGov`](@ref) `Tf`."""
set_Tf!(value::HydroTurbineGov, val) = value.Tf = val
"""Set [`HydroTurbineGov`](@ref) `Tg`."""
set_Tg!(value::HydroTurbineGov, val) = value.Tg = val
"""Set [`HydroTurbineGov`](@ref) `VELM`."""
set_VELM!(value::HydroTurbineGov, val) = value.VELM = val
"""Set [`HydroTurbineGov`](@ref) `gate_position_limits`."""
set_gate_position_limits!(value::HydroTurbineGov, val) = value.gate_position_limits = val
"""Set [`HydroTurbineGov`](@ref) `Tw`."""
set_Tw!(value::HydroTurbineGov, val) = value.Tw = val
"""Set [`HydroTurbineGov`](@ref) `At`."""
set_At!(value::HydroTurbineGov, val) = value.At = val
"""Set [`HydroTurbineGov`](@ref) `D_T`."""
set_D_T!(value::HydroTurbineGov, val) = value.D_T = val
"""Set [`HydroTurbineGov`](@ref) `q_nl`."""
set_q_nl!(value::HydroTurbineGov, val) = value.q_nl = val
"""Set [`HydroTurbineGov`](@ref) `P_ref`."""
set_P_ref!(value::HydroTurbineGov, val) = value.P_ref = val
"""Set [`HydroTurbineGov`](@ref) `ext`."""
set_ext!(value::HydroTurbineGov, val) = value.ext = val
"""Set [`HydroTurbineGov`](@ref) `states_types`."""
set_states_types!(value::HydroTurbineGov, val) = value.states_types = val


"""
`module  JAC.MultiPhotonDoubleIon`
... a submodel of JAC that contains all methods for computing multi-photon two-electron (double) ionization cross sections.
"""
module MultiPhotonDoubleIon

    using ..Basics, ..ManyElectron, ..Radial

    """
    `struct  MultiPhotonDoubleIon.Settings`
        ... defines a type for the settings in estimating multi-photon two-electron (double) ionization cross sections.

        + multipoles              ::Array{EmMultipole}           ... Specifies the multipoles of the radiation field that are to be included.
        + gauges                  ::Array{UseGauge}              ... Specifies the gauges to be included into the computations.
        + NoPhotons               ::Int64                        ... Number of photons in the multi-photon ionization
        + photonEnergies          ::Array{Float64,1}             ... List of photon energies.
        + printBeforeComputation  ::Bool                         ... True, if all energies and lines are printed before their evaluation.
        + selectLines             ::Bool                         ... True, if lines are selected individually for the computations.
        + selectedLines           ::Array{Tuple{Int64,Int64},1}  ... List of lines, given by tupels (inital-level, final-level).

    """
    struct Settings
        multipoles                ::Array{EmMultipole}
        gauges                    ::Array{UseGauge}
        NoPhotons                 ::Int64
        photonEnergies            ::Array{Float64,1}
        printBeforeComputation    ::Bool
        selectLines               ::Bool
        selectedLines             ::Array{Tuple{Int64,Int64},1}
    end

    """
    `MultiPhotonDoubleIon.Settings()`  ... constructor for the default values of multi-photon two-electron (double) ionization estimates.
    """
    function Settings()
        Settings(EmMultipole[], UseGauge[], 0, Float64[], false, false, Tuple{Int64,Int64}[])
    end

    # `Base.show(io::IO, settings::MultiPhotonDoubleIon.Settings)`
    #		... prepares a proper printout of the variable settings::MultiPhotonDoubleIon.Settings.
    function Base.show(io::IO, settings::MultiPhotonDoubleIon.Settings)
        println(io, "multipoles:               $(settings.multipoles)  ")
        println(io, "gauges:                   $(settings.gauges)  ")
        println(io, "NoPhotons:                $(settings.NoPhotons)  ")
        println(io, "photonEnergies:           $(settings.photonEnergies)  ")
        println(io, "printBeforeComputation:   $(settings.printBeforeComputation)  ")
        println(io, "selectLines:              $(settings.selectLines)  ")
        println(io, "selectedLines:            $(settings.selectedLines)  ")
    end

    """
    `struct  MultiPhotonDoubleIon.Channel`
        ... defines a type for a multi-photon channel to help characterize multi-photon (single-electron) ionization with
            well-defined multipolarities.

        + NoPhotons      ::EmMultipole            ... Number of photons in the multi-photon ionization
        + multipoles     ::Array{EmMultipole,1}   ... Multipoles of all N incoming photons.
        + gauge          ::EmGauge                ... Gauge for dealing with the (coupled) radiation field.
        + kappa1         ::Int64                  ... partial-wave of the first free electron
        + kappa2         ::Int64                  ... partial-wave of the second free electron
        + phase1         ::Float64                ... phase of the first partial wave
        + phase2         ::Float64                ... phase of the second partial wave
        + symmetry1      ::LevelSymmetry          ... total angular momentum and parity of scattering state after coupling of first electron
        + symmetry2      ::LevelSymmetry          ... total angular momentum and parity of scattering state after coupling of second electron
        + amplitude      ::Complex{Float64}       ... multi-photon ionization amplitude associated with the given channel.
   """
    struct  Channel
        NoPhotons        ::EmMultipole
        multipoles       ::Array{EmMultipole,1}
        gauge            ::EmGauge
        kappa1           ::Int64
        kappa2           ::Int64
        phase1           ::Float64
        phase2           ::Float64
        symmetry1        ::LevelSymmetry
        symmetry2        ::LevelSymmetry
        amplitude        ::Complex{Float64}
    end

    # MultiPhotonDoubleIon line between initial and final (bound-state) levels
    """
    `struct  MultiPhotonDoubleIon.Line`  ... defines a type for a multi-photon ionization line that may include the definition of channels.

        + initialLevel     ::Level           ... initial-(state) level
        + finalLevel       ::Level           ... final-(state) level
        + NoPhotons        ::Int64           ... Number of photons in the multi-photon process
        + photonEnergy     ::Float64         ... Energy of the incoming photons; all photons are assumed to have equal energy.
        + crossSection     ::EmProperty      ... Cross section for this multi-photon ionization.
        + hasChannels      ::Bool            ... Determines whether the individual (sub-) channels are defined in terms of
                                                 their multipolarities, etc., or not.
        + channels         ::Array{MultiPhotonDoubleIon.Channel,1}  ... List of MultiPhotonDoubleIon.Channels of this line.
    """
    struct  Line
        initialLevel       ::Level
        finalLevel         ::Level
        NoPhotons          ::Int64
        photonEnergy       ::Float64
        crossSection       ::EmProperty
        hasChannels        ::Bool
        channels           ::Array{MultiPhotonDoubleIon.Channel,1}
    end

    # `Base.show(io::IO, line::MultiPhotonDoubleIon.Line)`  ... prepares a proper printout of the variable line::MultiPhotonDoubleIon.Line.
    function Base.show(io::IO, line::MultiPhotonDoubleIon.Line)
        println(io, "initialLevel:      $(line.initialLevel)  ")
        println(io, "finalLevel:        $(line.finalLevel)  ")
        println(io, "NoPhotons:         $(line.NoPhotons)  ")
        println(io, "photonEnergy:      $(line.photonEnergy)  ")
        println(io, "crossSection:      $(line.crossSection)  ")
        println(io, "hasChannels:       $(line.hasChannels)  ")
        println(io, "channels:          $(line.channels)  ")
    end

    """
    `MultiPhotonDoubleIon.computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                                       settings::MultiPhotonDoubleIon.Settings; output=true)`
        ... to compute the multiphoton transition amplitudes and all properties as requested by the given settings.
            A list of lines::Array{MultiPhotonDoubleIon.Lines} is returned.
    """
    function  computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                           settings::MultiPhotonDoubleIon.Settings; output=true)
        println("")
        printstyled("MultiPhotonDoubleIon.computeLines(): The computation of multi-photon double ionization amplitudes starts now ... \n", color=:light_green)
        printstyled("----------------------------------------------------- ---------------------------------------------------------- \n", color=:light_green)
        println("")
        #
        #
        pathways = "Not yet implemented !"
        #
        if    output    return( pathways )
        else            return( nothing )
        end
    end

end # module


function to_opengl_mesh!(result, mesh_obs::TOrSignal{<: GeometryBasics.Mesh})
    m_attr = map(convert(Observable, mesh_obs)) do m
        return (m, GeometryBasics.attributes(m))
    end

    result[:faces] = indexbuffer(map(((m,_),)-> faces(m), m_attr))
    result[:vertices] = GLBuffer(map(((m,_),)-> decompose(Point, m), m_attr))

    attribs = m_attr[][2]

    function to_buffer(name, target)
        if haskey(attribs, name)
            val = attribs[name]
            if mesh_obs isa Observable
                val = map(((m, a),)-> a[name], m_attr)
            end
            if val[] isa AbstractVector
                result[target] = GLBuffer(map(metafree, val))
            elseif val[] isa AbstractMatrix
                result[target] = Texture(val)
            else
                error("unsupported attribute: $(name)")
            end
        end
    end
    to_buffer(:color, :vertex_color)
    to_buffer(:uv, :texturecoordinates)
    to_buffer(:uvw, :texturecoordinates)
    to_buffer(:normals, :normals)
    to_buffer(:attribute_id, :attribute_id)

    return result
end

function draw_mesh(@nospecialize(mesh), data::Dict)
    to_opengl_mesh!(data, mesh)
    @gen_defaults! data begin
        shading = true
        backlight = 0f0
        vertex_color = Vec4f(0)
        texturecoordinates = Vec2f(0)
        image = nothing => Texture
        matcap = nothing => Texture
        color_map = nothing => Texture
        color_norm = nothing
        fetch_pixel = false
        uv_scale = Vec2f(1)
        transparency = false
        shader = GLVisualizeShader(
            "util.vert", "mesh.vert", "mesh.frag", "fragment_output.frag",
            view = Dict(
                "light_calc" => light_calc(shading),
                "buffers" => output_buffers(to_value(transparency)),
                "buffer_writes" => output_buffer_writes(to_value(transparency))
            )
        )
    end
    return assemble_shader(data)
end


### THIS IS WIP, runner under construction

### Serial tempering Monte Carlo runner

export SerialTempMC

immutable SerialTemperingMCRunner <: SerialMCRunner
  burnin::Int
  thinning::Int
  nsteps::Int
  r::Range{Int}
  swapperiod::Int

  function SerialTemperingMCRunner(r::Range{Int}, p::Int)
    burnin = first(r)-1
    thinning = r.step
    nsteps = last(r)
    @assert burnin >= 0 "Number of burn-in iterations should be non-negative."
    @assert thinning >= 1 "Thinning should be >= 1."
    @assert nsteps > burnin "Total number of MCMC iterations should be greater than number of burn-in iterations."
    @assert p > 0 "Swap period should be positive."
    new(burnin, thinning, nsteps, r, p)
  end
end

SerialTemperingMCRunner(r::UnitRange{Int}, p::Int) = SerialTemperingMCRunner(first(r):1:last(r), p)

SerialTemperingMCRunner(; burnin::Int=0, thinning::Int=1, nsteps::Int=100, swapperiod::Int=5) =
  SerialTemperingMCRunner((burnin+1):thinning:nsteps, swapperiod)

typealias SerialTempMC SerialTemperingMCRunner

function run(systems::Vector{MCSystem})
  tic()

  local nchains = length(systems)
  local modelsize = systems[1].model.size
  local runner = systems[1].runner
  local jobtype = typeof(systems[1].job)

  @assert all(map(s->s.model.size, systems) .== modelsize) "Models must have the same parameter vector size."
  @assert isa(runner, SerialTempMC) "Runners must be of SerialTempMC type."
  @assert all(map(s->s.runner, systems) .== runner) "The runners in the vector of MC systems must be identical."
  @assert all(map(s->s.job, systems) .== jobtype) "Jobs must be of the same type."

	local logweights = zeros(nchains)  # log of task weights that will be adapted
  local at1 = 1  # Pick starting task

  # Pre-allocation for storing results
  mcchain::MCChain = MCChain(modelsize, length(runner.r))

  local mcstate1 = systems[at1].job.receive()
  local ppars, logtarget, pars
  ppars, logtarget, pars = mcstate1.ppars, mcstate1.logtarget, mcstate1.pars

  # Sampling loop
  i::Int = 1
  for j in 1:runner.nsteps
    if j % swapperiod == 0  # Attempt a task switch
      at2 = rand(1:(nchains-1)) # Pick another task at1 random
      at2 = at2 >= at1 ? at2+1 : at2

      mcstate2 = systems[at2].job.reset(pars)
      if rand() < exp(logtarget - mcstate2.logtarget + logweights[at2] - logweights[at1]) # accept swap ?
        at1, mcstate1 = at2, mcstate2
      end
    else
      mcstate1 = consume(systems[at1].task)
    end

    # TODO : add logweights adaptation
    ppars, logtarget, pars = mcstate1.ppars, mcstate1.logtarget, mcstate1.pars

    if in(j, runner.r)
      mcchain.samples[i, :] = ppars
      mcchain.logtargets[i] = logtarget

      i += 1
      #mcchain.misc[:mod][pos] = at1
    end
  end

  mcchain.runtime = toq()
  mcchain
end

# TODO: check that all elements of array contain MCMCTasks of the same type
function resume(systems::Array{MCMCTask}; nsteps::Int=100)
  run(MCMCTask[systems[j].model * systems[j].sampler * SerialTempMC(nsteps=nsteps, swapperiod=systems[j].runner.swapperiod)
    for j in 1:length(systems)])
end


@inline _sortperm(X::AbstractVecOrMat) = size(X,2)==1 ? sortperm(ndims(X)==1 ? X : vec(X), alg=RadixSort) : sortperm(collect(eachrow(X)))  # sort a data matrix

function Init!(o::StrBootTest{T}) where T  # for efficiency when varying r repeatedly to make CI, do stuff once that doesn't depend on r
	o.kX = o.kX₁ + o.kX₂
	o.kX₂==0 && (o.X₂ = zeros(T,o.Nobs,0))
  iszero(o.kY₂) && (o.Y₂ = zeros(T,o.Nobs,0))
  o.kZ = o.kX₁ + o.kY₂
  if o.LIML && o.kX₂==o.kY₂  # exactly identified LIML = 2SLS
  	o.κ = one(T)
  	o.LIML = false
  end
  if !(o.REst = length(o.R₁)>0)  # base model contains no restrictions?
    o.R₁ = zeros(T,0,o.kZ)
    o.r₁ = zeros(T,0)
  end
  isnan(o.κ) && (o.κ = o.kX₂>0 ? one(T) : zero(T))  # if κ in κ-class estimation not specified, it's 0 or 1 for OLS or 2SLS
  o.WRE = !(iszero(o.κ) || o.scorebs) || o.ARubin

  iszero(o.B) && (o.scorebs = true)

  o.haswt = !iszero(nrows(o.wt))
  o.sumwt = o.haswt ? sum(o.wt) : one(T)
  o._Nobs = o.haswt && o.fweights ? o.sumwt : T(o.Nobs)

  if !(iszero(o.NClustVar)  || o.issorted)
		o.subcluster = o.NClustVar - o.nerrclustvar
		p = _sortperm(view(o.ID, :, [collect(o.subcluster+1:o.NClustVar); collect(1:o.subcluster)]))  # sort by err cluster vars, then remaining boot cluster vars
		o.ID = ndims(o.ID)==1 ? o.ID[p] : o.ID[p,:]
		o.X₁ = ndims(o.X₁)==1 ? o.X₁[p] : o.X₁[p,:]
		o.X₂ = ndims(o.X₂)==1 ? o.X₂[p] : o.X₂[p,:]
		o.y₁ = o.y₁[p]
		o.Y₂ = ndims(o.Y₂)==1 ? o.Y₂[p] : o.Y₂[p,:]
		o.haswt && (o.wt = o.wt[p])
		isdefined(o, :FEID) && nrows(o.FEID)>0 && (o.FEID = o.FEID[p])
  end

  if o.WREnonARubin
	  if !iszero(o.NClustVar)
	    o.info✻, o.ID✻ = panelsetupID(o.ID, 1:o.nbootclustvar)
	  else
	    o.info⋂ = o.info✻ = Vector{UnitRange{Int64}}(undef, o.Nobs, 0)  # no clustering, so no collapsing by cluster
	  end
  elseif !iszero(o.NClustVar)
	  o.info✻ = panelsetup(o.ID, 1:min(o.NClustVar,o.nbootclustvar))  # bootstrap cluster grouping defs rel to original data
  else
	  info⋂ = info✻⋂ = o.info✻ = Vector{UnitRange{Int64}}(undef, o.Nobs, 0)  # causes no collapsing of data in panelsum() calls, only multiplying by weights if any
  end
  o.N✻ = nrows(o.info✻)

  if o.bootstrapt
    if o.NClustVar>0
	    minN = T(Inf)

	    combs = [x & 2^y > 0 for x in 2^o.nerrclustvar-1:-1:1, y in o.nerrclustvar-1:-1:0]  # represent all error clustering combinations. First is intersection of all error clustering vars
	    o.clust = Vector{StrClust{T}}(undef, nrows(combs))  # leave out no-cluster combination
	    o.NErrClustCombs = length(o.clust)

	    if o.NClustVar > o.nbootclustvar  # info for grouping by intersections of all bootstrap && clustering vars wrt data; used to speed crosstab UXAR wrt bootstrapping cluster && intersection of all error clusters
		    if o.WREnonARubin && !o.granular
		      o.info✻⋂, ID✻⋂ = panelsetupID(o.ID, 1:o.NClustVar)
		    else
		      o.info✻⋂           = panelsetup(o.ID, 1:o.NClustVar)
		    end
				if o.subcluster>0 && nrows(o.info✻) ≠ nrows(o.info✻⋂)
					throw(ErrorException("\nThis program can only perform the subcluster bootstrap when the bootstrap clusters are nested within the (intersections of the) error clusters.\n"))
				end
			else
		    o.info✻⋂ = o.info✻  # info for grouping by intersections of all bootstrap && clustering vars wrt data; used to speed crosstab UXAR wrt bootstrapping cluster && intersection of all error clusters
		    o.WREnonARubin && !o.granular && (ID✻⋂ = o.ID✻)
	    end

			Nall = length(o.info✻⋂)

	    if o.NClustVar > o.nerrclustvar  # info for intersections of error clustering wrt data
		    if o.WREnonARubin && !o.granular
		      o.info⋂, ID⋂ = panelsetupID(o.ID, o.subcluster+1:o.NClustVar)
		    else
		      o.info⋂      = panelsetup(o.ID, o.subcluster+1:o.NClustVar)
		    end
		    ID⋂_✻⋂ = length(o.info⋂)==o.Nobs ? o.ID : o.ID[first.(o.info⋂),:]  # version of ID matrix with one row for each all-error-cluster-var intersection instead of 1 row for each obs; gets resorted
		    o.ID✻⋂ = Nall==o.Nobs ? o.ID : o.ID[first.(o.info✻⋂),:]  # version of ID matrix with one row for each all-bootstrap && error cluster-var intersection instead of 1 row for each obs
	    else
		    o.info⋂ = o.info✻⋂  # info for intersections of error clustering wrt data
		    o.WREnonARubin && !o.granular && (ID⋂ = ID✻⋂)
		    o.ID✻⋂ = ID⋂_✻⋂ = nrows(o.info⋂)==o.Nobs ? o.ID : o.ID[first.(o.info⋂),:]  # version of ID matrix with one row for each all-error-cluster-var intersection instead of 1 row for each obs; gets resorted
	    end

	    o.BootClust = 2^(o.NClustVar - o.nbootclustvar)  # location of bootstrap clustering within list of cluster combinations

	    @inbounds for c ∈ 1:o.NErrClustCombs  # for each error clustering combination
		    ClustCols = o.subcluster .+ findall(@view combs[c,:])
		    even = isodd(length(ClustCols))  # not a typo

		    if isone(c)
		      if iszero(o.subcluster)
		    	  order = Vector{Int64}(undef,0)
		    	  info  = Vector{UnitRange{Int64}}(undef, 0)  # causes no collapsing of data in panelsum() calls
						N = Nall
		      else
		    	  order = _sortperm(@view ID⋂_✻⋂[:,ClustCols])
		    	  info  = panelsetup(view(ID⋂_✻⋂,order,:), ClustCols)
						N = nrows(info)
		      end
		    else
		      if any(combs[c, minimum(findall(combs[c,:] .≠ combs[c-1,:])):end])  # if this sort ordering same as last to some point and missing thereafter, no need to re-sort
		    	  order = _sortperm(@view ID⋂_✻⋂[:,ClustCols])
		    	  info = panelsetup(view(ID⋂_✻⋂,order,:), ClustCols)
		      else
		    	  order = Vector{Int64}(undef,0)
						info = panelsetup(ID⋂_✻⋂, ClustCols)
		      end
					N = nrows(info)
		    end

		    if o.small
					multiplier = T(N / (N-1))
					minN = min(minN,N)
				else
					multiplier = one(T)
				end
				o.clust[c] = StrClust{T}(N, multiplier, even, order, info)
	    end

	    (o.scorebs || !o.WREnonARubin) &&
		  	(o.ClustShare = o.haswt ? @panelsum(o, o.wt, o.info⋂)/o.sumwt : T.(length.(o.info⋂)./o.Nobs)) # share of observations by group

    else  # if no clustering, cast "robust" as clustering by observation
      o.clust = StrClust{T}(Nobs, small ? o._Nobs / (o._Nobs - one(T)) : one(T), true, Vector{Int64}(undef,0), Vector{UnitRange{Int64}}(undef,0))
      o.NErrClustCombs = one(Int16)
      (o.scorebs || !o.WREnonARubin) &&
    		(o.ClustShare = o.haswt ? o.wt/o.sumwt : fill(one(T)/o._Nobs,o._Nobs))
    end

		o.purerobust = o.robust && !o.scorebs && iszero(o.subcluster) && o.N✻==o.Nobs  # do we ever error-cluster *and* bootstrap-cluster by individual?
		o.granular   = o.WREnonARubin ? 2*o.Nobs*o.B*(2*o.N✻+1) < o.N✻*(o.N✻*o.Nobs+o.clust[1].N*o.B*(o.N✻+1)) :
											o.NClustVar>0 && !o.scorebs && (o.purerobust || (o.clust[1].N+o.N✻)*o.kZ*o.B + (o.clust[1].N-o.N✻)*o.B + o.kZ*o.B < o.clust[1].N*o.kZ^2 + o.Nobs*o.kZ + o.clust[1].N * o.N✻ * o.kZ + o.clust[1].N * o.N✻)

		if o.robust && !o.purerobust
			(o.subcluster>0 || o.granular) && !o.WREnonARubin &&
				(o.info⋂_✻⋂ = panelsetup(o.ID✻⋂, o.subcluster+1:o.NClustVar))  # info for error clusters wrt data collapsed to intersections of all bootstrapping && error clusters; used to speed crosstab UXAR wrt bootstrapping cluster && intersection of all error clusterings
			((o.scorebs && o.B>0) || (o.WREnonARubin && !o.granular && o.bootstrapt)) &&
				(o.JN⋂N✻ = zeros(T, o.clust[1].N, o.N✻))
		end

		if o.WREnonARubin && o.robust && o.bootstrapt && !o.granular
			if iszero(length(ID✻⋂))
				_, ID✻⋂ = panelsetupID(o.ID,              1:o.NClustVar)
				_, ID⋂  = panelsetupID(o.ID, o.subcluster+1:o.NClustVar)
			end
			o.IDCT⋂✻   = Vector{Vector{Int64}}(undef, o.N✻)
			o.infoCT⋂✻ = Vector{Vector{UnitRange{Int64}}}(undef, o.N✻)
			for i ∈ 1:o.N✻
				tmp = ID✻⋂[o.info✻[i]]                    # ID numbers w.r.t. intersection of all bootstrap/error clusterings contained in bootstrap cluster i
				o.infoCT⋂✻[i] = o.info✻⋂[tmp[1]:tmp[end]] # for each of those ID's, panel info for the all-bootstrap/error-clusterings data row groupings
				o.IDCT⋂✻[i] = ID⋂[first.(o.infoCT⋂✻[i])]  # ID numbers of those groupings w.r.t. the all-error-clusterings grouping
			end
		end
  else
	  minN = T(nrows(o.info✻))
	end
o._ID✻⋂=ID✻⋂
	InitFEs(o)

	if o.B>0 && o.robust && o.granular && !o.purerobust && o.bootstrapt && !o.WREnonARubin
		if o.NFE>0 && !o.FEboot
			_, o.ID✻ = panelsetupID(o.ID   , 1:o.nbootclustvar)
		else
			_, o.ID✻_✻⋂  = panelsetupID(o.ID✻⋂, 1:o.nbootclustvar)
		end
	end

	o.enumerate = o.B>0 && o.auxtwtype==rademacher && o.N✻*log(2) < log(o.B)+1e-6  # generate full Rademacher set?
	o.enumerate && (o.maxmatsize = 0)

	o.Nw = iszero(o.maxmatsize) ? one(Int64) : ceil(Int64, (o.B+1) * Float64(max(nrows(o.ID✻), length(o.ID✻_✻⋂), o.N✻) * sizeof(T)) / o.maxmatsize / 1073741824) # 1073741824 = giga(byte)
	if isone(o.Nw)
		MakeWildWeights!(o, o.B, first=true)  # make all wild weights, once
		o.enumerate && (o.B = ncols(o.v) - 1)  # replications reduced to 2^G
		o.WeightGrp = [1:ncols(o.v)]
	else
    o.seed = rand(o.rng, UInt64)
		_B = ceil(Int64, (o.B+1) / o.Nw)
		o.Nw = ceil(Int64, (o.B+1) / _B)
		o.WeightGrp = [(i-1)*_B+1:i*_B for i ∈ 1:o.Nw]
		o.WeightGrp[end] = first(o.WeightGrp[end]):o.B+1
	end

	if o.ML
		o.dof = nrows(o.R)
	else
		if o.ARubin
			o.R = hcat(zeros(o.kX₂,o.kX₁), Matrix(I(o.kX₂)))  # attack surface is all endog vars
			o.R₁ = o.kX₁>0 && nrows(o.R₁)>0 ? hcat(o.R₁[:,1:o.kX₁], zeros(nrows(o.R₁),o.kX₂)) : zeros(0, o.kX)  # and convert model constraints from referring to X₁, Y₂ to X₁, X₂
		end
		o.dof = nrows(o.R)

		if !o.WRE && iszero(o.κ)  # regular OLS
			o.DGP = StrEstimator{T}(true, o.LIML, o.Fuller, o.κ)
			o.Repl = StrEstimator{T}(true, false, zero(T), zero(T))  # XXX isDGP=1 for Repl? doesn't matter?
			setR!(o.DGP, o, o.null ? [o.R₁ ; o.R] : o.R₁)  # DGP constraints: model constraints + null if imposed
			setR!(o.Repl, o, o.R₁)  # model constraints only
			InitVarsOLS!(o.DGP, o, o.Repl.R₁perp)
			InitTestDenoms!(o.DGP, o)
			o.M = o.DGP  # StrEstimator object from which to get A, AR, XAR
		elseif o.ARubin
			if o.willplot  # for plotting/CI purposes get original point estimate since not normally generated
				o.DGP = StrEstimator{T}(true, o.LIML, o.Fuller, o.κ)
				setR!(o.DGP, o, o.R₁, zeros(T,0,o.kZ))  # no-null model
				InitVarsIV!(o.DGP, o)
				EstimateIV!(o.DGP, o, o.r₁)
				o.confpeak = o.DGP.β̈  # estimated coordinate of confidence peak
			end

			o.DGP = StrEstimator{T}(true, false, zero(T), zero(T))
			setR!(o.DGP, o, o.R₁)
			InitVarsARubin!(o.DGP, o)
			InitTestDenoms!(o.DGP, o)
			o.M = o.DGP  # StrEstimator object from which to get A, AR, XAR
			o.kZ = o.kX

		elseif o.WREnonARubin

			o.DGP = StrEstimator{T}(true, T(o.kX₂ ≠ o.kY₂), zero(T), one(T))
			setR!(o.DGP, o, o.null ? [o.R₁ ; o.R] : o.R₁, zeros(T,0,o.kZ))  # DGP constraints: model constraints + null if imposed
			InitVarsIV!(o.DGP, o)
			if !o.null  # if not imposing null, then DGP constraints, κ, Hessian, etc. do not vary with r and can be set now
				EstimateIV!(o.DGP, o, o.r₁)
				MakeResidualsIV!(o.DGP, o)
			end
			o.Repl = StrEstimator{T}(false, o.LIML, o.Fuller, o.κ)
			setR!(o.Repl, o, o.R₁, o.R)
			InitVarsIV!(o.Repl, o)
			EstimateIV!(o.Repl, o, o.r₁)
			InitWRE!(o)

			if o.bootstrapt
				o.δdenom_b = zeros(o.Repl.kZ, o.Repl.kZ)
				o.S✻UMZperp = Vector{Matrix{T}}(undef, o.Repl.kZ+1)
				o.S✻UPX     = Vector{Matrix{T}}(undef, o.Repl.kZ+1)
				o._J⋂ = zeros(o.clust[1].N, o.Repl.kZ)
				!o.granular && (o.S✻⋂u₁XinvXX = Matrix{Matrix{T}}(undef, o.Repl.kZ+1, o.N✻))
				if o.LIML || !o.robust
					o.YY✻_b   = zeros(o.Repl.kZ+1, o.Repl.kZ+1)
					o.YPXY✻_b = zeros(o.Repl.kZ+1, o.Repl.kZ+1)
				end
				o.NFE>0 && (o.bootstrapt || !isone(o.κ) || o.LIML) && (o.CTFEU = Vector{Matrix{T}}(undef, o.Repl.kZ+1))
			end

		else  # the score bootstrap for IV/GMM uses a IV/GMM DGP but then masquerades as an OLS test because most factors are fixed during the bootstrap. To conform, need DGP and Repl objects with different R, R₁, one with FWL, one not

			o.DGP = StrEstimator{T}(true, o.LIML, o.Fuller, o.κ)
			setR!(o.DGP, o, o.null ? [o.R₁ ; o.R] : o.R₁, zeros(T,0,o.kZ))  # DGP constraints: model constraints + null if imposed
			InitVarsIV!(o.DGP, o)
			o.Repl = StrEstimator{T}(true, o.LIML, o.Fuller, o.κ)
			setR!(o.Repl, o, o.R₁, I)  # process replication restraints = model constraints only
			InitVarsIV!(o.Repl, o, o.Repl.R₁perp)
			EstimateIV!(o.Repl, o, o.r₁)  # bit inefficient to estimate in both objects, but maintains the conformity
			InitTestDenoms!(o.Repl, o)
			o.M = o.Repl  # StrEstimator object from which to get A, AR, XAR; DGP follows WRE convention of using FWL, Repl follows OLS convention of not; scorebs for IV/GMM mixes the two
			if !o.null  # if not imposing null, then DGP constraints, κ, Hessian, etc. do not vary with r and can be set now
				EstimateIV!(o.DGP, o, o.r₁)
				MakeResidualsIV!(o.DGP, o)
			end
		end
  end

  if !o.WREnonARubin && o.bootstrapt
		o.denom = [Matrix{T}(undef,0,0) for _ in 1:o.dof, _ in 1:o.dof]
		if o.robust
			o.Kcd =                       Matrix{Matrix{T}}(undef, o.NErrClustCombs, o.dof)
			o.Jcd = iszero(o.B) ? o.Kcd : Matrix{Matrix{T}}(undef, o.NErrClustCombs, o.dof)  # if B = 0, Kcd will be multiplied by v, which is all 1's, and will constitute Jcd
		end

		if o.robust && o.granular<o.NErrClustCombs && o.B>0
			inds = o.subcluster>0 ?
				        [CartesianIndex(j, i) for (j,v) ∈ enumerate(o.info⋂_✻⋂) for i ∈ v] :  # crosstab c,c* is wide
								o.NClustVar == o.nbootclustvar ?
										[CartesianIndex(i, i) for i ∈ 1:Nall] :  # crosstab c,c* is square
										[CartesianIndex(i, j) for (j,v) ∈ enumerate(o.clust[o.BootClust].info) for i ∈ v]  # crosstab c,c* is tall
			o.crosstab⋂✻ind = LinearIndices(FakeArray(Tuple(max(inds...))...))[inds]
		end
	end

  o.small && (o.dof_r = o.NClustVar>0 ? minN - one(T) : o._Nobs - (o.kZ + o.NFE))  # floating-point dof_r allows for fractional fweights, FWIW...

  o.sqrt = isone(o.dof)  # work with t/z stats instead of F/chi2?

  if o.small
		o.multiplier = (o.smallsample = (o._Nobs - o.kZ - o.FEdfadj * o.NFE) / (o._Nobs - o.robust)) / o.dof  # divide by # of constraints because F stat is so defined
  else
		o.multiplier = o.smallsample = 1
  end

  !(o.robust || o.ML) && (o.multiplier *= o._Nobs)  # will turn sum of squared errors in denom of t/z into mean
  o.sqrt && (o.multiplier = √o.multiplier)

  o.dist = fill(T(NaN), 1, o.B+1)
  (o.Nw>1 || o.WREnonARubin || (!o.null && o.dof≤2)) && (o.numer = fill(T(NaN), o.dof, o.B+1))

  if o.WREnonARubin
		if o.Repl.kZ>1
			o.bootstrapt && o.robust &&
      	(o.Zyg = Vector{Matrix{T}}(undef,o.Repl.kZ))
			o.numer_b = Vector{T}(undef,nrows(o.Repl.RRpar))
		end
		o.bootstrapt && o.robust &&
    	(o.crosstabBootind = o.Nobs==o.N✻ ? Vector(diagind(FakeArray(o.N✻,o.N✻))) :
			                                    LinearIndices(FakeArray(o.Nobs,o.N✻))[CartesianIndex.(1:o.Nobs, o.ID✻)])
	else
		o.poles = o.anchor = zeros(T,0)
		o.interpolable = o.bootstrapt && o.B>0 && o.null && o.Nw==1 && (iszero(o.κ) || o.ARubin)
		if o.interpolable
			o.∂numer∂r = Vector{Matrix{T}}(undef, o.q)
			o.interpolate_u = !(o.robust || o.ML)
			o.interpolate_u && (o.∂u∂r = Vector{Matrix{T}}(undef, o.q))
			if o.robust
				o.∂denom∂r   = [Matrix{Matrix{T}}(undef, o.dof, o.dof) for _ in 1:o.q]
				o.∂²denom∂r² = [Matrix{Matrix{T}}(undef, o.dof, o.dof) for _ in 1:o.q, _ in 1:o.q]
				o.∂Jcd∂r     = [Matrix{Matrix{T}}(undef, o.NErrClustCombs, o.dof) for _ in 1:o.q]
			end
		end
  end
	o.initialized = true
	nothing
end

function InitFEs(o::StrBootTest{T}) where T
	if isdefined(o, :FEID) && nrows(o.FEID)>0
		p = _sortperm(o.FEID)
		sortID = o.FEID[p]
		i_FE = 1; o.FEboot = o.B>0 && !o.WREnonARubin && o.NClustVar>0; j = o.Nobs; o._FEID = ones(Int64, o.Nobs)
		o.invFEwt = zeros(T, o.NFE>0 ? o.NFE : o.Nobs)
		o.FEs = Vector{StrFE{T}}(undef, o.NFE>0 ? o.NFE : o.Nobs)
		@inbounds for i ∈ o.Nobs-1:-1:1
			if sortID[i] ≠ sortID[i+1]
				is = @view p[i+1:j]
				if o.haswt
					tmp  = o.wt[is]
					wt = tmp / (sumFEwt = sum(tmp))
				else
					sumFEwt = T(j - i)
					wt = fill(1/sumFEwt, j-i)
				end
				o.FEs[i_FE] = StrFE{T}(is, wt)
				if (o.B>0 && o.robust && o.granular < o.nerrclustvar) || (o.WREnonARubin && o.robust && o.granular && o.bootstrapt)
					o.invFEwt[i_FE] = one(T) / sumFEwt
				end

				j = i

				if o.FEboot  # are all of this FE's obs in same bootstrapping cluster? (But no need to check if B=0 for then CT_WE in 2nd term of (62) orthogonal to v = col of 1's)
					tmp = o.ID[is, 1:o.nbootclustvar]
					o.FEboot = all(tmp .== view(tmp, 1,:)')
				end
				i_FE += 1
			end
			o._FEID[p[i]] = i_FE
		end
		is = @view p[1:j]
		if o.haswt
			tmp = o.wt[is]
			wt = tmp / (sumFEwt = sum(tmp))
		else
			sumFEwt = T(j)
			wt = fill(1/sumFEwt, j)
		end
		o.FEs[i_FE] = StrFE{T}(is, wt)
		o.robust && ((o.B>0 && o.granular < o.nerrclustvar) || (o.WREnonARubin && o.granular && o.bootstrapt)) &&
			(o.invFEwt[i_FE] = 1 / sumFEwt)
		o.NFE = i_FE
		resize!(o.invFEwt, o.NFE)
		resize!(o.FEs    , o.NFE)
		if o.FEboot  # are all of this FE's obs in same bootstrapping cluster?
			tmp = o.ID[is, 1:o.nbootclustvar]
			o.FEboot = all(tmp .== @view tmp[1,:])
		end

		o.X₁ = partialFE(o, o.X₁)  # don't overwrite caller's data
		o.X₂ = partialFE(o, o.X₂)
		o.y₁ = partialFE(o, o.y₁)
		o.Y₂ = partialFE(o, o.Y₂)
	end
	nothing
end

# draw wild weight matrix of width _B. If first=true, insert column of 1s at front. For non-Anderson-Rubin WRE, subtract 1 from all
const ϕ = (1 + √5)/2

function MakeWildWeights!(o::StrBootTest{T}, _B::Integer; first::Bool=true) where T
  if _B>0  # in scoretest or waldtest WRE, still make v a col of 1's
    if o.enumerate
			o.v = o.WREnonARubin ? [zeros(o.N✻) count_binary(o.N✻, -2, 0)] :  # complete Rademacher set
								[ones( o.N✻) count_binary(o.N✻, -1, 1)]
		elseif o.auxtwtype==normal
			o.v = randn(o.rng, T, o.N✻, _B+first)
			o.WREnonARubin && (o.v .-= one(T))
		elseif o.auxtwtype==gamma
			tmp = quantile.(Gamma(4,.5), rand(o.rng, o.N✻, _B+first))
			o.v = T==Float64 ? tmp : T.(tmp)
			o.WREnonARubin && (o.v .-= one(T))
		elseif o.auxtwtype==webb
			o.v = rand(o.rng, T.([-√1.5, -1, -√.5, √.5, 1, √1.5] .- o.WREnonARubin), o.N✻, _B+first)
		elseif o.auxtwtype == mammen
			o.v = getindex.(Ref(T.([1-ϕ; ϕ] .- o.WREnonARubin)), ceil.(Int16, rand(o.rng, o.N✻, _B+first) ./ (ϕ/√5)))
		elseif o.WREnonARubin  # Rademacher
			o.v = -2rand(o.rng, Bool, o.N✻, _B+first)
		else
			o.v = rand(o.rng, Bool, o.N✻, _B+first) .- T(.5)  # rand(o.rng, Bool, o.N✻, _B+first) .- T(.5)
			o.v_sd = .5
		end

		first && !(o.enumerate && isone(o.v_sd)) && (o.v[:,1] .= o.WREnonARubin ? zero(T) : o.v_sd)  # keep original residuals in first entry to compute base model stat
  else
		o.v = Matrix{T}(undef,0,1)  # in places, ncols(v) indicates B -- 1 for classical tests
  end
	nothing
end


# Mutated

@inline function mask_unambiguous(x::UInt64)
    return nibble_mask(enumerate_nibbles(x), 0x1111111111111111)
end

@inline function mask_unambiguous(a::UInt64, b::UInt64)
    return mask_unambiguous(a) & mask_unambiguous(b)
end

@inline function mutation_bitcount(x::UInt64, y::UInt64, ::DNAAlphabet{4})
    m = mask_unambiguous(x, y)
    d = (x ⊻ y) & m
    return count_nonzero_nibbles(d), count_1111_nibbles(m)
end

@inline mutation_bitcount(x::UInt64, y::UInt64, ::DNAAlphabet{2}) = BioSequences.mismatch_bitcount(x, y, DNAAlphabet{2}())

# Counting mutations
let
    @info "Compiling bit-parallel mutation counter for LongSequence{<:NucleicAcidAlphabet}"

    counter4 = quote
        c, n = mutation_bitcount(x, y, A())
        count += c
        N += n
    end

    counter2 = :(count += mutation_bitcount(x, y, A()))

    BioSequences.compile_2seq_bitpar(
        :count_mutations_bitpar,
        arguments = (:(seqa::LongSequence{A}), :(seqb::LongSequence{A})),
        parameters = (:(A<:NucleicAcidAlphabet{4}),),
        init_code = :(count = 0; N = 0),
        head_code = counter4,
        body_code = counter4,
        tail_code = counter4,
        return_code = :(return count, N)
    ) |> eval

    BioSequences.compile_2seq_bitpar(
        :count_mutations_bitpar,
        arguments = (:(seqa::LongSequence{A}), :(seqb::LongSequence{A})),
        parameters = (:(A<:NucleicAcidAlphabet{2}),),
        init_code = :(count = 0),
        head_code = counter2,
        body_code = counter2,
        tail_code = counter2,
        return_code = :(return count)
    ) |> eval
end

count_n_mutations(seqa::LongSequence{<:NucleicAcidAlphabet}, seqb::LongSequence{<:NucleicAcidAlphabet}) = count_n_mutations(promote(seqa, seqb)...)
count_n_mutations(seqa::LongSequence{A}, seqb::LongSequence{A}) where {A<:NucleicAcidAlphabet} = count_mutations_bitpar(seqa, seqb)

struct MutationCount
    n_mutations::Int
    n_sites::Int
end

MutationCount() = MutationCount(0, 0)

count_mutations(seqa::LongSequence{<:NucleicAcidAlphabet}, seqb::LongSequence{<:NucleicAcidAlphabet}, args...) = count_mutations(promote(seqa, seqb)..., args...)
count_mutations(seqa::LongSequence{DNAAlphabet{2}}, seqb::LongSequence{DNAAlphabet{2}}) = MutationCount(count_n_mutations(seqa, seqb), min(length(seqa), length(seqb)))
count_mutations(seqa::LongSequence{DNAAlphabet{4}}, seqb::LongSequence{DNAAlphabet{4}}) = MutationCount(count_n_mutations(seqa, seqb)...)

function count_mutations(A::LongSequence{T}, B::LongSequence{T}, win_size::Int) where {T<:DNAAlphabet}
    L = min(length(A), length(B))
    last = win_size
    V = Vector{MutationCount}()
    for i in last:L
        first = last - win_size + 1
        a = A[first:last]
        b = B[first:last]
        push!(V, count_mutations(a, b))
    end
    return V
end

function count_mutations(seqs::Vector{<:LongSequence}, args...)
    N = length(seqs)
    M = [count_mutations(seqs[i], seqs[j], args...) for i in 1:N for j in (i + 1):N]
    return PairwiseListMatrix(M, false, eltype(M)())
end

function count_mutations(seqs::Vector{<:LongSequence}, names::Vector{String}, args...)
    return setlabels(count_mutations(seqs, args...), names)
end

function count_mutations(seqs::Dict{String,L}, args...) where {L<:LongSequence}
    seqsvec = Vector{L}(undef, length(seqs))
    names = Vector{String}(undef, length(seqs))
    i = 1
    for (k, v) in seqs
        seqsvec[i] = v
        names[i] = k
        i += 1
    end
    return count_mutations(seqsvec, names, args...)
end


import Base.Broadcast: Broadcasted, Extruded, BroadcastStyle, ArrayStyle

BroadcastStyle(::Type{<:CuArray}) = ArrayStyle{CuArray}()

function Base.similar(bc::Broadcasted{ArrayStyle{CuArray}}, ::Type{T}) where T
    similar(CuArray{T}, axes(bc))
end

# replace base functions with libdevice alternatives
# TODO: do this with Cassette.jl

cufunc(f) = f
cufunc(::Type{T}) where T = (x...) -> T(x...) # broadcasting type ctors isn't GPU compatible

Broadcast.broadcasted(::ArrayStyle{CuArray}, f, args...) =
  Broadcasted{ArrayStyle{CuArray}}(cufunc(f), args, nothing)

libdevice = :[
  cos, cospi, sin, sinpi, tan, acos, asin, atan,
  cosh, sinh, tanh, acosh, asinh, atanh,
  log, log10, log1p, log2, logb, ilogb,
  exp, exp2, exp10, expm1, ldexp,
  erf, erfinv, erfc, erfcinv, erfcx,
  brev, clz, ffs, byte_perm, popc,
  isfinite, isinf, isnan, nearbyint,
  nextafter, signbit, copysign, abs,
  sqrt, rsqrt, cbrt, rcbrt, pow,
  ceil, floor, saturate,
  lgamma, tgamma,
  j0, j1, jn, y0, y1, yn,
  normcdf, normcdfinv, hypot,
  fma, sad, dim, mul24, mul64hi, hadd, rhadd, scalbn].args

for f in libdevice
  isdefined(Base, f) || continue
  @eval cufunc(::typeof(Base.$f)) = CUDAnative.$f
end

using MacroTools

const _cufuncs = copy(libdevice)
cufuncs() = (global _cufuncs; _cufuncs)

function replace_device(ex)
  global _cufuncs
  MacroTools.postwalk(ex) do x
    x in _cufuncs ? :(CuArrays.cufunc($x)) : x
  end
end

macro cufunc(ex)
  global _cufuncs
  def = MacroTools.splitdef(ex)
  f = def[:name]
  def[:name] = Symbol(:cu, f)
  def[:body] = replace_device(def[:body])
  push!(_cufuncs, f)
  quote
    $(esc(MacroTools.combinedef(def)))
    CuArrays.cufunc(::typeof($(esc(f)))) = $(esc(def[:name]))
  end
end

# ForwardDiff Integration
using ForwardDiff: Dual, value, partials, unary_dual_definition
using DiffRules

for f in libdevice
  if haskey(DiffRules.DEFINED_DIFFRULES, (:Base,f,1))
    f == :tanh && continue
    diffrule = DiffRules.DEFINED_DIFFRULES[(:Base,f,1)]
    DiffRules.DEFINED_DIFFRULES[(:CUDAnative,f,1)] =
      (args...) -> replace_device(diffrule(args...))
    eval(unary_dual_definition(:CUDAnative, f))
  end
end

DiffRules.DEFINED_DIFFRULES[(:CUDAnative, :tanh, 1)] = x ->
  replace_device(:(1-tanh(x)^2))
eval(unary_dual_definition(:CUDAnative, :tanh))


using RadeonProRender, GeometryBasics, Colors
const RPR = RadeonProRender

function translationmatrix(t::Vec{3,T}) where {T}
    T0, T1 = zero(T), one(T)
    return Mat{4}(T1, T0, T0, T0, T0, T1, T0, T0, T0, T0, T1, T0, t[1], t[2], t[3], T1)
end

context = RPR.Context()
scene = RPR.Scene(context)
matsys = RPR.MaterialSystem(context, 0)

camera = RPR.Camera(context)
lookat!(camera, Vec3f(8, 0, 5), Vec3f(2, 0, 2), Vec3f(0, 0, 1))
RPR.rprCameraSetFocalLength(camera, 45.0)
set!(scene, camera)
set!(context, scene)

env_light = RadeonProRender.EnvironmentLight(context)
image_path = RPR.assetpath("studio026.exr")
img = RPR.Image(context, image_path)
set!(env_light, img)
setintensityscale!(env_light, 1.1)
push!(scene, env_light)

light = RPR.PointLight(context)
transform!(light, translationmatrix(Vec3f(2, 0, 5)))
f = 16
RPR.setradiantpower!(light, 700 / f, 641 / f, 630 / f)
push!(scene, light)

function add_shape!(scene, context, matsys, mesh; material=RPR.RPR_MATERIAL_NODE_DIFFUSE,
                    color=colorant"green", roughness=0.01)
    rpr_mesh = RPR.Shape(context, mesh)
    push!(scene, rpr_mesh)
    m = RPR.MaterialNode(matsys, material)
    set!(m, RPR.RPR_MATERIAL_INPUT_COLOR, color)
    set!(m, RPR.RPR_MATERIAL_INPUT_ROUGHNESS, roughness, roughness, roughness, roughness)
    set!(rpr_mesh, m)
    return rpr_mesh, m
end

mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(-10, -10, -1), Vec3f(20, 20, 1));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, -10, 0), Vec3f(0.1, 20, 5));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, -2, 0), Vec3f(5, 0.1, 5));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, 2, 0), Vec3f(5, 0.1, 5));
                       color=colorant"white")

mesh, mat = add_shape!(scene, context, matsys, Tesselation(Sphere(Point3f(2, 0, 2), 0.95f0), 100);
                       color=colorant"white")
rpr_mesh = RPR.Shape(context, Tesselation(Sphere(Point3f(2, 0, 2), 1.0f0), 100))
push!(scene, rpr_mesh)
material = RPR.MaterialNode(matsys, RPR.RPR_MATERIAL_NODE_UBERV2)

# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_DIFFUSE_COLOR, 0.358878, 0.497024 , 1.07771, 1.0)
# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_DIFFUSE_WEIGHT, Vec4f(0.0)...)
# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_DIFFUSE_ROUGHNESS, Vec4f(0.1)...)

# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFLECTION_COLOR, 0.01f0, 0.01f0, 0.01f0, 0.01f0)
# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFLECTION_WEIGHT, Vec4f(0.01)...)
# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_TRANSPARENCY, Vec4f(0.99)...)

set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFRACTION_COLOR, Vec4f(0.358878, 0.497024 , 1.07771, 1.0)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFRACTION_WEIGHT, Vec4f(0.8)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFRACTION_ROUGHNESS, Vec4f(0.01)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFRACTION_IOR, Vec4f(1.0)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFRACTION_ABSORPTION_COLOR, Vec4f(0.7, 0.6, 0.09, 1.0)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFRACTION_ABSORPTION_DISTANCE, Vec4f(3.0)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFRACTION_CAUSTICS, Vec4f(0.0)...)

# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFLECTION_ROUGHNESS, 0.1f0, 0f0, 0f0, 0f0)
# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFLECTION_MODE, Int(RPR.RPR_UBER_MATERIAL_IOR_MODE_PBR))
# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_REFLECTION_METALNESS, 0.f0, 0f0, 0f0, 0f0)

set!(material, RPR.RPR_MATERIAL_INPUT_UBER_SSS_SCATTER_COLOR, 0.358878, 0.497024 , 1.07771, 1.0)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_SSS_SCATTER_DISTANCE, Vec4f(3)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_SSS_SCATTER_DIRECTION, Vec4f(0)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_SSS_WEIGHT, Vec4f(1)...)
set!(material, RPR.RPR_MATERIAL_INPUT_UBER_SSS_MULTISCATTER, 1)
# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_BACKSCATTER_WEIGHT, Vec4f(1.0)...)
# set!(material, RPR.RPR_MATERIAL_INPUT_UBER_BACKSCATTER_COLOR, 0.4, 0.4, 0.5, 0.1)

set!(rpr_mesh, material)

fb_size = (800, 600)
frame_buffer = RPR.FrameBuffer(context, RGBA, fb_size)
frame_bufferSolved = RPR.FrameBuffer(context, RGBA, fb_size)
set!(context, RPR.RPR_AOV_COLOR, frame_buffer)

set_standard_tonemapping!(context)

begin
    clear!(frame_buffer)
    RPR.rprContextSetParameterByKey1u(context, RPR.RPR_CONTEXT_ITERATIONS, 200)
    RPR.render(context)
    RPR.rprContextResolveFrameBuffer(context, frame_buffer, frame_bufferSolved, false)
    RPR.save(frame_bufferSolved, "test.png")
end


using .Nodes
using ..Partitioners

# TODO: Allow the instantiation without the specification of width (Could be
#       determined from the training data)
struct Discriminator{P <: AbstractPartitioner,T <: AbstractNode} <: AbstractModel
    partitioner::P
    nodes::Vector{T}
end

function Discriminator{P,T}(partitioner::P; kargs...) where {P <: AbstractPartitioner,T <: AbstractNode}
    nodes = [T(;kargs...) for _ in 1:length(partitioner)]

    Discriminator{P,T}(partitioner, nodes)
end

function Discriminator{RandomPartitioner,T}(width::Int, n::Int; seed::Union{Nothing,Int}=nothing, kargs...) where {T <: AbstractNode}
    partitioner = RandomPartitioner(width, n; seed)

    Discriminator{RandomPartitioner,T}(partitioner; kargs...)
end

function Discriminator{P,T}(X::U, partitioner::P; kargs...) where {P <: AbstractPartitioner,T <: AbstractNode,U <: AbstractVecOrMat{Bool}}
    d = Discriminator{P,T}(partitioner; kargs...)

    train!(d, X)

    return d
end

function Discriminator{RandomPartitioner,T}(X::U, n::Int; seed::Union{Nothing,Int}=nothing, kargs...) where {T <: AbstractNode,U <: AbstractVecOrMat{Bool}}
    d = Discriminator{RandomPartitioner,T}(size(X)[end], n; seed, kargs...)

    train!(d, X)

    return d
end

## Aliases and convenience constructors
# Default node is DictNode with random partitioning
Discriminator(args...; kargs...) = Discriminator{RandomPartitioner,DictNode}(args...; kargs...)

const BitDiscriminator        = Discriminator{RandomPartitioner,DictNode{BitVector}}
const BleachingDiscriminator  = Discriminator{RandomPartitioner,AccNode} # This should be the default classification discriminator

RegressionDiscriminator = Discriminator{RandomPartitioner,RegressionNode}
GeneralizedRegressionDiscriminator = Discriminator{RandomPartitioner,GeneralizedRegressionNode}

# Generic functions
function train!(d::Discriminator{<:AbstractPartitioner,<:AbstractClassificationNode}, X::T) where {T <: AbstractVecOrMat{Bool}}
    for (node, x) in Iterators.zip(d.nodes, partition(d.partitioner, X))
        Nodes.train!(node, x)
    end

    return nothing
end

function train!(d::Discriminator{<:AbstractPartitioner,<:AbstractRegressionNode}, X, y)
    for (node, x) in Iterators.zip(d.nodes, partition(d.partitioner, X))
        Nodes.train!(node, x, y)
    end

    return nothing
end

initial_response(::AbstractVector{Bool}) = zero(Int)
initial_response(X::AbstractMatrix{Bool}) = zeros(Int, size(X, 1))

# TODO: Output response in relative terms? (i.e. divide by the number of nodes)
function predict(d::Discriminator{<:AbstractPartitioner,<:AbstractClassificationNode}, X::AbstractVecOrMat{Bool}; b::Int=0)
    response = initial_response(X)

    for (node, x) in Iterators.zip(d.nodes, partition(d.partitioner, X))
        response += predict(node, x; b)
    end

    return response
end

function predict(d::Discriminator{<:AbstractPartitioner,<:AbstractRegressionNode}, X::AbstractMatrix{Bool})
    [predict(d, x) for x in eachrow(X)]
end

# TODO: In the context of discriminators, predict and predict_response are the same thing.
#       Implement predict_response as an alias to predict

################################################################################
# Specialized training and prediction methods for the regresion variant

function predict(d::Discriminator{P,RegressionNode}, X::AbstractVector{Bool}) where {P <: AbstractPartitioner}
    partial_count = zero(Int)
    estimate = zero(Float64)

    for (node, x) in Iterators.zip(d.nodes, partition(d.partitioner, X))
        count, sum = predict(node, x)

        if node.γ != 1.0
            count = (1 - node.γ^count) / (1 - node.γ)
        end

        partial_count += count
        estimate += sum
    end

    return partial_count == 0 ? estimate : estimate / partial_count
end

################################################################################
# Specialized training and prediction methods for the generalized regresion discriminator

# This expects that α is a constant
function predict(d::Discriminator{P,GeneralizedRegressionNode}, X::AbstractVector{Bool}) where {P <: AbstractPartitioner}
    running_numerator = zero(Float64)
    running_denominator = zero(Float64)

    for (node, x) in Iterators.zip(d.nodes, partition(d.partitioner, X))
        count, estimate = predict(node, x)

        α = node.α
        if count != 0
            running_numerator += estimate / α
            running_denominator += 1 / α
        end
    end

    return running_denominator == 0 ? zero(Float64) : running_numerator / running_denominator
end


# Methods related to online control of dynamic systems

"""
Objects relevant to online control problem
"""
struct OnlineControlProblem
    initial_state
    control_modes
    # TODO: We don't need this IMDP here, get rid of it.
    imdp::IMDP
    pimdp::PIMDP
    dfa::DFA
    extents
    post_image_data
    control_policy
    local_gps::Dict
    global_gps::Dict # stores the paths to the full GPs generated offline
    reward_type::Tuple
    # TODO: There is probably a better place for this somewhere else!
    primary_reward_history
    secondary_reward_history
    online_data # Stores the data collected online
    global_gp_infos # Stores the persistent info dicts
end

"""
Setup the online control problem.
"""
function setup_online_control_problem(x0, modes, imdp, pimdp, dfa, extents, posts, policy, local_gps, global_gps, reward_type)
    reset_pimdp(x0, imdp, dfa, pimdp, extents, policy)

    # Setup local gp data
    online_data = Dict()
    for i=1:length(modes)
        online_data[i] = Dict("x_online" => [], "y_online" => Dict(i=>[] for i in keys(extents[1])))
    end
    # online_data["x_online"] = []
    # online_data["y_online"] = Dict(i=>[] for i in keys(extents[1]))

    global_gp_infos = Dict()

    return OnlineControlProblem(x0, modes, imdp, pimdp, dfa, extents, posts, policy, local_gps, global_gps, reward_type, [], [], online_data, global_gp_infos)
end

"""
One step of a control loop with an offline policy.
"""
function online_control_loop(ocp::OnlineControlProblem)
    res = check_pimdp_exit_conditions(ocp.pimdp, length(ocp.dfa.states))
    if !isnothing(res)
        return res
    end
    x_c = ocp.pimdp.trajectory[end]
    mode = ocp.pimdp.policy[map_pimdp_state_to_index(ocp.pimdp.state_history[end], length(ocp.dfa.states) )]
    x_n = ocp.control_modes[Int(mode)](x_c)
    propogate_pimdp_trajectory(ocp.pimdp, ocp.dfa, ocp.extents, x_n)
    return nothing
end

"""
Concatonate Online and Offline Data
"""
function concatonate_data_points(x_data_offline, y_data_offline, x_data_online, y_data_online)
    x_cat = [x_data_offline x_data_online...]
    y_cat = [y_data_offline; y_data_online]
    return x_cat, y_cat
end

using NearestNeighbors
"""
Create local GP for an extent.
"""
function get_local_data_knn(center, x_data, y_data; num_neighbors = 50)
    num_neighbors = minimum([num_neighbors, size(x_data,2)])
    kdtree = KDTree(x_data);
    sub_idx, _ = knn(kdtree, center, num_neighbors, true)

    # WHY THE HELL IS IT DOING THIS
    if typeof(sub_idx) == Array{Array{Int64,1},1}
        sub_idx = sub_idx[1]
    end

    sub_x_data = x_data[:, sub_idx]
    sub_y_data = y_data[sub_idx]
    return sub_x_data, sub_y_data
end

"""
One step optimal control loop with an offline policy and GP update.
"""
function online_control_loop_optimal(ocp::OnlineControlProblem;
                                     update_gp_flag=false, gp_use_key="global", update_pimdp_ring_states=false,
                                     pimdp_reward_dict=nothing)
    etime = @elapsed begin
        # Check if an exit condition has been reached
        res = check_pimdp_exit_conditions(ocp.pimdp, length(ocp.dfa.states))
        if !isnothing(res)

            return res
        end

        x_c = ocp.pimdp.trajectory[end]     # Current real state
        q_c = ocp.pimdp.state_history[end]
        pimdp_state = map_pimdp_state_to_index(q_c, length(ocp.dfa.states))  # Current PIMDP state index

        best_mode, best_reward, gp_info_dict = get_best_action(ocp, gp_use_key=gp_use_key, pimdp_reward_dict=pimdp_reward_dict)
        x_n = ocp.control_modes[best_mode](x_c)

        propogate_pimdp_trajectory(ocp.pimdp, ocp.dfa, ocp.extents, x_n)

        @debug "BEST ACTION: " best_mode
        @debug "BEST REWARD: " best_reward

        push!(ocp.primary_reward_history, best_reward[1])
        push!(ocp.secondary_reward_history, best_reward[2])

        # TODO: Move this to its own fcn / subfcns
        if update_gp_flag && isnothing(check_pimdp_exit_conditions(ocp.pimdp, length(ocp.dfa.states)))
            dim_keys = keys(gp_info_dict)
            push!(ocp.online_data[best_mode]["x_online"], x_c)
            # Update the GP with a new datapoint
            for (i, dim_key) in enumerate(dim_keys)
                y_meas = (x_n[i] - x_c[i])
                push!(ocp.online_data[best_mode]["y_online"][dim_key], y_meas)
                add_datapoint_to_gp(gp_info_dict[dim_key].gp, (x_c, y_meas))
                update_gp_info(gp_info_dict[dim_key])
            end

            pimdp_states_to_update = [q_c, ]
            if update_pimdp_ring_states
                new_states = get_pimdp_ring_states(q_c)

                for new_state in new_states
                    push!(pimdp_states_to_update, new_state)
                end
            end

            for pimdp_state in pimdp_states_to_update

                # Update transition probability to some states
                extent_id = pimdp_state[1]
                regions_to_update = []
                # pimdp_row = map_pimdp_state_to_Pmat_row(q_c, length(ocp.dfa.states), , best_mode)
                pimdp_rows = get_pimdp_row_idx_same_imdp(pimdp_state, ocp.dfa, length(ocp.pimdp.actions), best_mode)
                @assert length(pimdp_rows) > 0

                # og_post = ocp.post_image_data[best_mode][pimdp_state[1]]

                # Get the new post-image of the GP
                current_extent = ocp.extents[extent_id]
                lb = [current_extent[dim_key][1] for dim_key in dim_keys]
                ub = [current_extent[dim_key][2] for dim_key in dim_keys]

                ## TODO: THIS IS MESSY
                region_post_dict_new = bound_extent(current_extent, lb, ub, gp_info_dict, dim_keys, Dict("x1" => [1., 1.], "x2" => [1., 1.]); known_part_flag=true)
                dummy_region_post_dict = Dict(extent_id => region_post_dict_new)

                for pimdp_row_idx in pimdp_rows
                    prow_min = ocp.pimdp.Pmin[pimdp_row_idx, :]
                    prow_max = ocp.pimdp.Pmax[pimdp_row_idx, :]

                    # Find all states with interval bounds
                    regions_to_update = findall(x->x>0, prow_min .!= prow_max)
                    pimdp_row_vals_min = prow_min[regions_to_update]
                    pimdp_row_vals_max = prow_max[regions_to_update]

                    noise_proc = Truncated(Normal(0, 0.01), -0.01, 0.01)
                    η = 0.01
                    ϵ = -1

                    for target in regions_to_update
                        target_id = pimdp_col_to_imdp_state(target, length(ocp.dfa.states))
                        region_pair = (extent_id, target_id)
                        frame_row = region_pair_transitions(region_pair, ocp.extents, dummy_region_post_dict, ϵ, gp_info_dict, noise_proc, η)

                        if target_id == 1025.
                            minPrn = 1 - frame_row[4]
                            maxPrn = 1 - frame_row[3]

                            if minPrn > prow_min[target]
                                prow_min[target] = minPrn
                            end
                            if maxPrn < prow_max[target]
                                prow_max[target] = maxPrn
                            end
                        else
                            if frame_row[4] < prow_max[target]
                                prow_max[target] = frame_row[4]
                            end
                            if frame_row[3] > prow_min[target]
                                prow_min[target] = frame_row[3]
                            end
                        end
                    end

                    ocp.pimdp.Pmin[pimdp_row_idx, :] = prow_min
                    ocp.pimdp.Pmax[pimdp_row_idx, :] = prow_max
                end
            end
        end
    end
    @debug "Total time: $etime seconds"
end

"""
Update PIMDP row
"""
function update_pimdp_transition_bound_row()
    nothing
end

"""
Get the relevant GP info
"""
function get_gp_info(ocp, mode, current_state; gp_use_case="global")
    # Begin all of the fancy local GP processing
    global_gp_key = "global-$mode"

    # gp_use_case = "use-local-fresh"
    if gp_use_case == "global"
        if global_gp_key in keys(ocp.global_gp_infos)
            global_gp_info_dict = ocp.global_gp_infos[global_gp_key]
        else
            gp_filename = ocp.global_gps[mode]
            gp_info = load_gp("$gp_filename")
            global_gp_info_dict = gp_info[:gp_info]
            ocp.global_gp_infos[global_gp_key] = global_gp_info_dict
        end
        gp_info_dict = global_gp_info_dict
    elseif gp_use_case == "local-persistent"
        nothing
        # gp_key = "$imdp_state-$best_mode"
        # if gp_key in keys(ocp.local_gps)
        #     gp_info_dict = ocp.local_gps[gp_key]
        #     # gp_info_dict
        # else
        #     #gp_filename = filter(x -> occursin("-$imdp_state.bin", x), readdir(gp_dir))[1]
        #     # Load the global GP this could be optimized
        #     gp_filename = ocp.global_gps[best_mode]
        #     gp_info = load_gp("$gp_filename")
        #     # gp_info = load_gp("$gp_dir/$gp_filename")
        #     # gp_local = create_local_gp_knn(x_c, gp_info; num_neighbors = 100)
        #     gp_info_dict = gp_info[:gp_info]
        # end

    elseif gp_use_case == "local-fresh"
        if global_gp_key in keys(ocp.global_gp_infos)
            global_gp_info_dict = ocp.global_gp_infos[global_gp_key]
        else
            gp_filename = ocp.global_gps[mode]
            gp_info = load_gp("$gp_filename")
            global_gp_info_dict = gp_info[:gp_info]
            ocp.global_gp_infos[global_gp_key] = global_gp_info_dict
        end

        # Create a fresh local GP from the global GP
        gp_info_dict = copy(global_gp_info_dict)
        x_offline = global_gp_info_dict["x1"].gp.x
        x_online = ocp.online_data[mode]["x_online"]

        for dim_key in keys(gp_info_dict)
            y_offline = global_gp_info_dict[dim_key].gp.y
            y_online = ocp.online_data[mode]["y_online"][dim_key]
            x_cat, y_cat = concatonate_data_points(x_offline, y_offline, x_online, y_online)
            x_nn, y_nn = get_local_data_knn(current_state, x_cat, y_cat, num_neighbors=75)

            gp_info_dict[dim_key].gp = train_gp_1dim(x_nn, y_nn; se_params=[0., 0.65])
            update_gp_info(gp_info_dict[dim_key])
        end
    else
        @error "Unknown GP use case: " + gp_use_case
    end

    return gp_info_dict
end

"""
Get the policy from a verification result.
"""
function get_policy_from_offline_synthesis(filename)
    policy = Dict()
    open(filename, "r") do f
        while !eof(f)
            row_split = split(readline(f))
            if length(row_split) > 0
                q0 = parse(Int, row_split[1]) + 1
                u = parse(Int, row_split[2]) + 1
                policy[q0] = u
            end
        end
    end
    return policy
end

"""
Gets the best action at each state.
"""
function get_best_action(ocp::OnlineControlProblem; gp_use_key="global", pimdp_reward_dict=nothing)
    # Get the current state and PIMDP state
    x_c = ocp.pimdp.trajectory[end]     # Current real state
    q_c = ocp.pimdp.state_history[end]  # Current IMDP (extent ID) / DFA state
    pimdp_state = map_pimdp_state_to_index(q_c, length(ocp.dfa.states))  # Current PIMDP state index

    # Initialize the rewards and best control mode
    best_reward = [-Inf, -Inf]
    best_ub = -Inf
    best_mode = nothing
    best_gp_info_dict = nothing

    # Check if the current policy gives us a shortcut
    if ocp.control_policy[pimdp_state, 2] == 1.0
        best_mode = Int(ocp.control_policy[pimdp_state, 1])
        best_reward[1] = 2
        best_gp_info_dict = get_gp_info(ocp, best_mode, x_c, gp_use_case=gp_use_key)
    else
        for mode = 1:length(ocp.control_modes)

            gp_info_dict = get_gp_info(ocp, mode, x_c, gp_use_case=gp_use_key)
            # Extract the relevant row of the PIMDP transition matrix
            P_row_idx = map_pimdp_state_to_Pmat_row(q_c, length(ocp.dfa.states), length(ocp.pimdp.actions), mode)
            prow_max = ocp.pimdp.Pmax[P_row_idx, :]

            # Find all states with interval bounds with non-zero UB
            regions_with_nonzero_ub = findall(x->x>0., prow_max)
            @debug regions_with_nonzero_ub

            ##### Test out the point region transition fcn
            ϵ = -1.
            μ, σ = predict_one_step_full(gp_info_dict, x_c)
            # TODO: get rid of this noise dist implicit
            noise_sigma = 0.01
            dist = Truncated(Normal(0., noise_sigma), -noise_sigma, noise_sigma)

            non_trivial_regions = []
            prob_intervals = Dict()
            imdp_extent_ids = []
            nt_pimdp_states = []

            for reg in regions_with_nonzero_ub
                region_idx = pimdp_col_to_imdp_state(reg, length(ocp.dfa.states))
                push!(nt_pimdp_states, (region_idx, q_c[2]))
                push!(imdp_extent_ids, region_idx)
                res = point_region_transitions(μ, σ, region_idx, ocp.extents, ϵ, gp_info_dict; process_noise_dist=dist)
                if res[2] > 0.
                    push!(non_trivial_regions, map_pimdp_state_to_index((Int(region_idx), q_c[2]), length(ocp.dfa.states) ))
                    prob_intervals[map_pimdp_state_to_index((Int(region_idx), q_c[2]), length(ocp.dfa.states) )] = [res[1], res[2]]
                end
            end

            # Extract the potential actions based on the potential regions we could transition to.
            potential_actions = zeros(length(non_trivial_regions), 6)
            total_reward = 0.

            max_reward = 0.
            for (i,region) in enumerate(non_trivial_regions)
                potential_actions[i, 1] = region
                potential_actions[i, 2:4] = ocp.control_policy[region,:]
                potential_actions[i, 5] = prob_intervals[region][1]
                potential_actions[i, 6] = prob_intervals[region][2]
                # total_reward += ocp.control_policy[region,2]
                # total_reward += ocp.control_policy[region,3]
            end
            # total_reward *= 1. / length(non_trivial_regions)

            # Calculate the worst expected lower bound of satisfaction
            sorted_actions = potential_actions[sortperm(potential_actions[:,3]), :]
            region_idx = sorted_actions[:,1]
            minPrs = sorted_actions[:,5]
            maxPrs = sorted_actions[:,6]
            worst_probs = get_true_transition_probabilities(minPrs, maxPrs, collect(1:1:length(region_idx)))

            worst_exp_lb = sum(worst_probs .* sorted_actions[:, 3]')

            sorted_actions2 = potential_actions[sortperm(potential_actions[:,4]), :]
            region_idx = sorted_actions2[:,1]
            minPrs = sorted_actions2[:,5]
            maxPrs = sorted_actions2[:,6]
            worst_probs2 = get_true_transition_probabilities(minPrs, maxPrs, collect(1:1:length(region_idx)))
            # worst_exp_ub = sum(worst_probs2 .* sorted_actions[:, 4]')
            worst_exp_ub = sum(sorted_actions[:,4]')/length(sorted_actions[:,4]')

            if ocp.reward_type[1] == "worst-exp-lb"
                total_reward = worst_exp_lb
            elseif ocp.reward_type[1] == "worst-exp-sum"
                sorted_actions2 = potential_actions[sortperm(potential_actions[:,4]), :]
                region_idx = sorted_actions2[:,1]
                minPrs = sorted_actions2[:,5]
                maxPrs = sorted_actions2[:,6]
                worst_probs2 = get_true_transition_probabilities(minPrs, maxPrs, collect(1:1:length(region_idx)))
                total_reward = 0.5*(sum(worst_probs .* sorted_actions[:, 3]') + sum(worst_probs2 .* sorted_actions2[:, 4]'))
            elseif ocp.reward_type[1] == "worst-exp-lb-ub"
                # TOdo: Move all logic for this reward here
                total_reward = worst_exp_lb
            else
                @error "Unknown primary reward type!!!"
            end

            ###### Defining secondary reward just on DFA can result in weird behavior!!
            if ocp.reward_type[2] == "expected-dfa"
                x_new = [μ["x1"], μ["x2"]]
                x_new = reshape(x_new, 2, 1)
                secondary_reward = propogate_pimdp_trajectory_test(ocp.pimdp, ocp.dfa, ocp.extents, x_new)
            elseif ocp.reward_type[2] == "worst-dfa"
                secondary_reward = calculate_worst_dfa_reward(ocp.pimdp, ocp.dfa, imdp_extent_ids)
            elseif ocp.reward_type[2] == "best-dfa"
                secondary_reward = calculate_best_dfa_reward(ocp.pimdp, ocp.dfa, imdp_extent_ids)
            elseif ocp.reward_type[2] == "exp-probs"
                secondary_reward = (sum(minPrs) + sum(maxPrs))/length(maxPrs)
            elseif ocp.reward_type[2] == "combo"
                dfa_metric = calculate_worst_dfa_reward(ocp.pimdp, ocp.dfa, imdp_extent_ids)
                if dfa_metric > -Inf
                    secondary_reward = sum(minPrs) + sum(maxPrs)
                else
                    secondary_reward = dfa_metric
                end
            elseif ocp.reward_type[2] == "combo2"
                x_new = [μ["x1"], μ["x2"]]
                x_new = reshape(x_new, 2, 1)
                dfa_metric = propogate_pimdp_trajectory_test(ocp.pimdp, ocp.dfa, ocp.extents, x_new)
                if dfa_metric > -Inf
                    if isnothing(pimdp_reward_dict)
                        secondary_reward = calculate_pimdp_metric(ocp, mode)
                    else
                        secondary_reward = -pimdp_reward_dict[q_c]
                    end
                else
                    secondary_reward = dfa_metric
                end
            elseif ocp.reward_type[2] == "worst-exp-ub"
                sorted_actions = potential_actions[sortperm(potential_actions[:,4]), :]
                region_idx = sorted_actions[:,1]
                minPrs = sorted_actions[:,5]
                maxPrs = sorted_actions[:,6]
                worst_probs = get_true_transition_probabilities(minPrs, maxPrs, collect(1:1:length(region_idx)))
                secondary_reward = sum(worst_probs .* sorted_actions[:, 4]')
            elseif ocp.reward_type[2] == "worst-exp-ub-same"
                sorted_actions = potential_actions[sortperm(potential_actions[:,4]), :]
                secondary_reward = sum(worst_probs .* sorted_actions[:, 4]')
            elseif ocp.reward_type[2] == "none"
                secondary_reward = -Inf
            elseif ocp.reward_type[2] == "exp-pimdp-metric1"
                # x_new = [μ["x1"], μ["x2"]]
                # x_new = reshape(x_new, 2, 1)
                # pimdp_state_ex = propogate_pimdp_trajectory_test(ocp.pimdp, ocp.dfa, ocp.extents, x_new)

                secondary_reward = 0
                for ps in nt_pimdp_states
                    val = pimdp_reward_dict[ps]
                    if val < Inf
                        secondary_reward -= pimdp_reward_dict[ps]
                    end
                end
                secondary_reward *= 1/length(nt_pimdp_states)

                # if isnothing(pimdp_reward_dict)
                #     secondary_reward = calculate_pimdp_metric(ocp, mode)
                # else
                #     secondary_reward = -pimdp_reward_dict[pimdp_state_ex]
                # end
            elseif ocp.reward_type[2] == "worst-dfa-exp-pimdp-metric1"
                dfa_metric = calculate_worst_dfa_reward(ocp.pimdp, ocp.dfa, imdp_extent_ids)
                if dfa_metric > -Inf
                    # x_new = [μ["x1"], μ["x2"]]
                    # x_new = reshape(x_new, 2, 1)
                    # pimdp_state_ex = propogate_pimdp_trajectory_test(ocp.pimdp, ocp.dfa, ocp.extents, x_new)
                    secondary_reward = 0
                    for ps in nt_pimdp_states
                        val = pimdp_reward_dict[ps]
                        if val < Inf
                            secondary_reward -= pimdp_reward_dict[ps]
                        end
                    end
                    secondary_reward *= 1/length(nt_pimdp_states)
                    # mean([-pimdp_reward_dict[ps] for ps in nt_pimdp_states])
                    # if isnothing(pimdp_reward_dict)
                    #     secondary_reward = calculate_pimdp_metric(ocp, mode)
                    # else
                    #     secondary_reward = -pimdp_reward_dict[pimdp_state_ex]
                    # end
                else
                    secondary_reward = dfa_metric
                end
            else
                @error "Unknown sec reward type!!!"
            end

            if total_reward > best_reward[1]
                best_mode = mode
                best_reward[1] = total_reward
                best_reward[2] = secondary_reward
                best_gp_info_dict = gp_info_dict
                best_ub = worst_exp_ub
            elseif total_reward == best_reward[1]
                if worst_exp_ub > best_ub && ocp.reward_type[1] == "worst-exp-lb-ub"
                    best_mode = mode
                    best_ub = worst_exp_ub
                    best_reward[2] = secondary_reward
                    best_gp_info_dict = gp_info_dict
                elseif secondary_reward > best_reward[2]
                    best_reward[2] = secondary_reward
                    best_mode = mode
                    best_gp_info_dict = gp_info_dict
                elseif secondary_reward == best_reward[2]
                    # This chooses a random action among a final tie
                    if rand(1)[1] > 0.5
                        best_reward[2] = secondary_reward
                        best_mode = mode
                        best_gp_info_dict = gp_info_dict
                    end
                end
            end
        end
    end
    return best_mode, best_reward, best_gp_info_dict
end

"""
Get pimdp states with non-trivial transition bounds
"""
function get_nontrivial_next_pimdp_states(pimdp, dfa, pimdp_state, mode, foo)
    P_row_idx = map_pimdp_state_to_Pmat_row(pimdp_state, length(dfa.states), length(pimdp.actions), mode)
    # @info pimdp_state
    prow_max = pimdp.Pmax[P_row_idx, :]

    # Find all states with interval bounds with non-zero UB
    regions_with_nonzero_ub = findall(x->x>0., prow_max)
    return regions_with_nonzero_ub
end

"""
Get pimdp states with non-trivial transition bounds for all actions
"""
function get_nontrivial_next_pimdp_states(pimdp, dfa, pimdp_state, num_actions)
    pimdp_states = []
    for mode=1:num_actions
        pimdp_states = pimdp_states ∪ get_nontrivial_next_pimdp_states(pimdp, dfa, pimdp_state, mode, "foo")
    end
    # Get unique states
    unique!(pimdp_states)
    return pimdp_states
end

"""
PIMDP Distance Metric
"""
function calculate_pimdp_metric(pimdp, pimdp_state, dfa, num_actions; max_depth=Inf, helper_dict=nothing)
    # Get the current pimdp state
    # pimdp_state = pimdp.state_history[end]

    # Get the labels that will induce a transition
    target_labels = get_next_dfa_labels(dfa, pimdp_state[2])
    sink_labels = get_dfa_sink_labels(dfa)

    distance = 0
    # From the current PIMDP state, get the next possible states
    pimdp_states = get_nontrivial_next_pimdp_states(pimdp, dfa, pimdp_state, 1, "foo")

    pimdp_states_tried = []
    searching_flag = true

    best_dist = Inf
    while !isempty(pimdp_states) && searching_flag && distance < max_depth
        distance += 1
        next_states = []

        for state in pimdp_states
            imdp_state_idx = Int(pimdp_col_to_imdp_state(state, length(dfa.states)))
            pimdp_statep =  (imdp_state_idx, pimdp_state[2])
            if pimdp_statep ∈ keys(helper_dict)
                searching_flag = false
                subdistance = helper_dict[pimdp_statep] + distance
                if subdistance < best_dist
                    best_dist = subdistance
                end
            elseif pimdp.imdp_label_dict[imdp_state_idx] ∈ target_labels
                searching_flag = false
                if distance < best_dist
                    best_dist = distance
                end
            end

            push!(pimdp_states_tried, state)
            if searching_flag
                next_states = next_states ∪ get_nontrivial_next_pimdp_states(pimdp, dfa, pimdp_statep, num_actions)
            end
        end

        pimdp_states = setdiff(next_states, pimdp_states_tried)
        @assert isempty(pimdp_states ∩ pimdp_states_tried)
    end
        # For each next possible state, check the imdp label
    return best_dist
end

"""
Calculate Min Distance to Accepting state
"""
function calculate_pimdp_metric_all_offline(pimdp, dfa, num_actions; max_depth=Inf)
    # Get the current pimdp state
    pimdp_dist_dict = Dict()

    # Get the labels that will induce a transition
    states_finished = []
    sink_labels = get_dfa_sink_labels(dfa)
    for imdp_idx in keys(pimdp.imdp_label_dict), dfa_state=1:length(dfa.states)
    # for imdp_idx in [1], dfa_state=1:length(dfa.states)
        pimdp_state = (imdp_idx, dfa_state)
        imdp_label = pimdp.imdp_label_dict[imdp_idx]

        if pimdp_state[2] == dfa.sink_state || imdp_label ∈ sink_labels
            pimdp_dist_dict[pimdp_state] = Inf
        elseif pimdp_state[2] == dfa.accepting_state || possible_accept_transition(dfa, dfa_state, imdp_label)
            pimdp_dist_dict[pimdp_state] = 0
        elseif pimdp_state ∈ keys(pimdp_dist_dict)
            continue
        else
            val = calculate_pimdp_metric(pimdp, pimdp_state, dfa, num_actions, helper_dict=pimdp_dist_dict)
            pimdp_dist_dict[pimdp_state] = val
        end
        #     # From the current PIMDP state, get the next possible states
        #     pimdp_states = get_nontrivial_next_pimdp_states(pimdp, dfa, pimdp_state, num_actions)
        #     @info pimdp_states
        #     pimdp_states_tried = []
        #     searching_flag = true
        #     distance = 0
        #     while !isempty(pimdp_states) && searching_flag && distance < max_depth
        #         distance += 1
        #         @info distance
        #         next_states = []
        #         for state in pimdp_states

        #             imdp_idxp = pimdp_col_to_imdp_state(state, length(dfa.states))
        #             imdp_labelp = pimdp.imdp_label_dict[imdp_idxp]
        #             dfa_statep = state % length(dfa.states)
        #             pimdp_statep = (Int(imdp_idxp), dfa_statep)

        #             if (pimdp_statep[2] == 0 && imdp_labelp ∉ sink_labels) || possible_accept_transition(dfa, dfa_statep, imdp_labelp)
        #                 @info "STATE FOUND --------"
        #                 @info pimdp_statep
        #                 searching_flag = false
        #                 pimdp_dist_dict[pimdp_state] = distance
        #             elseif pimdp_statep in keys(pimdp_dist_dict)
        #                 searching_flag = false
        #                 pimdp_dist_dict[pimdp_state] = pimdp_dist_dict[pimdp_statep] + distance
        #             end
        #             push!(pimdp_states_tried, state)
        #             next_states = next_states ∪ get_nontrivial_next_pimdp_states(pimdp, dfa, pimdp_statep, num_actions)
        #             @info next_states
        #         end
        #         pimdp_states = setdiff(next_states, pimdp_states_tried)
        #         # Check for states already procced here
        #         @assert isempty(pimdp_states ∩ pimdp_states_tried)
        #     end
        # end
        @assert pimdp_state ∈ keys(pimdp_dist_dict)
    end
    return pimdp_dist_dict
end

function get_distance_to_dfa_label(ocp, mode)
    # TODO: Finish this
   # Get the current pimdp state
   pimdp_state = ocp.pimdp.state_history[end]

   # Get the labels that will induce a transition
   target_labels = get_next_dfa_labels(ocp.dfa, pimdp_state[2])

   # From the current PIMDP state, get the next possible states
   pimdp_states = get_nontrivial_next_pimdp_states(ocp, pimdp_state, mode)

   distance = Inf
   for state in pimdp_states
        imdp_state_idx = Int(pimdp_col_to_imdp_state(state, length(ocp.dfa.states)))

   end
end

"""
Get the policy from a synthesis result saved as a matrix.
"""
function get_policy_from_offline_synthesis_mat(filename)
    policy = Dict()
    synth_dict = matread(filename)

    policy = synth_dict["policy"]
    minPr = synth_dict["indVmin"]
    maxPr = synth_dict["indVmax"]

    for (i, action) in enumerate(synth_dict["policy"])
        policy[i] = Int(action)
    end
    return policy, minPr, maxPr
end


function Power(arg0::jdouble)
    return Power((jdouble,), arg0)
end

function value(obj::Power, arg0::DerivativeStructure)
    return jcall(obj, "value", DerivativeStructure, (DerivativeStructure,), arg0)
end

function value(obj::Power, arg0::jdouble)
    return jcall(obj, "value", jdouble, (jdouble,), arg0)
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

using Test, Random
using SuiteSparse, LinearAlgebra, SparseArrays

if Base.USE_GPL_LIBS
    include("umfpack.jl")
    include("cholmod.jl")
    include("spqr.jl")

    # Test multithreaded execution
    let p, cmd = `$(Base.julia_cmd()) --depwarn=error --startup-file=no threads.jl`
        # test both nthreads==1 and nthreads>1. spawn a process to test whichever
        # case we are not running currently.
        other_nthreads = Threads.nthreads() == 1 ? 4 : 1
        p = run(
                pipeline(
                    setenv(
                        cmd,
                        "JULIA_NUM_THREADS" => other_nthreads,
                        dir=@__DIR__()),
                    stdout = stdout,
                    stderr = stderr),
                wait = false)
        include("threads.jl")
        if !success(p)
            error("SuiteSparse threads test failed with nthreads == $other_nthreads")
        end
    end
end


using Base.Filesystem

function getDiskSpace( diskName  )
    cmdstr =  `cmd /c dir $diskName`
    # println( cmds )
    aa=read( cmdstr,String)
    bb= split( aa,'\n')
    for b1 in ( bb[end-1],)
        c= split(strip(b1), " ");
        for c1 in (c[3],)
            a11= c1
            # println( a11 )

            c0 = split(a11,",")
            w=""
            for c11 in c0
                w = string(w,c11)
            end
            # println( w ,"  ,    ",a11)
            space=  parse(UInt, w)
            return space
        end
    end
end
function getSizes( root,RAW1MaxSpace=3 , RAW2MaxSpace=3, RAW3MaxSpace=1, RAW4MaxSpace=1,Unit= 1024*1024*1024*1024)
    row1 = 0x0000
    row2= 0x0000
    row3 = 0x0000
    row4 = 0x0000
    row1fs=[]
    row2fs=[]
    row3fs=[]
    row4fs=[]
    a = Filesystem.readdir(root)
    for b in a
        m,e = splitext( b  )
        fsize = Base.Filesystem.filesize( joinpath(root,b) )
        # fsize =
        # println(m ," === ",e ,"  ", fsize)
        if occursin( "1",e)
            row1 += fsize
            append!(row1fs ,[b])
        elseif occursin( "2",e)
            row2 += fsize
            append!(row2fs ,[b])
        elseif occursin( "3",e)
            append!(row3fs ,[b])
            row3 += fsize
        elseif occursin( "4",e)
            append!(row4fs ,[b])
            row4 += fsize
        end
    end
    println("---------------------------")

    for (row,fs,max ) in ( (row1,row1fs,RAW1MaxSpace),(row2,row2fs,RAW2MaxSpace),(row3,row3fs,RAW3MaxSpace) ,(row4,row4fs,RAW4MaxSpace)   )
        sort!( fs,rev=true )
        # row *= Unit #in kbytes
        max *= Unit # in kbytes
        while row > max && length( fs ) >2
            if length( fs ) >1
                fs1 = pop!(fs)
                fullnamefs = joinpath(root,fs1)
                fsize = Base.Filesystem.filesize( fullnamefs )
                # fsize =round(UInt, s  )
                Filesystem.rm(fullnamefs )
                row -= fsize

                println( "$fullnamefs, Cursize(kbytes):$(row/Unit) ,Max(kbytes):$(max/Unit)")
            end
        end
        println()
    end
    # return ( row1,row2,row3,row4 )
end

function testMain(root)
    RAW1MaxSpace=0
    RAW2MaxSpace=0
    RAW3MaxSpace=0
    RAW4MaxSpace=1
    Unit= 1024*1024*1024
    while true
        rowsizes  = getSizes( root ,RAW1MaxSpace, RAW2MaxSpace, RAW3MaxSpace, RAW4MaxSpace,Unit)
        exit(1)
        sleep(10)
        println( "start clear expired files ")
    end
end

# testMain()


cd("/Users/federicoclaudi/Documents/Github/LocomotionControl/analysis/behavior")
import Pkg;
Pkg.activate(".")

include("/Users/federicoclaudi/Documents/Github/LocomotionControl/analysis/behavior/multiple_horizons/simulate.jl")

"""
Run a simulation in which the model can only plan for `planning_horizon` seconds ahead.
"""
function run_simulation_and_plot(; planning_horizon::Float64=.5, n_iter=480, Δt=.01)
    # run global solution
    track = Track(;start_waypoint=2, keep_n_waypoints=-1)

    _, bike, _, globalsolution = run_mtm(
        :dynamics,
        3;
        showtrials=nothing,
        track=track,
        n_iter=5000,
        fcond=State(; u=30, ω=0),
        timed=false,
        showplots=false,
    )

    # plot background & global trajectory
    colors = [
            range(HSL(colorant"red"), stop=HSL(colorant"green"), length=(Int ∘ floor)(n_iter/2))...,
            range(HSL(colorant"green"), stop=HSL(colorant"blue"), length=(Int ∘ ceil)(n_iter/2))...
    ]

    simtracker = SimTracker(Δt=Δt)
    p1 = draw(:arena)
    draw!(FULLTRACK; border_alpha=.25, alpha=0.0)

    for i in pbar(1:n_iter, redirectstdout=false, description="Running horizon: $planning_horizon seconds", expand=true)
        simtracker.iter = i

        # run simulation and store results
        initial_state, solution, shouldstop, track = step(simtracker, globalsolution, planning_horizon)
        shouldstop && break

        # plot stuff

        # draw!(track; color=colors[i], border_lw=5, alpha=0.0)
        i%5 == 0 && plot_bike_trajectory!(solution, bike; showbike=false, label=nothing, color=colors[i], alpha=.8, lw=4)

        if i%10 == 0
            draw!(initial_state; color=colors[i], alpha=1)
        end
        # plot!(; title="T: $(round(simtracker.t; digits=2))s | horizon: $planning_horizon s")

        simtracker.s > 260 && break
    end

    display(p1)
end

run_simulation_and_plot(planning_horizon=.75)
println("done")


@testset "smarts.smarts" begin

@testset "chain" begin
    nullmol = parse(SMARTS, "")
    @test nodecount(nullmol) == 0
    @test edgecount(nullmol) == 0

    aliphatic = parse(SMARTS, "C")
    @test nodeattr(aliphatic, 1).query == (
        :and => (:atomsymbol => :C, :isaromatic => false)
    )
    @test edgecount(aliphatic) == 0

    carbonyl = parse(SMARTS, "[CX3]=[OX1]")
    @test nodeattr(carbonyl, 1).query == (:and => (
        :and => (:atomsymbol => :C, :isaromatic => false),
        :connectivity => 3
    ))
    @test nodeattr(carbonyl, 2).query == (:and => (
        :and => (:atomsymbol => :O, :isaromatic => false),
        :connectivity => 1
    ))

    ether = parse(SMARTS, "[#6][OD2][#6]")
    @test nodeattr(ether, 1).query == (:atomsymbol => :C)
    @test nodeattr(ether, 2).query == (:and => (
        :and => (:atomsymbol => :O, :isaromatic => false),
        :nodedegree => 2
    ))

    notH = parse(SMARTS, "[!#1]")
    @test nodeattr(notH, 1).query == (:not => (:atomsymbol => :H))
end

end # smarts.SMARTS