function GLOM_plot(fig_loc::String, plot_xs::AbstractVector, obs_xs::AbstractVector, feat::AbstractVector, err::AbstractVector, GLOM_feat::AbstractVector, GLOM_err::AbstractVector, ylabel::String; label_feat::String="obs", label_GLOM::String="GLOM")
    plt = scatter(obs_xs, feat; yerror=err, label=label_feat,
        title=star_str * ", obs RV std: $(round(std(obs_rvs); digits=3))",
        xlabel = "Time (days)",
        ylabel = ylabel)
    plot!(plt, plot_xs, GLOM_feat; ribbons=GLOM_err, fillalpha=0.3, label=label_GLOM)
    png(plt, fig_loc)
end
function GLOM_ind_plots(name_prefix::String, plot_xs::AbstractVector, obs_xs::AbstractVector, rvs_and_inds::AbstractVector{<:AbstractVector}, rvs_and_inds_err::AbstractVector{<:AbstractVector}, GLOM_at_plot_xs::AbstractVector{<:Vector}, GLOM_err_at_plot_xs::AbstractVector{<:AbstractVector})
    for i in 2:length(rvs_and_inds)
        GLOM_plot(fig_dir * name_prefix * "ind$(i-1)", plot_xs, obs_xs, rvs_and_inds[i], rvs_and_inds_err[i], GLOM_at_plot_xs[i], GLOM_err_at_plot_xs[i], "ind $(i-1)")
    end
end
function GLOM_plots(name_prefix::String, plot_xs::AbstractVector, obs_xs::AbstractVector, rvs_and_inds::AbstractVector{<:AbstractVector}, rvs_and_inds_err::AbstractVector{<:AbstractVector}, GLOM_at_plot_xs::AbstractVector{<:Vector}, GLOM_err_at_plot_xs::AbstractVector{<:AbstractVector})
    GLOM_plot(fig_dir * name_prefix * "rv", plot_xs, obs_xs, rvs_and_inds[1], rvs_and_inds_err[1], GLOM_at_plot_xs[1], GLOM_err_at_plot_xs[1], "RV (m/s)")
    GLOM_ind_plots(name_prefix, plot_xs, obs_xs, rvs_and_inds, rvs_and_inds_err, GLOM_at_plot_xs, GLOM_err_at_plot_xs)
end
function GLOM_plots(name_prefix::String, plot_xs::AbstractVector, obs_xs::AbstractVector, rvs_and_inds::AbstractVector{<:AbstractVector}, rvs_and_inds_err::AbstractVector{<:AbstractVector}, GLOM_at_plot_xs::AbstractVector{<:Vector}, GLOM_err_at_plot_xs::AbstractVector{<:AbstractVector}, ks::GLOM_RV.KeplerSignal)
    rvs_np = GLOM_RV.remove_kepler(GLOM_RV.get_rv(glo_rv), glo_rv.time, ks)
    GLOM_plot(fig_dir * name_prefix * "rv_nokep", plot_xs, obs_xs, rvs_np, rvs_and_inds_err[1], GLOM_at_plot_xs[1], GLOM_err_at_plot_xs[1], "RV (m/s)"; label_feat="obs-Kep")
    GLOM_plot(fig_dir * name_prefix * "rv", plot_xs, obs_xs, rvs_and_inds[1], rvs_and_inds_err[1], GLOM_at_plot_xs[1] + ustrip.(ks.(plot_xs.*u"d")), GLOM_err_at_plot_xs[1], "RV (m/s)"; label_GLOM="GLOM+Kep")
    GLOM_ind_plots(name_prefix, plot_xs, obs_xs, rvs_and_inds, rvs_and_inds_err, GLOM_at_plot_xs, GLOM_err_at_plot_xs)
end
function plot_helper(glo_rv::GLOM_RV.GLO_RV, plot_xs::AbstractVector, rvs_and_inds::AbstractVector{<:AbstractVector}, rvs_and_inds_err::AbstractVector{<:AbstractVector}, prefix::String, ks::GLOM_RV.KeplerSignal, fit_total_hyperparameters::AbstractVector)
    GLOM_at_plot_xs, GLOM_err_at_plot_xs, GLOM_at_obs_xs = GLOM_RV.GLOM_posteriors(glo_rv.glo, plot_xs, fit_total_hyperparameters; y_obs = GLOM_RV.remove_kepler(glo_rv, ks))
    GLOM_plots(prefix, plot_xs, obs_xs, rvs_and_inds, rvs_and_inds_err, GLOM_at_plot_xs, GLOM_err_at_plot_xs, ks)

    plot_kep_xs = collect(LinRange(0, ustrip(best_period), 1000))
    scatter(remainder(glo.x_obs, ustrip(best_period)), ustrip.(GLOM_RV.get_rv(glo_rv)) - GLOM_at_obs_xs[1]; yerror=ustrip.(GLOM_RV.get_rv_noise(glo_rv)), label="data")
    plot!(plot_kep_xs, ustrip.(ks.(plot_kep_xs.*u"d")); label="kep")
    png(fig_dir * prefix * "rv_phase")
end
function plot_helper(glo::GLOM.GLO, plot_xs::AbstractVector, rvs_and_inds::AbstractVector{<:AbstractVector}, rvs_and_inds_err::AbstractVector{<:AbstractVector}, prefix::String, fit_total_hyperparameters::AbstractVector)
    GLOM_at_plot_xs, GLOM_err_at_plot_xs, GLOM_at_obs_xs = GLOM_RV.GLOM_posteriors(glo, plot_xs, fit_total_hyperparameters)
    GLOM_plots(prefix, plot_xs, obs_xs, rvs_and_inds, rvs_and_inds_err, GLOM_at_plot_xs, GLOM_err_at_plot_xs)
    GLOM_at_obs_xs_inf = GLOM_RV.GLOM_posteriors(glo, fit_total_hyperparameters; inflate_errors=1)
    scatter(obs_rvs, GLOM_at_obs_xs[1];
        title = star_str * ", obs RV std: $(round(std(obs_rvs); digits=3))",
        xlabel = "Observed RVs (m/s)",
        ylabel = "GLOM RVs (m/s)",
        label =  "GLOM posterior, resid std        : $(round(std(obs_rvs - GLOM_at_obs_xs[1]); digits=3)), ρ: $(round(cor(obs_rvs, GLOM_at_obs_xs[1]); digits=3))",
        legend=:topleft)
    scatter!(obs_rvs, GLOM_at_obs_xs_inf[1];
        label = "\" with inflated errors, resid std: $(round(std(obs_rvs - GLOM_at_obs_xs_inf[1]); digits=3)), ρ: $(round(cor(obs_rvs, GLOM_at_obs_xs_inf[1]); digits=3))")
    png(fig_dir * prefix * "inflated_errs")
end


# Adds the unprimed model states to the dictionary of ranges to properly
# fill the Jacobian matrix
# Note: This is not augmentation in the usual sense, which incorporates lags of
# model state variables after the transition equation has been solved for
# but rather, augmenting the model states with lags prior to solution since the
# solution method requires lags
function augment_model_states(endo::AbstractDict{Symbol, UnitRange}, n_model_states::Int64)
    endo_aug = deepcopy(endo)
    for (state::Symbol, inds::UnitRange) in endo
        unprimed_state = unprime(state)
        unprimed_inds::UnitRange  = reindex_unprimed_model_states(inds, n_model_states)
        endo_aug[unprimed_state] = unprimed_inds
    end

    # Ensure all ranges are consecutive
    endo_ranges = endo_aug.vals
  #=  for i in 1:length(endo_ranges)
        if i == 1
            @assert endo_ranges[1].start == 1
        else
            @assert endo_ranges[i].start == endo_ranges[i-1].stop + 1
        end
    end =#

    return endo_aug
end

# For the RANK equivalent of the HANK models
function augment_model_states(endo::AbstractDict{Symbol, Int64}, n_model_states::Int64)
    endo_aug = deepcopy(endo)
    for (state::Symbol, ind::Int) in endo
        unprimed_state           = unprime(state)
        endo_aug[unprimed_state] = ind + n_model_states
    end

    return endo_aug
end

# The reason for this function is that the canonical form for the Klein solution method is
# as follows:
# E_t A([x_{t+1}, y_{t+1}]) = B[x_t, y_t]
# Where we only need to track [x_{t+1}, y_{t+1}] as states when
# we transform the system from canonical form to state space form
# Hence because we only keep track of the t+1 indexed variables, which we denote
# with a ′, we need a way of calculating the indices in the Jacobian corresponding
# to the t indexed variables.
function reindex_unprimed_model_states(inds::UnitRange, n_model_states::Int64)
    return inds .+ n_model_states
end

# Returns a model state variable symbol without the prime
function unprime(state::Symbol)
    return Symbol(replace(string(state), "′" => ""))
end

######################################################################################
# For normalizing the states
function normalize_model_state_indices!(m::AbstractModel)
    normalized_model_states = m.normalized_model_states

    endo                     = m.endogenous_states
    model_state_keys         = collect(keys(endo))
    normalized_model_states = intersect(model_state_keys, normalized_model_states) #findall((in)(normalized_model_states), model_state_keys)

    state_inds                  = get_setting(m, :state_indices)
    jump_inds                   = get_setting(m, :jump_indices)
    state_normalization_factor  = get_setting(m, :backward_looking_states_normalization_factor)
    jump_normalization_factor   = get_setting(m, :jumps_normalization_factor)

    m.endogenous_states = normalize(m, endo, model_state_keys, normalized_model_states,
                                    state_inds, jump_inds,
                                    state_normalization_factor, jump_normalization_factor)
end

# Shift a UnitRange type down by an increment
# If this UnitRange is the first_range in the group being shifted, then do not
# subtract the increment from the start
# e.g.
# If you have 1:80 as your first range, then shift(1:80, -1; first_range = true)
# should return 1:79
# However, if you have 81:160 as range, that is not your first range, then the
# function should return 80:159
function shift(inds::UnitRange, increment::Int64; first_range::Bool = false)
    if first_range
        return UnitRange(inds.start, inds.stop + increment)
    else
        return UnitRange(inds.start + increment, inds.stop + increment)
    end
end

# normalize the distributional states located in indices specified by normalized_state_inds
function normalize(m::AbstractModel,
                   endo::AbstractDict{Symbol, UnitRange},
                   model_state_keys::Vector{Symbol},
                   normalized_model_states::Vector{Symbol},
                   state_inds::AbstractArray{Int64},
                   jump_inds::AbstractArray{Int64},
                   state_normalization_factor::Int64,
                   jump_normalization_factor::Int64)
    n_model_state_vars = length(endo)

    # For each model state that needs to be normalized...
    for state in normalized_model_states
        normalization_factor =
        if state in get_setting(m, :states)
            state_normalization_factor
        elseif state in get_setting(m, :jumps)
            jump_normalization_factor
        end
        # Subtract 1 from both the beginning and end of the UnitRange...
        for state in model_state_keys
            inds = endo[state]
            # except for first UnitRange in a given normalization
            first_range = (state==get_setting(m, :states)[1])
            endo[state] = shift(inds, -normalization_factor, first_range = first_range)
        end
        #=for j in i:n_model_state_vars
            inds = endo[model_state_keys[j]]
            # Except for the first UnitRange in a given normalization
            first_range = (j == i)
            endo[model_state_keys[j]] = shift(inds, -normalization_factor, first_range = first_range)
        end=#
    end

    # Ensure all ranges are consecutive
    endo_ranges = endo.vals
   #= for i in 1:n_model_state_vars
        if i == 1
            @assert endo_ranges[1].start == 1
        else
            @assert endo_ranges[i].start == endo_ranges[i-1].stop + 1
        end
    end =#

    return endo
end

"""
```
stack_indices(key_dict::OrderedDict, keys::Vector{Symbol})
```
Stacks the indices of a OrderedDict{Symbol, AbstractRange} into a Vector of Ints defined by the Ranges in the dictionary.
"""
function stack_indices(key_dict::AbstractDict, keys::Vector{Symbol})
    indices = Vector{Int64}(undef, 0)
    for i in getindex.(Ref(key_dict), keys)
        indices = [indices; i]
    end
    return indices
end


include("learners/learners.jl")
include("explorers/explorers.jl")
include("q_based_policy.jl")


# flagsSet:   flags that must be set
# flagsUnset: flags that must be unset
function processbam!(bamFile::BamFile, ra; mappingQualityThreshold::Int=30,
                     flagsSet::Int=0, flagsUnset::Int=0, ignoreChimeric::Bool=true)
    for r in bamFile
        f = flag(r)
        f&flagsSet   != flagsSet && continue
        f&flagsUnset != 0        && continue

        mapq(r)<mappingQualityThreshold && continue

        ignoreChimeric && hastag(r,tag"SA") && continue

        processread!(ra,r)
    end
end

phred2prob(q::Integer) = min(10.0^(-Float64(q)/10.0),0.75) # cap at 0.75, since more gives lower chances of this base than the others!


using HorizonSideRobots
function main(r::Robot)
    ugol(Nord,Ost)
end

function ugol(west1,west2)

    while (isborder(r,west1)==false) || (isborder(r,west2)==false)
        while isborder(r,west1)==false
            moves!(r,west1)
        end
        while isborder(r,west2)==false
            moves!(r,west2)
        end
        #move!(r,west2)
        #move!(r,west1)
        #move!(r,inverse(west1))
    end
end

function moves!(r::Robot, side::HorizonSide)
    num_steps=0
    while isborder(r,side)==false
        move!(r,side)
        num_steps+=1
    end
    return num_steps
end

r=Robot("5.sit", animate=true)


import Base: ==, convert

function packVariable(dfg::G, v::DFGVariable)::Dict{String, Any} where G <: AbstractDFG
    props = Dict{String, Any}()
    props["label"] = string(v.label)
    props["timestamp"] = string(v.timestamp)
    props["tags"] = JSON2.write(v.tags)
    props["estimateDict"] = JSON2.write(v.estimateDict)
    props["solverDataDict"] = JSON2.write(Dict(keys(v.solverDataDict) .=> map(vnd -> pack(dfg, vnd), values(v.solverDataDict))))
    props["smallData"] = JSON2.write(v.smallData)
    props["ready"] = v.ready
    props["backendset"] = v.backendset
    return props
end

function unpackVariable(dfg::G, packedProps::Dict{String, Any})::DFGVariable where G <: AbstractDFG
    label = Symbol(packedProps["label"])
    timestamp = DateTime(packedProps["timestamp"])
    tags =  JSON2.read(packedProps["tags"], Vector{Symbol})
    #TODO this will work for some time, but unpacking in an <: AbstractPointParametricEst would be lekker.
    estimateDict = JSON2.read(packedProps["estimateDict"], Dict{Symbol, MeanMaxPPE})
    smallData = nothing
    smallData = JSON2.read(packedProps["smallData"], Dict{String, String})

    packed = JSON2.read(packedProps["solverDataDict"], Dict{String, PackedVariableNodeData})
    solverData = Dict(Symbol.(keys(packed)) .=> map(p -> unpack(dfg, p), values(packed)))

    # Rebuild DFGVariable
    variable = DFGVariable(Symbol(packedProps["label"]))
    variable.timestamp = timestamp
    variable.tags = tags
    variable.estimateDict = estimateDict
    variable.solverDataDict = solverData
    variable.smallData = smallData
    variable.ready = packedProps["ready"]
    variable.backendset = packedProps["backendset"]

    return variable
end

function pack(dfg::G, d::VariableNodeData)::PackedVariableNodeData where G <: AbstractDFG
  @debug "Dispatching conversion variable -> packed variable for type $(string(d.softtype))"
  return PackedVariableNodeData(d.val[:],size(d.val,1),
                                d.bw[:], size(d.bw,1),
                                d.BayesNetOutVertIDs,
                                d.dimIDs, d.dims, d.eliminated,
                                d.BayesNetVertID, d.separator,
                                d.softtype != nothing ? string(d.softtype) : nothing, d.initialized, d.inferdim, d.ismargin, d.dontmargin)
end

function unpack(dfg::G, d::PackedVariableNodeData)::VariableNodeData where G <: AbstractDFG
  r3 = d.dimval
  c3 = r3 > 0 ? floor(Int,length(d.vecval)/r3) : 0
  M3 = reshape(d.vecval,r3,c3)

  r4 = d.dimbw
  c4 = r4 > 0 ? floor(Int,length(d.vecbw)/r4) : 0
  M4 = reshape(d.vecbw,r4,c4)

  # TODO -- allow out of module type allocation (future feature, not currently in use)
  @debug "Dispatching conversion packed variable -> variable for type $(string(d.softtype))"
  st = nothing #IncrementalInference.ContinuousMultivariate # eval(parse(d.softtype))
  mainmod = getSerializationModule(dfg)
  mainmod == nothing && error("Serialization module is null - please call setSerializationNamespace!(\"Main\" => Main) in your main program.")
  try
      if d.softtype != ""
          unpackedTypeName = split(d.softtype, "(")[1]
          unpackedTypeName = split(unpackedTypeName, '.')[end]
          @debug "DECODING Softtype = $unpackedTypeName"
          st = getfield(mainmod, Symbol(unpackedTypeName))()
      end
  catch ex
      @error "Unable to deserialize soft type $(d.softtype)"
      io = IOBuffer()
      showerror(io, ex, catch_backtrace())
      err = String(take!(io))
      @error err
  end
  @debug "Net conversion result: $st"

  if st == nothing
      error("The variable doesn't seem to have a softtype. It needs to set up with an InferenceVariable from IIF. This will happen if you use DFG to add serialized variables directly and try use them. Please use IncrementalInference.addVariable().")
  end

  return VariableNodeData(M3,M4, d.BayesNetOutVertIDs,
    d.dimIDs, d.dims, d.eliminated, d.BayesNetVertID, d.separator,
    st, d.initialized, d.inferdim, d.ismargin, d.dontmargin )
end

function compare(a::VariableNodeData, b::VariableNodeData)
    a.val != b.val && @debug("val is not equal")==nothing && return false
    a.bw != b.bw && @debug("bw is not equal")==nothing && return false
    a.BayesNetOutVertIDs != b.BayesNetOutVertIDs && @debug("BayesNetOutVertIDs is not equal")==nothing && return false
    a.dimIDs != b.dimIDs && @debug("dimIDs is not equal")==nothing && return false
    a.dims != b.dims && @debug("dims is not equal")==nothing && return false
    a.eliminated != b.eliminated && @debug("eliminated is not equal")==nothing && return false
    a.BayesNetVertID != b.BayesNetVertID && @debug("BayesNetVertID is not equal")==nothing && return false
    a.separator != b.separator && @debug("separator is not equal")==nothing && return false
    a.initialized != b.initialized && @debug("initialized is not equal")==nothing && return false
    abs(a.inferdim - b.inferdim) > 1e-14 && @debug("inferdim is not equal")==nothing && return false
    a.ismargin != b.ismargin && @debug("ismargin is not equal")==nothing && return false
    a.dontmargin != b.dontmargin && @debug("dontmargin is not equal")==nothing && return false
    typeof(a.softtype) != typeof(b.softtype) && @debug("softtype is not equal")==nothing && return false
    return true
end

"""
    $(SIGNATURES)
Equality check for VariableNodeData.
"""
function ==(a::VariableNodeData,b::VariableNodeData, nt::Symbol=:var)
  return DistributedFactorGraphs.compare(a,b)
end

"""
    ==(x::T, y::T) where T <: AbstractPointParametricEst
Equality check for AbstractPointParametricEst.
"""
@generated function ==(x::T, y::T) where T <: AbstractPointParametricEst
    mapreduce(n -> :(x.$n == y.$n), (a,b)->:($a && $b), fieldnames(x))
end

@generated function Base.:(==)(x::T, y::T) where T <: Union{DFGFactorSummary, DFGVariableSummary, SkeletonDFGVariable, SkeletonDFGFactor}
    mapreduce(n -> :(x.$n == y.$n), (a,b)->:($a && $b), fieldnames(x))
end

"""
    $(SIGNATURES)
Equality check for DFGVariable.
"""
function ==(a::DFGVariable, b::DFGVariable)::Bool
    return compareVariable(a, b)
end

"""
    $(SIGNATURES)
Convert a DFGVariable to a DFGVariableSummary.
"""
function convert(::Type{DFGVariableSummary}, v::DFGVariable)
    return DFGVariableSummary(v.label, v.timestamp, deepcopy(v.tags), deepcopy(v.estimateDict), Symbol(typeof(getSofttype(v))), v._internalId)
end


""" Lay out and draw wiring diagrams using Graphviz.
"""
module GraphvizWiringDiagrams
export to_graphviz, graphviz_layout

import JSON
using LightGraphs, MetaGraphs
using LinearAlgebra: normalize
using StaticArrays

import ...Theories: HomExpr
using ...WiringDiagrams, ...WiringDiagrams.WiringDiagramSerialization
import ..Graphviz
import ..Graphviz: to_graphviz
using ..Graphviz: parse_graphviz, run_graphviz
using ..WiringDiagramLayouts: BoxLayout, PortLayout, WirePoint,
  LayoutOrientation, LeftToRight, RightToLeft, TopToBottom, BottomToTop,
  is_horizontal, is_vertical, box_label, wire_label, port_sign, svector

# Drawing
#########

# Default Graphviz font. Reference: http://www.graphviz.org/doc/fontfaq.txt
const default_font = "Serif"

# Default graph, node, edge, and cell attributes.
const default_graph_attrs = Graphviz.Attributes(
  :fontname => default_font,
)
const default_node_attrs = Graphviz.Attributes(
  :fontname => default_font,
  :shape => "none",
  :width => "0",
  :height => "0",
  :margin => "0",
)
const default_edge_attrs = Graphviz.Attributes(
  :arrowsize => "0.5",
  :fontname => default_font,
)
const default_cell_attrs = Graphviz.Attributes(
  :border => "1",
  :cellpadding => "4",
)

struct GraphvizBox
  stmts::Vector{Graphviz.Statement} # Usually Graphviz.Node
  input_ports::Vector{Graphviz.NodeID}
  output_ports::Vector{Graphviz.NodeID}
end

""" Draw a wiring diagram using Graphviz.

The input `f` can also be a morphism expression, which is converted into a
wiring diagram.

# Arguments
- `graph_name="G"`: name of Graphviz digraph
- `orientation=TopToBottom`: orientation of layout.
  One of `LeftToRight`, `RightToLeft`, `TopToBottom`, or `BottomToTop`.
- `node_labels=true`: whether to label the nodes
- `labels=false`: whether to label the edges
- `label_attr=:label`: what kind of edge label to use (if `labels` is true).
  One of `:label`, `:xlabel`, `:headlabel`, or `:taillabel`.
- `port_size="24"`: minimum size of ports on box, in points
- `junction_size="0.05"`: size of junction nodes, in inches
- `outer_ports=true`: whether to display the outer box's input and output ports.
  If disabled, no incoming or outgoing wires will be shown either!
- `anchor_outer_ports=true`: whether to enforce ordering of the outer box's
  input and output, i.e., ordering of the incoming and outgoing wires
- `graph_attrs=default_graph_attrs`: top-level graph attributes
- `node_attrs=default_node_attrs`: top-level node attributes
- `edge_attrs=default_edge_attrs`: top-level edge attributes
- `cell_attrs=default_cell_attrs`: main cell attributes in node HTML-like label
"""
function to_graphviz(f::WiringDiagram;
    graph_name::String="G", orientation::LayoutOrientation=TopToBottom,
    node_labels::Bool=true, labels::Bool=false, label_attr::Symbol=:label,
    port_size::String="24", junction_size::String="0.05",
    outer_ports::Bool=true, anchor_outer_ports::Bool=true,
    graph_attrs::AbstractDict=Dict(), node_attrs::AbstractDict=Dict(),
    edge_attrs::AbstractDict=Dict(), cell_attrs::AbstractDict=Dict())::Graphviz.Graph
  @assert label_attr in (:label, :xlabel, :headlabel, :taillabel)

  # State variables.
  stmts = Graphviz.Statement[]
  port_map = Dict{Port,Graphviz.NodeID}()
  update_port_map! = (v::Int, kind::PortKind, node_ids) -> begin
    for (i, node_id) in enumerate(node_ids)
      port_map[Port(v,kind,i)] = node_id
    end
  end

  # Invisible nodes for incoming and outgoing wires.
  if outer_ports
    gv_box = graphviz_outer_box(f;
      anchor=anchor_outer_ports, orientation=orientation)
    append!(stmts, gv_box.stmts)
    update_port_map!(input_id(f), OutputPort, gv_box.input_ports)
    update_port_map!(output_id(f), InputPort, gv_box.output_ports)
  end

  # Visible nodes for boxes.
  cell_attrs = merge(default_cell_attrs, Graphviz.as_attributes(cell_attrs))
  for v in box_ids(f)
    gv_box = graphviz_box(box(f,v), box_id([v]),
      orientation=orientation, labels=node_labels, port_size=port_size,
      junction_size=junction_size, cell_attrs=cell_attrs)
    append!(stmts, gv_box.stmts)
    update_port_map!(v, InputPort, gv_box.input_ports)
    update_port_map!(v, OutputPort, gv_box.output_ports)
  end

  # Edges.
  for (i, wire) in enumerate(wires(f))
    source, target = wire.source, wire.target
    if !(haskey(port_map, source) && haskey(port_map, target))
      continue
    end
    # Use the port value to label the wire. We take the source port.
    # In most wiring diagrams, the source and target ports should yield the
    # same label, but that is not guaranteed. An advantage of choosing the
    # source port over the target port is that it will carry the
    # "more specific" type when implicit conversions are allowed.
    port = port_value(f, source)
    attrs = Graphviz.Attributes(
      :id => wire_id(Int[], i),
      :comment => edge_label(port),
    )
    if labels
      attrs[label_attr] = edge_label(port)
    end
    edge = Graphviz.Edge(port_map[source], port_map[target]; attrs...)
    push!(stmts, edge)
  end

  # Graph.
  Graphviz.Digraph(graph_name, stmts;
    graph_attrs=merge(default_graph_attrs, Graphviz.as_attributes(graph_attrs),
      Graphviz.Attributes(
        :rankdir => rank_dir(orientation)
      )),
    node_attrs=merge(default_node_attrs, Graphviz.as_attributes(node_attrs)),
    edge_attrs=merge(default_edge_attrs, Graphviz.as_attributes(edge_attrs)))
end

function to_graphviz(f::HomExpr; kw...)::Graphviz.Graph
  to_graphviz(to_wiring_diagram(f); kw...)
end

""" Graphviz box for a generic box.
"""
function graphviz_box(box::AbstractBox, node_id;
    orientation::LayoutOrientation=TopToBottom,
    labels::Bool=true, port_size::String="0",
    cell_attrs::AbstractDict=Dict(), kw...)
  # Main node.
  nin, nout = length(input_ports(box)), length(output_ports(box))
  text_label = labels ? node_label(box.value) : ""
  html_label = node_html_label(nin, nout, text_label;
    orientation=orientation, port_size=port_size, attrs=cell_attrs)
  # Note: The `id` attribute is included in the Graphviz output but is not used
  # internally by Graphviz. It is for use by downstream applications.
  # Reference: http://www.graphviz.org/doc/info/attrs.html#d:id
  node = Graphviz.Node(node_id,
    id = node_id,
    comment = node_label(box.value),
    label = html_label,
  )

  # Input and output ports.
  graphviz_port = (kind::PortKind, port::Int) -> begin
    Graphviz.NodeID(node_id, port_name(kind, port), port_anchor(kind, orientation))
  end
  inputs = [ graphviz_port(InputPort, i) for i in 1:nin ]
  outputs = [ graphviz_port(OutputPort, i) for i in 1:nout ]

  GraphvizBox([node], inputs, outputs)
end

""" Graphviz box for a junction.
"""
function graphviz_box(junction::Junction, node_id;
    junction_size::String="0", kw...)
  graphviz_junction(junction, node_id;
    comment = "junction",
    label = "",
    shape = "circle",
    style = "filled",
    fillcolor = "black",
    width = junction_size,
    height = junction_size,
  )
end

function graphviz_junction(junction::Junction, node_id; kw...)
  node = Graphviz.Node(node_id; id=node_id, kw...)
  nin, nout = length(input_ports(junction)), length(output_ports(junction))
  GraphvizBox([node],
    repeat([Graphviz.NodeID(node_id)], nin),
    repeat([Graphviz.NodeID(node_id)], nout))
end

""" Create "HTML-like" node label for a box.
"""
function node_html_label(nin::Int, nout::Int, text_label::String;
    orientation::LayoutOrientation=TopToBottom, port_size::String="0",
    attrs::AbstractDict=Dict())::Graphviz.Html
  if is_vertical(orientation)
    input_label = ports_horizontal_html_label(InputPort, nin, port_size)
    output_label = ports_horizontal_html_label(OutputPort, nout, port_size)
    Graphviz.Html("""

$(orientation == TopToBottom ? input_label : output_label)

$(escape_html(text_label))

$(orientation == BottomToTop ? input_label : output_label)

""")
  else
    input_label = ports_vertical_html_label(InputPort, nin, port_size)
    output_label = ports_vertical_html_label(OutputPort, nout, port_size)
    Graphviz.Html("""

      $(orientation == LeftToRight ? input_label : output_label)
      $(escape_html(text_label))
      $(orientation == RightToLeft ? input_label : output_label)

""")
  end
end

""" Create horizontal "HTML-like" label for the input or output ports of a box.
"""
function ports_horizontal_html_label(kind::PortKind, nports::Int,
    port_size::String="0")::Graphviz.Html
  cols = if nports > 0
    join("""	"""
         for i in 1:nports)
  else
    """	"""
  end
  Graphviz.Html("""

$cols

""")
end

""" Create vertical "HTML-like" label for the input or output ports of a box.
"""
function ports_vertical_html_label(kind::PortKind, nports::Int,
    port_size::String="0")::Graphviz.Html
  rows = if nports > 0
    join("""

"""
         for i in 1:nports)
  else
    """

"""
  end
  Graphviz.Html("""

$rows
""")
end

""" Graphviz box for the outer box of a wiring diagram.
"""
function graphviz_outer_box(f::WiringDiagram;
    anchor::Bool=true, orientation::LayoutOrientation=TopToBottom)
  # Subgraphs containing invisible nodes.
  stmts = Graphviz.Statement[]
  ninputs, noutputs = length(input_ports(f)), length(output_ports(f))
  if ninputs > 0
    push!(stmts, graphviz_outer_ports(input_id(f), InputPort, ninputs;
      anchor=anchor, orientation=orientation))
  end
  if noutputs > 0
    push!(stmts, graphviz_outer_ports(output_id(f), OutputPort, noutputs;
      anchor=anchor, orientation=orientation))
  end

  # Input and output ports.
  graphviz_port = (port::Port) -> Graphviz.NodeID(
    port_node_name(port.box, port.port),
    port_anchor(port.kind, orientation)
  )
  inputs = [ graphviz_port(Port(input_id(f), OutputPort, i)) for i in 1:ninputs ]
  outputs = [ graphviz_port(Port(output_id(f), InputPort, i)) for i in 1:noutputs ]

  GraphvizBox(stmts, inputs, outputs)
end

""" Create invisible nodes for the input or output ports of an outer box.
"""
function graphviz_outer_ports(v::Int, kind::PortKind, nports::Int;
    anchor::Bool=true, orientation::LayoutOrientation=TopToBottom)::Graphviz.Subgraph
  @assert nports > 0
  port_width = "$(round(24/72,digits=3))" # port width in inches
  nodes = [ port_node_name(v, i) for i in 1:nports ]
  stmts = Graphviz.Statement[
    Graphviz.Node(nodes[i], id=port_name(kind, i)) for i in 1:nports
  ]
  if anchor
    push!(stmts, Graphviz.Edge(nodes))
  end
  Graphviz.Subgraph(
    stmts,
    graph_attrs=Graphviz.Attributes(
      :rank => kind == InputPort ? "source" : "sink",
      :rankdir => is_vertical(orientation) ? "LR" : "TB"
    ),
    node_attrs=Graphviz.Attributes(
      :style => "invis",
      :shape => "none",
      :label => "",
      :width => is_vertical(orientation) ? port_width : "0",
      :height => is_horizontal(orientation) ? port_width : "0",
    ),
    edge_attrs=Graphviz.Attributes(
      :style => "invis",
    ),
  )
end
port_node_name(v::Int, port::Int) = string(box_id([v]), "p", port)

""" Graphviz rank direction (`rankdir`) for layout orientation.

Reference: https://www.graphviz.org/doc/info/attrs.html#k:rankdir
"""
rank_dir(orientation::LayoutOrientation) = rank_dirs[orientation]

const rank_dirs = Dict{LayoutOrientation,String}(
  TopToBottom => "TB",
  BottomToTop => "BT",
  LeftToRight => "LR",
  RightToLeft => "RL",
)

""" Graphviz anchor for port.
"""
function port_anchor(kind::PortKind, orientation::LayoutOrientation)
  port_anchors[(kind, orientation)]
end

const port_anchors = Dict{Tuple{PortKind,LayoutOrientation},String}(
  (InputPort, TopToBottom) => "n",
  (OutputPort, TopToBottom) => "s",
  (InputPort, BottomToTop) => "s",
  (OutputPort, BottomToTop) => "n",
  (InputPort, LeftToRight) => "w",
  (OutputPort, LeftToRight) => "e",
  (InputPort, RightToLeft) => "e",
  (OutputPort, RightToLeft) => "w",
)

""" Create a label for the main content of a box.
"""
node_label(box_value) = box_label(MIME("text/plain"), box_value)

""" Create a label for an edge.
"""
edge_label(port_value) = wire_label(MIME("text/plain"), port_value)

""" Encode attributes for Graphviz HTML-like labels.
"""
function html_attributes(attrs::AbstractDict)::String
  join([ "$(uppercase(string(k)))=\"$v\""
         for (k,v) in Graphviz.as_attributes(attrs) ], " ")
end

""" Escape special HTML characters: &, <, >, ", '

Adapted from [HttpCommon.jl](https://github.com/JuliaWeb/HttpCommon.jl).
"""
function escape_html(s::AbstractString)
  reduce(replace, [
    "&" => "&",
    "\"" => """,
    "'" => "'",
    "<" => "<",
    ">" => ">",
  ], init=s)
end

# Layout
########

""" Lay out wiring diagram using Graphviz.

Note: At this time, only the positions and sizes of the boxes, and the positions
of the outer ports, are used. The positions of the box ports and the splines for
the wires are ignored.
"""
function graphviz_layout(diagram::WiringDiagram; kw...)
  graph = to_graphviz(diagram; kw...)
  doc = JSON.parse(run_graphviz(graph, format="json0"))
  graphviz_layout(diagram, parse_graphviz(doc, multigraph=true))
end

function graphviz_layout(diagram::WiringDiagram, graph::MetaDiGraph)
  # Graphviz uses the standard Cartesian coordinate system (with origin in
  # bottom left corner), while our layout system uses a coordinate system with
  # centered origin and positive y-axis pointing downwards.
  orientation = inverse_rank_dirs[get_prop(graph, :rankdir)]
  bounds, pad = get_prop(graph, :bounds), get_prop(graph, :pad)
  diagram_size = bounds + 2*pad
  transform_point(p) = SVector(1,-1) .* (p - diagram_size/2)

  # Assumes vertices are in same order as created by `to_graphviz`:
  # 1. Input ports of outer box
  nin, nout = length(input_ports(diagram)), length(output_ports(diagram))
  main_dir = svector(orientation, 1.0, 0.0)
  inputs = map(enumerate(input_ports(diagram))) do (i, value)
    attrs = props(graph, i)
    pos = transform_point(attrs[:position])
    PortLayout(; value=value, position=pos, normal=-main_dir)
  end

  # 2. Output ports of outer box
  outputs = map(enumerate(output_ports(diagram))) do (i, value)
    attrs = props(graph, nin + i)
    pos = transform_point(attrs[:position])
    PortLayout(; value=value, position=pos, normal=main_dir)
  end

  # 3. Inner boxes
  # TODO: Use port positions from Graphviz layout. Obtain these from either
  # the HTML label or, more likely, an incident edge.
  layout = WiringDiagram(BoxLayout(size=diagram_size), inputs, outputs)
  node_map = Dict{Int,Int}()
  for (i, box) in enumerate(boxes(diagram))
    node = nin + nout + i
    attrs = props(graph, node)
    shape = Symbol(get(attrs, :shape, :ellipse))
    if shape == :none; shape = :rectangle end # XXX: Assume HTML label.
    @assert shape == :rectangle "Only the rectangle shape is implemented so far"
    pos, size = transform_point(attrs[:position]), attrs[:size]
    node_map[node] = add_box!(layout, Box(
      BoxLayout(; value=box.value, shape=shape, position=pos, size=size),
      layout_linear_ports(InputPort, input_ports(box), size, orientation),
      layout_linear_ports(OutputPort, output_ports(box), size, orientation)
    ))
  end

  # Add wires using spline points from Graphviz edge layout.
  function map_port(node::Int, portname::Union{String,Nothing})
    if node <= nin
      Port(input_id(layout), OutputPort, node)
    elseif node <= nin + nout
      Port(output_id(layout), InputPort, node - nin)
    else
      # XXX: Parsing the port names is hacky. The alternative is parsing the
      # ordered ports in the HTML node label and matching the port names.
      kind = startswith(portname, "out") ? OutputPort : InputPort
      port = parse(Int, portname[findfirst(r"[0-9]+", portname)])
      Port(node_map[node], kind, port)
    end
  end
  for edge in edges(graph)
    for attrs in get_prop(graph, edge, :edges)
      # Chop off start and end points to obtain intermediate points.
      points = attrs[:spline][3:end-2]
      wire_layout = map(Iterators.partition(points, 3)) do (c1, p, c2)
        WirePoint(transform_point(p),
                  normalize(transform_point(c2) - transform_point(c1)))
      end
      add_wire!(layout, Wire(wire_layout,
        map_port(src(edge), get(attrs, :tailport, nothing)),
        map_port(dst(edge), get(attrs, :headport, nothing))
      ))
    end
  end
  layout
end

""" Lay out ports linearly, equispaced along a rectangular box.

FIXME: Should this be in `WiringDiagramLayouts` as an alternative layout method?
"""
function layout_linear_ports(port_kind::PortKind, port_values::Vector,
                             box_size::StaticVector{2}, orientation::LayoutOrientation)
  n = length(port_values)
  sign = port_sign(port_kind, orientation)
  normal = svector(orientation, sign, 0.0)
  map(port_values, range(-1,1,length=n+2)[2:end-1]) do value, coeff
    pos = svector(orientation, sign, coeff) .* box_size/2
    PortLayout(; value=value, position=pos, normal=normal)
  end
end

const inverse_rank_dirs = Dict{String,LayoutOrientation}(
  "TB" => TopToBottom,
  "BT" => BottomToTop,
  "LR" => LeftToRight,
  "RL" => RightToLeft,
)

end


include("vi_optimizers.jl")
include("analytic.jl")
include("analyticVI.jl")
include("numericalVI.jl")
include("sampling.jl")
include("optimisers.jl")

export RobbinsMonro, ALRSVI

post_process!(::AbstractGP) = nothing

# Utils to iterate over inference objects
Base.length(::AbstractInference) = 1

Base.iterate(i::AbstractInference) = (i, nothing)
Base.iterate(::AbstractInference, ::Any) = nothing

isStochastic(i::AbstractInference) = i.stoch

## Multiple accessors
conv_crit(i::AbstractInference) = i.ϵ

## Conversion from natural to standard distribution parameters ##
function global_update!(gp::AbstractLatent)
    gp.post.Σ .= -0.5 * inv(nat2(gp))
    return gp.post.μ .= cov(gp) * nat1(gp)
end

## For the online case, the size may vary and inplace updates are note valid
function global_update!(gp::OnlineVarLatent)
    gp.post.Σ = -0.5 * inv(nat2(gp))
    return gp.post.μ = cov(gp) * nat1(gp)
end

## Default function for getting a view on y
xview(inf::AbstractInference) = inf.xview
setxview!(inf::AbstractInference, xview) = inf.xview = xview

setyview!(inf::AbstractInference, yview) = inf.yview = yview
yview(inf::AbstractInference) = inf.yview

nMinibatch(inf::AbstractInference) = inf.nMinibatch
setnMinibatch!(inf::AbstractInference, n::Int) = inf.nMinibatch = n

nSamples(i::AbstractInference) = i.nSamples
setnSamples!(inf::AbstractInference, n::Int) = inf.nSamples = n

getρ(inf::AbstractInference) = inf.ρ
setρ!(inf::AbstractInference, ρ) = inf.ρ = ρ

MBIndices(inf::AbstractInference) = inf.MBIndices
function setMBIndices!(inf::AbstractInference, mbindices::AbstractVector)
    return inf.MBIndices .= mbindices
end

setHPupdated!(inf::AbstractInference, status::Bool) = inf.HyperParametersUpdated = status
isHPupdated(inf::AbstractInference) = inf.HyperParametersUpdated

nIter(inf::AbstractInference) = inf.nIter

get_opt(::AbstractInference) = nothing
get_opt(i::VariationalInference) = i.vi_opt
get_opt(i::SamplingInference) = i.opt
get_opt(i::AbstractInference, n::Int) = get_opt(i)[n]
opt_type(i::AbstractInference) = first(get_opt(i))

# Initialize the final version of the inference objects
# using the right parametrization and size
function tuple_inference(
    i::AbstractInference,
    nLatent::Int,
    nFeatures::Int,
    nSamples::Int,
    nMinibatch::Int,
    xview::AbstractVector,
    yview::AbstractVector,
)
    return tuple_inference(
        i, nLatent, fill(nFeatures, nLatent), nSamples, nMinibatch, xview, yview
    )
end


#!/usr/bin/env julia

module QJuliaSUNtest

#load path to qjulia home directory
push!(LOAD_PATH, joinpath(@__DIR__, "../core"))
push!(LOAD_PATH, joinpath(@__DIR__, "../main/fields"))
push!(LOAD_PATH, @__DIR__)

import QJuliaRegisters
import QJuliaEnums
import QJuliaGrid
import QJuliaFields
import QJuliaCUDAFields

using CUDAdrv
using CUDAnative
using CuArrays
using BenchmarkTools
using StaticArrays

#Float2 version
@inline function axpy_(a, x::Complex{T}, y::Complex{T})::Complex{T} where T <: AbstractFloat

   rres = CUDAnative.fma(a, real(x), real(y))
   ires = CUDAnative.fma(a, imag(x), imag(y))

   return Complex{T}(rres, ires)
 end

 @inline function caxpy_(a::Complex{T}, x::Complex{T}, y::Complex{T})::Complex{T} where T <: AbstractFloat

    rres = CUDAnative.fma(+real(a), real(x), real(y))
    rres = CUDAnative.fma(-imag(a), imag(x), rres   )

    ires = CUDAnative.fma(+imag(a), real(x), imag(y))
    ires = CUDAnative.fma(+real(a), imag(x), ires   )

    return Complex{T}(rres, ires)
  end

 @inline function axpy_(a::T, x::NTuple{2, VecElement{T}}, y::NTuple{2, VecElement{T}})::NTuple{2, VecElement{T}} where T <: AbstractFloat

   res1 = CUDAnative.fma(a, x[1].value, y[1].value)
   res2 = CUDAnative.fma(a, x[2].value, y[2].value)

   return NTuple{2, VecElement{T}}((res1, res2))
 end

 @inline function caxpy_(a::Complex{T}, x::NTuple{2, VecElement{T}}, y::NTuple{2, VecElement{T}})::NTuple{2, VecElement{T}} where T <: AbstractFloat

   rres = CUDAnative.fma(+real(a), x[1].value, y[1].value)
   rres = CUDAnative.fma(-imag(a), x[2].value, rres      )

   ires = CUDAnative.fma(+real(a), x[2].value, y[2].value)
   ires = CUDAnative.fma(+imag(a), x[1].value, ires      )

   return NTuple{2, VecElement{T}}((res1, res2))
 end

 @inline function axpy_(a::NTuple{2, VecElement{T}}, x::NTuple{2, VecElement{T}}, y::NTuple{2, VecElement{T}})::NTuple{2, VecElement{T}} where T <: AbstractFloat
   return (VecElement(CUDAnative.fma(a[1].value, x[1].value, y[1].value)), VecElement(CUDAnative.fma(a[2].value, x[2].value, y[2].value)))
   #return (VecElement(a[1].value*x[1].value + y[1].value), VecElement(a[2].value*x[1].value + y[2].value))
 end

 @inline function caxpy_(a::NTuple{2, VecElement{T}}, x::NTuple{2, VecElement{T}}, y::NTuple{2, VecElement{T}})::NTuple{2, VecElement{T}} where T <: AbstractFloat
   res  = (VecElement(CUDAnative.fma(+a[1].value, x[1].value, y[1].value)),
           VecElement(CUDAnative.fma(+a[1].value, x[2].value, y[2].value)))
   return (VecElement(CUDAnative.fma(-a[2].value, x[2].value, res[1].value)),
           VecElement(CUDAnative.fma(+a[2].value, x[1].value, res[2].value)))
 end

 @inline function caxpy_(a::NTuple{4, VecElement{T}}, x::NTuple{4, VecElement{T}}, y::NTuple{4, VecElement{T}})::NTuple{4, VecElement{T}} where T <: AbstractFloat
   res  = (VecElement(CUDAnative.fma(+a[1].value, x[1].value, y[1].value)),
           VecElement(CUDAnative.fma(+a[1].value, x[2].value, y[2].value)),
           VecElement(CUDAnative.fma(+a[3].value, x[3].value, y[3].value)),
           VecElement(CUDAnative.fma(+a[3].value, x[4].value, y[4].value)))
   return (VecElement(CUDAnative.fma(-a[2].value, x[2].value, res[1].value)),
           VecElement(CUDAnative.fma(+a[2].value, x[1].value, res[2].value)),
           VecElement(CUDAnative.fma(-a[4].value, x[4].value, res[3].value)),
           VecElement(CUDAnative.fma(+a[4].value, x[3].value, res[4].value)))
 end

 # a few helper methods:
 @inline function get_regN(meta::NTuple{N, VecElement{T}}) where T <: AbstractFloat where N; return N; end
 @inline function get_regT(meta::NTuple{N, VecElement{T}}) where T <: AbstractFloat where N; return T; end

function axpy(a, xaccessor, yaccessor)
    xcb = (blockIdx().x-1) * blockDim().x + threadIdx().x
    # get data arrays from the accessors
    xdata = xaccessor[1]; ydata = yaccessor[1]
    args  = xaccessor[2]
    #
    if xcb > args.length; return; end

#    if (Int64(blockIdx().x) == 1 && threadIdx().x == 1)
#      for s in 1:args.nSpin
#        for c in 1:args.nColor
#          @cuprintf("color %d  ", c)
#          @cuprintf("spin  %d\n", s)
#        end
#      end
#      @cuprintf("Info block %ld, thread %ld : %d  (length), %d \n", Int64(blockIdx().x), Int64(threadIdx().x), args.length, args.length)
#      @cuprintf("Info block %ld, thread %ld : %d  (colors) \n", Int64(blockIdx().x), Int64(threadIdx().x), args.nColor)
#    end

    y_site_view = view(ydata,xcb,:,:,1)
    x_site_view = view(xdata,xcb,:,:,1)

    for s in 1:args.nSpin
      for c in 1:args.nColor
        y_site_view[c, s] = axpy_(a, x_site_view[c, s], y_site_view[c, s])
      end
    end

  #  if (Int64(blockIdx().x) == 1 && threadIdx().x == 4)
  #    for s in 1:args.nSpin
  #      for c in 1:args.nColor
  #        @cuprintf("Info block %ld, thread %ld : %le, %le!\n", Int64(blockIdx().x), Int64(threadIdx().x), x_site_view[c,s][1].value, x_site_view[c,s][2].value)
  #      end
  #    end
  #  end
    return
end

function set_value(x, a)
    xcb = (blockIdx().x-1) * blockDim().x + threadIdx().x

    (xlength, xcolors, xspins, xvecs) = size(x)

    if xcb > xlength; return; end

    xview = view(x,:,:,:,1)

    for s in 1:xspins
      for c in 1:xcolors
        xview[xcb, c, s] =  NTuple{2, Float32}((Float32(s), Float32(c)))
      end
    end

    return
end

function sun_test(x, y)
    xcb = (blockIdx().x-1) * blockDim().x + threadIdx().x
    #    wid, lane = fldmod1(threadIdx().x, CUDAnative.warpsize())
    lid = Int64(threadIdx().x)

    (nElems, nCols, nRows, nDirs) = size(x)

    if xcb > nElems; return; end

    dir = 1

    y_site_view = view(y,xcb,:,:,dir)
    x_site_view = view(x,xcb,:,:,dir)

    #shared_dims = blockDim().x*nCols*nRows*get_regN(x_site_view[1,1])
    #tmp = @cuDynamicSharedMem(get_regT(x_site_view[1,1]), shared_dims)
    msize = nCols*nRows*4 #get_regN(x_site_view[1,1])

    shared_dims = blockDim().x*msize
    tmp = @cuDynamicSharedMem(Float32, shared_dims)

    #if (Int64(blockIdx().x) == 1)
    #  @cuprintf("Info block %ld, thread %ld : %d  (wid), %d (lane), %d (warp size) \n", Int64(blockIdx().x), Int64(threadIdx().x), wid, lane, CUDAnative.warpsize())
    #end

    for r in 1:nRows
      for c in 1:nCols
        temp = NTuple{4, VecElement{Float32}}((0.0f0, 0.0f0, 0.0f0, 0.0f0))
        for j in 1:nCols
          temp = caxpy_(x_site_view[r, j], y_site_view[j, c], temp)
        end
        @inbounds tmp[msize*(lid-1)+(r+c*3)*2+0] = temp[1].value
        @inbounds tmp[msize*(lid-1)+(r+c*3)*2+1] = temp[2].value
        @inbounds tmp[msize*(lid-1)+(r+c*3)*2+2] = temp[3].value
        @inbounds tmp[msize*(lid-1)+(r+c*3)*2+3] = temp[4].value
      end
    end

    for r in 1:nRows
      for c in 1:nCols
        temp = NTuple{4, VecElement{Float32}}((tmp[msize*(lid-1)+(r+c*3)*2+0], tmp[msize*(lid-1)+(r+c*3)*2+1], tmp[msize*(lid-1)+(r+c*3)*2+2], tmp[msize*(lid-1)+(r+c*3)*2+3]))
        y_site_view[r, c] = temp
      end
    end

#    if (Int64(blockIdx().x) == 1 && threadIdx().x == 4)
#      for c in 1:xrows
#        for r in 1:xcols
#          @cuprintf("Info block %ld, thread %ld : %le, %le!\n", Int64(blockIdx().x), Int64(threadIdx().x), y_site_view[r,c][1].value, y_site_view[r,c][2].value)
#        end
#      end
#    end

    return
end

function sun_test_pure_arg_templ(x::CuDeviceArray{NTuple{N,VecElement{T}},4}, y::CuDeviceArray{NTuple{N,VecElement{T}},4}) where T <: AbstractFloat where N
    xcb = (blockIdx().x-1) * blockDim().x + threadIdx().x
    #    wid, lane = fldmod1(threadIdx().x, CUDAnative.warpsize())
    tid  = Int64(threadIdx().x)
    bdim = Int64(blockDim().x)

    (nElems, nCols, nRows, nDirs) = size(x)

    if xcb > nElems; return; end

    dir = 1

    y_site_view = view(y,xcb,:,:,dir)
    x_site_view = view(x,xcb,:,:,dir)

    #shared_dims = blockDim().x*nCols*nRows*get_regN(x_site_view[1,1])
    #tmp = @cuDynamicSharedMem(get_regT(x_site_view[1,1]), shared_dims)

    shared_dims = bdim*nCols*nRows*N
    tmp = @cuDynamicSharedMem(T, shared_dims)

    #if (Int64(blockIdx().x) == 1)
    #  @cuprintf("Info block %ld, thread %ld : %d  (wid), %d (lane), %d (warp size) \n", Int64(blockIdx().x), Int64(threadIdx().x), wid, lane, CUDAnative.warpsize())
    #end

    for r in 1:nRows
      for c in 1:nCols
        temp = NTuple{N, VecElement{T}}(ntuple(i->T(0.0), N))
        for j in 1:nCols
          temp = caxpy_(x_site_view[r, j], y_site_view[j, c], temp)
        end
        offset = tid + ((r-1)+(c-1)*3)*bdim
        for l in 1:N
          @inbounds tmp[offset+(l-1)*9*bdim] = temp[l].value
        end
      end
    end

    for r in 1:nRows
      for c in 1:nCols
        offset = tid + ((r-1)+(c-1)*3)*bdim
        temp = N == 2 ? NTuple{N, VecElement{T}}( (tmp[offset+0*9*bdim], tmp[offset+1*9*bdim])) : NTuple{N, VecElement{T}}( (tmp[offset+0*9*bdim], tmp[offset+1*9*bdim],tmp[offset+2*9*bdim], tmp[offset+3*9*bdim]) )
        @inbounds y_site_view[r, c] = temp
      end
    end

#    if (Int64(blockIdx().x) == 1 && threadIdx().x == 4)
#      for c in 1:xrows
#        for r in 1:xcols
#          @cuprintf("Info block %ld, thread %ld : %le, %le!\n", Int64(blockIdx().x), Int64(threadIdx().x), y_site_view[r,c][1].value, y_site_view[r,c][2].value)
#        end
#      end
#    end

    return
end

function sun_extended(z::CuDeviceArray{NTuple{N, VecElement{T}},4},  x::CuDeviceArray{NTuple{N,VecElement{T}},4}, y::CuDeviceArray{NTuple{N,VecElement{T}},4}) where T <: AbstractFloat where N
    xcb = (blockIdx().x-1) * blockDim().x + threadIdx().x

    (nElems, nCols, nRows, nDirs) = size(x)

    if xcb > nElems; return; end

    dir = 1

    z_site_view = view(z,xcb,:,:,dir)
    y_site_view = view(y,xcb,:,:,dir)
    x_site_view = view(x,xcb,:,:,dir)

    #if (Int64(blockIdx().x) == 1)
    #  @cuprintf("Info block %ld, thread %ld : %d  (wid), %d (lane), %d (warp size) \n", Int64(blockIdx().x), Int64(threadIdx().x), wid, lane, CUDAnative.warpsize())
    #end

    for r in 1:nRows
      for c in 1:nCols
        temp = NTuple{N, VecElement{T}}(ntuple(i->T(0.0), N))
        for j in 1:nCols
          temp = caxpy_(x_site_view[r, j], y_site_view[j, c], temp)
        end
        z_site_view[r,c] = temp
      end
    end

    return
 end

function sun_test_templ(x::Tuple{CuDeviceArray{NTuple{N,VecElement{T}},4},QJuliaFields.AccessorArgs}, y::Tuple{CuDeviceArray{NTuple{N,VecElement{T}},4},QJuliaFields.AccessorArgs}) where T <: AbstractFloat where N
    xcb = (blockIdx().x-1) * blockDim().x + threadIdx().x
    #    wid, lane = fldmod1(threadIdx().x, CUDAnative.warpsize())
    tid  = Int64(threadIdx().x)
    bdim = Int64(blockDim().x)

    (nElems, nCols, nRows, nDirs) = size(x[1])

    if xcb > nElems; return; end

    dir = 1

    y_site_view = view(y[1],xcb,:,:,dir)
    x_site_view = view(x[1],xcb,:,:,dir)

    #shared_dims = blockDim().x*nCols*nRows*get_regN(x_site_view[1,1])
    #tmp = @cuDynamicSharedMem(get_regT(x_site_view[1,1]), shared_dims)

    shared_dims = bdim*nCols*nRows*N
    tmp = @cuDynamicSharedMem(T, shared_dims)

    #if (Int64(blockIdx().x) == 1)
    #  @cuprintf("Info block %ld, thread %ld : %d  (wid), %d (lane), %d (warp size) \n", Int64(blockIdx().x), Int64(threadIdx().x), wid, lane, CUDAnative.warpsize())
    #end

    for r in 1:nRows
      for c in 1:nCols
        temp = NTuple{N, VecElement{T}}(ntuple(i->T(0.0), N))
        for j in 1:nCols
          temp = caxpy_(x_site_view[r, j], y_site_view[j, c], temp)
        end
        offset = tid + ((r-1)+(c-1)*3)*bdim
        for l in 1:N
          @inbounds tmp[offset+(l-1)*9*bdim] = temp[l].value
        end
      end
    end

    for r in 1:nRows
      for c in 1:nCols
        offset = tid + ((r-1)+(c-1)*3)*bdim
        temp = N == 2 ? NTuple{N, VecElement{T}}( (tmp[offset+0*9*bdim], tmp[offset+1*9*bdim])) : NTuple{N, VecElement{T}}( (tmp[offset+0*9*bdim], tmp[offset+1*9*bdim],tmp[offset+2*9*bdim], tmp[offset+3*9*bdim]) )
        y_site_view[r, c] = temp
      end
    end

#    if (Int64(blockIdx().x) == 1 && threadIdx().x == 4)
#      for c in 1:xrows
#        for r in 1:xcols
#          @cuprintf("Info block %ld, thread %ld : %le, %le!\n", Int64(blockIdx().x), Int64(threadIdx().x), y_site_view[r,c][1].value, y_site_view[r,c][2].value)
#        end
#      end
#    end

    return
end

  const  nthreads = 128

  #data_type = ComplexF32
  #data_type = QJuliaRegisters.float2
  data_type = QJuliaRegisters.float4

  const N = 32
  println("Test lattice size ::" , N)

  csGrid   = QJuliaGrid.QJuliaGridDescr_qj{data_type}(QJuliaEnums.QJULIA_CUDA_FIELD_LOCATION, 0, (N,N,N,N,))

  # compilation time
  gaugeParam  = QJuliaGrid.CreateGaugeParams( csGrid )

  cuda_sun_m1 = QJuliaCUDAFields.CreateGenericField( gaugeParam )
  cuda_sun_m2 = QJuliaCUDAFields.CreateGenericField( gaugeParam )
  cuda_sun_m3 = QJuliaCUDAFields.CreateGenericField( gaugeParam )

  println("\n=====GAUGE FIELD INFO=====")
  QJuliaFields.field_info(cuda_sun_m1)

  println("====================")

  cuda_sun_m1_accessor = QJuliaFields.create_field_accessor(cuda_sun_m1)
  cuda_sun_m2_accessor = QJuliaFields.create_field_accessor(cuda_sun_m2)
  cuda_sun_m3_accessor = QJuliaFields.create_field_accessor(cuda_sun_m3)

  accessor_dims = size(cuda_sun_m1_accessor[1])

  println("  ::: ", typeof(cuda_sun_m1_accessor))

  len = accessor_dims[1]

  nblocks= ceil.(Int, len ./ nthreads)
  # calculate size of dynamic shared memory
  regN    = QJuliaRegisters.register_size(typeof(cuda_sun_m1_accessor[1][1]))
  regT    = QJuliaRegisters.register_type(typeof(cuda_sun_m1_accessor[1][1]))

  ncolors = accessor_dims[2]

  shmem_bytes = ncolors*ncolors*regN* nthreads * sizeof(regT)

  println("Execution config for length ", len, " : ", nthreads , " threads, ", nblocks, " blocks , shared mem ", shmem_bytes )

  # compilation options
#code inspections:
#@device_code_typed
#@device_code_llvm
#@device_code_ptx
#@device_code_sass
  @device_code_ptx @cuda blocks=nblocks threads=nthreads shmem=shmem_bytes sun_test_pure_arg_templ(cuda_sun_m1_accessor[1], cuda_sun_m2_accessor[1])
  #@time @cuda blocks=nblocks threads=nthreads shmem=shmem_bytes sun_test_templ(cuda_sun_m1_accessor, cuda_sun_m2_accessor)
  #@device_code_ptx @cuda blocks=nblocks threads=nthreads shmem=shmem_bytes sun_extended(cuda_sun_m3_accessor[1],  cuda_sun_m1_accessor[1], cuda_sun_m2_accessor[1])

  # execution time
  niters = 10000
  secs = CUDAdrv.@elapsed begin # begin profile
  for i in 1:niters
    @cuda blocks=nblocks threads=nthreads shmem=shmem_bytes sun_test_pure_arg_templ(cuda_sun_m1_accessor[1], cuda_sun_m2_accessor[1])
    #@cuda blocks=nblocks threads=nthreads shmem=shmem_bytes sun_test_templ(cuda_sun_m1_accessor, cuda_sun_m2_accessor)
    #@cuda blocks=nblocks threads=nthreads shmem=shmem_bytes sun_extended(cuda_sun_m3_accessor[1], cuda_sun_m1_accessor[1], cuda_sun_m2_accessor[1])
  end
  end # end profile

  vol  = regN*(len / 2.0) #complex array length
  ncol = 3
  bytes_per_link_dir = 3.0*vol*length(cuda_sun_m1.v) / 4.0
  footprint_per_link_dir = 2.0*vol*length(cuda_sun_m1.v) / 4.0
  flops = (8.0*ncol - 2.0)*ncol*ncol*vol

  secs_per_iter = (secs / Float64(niters))

  println("Bytes :: ", bytes_per_link_dir, " footprint :: ", footprint_per_link_dir, " flops ",  flops / secs_per_iter)

end # module


module TroposphereWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const Class = @jimport java.lang.Class
const DataProvidersManager = @jimport org.orekit.data.DataProvidersManager
const DateTimeComponents = @jimport org.orekit.time.DateTimeComponents
const DiscreteTroposphericModel = @jimport org.orekit.models.earth.troposphere.DiscreteTroposphericModel
const Enum = @jimport java.lang.Enum
const EstimatedTroposphericModel = @jimport org.orekit.models.earth.troposphere.EstimatedTroposphericModel
const Field = @jimport org.hipparchus.Field
const FieldAbsoluteDate = @jimport org.orekit.time.FieldAbsoluteDate
const FixedTroposphericDelay = @jimport org.orekit.models.earth.troposphere.FixedTroposphericDelay
const GlobalMappingFunctionModel = @jimport org.orekit.models.earth.troposphere.GlobalMappingFunctionModel
const InputStream = @jimport java.io.InputStream
const JString = @jimport java.lang.String
const List = @jimport java.util.List
const MappingFunction = @jimport org.orekit.models.earth.troposphere.MappingFunction
const MariniMurrayModel = @jimport org.orekit.models.earth.troposphere.MariniMurrayModel
const MendesPavlisModel = @jimport org.orekit.models.earth.troposphere.MendesPavlisModel
const NavigableSet = @jimport java.util.NavigableSet
const NiellMappingFunctionModel = @jimport org.orekit.models.earth.troposphere.NiellMappingFunctionModel
const Object = @jimport java.lang.Object
const Optional = @jimport java.util.Optional
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const SaastamoinenModel = @jimport org.orekit.models.earth.troposphere.SaastamoinenModel
const TimeScale = @jimport org.orekit.time.TimeScale
const TimeSpanEstimatedTroposphericModel = @jimport org.orekit.models.earth.troposphere.TimeSpanEstimatedTroposphericModel
const ViennaModelCoefficientsLoader = @jimport org.orekit.models.earth.troposphere.ViennaModelCoefficientsLoader
const ViennaModelType = @jimport org.orekit.models.earth.troposphere.ViennaModelType
const ViennaOneModel = @jimport org.orekit.models.earth.troposphere.ViennaOneModel
const ViennaThreeModel = @jimport org.orekit.models.earth.troposphere.ViennaThreeModel

export DiscreteTroposphericModel
export EstimatedTroposphericModel
export FixedTroposphericDelay
export GlobalMappingFunctionModel
export MappingFunction
export MariniMurrayModel
export MendesPavlisModel
export NiellMappingFunctionModel
export SaastamoinenModel
export TimeSpanEstimatedTroposphericModel
export ViennaModelCoefficientsLoader
export ViennaModelType
export ViennaOneModel
export ViennaThreeModel
export add_tropospheric_model_valid_after
export add_tropospheric_model_valid_before
export compare_to
export compute_zenith_delay
export describe_constable
export equals
export extract_parameters
export get_a
export get_class
export get_declaring_class
export get_default_model
export get_low_elevation_threshold
export get_parameters
export get_parameters_drivers
export get_standard_model
export get_supported_names
export get_transitions
export get_tropospheric_model
export get_zenith_delay
export hash_code
export load_data
export load_vienna_coefficients
export mapping_factors
export name
export notify
export notify_all
export ordinal
export path_delay
export set_low_elevation_threshold
export still_accepts_data
export to_string
export value_of
export values
export wait

include("discrete_tropospheric_model.jl")
include("estimated_tropospheric_model.jl")
include("fixed_tropospheric_delay.jl")
include("global_mapping_function_model.jl")
include("mapping_function.jl")
include("marini_murray_model.jl")
include("mendes_pavlis_model.jl")
include("niell_mapping_function_model.jl")
include("saastamoinen_model.jl")
include("time_span_estimated_tropospheric_model.jl")
include("vienna_model_coefficients_loader.jl")
include("vienna_model_type.jl")
include("vienna_one_model.jl")
include("vienna_three_model.jl")

end


using ArgParse
using EarthEngine
using EarthEngineREST
using ProgressMeter
using NCDatasets
using Dates
using ArgParse

Initialize()

function persist_stack(outpath, stack, lons, lats, dates)
   ds = NCDataset(outpath, "c", attrib=Dict(
        "title" => "Sentinel-1 Observation Time Series",
        "history" => "date_created: $(now(Dates.UTC))+00:00",
    ))

   defDim(ds,"lon",length(lons))
   defDim(ds,"lat",length(lats))
   defDim(ds,"time",length(dates))

   v = defVar(ds, "VV", Float32, ("lon", "lat", "time"), fillvalue = -9999.,attrib = Dict(
      "long_name" => "SAR VV Backscatter",
      "units" => "dB"
   ))
   v[:,:,:] = stack[:,end:-1:1,:]

   lon = defVar(ds, "lon", Float32,("lon",),
      attrib = Dict(
         "long_name" => "Longitude",
         "units" => "degrees_east",
         "standard_name" => "longitude"
      )
   )
   lon[:] = lons[:]

   lat = defVar(ds, "lat", Float32,("lat",),
      attrib = Dict(
         "long_name" => "Latitude",
         "units" => "degrees_north",
         "standard_name" => "latitude"
      )
   )
   lat[:] = lats[:]

   time = defVar(ds, "time", Int64, ("time",),
      attrib = Dict(
         "calendar" => "gregorian",
         "units" => "days since 1970-01-01"
      )
   )
   time[:] = dates[:]

   close(ds)
end

function input_data(box,secretkey,outfile;resolution::Float64=0.0015, basinlvl::Int=7, ascending::Bool=true)

   session = EESession(secretkey)
   Initialize(session)

   box = Rectangle(bbox...);
   basin_geo = geometry(box)

   basinmask = mask(reduceToImage(basins,["HYBAS_ID"],firstNonNull()))

   jrc = toFloat(unmask(select(EE.Image("JRC/GSW1_2/GlobalSurfaceWater"), "occurrence"),0)) / 100
   permanentwater = jrc < 0.8

   hand = select(EE.Image("MERIT/Hydro/v1_0_1"), "hnd")
   handmask = hand < 15

   path = ascending ? "ASCENDING" : "DESCENDING"

   scenes = filter(
      filterBounds(EE.ImageCollection("COPERNICUS/S1_GRD"), basin_geo),
      And(EE.Filter(),
          eq(EE.Filter(),"orbitProperties_pass", path),
          eq(EE.Filter(),"instrumentMode", "IW"),
      )
   )

   # aoi = geometry(EE.Image(get(toList(scenes,2),1)));
   aoi = basin_geo

   pixelgrid = PixelGrid(session, aoi, resolution, "EPSG:4326")

   eedates = distinct(
      map(
         aggregate_array(
            filterDate(scenes, "2015-01-01", "2021-01-01"),
            "system:time_start"
         ),
         @eefunc x -> format(EE.Date(x),"YYYY-MM-dd") EE.ComputedObject
      )
   )

   dates = computevalue(session,eedates)

   datesp = s1_date_pairs(dates)

   eedatesp = EE.List(datesp)

   function aggs1(date::EE.ComputedObject)
      t1 = EE.Date(date)
      t2 = advance(date, 1, "day")
      comp = filterDate(scenes,t1,t2)
      # return updateMask(mosaic(comp), And(permanentwater, handmask))
      return updateMask(mosaic(comp), And(permanentwater,basinmask))
   end

   # functions to process the sentinel series into something usable by fier
   # aggregate the data by to prevent swath seams
   function sq_agg(i)
       i = EE.List(i)
       x = set(mean(filterDate(scenes, get(i,0),get(i,1))),"system:time_start",millis(EE.Date(get(i,0))))
       return updateMask(float(x), And(permanentwater,basinmask))
   end

   imglist = map(eedatesp, @eefunc sq_agg EE.ComputedObject)

   # nimgs = computevalue(session, size(imglist))
   nimgs = computevalue(session,size(imglist))
   x_size = pixelgrid.dimensions.width
   y_size = pixelgrid.dimensions.height

   stack = Array{Union{Missing, Float64}, 3}(undef, x_size, y_size, nimgs);
   slice = Array{Union{Missing,Float64}, 3}(undef, x_size, y_size, 1);

   @showprogress for i in 1:nimgs
      img = unmask(get(imglist,i-1),999)
      slice[:,:,:] = computepixels(session, pixelgrid, img, ["VV",])[:,:,:]
      slice[findall(slice .>= 999)] .= missing
      stack[:,:,i] = slice[:,:,1]
   end

   lons,lats = EarthEngineREST.extract_gridcoordinates(pixelgrid)

   basedate = DateTime("1970-01-01")

   finaldates = map(x -> x[1], datesp)

   ncdates = @. Dates.value(Day(DateTime(finaldates) - basedate))

   persist_stack(outfile, stack, lons, lats, ncdates)

   return

end

# look at adjacent swaths within aoi and aggregate
function s1_date_pairs(d)
    d = sort(d)
    n = length(d)
    i = 1
    date_pairs = []
    while i < n
        j=1
        t1 = DateTime(d[i])
        t2 = DateTime(d[i+1])
        daydiff = Day(t2-t1)
        if (Dates.value(daydiff % 6) != 0) & (i+1<n)
            still_looking = true
            while still_looking
                t2_ = DateTime(d[i+j])
                daydiff_ = Day(t2_-t1)
                if (Dates.value(daydiff_ % 6) != 0)
                    j+=1
                else
                    still_looking = false
               end
            end
            push!(date_pairs,(d[i],d[i+j]))
            i = i + j
        elseif i+1 > n
            push!(date_pairs, (d[i],Dates.format(t2+Day(1),"YYYY-mm-dd")))
            i+=1
        else
            push!(date_pairs,(d[i],d[i+1]))
            i += 1
         end
   end

   return date_pairs
end

function parse_commandline()
    s = ArgParseSettings(
        description = "Forecasting of Inundation Extents using REOF (FIER) Command Line Application",
        add_version = true,
        version = @project_version
    )

    @add_arg_table s begin
        "--bbox", "-b"
        nargs = 4
        arg_type = Float64
        help = "bounding box to extract data for"
        required = true

        "--secretkey", "-k"
        arg_type = String
        help = "toml file specifying the configuration for command"
        required = true

        "--outfile", "-o"
        arg_type = String
        help = "output file path"
        required = true

    end

    return parse_args(s)
end

function main()
   args = parse_commandline()
   println(args)
   input_data(
      args["point"],
      args["secretkey"],
      args["outfile"],
   )
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


using JudiLing
using CSV
using Test
using DataFrames

@testset "path finding for latin" begin
    latin_train =
        DataFrame(CSV.File(joinpath("data", "latin_mini.csv")))
    cue_obj_train = JudiLing.make_cue_matrix(
        latin_train,
        grams = 3,
        target_col = :Word,
        tokenized = false,
        keep_sep = false,
    )

    latin_val = latin_train[101:150, :]
    cue_obj_val = JudiLing.make_cue_matrix(
        latin_val,
        cue_obj_train,
        grams = 3,
        target_col = :Word,
        tokenized = false,
        keep_sep = false,
    )

    n_features = size(cue_obj_train.C, 2)

    S_train, S_val = JudiLing.make_S_matrix(
        latin_train,
        latin_val,
        ["Lexeme"],
        ["Person", "Number", "Tense", "Voice", "Mood"],
        ncol = n_features,
    )

    G_train = JudiLing.make_transform_matrix(S_train, cue_obj_train.C)

    Chat_train = S_train * G_train
    Chat_val = S_val * G_train

    F_train = JudiLing.make_transform_matrix(cue_obj_train.C, S_train)

    Shat_train = cue_obj_train.C * F_train
    Shat_val = cue_obj_val.C * F_train

    A = cue_obj_train.A

    max_t = JudiLing.cal_max_timestep(latin_train, latin_val, :Word)

    res_train, gpi_train = JudiLing.learn_paths(
        latin_train,
        latin_train,
        cue_obj_train.C,
        S_train,
        F_train,
        Chat_train,
        A,
        cue_obj_train.i2f,
        cue_obj_train.f2i,
        gold_ind = cue_obj_train.gold_ind,
        Shat_val = Shat_train,
        check_gold_path = true,
        max_t = max_t,
        max_can = 10,
        grams = 3,
        threshold = 0.1,
        tokenized = false,
        sep_token = "_",
        keep_sep = false,
        target_col = :Word,
        issparse = :dense,
        verbose = false,
    )

    res_val, gpi_val = JudiLing.learn_paths(
        latin_train,
        latin_val,
        cue_obj_train.C,
        S_val,
        F_train,
        Chat_val,
        A,
        cue_obj_train.i2f,
        cue_obj_train.f2i,
        gold_ind = cue_obj_val.gold_ind,
        Shat_val = Shat_val,
        check_gold_path = true,
        max_t = max_t,
        max_can = 10,
        grams = 3,
        threshold = 0.1,
        is_tolerant = true,
        tolerance = 0.1,
        max_tolerance = 0,
        tokenized = false,
        sep_token = "-",
        keep_sep = false,
        target_col = :Word,
        issparse = :dense,
        verbose = false,
    )

    res_train = JudiLing.learn_paths(
        latin_train,
        latin_train,
        cue_obj_train.C,
        S_train,
        F_train,
        Chat_train,
        A,
        cue_obj_train.i2f,
        cue_obj_train.f2i,
        max_t = max_t,
        max_can = 10,
        grams = 3,
        threshold = 0.1,
        tokenized = false,
        sep_token = "_",
        keep_sep = false,
        target_col = :Word,
        issparse = :dense,
        verbose = false,
    )

    res_val = JudiLing.learn_paths(
        latin_train,
        latin_val,
        cue_obj_train.C,
        S_val,
        F_train,
        Chat_val,
        A,
        cue_obj_train.i2f,
        cue_obj_train.f2i,
        max_t = max_t,
        max_can = 10,
        grams = 3,
        threshold = 0.1,
        is_tolerant = true,
        tolerance = 0.1,
        max_tolerance = 0,
        tokenized = false,
        sep_token = "-",
        keep_sep = false,
        target_col = :Word,
        issparse = :dense,
        verbose = false,
    )

    res_train = JudiLing.learn_paths(
        latin_train,
        cue_obj_train,
        S_train,
        F_train,
        Chat_train,
        verbose = false,
    )

    res_train, ts_train = JudiLing.learn_paths(
        latin_train,
        latin_train,
        cue_obj_train.C,
        S_train,
        F_train,
        Chat_train,
        A,
        cue_obj_train.i2f,
        cue_obj_train.f2i,
        max_t = max_t,
        max_can = 10,
        grams = 3,
        threshold = 0.1,
        tokenized = false,
        sep_token = "_",
        keep_sep = false,
        target_col = :Word,
        issparse = :dense,
        check_threshold_stat = true,
        verbose = false,
    )

    res_val, ts_val = JudiLing.learn_paths(
        latin_train,
        latin_val,
        cue_obj_train.C,
        S_val,
        F_train,
        Chat_val,
        A,
        cue_obj_train.i2f,
        cue_obj_train.f2i,
        max_t = max_t,
        max_can = 10,
        grams = 3,
        threshold = 0.1,
        is_tolerant = true,
        tolerance = 0.1,
        max_tolerance = 0,
        tokenized = false,
        sep_token = "-",
        keep_sep = false,
        target_col = :Word,
        issparse = :dense,
        check_threshold_stat = true,
        verbose = false,
    )

    res_train, gpi_train, ts_train = JudiLing.learn_paths(
        latin_train,
        latin_train,
        cue_obj_train.C,
        S_train,
        F_train,
        Chat_train,
        A,
        cue_obj_train.i2f,
        cue_obj_train.f2i,
        gold_ind = cue_obj_train.gold_ind,
        Shat_val = Shat_train,
        check_gold_path = true,
        max_t = max_t,
        max_can = 10,
        grams = 3,
        threshold = 0.1,
        tokenized = false,
        sep_token = "_",
        keep_sep = false,
        target_col = :Word,
        issparse = :dense,
        check_threshold_stat = true,
        verbose = false,
    )

    res_val, gpi_val, ts_val = JudiLing.learn_paths(
        latin_train,
        latin_val,
        cue_obj_train.C,
        S_val,
        F_train,
        Chat_val,
        A,
        cue_obj_train.i2f,
        cue_obj_train.f2i,
        gold_ind = cue_obj_val.gold_ind,
        Shat_val = Shat_val,
        check_gold_path = true,
        max_t = max_t,
        max_can = 10,
        grams = 3,
        threshold = 0.1,
        is_tolerant = true,
        tolerance = 0.1,
        max_tolerance = 0,
        tokenized = false,
        sep_token = "-",
        keep_sep = false,
        target_col = :Word,
        issparse = :dense,
        check_threshold_stat = true,
        verbose = false,
    )
end


@testset "Trajectory" begin

function get_states_to_indices(game)
    state_index = StdMap{StdString, Cint}()
    to_visit = []
    push!(to_visit, new_initial_state(game))
    i = 0
    while length(to_visit) != 0
        state = pop!(to_visit)
        if (!is_chance_node(state)) && (!is_terminal(state))
            state_index[information_state_string(state)] = Cint(i)
        end
        i += 1
        for action in legal_actions(state)
            push!(to_visit, child(state, action))
        end
    end
    state_index
end

@testset "BatchedTrajectory" begin
    for game_name in ["kuhn_poker", "leduc_poker", "liars_dice"]
        game = load_game(game_name)
        batch_size = 32
        states_to_inds = get_states_to_indices(game)
        policies = StdVector([get_uniform_policy(game) for _ in 1:2])
        t = record_batched_trajectories(game, policies, states_to_inds, batch_size, false, 123, -1)
        @test length(legal_actions(t)) == batch_size
        @test length(actions(t)) == batch_size
        @test length(player_policies(t)) == batch_size
        @test length(player_ids(t)) == batch_size
        @test length(next_is_terminal(t)) == batch_size
    end
end

end


using StaticArrays: SMatrix

using ..Continuations: AbstractContinuationProblem, AbstractContinuationSolver,
    AbstractContinuationCache,
    as, SweepSetup, ContinuationSweep, ContinuationSolution,
    ContinuationCache, ContinuationOptions, ContinuationSolver,
    residual_jacobian!

using ..BifurcationsBase: timekind, Continuous, Discrete
import ..BifurcationsBase: TimeKind

using ..BifurcationsBase
using ..BifurcationsBase:
    AbstractSpecialPoint, SpecialPoint, SpecialPointInterval,
    BifurcationSweep, BifurcationSolution, BifurcationSolver,
    BifurcationProblem, BifurcationCache,
    eigvals_prototype, allocate_sweep!, check_sweep_length, record!
import ..BifurcationsBase: analyze!, re_analyze!

abstract type Codim1Problem{skind, tkind} <: BifurcationProblem{skind, tkind} end

module PointTypes
@enum(
    PointType,
    none,
    # initial,
    # simple_bifurcation,
    saddle_node,
    hopf,
    period_doubling,
)
end  # module
using .PointTypes: PointType

const AbstractCodim1SpecialPoint{tkind <: TimeKind} =
    AbstractSpecialPoint{tkind, PointType}

const Codim1Sweep{tkind <: TimeKind, ptType <: PointType} =
    BifurcationSweep{tkind, ptType}
const Codim1Solution{S <: ContinuationSolution, W <: Codim1Sweep} =
    BifurcationSolution{S, W}
# The constraints `<: TimeKind` and `<: ContinuationSolution` are
# important.  For example, the compiler didn't infer that
# `Codim1Sweep` was more specific than `BifurcationSweep` when `tkind`
# type parameter for `Codim1Sweep` was unconstrained.

BifurcationsBase.point_type_type(::Codim1Problem) = PointType

BifurcationsBase.regular_point(::Type{PointType}) = PointTypes.none

BifurcationsBase.eigvals_prototype(prob::Codim1Problem,
                                   cache::ContinuationCache) =
    cache.u[1:end - 1] .+ 1im
# TODO: improve it for SVector

mutable struct Codim1Cache{P, C <: ContinuationCache{P},
                           JType, eType,
                           } <: BifurcationCache{P}
    super::C
    J::JType
    prev_det::Float64
    eigvals::eType
    point_type::PointType
end

function Codim1Cache(super::C,
                     J::JType,
                     eigvals::eType,
                     point_type = PointTypes.none,
                     ) where {P, C <: ContinuationCache{P},
                              JType, eType}
    return Codim1Cache{P, C, JType, eType}(
        super, J, NaN, eigvals, point_type)
end
# TODO: Remove this constructor after removing the type parameter `P`.

Codim1Cache(prob::Codim1Problem, super::ContinuationCache) =
    Codim1Cache(
        super,
        ds_jacobian(super),
        copy(eigvals_prototype(prob, super)),
    )

BifurcationsBase.BifurcationCache(prob::Codim1Problem,
                                  super::ContinuationCache) =
    Codim1Cache(prob, super)

const Codim1Solver{
        R <: ContinuationSolver,
        P <: Codim1Problem,
        C <: Codim1Cache,
        S <: Codim1Solution,
        } =
    BifurcationSolver{R, P, C, S}

function re_analyze!(solver::Codim1Solver, u::AbstractVector)
    residual_jacobian!(as(solver.cache, ContinuationCache), u)
    analyze!(solver.cache, solver.opts)

    # Suppress special point recording:
    # It's a bit ugly hack... (communicate by sharing!)
    solver.cache.point_type = PointTypes.none  # TODO: FIX!

    record!(solver.sol, solver.cache)
end

function analyze!(cache::Codim1Cache, opts)
    cache.J = J = ds_jacobian(cache)
    eigvals = ds_eigvals(timekind(cache), J)
    cache.point_type = guess_point_type(timekind(cache), cache,
                                        J, eigvals, opts)
    cache.eigvals = eigvals
end

ds_jacobian(solver) = ds_jacobian(as(solver, ContinuationSolver).cache)
ds_jacobian(cache::Codim1Cache) = ds_jacobian(as(cache, ContinuationCache))
ds_jacobian(cache::ContinuationCache) = ds_jacobian(cache.J)
ds_jacobian(HJ::AbstractArray) = @view HJ[:, 1:end-1]
ds_jacobian(HJ::SMatrix) = HJ[:, 1:end-1]
# TOOD: optimize it for StaticArrays using generated functions


### BasicDiscMuvParameter

type BasicDiscMuvParameter <: Parameter{Discrete, Multivariate}
  key::Symbol
  index::Integer
  pdf::Union{DiscreteMultivariateDistribution, Void}
  prior::Union{DiscreteMultivariateDistribution, Void}
  setpdf::Union{Function, Void}
  setprior::Union{Function, Void}
  loglikelihood!::Union{Function, Void}
  logprior!::Union{Function, Void}
  logtarget!::Union{Function, Void}
  states::VariableStateVector

  function BasicDiscMuvParameter(
    key::Symbol,
    index::Integer,
    pdf::Union{DiscreteMultivariateDistribution, Void},
    prior::Union{DiscreteMultivariateDistribution, Void},
    setpdf::Union{Function, Void},
    setprior::Union{Function, Void},
    ll::Union{Function, Void},
    lp::Union{Function, Void},
    lt::Union{Function, Void},
    states::VariableStateVector
  )
    instance = new()

    instance.key = key
    instance.index = index
    instance.pdf = pdf
    instance.prior = prior

    args = (setpdf, setprior, ll, lp, lt)
    fnames = fieldnames(BasicDiscMuvParameter)[5:9]

    # Check that all generic functions have correct signature
    for i in 1:5
      if isa(args[i], Function) &&
        isgeneric(args[i]) &&
        !method_exists(args[i], (BasicDiscMuvParameterState, VariableStateVector))
        error("$(fnames[i]) has wrong signature")
      end
    end

    # Define setpdf (i = 1) and setprior (i = 2)
    for (i, setter, distribution) in ((1, :setpdf, :pdf), (2, :setprior, :prior))
      setfield!(
        instance,
        setter,
        if isa(args[i], Function)
          eval(codegen_setfield(instance, distribution, args[i]))
        else
          nothing
        end
      )
    end

    # Define loglikelihood!
    setfield!(
      instance,
      :loglikelihood!,
      if isa(args[3], Function)
        eval(codegen_closure(instance, args[3]))
      else
        nothing
      end
    )

    # Define logprior!
    setfield!(
      instance,
      :logprior!,
      if isa(args[4], Function)
        eval(codegen_closure(instance, args[4]))
      else
        if (isa(prior, DiscreteMultivariateDistribution) && method_exists(logpdf, (typeof(prior), Vector{eltype(prior)}))) ||
          isa(args[2], Function)
          eval(codegen_target_closure_via_distribution(instance, :prior, logpdf, :logprior))
        else
          nothing
        end
      end
    )

    # Define logtarget!
    setfield!(
      instance,
      :logtarget!,
      if isa(args[5], Function)
        eval(codegen_closure(instance, args[5]))
      else
        if isa(args[3], Function) && isa(getfield(instance, :logprior!), Function)
          eval(codegen_sumtarget_closure(instance, :loglikelihood!, :logprior!, :logtarget, :loglikelihood, :logprior))
        elseif (isa(pdf, DiscreteMultivariateDistribution) && method_exists(logpdf, (typeof(pdf), Vector{eltype(pdf)}))) ||
          isa(args[1], Function)
          eval(codegen_target_closure_via_distribution(instance, :pdf, logpdf, :logtarget))
        else
          nothing
        end
      end
    )

    instance
  end
end

BasicDiscMuvParameter(key::Symbol, index::Integer=0; signature::Symbol=:high, args...) =
  BasicDiscMuvParameter(key, Val{signature}, index; args...)

BasicDiscMuvParameter(
  key::Symbol,
  ::Type{Val{:low}},
  index::Integer=0;
  pdf::Union{DiscreteMultivariateDistribution, Void}=nothing,
  prior::Union{DiscreteMultivariateDistribution, Void}=nothing,
  setpdf::Union{Function, Void}=nothing,
  setprior::Union{Function, Void}=nothing,
  loglikelihood::Union{Function, Void}=nothing,
  logprior::Union{Function, Void}=nothing,
  logtarget::Union{Function, Void}=nothing,
  states::VariableStateVector=VariableState[]
) =
  BasicDiscMuvParameter(key, index, pdf, prior, setpdf, setprior, loglikelihood, logprior, logtarget, states)

function BasicDiscMuvParameter(
  key::Symbol,
  ::Type{Val{:high}},
  index::Integer=0;
  pdf::Union{DiscreteMultivariateDistribution, Void}=nothing,
  prior::Union{DiscreteMultivariateDistribution, Void}=nothing,
  setpdf::Union{Function, Void}=nothing,
  setprior::Union{Function, Void}=nothing,
  loglikelihood::Union{Function, Expr, Void}=nothing,
  logprior::Union{Function, Expr, Void}=nothing,
  logtarget::Union{Function, Expr, Void}=nothing,
  states::VariableStateVector=VariableState[],
  nkeys::Integer=0,
  vfarg::Bool=false
)
  inargs = (setpdf, setprior, loglikelihood, logprior, logtarget)

  fnames = Array(Any, 5)
  fnames[1:2] = fill(Symbol[], 2)
  fnames[3:5] = [Symbol[f] for f in fieldnames(BasicDiscMuvParameterState)[2:4]]

  @assert nkeys >= 0 "nkeys must be non-negative, got $nkeys"

  outargs = Union{Function, Void}[nothing for i in 1:5]

  for i in 1:5
    if isa(inargs[i], Function)
      outargs[i] = eval(
        codegen_lowlevel_variable_method(inargs[i], :BasicDiscMuvParameterState, true, fnames[i], nkeys, vfarg)
      )
    end
  end

  BasicDiscMuvParameter(key, index, pdf, prior, outargs..., states)
end

value_support(::Type{BasicDiscMuvParameter}) = Discrete
value_support(::BasicDiscMuvParameter) = Discrete

variate_form(::Type{BasicDiscMuvParameter}) = Multivariate
variate_form(::BasicDiscMuvParameter) = Multivariate

default_state_type(::BasicDiscMuvParameter) = BasicDiscMuvParameterState

default_state{N<:Integer}(variable::BasicDiscMuvParameter, value::Vector{N}, outopts::Dict) =
  BasicDiscMuvParameterState(
    value, (haskey(outopts, :diagnostics) && in(:accept, outopts[:diagnostics])) ? [:accept] : Symbol[]
  )

Base.show(io::IO, ::Type{BasicDiscMuvParameter}) = print(io, "BasicDiscMuvParameter")
Base.writemime(io::IO, ::MIME"text/plain", t::Type{BasicDiscMuvParameter}) = show(io, t)


"""
    Code for convenience functions for calculating seprate CCFs for each chunk in an AbstractChunkList (see RvSpectMLBase)
Author: Eric Ford
Created: August 2020
"""

"""  `calc_order_ccfs_chunklist ( chunklist_timeseries, list_of_ccf_plans )`
Convenience function to compute separate CCFs for each chunk (potentially an order or view around one or two lines) in a spectrum.
CCF is evaluated using line list and mask_shape provided by the ccf plan for each chunk.
# Inputs:
- `chunklist_timeseries`:
- `list_of_ccf_plans`: ccf plans (one for each chunk)
# Optional Arguments:
- `assume_sorted`:  if true, skips checking the line_list is sorted by wavelength
# Return:
A 2-d array containing the CCF at each (velocity, chunk)
"""
function calc_order_ccfs_chunklist!(chunk_ccfs_out::AbstractArray{T1,2}, chunk_list::AbstractChunkList,
    plan_for_chunk::AbstractVector{PlanT} = BasicCCFPlan(); assume_sorted::Bool = false  ) where {
                        T1<:Real, PlanT<:AbstractCCFPlan }
    num_chunks = length(chunk_list)
    @assert length(plan_for_chunk) == num_chunks
    @assert size(chunk_ccfs_out,2) == num_chunks
    @assert assume_sorted || issorted( first(plan_for_chunk).line_list.λ )
    num_vels = maximum(map(chid->calc_length_ccf_v_grid(plan_for_chunk[chid]), 1:num_chunks ))
    @assert size(chunk_ccfs_out,1) == num_vels
    #flush(stdout); println("num_vels = ",num_vels," num_chunks = ", num_chunks, " length(chunk_list.data)=", length(chunk_list.data));     flush(stdout);
    for chid in 1:num_chunks
        # chunk_ccfs_out[:,i] .=
        #flush(stdout);         println("chid = ", chid);         flush(stdout);
        calc_ccf_chunk!(view(chunk_ccfs_out, 1:num_vels, chid ), chunk_list.data[chid], plan_for_chunk[chid], assume_sorted=assume_sorted )
    end
    return chunk_ccfs_out
end

"""  `calc_order_ccfs_chunklist ( chunklist_timeseries, list_of_ccf_plans )`
Convenience function to compute separate CCFs for each chunk (potentially an order or view around one or two lines) in a spectrum.
CCF is evaluated using line list and mask_shape provided by the ccf plan for each chunk.
# Inputs:
- `chunklist_timeseries`:
- `list_of_ccf_plans`: ccf plans (one for each chunk)
# Optional Arguments:
- `assume_sorted`:  if true, skips checking the line_list is sorted by wavelength
# Return:
A 2-d array containing the CCF at each (velocity, chunk)
"""
function calc_order_ccfs_chunklist(chunk_list::AbstractChunkList,
                plan_for_chunk::AbstractVector{PlanT} = BasicCCFPlan();
                Δfwhm::Real = 0, assume_sorted::Bool = false ) where {
                        PlanT<:AbstractCCFPlan }
    num_chunks = length(chunk_list)
    @assert length(plan_for_chunk) == num_chunks
    @assert assume_sorted || issorted( first(plan_for_chunk).line_list.λ )
    num_vels = maximum(map(chid->calc_length_ccf_v_grid(plan_for_chunk[chid]), 1:num_chunks ))
    if Δfwhm > 0
       this_plan_for_chunk = copy(plan_for_chunk)
       #increase_mask_fwhm!(this_plan_for_chunk,Δfwhm)
       map(chid->increase_mask_fwhm!(this_plan_for_chunk[chid],Δfwhm), 1:num_chunks )
    else
       this_plan_for_chunk = plan_for_chunk
    end

    chunk_ccfs_out = zeros( num_vels, num_chunks )
    calc_order_ccfs_chunklist!(chunk_ccfs_out, chunk_list, this_plan_for_chunk, assume_sorted=assume_sorted )
    return chunk_ccfs_out
end

"""  `calc_order_ccfs_chunklist ( chunklist_timeseries, list_of_ccf_plans )`
Convenience function to compute separate CCFs for each chunk (potentially an order or view around one or two lines) in a spectrum.
CCF is evaluated using line list and mask_shape provided by the ccf plan for each chunk.
# Inputs:
- `chunklist_timeseries`:
- `list_of_ccf_plans`: ccf plans (one for each chunk)
# Optional Arguments:
- `assume_sorted`:  if true, skips checking the line_list is sorted by wavelength
# Return:
A 2-d array containing the CCF at each (velocity, chunk)
"""
function calc_order_ccf_and_vars_chunklist!(chunk_ccfs_out::AbstractArray{T1,2}, chunk_ccf_vars_out::AbstractArray{T1,2}, chunk_list::AbstractChunkList,
    plan_for_chunk::AbstractVector{PlanT} = BasicCCFPlan(); assume_sorted::Bool = false ) where {
                        T1<:Real, PlanT<:AbstractCCFPlan }
    num_chunks = length(chunk_list)
    @assert length(plan_for_chunk) == num_chunks
    @assert size(chunk_ccfs_out,2) == num_chunks
    @assert assume_sorted || issorted( first(plan_for_chunk).line_list.λ )
    num_vels = maximum(map(chid->calc_length_ccf_v_grid(plan_for_chunk[chid]), 1:num_chunks ))
    @assert size(chunk_ccfs_out,1) == num_vels

    for chid in 1:num_chunks
        calc_ccf_and_var_chunk!(view(chunk_ccfs_out, 1:num_vels, chid ), view(chunk_ccf_vars_out, 1:num_vels, chid ),
            chunk_list.data[chid], plan_for_chunk[chid], assume_sorted=assume_sorted )
    end
    return (ccfs=chunk_ccfs_out, ccf_vars=chunk_ccf_vars_out)
end

"""  `calc_order_ccfs_chunklist ( chunklist_timeseries, list_of_ccf_plans )`
Convenience function to compute separate CCFs for each chunk (potentially an order or view around one or two lines) in a spectrum.
CCF is evaluated using line list and mask_shape provided by the ccf plan for each chunk.
# Inputs:
- `chunklist_timeseries`:
- `list_of_ccf_plans`: ccf plans (one for each chunk)
# Optional Arguments:
- `assume_sorted`:  if true, skips checking the line_list is sorted by wavelength
# Return:
A 2-d array containing the CCF at each (velocity, chunk)
"""
function calc_order_ccf_and_vars_chunklist(chunk_list::AbstractChunkList,
            plan_for_chunk::AbstractVector{PlanT} = BasicCCFPlan();
            Δfwhm::Real = 0, assume_sorted::Bool = false ) where {
                        PlanT<:AbstractCCFPlan }
    num_chunks = length(chunk_list)
    @assert length(plan_for_chunk) == num_chunks
    @assert assume_sorted || issorted( first(plan_for_chunk).line_list.λ )

    if Δfwhm > 0
       this_plan_for_chunk = copy(plan_for_chunk)
       #increase_mask_fwhm!(this_plan_for_chunk,Δfwhm)
       map(chid->increase_mask_fwhm!(this_plan_for_chunk[chid],Δfwhm), 1:num_chunks )
    else
       this_plan_for_chunk = plan_for_chunk
    end

    num_vels = maximum(map(chid->calc_length_ccf_v_grid(plan_for_chunk[chid]), 1:num_chunks ))
    chunk_ccfs_out = zeros( num_vels, num_chunks )
    chunk_ccf_vars_out = zeros( num_vels, num_chunks )

    calc_order_ccf_and_vars_chunklist!(chunk_ccfs_out, chunk_ccf_vars_out, chunk_list, this_plan_for_chunk, assume_sorted=assume_sorted )
    return (ccfs=chunk_ccfs_out, ccf_vars=chunk_ccf_vars_out)
end


include("DataDictionary.jl")
include("FluxDriver.jl")
include("FVA.jl")

using GLPK
using DelimitedFiles


@testset "not equal" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, y[1:3], CS.Integers([-3, 1, 2, 3]))
    @constraint(m, sum(y) + 1 != 5)
    optimize!(m)
    com = CS.get_inner_model(m)
    constraint = get_constraints_by_type(com, CS.LinearConstraint)[1]

    # doesn't check the length
    # 1+2+1 + constant (1) == 5
    @test !CS.is_constraint_solved(constraint, constraint.fct, constraint.set, [1, 2, 1])
    @test CS.is_constraint_solved(constraint, constraint.fct, constraint.set, [1, 2, 2])

    constr_indices = constraint.indices
    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        -3,
    )
    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        1,
    )
    CS.fix!(com, com.search_space[constr_indices[1]], 2)
    CS.fix!(com, com.search_space[constr_indices[2]], 1)
    @test !CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[3],
        1,
    )

    # need to create a backtrack_vec to reverse pruning
    dummy_backtrack_obj = CS.BacktrackObj(com)
    push!(com.backtrack_vec, dummy_backtrack_obj)
    # reverse previous fix
    CS.reverse_pruning!(com, 1)
    com.c_backtrack_idx = 1
    # now setting it to 1 should be feasible
    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[3],
        1,
    )

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, y[1:3], CS.Integers([-3, 1, 2, 3]))
    @constraint(m, sum(y) + 1 != 5)
    optimize!(m)

    com = CS.get_inner_model(m)
    constraint = get_constraints_by_type(com, CS.LinearConstraint)[1]

    CS.fix!(com, com.search_space[constr_indices[1]], 2)
    CS.fix!(com, com.search_space[constr_indices[2]], 1)
    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)
    @test sort(CS.values(com.search_space[constr_indices[3]])) == [-3, 2, 3]
end

@testset "not equal is_constraint_violated test" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, -5 <= x[1:2] <= 5, Int)
    @constraint(m, sum(x) != 7)
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 5; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 2; check_feasibility = false)
    @test CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, -5 <= x[1:2] <= 5, Int)
    @constraint(m, sum(x) != 7)
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 5; check_feasibility = false)
    @test !CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)
end


__precompile__(true)

module FixedEffectModels

##############################################################################
##
## Dependencies
##
##############################################################################
using Compat
import Base.BLAS: axpy!
import Base: A_mul_B!, Ac_mul_B!, size, sumabs2, copy!, getindex, length, fill!, norm, scale!, eltype, length, view, start, next, done
import Distributions: TDist, ccdf, FDist, Chisq, AliasTable, Categorical
import DataArrays: RefArray, PooledDataArray, PooledDataVector, DataArray, DataVector, compact, NAtype
import DataFrames: @~, DataFrame, AbstractDataFrame, ModelMatrix, ModelFrame, Terms, coefnames, Formula, complete_cases, names!, pool
import StatsBase: coef, nobs, coeftable, vcov, predict, residuals, var, RegressionModel, model_response, stderr, confint, fit, CoefTable, df_residual
##############################################################################
##
## Exported methods and types
##
##############################################################################

export group,
reg,
partial_out,
residualize!,
getfe!,
decompose!,
allvars,

Ones,
FixedEffect,
FixedEffectProblem,

AbstractRegressionResult,
title,
top,
RegressionResult,
RegressionResultIV,
RegressionResultFE,
RegressionResultFEIV,

AbstractVcovMethod,
AbstractVcovMethodData,
vcov!,
shat!,
VcovMethodData,
VcovData,
VcovSimple,
VcovWhite,
VcovCluster

##############################################################################
##
## Load files
##
##############################################################################
include("utils/group.jl")
include("utils/formula.jl")
include("utils/lsmr.jl")
include("utils/basecol.jl")
include("utils/Ones.jl")
include("utils/combinations.jl")

include("RegressionResult.jl")

include("fixedeffect/FixedEffect.jl")
include("fixedeffect/FixedEffectProblem.jl")
include("fixedeffect/FixedEffectProblem_LSMR.jl")
include("fixedeffect/FixedEffectProblem_Factorization.jl")

include("vcov/types.jl")
include("vcov/vcovsimple.jl")
include("vcov/vcovwhite.jl")
include("vcov/vcovcluster.jl")
include("vcov/ranktest.jl")

include("reg.jl")
include("partial_out.jl")

end  # module FixedEffectModels


# SPDX-License-Identifier: MIT

using CUDA
using CUDA.CUSPARSE
using JLD2
using MarkovModels
using SparseArrays
using LinearAlgebra

function main(fsmfile, T, B, N; warmup)
    # Load the Finite State Machine (FSM).
    fsm = jldopen(fsmfile, "r") do f
       f["fsm"]
    end

    # Number states in the FSM.
    S = length(fsm)

    # Create a batch of the same FSM.
    fsm = union(repeat([fsm], B)...)

    # Number of pdfs associated with the FSM. This is specific to the FSM
    # and should not be changed unless you use your own FSM.
    npdfs = 84

    # Length of the sequences of the batch. Here we assume that all
    # the sequence have the same length.
    lengths = repeat([N], B)

    # Generate some pseudo-(log)likelihood
    lhs = ones(T, npdfs, N, B)

    # Computing precision
    precision = T == Float64 ? "double" : "single"

    # Dummy run to force compilation
    if warmup pdfposteriors(fsm, lhs, lengths) end

    etime = time_ns()
    pdfposteriors(fsm, lhs, lengths)
    etime = (time_ns() - etime) / 1e9
    println("julia\t$precision\t$N\t$S\t$B\tsemifield\tcpu\t$etime")

    # Move the data and the FSM to the GPU
    fsm = fsm |> MarkovModels.gpu
    lhs = CuArray(lhs)

    # Dummy run to force compilation.
    if warmup pdfposteriors(fsm, lhs, lengths) end

    etime = time_ns()
    pdfposteriors(fsm, lhs, lengths)
    etime = (time_ns() - etime) / 1e9
    println("julia\t$precision\t$N\t$S\t$B\tsemifield\tgpu\t$etime")
end

main(
    "denominator_fsm.jld2", # "numerator_fsm.jld2"
    Float32,    # computing precision
    128,        # batch size
    700;        # sequence length
    warmup=true # if true, make a dummy run first to avoid compilation time
)