# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    HullMethod

A method for computing hulls of point sets or
collections of geometries.
"""
abstract type HullMethod end

"""
    hull(object, method)

Compute the hull of `object` with given `method`.
"""
function hull end

# ----------------
# IMPLEMENTATIONS
# ----------------

include("hulls/graham.jl")


using HTTP
using DataFrames

"""
    network(;apikey, bbox)

Returns shape for every street of the network within the bounding box.

# Arguments
apikey::String   : API key to access HERE
bbox::String     : Bounding box defined as `"lat₁,long₁;lat₂,long₂"`

For more information visit: https://developer.here.com/documentation/traffic/dev_guide/topics/guide.html

Example request: https://developer.here.com/documentation/traffic/dev_guide/topics_v6.1/example-flow-bounding-box.html
"""
function network(;apikey, bbox)
    # Request
    baseurl = "https://traffic.ls.hereapi.com"
    path = "traffic/6.2"
    resource = "flow.xml"
    url = "$baseurl/$path/$resource?apiKey=$apikey&bbox=$bbox&responseattributes=shape"
    r = HTTP.request(:GET, url)
    f = parsexml(String(r.body))
    rootNode = root(f)

    # Build network
    df = DataFrame(lid = Int64[], name = String[], length = Float64[], ffs = Float64[], shape = String[])
    for node in iterate(rootNode)
        if nodename(node) == "FI"
            shape = ""
            for childnode in eachelement(node)
                # Fetch name and length
                if childnode == firstelement(node)
                    push!(df[!, :lid], length(df[!, :lid]) + 1)
                    push!(df[!, :name], childnode["DE"])
                    push!(df[!, :length], parse(Float64, childnode["LE"]))
                # Fetch free flow speed
                elseif childnode == lastelement(node)
                    push!(df[!, :ffs], parse(Float64, childnode["FF"]))
                # Fetch link shape (Note: Link shape is of the same format as the bbox)
                else
                    shape *= replace(nodecontent(childnode), " " => ";")
                end
            end
            push!(df[!, :shape], shape)
        end
    end

    return df
end


export Model2Fit_Empty

"""

An empty model without parameters. This is useful to initialize
summation:

```julia
model = Model2Fit_Empty()

for i in 1:n
    model = model + GaussianPeak()
end
```
"""
struct Model2Fit_Empty <: Abstract_Model2Fit
end

parameter_size(m::Model2Fit_Empty) = 0
accumulate_y!(m::Model2Fit_Empty,Y::AbstractVector,X::AbstractVector,θ::AbstractVector) = Y


import GridGames
using PrettyTables
using Dates

mutable struct Counter
    n::Int32
end
Base.:(==)(a::Counter, b::Counter) = a.n == b.n
set!(c::Counter, n) = c.n = n
increment!(c::Counter) = c.n += 1

mutable struct RunningAvg
    count::Int
    avg::Float32
    RunningAvg() = new(0, 0.)
end

function update_avg!(x::RunningAvg, amount)
    x.count += 1
    x.avg += (1 ./ x.count)*(amount - x.avg)
end

# @todo rename to MuZeroNetwork
abstract type Network end

abstract type AbstractAction end
abstract type AbstractGame end
abstract type AbstractState end
abstract type AbstractPlayer end

AbstractEnv = Union{AbstractGame,GridGames.AbstractGridGame}
AbstractEnvState = Union{AbstractState,GridGames.AbstractGridState}
AbstractEnvAction = Union{AbstractAction,GridGames.AbstractGridAction}
AbstractEnvPlayer = Union{AbstractPlayer,GridGames.AbstractGridPlayer}

abstract type AbstractBoardGame <: AbstractGame end
abstract type AbstractBoardState <: AbstractState end
abstract type AbstractBoardGameSVG <: AbstractBoardGame end
abstract type AbstractBoardStateSVG <: AbstractBoardState end
abstract type TrainingTest end
abstract type TrainingTests end
abstract type AbstractImage end
abstract type AbstractMuZeroGame end
abstract type AbstractLearner end

struct BoardGameSVG <: AbstractBoardGameSVG
    x::AbstractBoardGame
end

avgx(x,size) = mean.(collect(Iterators.partition(x, size)))

mutable struct MuZeroGame <: AbstractMuZeroGame
    game::AbstractEnv
    child_visits::Vector{Vector{Float16}}
    root_values::Vector{Float16}
    net_step::Integer
    function MuZeroGame(game::AbstractEnv, net_step=0)
        new(
            game,
            Vector{Int16}(),
            [],
            net_step,
        )
    end
end

# Sample position from game either uniformly or according to some priority.
function sample_position(game::MuZeroGame)
    to = length(history(game)) + 1
    rand(1:to)
end

ensuredirexists(dir) = mkpath(expanduser(dir))

function ensure_pathto_exists(path)
    a = split(expanduser(path), "/")
    join(a[1:(end-1)], "/") |> mkpath
end

valornothing(d, k) = valornothing(d, (x)->x, k)
valornothing(d, fn, k) = haskey(d, k) ? fn(d[k]) : nothing
valormissing(d, k) = valormissing(d, (x)->x, k)
valormissing(d, fn, k) = haskey(d, k) ? fn(d[k]) : missing

lastbykeyornothing(d::Dict) = lastbykeyornothing(d,(x)->x)
function lastbykeyornothing(d::Dict, fn)
    if isempty(d)
        return nothing
    end
    last_key = keys(d) |> collect |> sort |> last
    fn(d[last_key])
end

isfileexpanduser(path) = isfile(expanduser(path))
isdirexpanduser(path) = isdir(expanduser(path))
readdirexpanduser(path) = readdir(expanduser(path), join=true)
serializeexpanduser(path, x) = serialize(expanduser(path), x)
deserializeexpanduser(path) = deserialize(expanduser(path))

getcheckpointstr(dt::DateTime) = getcheckpointstr(string(dt))

function getcheckpointstr(datetime_str::String)
    replace(
        datetime_str,
        r"[\:]"=>"-"
    )[1:(end-4)]
end


"""
Single variable bound constraints
"""
# MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{AbstractVector{<:MOI.AbstractScalarSet}}) = true
MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{MOI.LessThan{T}}) where {T<:Real} =
    true

MOI.supports_constraint(
    ::Optimizer,
    ::Type{SVF},
    ::Type{MOI.GreaterThan{T}},
) where {T<:Real} = true

MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{MOI.EqualTo{T}}) where {T<:Real} =
    true

MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{MOI.Interval{T}}) where {T<:Real} =
    true

MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{Integers}) = true

"""
Binary/Integer variable support
"""
MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{<:VAR_TYPES}) = true

"""
    check_var_bounds(model)

Checks whether all variables are integral and bounded on both sides. If not an error is throws
"""
function check_var_bounds(model::Optimizer)
    for var in model.variable_info
        if !var.has_lower_bound || !var.has_upper_bound || !var.is_integer
            throw(ErrorException("Each variable must be an integer and bounded. Currently the variable index $(var.idx) doesn't fulfill this requirements."))
        end
    end
end

"""
    addupd_var_in_inner_model(model::Optimizer, vidx::Int)

Adds all variables to model.inner.search_space
"""
function addupd_var_in_inner_model(model::Optimizer, vidx::Int)
    if length(model.inner.search_space) < vidx
        push!(model.inner.search_space, model.variable_info[vidx])
    else
        model.inner.search_space[vidx] = model.variable_info[vidx]
    end
end

"""
    MOI.add_variable(model::Optimizer)
"""
function MOI.add_variable(model::Optimizer)
    vidx = length(model.variable_info) + 1
    push!(model.variable_info, Variable(vidx))
    changes = changes = Vector{Vector{Tuple{Symbol,Int,Int,Int}}}()
    push!(changes, Vector{Tuple{Symbol,Int,Int,Int}}())
    model.variable_info[vidx].changes = changes
    push!(model.inner.subscription, Int[])
    push!(model.inner.bt_infeasible, 0)
    push!(model.inner.var_in_obj, false)
    addupd_var_in_inner_model(model, vidx)
    return MOI.VariableIndex(vidx)
end

MOI.add_variables(model::Optimizer, n::Int) = [MOI.add_variable(model) for i in 1:n]

function check_inbounds(model::Optimizer, vi::VI)
    num_variables = length(model.variable_info)
    if !(1 <= vi.value <= num_variables)
        @error "Invalid variable index $vi. ($num_variables variables in the model.)"
    end
    return
end

check_inbounds(model::Optimizer, var::SVF) = check_inbounds(model, var.variable)

has_upper_bound(model::Optimizer, vi::VI) = model.variable_info[vi.value].has_upper_bound

has_lower_bound(model::Optimizer, vi::VI) = model.variable_info[vi.value].has_lower_bound

is_fixed(model::Optimizer, vi::VI) = model.variable_info[vi.value].is_fixed

function MOI.add_constraint(model::Optimizer, v::SVF, t::MOI.Integer)
    vi = v.variable
    model.variable_info[vi.value].is_integer = true

    cidx = length(model.var_constraints) + 1
    push!(model.var_constraints, (vi.value, :int, typemin(Int), typemax(Int)))

    addupd_var_in_inner_model(model, vi.value)
    return MOI.ConstraintIndex{SVF,MOI.Integer}(cidx)
end

function MOI.add_constraint(model::Optimizer, v::SVF, t::MOI.ZeroOne)
    vi = v.variable
    model.variable_info[vi.value].is_integer = true

    # this gets called after setting lower and upper bound
    # => make sure that it's not already set
    if !has_upper_bound(model, vi) || model.variable_info[vi.value].upper_bound > 1
        model.variable_info[vi.value].upper_bound = 1
        model.variable_info[vi.value].max = 1
        model.variable_info[vi.value].has_upper_bound = true
    end
    if !has_lower_bound(model, vi) || model.variable_info[vi.value].lower_bound < 0
        model.variable_info[vi.value].lower_bound = 0
        model.variable_info[vi.value].min = 0
        model.variable_info[vi.value].has_lower_bound = true
    end
    min_val, max_val = model.variable_info[vi.value].lower_bound,
                       model.variable_info[vi.value].upper_bound
    values = collect(min_val:max_val)
    model.variable_info[vi.value].values = values
    model.variable_info[vi.value].init_vals = copy(values)
    model.variable_info[vi.value].init_val_to_index = 1:length(values)
    model.variable_info[vi.value].offset = 1
    model.variable_info[vi.value].indices = 1:length(values)
    model.variable_info[vi.value].first_ptr = 1
    model.variable_info[vi.value].last_ptr = length(values)
    addupd_var_in_inner_model(model, vi.value)

    cidx = length(model.var_constraints) + 1
    push!(model.var_constraints, (vi.value, :bin, typemin(Int), typemax(Int)))
    return MOI.ConstraintIndex{SVF,MOI.ZeroOne}(cidx)
end

function MOI.add_constraint(model::Optimizer, v::SVF, t::Integers)
    vi = v.variable
    model.variable_info[vi.value].is_integer = true

    set_vals = t.values
    set_variable_from_integers!(model.variable_info[vi.value], set_vals)

    addupd_var_in_inner_model(model, vi.value)
    min_val = model.variable_info[vi.value].min
    max_val = model.variable_info[vi.value].max

    cidx = length(model.var_constraints) + 1
    # TODO: If we want to do something with var_constraints later we need to save the actual input values
    push!(model.var_constraints, (vi.value, :integers, min_val, max_val))
    return MOI.ConstraintIndex{SVF,Integers}(cidx)
end

#=
Populating Variable bounds
=#
function MOI.add_constraint(
    model::Optimizer,
    v::SVF,
    interval::MOI.Interval{T},
) where {T<:Real}
    vi = v.variable
    check_inbounds(model, vi)
    isnan(interval.upper) &&
        throw(ErrorException("The interval bounds can not contain NaN and must be an Integer. Currently it has an upper bound of $(interval.upper)"))
    has_upper_bound(model, vi) && @error "Upper bound on variable $vi exists already."
    isnan(interval.lower) &&
        throw(ErrorException("The interval bounds can not contain NaN and must be an Integer. Currently it has a lower bound of $(interval.lower)"))
    has_lower_bound(model, vi) && @error "Lower bound on variable $vi exists already."
    is_fixed(model, vi) && @error "Variable $vi is fixed. Cannot also set upper bound."

    model.variable_info[vi.value].upper_bound = interval.upper
    model.variable_info[vi.value].max = interval.upper
    model.variable_info[vi.value].has_upper_bound = true
    model.variable_info[vi.value].lower_bound = interval.lower
    model.variable_info[vi.value].offset = 1 - interval.lower
    model.variable_info[vi.value].min = interval.lower
    model.variable_info[vi.value].has_lower_bound = true

    model.variable_info[vi.value].values =
        (model.variable_info[vi.value].min):(model.variable_info[vi.value].max)
    num_vals = model.variable_info[vi.value].max - model.variable_info[vi.value].min + 1
    model.variable_info[vi.value].indices = 1:num_vals
    model.variable_info[vi.value].first_ptr = 1
    model.variable_info[vi.value].last_ptr = num_vals
    model.variable_info[vi.value].init_vals = copy(model.variable_info[vi.value].values)
    model.variable_info[vi.value].init_val_to_index =
        copy(model.variable_info[vi.value].indices)

    addupd_var_in_inner_model(model, vi.value)

    cidx = length(model.var_constraints) + 1
    push!(model.var_constraints, (vi.value, :interval, interval.lower, interval.upper))
    return MOI.ConstraintIndex{SVF,MOI.Interval{T}}(cidx)
end

function MOI.add_constraint(model::Optimizer, v::SVF, lt::MOI.LessThan{T}) where {T<:Real}
    vi = v.variable
    check_inbounds(model, vi)
    isnan(lt.upper) &&
        throw(ErrorException("The variable bounds can not contain NaN and must be an Integer. Currently it has an upper bound of $(lt.upper)"))
    has_upper_bound(model, vi) && @error "Upper bound on variable $vi already exists."
    is_fixed(model, vi) && @error "Variable $vi is fixed. Cannot also set upper bound."

    model.variable_info[vi.value].upper_bound = lt.upper
    model.variable_info[vi.value].max = lt.upper
    model.variable_info[vi.value].has_upper_bound = true

    if has_lower_bound(model, vi)
        model.variable_info[vi.value].values =
            (model.variable_info[vi.value].min):(model.variable_info[vi.value].max)
        num_vals = model.variable_info[vi.value].max - model.variable_info[vi.value].min + 1
        model.variable_info[vi.value].indices = 1:num_vals
        model.variable_info[vi.value].first_ptr = 1
        model.variable_info[vi.value].last_ptr = num_vals
        model.variable_info[vi.value].init_vals = copy(model.variable_info[vi.value].values)
        model.variable_info[vi.value].init_val_to_index =
            copy(model.variable_info[vi.value].indices)
    end

    addupd_var_in_inner_model(model, vi.value)

    cidx = length(model.var_constraints) + 1
    push!(model.var_constraints, (vi.value, :leq, typemin(Int), lt.upper))
    return MOI.ConstraintIndex{SVF,MOI.LessThan{T}}(cidx)
end

function MOI.add_constraint(
    model::Optimizer,
    v::SVF,
    gt::MOI.GreaterThan{T},
) where {T<:Real}
    vi = v.variable
    check_inbounds(model, vi)
    isnan(gt.lower) &&
        throw(ErrorException("The variable bounds can not contain NaN and must be an Integer. Currently it has an upper bound of $(gt.upper)"))
    has_lower_bound(model, vi) && @error "Lower bound on variable $vi already exists."
    is_fixed(model, vi) && @error "Variable $vi is fixed. Cannot also set lower bound."

    model.variable_info[vi.value].lower_bound = gt.lower
    model.variable_info[vi.value].min = gt.lower
    model.variable_info[vi.value].offset = 1 - gt.lower
    model.variable_info[vi.value].has_lower_bound = true

    if has_upper_bound(model, vi)
        model.variable_info[vi.value].values =
            (model.variable_info[vi.value].min):(model.variable_info[vi.value].max)
        num_vals = model.variable_info[vi.value].max - model.variable_info[vi.value].min + 1
        model.variable_info[vi.value].indices = 1:num_vals
        model.variable_info[vi.value].first_ptr = 1
        model.variable_info[vi.value].last_ptr = num_vals
        model.variable_info[vi.value].init_vals = copy(model.variable_info[vi.value].values)
        model.variable_info[vi.value].init_val_to_index =
            copy(model.variable_info[vi.value].indices)
    end
    addupd_var_in_inner_model(model, vi.value)

    cidx = length(model.var_constraints) + 1
    push!(model.var_constraints, (vi.value, :geq, gt.lower, typemax(Int)))
    return MOI.ConstraintIndex{SVF,MOI.GreaterThan{T}}(cidx)
end

function MOI.add_constraint(model::Optimizer, v::SVF, eq::MOI.EqualTo{T}) where {T<:Real}
    vi = v.variable
    check_inbounds(model, vi)
    isnan(eq.value) && @error "Invalid fixed value $(eq.value)."
    has_lower_bound(model, vi) && @error "Variable $vi has a lower bound. Cannot be fixed."
    has_upper_bound(model, vi) && @error "Variable $vi has an upper bound. Cannot be fixed."
    is_fixed(model, vi) && @error "Variable $vi is already fixed."
    model.variable_info[vi.value].lower_bound = eq.value
    model.variable_info[vi.value].upper_bound = eq.value
    model.variable_info[vi.value].is_fixed = true
    model.variable_info[vi.value].has_lower_bound = true
    model.variable_info[vi.value].has_upper_bound = true
    model.variable_info[vi.value].min = eq.value
    model.variable_info[vi.value].max = eq.value
    model.variable_info[vi.value].values = [eq.value]
    model.variable_info[vi.value].indices = [1]
    model.variable_info[vi.value].first_ptr = 1
    model.variable_info[vi.value].last_ptr = 1
    model.variable_info[vi.value].init_vals = copy(model.variable_info[vi.value].values)
    model.variable_info[vi.value].init_val_to_index =
        copy(model.variable_info[vi.value].indices)
    addupd_var_in_inner_model(model, vi.value)

    cidx = length(model.var_constraints) + 1
    push!(model.var_constraints, (vi.value, :eq, eq.value, eq.value))
    return MOI.ConstraintIndex{SVF,MOI.EqualTo{T}}(cidx)
end

function set_variable_from_integers!(var::Variable, set_vals)
    min_val, max_val = extrema(set_vals)
    range = max_val - min_val + 1
    vals = copy(set_vals)

    # fill the indices such that 1:length(set_vals) map to set_vals
    indices = zeros(Int, range)
    # .- needed for the offset
    # values that don't exist get an index out of the range of the values array
    indices[set_vals .- (min_val - 1)] = 1:length(set_vals)
    j = length(set_vals) + 1
    for i in 1:range
        if indices[i] == 0
            indices[i] = j
            j += 1
        end
    end

    var.upper_bound = max_val
    var.max = max_val
    var.has_upper_bound = true
    var.lower_bound = min_val
    var.min = min_val
    var.has_lower_bound = true
    var.values = vals
    var.init_vals = copy(vals)
    var.offset = 1 - min_val
    var.indices = indices
    # needs copy to be different
    var.init_val_to_index = copy(indices)
    var.first_ptr = 1
    var.last_ptr = length(set_vals)
end

"""
    set_init_fixes!(com::CS.CoM)

Redefines variables if they are fixed to a specific value.
Return feasibility
"""
function set_init_fixes!(com::CS.CoM)
    for fixes in com.init_fixes
        vidx = fixes[1]
        val = fixes[2]
        var = com.search_space[vidx]
        !fix!(com, var, val) && return false
    end
    return true
end


function cc_add(x::Neuro1DRealSignal, y::Neuro1DRealSignal)::Neuro1DRealSignal
    if x.sample_rate != y.sample_rate
        throw(ArgumentError("the sample rate of x and y should be identical"))
    end

    if size(x.signal) != size(x.signal)
        throw(ArgumentError("the size of x and y should be identical"))
    end

    Neuro1DRealSignal(
        signal = x.signal .+ y.signal,
        sample_rate = x.sample_rate
    )
end


"""
Phenotype
"""

mutable struct Network
    nodes::Dict{Int, Node}
    input_id::Array
    hidden_id::Array
    output_id::Array

    function Network(config)
        println("build")

        nodes = Dict{Int, Node}()

        num_inputs = config["num_inputs"]
        num_outputs = config["num_outputs"]

        net = new(nodes, [], [], [])

        for i = 1:num_inputs
            add_node!(net, Node(i, INPUT::NodeType))
            for j = num_inputs+1:num_outputs+num_inputs
                if i == 1
                    add_node!(net, Node(j, OUTPUT::NodeType))
                end
                add_connection!(net, Connection(i, j))
            end
        end
        return net
    end

end

function add_node!(net::Network, node::Node)
    @assert !haskey(net.nodes, node.id) "ID already used"
    net.nodes[node.id] = node

    if node.dtype == INPUT
        push!(net.input_id, node.id)
    elseif node.dtype == OUTPUT
        push!(net.output_id, node.id)
    else
        push!(net.hidden_id, node.id)
    end
end

function add_connection!(net::Network, conn::Connection)
    push!(net.nodes[conn.input].connections["in"], conn)
    push!(net.nodes[conn.output].connections["out"], conn)
end


include("../NativeInterface/native.jl")

function call_basic_triangulation(flat_vertices::Vector{Cdouble}, flat_vertices_map::Vector{Cint})
    return NativeInterface.basic_triangulation(flat_vertices, flat_vertices_map)
end

function call_constrained_triangulation(flat_vertices::Vector{Cdouble}, flat_vertices_map::Vector{Cint}, flat_vertices_edge::Vector{Cint})
    return NativeInterface.constrained_triangulation(flat_vertices, flat_vertices_map, flat_vertices_edge)
end

function call_constrained_triangulation_bounded(flat_vertices::Vector{Cdouble}, flat_vertices_map::Vector{Cint}, flat_vertices_edge::Vector{Cint}, flat_boundary_edges::Vector{Cint})
    options = NativeInterface.TriangulateOptions()
    options.pslg = true

    return NativeInterface.constrained_triangulation_bounded(flat_vertices, flat_vertices_map, flat_vertices_edge, flat_boundary_edges, options)
end

function call_constrained_triangulation_bounded(flat_vertices::Vector{Cdouble}, flat_vertices_map::Vector{Cint}, flat_vertices_edge::Vector{Cint}, flat_boundary_edges::Vector{Cint}, flat_holes::Vector{Cdouble})
    options = NativeInterface.TriangulateOptions()
    options.pslg = true

    return NativeInterface.constrained_triangulation_bounded(flat_vertices, flat_vertices_map, flat_vertices_edge, flat_boundary_edges, flat_holes, options)
end


### Generic MCBaseSample type used by most samplers

type MCBaseSample <: MCSample{NullOrder}
  sample::Vector{Float64}
  logtarget::Float64
end

MCBaseSample(s::Vector{Float64}) = MCBaseSample(s, NaN)
MCBaseSample() = MCBaseSample(Float64[], NaN)

MCBaseSample(l::Int) = MCBaseSample(fill(NaN, l), NaN)

logtarget!(s::MCSample, f::Function) = (s.logtarget = f(s.sample))

### MCGradSample is used by samplers that compute (up to) the gradient of the log-target (for ex HMC and MALA)

type MCGradSample <: MCSample{FirstOrder}
  sample::Vector{Float64}
  logtarget::Float64
  gradlogtarget::Vector{Float64}
end

MCGradSample(s::Vector{Float64}) = MCGradSample(s, NaN, Float64[])
MCGradSample() = MCGradSample(Float64[], NaN, Float64[])

MCGradSample(l::Int) = MCGradSample(fill(NaN, l), NaN, fill(NaN, l))

gradlogtargetall!(s::MCSample, f::Function) = ((s.logtarget, s.gradlogtarget) = f(s.sample))

### MCMCTensorSample is used by samplers that compute (up to) the tensor of the log-target (for ex SMMALA)

type MCTensorSample <: MCSample{SecondOrder}
  sample::Vector{Float64}
  logtarget::Float64
  gradlogtarget::Vector{Float64}
  tensorlogtarget::Matrix{Float64}
end

MCTensorSample(s::Vector{Float64}) = MCTensorSample(s, NaN, Float64[], Array(Float64, 0, 0))
MCTensorSample() = MCTensorSample(Float64[], NaN, Float64[], Array(Float64, 0, 0))

MCTensorSample(l::Int) = MCTensorSample(fill(NaN, l), NaN, fill(NaN, l), fill(NaN, l, l))

tensorlogtargetall!(s::MCSample, f::Function) = ((s.logtarget, s.gradlogtarget, s.tensorlogtarget) = f(s.sample))

### MCMCDTensorSample is used by samplers that compute (up to) the derivative of the tensor of the log-target (for ex
### RMHMC and PMALA)

type MCDTensorSample <: MCSample{ThirdOrder}
  sample::Vector{Float64}
  logtarget::Float64
  gradlogtarget::Vector{Float64}
  tensorlogtarget::Matrix{Float64}
  dtensorlogtarget::Array{Float64, 3}
end

MCDTensorSample(s::Vector{Float64}) = MCDTensorSample(s, NaN, Float64[], Array(Float64, 0, 0), Array(Float64, 0, 0, 0))
MCDTensorSample() = MCDTensorSample(Float64[], NaN, Float64[], Array(Float64, 0, 0), Array(Float64, 0, 0, 0))

MCDTensorSample(l::Int) = MCDTensorSample(fill(NaN, l), NaN, fill(NaN, l), fill(NaN, l, l), fill(NaN, l, l, l))

dtensorlogtargetall!(s::MCSample, f::Function) =
  ((s.logtarget, s.gradlogtarget, s.tensorlogtarget, s.dtensorlogtarget) = f(s.sample))


# ------------ Factor ------------ #

@inline function (ctx::RegenerateContext)(fn::typeof(factor), arg::A) where A <: Number
    ctx.score += arg
    ctx.weight += arg
    return arg
end


using Distributed
using Dates
# addprocs(6) # Call instead with: julia -O3 -p18
using RoME
using RoMEPlotting
using Gadfly
@everywhere using RoME, RoMEPlotting, Gadfly

# Make sure plots look a bit nicer.
latex_fonts = Theme(major_label_font="CMU Serif", major_label_font_size=16pt,
                    minor_label_font="CMU Serif", minor_label_font_size=14pt,
                    key_title_font="CMU Serif", key_title_font_size=12pt,
                    key_label_font="CMU Serif", key_label_font_size=10pt)
Gadfly.push_theme(latex_fonts)

# Parse the arguments.
qfl_length = parse(Int, ARGS[1])

# Let's load the Manhattan scenario using the g2o file.
file = (normpath(Base.find_package("RoME"), "../..", "examples", "manhattan_incremental.g2o"))
global instructions = importG2o(file)

# Let's also get the file for batch factor graph at step 500.
fg_name = "fg-after-solve.tar.gz"
global fg_file = (normpath(Base.find_package("RoME"), "../..", "examples", fg_name))

function go_fixedlag_frombatch(qfl_length_arg::Integer)
    # Choose where to save the step's data.
    qfl_length = qfl_length_arg # Fixed lag window size.
    data_logpath = ENV["HOME"]*"/Documents/wafr/manhattan-frombatch-b$(qfl_length)-$(now())"

    # Instead of creating new graph, load the batch one.
    fg = LightDFG{SolverParams}(solverParams=SolverParams(logpath=data_logpath))
    loadDFG(fg_file, Main, fg)
    tree = emptyBayesTree()

    # Set up the fixed lag smoothing.
    getSolverParams(fg).logpath = data_logpath
    getSolverParams(fg).isfixedlag = true
    getSolverParams(fg).qfl = qfl_length
    getSolverParams(fg).limitfixeddown = true
    getSolverParams(fg).dbg = true
    # CHECK I think we need to manually set all variables to not solvable.
    for vsym in ls(fg)
        getVariable(fg, vsym).solvable = 0
    end
    # But release the last 10 poses.
    activevars = []
    for i in 350:360
        push!(activevars, Symbol("x", i))
        getVariable(fg, Symbol("x", i)).solvable = 1
    end
    # CHECK As well as all of the factors.
    for fsym in lsf(fg)
        getFactor(fg, fsym).solvable = 0
    end
    # And activate the factors corresponding to the active variables.
    sym_act_vars = Array{Symbol, 1}(activevars)
    for fsym in getFactorsAmongVariablesOnly(fg, sym_act_vars)
        getFactor(fg, fsym).solvable = 1
    end

    # Add the next---or initial offset of---measurements to the graph.
    padded_step = lpad(501, 4, "0")
    parseG2oInstruction!(fg, instructions[501])

    # Solve the graph, and save a copy of the tree.
    saveDFG(fg, "$(getLogPath(fg))/fg-before-solve$(padded_step)")
    getSolverParams(fg).maxincidence = 1000
    tree = solveTree!(fg)
    saveDFG(fg, "$(getLogPath(fg))/fg-after-solve$(padded_step)")
    saveTree(tree, "$(getLogPath(fg))/tree$(padded_step).jld2")
    drawTree(tree, show=false, filepath="$(getLogPath(fg))/bt$(padded_step).pdf")

    # Analyze clique counts.
    fid = open("$(getLogPath(fg))/clique-counts.txt", "w")
    nCliqs, nMarg, nReused, nBoth = calcCliquesRecycled(tree)
    println(fid, "$(padded_step), $(nCliqs), $(nMarg), $(nReused), $(nBoth)")

    # Just store some quick plots.
    pl1 = drawPoses(fg, spscale=0.6)
    Gadfly.draw(PDF("$(getLogPath(fg))/poses$(padded_step).pdf", 20cm, 10cm), pl1)

    # Solver stride.
    solveStride = 0
    # Run the loop for the remaining time steps.
    for i in 502:5453
        # Add the next measurement to the graph.
        parseG2oInstruction!(fg, instructions[i])
        padded_step = lpad(i, 4, "0")

        # Store each graph.
        saveDFG(fg, "$(getLogPath(fg))/fg-before-solve$(padded_step)")

        # Just store some quick plots, on another process
        remotecall((fgl, padded_stepl) -> begin
          @info "drawPoses, $(padded_stepl), for fg num variables=$(length(ls(fgl)))."
          pl1 = plotSLAM2DPoses(fgl, dyadScale=0.6, lbls=false)
          pl1 |> PDF("$(getLogPath(fgl))/poses$(padded_stepl).pdf", 20cm, 10cm)
        end, rand(Categorical(nprocs()-1))+1, fg, padded_step)

        # Only solve every 10th instruction.
        solveStride += 1
        if solveStride % 10 != 0
            @info "solveStride=$solveStride"
            continue
        end
        @info "Going for solve"

        # Solve the graph, and save a copy of the tree.
        getSolverParams(fg).maxincidence = 1000
        tree = solveTree!(fg, tree)
        saveDFG(fg, "$(getLogPath(fg))/fg-after-solve$(padded_step)")
        saveTree(tree, "$(getLogPath(fg))/tree$(padded_step).jld2")
        drawTree(tree, show=false, filepath="$(getLogPath(fg))/bt$(padded_step).pdf")

        # Analyze clique number.
        nCliqs, nMarg, nReused, nBoth = calcCliquesRecycled(tree)
        println(fid, "$(padded_step), $(nCliqs), $(nMarg), $(nReused), $(nBoth)")
        flush(fid)

        # Run the garbage collector.
        GC.gc()
    end
    # Final plot.
    padded_step = lpad(final_timestep+1, 4, "0")
    pl1 = drawPoses(fg, spscale=0.6)
    pl1 |> PDF("$(getLogPath(fg))/poses$(padded_step).pdf", 20cm, 10cm)

    close(fid)
end

# Run within a function to avoid undefined variable errors.
go_fixedlag_frombatch(qfl_length)


using DataStructures
using Random

#############
# Trimmed SDD types and structs
#############

"Root of the trimmed SDD manager node hierarchy"
abstract type TrimMgrNode <: SddMgrNode end

# alias structured logical nodes with a trimmed sdd manager vtree
const TrimTrue = SddTrueNode{TrimMgrNode}
const trimtrue = TrimTrue()
const TrimFalse = SddFalseNode{TrimMgrNode}
const trimfalse = TrimFalse()

const TrimConstant = SddConstantNode{TrimMgrNode}
const Trim⋁ = Sdd⋁Node{TrimMgrNode}
const Trim⋀ = Sdd⋀Node{TrimMgrNode}
const TrimNode = SddNode{<:TrimMgrNode} # would this be better?: Union{TrimTrue,TrimFalse,TrimConstant,Trim⋁,Trim⋀}

# alias SDD terminology
const Element = Tuple{TrimNode,TrimNode}
Element(prime::TrimNode, sub::TrimNode)::Element = (prime, sub)
const XYPartition = Set{Element}

const Unique⋁Cache = Dict{XYPartition,Trim⋁}
const ApplyCache = Dict{Tuple{TrimNode,TrimNode},TrimNode}

mutable struct TrimSddMgrInnerNode <: TrimMgrNode

    left::TrimMgrNode
    right::TrimMgrNode

    parent::Union{TrimSddMgrInnerNode, Nothing}

    variables::BitSet
    unique⋁cache::Unique⋁Cache

    conjoin_cache::ApplyCache
    disjoin_cache::ApplyCache

    TrimSddMgrInnerNode(left::TrimMgrNode, right::TrimMgrNode) = begin
        # @assert isempty(intersect(variables(left), variables(right)))
        this = new(left, right,
            nothing,
            union(variables(left), variables(right)),
            Unique⋁Cache(), ApplyCache(), ApplyCache())
        left.parent = this
        right.parent = this
        this
    end

end

mutable struct TrimSddMgrLeafNode <: TrimMgrNode

    var::Var
    parent::Union{TrimSddMgrInnerNode, Nothing}

    positive_literal::SddLiteralNode{TrimSddMgrLeafNode} # aka TrimLiteral
    negative_literal::SddLiteralNode{TrimSddMgrLeafNode} # aka TrimLiteral

    TrimSddMgrLeafNode(v::Var) = begin
        this = new(v, nothing)
        this.positive_literal = SddLiteralNode(var2lit(v), this)
        this.negative_literal = SddLiteralNode(-var2lit(v), this)
        this
    end

end

const TrimLiteral = SddLiteralNode{TrimSddMgrLeafNode}

const TrimSddMgr = AbstractVector{<:TrimMgrNode}
const TrimSdd = AbstractVector{<:TrimNode}

Base.eltype(::Type{TrimSdd}) = TrimNode

#####################
# Constructor
#####################

TrimMgrNode(v::Var) = TrimSddMgrLeafNode(v)
TrimMgrNode(left::TrimMgrNode, right::TrimMgrNode) = TrimSddMgrInnerNode(left, right)

#####################
# Traits
#####################

@inline NodeType(::TrimSddMgrLeafNode) = Leaf()
@inline NodeType(::TrimSddMgrInnerNode) = Inner()

#####################
# Methods
#####################

@inline children(n::TrimSddMgrInnerNode) = [n.left, n.right]

@inline variables(n::TrimSddMgrLeafNode)::BitSet = BitSet(n.var)
@inline variables(n::TrimSddMgrInnerNode)::BitSet = n.variables

import ..Utils: parent, descends_from, lca # make available for extension

@inline prime(e::Element) = e[1]
@inline sub(e::Element) = e[2]

@inline parent(n::TrimMgrNode)::Union{TrimSddMgrInnerNode, Nothing} = n.parent

@inline pointer_sort(s,t) = (hash(s) <= hash(t)) ? (s,t) : (t,s)

#TODO replace this by a bitset subset check on the set of variables
@inline descends_from(n::TrimNode, m::TrimNode) = descends_from(vtree(n), vtree(m))
@inline descends_from(::TrimMgrNode, ::TrimSddMgrLeafNode) = false
@inline descends_from(n::TrimMgrNode, m::TrimSddMgrInnerNode) = variables(n) ⊆ variables(m)

@inline descends_left_from(n::TrimNode, m::TrimNode)::Bool = descends_left_from(vtree(n), vtree(m))
@inline descends_left_from(n::TrimMgrNode, m::TrimSddMgrInnerNode)::Bool = (n === m.left) || descends_from(n, m.left)
@inline descends_left_from(::TrimMgrNode, ::TrimSddMgrLeafNode)::Bool = false

@inline descends_right_from(n::TrimNode, m::TrimNode)::Bool = descends_right_from(vtree(n), vtree(m))
@inline descends_right_from(n::TrimMgrNode, m::TrimSddMgrInnerNode)::Bool = (n === m.right) || descends_from(n,m.right)
@inline descends_right_from(::TrimMgrNode, ::TrimSddMgrLeafNode)::Bool = false

function lca(xy::XYPartition)::TrimSddMgrInnerNode
    # @assert !isempty(xy)
    # @assert all(e -> (prime(e) !== trimfalse), xy)
    element_vtrees = [parentlca(prime(e),sub(e)) for e in xy]
    return lca(element_vtrees...)
end

parentlca(p::TrimNode, s::TrimNode)::TrimSddMgrInnerNode =
    lca(parent(vtree(p)), parent(vtree(s)))
parentlca(p::TrimNode, ::TrimConstant)::TrimSddMgrInnerNode =
    parent(vtree(p))
parentlca(::TrimConstant, s::TrimNode)::TrimSddMgrInnerNode =
    parent(vtree(s))
parentlca(p::TrimConstant, s::TrimConstant)::TrimSddMgrInnerNode =
    error("This XY partition should have been trimmed to remove ($p,$s)!")

"""
Get the canonical compilation of the given XY Partition
"""
function canonicalize(xy::XYPartition)::TrimNode
    # @assert !isempty(xy)
    return canonicalize_compressed(compress(xy))
end

"""
Compress a given XY Partition (merge elements with identical subs)
"""
function compress(xy::XYPartition)::XYPartition
    # @assert !isempty(xy)
    sub2elems = groupby(e -> sub(e), xy)
    #TODO avoid making a new partition if existing one is unchanged
    compressed_elements = XYPartition()
    for (subnode,elements) in sub2elems
        primenode = mapreduce(e -> prime(e), (p1, p2) -> disjoin(p1, p2), elements)
        push!(compressed_elements, (primenode, subnode))
    end
    return compressed_elements
end

"""
Get the canonical compilation of the given compressed XY Partition
"""
function canonicalize_compressed(xy::XYPartition)::TrimNode
    # @assert !isempty(xy)
    # trim
    if length(xy) == 1 && (prime(first(xy)) === trimtrue)
        return sub(first(xy))
    elseif length(xy) == 2
        l = [xy...]
        if (sub(l[1]) === trimtrue) && (sub(l[2]) === trimfalse)
            return prime(l[1])
        elseif (sub(l[2]) === trimtrue) && (sub(l[1]) === trimfalse)
            return prime(l[2])
        end
    end
    # get unique node representation
    return unique⋁(xy)
end

@inline function remove_false_primes(xy)
    return filter(e -> (prime(e) !== trimfalse), xy)
end

"""
Construct a unique decision gate for the given vtree
"""
function unique⋁(xy::XYPartition, mgr::TrimSddMgrInnerNode = lca(xy))::Trim⋁
    #TODO add finalization trigger to remove from the cache when the node is gc'ed + weak value reference
    get!(mgr.unique⋁cache, xy) do
        node = Trim⋁(xy2ands(xy, mgr), mgr)
        not_xy = negate(xy)
        not_node = Trim⋁(xy2ands(not_xy, mgr), mgr, node)
        node.negation = not_node
        mgr.unique⋁cache[not_xy] = not_node
        node
    end
end

@inline negate(xy::XYPartition) = XYPartition([Element(prime(e), !sub(e)) for e in xy])
@inline xy2ands(xy::XYPartition, mgr::TrimSddMgrInnerNode) = [Trim⋀(prime(e), sub(e), mgr) for e in xy]

"""
Compile a given variable, literal, or constant
"""
compile(mgr::TrimSddMgr, x::Union{Var,Lit})::TrimLiteral = compile(mgr[end], x)

function compile(n::TrimMgrNode, v::Var)::TrimLiteral
    compile(n,var2lit(v))
end

function compile(n::TrimSddMgrLeafNode, l::Lit)::TrimLiteral
    # @assert n.var == lit2var(l)
    if l>0 # positive literal
        n.positive_literal
    else
        n.negative_literal
    end
end

function compile(n::TrimSddMgrInnerNode, l::Lit)::TrimLiteral
    if lit2var(l) in variables(n.left)
        compile(n.left, l)
    elseif lit2var(l) in variables(n.right)
        compile(n.right, l)
    else
        error("$v is not contained in this vtree")
    end
end

function compile(constant::Bool)::TrimConstant
    if constant == true
        trimtrue
    else
        trimfalse
    end
end

"""
Conjoin two SDDs
"""
@inline conjoin(::TrimFalse, ::TrimTrue)::TrimFalse = trimfalse
@inline conjoin(::TrimTrue, ::TrimFalse)::TrimFalse = trimfalse
@inline conjoin(s::TrimNode, ::TrimTrue)::TrimNode = s
@inline conjoin(::TrimNode, ::TrimFalse)::TrimFalse = trimfalse
@inline conjoin(::TrimTrue, s::TrimNode)::TrimNode = s
@inline conjoin(::TrimFalse, ::TrimNode)::TrimFalse = trimfalse
@inline conjoin(::TrimTrue, s::TrimTrue)::TrimNode = trimtrue
@inline conjoin(::TrimFalse, s::TrimFalse)::TrimNode = trimfalse

function conjoin(s::TrimLiteral, t::TrimLiteral)::TrimNode
    if vtree(s) == vtree(t)
        (s === t) ? s : trimfalse
    else
        conjoin_indep(s,t)
    end
end

function conjoin(s::TrimNode, t::TrimNode)::TrimNode
    if vtree(s) == vtree(t)
        conjoin_cartesian(t,s)
    elseif descends_from(s,t)
        conjoin_descendent(s,t)
    elseif descends_from(t,s)
        conjoin_descendent(t,s)
    else
        conjoin_indep(s,t)
    end
end

"""
Conjoin two SDDs when they respect the same vtree node
"""
function conjoin_cartesian(n1::TrimNode, n2::TrimNode)::TrimNode
    if n1 === n2
        return n1
    elseif n1 == !n2
        return trimfalse
    end
    (n1,n2) = pointer_sort(n1,n2)

    get!(vtree(n1).conjoin_cache, (n1,n2)) do
        elems_prod = Vector{Element}()
        elems1 = copy(children(n1))
        elems2 = copy(children(n2))
        for e1 in elems1
            for e2 in elems2
                if prime(e1) === prime(e2)
                    push!(elems_prod, Element(prime(e1), conjoin(sub(e1),sub(e2))))
                    filter!(e -> prime(e) !== prime(e1), elems1)
                    filter!(e -> prime(e) !== prime(e1), elems2)
                    break # go to next e1
                elseif prime(e1) === !prime(e2)
                    for e3 in elems2
                        if e3 !== e2
                            push!(elems_prod, Element(prime(e3), conjoin(sub(e3),sub(e1))))
                        end
                    end
                    for e4 in elems1
                        if e4 !== e1
                            push!(elems_prod, Element(prime(e4), conjoin(sub(e4),sub(e2))))
                        end
                    end
                    filter!(e -> prime(e) !== prime(e1), elems1)
                    filter!(e -> prime(e) !== prime(e2), elems2)
                    break # go to next e1
                end
            end
        end
        product = vec([(e1,e2) for e1 in elems1, e2 in elems2])
        #TODO sort product by probability of subsumes
        while !isempty(product)
            (e1, e2) = pop!(product)
            newprime = conjoin(prime(e1),prime(e2))
            if newprime != trimfalse
                push!(elems_prod, Element(newprime, conjoin(sub(e1),sub(e2))))
            end
            if newprime === prime(e1)
                # p1 |= p2 and therefore p1 will be mutex with all other p2-primes
                filter!(p -> prime(p[1]) !== prime(e1), product)
            elseif newprime === prime(e2)
                # p2 |= p1 and therefore p2 will be mutex with all other p1-primes
                filter!(p -> prime(p[2]) !== prime(e2), product)
            end
        end
        canonicalize(XYPartition(elems_prod))
    end
end

"""
Conjoin two SDDs when one descends from the other
"""
function conjoin_descendent(d::TrimNode, n::TrimNode)::TrimNode
    get!(vtree(n).conjoin_cache, (d,n)) do
        if descends_left_from(d, n)
            elements = Element[Element(conjoin(prime(e),d), sub(e)) for e in children(n)]
            elements = remove_false_primes(elements)
            push!(elements, Element(!d, trimfalse))
        else
            # @assert descends_right_from(d, n)
            elements = Element[Element(prime(e),conjoin(sub(e),d)) for e in children(n)]
        end
        #TODO are there cases where we don't need all of compress-trim-unique?
        canonicalize(XYPartition(elements))
    end
end

"""
Conjoin two SDDs in separate parts of the vtree
"""
function conjoin_indep(s::TrimNode, t::TrimNode)::Trim⋁
    # @assert GateType(s)!=ConstantGate() && GateType(t)!=ConstantGate()
    mgr = parentlca(s,t)
    # @assert vtree(s) != mgr && vtree(t) != mgr
    (s,t) = pointer_sort(s,t)
    get!(mgr.conjoin_cache, (s,t)) do
        if descends_left_from(vtree(s), mgr)
            # @assert descends_right_from(vtree(t), mgr)
            elements = Element[Element(s,t),Element(!s,trimfalse)]
        else
            # @assert descends_left_from(vtree(t), mgr)
            # @assert descends_right_from(vtree(s), mgr)
            elements = Element[Element(t,s),Element(!t,trimfalse)]
        end
        # TODO: the XY partition must already be compressed and trimmed
        unique⋁(XYPartition(elements), mgr)
    end
end

@inline Base.:&(s,t) = conjoin(s,t)

"""
Disjoin two SDDs
"""
@inline disjoin(::TrimFalse, ::TrimTrue)::TrimTrue = trimtrue
@inline disjoin(::TrimTrue, ::TrimFalse)::TrimTrue = trimtrue
@inline disjoin(::TrimNode, ::TrimTrue)::TrimTrue = trimtrue
@inline disjoin(s::TrimNode, ::TrimFalse)::TrimNode = s
@inline disjoin(::TrimTrue, ::TrimNode)::TrimTrue = trimtrue
@inline disjoin(::TrimFalse, s::TrimNode)::TrimNode = s
@inline disjoin(::TrimTrue, s::TrimTrue)::TrimNode = trimtrue
@inline disjoin(::TrimFalse, s::TrimFalse)::TrimNode = trimfalse

function disjoin(s::TrimLiteral, t::TrimLiteral)::TrimNode
    if vtree(s) == vtree(t)
        (s === t) ? s : trimtrue
    else
        disjoin_indep(s,t)
    end
end

function disjoin(s::TrimNode, t::TrimNode)::TrimNode
    if vtree(s) == vtree(t)
        disjoin_cartesian(t,s)
    elseif descends_from(s,t)
        disjoin_descendent(s,t)
    elseif descends_from(t,s)
        disjoin_descendent(t,s)
    else
        disjoin_indep(s,t)
    end
end

"""
Disjoin two SDDs when they respect the same vtree node
"""
function disjoin_cartesian(n1::TrimNode, n2::TrimNode)::TrimNode
    if n1 === n2
        return n1
    elseif n1 == !n2
        return trimtrue
    end
    (n1,n2) = pointer_sort(n1,n2)

    get!(vtree(n1).disjoin_cache, (n1,n2)) do
        elems_prod = Vector{Element}()
        elems1 = copy(children(n1))
        elems2 = copy(children(n2))
        for e1 in elems1
            for e2 in elems2
                if prime(e1) === prime(e2)
                    push!(elems_prod, Element(prime(e1), disjoin(sub(e1),sub(e2))))
                    filter!(e -> prime(e) !== prime(e1), elems1)
                    filter!(e -> prime(e) !== prime(e2), elems2)
                    break # go to next e1
                elseif prime(e1) === !prime(e2)
                    for e3 in elems2
                        if e3 !== e2
                            push!(elems_prod, Element(prime(e3), disjoin(sub(e3),sub(e1))))
                        end
                    end
                    for e4 in elems1
                        if e4 !== e1
                            push!(elems_prod, Element(prime(e4), disjoin(sub(e4),sub(e2))))
                        end
                    end
                    filter!(e -> prime(e) !== prime(e1), elems1)
                    filter!(e -> prime(e) !== prime(e2), elems2)
                    break # go to next e1
                end
            end
        end
        product = vec([(e1,e2) for e1 in elems1, e2 in elems2])
        while !isempty(product)
            (e1, e2) = pop!(product)
            newprime = conjoin(prime(e1),prime(e2))
            if newprime != trimfalse
                push!(elems_prod, Element(newprime, disjoin(sub(e1),sub(e2))))
            end
            if newprime === prime(e1)
                # p1 |= p2 and therefore p1 will be mutex with all other p2-primes
                filter!(p -> prime(p[1]) !== prime(e1), product)
            elseif newprime === prime(e2)
                # p2 |= p1 and therefore p2 will be mutex with all other p1-primes
                filter!(p -> prime(p[2]) !== prime(e2), product)
            end
        end
        canonicalize(XYPartition(elems_prod))
    end
end

"""
Disjoin two SDDs when one descends from the other
"""
function disjoin_descendent(d::TrimNode, n::TrimNode)::TrimNode
    get!(vtree(n).disjoin_cache, (d,n)) do
        if descends_left_from(d, n)
            not_d = !d
            elements = Element[Element(conjoin(prime(e),not_d), sub(e)) for e in children(n)]
            elements = remove_false_primes(elements)
            push!(elements,Element(d, trimtrue))
        else
            # @assert descends_right_from(d, n)
            elements = Element[Element(prime(e),disjoin(sub(e),d)) for e in children(n)]
        end
        #TODO are there cases where we don't need all of compress-trim-unique?
        canonicalize(XYPartition(elements))
    end
end

"""
Disjoin two SDDs in separate parts of the vtree
"""
function disjoin_indep(s::TrimNode, t::TrimNode)::Trim⋁
    # @assert GateType(s)!=ConstantGate() && GateType(t)!=ConstantGate()
    mgr = parentlca(s,t)
    # @assert vtree(s) != mgr && vtree(t) != mgr
    (s,t) = pointer_sort(s,t)
    get!(mgr.disjoin_cache, (s,t)) do
        if descends_left_from(vtree(s), mgr)
            # @assert descends_right_from(vtree(t), mgr)
            elements = Element[Element(s,trimtrue),Element(!s,t)]
        else
            # @assert descends_left_from(vtree(t), mgr)
            # @assert descends_right_from(vtree(s), mgr)
            elements = Element[Element(t,trimtrue),Element(!t,s)]
        end
        # TODO: the XY partition must already be compressed and trimmed
        unique⋁(XYPartition(elements), mgr)
    end
end

@inline Base.:|(s,t) = disjoin(s,t)

"""
Negate an SDD
"""
@inline negate(::TrimFalse)::TrimTrue = trimtrue
@inline negate(::TrimTrue)::TrimFalse = trimfalse

function negate(s::TrimLiteral)::TrimLiteral
    if positive(s)
        vtree(s).negative_literal
    else
        vtree(s).positive_literal
    end
end

negate(node::Trim⋁)::Trim⋁ = node.negation

@inline Base.:!(s) = negate(s)


using LinearAlgebra
using BenchmarkTools
using Keldysh

suite = BenchmarkGroup()

function setup_indexing_benchmark()
  suite["indexing"] = BenchmarkGroup()

  c = FullContour(tmax=2.0, β=5.0)
  grid = FullTimeGrid(c, 41, 101)
  N = length(grid)

  X = zeros(1,1,N,N)
  Y = Keldysh.GenericStorage(N,N,1,true)

  G1 = GenericTimeGF(grid, 1, true)
  G2 = FullTimeGF(grid, 1, fermionic, true)

  b = @benchmarkable $X[1,1,i,j] setup=(i=rand(1:$N); j=rand(1:$N))
  suite["indexing"]["Array"] = b

  b = @benchmarkable $Y[i,j] setup=(i=rand(1:$N); j=rand(1:$N))
  suite["indexing"]["GenericStorage"] = b

  b = @benchmarkable $G1[t1,t2] setup=(t1=$grid[rand(1:$N)]; t2=$grid[rand(1:$N)])
  suite["indexing"]["GenericTimeGF"] = b

  b = @benchmarkable $G2[t1,t2] setup=(t1=$grid[rand(1:$N)]; t2=$grid[rand(1:$N)])
  suite["indexing"]["FullTimeGF"] = b

  tune!(suite["indexing"])
end

function setup_interp_benchmark()
  suite["interp"] = BenchmarkGroup()

  c = FullContour(tmax=2.0, β=5.0)
  l = length(c)
  grid = FullTimeGrid(c, 41, 101)

  G1 = GenericTimeGF(grid, 1, true)

  b = @benchmarkable $G1(t1,t2) setup=(t1=Keldysh.get_point($c, $l*rand()); t2=Keldysh.get_point($c, $l*rand()))
  suite["interp"]["GenericTimeGF"] = b

  tune!(suite["interp"])
end

function main()
  setup_indexing_benchmark()
  setup_interp_benchmark()

  compare = false
  if length(ARGS) > 0

    if length(ARGS) != 2
      error("usage: julia test/benchmark.jl test/params.ref.json test/results.ref.json")
    end

    params_file = ARGS[1]
    p = BenchmarkTools.load(params_file)[1]
    loadparams!(suite, p, :evals, :samples)

    ref_file = ARGS[2]
    ref = BenchmarkTools.load(ref_file)[1]
    compare = true
  end

  results = run(suite, verbose=true)
  display(median(results))
  println()

  BenchmarkTools.save("params.json", params(suite))
  BenchmarkTools.save("results.json", results)

  if compare
    cmp = judge(median(results), median(ref))
    display(cmp)
  end
end

main()


#=
Final parsing of preparsing structures

Numbers and variables will finaly
=#
include("ofMeshDictPreParser.jl")

abstract type OfEntity end
abstract type AbstractOfVar <: OfEntity end
abstract type AbstractOfVal <: OfEntity end

struct OfDictionary <: OfEntity
    name::String
    entries::Vector{OfEntity}
end

# array list of same entities with length
struct OfArray <: OfEntity
    length::Int64
    entries::Vector{OfEntity}
end

struct OfList<: OfEntity
    length::Int64
    entries::Vector{OfEntity}
end

struct OfVariable{T} <: AbstractOfVar
    name::String
    value::T
end

struct OfValue{T} <: AbstractOfVal
    value::T
end

"""
    dfs_parse(parse_entities::Vector{PreParsedEntity}, tokens::Vector{Token})

parses pre parsed vector structure to finally parsed file structure
"""
function dfs_parse(parse_entities::Vector{PreParsedEntity}, tokens::Vector{Token})

    of_entities = Vector{OfEntity}(undef, 0)

    for parse_elem in parse_entities
        if typeof(parse_elem) == PreParsedVariable
            push!(of_entities, parse_of_variable(parse_elem, tokens))
        elseif typeof(parse_elem) == PreParsedFNum
            push!(of_entities, OfValue(parse(Float64,parse_elem.string)))
        elseif typeof(parse_elem) == PreParsedINum
            push!(of_entities, OfValue(parse(Int64,parse_elem.string)))
        elseif typeof(parse_elem) == PreParsedDictonary
            mypush!(of_entities, OfDictionary(parse_elem.name,
                        dfs_parse(parse_elem.entries, tokens)))
        elseif typeof(parse_elem) == PreParsedArray
            mypush!(of_entities, OfArray(parse_elem.length,
                        dfs_parse(parse_elem.entries, tokens)))
        elseif typeof(parse_elem) == PreParsedList
            parsed_list = dfs_parse(parse_elem.entries, tokens)
            mypush!(of_entities, OfList(length(parsed_list), parsed_list))
        end
    end

    return of_entities
end

function mypush!(of_entities::Vector{OfEntity}, OfEntity::T) where T <: OfEntity
    push!(of_entities, OfEntity)
end

function mypush!(of_entities::Vector{OfEntity},
                    entities::Vector{T}) where T <: OfEntity
    for e in entities
        push!(of_entities, e)
    end
end

"""
    parse_of_variable(pe::PreParsedEntity, tokens::Vector{Token})

Parsing of a PreParsedEntity to final type OfVariable, which completes parsing.
"""
# TODO: only works for numbers and strings by now
function parse_of_variable(pe::PreParsedEntity, tokens::Vector{Token})::OfVariable

    name = ""
    value = []

    if tokens[pe.id_token_from].type_id == TokenTypeID_Dict[:word]
        name = tokens[pe.id_token_from].string

        if ( (pe.id_token_to - pe.id_token_from) == 2 )
            if  (tokens[pe.id_token_from+1].type_id ==
                 TokenTypeID_Dict[:int_number])

                value = parse(Int64, tokens[pe.id_token_from+1].string)

            elseif (tokens[pe.id_token_from+1].type_id ==
                    TokenTypeID_Dict[:float_number])

                value = parse(Float64, tokens[pe.id_token_from+1].string)

            elseif (tokens[pe.id_token_from+1].type_id
                    == TokenTypeID_Dict[:word])

                value = tokens[pe.id_token_from+1].string

            elseif (tokens[pe.id_token_from+1].type_id
                == TokenTypeID_Dict[:string])

                value = tokens[pe.id_token_from+1].string

            else
                value = nothing
            end
        elseif ( ((pe.id_token_to - pe.id_token_from) == 3) &&
                 (tokens[pe.id_token_from+1].type_id ==
                        TokenTypeID_Dict[:minus]) )

            if  (tokens[pe.id_token_from+2].type_id ==
                TokenTypeID_Dict[:int_number])

                value = parse(Int64, tokens[pe.id_token_from+1].string *
                                        tokens[pe.id_token_from+2].string)

            elseif (tokens[pe.id_token_from+2].type_id ==
                    TokenTypeID_Dict[:float_number])

                 value = parse(Float64, tokens[pe.id_token_from+1].string *
                                        tokens[pe.id_token_from+2].string)
            else
                value =  nothing
            end

        else
            value = 0
        end
    else
        error("First Token of Variable should be a word!")
    end

    return OfVariable(name, value)
end

"""
    ofprint(ofvar::OfEntity, depth::Int64=0)

dfs print in case for single OfEntity
"""
function ofprint(ofvar::OfEntity, depth::Int64=0)
    ofprint([ofvar], depth)
end

"""
    ofprint(ofvars::Vector{OfEntity}, depth::Int64=0)

dfs print of Vector{OfEntity} to get a nicer view of finally parsed entities
"""
function ofprint(ofvars::Vector{OfEntity}, depth::Int64=0)
    for ofvar in ofvars
        if typeof(ofvar) <: AbstractOfVar
            println("\t"^depth,ofvar.name, ": ", ofvar.value)

        elseif typeof(ofvar) <: AbstractOfVal
            print(ofvar.value, " ")

        elseif typeof(ofvar) == OfDictionary
            println()
            println("\t"^depth,ofvar.name, " {")
            depth += 1
            ofprint(ofvar.entries, depth)
            depth -= 1
            println("\t"^depth,"}")
        elseif typeof(ofvar) == OfArray
            print("\t"^depth,ofvar.length, " (")
            depth += 1
            ofprint(ofvar.entries,depth)
            print(")\n")
            depth -= 1
        elseif typeof(ofvar) == OfList
            print("\t"^depth,"(")
            ofprint(ofvar.entries, depth)
            print(")\n")
        end
    end

end

"""
    get_dictionary(ofvars::Vector{OfEntity}, dict_name::string)

Returns Dictionary in first layer of ofvars with .name == dict_name
"""
function get_dictionary(ofvars::Vector{OfEntity}, dict_name::String)

    for f in ofvars
        if typeof(f) == OfDictionary && f.name == dict_name
            return f
        end
    end

    return nothing
end

"""
    get_variable(ofvars::Vector{OfEntity}, var_name::string)

Returns Variable in first layer of ofvars with .name == var_name
"""
function get_variable(ofvars::Vector{OfEntity}, var_name::String)

    for f in ofvars
        if typeof(f) <: AbstractOfVar && f.name == var_name
            return f
        end
    end

    return nothing
end

"""
    get_variable_value(ofvars::Vector{OfEntity}, var_name::string)

Returns Variable value in first layer of ofvars with .name == var_name
"""
function get_variable_value(ofvars::Vector{OfEntity}, var_name::String)

    var = get_variable(ofvars::Vector{OfEntity}, var_name::String)

    if !isnothing(var)
        return var.value
    else
        error("Variable ", var_name, " cannot be found")
    end

end


module SystemStructures

using DataStructures
using Symbolics: linear_expansion, unwrap
using SymbolicUtils: istree, operation, arguments, Symbolic
using SymbolicUtils: quick_cancel, similarterm
using ..ModelingToolkit
import ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!, flatten,
    value, InvalidSystemException, isdifferential, _iszero, isparameter,
    independent_variables, isinput, SparseMatrixCLIL
using ..BipartiteGraphs
import ..BipartiteGraphs: invview, complete
using Graphs
using UnPack
using Setfield
using SparseArrays

quick_cancel_expr(expr) = Rewriters.Postwalk(
    quick_cancel,
    similarterm=(x, f, args; kws...) -> similarterm(
        x, f, args, SymbolicUtils.symtype(x);
        metadata=SymbolicUtils.metadata(x), kws...
    )
)(expr)

#=
When we don't do subsitution, variable information is split into two different
places, i.e. `states` and the right-hand-side of `observed`.

eqs = [0 ~ z + x; 0 ~ y + z^2]
states = [y, z]
observed = [x ~ sin(y) + z]
struct Reduced
    var
    expr
    idxs
end
fullvars = [Reduced(x, sin(y) + z, [2, 3]), y, z]
active_𝑣vertices = [false, true, true]
      x   y   z
eq1:  1       1
eq2:      1   1

      x   y   z
eq1:      1   1
eq2:      1   1

for v in 𝑣vertices(graph); active_𝑣vertices[v] || continue

end
=#

export SystemStructure, SystemPartition
export initialize_system_structure, find_linear_equations, linear_subsys_adjmat
export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq
export dervars_range, diffvars_range, algvars_range
export DiffGraph

@enum VariableType::Int8 DIFFERENTIAL_VARIABLE ALGEBRAIC_VARIABLE DERIVATIVE_VARIABLE

Base.@kwdef struct SystemPartition
    e_solved::Vector{Int}
    v_solved::Vector{Int}
    e_residual::Vector{Int}
    v_residual::Vector{Int}
end

function Base.:(==)(s1::SystemPartition, s2::SystemPartition)
    tup1 = (s1.e_solved, s1.v_solved, s1.e_residual, s1.v_residual)
    tup2 = (s2.e_solved, s2.v_solved, s2.e_residual, s2.v_residual)
    tup1 == tup2
end

struct DiffGraph <: Graphs.AbstractGraph{Int}
    primal_to_diff::Vector{Union{Int, Nothing}}
    diff_to_primal::Union{Nothing, Vector{Union{Int, Nothing}}}
end
DiffGraph(n::Integer, with_badj::Bool=false) = DiffGraph(Union{Int, Nothing}[nothing for _=1:n],
    with_badj ? Union{Int, Nothing}[nothing for _=1:n] : nothing)

@noinline require_complete(dg::DiffGraph) = dg.diff_to_primal === nothing &&
    error("Not complete. Run `complete` first.")

Graphs.is_directed(dg::DiffGraph) = true
Graphs.edges(dg::DiffGraph) = (i => v for (i, v) in enumerate(dg.primal_to_diff) if v !== nothing)
Graphs.nv(dg::DiffGraph) = length(dg.primal_to_diff)
Graphs.ne(dg::DiffGraph) = count(x->x !== nothing, dg.primal_to_diff)
Graphs.vertices(dg::DiffGraph) = Base.OneTo(nv(dg))
function Graphs.outneighbors(dg::DiffGraph, var::Integer)
    diff = dg.primal_to_diff[var]
    return diff === nothing ? () : (diff,)
end
function Graphs.inneighbors(dg::DiffGraph, var::Integer)
    require_complete(dg)
    diff = dg.diff_to_primal[var]
    return diff === nothing ? () : (diff,)
end
function Graphs.add_vertex!(dg::DiffGraph)
    push!(dg.primal_to_diff, nothing)
    if dg.diff_to_primal !== nothing
        push!(dg.diff_to_primal, nothing)
    end
    return length(dg.primal_to_diff)
end

function Graphs.add_edge!(dg::DiffGraph, var::Integer, diff::Integer)
    dg[var] = diff
end

# Also pass through the array interface for ease of use
Base.:(==)(dg::DiffGraph, v::AbstractVector) = dg.primal_to_diff == v
Base.:(==)(dg::AbstractVector, v::DiffGraph) = v == dg.primal_to_diff
Base.eltype(::DiffGraph) = Union{Int, Nothing}
Base.size(dg::DiffGraph) = size(dg.primal_to_diff)
Base.length(dg::DiffGraph) = length(dg.primal_to_diff)
Base.getindex(dg::DiffGraph, var::Integer) = dg.primal_to_diff[var]
function Base.setindex!(dg::DiffGraph, val::Union{Integer, Nothing}, var::Integer)
    if dg.diff_to_primal !== nothing
        old_pd = dg.primal_to_diff[var]
        if old_pd !== nothing
            dg.diff_to_primal[old_pd] = nothing
        end
        if val !== nothing
            old_dp = dg.diff_to_primal[val]
            old_dp === nothing || error("Variable already assigned.")
            dg.diff_to_primal[val] = var
        end
    end
    return dg.primal_to_diff[var] = val
end
Base.iterate(dg::DiffGraph, state...) = iterate(dg.primal_to_diff, state...)

function complete(dg::DiffGraph)
    dg.diff_to_primal !== nothing && return dg
    diff_to_primal = zeros(Int, length(dg.primal_to_diff))
    for (var, diff) in edges(dg)
        diff_to_primal[diff] = var
    end
    return DiffGraph(dg.primal_to_diff, diff_to_primal)
end

function invview(dg::DiffGraph)
    require_complete(dg)
    return DiffGraph(dg.diff_to_primal, dg.primal_to_diff)
end

Base.@kwdef struct SystemStructure
    fullvars::Vector
    vartype::Vector{VariableType}
    # Maps the (index of) a variable to the (index of) the variable describing
    # its derivative.
    var_to_diff::DiffGraph
    algeqs::BitVector
    graph::BipartiteGraph{Int,Nothing}
    solvable_graph::BipartiteGraph{Int,Nothing}
    var_eq_matching::Matching
    scc::Vector{Vector{Int}}
    partitions::Vector{SystemPartition}
end

isdervar(s::SystemStructure, var::Integer) = s.vartype[var] === DERIVATIVE_VARIABLE
isdiffvar(s::SystemStructure, var::Integer) = s.vartype[var] === DIFFERENTIAL_VARIABLE
isalgvar(s::SystemStructure, var::Integer) = s.vartype[var] === ALGEBRAIC_VARIABLE

dervars_range(s::SystemStructure) = Iterators.filter(Base.Fix1(isdervar, s), eachindex(s.vartype))
diffvars_range(s::SystemStructure) = Iterators.filter(Base.Fix1(isdiffvar, s), eachindex(s.vartype))
algvars_range(s::SystemStructure) = Iterators.filter(Base.Fix1(isalgvar, s), eachindex(s.vartype))

isalgeq(s::SystemStructure, eq::Integer) = s.algeqs[eq]
isdiffeq(s::SystemStructure, eq::Integer) = !isalgeq(s, eq)

function initialize_system_structure(sys; quick_cancel=false)
    sys = flatten(sys)
    ivs = independent_variables(sys)
    eqs = copy(equations(sys))
    neqs = length(eqs)
    algeqs = trues(neqs)
    dervaridxs = OrderedSet{Int}()
    var2idx = Dict{Any,Int}()
    symbolic_incidence = []
    fullvars = []
    var_counter = Ref(0)
    addvar! = let fullvars=fullvars, var_counter=var_counter
        var -> begin
            get!(var2idx, var) do
                push!(fullvars, var)
                var_counter[] += 1
            end
        end
    end

    vars = OrderedSet()
    for (i, eq′) in enumerate(eqs)
        if _iszero(eq′.lhs)
            rhs = quick_cancel ? quick_cancel_expr(eq′.rhs) : eq′.rhs
            eq = eq′
        else
            lhs = quick_cancel ? quick_cancel_expr(eq′.lhs) : eq′.lhs
            rhs = quick_cancel ? quick_cancel_expr(eq′.rhs) : eq′.rhs
            eq = 0 ~ rhs - lhs
        end
        vars!(vars, eq.rhs)
        isalgeq = true
        statevars = []
        for var in vars
            any(isequal(var), ivs) && continue
            if isparameter(var) || (istree(var) && isparameter(operation(var)))
                continue
            end
            varidx = addvar!(var)
            push!(statevars, var)

            dvar = var
            idx = varidx
            while isdifferential(dvar)
                if !(idx in dervaridxs)
                    push!(dervaridxs, idx)
                end
                isalgeq = false
                dvar = arguments(dvar)[1]
                idx = addvar!(dvar)
            end
        end
        push!(symbolic_incidence, copy(statevars))
        empty!(statevars)
        empty!(vars)
        algeqs[i] = isalgeq
        if isalgeq
            eqs[i] = eq
        else
            eqs[i] = eqs[i].lhs ~ rhs
        end
    end

    # sort `fullvars` such that the mass matrix is as diagonal as possible.
    dervaridxs = collect(dervaridxs)
    sorted_fullvars = OrderedSet(fullvars[dervaridxs])
    for dervaridx in dervaridxs
        dervar = fullvars[dervaridx]
        diffvar = arguments(dervar)[1]
        if !(diffvar in sorted_fullvars)
            push!(sorted_fullvars, diffvar)
        end
    end
    for v in fullvars
        if !(v in sorted_fullvars)
            push!(sorted_fullvars, v)
        end
    end
    fullvars = collect(sorted_fullvars)
    var2idx = Dict(fullvars .=> eachindex(fullvars))
    dervaridxs = 1:length(dervaridxs)

    nvars = length(fullvars)
    diffvars = []
    vartype = fill(DIFFERENTIAL_VARIABLE, nvars)
    var_to_diff = DiffGraph(nvars, true)
    for dervaridx in dervaridxs
        vartype[dervaridx] = DERIVATIVE_VARIABLE
        dervar = fullvars[dervaridx]
        diffvar = arguments(dervar)[1]
        diffvaridx = var2idx[diffvar]
        push!(diffvars, diffvar)
        var_to_diff[diffvaridx] = dervaridx
    end

    algvars = setdiff(states(sys), diffvars)
    for algvar in algvars
        # it could be that a variable appeared in the states, but never appeared
        # in the equations.
        algvaridx = get(var2idx, algvar, 0)
        algvaridx == 0 && throw(InvalidSystemException("The system is missing "
            * "an equation for $algvar."
        ))
        vartype[algvaridx] = ALGEBRAIC_VARIABLE
    end

    graph = BipartiteGraph(neqs, nvars, Val(false))
    for (ie, vars) in enumerate(symbolic_incidence), v in vars
        jv = var2idx[v]
        add_edge!(graph, ie, jv)
    end

    @set! sys.eqs = eqs
    @set! sys.structure = SystemStructure(
        fullvars = fullvars,
        vartype = vartype,
        var_to_diff = var_to_diff,
        algeqs = algeqs,
        graph = graph,
        solvable_graph = BipartiteGraph(nsrcs(graph), ndsts(graph), Val(false)),
        var_eq_matching = Matching(ndsts(graph)),
        scc = Vector{Int}[],
        partitions = SystemPartition[],
    )
    return sys
end

function linear_subsys_adjmat(sys)
    s = structure(sys)
    @unpack fullvars, graph = s
    is_linear_equations = falses(nsrcs(graph))
    eqs = equations(sys)
    eadj = Vector{Int}[]
    cadj = Vector{Int}[]
    coeffs = Int[]
    for (i, eq) in enumerate(eqs); isdiffeq(eq) && continue
        empty!(coeffs)
        linear_term = 0
        all_int_vars = true

        term = value(eq.rhs - eq.lhs)
        for j in 𝑠neighbors(graph, i)
            var = fullvars[j]
            a, b, islinear = linear_expansion(term, var)
            a = unwrap(a)
            if islinear && !(a isa Symbolic) && a isa Number && !isinput(var)
                if a == 1 || a == -1
                    a = convert(Integer, a)
                    linear_term += a * var
                    push!(coeffs, a)
                else
                    all_int_vars = false
                end
            end
        end

        # Check if all states in the equation is both linear and homogeneous,
        # i.e. it is in the form of
        #
        #       ``∑ c_i * v_i = 0``,
        #
        # where ``c_i`` ∈ ℤ and ``v_i`` denotes states.
        if all_int_vars && isequal(linear_term, term)
            is_linear_equations[i] = true
            push!(eadj, copy(𝑠neighbors(graph, i)))
            push!(cadj, copy(coeffs))
        else
            is_linear_equations[i] = false
        end
    end

    linear_equations = findall(is_linear_equations)
    return SparseMatrixCLIL(nsrcs(graph),
                           ndsts(graph),
                           linear_equations, eadj, cadj)
end

function Base.show(io::IO, mime::MIME"text/plain", s::SystemStructure)
    @unpack graph = s
    S = incidence_matrix(graph, Num(Sym{Real}(:×)))
    print(io, "Incidence matrix:")
    show(io, mime, S)
end

end # module


using MLStyle

function allow_patterns(node)
    @match node begin
        :(case($(targets...)) do
            $(cases...)
        end) => let target = length(targets) === 1 ? targets[1] : Expr(:tuple, targets...),
                   cases = map(allow_patterns, cases)
                   quote
                     $MLStyle.@match $target begin
                         $(cases...)
                     end
                   end
               end
        Expr(hd, tl...) => Expr(hd, map(allow_patterns, tl)...)
        a => a
    end
end

macro allow_patterns(node)
    allow_patterns(node) |> esc
end

# @allow_patterns module SyntaxExtended
#      case(1) do
#          2        => error(1)
#          5:10     => error(2)
#          1:4 && a =>
#          case(a + 1) do
#              a => 5a
#          end
#      end |> println
# end


"""
this script fetches sequential pairs of poses, fetches the big data (radar pings) tied to those poses, and then determines the pairwise factors that should be added between these sequential pairs
"""

using Images
using Caesar
using JSON2
using Manifolds

using ImageView

# import Rotations as _Rot

# using Caesar: ScatterAlignPose2

##

# Where to fetch data
# dfgDataFolder = ENV["HOME"]*"/data/rex";
dfgDataFolder = "/tmp/caesar/rex"

# Load the graph
fg = loadDFG("$dfgDataFolder/dfg")

# add the datastore locations
ds = FolderStore{Vector{UInt8}}(:radar, "$dfgDataFolder/data/radar")
addBlobStore!(fg, ds)

ds = FolderStore{Vector{UInt8}}(:gps_fix, "$dfgDataFolder/data/gps")
addBlobStore!(fg, ds)

ds = FolderStore{Vector{UInt8}}(:lidar, "$dfgDataFolder/data/lidar")
addBlobStore!(fg, ds)

##

# fetch variables containing a full sweep
allSweepVariables = filter(v -> :RADARSWEEP in listDataEntries(v), getVariables(fg)) |> sortDFG

fsvars = allSweepVariables .|> getLabel

# helper function to retrieve the radar sweep for a given variable
function fetchSweep(dfg::AbstractDFG, varlabel::Symbol)
  #
  entry,rawData = getData(dfg, varlabel, :RADARSWEEP)
  rawdata = Vector{Float64}(JSON2.read(IOBuffer(rawData)))
  n = Int(sqrt(length(rawdata)))
  sweep = reshape(rawdata,(n,n))
  return sweep # That's pretty sweep if i say so myself...
end

##

# fetch all radar pings
sweeps = fetchSweep.(fg, fsvars);

# Filter the images
kg = Kernel.gaussian(7)
sweeps = map(s -> imfilter(s, kg), sweeps)
# Normalize
sweeps = map(s -> s/maximum(s), sweeps);

# At this point we can load the sweeps; let's work on registration
# First step is to have a function that evaluates the cost of a given transform
# between two subsequent images.

## Building the graph

# maybe one exists, or build a new one below
# newfg = loadDFG("$dfgDataFolder/newfg")

startsweep = 5
endsweep = 26
graphinit = true

len = size(sweeps[5],1)
domain = (range(-100,100,length=976),range(-100,100,length=976))

## use a standard builder

STEP = 20
sweeps_ = sweeps[startsweep:STEP:endsweep]

# build the graph
newfg = buildGraphChain!( sweeps_,
                          ScatterAlignPose2,
                          (_, data) -> (data.currData,data.nextData,domain),
                          stopAfter=1,
                          doRef = false,
                          inflation_fct=0.0,
                          solverParams=SolverParams(graphinit=true, inflateCycles=1) )
#

##

# Save the graph
saveDFG("$dfgDataFolder/segment_test.tar.gz", newfg);

##

lsf(newfg)
# this should run the radar alignment
X1_ = approxConvBelief(newfg, :x0x1f1, :x1)
pts = getPoints(X1_)

# solving will internally call initAll!(newfg)
tree = solveTree!(newfg)

##

using Cairo, Gadfly
Gadfly.set_default_plot_size(40cm,20cm)

#


export Node, CGPInd, get_genes, set_genes!, reset!, forward_connections, get_output_trace
import Base.copy, Base.String, Base.show, Base.summary
import Cambrian.print

"default function for nodes, will cause error if used as a function node"
function f_null(args...)::Nothing
    nothing
end

struct Node
    x::Int16
    y::Int16
    f::Function
    p::Array{Float64}
    active::Bool
end

struct CGPInd <: Cambrian.Individual
    n_in::Int16
    n_out::Int16
    n_parameters::Int16
    chromosome::Array{Float64}
    genes::Array{Float64}
    outputs::Array{Int16}
    nodes::Array{Node}
    buffer::AbstractArray
    fitness::Array{Float64}
end

function recur_active(active::BitArray, n_in::Integer, ind::Int16,
                      xs::Array{Int16}, ys::Array{Int16}, fs::Array{Int16},
                      two_arity::BitArray)::BitArray
    if ind > 0 && ~active[ind]
        active[ind] = true
        active = recur_active(active, n_in, xs[ind], xs, ys, fs, two_arity)
        if two_arity[fs[ind]]
            active = recur_active(active, n_in, ys[ind], xs, ys, fs, two_arity)
        end
    end
    active
end

function find_active(cfg::NamedTuple, genes::Array{Float64},
                     outputs::Array{Int16})::BitArray
    R = cfg.rows
    C = cfg.columns
    n = Int16(cfg.n_in)
    active = falses(R, C)
    xs = Int16.(genes[:, :, 1] .- n)
    ys = Int16.(genes[:, :, 2] .- n)
    fs = Int16.(genes[:, :, 3])
    for i in eachindex(outputs)
        active = recur_active(active, cfg.n_in, outputs[i] - n, xs, ys, fs,
                              cfg.two_arity)
    end
    active
end

function CGPInd(cfg::NamedTuple, chromosome::Array{Float64},
                genes::Array{Float64}, outputs::Array{Int16}; kwargs...)::CGPInd
    R = cfg.rows
    C = cfg.columns
    nodes = Array{Node}(undef, R * C + cfg.n_in)
    p = Float64[]
    for i in 1:cfg.n_in
        nodes[i] = Node(0, 0, f_null, p, false)
    end
    i = cfg.n_in
    active = find_active(cfg, genes, outputs)
    for y in 1:C
        for x in 1:R
            i += 1
            if cfg.n_parameters > 0
                p = genes[x, y, 4:end]
            end
            nodes[i] = Node(Int16(genes[x, y, 1]), Int16(genes[x, y, 2]),
                            cfg.functions[Int16(genes[x, y, 3])], p,
                            active[x, y])
        end
    end
    kwargs_dict = Dict(kwargs)
    # Use given input buffer or default to Array{Float64, 1} type
    if haskey(kwargs_dict, :buffer)
        buffer = kwargs_dict[:buffer]
    else
        buffer = zeros(R * C + cfg.n_in)
    end
    fitness = -Inf .* ones(cfg.d_fitness)
    CGPInd(cfg.n_in, cfg.n_out, cfg.n_parameters, chromosome, genes, outputs,
           nodes, buffer, fitness)
end

function CGPInd(cfg::NamedTuple, chromosome::Array{Float64}; kwargs...)::CGPInd
    R = cfg.rows
    C = cfg.columns
    P = cfg.n_parameters
    # chromosome: node genes, output genes
    genes = reshape(chromosome[1:(R*C*(3+P))], R, C, 3+P)
    # TODO: recurrency is ugly and slow
    maxs = collect(1:R:R*C)
    maxs = round.((R*C .- maxs) .* cfg.recur .+ maxs)
    maxs = min.(R*C + cfg.n_in, maxs .+ cfg.n_in)
    maxs = repeat(maxs, 1, R)'
    genes[:, :, 1] .*= maxs
    genes[:, :, 2] .*= maxs
    genes[:, :, 3] .*= length(cfg.functions)
    genes[:, :, 1:3] = Int16.(ceil.(genes[:, :, 1:3])) # ceil all genes except parameters that stay in [0,1]
    outputs = Int16.(ceil.(chromosome[(R*C*(3+P)+1):end] .* (R * C + cfg.n_in)))
    CGPInd(cfg, chromosome, genes, outputs; kwargs...)
end

function CGPInd(cfg::NamedTuple; kwargs...)::CGPInd
    chromosome = rand(cfg.rows * cfg.columns * (3 + cfg.n_parameters) + cfg.n_out)
    CGPInd(cfg, chromosome; kwargs...)
end

function CGPInd(cfg::NamedTuple, ind::String)::CGPInd
    dict = JSON.parse(ind)
    CGPInd(cfg, Array{Float64}(dict["chromosome"]))
end

function copy(n::Node)
    Node(n.x, n.y, n.f, n.active)
end

function copy(ind::CGPInd)
    buffer = zeros(length(ind.buffer))
    nodes = Array{Node}(undef, length(ind.nodes))
    for i in eachindex(ind.nodes)
        nodes[i] = copy(ind.nodes[i])
    end
    CGPInd(ind.n_in, ind.n_out, copy(ind.chromosome), copy(ind.genes),
           copy(ind.outputs), nodes, buffer, copy(ind.fitness))
end

function String(n::Node)
    JSON.json(Dict(:x=>n.x, :y=>n.y, :f=>string(n.f),  :p=>string(n.p),
              :active=>n.active))
end

function show(io::IO, n::Node)
    print(io, String(n))
end

function String(ind::CGPInd)
    JSON.json(Dict("chromosome"=>ind.chromosome, "fitness"=>ind.fitness,
                   "xs"=>ind.genes[:, :, 1][:], "ys"=>ind.genes[:, :, 2][:],
                   "fs"=>ind.genes[:, :, 3][:]))
end

function show(io::IO, ind::CGPInd)
    print(io, String(ind))
end

function get_active_nodes(ind::CGPInd)
    ind.nodes[[n.active for n in ind.nodes]]
end

function summary(io::IO, ind::CGPInd)
    print(io, string("CGPInd(", get_active_nodes(ind), ", ",
                     findall([n.active for n in ind.nodes]), ", ",
                     ind.outputs, " ,",
                     ind.fitness, ")"))
end

function interpret(i::CGPInd)
    x::AbstractArray->process(i, x)
end

function reset!(c::CGPInd)
    c.buffer .= 0.0
end

"""
    get_gene_indexes(ind::CGPInd, node_id::Integer)

Given an individual and the index of one of its nodes, return the indexes of the
chromosome used to encode this particular node.
"""
function get_gene_indexes(c::CGPInd, node_id::Integer)
    index_start = node_id - c.n_in
    step = length(c.nodes) - c.n_in
    index_end = (2 + c.n_parameters) * (length(c.nodes) - c.n_in) + node_id - c.n_in
    return index_start:step:index_end
end

"""
    get_genes(ind::CGPInd, node_id::Integer)::Array{Float64}

Given an individual and the index of one of its nodes, return the chromosome
used to encode this particular node.

Example:
    get_genes(ind, 42)
"""
function get_genes(c::CGPInd, node_id::Integer)::Array{Float64}
    if node_id > c.n_in
        return c.chromosome[get_gene_indexes(c, node_id)]
    else
        return zeros(3 + c.n_parameters)
    end
end

"""
    get_genes(c::CGPInd, nodes::Array{<:Integer})::Array{Float64}

Given an individual and an array of indexes of one of its nodes, return the
chromosomes used to encode these particular nodes.

Example:
    get_genes(ind, [7, 42])
"""
function get_genes(c::CGPInd, nodes::Array{<:Integer})::Array{Float64}
    if length(nodes) > 0
        return reduce(vcat, map(x->get_genes(c, x), nodes))
    else
        return Array{Float64}(0)
    end
end

"set the genes of node_id to genes"
function set_genes!(c::CGPInd, node_id::Integer, genes::Array{Float64})
    if node_id > c.n_in
        @assert length(genes) == 3 + c.n_parameters
        c.chromosome[get_gene_indexes(c, node_id)] = genes
    end
end

"a list for each node i of a list of all the nodes which have the node i as an input, and i"
function forward_connections(c::CGPInd)
    connections = [[i] for i in 1:length(c.nodes)]
    for ci in eachindex(c.nodes)
        conns = [c.nodes[ci].x, c.nodes[ci].y]
        for conn in conns
            if conn > 0
                push!(connections[conn], ci)
            end
        end
    end
    map(unique, connections)
end

"recursively walk back through inputs"
function recur_output_trace(c::CGPInd, ind::Int16, visited::Array{Int16})
    if ~(ind in visited)
        push!(visited, ind)
        if ind > c.n_in
            recur_output_trace(c, c.nodes[ind].x, visited)
            if CGPFunctions.arity[string(c.nodes[ind].f)] == 2
                recur_output_trace(c, c.nodes[ind].y, visited)
            end
        end
    end
    visited
end

"return a list of node indices which determine the output"
function get_output_trace(c::CGPInd, output_ind::Integer)
    recur_output_trace(c, c.outputs[output_ind], Array{Int16}(undef, 0))
end

"return a list of node indices which determine the output for multiple outputs"
function get_output_trace(c::CGPInd, outputs::Array{<:Integer})
    if length(outputs) > 0
        return unique(reduce(vcat, map(x->get_output_trace(c, x), outputs)))
    else
        return Array{Int16}(undef, 0)
    end
end

"return a list of node indices which determine the output for all outputs"
function get_output_trace(c::CGPInd)
    get_output_trace(c, collect(1:c.n_out))
end


#####
##### `AbstractLayout`
#####

abstract type AbstractLayout end

struct CPUDevice end

struct Layout{L, S, D} <: AbstractLayout
    length::L
    size::S
    device::D
    ismutable::Bool
end

Base.length(layout::Layout) = layout.length

Base.size(layout::Layout) = layout.size

device(layout::Layout) = layout.device

ismutable(layout::Layout) = layout.ismutable

#####
##### `AbstractDomain`
#####

abstract type AbstractDomain end

struct RealDomain <: AbstractDomain end

struct ComplexDomain <: AbstractDomain end

#####
##### `AbstractArgument`
#####

abstract type AbstractArgument end

#== `Signature` ==#

struct Signature{I <: Tuple{Vararg{AbstractArgument}},
                 O <: Tuple{Vararg{AbstractArgument}}}
    input::I
    output::O
end

Signature(input, output) = Signature(markupify.(tuplify(input)),
                                     markupify.(tuplify(output)))

tuplify(x) = tuple(x)
tuplify(x::Tuple) = x

markupify(x::AbstractArgument) = x
markupify(x) = markup(x)

#== `Ignore` ==#

struct Ignore <: AbstractArgument end

Base.length(::Ignore) = 0

#== `AbstractVariable` ==#

abstract type AbstractVariable <: AbstractArgument end

struct Scalar{D <: AbstractDomain} <: AbstractVariable
    domain::D
end

const RealScalar = Scalar{RealDomain}

RealScalar() = Scalar(RealDomain())

const ComplexScalar = Scalar{ComplexDomain}

ComplexScalar() = Scalar(ComplexDomain())

struct Tensor{D <: AbstractDomain, L <: AbstractLayout} <: AbstractVariable
    domain::D
    layout::L
end

const RealTensor = Tensor{RealDomain}

RealTensor(layout) = Tensor(RealDomain(), layout)

const ComplexTensor = Tensor{ComplexDomain}

ComplexTensor(layout) = Tensor(ComplexDomain(), layout)

Base.length(::Scalar) = 1
Base.length(t::Tensor) = length(t.layout)

#####
##### `@sig`
#####

macro sig(expr)
    signature_type_from_expr(expr)
end

const MALFORMED_SIG_ERROR_MESSAGE = "Malformed expression given to `@sig`; see `@sig` docstring for proper format."

function signature_type_from_expr(expr)
    @assert(Meta.isexpr(expr, :call) && expr.args[1] === :→ && length(expr.args) === 3, MALFORMED_SIG_ERROR_MESSAGE)
    input_types = map(parse_into_markup_type, split_infix_args(expr.args[2], :⊕))
    output_types = map(parse_into_markup_type, split_infix_args(expr.args[3], :⊕))
    return :(Signature{<:Tuple{$(input_types...)}, <:Tuple{$(output_types...)}})
end

split_infix_args(invocation::Symbol, ::Symbol) = (invocation,)

function split_infix_args(invocation::Expr, op::Symbol)
    if Meta.isexpr(invocation, :call) && invocation.args[1] === op
        return (split_infix_args(invocation.args[2], op)..., invocation.args[3])
    end
    return (invocation,)
end

function parse_into_markup_type(x)
    if x === :R
        return :(Scalar{RealDomain})
    elseif x === :C
        return :(Scalar{ComplexDomain})
    elseif x === :_
        return :(Ignore)
    elseif isa(x, Expr) && length(x.args) === 1
        if Meta.isexpr(x, :vect)
            domain = x.args[1]
            if domain === :R
                return :(Tensor{RealDomain})
            elseif domain === :C
                return :(Tensor{ComplexDomain})
            end
        elseif Meta.isexpr(x, :braces)
            vararg_type = parse_into_markup_type(x.args[1])
            return :(Vararg{$vararg_type})
        end
    end
    error(string("Encountered unparseable signature element `", x, "`. ",
                 MALFORMED_SIG_ERROR_MESSAGE))
end


#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

using OrdinaryDiffEq: ODESolution
using FMIImport: FMU2Solution

"""
Plots data from a ME-FMU.

Optional `t_in_solution` controls if the first state in the solution is interpreted as t(ime).
Optional keyword argument `maxLabelLength` controls the maximum length for legend labels (too long labels are cut from front).
"""
function fmiPlot(solution::FMU2Solution; kwargs...)
    fig = Plots.plot(; xlabel="t [s]")
    fmiPlot!(fig, solution; kwargs...)
    return fig
end
function fmiPlot!(fig, solution::FMU2Solution;
    states::Union{Bool, Nothing}=nothing,
    values::Union{Bool, Nothing}=nothing,
    stateEvents::Union{Bool, Nothing}=nothing,
    timeEvents::Union{Bool, Nothing}=nothing,
    stateIndices=nothing,
    valueIndices=nothing,
    maxLabelLength=64,
    plotkwargs...)

    numStateEvents = 0
    numTimeEvents = 0
    for e in solution.events
        if e.indicator > 0
            numStateEvents += 1
        else
            numTimeEvents += 1
        end
    end

    if states === nothing
        states = (solution.states !== nothing)
    end

    if values === nothing
        values = (solution.values !== nothing)
    end

    if stateEvents === nothing
        stateEvents = false
        for e in solution.events
            if e.indicator > 0
                stateEvents = true
                break
            end
        end

        if numStateEvents > 100
            @info "fmiPlot(...): Number of state events ($(numStateEvents)) exceeding 100, disabling automatic plotting of state events (can be forced with keyword `stateEvents=true`)."
            stateEvents = false
        end
    end

    if timeEvents === nothing
        timeEvents = false
        for e in solution.events
            if e.indicator == 0
                timeEvents = true
                break
            end
        end

        if numTimeEvents > 100
            @info "fmiPlot(...): Number of time events ($(numTimeEvents)) exceeding 100, disabling automatic plotting of time events (can be forced with keyword `timeEvents=true`)."
            timeEvents = false
        end
    end

    if stateIndices === nothing
        stateIndices = 1:length(solution.fmu.modelDescription.stateValueReferences)
    end

    if valueIndices === nothing
        if solution.values !== nothing
            valueIndices = 1:length(solution.values.saveval[1])
        end
    end

    plot_min = Inf
    plot_max = -Inf

    # plot states
    if states
        t = solution.states.t
        numValues = length(solution.states.u[1])

        for v in 1:numValues
            if v ∈ stateIndices
                vr = solution.fmu.modelDescription.stateValueReferences[v]
                vrNames = fmi2ValueReferenceToString(solution.fmu, vr)
                vrName = vrNames[1]

                vals = collect(data[v] for data in solution.states.u)

                plot_min = min(plot_min, vals...)
                plot_max = max(plot_max, vals...)

                # prevent legend labels from getting too long
                label = "$vrName ($vr)"
                labelLength = length(label)
                if labelLength > maxLabelLength
                    label = "..." * label[labelLength-maxLabelLength:end]
                end

                Plots.plot!(fig, t, vals; label=label, plotkwargs...)
            end
        end
    end

    # plot recorded values
    if values
        t = solution.values.t
        numValues = length(solution.values.saveval[1])

        for v in 1:numValues
            if v ∈ valueIndices
                vr = solution.valueReferences[v]
                vrNames = fmi2ValueReferenceToString(solution.fmu, vr)
                vrName = vrNames[1]

                vals = collect(data[v] for data in solution.values.saveval)

                plot_min = min(plot_min, vals...)
                plot_max = max(plot_max, vals...)

                # prevent legend labels from getting too long
                label = "$vrName ($vr)"
                labelLength = length(label)
                if labelLength > maxLabelLength
                    label = "..." * label[labelLength-maxLabelLength:end]
                end

                Plots.plot!(fig, t, vals; label=label, plotkwargs...)
            end
        end
    end

    if stateEvents
        first = true
        for e in solution.events
            if e.indicator > 0
                Plots.plot!(fig, [e.t, e.t], [plot_min, plot_max]; label=(first ? "State event(s)" : nothing), style=:dash, color=:blue)
                first = false
            end
        end
    end

    if timeEvents
        first = true
        for e in solution.events
            if e.indicator == 0
                Plots.plot!(fig, [e.t, e.t], [plot_min, plot_max]; label=(first ? "Time event(s)" : nothing), style=:dash, color=:red)
                first = false
            end
        end
    end

    return fig
end

"""
Extended the original plot-command by plotting FMUs.
"""
function Plots.plot(fmu::FMU2, args...; kwargs...)
    fmiPlot(fmu, args...; kwargs...)
end
function Plots.plot!(fig, fmu::FMU2, args...; kwargs...)
    fmiPlot!(fig, fmu, args...; kwargs...)
end


module GEAS

using LinearAlgebra, Dates, HTTP, DataFrames, Random, Distributions,
    Serialization, Statistics, StatsPlots, JLD2, LaTeXStrings
import Octavian: matmul, matmul!

const prj_dir = begin
    t = splitdir(pathof(GEAS))[1]
    l = findlast('/', t) - 1
    t[1:l]
end

dat_dir = joinpath(prj_dir, "dat")
bin_dir = joinpath(prj_dir, "bin")
plink   = joinpath(bin_dir, "plink")
beagle  = joinpath(bin_dir, "beagle.jar")
macs    = joinpath(bin_dir, "macs")
qmsim   = joinpath(bin_dir, "qmsim")

include("breeding.jl")
include("b-Update.jl")
include("dropping.jl")
include("evaluation.jl")
include("gedit.jl")
include("genotypes.jl")
include("inbreeding.jl")
include("phenotypes.jl")
include("sim-base.jl")
include("sim-QTL.jl")
include("structs.jl")
include("summarize.jl")
include("workflow.jl")
include("check-par.jl")
include("utils.jl")
include("gwas.jl")
include("bayes-alphabet.jl")

copyright()

end # module


using Oxygen

println("Testing SMILES Parser")

function sameatom(atoms::Tuple{OxygenAtom,OxygenAtom})::Bool
    (a,b) = atoms
    a.number == b.number
end

function sameatoms(a::Array{OxygenAtom,1}, b::Array{OxygenAtom,1})::Bool
    all(map(sameatom, zip(a,b)))
end

@testset "SmilesToMol" begin
    ethyl_acetate = "CC(=O)OCC"
    ethyl_acetate_mol = smilestomol(ethyl_acetate)
    @test typeof(ethyl_acetate_mol) == OxygenMol
    @test sameatoms(ethyl_acetate_mol.atoms, [OxygenAtom("C"), OxygenAtom("C"), OxygenAtom("O"), OxygenAtom("O"), OxygenAtom("C"), OxygenAtom("C")])
    @test ethyl_acetate_mol.adj ==
          [[(2, 1.0)], [(1, 1.0), (3, 2.0), (4, 1.0)], [(2, 2.0)], [(2, 1.0), (5, 1.0)], [(4, 1.0), (6, 1.0)], [(5, 1.0)]]

    benzene = "c1ccccc1"
    benzene_mol = smilestomol(benzene)
    @test typeof(benzene_mol) == OxygenMol
    @test sameatoms(benzene_mol.atoms, [OxygenAtom("C"), OxygenAtom("C"), OxygenAtom("C"), OxygenAtom("C"), OxygenAtom("C"), OxygenAtom("C")])
    @test benzene_mol.adj == [
        [(2, 1.5), (6, 1.5)],
        [(1, 1.5), (3, 1.5)],
        [(2, 1.5), (4, 1.5)],
        [(3, 1.5), (5, 1.5)],
        [(4, 1.5), (6, 1.5)],
        [(5, 1.5), (1, 1.5)],
    ]

    cipro = "C1CC1N2C=C(C(=O)C3=CC(=C(C=C32)N4CCNCC4)F)C(=O)O"
    cipro_mol = smilestomol("C1CC1N2C=C(C(=O)C3=CC(=C(C=C32)N4CCNCC4)F)C(=O)O")
    @test typeof(cipro_mol) == OxygenMol
    @test sameatoms(cipro_mol.atoms[7:11], [OxygenAtom("C"), OxygenAtom("O"), OxygenAtom("C"), OxygenAtom("C"), OxygenAtom("C")])
    @test cipro_mol.adj[7:11] == [
        [(6, 1.0), (8, 2.0), (9, 1.0)],
        [(7, 2.0)],
        [(7, 1.0), (10, 2.0), (14, 1.0)],
        [(9, 2.0), (11, 1.0)],
        [(10, 1.0), (12, 2.0), (21, 1.0)],
    ]
end


module PathWeightSampling

using Distributions
import DiffEqBase
using DiffEqJump
include("trajectories/trajectory.jl")
include("EventIter.jl")
include("MetropolisSampler.jl")
include("marginal_strategies/strategies.jl")
include("GaussianSystem.jl")
include("trajectories/distribution.jl")
include("EmpiricalDistribution.jl")
include("model_setup/DrivenJumpProblem.jl")
include("JumpNetwork.jl")
include("write_hdf5.jl")

include("example_systems.jl")

include("ParallelRun.jl")

using .DrivenJumpProblems

export TIEstimate, AnnealingEstimate, DirectMCEstimate, SMCEstimate,
    generate_configuration, mutual_information

end # module


abstract type AbstractMeasureSpec end

#holds the parameters for forming the weights
#not meant to hold any meaningful data- just parameters
struct MeasureSpec{TFvol<:NSymbol, TFDvol<:NSymbol} <: AbstractMeasureSpec
  Fξs::Vector{Symbol}
  ξs::Vector{Characteristic}
  ξidx::Dict{Symbol, Characteristic}
  Fvol::TFvol
  FDvol::TFDvol

  FW::Vector{Symbol} #these are global controls (use case: intercept, controls by k or n)
  FWbyt::Vector{Symbol} #these are controls grouped by t only (use case: time fixed effects)
  FWbykt::Dict{Symbol, Vector{Symbol}} #these are controls grouped by t AND k

  K::Int
end

Fweights(ms::MeasureSpec, s::Symbol) = (ξ->ξ.X[s]).(ms.ξs)
Fgroupedcontrols(ms::MeasureSpec) = [ms.FWbyt...;(Fξ->ms.FWbykt[Fξ]).(ms.Fξs)...;]
Fcontrols(ms::MeasureSpec) = [ms.FW...; Fgroupedcontrols(ms)...;]
grossleverages(ms::MeasureSpec) = (ξ->ξ.ws.gross).(ms.ξs)

#creates the control names as needed
function getcontrolnames(Fξs::Vector{Symbol};
    controlsprefix::Symbol = PARAM[:controlsprefix],
    controlsglobal::Vector{Symbol} = PARAM[:controlsglobal],
    controlsbykt::Vector{Symbol} = PARAM[:controlsbykt],
    controlsbyt::Vector{Symbol} = PARAM[:controlsbyt])

  FW = (s->Symbol(controlsprefix,s)).(controlsglobal)
  FWbyt = (s->Symbol(controlsprefix,s)).(controlsbyt)
  FWbykt::Dict{Symbol,Vector{Symbol}} = Dict{Symbol,Vector{Symbol}}()
  for Fξ ∈ Fξs
    FWbykt[Fξ] = [Symbol(controlsprefix, s, Fξ) for s ∈ controlsbykt]
  end

  ntout = (;FW=FW, FWbyt=FWbyt, FWbykt=FWbykt)
  #println("ntout: $(ntout), typeof: $(typeof(ntout))")
  return ntout
end

function MeasureSpec(ξstuple::Characteristic...;
  Fvol::NSymbol = PARAM[:Fvol],
  FDvol::NSymbol = nothing, #set the keyword to Symbol(:D, Fvol) to enable
  ξs::Vector{Characteristic} = collect(ξstuple),
  Fξs::Vector{Symbol} = (ξ->ξ.Fξ).(ξs),
  controlnames::NamedTuple=getcontrolnames(Fξs)
  )

  ξidx::Dict{Symbol, Characteristic} = Dict(zip(Fξs, ξs))

  #permnodate::Vector{Tuple{Int, Date}} = (Tuple).(zip(panel.permno, panel.date))
  #println("controlnames: $(controlnames), typeof: $(typeof(controlnames))")

  return MeasureSpec(Fξs, ξs, ξidx, Fvol, FDvol,
    controlnames.FW, controlnames.FWbyt, controlnames.FWbykt, length(Fξs))
end

#set up a get index notation
Base.getindex(ms::MeasureSpec, Fξ::Symbol) = ms.ξidx[Fξ]
Base.getindex(ms::MeasureSpec, Fξs::Symbol...) = (Fξ->ms.idx[Fξ]).(Fξs)
Base.getindex(ms::MeasureSpec, i::Int) = ms.ξiBase.dx[i]
Base.getindex(ms::MeasureSpec, is::Int...) = ms.ξidx[is]

Base.length(ms::MeasureSpec) = length(Fξs)

#make iteration work
function Base.iterate(ms::MeasureSpec, state::Int = 1)
  (state > ms.K) && return nothing
  return (ms[state], state + 1)
end

#some convenience methods
getT(panel::AbstractDataFrame, ms::MeasureSpec) = length(unique(panel.date))
getN(panel::AbstractDataFrame, ms::MeasureSpec) = length(unique(panel.permno))
getdims(panel::AbstractDataFrame, ms::MeasureSpec) = (
  T=getT(panel, ms), N = getN(panel, ms), K = ms.K)

function gettidx(panel::AbstractDataFrame, ms::MeasureSpec)
  date::Vector{Date} = sort(unique(panel.date))
  tidx::Dict{Date,Int} = Dict(date[t]=>t for t ∈ 1:length(date))

  return tidx
end

#=makes sure the MeasureSpec matches the ordering in the dataframe
function verifymeasureintegrity(panel::DataFrame, ms::MeasureSpec)
  Nrows::Int = size(panel, 1)

  permno::Vector{Int} = (t->t[1]).(ms.permnodate)
  date::Vector{Date} = (t->t[2]).(ms.permnodate)

  (sum(permno .== panel.permno) == Nrows) || error("permno integrity check failed")
  (sum(date .== panel.date) == Nrows) || error("date integrity check failed")
end=#


export
    WindowFilter,
    SyncPF,
    AsyncPF,
    CoherentPF

"""
    struct WindowFilter <: Filter

    General construction for a filter based on a flexible windowing scheme.

#arguments
    window_creator Defines the method used to create windows
    inference_algorithm Defines the algorithm to use on a window
    postprocess! A postprocessing function, that takes the runtime and does any additional processing needed to carry to the next iteration. Defaults to doing nothing.

"""
mutable struct WindowFilter <: Filter
    window_creator :: WindowCreator
    inference_algorithm :: InstantAlgorithm
    postprocess! :: Function
    latest_window :: Union{InstantRuntime, Nothing}
    WindowFilter(wc, ia, pp) = new(wc, ia, pp, nothing)
    WindowFilter(wc, ia) = new(wc, ia, run -> nothing, nothing)
end

function init_filter(wf::WindowFilter, dynrun::DynamicRuntime)
    ensure_all!(dynrun, current_time(dynrun))
    instrun = initial_instant_runtime(dynrun)
    # We assume no evidence or interventions at time 0
    # TODO: Handle placeholder beliefs
    infer(wf.inference_algorithm, instrun)
    wf.latest_window = instrun
    retrieve_values_from_instant_runtime!(dynrun, instrun)
    _store_beliefs(wf, dynrun, instrun)
end

function _store_beliefs(wf::WindowFilter, dynrun::DynamicRuntime{T}, instrun::InstantRuntime) where T
    dynnet = get_network(dynrun)
    for instvar in get_variables(get_network(instrun))
        instinst = current_instance(instrun, instvar)
        belief = marginal(wf.inference_algorithm, instrun, instinst)
        (dynname, t) = dynamic_name_and_time(instvar, T)
        dyninst = get_instance(dynrun, get_node(dynnet, dynname), t)
        set_value!(dynrun, dyninst, :belief, belief)
    end
end

function filter_step(wf::WindowFilter, dynrun::DynamicRuntime{T}, variables::Vector{Variable}, time::T, evidence::Dict{Symbol, Score}) where T
    dynnet = get_network(dynrun)
    insts = create_window(wf.window_creator, dynrun, variables, time)
    instrun = instant_runtime_from_instances(dynrun, insts)
    # Apply the dynamic evidence to the instant runtime
    instev = Dict{Symbol, Score}()
    for (name, sc) in evidence
        instev[instant_name(name, time)] = sc
    end
    # Apply beliefs in the dynamic network as placeholder beliefs in the instant network.
    placeholder_beliefs = Dict{Symbol,Dist}()
    inst_phs = get_placeholders(get_network(instrun))
    for instnode in inst_phs
        (dynname, t) = dynamic_name_and_time(instnode, T)
        dynnode = get_node(dynnet, dynname)
        dyninst = get_instance(dynrun, dynnode, t)
        belief = get_value(dynrun, dyninst, :belief)
        placeholder_beliefs[get_name(instnode)] = belief
    end
    # TODO: Handle interventions
    instinterv = Dict{Symbol,Dist}()
    infer(wf.inference_algorithm, instrun, instev, instinterv, placeholder_beliefs)
    wf.latest_window = instrun
    retrieve_values_from_instant_runtime!(dynrun, instrun)
    set_time!(dynrun, time)
    _store_beliefs(wf, dynrun, instrun)
end

function answer(::Marginal, ::WindowFilter, dynrun::Runtime, target::VariableInstance)
    return get_value(dynrun, target, :belief)
end

# function answer(query::Query, wf::WindowFilter, dynrun::Runtime, targets::Vector{VariableInstance})
#     # TODO: This code assumes that targets are in the latest_window, which might not be true
#     # for an asynchronous filter. We need to construct an instant window for the targets.
#     # For that, we need to put information from the dynamic window into the instant window,
#     # which is not done yet.
#     instrun = wf.latest_window
#     insttargets = VariableInstance[]
#     for target in targets
#         instname = instant_name(get_name(target), current_time(dynrun))
#         instnode = get_node(get_network(instrun), instname)
#         insttarget = current_instance(instrun, instnode)
#         push!(insttargets, insttarget)
#     end
#     answer(query, wf.inference_algorithm, wf.latest_window, insttargets)
# end