## reductions ##

###### Functors ######

# Note that functors are merely used as internal machinery to enhance code reuse.
# They are not exported.
# When function arguments can be inlined, the use of functors can be removed.

abstract Func{N}

type IdFun <: Func{1} end
type AbsFun <: Func{1} end
type Abs2Fun <: Func{1} end
type ExpFun <: Func{1} end
type LogFun <: Func{1} end

type AddFun <: Func{2} end
type MulFun <: Func{2} end
type AndFun <: Func{2} end
type OrFun <: Func{2} end
type MaxFun <: Func{2} end
type MinFun <: Func{2} end

evaluate(::IdFun, x) = x
evaluate(::AbsFun, x) = abs(x)
evaluate(::Abs2Fun, x) = abs2(x)
evaluate(::ExpFun, x) = exp(x)
evaluate(::LogFun, x) = log(x)
evaluate(f::Callable, x) = f(x)

evaluate(::AddFun, x, y) = x + y
evaluate(::MulFun, x, y) = x * y
evaluate(::AndFun, x, y) = x & y
evaluate(::OrFun, x, y) = x | y
evaluate(::MaxFun, x, y) = scalarmax(x, y)
evaluate(::MinFun, x, y) = scalarmin(x, y)
evaluate(f::Callable, x, y) = f(x, y)

###### Generic (map)reduce functions ######

# r_promote: promote x to the type of reduce(op, [x])
r_promote(op, x) = x
r_promote(::AddFun, x) = x + zero(x)
r_promote(::MulFun, x) = x * one(x)

## foldl && mapfoldl

function mapfoldl_impl(f, op, v0, itr, i)
    if done(itr, i)
        return v0
    else
        (x, i) = next(itr, i)
        v = evaluate(op, v0, evaluate(f, x))
        while !done(itr, i)
            (x, i) = next(itr, i)
            v = evaluate(op, v, evaluate(f, x))
        end
        return v
    end
end

mapfoldl(f, op, v0, itr) = mapfoldl_impl(f, op, v0, itr, start(itr))

function mapfoldl(f, op::Function, v0, itr)
    is(op, +) ? mapfoldl(f, AddFun(), v0, itr) :
    is(op, *) ? mapfoldl(f, MulFun(), v0, itr) :
    is(op, &) ? mapfoldl(f, AndFun(), v0, itr) :
    is(op, |) ? mapfoldl(f, OrFun(), v0, itr) :
    mapfoldl_impl(f, op, v0, itr, start(itr))
end

function mapfoldl(f, op, itr)
    i = start(itr)
    done(itr, i) && error("Argument is empty.")
    (x, i) = next(itr, i)
    v0 = evaluate(f, x)
    mapfoldl_impl(f, op, v0, itr, i)
end

foldl(op, v0, itr) = mapfoldl(IdFun(), op, v0, itr)
foldl(op, itr) = mapfoldl(IdFun(), op, itr)

## foldr & mapfoldr

function mapfoldr_impl(f, op, v0, itr, i::Integer)
    if i == 0
        return v0
    else
        x = itr[i]
        v = evaluate(op, evaluate(f, x), v0)
        while i > 1
            x = itr[i -= 1]
            v = evaluate(op, evaluate(f, x), v)
        end
        return v
    end
end

mapfoldr(f, op, v0, itr) = mapfoldr_impl(f, op, v0, itr, endof(itr))
mapfoldr(f, op, itr) = (i = endof(itr); mapfoldr_impl(f, op, evaluate(f, itr[i]), itr, i-1))

foldr(op, v0, itr) = mapfoldr(IdFun(), op, v0, itr)
foldr(op, itr) = mapfoldr(IdFun(), op, itr)

## reduce & mapreduce

# mapreduce_***_impl require ifirst < ilast
function mapreduce_seq_impl(f, op, A::AbstractArray, ifirst::Int, ilast::Int)
    @inbounds fx1 = evaluate(f, A[ifirst])
    @inbounds fx2 = evaluate(f, A[ifirst+=1])
    @inbounds v = evaluate(op, fx1, fx2)
    while ifirst < ilast
        @inbounds fx = evaluate(f, A[ifirst+=1])
        v = evaluate(op, v, fx)
    end
    return v
end

function mapreduce_pairwise_impl(f, op, A::AbstractArray, ifirst::Int, ilast::Int, blksize::Int)
    if ifirst + blksize > ilast
        return mapreduce_seq_impl(f, op, A, ifirst, ilast)
    else
        imid = (ifirst + ilast) >>> 1
        v1 = mapreduce_seq_impl(f, op, A, ifirst, imid)
        v2 = mapreduce_seq_impl(f, op, A, imid+1, ilast)
        return evaluate(op, v1, v2)
    end
end

mapreduce(f, op, itr) = mapfoldl(f, op, itr)
mapreduce(f, op, v0, itr) = mapfoldl(f, op, v0, itr)
mapreduce_impl(f, op, A::AbstractArray, ifirst::Int, ilast::Int) =
    mapreduce_seq_impl(f, op, A, ifirst, ilast)

# handling empty arrays
mr_empty(f, op, T) = error("Reducing over an empty array is not allowed.")
mr_empty(::IdFun, op::AddFun, T) = r_promote(op, zero(T))
mr_empty(::AbsFun, op::AddFun, T) = r_promote(op, abs(zero(T)))
mr_empty(::Abs2Fun, op::AddFun, T) = r_promote(op, abs2(zero(T)))
mr_empty(::IdFun, op::MulFun, T) = r_promote(op, one(T))
mr_empty(::AbsFun, op::MaxFun, T) = abs(zero(T))
mr_empty(::Abs2Fun, op::MaxFun, T) = abs2(zero(T))
mr_empty(f, op::AndFun, T) = true
mr_empty(f, op::OrFun, T) = false

function _mapreduce{T}(f, op, A::AbstractArray{T})
    n = length(A)
    if n == 0
        return mr_empty(f, op, T)
    elseif n == 1
        return r_promote(op, evaluate(f, A[1]))
    elseif n < 16
        @inbounds fx1 = evaluate(f, A[1])
        @inbounds fx2 = evaluate(f, A[2])
        s = evaluate(op, fx1, fx2)
        i = 2
        while i < n
            @inbounds fx = evaluate(f, A[i+=1])
            s = evaluate(op, s, fx)
        end
        return s
    else
        return mapreduce_impl(f, op, A, 1, n)
    end
end

mapreduce(f, op, A::AbstractArray) = _mapreduce(f, op, A)
mapreduce(f, op, a::Number) = evaluate(f, a)

function mapreduce(f, op::Function, A::AbstractArray)
    is(op, +) ? _mapreduce(f, AddFun(), A) :
    is(op, *) ? _mapreduce(f, MulFun(), A) :
    is(op, &) ? _mapreduce(f, AndFun(), A) :
    is(op, |) ? _mapreduce(f, OrFun(), A) :
    _mapreduce(f, op, A)
end

reduce(op, v0, itr) = mapreduce(IdFun(), op, v0, itr)
reduce(op, itr) = mapreduce(IdFun(), op, itr)
reduce(op, a::Number) = a

###### Specific reduction functions ######

## sum

function mapreduce_seq_impl(f, op::AddFun, a::AbstractArray, ifirst::Int, ilast::Int)
    @inbounds if ifirst + 6 >= ilast  # length(a) < 8
        i = ifirst
        s = evaluate(f, a[i]) + evaluate(f, a[i+1])
        i = i+1
        while i < ilast
            s += evaluate(f, a[i+=1])
        end
        return s

    else # length(a) >= 8, manual unrolling
        s1 = evaluate(f, a[ifirst]) + evaluate(f, a[ifirst + 4])
        s2 = evaluate(f, a[ifirst + 1]) + evaluate(f, a[ifirst + 5])
        s3 = evaluate(f, a[ifirst + 2]) + evaluate(f, a[ifirst + 6])
        s4 = evaluate(f, a[ifirst + 3]) + evaluate(f, a[ifirst + 7])
        i = ifirst + 8
        il = ilast - 3
        while i <= il
            s1 += evaluate(f, a[i])
            s2 += evaluate(f, a[i+1])
            s3 += evaluate(f, a[i+2])
            s4 += evaluate(f, a[i+3])
            i += 4
        end
        while i <= ilast
            s1 += evaluate(f, a[i])
            i += 1
        end
        return s1 + s2 + s3 + s4
    end
end

# Note: sum_seq uses four accumulators, so each accumulator gets at most 256 numbers
sum_pairwise_blocksize(f) = 1024

# This appears to show a benefit from a larger block size
sum_pairwise_blocksize(::Abs2Fun) = 4096

mapreduce_impl(f, op::AddFun, A::AbstractArray, ifirst::Int, ilast::Int) =
    mapreduce_pairwise_impl(f, op, A, ifirst, ilast, sum_pairwise_blocksize(f))

sum(f::Union(Function,Func{1}), a) = mapreduce(f, AddFun(), a)
sum(a) = mapreduce(IdFun(), AddFun(), a)
sum(a::AbstractArray{Bool}) = countnz(a)
sumabs(a) = mapreduce(AbsFun(), AddFun(), a)
sumabs2(a) = mapreduce(Abs2Fun(), AddFun(), a)

# Kahan (compensated) summation: O(1) error growth, at the expense
# of a considerable increase in computational expense.
function sum_kbn{T<:FloatingPoint}(A::AbstractArray{T})
    n = length(A)
    if n == 0
        return sumzero(T)
    end
    c = zero(T)
    s = A[1] + c
    for i in 2:n
        @inbounds Ai = A[i]
        t = s + Ai
        if abs(s) >= abs(Ai)
            c += ((s-t) + Ai)
        else
            c += ((Ai-t) + s)
        end
        s = t
    end
    s + c
end

## prod

prod(f::Union(Function,Func{1}), a) = mapreduce(f, MulFun(), a)
prod(a) = mapreduce(IdFun(), MulFun(), a)

prod(A::AbstractArray{Bool}) =
    error("use all() instead of prod() for boolean arrays")

## maximum & minimum

function mapreduce_impl(f, op::MaxFun, A::AbstractArray, first::Int, last::Int)
    # locate the first non NaN number
    v = evaluate(f, A[first])
    i = first + 1
    while v != v && i <= last
        @inbounds v = evaluate(f, A[i])
        i += 1
    end
    while i <= last
        @inbounds x = evaluate(f, A[i])
        if x > v
            v = x
        end
        i += 1
    end
    v
end

function mapreduce_impl(f, op::MinFun, A::AbstractArray, first::Int, last::Int)
    # locate the first non NaN number
    v = evaluate(f, A[first])
    i = first + 1
    while v != v && i <= last
        @inbounds v = evaluate(f, A[i])
        i += 1
    end
    while i <= last
        @inbounds x = evaluate(f, A[i])
        if x < v
            v = x
        end
        i += 1
    end
    v
end

maximum(f::Union(Function,Func{1}), a) = mapreduce(f, MaxFun(), a)
minimum(f::Union(Function,Func{1}), a) = mapreduce(f, MinFun(), a)

maximum(a) = mapreduce(IdFun(), MaxFun(), a)
minimum(a) = mapreduce(IdFun(), MinFun(), a)

maxabs(a) = mapreduce(AbsFun(), MaxFun(), a)
minabs(a) = mapreduce(AbsFun(), MinFun(), a)

## extrema

extrema(r::Range) = (minimum(r), maximum(r))
extrema(x::Real) = (x, x)

function extrema(itr)
    s = start(itr)
    done(itr, s) && error("argument is empty")
    (v, s) = next(itr, s)
    while v != v && !done(itr, s)
        (x, s) = next(itr, s)
        v = x
    end
    vmin = v
    vmax = v
    while !done(itr, s)
        (x, s) = next(itr, s)
        if x > vmax
            vmax = x
        elseif x < vmin
            vmin = x
        end
    end
    return (vmin, vmax)
end

## all & any

function mapfoldl(f, ::AndFun, itr)
    for x in itr
        if !evaluate(f, x)
            return false
        end
    end
    return true
end

function mapfoldl(f, ::OrFun, itr)
    for x in itr
        if evaluate(f, x)
            return true
        end
    end
    return false
end

function mapreduce_impl(f, op::AndFun, A::AbstractArray, ifirst::Int, ilast::Int)
    while ifirst <= ilast
        @inbounds x = A[ifirst]
        if !evaluate(f, x)
            return false
        end
        ifirst += 1
    end
    return true
end

function mapreduce_impl(f, op::OrFun, A::AbstractArray, ifirst::Int, ilast::Int)
    while ifirst <= ilast
        @inbounds x = A[ifirst]
        if evaluate(f, x)
            return true
        end
        ifirst += 1
    end
    return false
end

all(a) = mapreduce(IdFun(), AndFun(), a)
any(a) = mapreduce(IdFun(), OrFun(), a)

all(pred::Union(Function,Func{1}), a) = mapreduce(pred, AndFun(), a)
any(pred::Union(Function,Func{1}), a) = mapreduce(pred, OrFun(), a)

## in & contains

immutable EqX{T} <: Func{1}
    x::T
end
EqX{T}(x::T) = EqX{T}(x)
evaluate(f::EqX, y) = (y == f.x)

in(x, itr) = any(EqX(x), itr)

const ∈ = in
∉(x, itr)=!∈(x, itr)
∋(itr, x)= ∈(x, itr)
∌(itr, x)=!∋(itr, x)

function contains(itr, x)
    depwarn("contains(collection, item) is deprecated, use in(item, collection) instead", :contains)
    in(x, itr)
end

function contains(eq::Function, itr, x)
    for y in itr
        if eq(y, x)
            return true
        end
    end
    return false
end

## countnz & count

function count(pred::Union(Function,Func{1}), itr)
    n = 0
    for x in itr
        if evaluate(pred, x)
            n += 1
        end
    end
    return n
end

function count(pred::Union(Function,Func{1}), a::AbstractArray)
    n = 0
    for i = 1:length(a)
        @inbounds if evaluate(pred, a[i])
            n += 1
        end
    end
    return n
end

type NotEqZero <: Func{1} end
evaluate(NotEqZero, x) = (x != 0)

countnz(a) = count(NotEqZero(), a)


mutable struct DelayMNRMJumpAggregation{T,S,F1,F2,RNG,DG,PQ} <: AbstractDSSAJumpAggregator
    next_jump::Int
    prev_jump::Int
    next_jump_time::T
    end_time::T
    cur_rates::Vector{T}
    sum_rate::T
    ma_jumps::S
    rates::F1
    affects!::F2
    save_positions::Tuple{Bool,Bool}
    rng::RNG
    dep_gr::DG
    pq::PQ
    next_delay::Union{Nothing,Vector{Int}}
    num_next_delay::Union{Nothing,Vector{Int}}
    time_to_next_jump::T
    dt_delay::T
    dep_gr_delay::Union{Nothing,Dict{Int,Vector{Int}}}
end

function DelayMNRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool,Bool}, rng::RNG; num_specs, dep_graph=nothing, kwargs...) where {T,S,F1,F2,RNG}

    # a dependency graph is needed and must be provided if there are constant rate jumps
    if dep_graph === nothing
        if (get_num_majumps(maj) == 0) || !isempty(rs)
            error("To use ConstantRateJumps with the Delay Modified Next Reaction Method (Delay MNRM) algorithm a dependency graph must be supplied.")
        else
            dg = make_dependency_graph(num_specs, maj)
        end
    else
        dg = dep_graph
        # make sure each jump depends on itself
        add_self_dependencies!(dg)
    end

    pq = MutableBinaryMinHeap{T}()

    nd = nothing
    nnd = nothing
    ttnj = zero(et)
    dt_delay = zero(et)
    dg_delay = nothing
    DelayMNRMJumpAggregation{T,S,F1,F2,RNG,typeof(dg),typeof(pq)}(nj, nj, njt, et, crs, sr, maj, rs, affs!, sps, rng, dg, pq, nd, nnd, ttnj, dt_delay, dg_delay)
end

############################# Required Functions ##############################
# creating the JumpAggregation structure (function wrapper-based constant jumps)
function aggregate(aggregator::DelayMNRM, u, p, t, end_time, constant_jumps,
    ma_jumps, save_positions, rng; kwargs...)

    # handle constant jumps using function wrappers
    rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps)
    build_jump_aggregation(DelayMNRMJumpAggregation, u, p, t, end_time, ma_jumps,
        rates, affects!, save_positions, rng; num_specs=length(u), kwargs...)
end

# set up a new simulation and calculate the first jump / jump time
function initialize!(p::DelayMNRMJumpAggregation, integrator, u, params, t)
    fill_rates_and_get_times!(p, u, params, t)
    p.dep_gr_delay = dep_gr_delay(p, integrator)
    find_next_delay_dt!(p, integrator)
    generate_jumps!(p, integrator, u, params, t)
    nothing
end

# execute one jump, changing the system state
function execute_jumps!(p::DelayMNRMJumpAggregation, integrator, u, params, t)
    # execute jump
    update_state_delay!(p, integrator, u, t)
    # update current jump rates and times
    update_dependent_rates_delay!(p, integrator, u, params, t)
    nothing
end

# calculate the next jump / jump time
# just the top of the priority queue
function generate_jumps!(p::DelayMNRMJumpAggregation, integrator, u, params, t)
    next_jump_time_reaction, next_jump = top_with_handle(p.pq)
    dt_reaction = next_jump_time_reaction - t
    dt_delay_generation!(p, integrator)
    compare_delay!(p, integrator.de_chan, p.dt_delay, dt_reaction, t)
    if p.next_delay != nothing
        p.next_jump = 0
    else
        p.next_jump = next_jump
    end
    nothing
end

# recalculate jump rates for jumps that depend on the just executed jump (p.next_jump)
function update_dependent_rates_delay!(p::DelayMNRMJumpAggregation, integrator, u, params, t)
    if p.next_delay == nothing  # if next reaction is not delay reaction
        @inbounds dep_rxs = p.dep_gr[p.next_jump]
    else
        # find the dep_rxs w.r.t next_delay vectors
        dep_rxs_ = [p.dep_gr_delay[p.next_delay[i]] for i in eachindex(p.next_delay)]
        dep_rxs = reduce(vcat, dep_rxs_)
    end
    @unpack cur_rates, rates, ma_jumps = p
    num_majumps = get_num_majumps(ma_jumps)
    @inbounds for rx in dep_rxs
        oldrate = cur_rates[rx]

        # update the jump rate
        @inbounds cur_rates[rx] = calculate_jump_rate(ma_jumps, num_majumps, rates, u, params, t, rx)

        # calculate new jump times for dependent jumps
        if rx != p.next_jump && oldrate > zero(oldrate)
            if cur_rates[rx] > zero(eltype(cur_rates))
                update!(p.pq, rx, t + oldrate / cur_rates[rx] * (p.pq[rx] - t))
            else
                update!(p.pq, rx, typemax(t))
            end
        else
            if cur_rates[rx] > zero(eltype(cur_rates))
                update!(p.pq, rx, t + randexp(p.rng) / cur_rates[rx])
            else
                update!(p.pq, rx, typemax(t))
            end
        end
    end
    nothing
end

# reevaulate all rates, recalculate all jump times, and reinit the priority queue
function fill_rates_and_get_times!(p::DelayMNRMJumpAggregation, u, params, t)

    # mass action jumps
    majumps = p.ma_jumps
    cur_rates = p.cur_rates
    # num_de_chan = length(integrator.de_chan) # num_de_chan
    pqdata = Vector{typeof(t)}(undef, length(cur_rates))
    @inbounds for i in 1:get_num_majumps(majumps)
        cur_rates[i] = evalrxrate(u, i, majumps)
        pqdata[i] = t + randexp(p.rng) / cur_rates[i]
    end

    # constant rates
    rates = p.rates
    idx = get_num_majumps(majumps) + 1
    @inbounds for rate in rates
        cur_rates[idx] = rate(u, params, t)
        pqdata[idx] = t + randexp(p.rng) / cur_rates[idx]
        idx += 1
    end

    # setup a new indexed priority queue to storing rx times
    p.pq = MutableBinaryMinHeap(pqdata)
    nothing
end


""" Data structure for (directed) wiring diagrams, aka string diagrams.

A (directed) wiring diagram consists of a collection of boxes with input and
output ports connected by wires. A box can be atomic (possessing no internal
structure) or can itself be a wiring diagram. Thus, wiring diagrams can be
nested recursively. Wiring diagrams are closely related to what the CS
literature calls "directed graphs with ports" or more simply "port graphs". The
main difference is that a wiring diagram has an "outer box": a wiring diagram
has its own ports that can be connected to the ports of its boxes.

This module provides a generic data structure for wiring diagrams. Arbitrary
data can be attached to the boxes, ports, and wires of a wiring diagram. The
diagrams are "abstract" in the sense that they cannot be directly rendered as
raster or vector graphics. However, they form a useful intermediate
representation that can be serialized to and from GraphML or translated into
Graphviz or other declarative diagram languages.
"""
module DirectedWiringDiagrams
export AbstractBox, Box, WiringDiagram, Wire, Port, PortKind,
  InputPort, OutputPort, input_ports, output_ports, input_id, output_id,
  outer_ids, boxes, box_ids, nboxes, nwires, box, wires, has_wire, graph,
  add_box!, add_boxes!, add_wire!, add_wires!, rem_box!, rem_boxes!, rem_wire!,
  rem_wires!, port_value, validate_ports, is_permuted_equal,
  all_neighbors, neighbors, outneighbors, inneighbors, in_wires, out_wires,
  singleton_diagram, induced_subdiagram, encapsulated_subdiagram,
  ocompose, substitute, encapsulate

using Compat
using AutoHashEquals

using ...Present, ...CSetDataStructures, ...Graphs.BasicGraphs
using ...Graphs.BasicGraphs: TheoryGraph
import ...Graphs: all_neighbors, neighbors, outneighbors, inneighbors

# Data types
############

""" Kind of port: input or output.
"""
@enum PortKind InputPort OutputPort

""" A port on a box to which wires can be connected.
"""
@auto_hash_equals struct Port
  box::Int
  kind::PortKind
  port::Int
end
set_box(port::Port, box::Int) = Port(box, port.kind, port.port)

function Base.isless(p1::Port, p2::Port)::Bool
  # Lexicographic order.
  p1.box < p2.box ||
    (p1.box == p2.box &&
      (p1.kind < p2.kind || (p1.kind == p2.kind && p1.port < p2.port)))
end

""" A wire connecting one port to another.
"""
@auto_hash_equals struct Wire{Value}
  value::Value
  source::Port
  target::Port
end

Wire(value, src::Tuple{Int,PortKind,Int}, tgt::Tuple{Int,PortKind,Int}) =
  Wire(value, Port(src[1],src[2],src[3]), Port(tgt[1],tgt[2],tgt[3]))
Wire(value, src::Tuple{Int,Int}, tgt::Tuple{Int,Int}) =
  Wire(value, Port(src[1],OutputPort,src[2]), Port(tgt[1],InputPort,tgt[2]))
Wire(value, pair::Pair) = Wire(value, first(pair), last(pair))

Wire(src::Port, tgt::Port) = Wire(nothing, src, tgt)
Wire(src::Tuple, tgt::Tuple) = Wire(nothing, src, tgt)
Wire(pair::Pair) = Wire(nothing, first(pair), last(pair))

function Base.show(io::IO, wire::Wire)
  skip_kind = wire.source.kind == OutputPort && wire.target.kind == InputPort
  show_port = (io::IO, port::Port) -> begin
    if skip_kind
      print(io, "($(port.box),$(port.port))")
    else
      print(io, "($(port.box),$(string(port.kind)),$(port.port)")
    end
  end
  print(io, "Wire(")
  if !isnothing(wire.value)
    show(io, wire.value)
    print(io, ", ")
  end
  show_port(io, wire.source)
  print(io, " => ")
  show_port(io, wire.target)
  print(io, ")")
end

function Base.isless(w1::Wire, w2::Wire)::Bool
  # Lexicographic order.
  isless(w1.source, w2.source) ||
    (w1.source == w2.source && isless(w1.target, w2.target))
end

""" Internal wiring diagram data corresponding to `Port`. Do not use directly.
"""
@auto_hash_equals struct PortData
  kind::PortKind
  port::Int
end
PortData(port::Port) = PortData(port.kind, port.port)
Port(port::PortData, v::Int) = Port(v, port.kind, port.port)

""" Internal wiring diagram data corresponding to `Wire`. Do not use directly.
"""
@auto_hash_equals struct WireData{Value}
  value::Value
  source::PortData
  target::PortData
end
function WireData(wire::Wire)
  WireData(wire.value, PortData(wire.source), PortData(wire.target))
end
function Wire(wire::WireData, src::Int, tgt::Int)
  Wire(wire.value, Port(wire.source, src), Port(wire.target, tgt))
end

""" Base type for any box (node) in a wiring diagram.

This type represents an arbitrary black box with inputs and outputs.
"""
abstract type AbstractBox end

input_ports(box::AbstractBox)::Vector = box.input_ports
output_ports(box::AbstractBox)::Vector = box.output_ports

""" An atomic box in a wiring diagram.

These boxes have no internal structure.
"""
@auto_hash_equals struct Box{Value} <: AbstractBox
  value::Value
  input_ports::Vector
  output_ports::Vector
end

Box(inputs::Vector, outputs::Vector) = Box(nothing, inputs, outputs)

function Base.show(io::IO, box::Box)
  print(io, "Box(")
  if !isnothing(box.value)
    show(io, box.value)
    print(io, ", ")
  end
  print(io, "[")
  join(io, [sprint(show, port) for port in box.input_ports], ",")
  print(io, "], [")
  join(io, [sprint(show, port) for port in box.output_ports], ",")
  print(io, "])")
end

@present TheoryWiringDiagramGraph <: TheoryGraph begin
  Box::Data
  WireData::Data

  box::Attr(V,Box)
  wire::Attr(E,WireData)
end

const WiringDiagramGraphUnionAll =
  ACSetType(TheoryWiringDiagramGraph, index=[:src, :tgt])

""" Internal datatype for graph underlying a directed wiring diagram.

Boxes and wires are attached to vertices and edges, respectively.
"""
const WiringDiagramGraph = WiringDiagramGraphUnionAll{
  Union{AbstractBox,Nothing},WireData}

""" A directed wiring diagram, also known as a string diagram.

The wiring diagram is implemented using the following internal data structure.
The "skeleton" of the diagram is an instance of `Catlab.Graphs.AbstractGraph`: a
directed multigraph whose vertices correspond to boxes and whose edges
correspond to wires. There are two special vertices, accessible via `input_id`
and `output_id`, that represent the input and output ports of the outer box.
"""
mutable struct WiringDiagram{Theory} <: AbstractBox
  graph::WiringDiagramGraph
  value::Any
  input_ports::Vector
  output_ports::Vector

  function WiringDiagram{T}(value, inputs::Vector, outputs::Vector) where T
    graph = WiringDiagramGraph()
    diagram = new{T}(graph, value, inputs, outputs)
    add_vertices!(graph, 2, box=nothing)
    return diagram
  end
  function WiringDiagram(d::WiringDiagram{T}) where T
    # Copy constructor for shallow copy.
    new{T}(copy(d.graph), d.value, d.input_ports, d.output_ports)
  end
end

function WiringDiagram{T}(inputs::Vector, outputs::Vector) where T
  WiringDiagram{T}(nothing, inputs, outputs)
end
WiringDiagram(args...) = WiringDiagram{Any}(args...)

input_id(::WiringDiagram) = 1
output_id(::WiringDiagram) = 2
outer_ids(::WiringDiagram) = (1,2)

""" Check equality of wiring diagrams.

Warning: This method checks equality of the underlying graph representation, not
mathematical equality which involves graph isomorphism.

See also: `is_permuted_equal`
"""
function Base.:(==)(d1::WiringDiagram, d2::WiringDiagram)
  (input_ports(d1) == input_ports(d2) &&
   output_ports(d1) == output_ports(d2) && d1.value == d2.value &&
   boxes(d1) == boxes(d2) && sort!(wires(d1)) == sort!(wires(d2)))
end

""" Check equality of wiring diagram under permutation of boxes.

When the boxes in the first diagram `d1` are permuted according to `σ`,
does it become identical to the second diagram `d2`?
"""
function is_permuted_equal(d1::WiringDiagram, d2::WiringDiagram, σ::Vector{Int})
  @assert nboxes(d1) == length(σ) && nboxes(d2) == length(σ)
  d1_ids, d2_ids = box_ids(d1), box_ids(d2)
  box_map = Dict{Int,Int}(d1_ids[σ[i]] => d2_ids[i] for i in eachindex(σ))
  is_induced_equal(d1, d2, box_map)
end
function is_induced_equal(d1::WiringDiagram, d2::WiringDiagram, box_map::Dict{Int,Int})
  box_map[input_id(d1)] = input_id(d2)
  box_map[output_id(d1)] = output_id(d2)
  map_wire = wire::Wire -> Wire(wire.value,
    set_box(wire.source, box_map[wire.source.box]),
    set_box(wire.target, box_map[wire.target.box]))
  (input_ports(d1) == input_ports(d2) && output_ports(d1) == output_ports(d2) &&
   all(box(d1,v) == box(d2,box_map[v]) for v in box_ids(d1)) &&
   sort!(map(map_wire, wires(d1))) == sort!(wires(d2)))
end

Base.copy(diagram::WiringDiagram) = WiringDiagram(diagram)

function Base.show(io::IO, diagram::WiringDiagram{T}) where T
  sshowcompact = x -> sprint(show, x, context=:compact => true)
  print(io, "WiringDiagram")
  if T != Any
    print(io, "{$T}")
  end
  print(io, "(")
  if !isnothing(diagram.value)
    show(io, diagram.value)
    print(io, ", ")
  end
  print(io, "[")
  join(io, map(sshowcompact, input_ports(diagram)), ",")
  print(io, "], [")
  join(io, map(sshowcompact, output_ports(diagram)), ",")
  print(io, "], ")
  if get(io, :compact, false)
    print(io, "{$(nboxes(diagram)) boxes}, {$(nwires(diagram)) wires}")
  else
    print(io, "\n[ $(input_id(diagram)) => {inputs},\n  ")
    print(io, "$(output_id(diagram)) => {outputs},\n  ")
    join(io, [ "$v => $(sshowcompact(box(diagram, v)))"
               for v in box_ids(diagram) ], ",\n  ")
    print(io, " ],\n[ ")
    join(io, map(sshowcompact, wires(diagram)), ",\n  ")
    print(io, " ]")
  end
  print(io, ")")
end

# Imperative interface
######################

# Basic accessors.

box(f::WiringDiagram, v::Int) = subpart(f.graph, v, :box)

function boxes(f::WiringDiagram)
  collect(AbstractBox, subpart(f.graph, box_ids(f), :box))
end

nboxes(f::WiringDiagram) = nv(f.graph) - 2

function box_ids(f::WiringDiagram)
  Int[ v for v in 1:nv(f.graph) if !(v in outer_ids(f)) ]
end

function wires(f::WiringDiagram, src::Int, tgt::Int)
  [ Wire(subpart(f.graph, e, :wire), src, tgt)
    for e in edges(f.graph, src, tgt) ]
end
function wires(f::WiringDiagram)
  g = f.graph
  [ Wire(subpart(g, e, :wire), src(g, e), tgt(g, e)) for e in edges(g) ]
end
nwires(f::WiringDiagram) = ne(f.graph)

has_wire(f::WiringDiagram, src::Int, tgt::Int) =
  !isempty(edges(f.graph, src, tgt))
has_wire(f::WiringDiagram, wire::Wire) =
  wire in wires(f, wire.source.box, wire.target.box)
has_wire(f::WiringDiagram, pair::Pair) = has_wire(f, Wire(pair))

function input_ports(f::WiringDiagram, v::Int)
  if v == input_id(f)
    error("Input vertex does not have input ports within wiring diagram")
  elseif v == output_id(f)
    output_ports(f)
  else
    input_ports(box(f, v))
  end
end

function output_ports(f::WiringDiagram, v::Int)
  if v == input_id(f)
    input_ports(f)
  elseif v == output_id(f)
    error("Output vertex does not have output ports within wiring diagram")
  else
    output_ports(box(f, v))
  end
end

function port_value(f::WiringDiagram, port::Port)
  get_ports = port.kind == InputPort ? input_ports : output_ports
  get_ports(f, port.box)[port.port]
end

# Graph mutation.

function add_box!(f::WiringDiagram, box::AbstractBox)
  add_vertex!(f.graph, box=box)
  return nv(f.graph)
end

function add_boxes!(f::WiringDiagram, boxes)
  boxes = collect(boxes)
  add_vertices!(f.graph, length(boxes), box=boxes)
end

function rem_box!(f::WiringDiagram, v::Int)
  @assert v ∉ outer_ids(f)
  rem_vertex!(f.graph, v)
end

function rem_boxes!(f::WiringDiagram, vs)
  @assert all(v ∉ outer_ids(f) for v in vs)
  rem_vertices!(f.graph, sort!(collect(vs)))
end

function add_wire!(f::WiringDiagram, wire::Wire)
  validate_ports(port_value(f, wire.source), port_value(f, wire.target))
  add_edge!(f.graph, wire.source.box, wire.target.box, wire=WireData(wire))
end
add_wire!(f::WiringDiagram, pair::Pair) = add_wire!(f, Wire(pair))

function add_wires!(f::WiringDiagram, wires)
  for wire in wires
    add_wire!(f, wire)
  end
end

function rem_wire!(f::WiringDiagram, wire::Wire)
  g, wire_data = f.graph, WireData(wire)
  for e in edges(g, wire.source.box, wire.target.box)
    subpart(g, e, :wire) == wire_data && return rem_edge!(g, e)
  end
  error("Wire $wire does not exist, so cannot be removed")
end
rem_wire!(f::WiringDiagram, pair::Pair) = rem_wire!(f, Wire(pair))

function rem_wires!(f::WiringDiagram, wires)
  for wire in wires
    rem_wire!(f, wire)
  end
end

function rem_wires!(f::WiringDiagram, src::Int, tgt::Int)
  rem_edges!(f.graph, sort!(collect(edges(f.graph, src, tgt))))
end

""" Check compatibility of source and target ports.

The default implementation is a no-op.
"""
function validate_ports(source_port, target_port) end

# Graph properties.

""" Retrieve the graph underlying the wiring diagram.

The graph is an instance of `Catlab.Graphs.AbstractGraph`. Do not mutate it! All
mutations should use the wiring diagrams API: `add_box!`, `rem_box!`, and so on.
"""
graph(diagram::WiringDiagram) = diagram.graph

# Convenience methods delegated to underlying graph.
all_neighbors(d::WiringDiagram, v::Int) = all_neighbors(graph(d), v)
neighbors(d::WiringDiagram, v::Int) = neighbors(graph(d), v)
outneighbors(d::WiringDiagram, v::Int) = outneighbors(graph(d), v)
inneighbors(d::WiringDiagram, v::Int) = inneighbors(graph(d), v)

""" Get all wires coming into or out of the box.
"""
function wires(d::WiringDiagram, v::Int)
  g = graph(d)
  [ Wire(subpart(g, e, :wire), src(g, e), tgt(g, e))
    for e in unique!(sort!([incident(g, v, :src); incident(g, v, :tgt)])) ]
end

""" Get all wires coming into the box.
"""
function in_wires(d::WiringDiagram, v::Int)
  g = graph(d)
  [ Wire(subpart(g, e, :wire), src(g, e), v) for e in incident(g, v, :tgt) ]
end

""" Get all wires coming into the port.
"""
function in_wires(d::WiringDiagram, port::Port)
  filter(wire -> wire.target == port, in_wires(d, port.box))
end
function in_wires(d::WiringDiagram, v::Int, port::Int)
  in_wires(d, Port(v, InputPort, port))
end

""" Get all wires coming out of the box.
"""
function out_wires(d::WiringDiagram, v::Int)
  g = graph(d)
  [ Wire(subpart(g, e, :wire), v, tgt(g, e)) for e in incident(g, v, :src) ]
end

""" Get all wires coming out of the port.
"""
function out_wires(d::WiringDiagram, port::Port)
  filter(wire -> wire.source == port, out_wires(d, port.box))
end
function out_wires(d::WiringDiagram, v::Int, port::Int)
  out_wires(d, Port(v, OutputPort, port))
end

# Other constructors
#-------------------

""" Wiring diagram with a single box connected to the outer ports.
"""
function singleton_diagram(T::Type, box::AbstractBox)
  inputs, outputs = input_ports(box), output_ports(box)
  d = WiringDiagram{T}(inputs, outputs)
  v = add_box!(d, box)
  add_wires!(d, ((input_id(d),i) => (v,i) for i in eachindex(inputs)))
  add_wires!(d, ((v,i) => (output_id(d),i) for i in eachindex(outputs)))
  return d
end
singleton_diagram(box::AbstractBox) = singleton_diagram(Any, box)

""" The wiring diagram induced by a subset of its boxes.

See also `encapsulated_subdiagram`.
"""
function induced_subdiagram(d::WiringDiagram{T}, vs::Vector{Int}) where T
  sub = WiringDiagram{T}(input_ports(d), output_ports(d))
  vmap = Dict(input_id(d) => input_id(sub), output_id(d) => output_id(sub))
  for v in vs
    vmap[v] = add_box!(sub, box(d, v))
  end
  for wire in wires(d)
    src, tgt = wire.source, wire.target
    if haskey(vmap, src.box) && haskey(vmap, tgt.box)
      add_wire!(sub,
        Wire(set_box(src, vmap[src.box]), set_box(tgt, vmap[tgt.box])))
    end
  end
  return sub
end

# Operadic interface
####################

""" Operadic composition of wiring diagrams.

This generic function has two different signatures, corresponding to the "full"
and "partial" notions of operadic composition (Yau, 2018, *Operads of Wiring
Diagrams*, Definitions 2.3 and 2.10).

This operation is a simple wrapper around [`substitute`](@ref).
"""
function ocompose(f::WiringDiagram, gs::Vector{<:WiringDiagram})
  @assert length(gs) == nboxes(f)
  substitute(f, box_ids(f), gs)
end
function ocompose(f::WiringDiagram, i::Int, g::WiringDiagram)
  @assert 1 <= i <= nboxes(f)
  substitute(f, box_ids(f)[i], g)
end

# Substitution
##############

""" Substitute wiring diagrams for boxes.

Performs one or more substitutions. When performing multiple substitutions, the
substitutions are simultaneous.

This operation implements the operadic composition of wiring diagrams, see also
[`ocompose`](@ref).
"""
function substitute(d::WiringDiagram; kw...)
  substitute(d, filter(v -> box(d,v) isa WiringDiagram, box_ids(d)); kw...)
end
function substitute(d::WiringDiagram, v::Int; kw...)
  substitute(d, v, box(d,v)::WiringDiagram; kw...)
end
function substitute(d::WiringDiagram, vs::AbstractVector{Int}; kw...)
  substitute(d, vs, WiringDiagram[ box(d,v) for v in vs ]; kw...)
end
function substitute(d::WiringDiagram, v::Int, sub::WiringDiagram; kw...)
  substitute(d, [v], [sub]; kw...)
end

function substitute(d::WiringDiagram{T}, vs::AbstractVector{Int},
                    subs::Vector{<:WiringDiagram};
                    merge_wire_values=default_merge_wire_values) where T
  # In outline, the algorithm is:
  #
  # 1. Create an empty wiring diagram.
  # 2. Add *all* boxes of original diagram and the diagrams to be substituted
  #    (in the appropriate order).
  # 3. Add *all* wires of original diagram and the diagrams to be substituted.
  # 4. Remove the boxes that were substituted (hence also removing extraneous
  #    wires from step 3).
  #
  # This may seem convoluted, but it is the simplest way I know to handle the
  # problem of *instantaneous wires*. Some authors ban instantaneous wires, but
  # we need them to represent identities, braidings, etc.
  @assert length(vs) == length(subs)
  result = WiringDiagram{T}(d.value, input_ports(d), output_ports(d))

  # Add boxes by interleaving, in the correct order, the non-substituted boxes
  # of the original diagram and the internal boxes of the substituted diagrams.
  # At the very end, add the substituted boxes too.
  vmap = Dict(input_id(d) => input_id(result), output_id(d) => output_id(result))
  sub_maps = Dict(zip(vs, ((sub, Dict{Int,Int}()) for sub in subs)))
  for v in box_ids(d)
    if haskey(sub_maps, v)
      sub, sub_map = sub_maps[v]
      for u in box_ids(sub)
        sub_map[u] = add_box!(result, box(sub, u))
      end
    else
      vmap[v] = add_box!(result, box(d, v))
    end
  end
  for v in vs
    vmap[v] = add_box!(result, box(d, v))
  end

  # Add the wires of the original diagram, then add the internal wires of the
  # substituted diagrams.
  for wire in wires(d)
    add_wire!(result, Wire(wire.value,
      set_box(wire.source, vmap[wire.source.box]),
      set_box(wire.target, vmap[wire.target.box])))
  end
  for v in vs
    substitute_wires!(result, vmap[v], sub_maps[v]...;
      merge_wire_values=merge_wire_values)
  end

  # Finally, remove the substituted boxes. Because they were added last, this
  # will not change the IDs of the other boxes.
  rem_boxes!(result, (vmap[v] for v in vs))
  result
end

""" Substitute wires of sub-diagram into containing wiring diagram.
"""
function substitute_wires!(d::WiringDiagram, v::Int,
                           sub::WiringDiagram, sub_map::Dict{Int,Int};
                           merge_wire_values=default_merge_wire_values)
  for wire in wires(sub)
    src = get(sub_map, wire.source.box, 0)
    tgt = get(sub_map, wire.target.box, 0)
    # Special case: wire from input port to output port.
    if wire.source.box == input_id(sub) && wire.target.box == output_id(sub)
      for in_wire in in_wires(d, v, wire.source.port)
        for out_wire in out_wires(d, v, wire.target.port)
          add_wire!(d, Wire(
            merge_wire_values(in_wire.value, wire.value, out_wire.value),
            in_wire.source, out_wire.target))
        end
      end
    # Special case: wire from input port to internal box.
    elseif wire.source.box == input_id(sub)
      for in_wire in in_wires(d, v, wire.source.port)
        add_wire!(d, Wire(
          merge_wire_values(in_wire.value, wire.value, nothing),
          in_wire.source, set_box(wire.target, tgt)))
      end
    # Special case: wire from internal box to output port.
    elseif wire.target.box == output_id(sub)
      for out_wire in out_wires(d, v, wire.target.port)
        add_wire!(d, Wire(
          merge_wire_values(nothing, wire.value, out_wire.value),
          set_box(wire.source, src), out_wire.target))
      end
    # Default case: wire between two internal boxes.
    else
      add_wire!(d, Wire(
        merge_wire_values(nothing, wire.value, nothing),
        set_box(wire.source, src), set_box(wire.target, tgt)))
    end
  end
end

default_merge_wire_values(::Any, middle::Any, ::Any) = middle

# Encapsulation
###############

""" Encapsulate multiple boxes within a single sub-diagram.

This operation is a (one-sided) inverse to subsitution, see
[`substitute`](@ref).
"""
function encapsulate(d::WiringDiagram, vs::Vector{Int}; value=nothing, kw...)
  encapsulate(d, [vs]; values=[value], kw...)
end

function encapsulate(d::WiringDiagram{T}, vss::Vector{Vector{Int}};
    discard_boxes::Bool=false, make_box=Box, values=nothing) where T
  if isempty(vss); return d end
  if any(isempty(vs) for vs in vss)
    error("Cannot encapsulate an empty set of boxes")
  end
  if !allunique(reduce(vcat, vss))
    error("Cannot encapsulate overlapping sets of boxes")
  end
  if isnothing(values)
    values = repeat([nothing], length(vss))
  end
  result = WiringDiagram{T}(d.value, input_ports(d), output_ports(d))

  # Add boxes, both encapsulated and non-encapsulated, to new wiring diagram.
  encapsulated_representatives = Dict(
    minimum(vs) => (vs, value) for (vs, value) in zip(vss, values))
  all_encapsulated = Set(v for vs in vss for v in vs)
  vmap = Dict(input_id(d) => input_id(result), output_id(d) => output_id(result))
  port_map = Dict{Port,Port}()
  for v in box_ids(d)
    if haskey(encapsulated_representatives, v)
      vs, value = encapsulated_representatives[v]
      sub, sub_map = encapsulated_subdiagram(d, vs;
        discard_boxes=discard_boxes, make_box=make_box, value=value)
      subv = add_box!(result, sub)
      merge!(port_map, Dict(port => Port(data, subv)
                            for (port, data) in sub_map))
    elseif v ∉ all_encapsulated
      vmap[v] = add_box!(result, box(d, v))
    end
  end

  # Add wires to new wiring diagram.
  for wire in wires(d)
    src, tgt = wire.source, wire.target
    new_src = if haskey(vmap, src.box); set_box(src, vmap[src.box])
      elseif haskey(port_map, src); port_map[src]
      else; continue end
    new_tgt = if haskey(vmap, tgt.box); set_box(tgt, vmap[tgt.box])
      elseif haskey(port_map, tgt); port_map[tgt]
      else; continue end
    add_wire!(result, Wire(new_src, new_tgt))
  end
  result
end

""" Create an encapsulating box for a set of boxes in a wiring diagram.

To a first approximation, the union of input ports of the given boxes will
become the inputs ports of the encapsulating box and likewise for the output
ports. However, when copies or merges occur, as in a cartesian or cocartesian
category, a simplification procedure may reduce the number of ports on the
encapsulating box.

Specifically:

1. Each input port of an encapsulated box will have at most one incoming wire
from the encapsulating outer box, and each output port of an encapsulated box
will have at most one outgoing wire to the encapsulating outer box.

2. A set of ports connected to the same outside (non-encapsulated) ports will be
simplified into a single port of the encapsulating box.

See also `induced_subdiagram`.
"""
function encapsulated_subdiagram(d::WiringDiagram{T}, vs::Vector{Int};
    discard_boxes::Bool=false, make_box=Box, value=nothing) where T
  # Add encapsulated box to new wiring diagram.
  inputs, outputs = [], []
  result = discard_boxes ? nothing : WiringDiagram{T}(value, inputs, outputs)
  vmap = if discard_boxes
    Dict(v => nothing for v in vs)
  else
    Dict(v => add_box!(result, box(d, v)) for v in vs)
  end

  # Process wires into, or out of, encapsulated boxes.
  port_map = Dict{Port,PortData}()
  inner_port_map = Dict{Tuple{Vector{Port},Any},Int}()
  for v in vs
    # Add input ports to encapsulating box and corresponding wire.
    for (port, value) in enumerate(input_ports(d, v))
      tgt = Port(v, InputPort, port)
      srcs = sort!([ wire.source for wire in in_wires(d, tgt)
                     if !haskey(vmap, wire.source.box) ])
      if isempty(srcs) continue end
      src = get!(inner_port_map, (srcs, value)) do
        push!(inputs, value)
        port_data = port_map[tgt] = PortData(InputPort, length(inputs))
        port_data.port
      end
      if discard_boxes; continue end
      add_wire!(result,
        Wire(Port(input_id(result), OutputPort, src), set_box(tgt, vmap[v])))
    end

    # Add output ports to encapsulating box and corresponding wire.
    for (port, value) in enumerate(output_ports(d, v))
      src = Port(v, OutputPort, port)
      tgts = sort([ wire.target for wire in out_wires(d, src)
                    if !haskey(vmap, wire.target.box) ])
      if isempty(tgts) continue end
      tgt = get!(inner_port_map, (tgts, value)) do
        push!(outputs, value)
        port_data = port_map[src] = PortData(OutputPort, length(outputs))
        port_data.port
      end
      if discard_boxes; continue end
      add_wire!(result,
        Wire(set_box(src, vmap[v]), Port(output_id(result), InputPort, tgt)))
    end

    # Add wires between two encapsulated boxes.
    if discard_boxes; continue end
    for wire in out_wires(d, v)
      src, tgt = wire.source, wire.target
      if haskey(vmap, src.box) && haskey(vmap, tgt.box) # Clause #1 always true.
        add_wire!(result,
          Wire(set_box(src, vmap[src.box]), set_box(tgt, vmap[tgt.box])))
      end
    end
  end

  # Yield input and output port lists with the tightest possible types.
  inputs, outputs = [ x for x in inputs ], [ x for x in outputs ]
  if discard_boxes
    result = make_box(value, inputs, outputs)
  else
    result.input_ports, result.output_ports = inputs, outputs
  end
  (result, port_map)
end

end


const MJSTATE_FIELDS = (:time, :qpos, :qvel, :act, :mocap_pos, :mocap_quat, :userdata, :qacc_warmstart)
const MJACTION_FIELDS = (:ctrl, :qfrc_applied, :xfrc_applied)

const DEFAULT_SKIP = 1
const DEFAULT_CAN_SHARE_MODEL = true

"""
    MJSim

The `MJSim` type couples a `jlModel` and `jlData` from MuJoCo.jl to provide a full simulation.

The following are the official/internal/minimum set of fields from `jlData` for
state, observation, and action in MuJoCo:

- State: `(time, qpos, qvel, act, mocap_pos, mocap_quat, userdata, qacc_warmstart)`
- Observation: `sensordata`
- Action: `(ctrl, qfrc_applied, xfrc_applied)`

MJSim follows this definition except for actions (e.g. `setaction!``), which is
composed of just `ctrl` by default.

For more information, see the "State and control" section of the
[MuJoCo Documentation](www.mujoco.org/book/programming.html)

# Fields

- `m::jlModel`: contains the model description and is expected to remain constant.
- `d::jlData`: contains all dynamic variables and intermediate results.
- `mn::Tuple`: named-access version of `MJSim.m` provided by AxisArrays.jl.
- `dn::Tuple`: named-access version of `MJSim.d` provided by AxisArrays.jl.
- `initstate::Vector{mjtNum}`: The initial state vector at the time when this MJSim was constructed.
- `skip::Int`: the number of times the simulation is integrated, yielding an
  effective simulation timestep of `skip * m.opt.timestep`.
"""
struct MJSim{MN, DN, S, SE, A}
    m::jlModel
    d::jlData
    mn::MN
    dn::DN
    initstate::Vector{mjtNum}
    skip::Int

    # internal
    statespace::S
    observationspace::SE
    actionspace::A

    @doc """
        $(TYPEDSIGNATURES)

    Construct an MJSim from `m` and `d`, with skip `skip`.
    """
    function MJSim(m::jlModel, d::jlData; skip::Integer = DEFAULT_SKIP)
        check_skip(skip)

        m_named, d_named = namify(m, d)

        nameshapes = map(MJSTATE_FIELDS) do name
            field = getproperty(d, name)
            name => Shape(eltype(field), size(field)...)
        end
        statespace = MultiShape(nameshapes...)

        if m.nsensor > 0
            nameshapes = map(1:m.nsensor) do id
              name = jl_id2name(m, MJCore.mjOBJ_SENSOR, id)
              name = isnothing(name) ? Symbol("unnamed_$id") : Symbol(name)
              dof = m.sensor_dim[id]
              shape = dof == 1 ? ScalarShape(mjtNum) : VectorShape(mjtNum, dof)
              name => shape
            end
            observationspace = MultiShape(nameshapes...)
        else
            observationspace = VectorShape(mjtNum, 0)
        end

        if m.nu > 0
            clampctrl = !jl_disabled(m, MJCore.mjDSBL_CLAMPCTRL)
            nameshapes = map(1:m.nu) do id
                name = jl_id2name(m, MJCore.mjOBJ_ACTUATOR, id)
                name = isnothing(name) ? Symbol("unnamed_$id") : Symbol(name)
                name => ScalarShape(mjtNum) #TODO boundedshape
            end
            actionspace = MultiShape(nameshapes...)
        else
            actionspace = VectorShape(mjtNum, 0)
        end

        initstate = mapreduce(vcat, MJSTATE_FIELDS) do name
            field = getproperty(d, name)
            ndims(field) > 1 ? vec(field) : field
        end

    	sim = new{typeof(m_named), typeof(d_named), typeof(statespace), typeof(observationspace), typeof(actionspace)}(m, d, m_named, d_named,
            initstate, skip,
            statespace, observationspace, actionspace)

        forward!(sim)
    end
end

"""
    $(TYPEDSIGNATURES)

Construct an `MJSim` the MJCF or MJB model file located at `modelpath`, with skip `skip`.
"""
function MJSim(modelpath::String; skip::Integer = DEFAULT_SKIP)
    m = jlModel(modelpath)
    d = jlData(m)
    MJSim(m, d, skip=skip)
end

function LyceumBase.tconstruct(::Type{MJSim}, N::Integer, m::jlModel; skip::Integer = DEFAULT_SKIP, can_share_model::Bool=DEFAULT_CAN_SHARE_MODEL)
    N > 0 || throw(ArgumentError("N must be > 0"))
    if can_share_model
        ntuple(_ -> MJSim(m, jlData(m), skip=skip), N)
    else
        ntuple(N) do _
            mcopy = copy(m)
            MJSim(mcopy, jlData(mcopy), skip=skip)
        end
    end
end

function LyceumBase.tconstruct(::Type{MJSim}, N::Integer, modelpath::AbstractString; kwargs...)
    tconstruct(MJSim, N, jlModel(modelpath); kwargs...)
end

"""
    $(TYPEDSIGNATURES)

Return a description of `sim`'s statespace.
"""
@inline LyceumBase.statespace(sim::MJSim) = sim.statespace

"""
    $(SIGNATURES)

Copy the following state fields from `sim.d` into `state`:
```
($(join(map(string, MJSTATE_FIELDS), ", ")))
```
"""
@propagate_inbounds function LyceumBase.getstate!(state::RealVec, sim::MJSim)
    @boundscheck checkaxes(statespace(sim), state)
    shaped = statespace(sim)(state)
    @uviews shaped begin _copyshaped!(shaped, sim.d) end
    state
end

"""
    $(SIGNATURES)

Return a flattened vector of the following state fields from `sim.d`:
```
($(join(map(string, MJSTATE_FIELDS), ", ")))
```
"""
@propagate_inbounds LyceumBase.getstate(sim::MJSim) = getstate!(allocate(statespace(sim)), sim)

# Two things to note about MuJoCo state:
# 1. MJSTATE_FIELDS are theoretically the minimal fields/sufficient statistics, as in
#    all other fields in mjData are a function of MJSTATE_FIELDS, but this is only true for
#    **dynamic** elements. For example, calling rand!(d.xpos) followed by mj_forward(m, d) will
#    overwrite the entries of d.xpos corresponding to dynamically changing elements, but NOT for
#    e.g. the root world body (d.xpos[:, 1]). If the user inadvertantly overwrites these fields,
#    mjData will be left in a corrupt state. Because of this, we call mj_resetData(m, d) followed
#    by mj_forward(m, d) inside the exported setstate!, while providing the ability to
#    skip this using copystate! for users who know what they're doing and want to squeeze some
#    extra performance out.
# 2. qacc_warmstart is part of state in that it will control the final solution (because MuJoCo's
#    optimizer may not find the exact minima). This means that:
#       setstate!(sim, state)
#       @assert getstate(sim) == state
#    will fail if the optimizer terminated early. To avoid this, we first copy over the state
#    variables (MJSTATE_FIELDS, including qacc_warmstart), call mj_forward, and then re-copy
#    over qacc_warmstart so that the above equality holds true. This means that mjData is left
#    in a slightly inconsitent state (qacc_warmstart will be slightly out of sync).

"""
    $(SIGNATURES)

Copy the components of `state` to their respective fields in `sim.d`, namely:
```
($(join(map(string, MJSTATE_FIELDS), ", ")))
```
"""
@propagate_inbounds function LyceumBase.setstate!(sim::MJSim, state::RealVec)
    mj_resetData(sim.m, sim.d)
    copystate!(sim, state)
    forward!(sim)
    shaped = statespace(sim)(state)
    @uviews shaped begin copyto!(sim.d.qacc_warmstart, shaped.qacc_warmstart) end
    sim
end

@propagate_inbounds function copystate!(sim::MJSim, state::RealVec)
    @boundscheck checkaxes(statespace(sim), state)
    shaped = @inbounds statespace(sim)(state)
    @uviews shaped begin _copyshaped!(sim.d, shaped) end
    sim
end

@inline function _copyshaped!(dst, src)
    @inbounds dst.time = src.time
    @inbounds copyto!(dst.qpos, src.qpos)
    @inbounds copyto!(dst.qvel, src.qvel)
    @inbounds copyto!(dst.act, src.act)
    @inbounds copyto!(dst.mocap_pos, src.mocap_pos)
    @inbounds copyto!(dst.mocap_quat, src.mocap_quat)
    @inbounds copyto!(dst.userdata, src.userdata)
    @inbounds copyto!(dst.qacc_warmstart, src.qacc_warmstart)
    dst
end

"""
    $(TYPEDSIGNATURES)

Return a description of `sim`'s observation space.
"""
@inline LyceumBase.observationspace(sim::MJSim) = sim.observationspace

"""
    $(SIGNATURES)

Copy `sim.d.sensordata` into `obs`.
"""
@propagate_inbounds function LyceumBase.getobservation!(obs::RealVec, sim::MJSim)
    @boundscheck checkaxes(observationspace(sim), obs)
    @inbounds copyto!(obs, sim.d.sensordata)
end

"""
    $(TYPEDSIGNATURES)

Return a copy of `sim.d.sensordata`.
"""
@propagate_inbounds LyceumBase.getobservation(sim::MJSim) = getobservation!(allocate(sim.observationspace), sim)

"""
    $(TYPEDSIGNATURES)

Return a description of `sim`'s action space.
"""
@inline LyceumBase.actionspace(sim::MJSim) = sim.actionspace

"""
    $(TYPEDSIGNATURES)

Copy `sim.d.ctrl` into `action`.
"""
@propagate_inbounds function LyceumBase.getaction!(action::RealVec, sim::MJSim)
    @boundscheck checkaxes(actionspace(sim), action)
    @inbounds copyto!(action, sim.d.ctrl)
end

"""
    $(TYPEDSIGNATURES)

Return a copy of `sim.d.ctrl`.
"""
@propagate_inbounds LyceumBase.getaction(sim::MJSim, action::RealVec) = getaction!(allocate(sim.actionspace), sim)

"""
    $(TYPEDSIGNATURES)

Copy `action` into `sim.d.ctrl` into `action` and compute the new forward dynamics.
"""
@propagate_inbounds function LyceumBase.setaction!(sim::MJSim, action::RealVec)
    forwardskip!(setaction_nofwd!(sim, action), MJCore.mjSTAGE_VEL)
end

@propagate_inbounds function setaction_nofwd!(sim::MJSim, action::RealVec)
    @boundscheck checkaxes(actionspace(sim), action)
    @inbounds copyto!(sim.d.ctrl, action)
    sim
end

@inline LyceumBase.reset!(sim::MJSim) = forward!(reset_nofwd!(sim))
@inline reset_nofwd!(sim::MJSim) = (mj_resetData(sim.m, sim.d); sim)

# typically 2-3x faster than reset!
@inline fastreset!(sim::MJSim) = forward!(fastreset_nofwd!(sim))
@inline function fastreset_nofwd!(sim::MJSim)
    zerofullctrl_nofwd!(sim)
    @inbounds copystate!(sim, sim.initstate)
end

"""
    $(TYPEDSIGNATURES)

Step the simulation by `skip` steps, where `skip` defaults to `sim.skip`.

State-dependent controls (e.g. the `ctrl`, `xfrc_applied`, `qfrc_applied`
fields of `sim.d`) should be set before calling `step!`.
"""
function LyceumBase.step!(sim::MJSim, skip::Integer=sim.skip)
    check_skip(skip)
    # According to MuJoCo docs order must be:
    # 1. mj_step1
    # 2. set_{ctrl, xfrc_applied, qfrc_applied}
    # 3. mj_step2
    # This implies user should call setaction! then step!
    for _=1:skip
        mj_step2(sim.m, sim.d)
        mj_step1(sim.m, sim.d)
    end
    sim
end

"""
    $(TYPEDSIGNATURES)

Zero out `sim.d.ctrl` and compute new forward dynamics.
"""
@inline zeroctrl!(sim::MJSim) = forwardskip!(zeroctrl_nofwd!(sim), MJCore.mjSTAGE_VEL)

@inline zeroctrl_nofwd!(sim::MJSim) = (fill!(sim.d.ctrl, zero(mjtNum)); sim)

"""
    $(TYPEDSIGNATURES)

Zero out all of the fields in `sim.d` that contribute to forward dynamics calculations,
namely `ctrl`, `qfrc_applied`, and `xfrc_applied`, and compute the forward dynamics.
"""
@inline zerofullctrl!(sim::MJSim) = forwardskip!(zerofullctrl_nofwd!(sim), MJCore.mjSTAGE_VEL)

@inline function zerofullctrl_nofwd!(sim::MJSim)
    fill!(sim.d.ctrl, zero(mjtNum))
    fill!(sim.d.qfrc_applied, zero(mjtNum))
    fill!(sim.d.xfrc_applied, zero(mjtNum))
    sim
end

"""
    $(TYPEDSIGNATURES)

Calculate the center of mass of all the bodies in simulation `sim` by computing the
mass-weighted sum of the Cartesian coordinates of each body in the world frame.
"""
function masscenter(sim::MJSim)
    mcntr = zeros(SVector{3, Float64})
    mtotal = 0.0
    body_mass, xipos = sim.m.body_mass, sim.d.xipos
    @uviews xipos for i=1:sim.m.nbody
        bmass = body_mass[i]
        bodycom = SVector{3, Float64}(uview(xipos, :, i))
        mcntr = mcntr + bmass * bodycom
        mtotal += bmass
    end
    mcntr = SVector(mcntr)
    mcntr / mtotal
end

"""
    $(TYPEDSIGNATURES)

Compute the forward dynamics of the simulation and store them in `sim.d`.
Equivalent to `mj_forward(sim.m, sim.d)`.
"""
@inline forward!(sim::MJSim) = (mj_forward(sim.m, sim.d); sim)

"""
    $(TYPEDSIGNATURES)

Compute the forward dynamics of the simulation and store them in `sim.d`, optionally
skipping parts of the dynamics calculation. Equivalent to `mj_forwardSkip(sim.m, sim.d,
skipstage, skipsensor)`.
"""
@inline function forwardskip!(sim::MJSim, skipstage::MJCore.mjtStage=MJCore.mjSTAGE_NONE, skipsensor::Bool=false)
    mj_forwardSkip(sim.m, sim.d, skipstage, skipsensor)
    sim
end

"""
    $(TYPEDSIGNATURES)

Return the effective timestep of `sim`. Equivalent to `sim.skip * sim.m.opt.timestep`.
"""
@inline LyceumBase.timestep(sim::MJSim) = sim.m.opt.timestep * sim.skip

"""
    $(TYPEDSIGNATURES)

Return the current simulation time, in seconds, of `sim`. Equivalent to sim.d.time.
"""
@inline Base.time(sim::MJSim) = sim.d.time

@inline getsim(sim::MJSim) = sim

# TODO(cxs):: Make more informative. This is a temp fix for those gawdawful error messages from sim.{mn, dn}.
Base.show(io::IO, ::MIME"text/plain", sim::Union{<:MJSim, Type{<:MJSim}}) = show(io, sim)
Base.show(io::IO, sim::Union{<:MJSim, Type{<:MJSim}}) = print(io, "MJSim")

@inline check_skip(skip) = skip > 0 || throw(ArgumentError("`skip` must be > 0"))


using Random
using Autumn
include("generativemodel.jl")

a = au"""(program
      (= GRID_SIZE 16)

      (object Particle (Cell 0 0 "blue"))

      (: particles (List Particle))
      (= particles
        (initnext (list)
                  (updateObj (prev particles) (--> obj (Particle (uniformChoice (adjPositions (.. obj origin))))))))

      (on clicked (= particles (addObj (prev particles) (Particle (Position (.. click x) (.. click y))))))
      )"""

mod = eval(compiletojulia(a))

rng = MersenneTwister(0)
state = mod.init(nothing, nothing, nothing, nothing, nothing, rng);

# add particles
for i in 1:3
  global state = mod.next(state, mod.Click(rand(0:15),rand(0:15)), nothing, nothing, nothing, nothing); mod.render(state.scene)
end

# let particles move
for i in 1:5
  global state = mod.next(state, mod.Click(rand(0:15),rand(0:15)), nothing, nothing, nothing, nothing); mod.render(state.scene)
end

# render
rendering = mod.render(state.scene)

# parse scene rendering into types and objects
types_and_objects = parsescene_autumn_singlecell(rendering)

# # generate random dynamics on object decomposition (types and objects)
# println(generateprogram_given_objects(types_and_objects, group=true))


module Behavior

export i_dt, i_dt_wrapper, identify_fixatedTargets, plot_eye_events

using DataFrames
using MAT
using JLD
using Dates
using LinearAlgebra
using DSP
using Plots
using Statistics
using KernelDensity
using StatsBase
using Images

using ..MyUtils

### functions

#################################################################################
#################################################################################

function i_dt_wrapper(
    df,
    sf,
    fix_duration_threshold,
    fix_dispersion_threshold,
    saccade_velocity_threshold,
    saccade_duration_threshold;
    plot_saccades=false,
    plot_fixations=false,
    which_trials="all",
    how_ = "mean",
    which="Fixations_clean"
    )

    ####
    # functions takes a dataframe containg trial data of an experimental session as input
    # and returns the dataframe with filled :Fixsation and :Saccades columns
    ####

    for trl in range(1, size(df, 1), step=1)

        xy_, fix_idx, sac_idx, blink_idx = i_dt( df[!, "EyeXY"][trl], sf, fix_duration_threshold, fix_dispersion_threshold, saccade_velocity_threshold, saccade_duration_threshold);#, blink_velocity_threshold );
        df.Fixations[trl] = fix_idx
        df.Saccades[trl] = sac_idx

    end

end

function plot_eye_events(
    df;
    what="both", # fixations / saccades
    how_="mean",
    which_fixations="all_fixations",
    p=plot(legend=false))

    if what == "both" || what == "saccades"
        ### plot saccades
        for (itrl,trls_saccades) in enumerate(df[:Saccades])
            for sac_idxs in trls_saccades
                xx = df[itrl,:EyeXY][1][sac_idxs]
                yy = df[itrl,:EyeXY][2][sac_idxs]
                plot!(xx,yy)
            end
        end
    end

    if which_fixations == "all_fixations"
        clean_fixations_bool = [[fix in [-1,0,1,2,3,4,5,6,7,8] for fix in el] for el in df.LookedAtTarget]
    elseif which_fixations == "clean_fixations"
        clean_fixations_bool = [[fix in [0,1,2,3,4,5,6,7,8] for fix in el] for el in df.LookedAtTarget]
    end

    if what == "both" || what == "fixations"
        for (itrl,trls_fix) in enumerate(df.Fixations)
            for (ifix, fix) in enumerate(trls_fix)
                if clean_fixations_bool[itrl][ifix]
                    scatter!( [mean( df.EyeXY[itrl][1][fix] )], [mean( df.EyeXY[itrl][2][fix] )] )
                end
            end
        end

        # for (itrl,trls_fixations) in enumerate(df[:Fixations])
        #     for fix_idxs in trls_fixations
        #         # print(df[itrl,:EyeXY][1][sacd])
        #         xx = df[itrl,:EyeXY][1][fix_idxs]
        #         yy = df[itrl,:EyeXY][2][fix_idxs]
        #         if how_ == "mean"
        #             scatter!([mean(xx)],[mean(yy)])
        #         elseif how_ == "raw"
        #             scatter!(xx,yy, markersize=0.5)
        #         end
        #     end
        # end
    end

    xlims!(-15,15)
    ylims!(-15,15)
    return p
end

#################################################################################
#################################################################################

function identify_fixatedTargets(
    df;
    nbins_length=150,
    nbins_angles=150,
    kde_bandwidth_length=1/3,
    kde_bandwidth_angle=1/15,
    sf=2000,
    plot_distributions=false,
    plot_only=false,
    target_rotation=22.5)

    #
    # also identifies Invalid trials, i.e. those that are inititated and Fritz went for the WRONG target
    #
    #

    if !isempty(df)

        ### rotate targets to prevent splitting a cluster between 0 and 360 degrees
        rot_angl = target_rotation
        rot_rad = 2*pi/360 * rot_angl ### rotate the data points by 22.5 degrees to get target3 as one cluster
        rot_matr = [cos(rot_rad) -sin(rot_rad); sin(rot_rad) cos(rot_rad) ]
        ref_dist = (2pi)/360 * 45 ### theoretical distance of targets in rad

        ### determin theoretical angles of targets
        angle_of_first_target = (2*pi)/360 * rot_angl
        angle_of_last_target = 7*(2*pi)/360*45 + (2*pi)/360*rot_angl
        theoretical_x_vals_targets_va = collect(range(angle_of_first_target, angle_of_last_target; step=ref_dist))

        ### get vector length and angles of fixations
        vec_length = Array{Any}(undef, 0)
        vec_angl = Array{Any}(undef, 0)

        for (itrl,trls_fix) in enumerate(df.Fixations)
            push!(vec_length, Array{Float64,1}(undef,0))
            push!(vec_angl, Array{Float64,1}(undef,0))

            for (ifix,fix) in enumerate(trls_fix)

                # vec = [mean(xy_nb[itrl][1][fix]), mean(xy_nb[itrl][2][fix])]
                vec = [mean(df.EyeXY[itrl][1][fix]), mean(df.EyeXY[itrl][2][fix])]

                ### rotate fixations to prevent the cluster at 0deg to be splitted between >0 and <360
                vec_rotated = rot_matr*vec

                ### length
                push!(vec_length[itrl], round(norm(vec_rotated), digits=1))

                ### angle relative to x axis
                if vec_rotated[2] >= 0
                    push!(vec_angl[itrl], acos(dot([1,0],vec_rotated)/(norm([1,0])*norm(vec_rotated))))
                else
                    push!(vec_angl[itrl], acos(dot([1,0],vec_rotated)/(norm([1,0])*norm(vec_rotated)))+pi)
                end

            end
        end

        ### fit histograms and kernel densitys to length
        vec_length_ = MyUtils.flatten(vec_length)
        h_vl = fit(Histogram, vec_length_, nbins=nbins_length)
        h_norm_vl = normalize(h_vl, mode=:pdf)
        vec_len_dens = kde(vec_length_, bandwidth=kde_bandwidth_length)
        ### find modes of fitted distributions
        x_idx_all_modes_vl = findlocalmaxima(vec_len_dens.density)
        x_idx_two_largest_modes_vl = x_idx_all_modes_vl[sortperm(vec_len_dens.density[x_idx_all_modes_vl], rev=true)][1:2]
        x_val_two_largest_modes_vl = vec_len_dens.x[x_idx_two_largest_modes_vl]
        y_val_two_largest_modes_vl = vec_len_dens.density[x_idx_two_largest_modes_vl]
        ### sort saccades according fixation vector length
        # min and max vector lengths of central fixations and target fixations
        c_fix = [x_val_two_largest_modes_vl[1]-1, x_val_two_largest_modes_vl[1]+1]
        margin_t_fix_i = 3
        margin_t_fix_o = 3
        t_fix = [x_val_two_largest_modes_vl[2]-margin_t_fix_i, x_val_two_largest_modes_vl[2]+margin_t_fix_o]

        fixation_types = Array{Any,1}(undef, 0)

        for (itrl,trls_fix) in enumerate(df.Fixations)
            clean_fixations = [] ### Fixations that fall into the defined areas of central fixation point and targets

            push!(fixation_types, Array{Float64,1}(undef,0))

            for (ifix,fix) in enumerate(trls_fix)

                ### catch central fixations
                if c_fix[1] < vec_length[itrl][ifix] < c_fix[2]
                    push!(fixation_types[itrl], 0)
                elseif t_fix[1] < vec_length[itrl][ifix] < t_fix[2]
                    push!(fixation_types[itrl], 10)
                else
                    push!(fixation_types[itrl], -1)
                end

            end
        end

        bool_list_central_fixations = [ el.==10.0 for el in fixation_types ]
        vec_angl_targets_only = [vec_angl[ii][el] for (ii,el) in enumerate(bool_list_central_fixations)]

        ### fit histograms and kernel densitys to angles distributions
        vec_angl_ = MyUtils.flatten(vec_angl_targets_only)
        h_va = fit(Histogram, vec_angl_, nbins=nbins_angles)
        h_norm_va = normalize(h_va, mode=:pdf)
        vec_angl_dens = kde(vec_angl_, bandwidth=kde_bandwidth_angle)
        ### find modes of fitted distributions
        x_idx_all_modes_va = findlocalmaxima(vec_angl_dens.density)
        x_val_all_modes_va = vec_angl_dens.x[x_idx_all_modes_va]
        which_modes = [argmin(abs.(tx .- x_val_all_modes_va)) for tx in theoretical_x_vals_targets_va] ### find modes closest to theoretical location of modes
        x_idx_eight_largest_modes_va = x_idx_all_modes_va[which_modes]
        x_val_eight_largest_modes_va = x_val_all_modes_va[which_modes]
        y_val_eight_largest_modes_va = vec_angl_dens.density[x_idx_eight_largest_modes_va]
        ### correct for potentially wrong modes: replace modes that deviate more than 0.5 from theoretical mode
        false_modes = x_val_eight_largest_modes_va .- theoretical_x_vals_targets_va .> 0.5
        x_val_eight_largest_modes_va[false_modes] .=  theoretical_x_vals_targets_va[false_modes]

        ### angular range of target fixations
        margin_ts = 0.2
        t1 = [x_val_eight_largest_modes_va[1]-margin_ts, x_val_eight_largest_modes_va[1]+margin_ts]
        t2 = [x_val_eight_largest_modes_va[2]-margin_ts, x_val_eight_largest_modes_va[2]+margin_ts]
        t3 = [x_val_eight_largest_modes_va[3]-margin_ts, x_val_eight_largest_modes_va[3]+margin_ts]
        t4 = [x_val_eight_largest_modes_va[4]-margin_ts, x_val_eight_largest_modes_va[4]+margin_ts]
        t5 = [x_val_eight_largest_modes_va[5]-margin_ts, x_val_eight_largest_modes_va[5]+margin_ts]
        t6 = [x_val_eight_largest_modes_va[6]-margin_ts, x_val_eight_largest_modes_va[6]+margin_ts]
        t7 = [x_val_eight_largest_modes_va[7]-margin_ts, x_val_eight_largest_modes_va[7]+margin_ts]
        t8 = [x_val_eight_largest_modes_va[8]-margin_ts, x_val_eight_largest_modes_va[8]+margin_ts]

        for (itrl,trls_fix) in enumerate(fixation_types)
            clean_fixations = Array{Array{Float64,1},1}(undef,0)

            for (ifix,fix) in enumerate(trls_fix)

                ### catch target fixations, could either be 0 for central fixations or -1 for noise fixations or 10 for target fixations
                if fix == 10.0

                    ### sort into traget IDs
                    if t1[1] < vec_angl[itrl][ifix] < t1[2]
                        fixation_types[itrl][ifix] = 3
                    elseif t2[1] < vec_angl[itrl][ifix] < t2[2]
                        fixation_types[itrl][ifix] = 2
                    elseif t3[1] < vec_angl[itrl][ifix] < t3[2]
                        fixation_types[itrl][ifix] = 1
                    elseif t4[1] < vec_angl[itrl][ifix] < t4[2]
                        fixation_types[itrl][ifix] = 8
                    elseif t5[1] < vec_angl[itrl][ifix] < t5[2]
                        fixation_types[itrl][ifix] = 4
                    elseif t6[1] < vec_angl[itrl][ifix] < t6[2]
                        fixation_types[itrl][ifix] = 5
                    elseif t7[1] < vec_angl[itrl][ifix] < t7[2]
                        fixation_types[itrl][ifix] = 6
                    elseif t8[1] < vec_angl[itrl][ifix] < t8[2]
                        fixation_types[itrl][ifix] = 7
                    else
                        fixation_types[itrl][ifix] = -1
                    end

                end

                ### store fixations belonging to target or central fixation dot
                if !(fixation_types[itrl][end] == -1)
                    push!( clean_fixations, df.Fixations[itrl][ifix] )
                end

            end # end iterate trials fixations

            if !plot_only
                df.LookedAtTarget[itrl] = fixation_types[itrl]

                ### add columns Values
                df.Valid[itrl] = (!df.WasMissed[itrl]) & (any(unique(df[itrl, :LookedAtTarget]) .> 0))
                df.Invalid[itrl] = (!df.WasMissed[itrl]) & (!any(unique(df[itrl, :LookedAtTarget]) .> 0))
            #     # df.Hit[itrl] = (df.Valid[itrl]) & (df[itrl, :LookedAtTarget][end] == df[itrl, :ValidTargetIndex])
                if !isempty(df[itrl, :LookedAtTarget])
                    df.FA[itrl] = ((df.Valid[itrl]) & (df[itrl, :LookedAtTarget][end] != df[itrl, :ValidTargetIndex])) & !(df.Hit[itrl])
                end

                ### Get reaction times
                targetFixation_idxs = findall([ !(el in [-1,0]) for el in df[itrl,:LookedAtTarget]]) ### list of all target fixations
                if !isempty(targetFixation_idxs)
                    firstTargetFixation_id = minimum(targetFixation_idxs) ### the first target fixation
                    firstTargetFixation_time = df[itrl,:Fixations][firstTargetFixation_id] / sf
                    RT_firstTarget = minimum(firstTargetFixation_time .- df[itrl, :LuminanceChangeTime])

                    if df[itrl,:LookedAtTarget][firstTargetFixation_id] == df[itrl,:ValidTargetIndex]
                        RT_CorrectTarget = RT_firstTarget
                    else
                        correctTargetFixation_idxs = findall(df[itrl,:LookedAtTarget] .== df[itrl,:ValidTargetIndex]) ### list of all fixations on correct target
                        if !isempty(correctTargetFixation_idxs)
                            correctTargetFixation_idx = minimum(correctTargetFixation_idxs) ### the first fixation on correct target
                            correctTargetFixation_time = df[itrl,:Fixations][correctTargetFixation_idx] / sf
                            RT_CorrectTarget = minimum(correctTargetFixation_time .- df[itrl, :LuminanceChangeTime])
                        else ### if correct target was not fixated
                            RT_CorrectTarget = NaN
                        end
                    end

                else ### if no target was fixated
                    RT_firstTarget = NaN
                    RT_CorrectTarget = NaN
                end
                df.RTfirstFixation[itrl] = RT_firstTarget
                df.RTcorrectFixation[itrl] = RT_CorrectTarget

            end

        end # end iterate trials

        if plot_distributions | plot_only
            p = plot(layout=2, size=(1000,500), xlabel="saccade length / AU")
            plot!(h_norm_vl, subplot=1, xlims=[0, 50])
            plot!(vec_len_dens.x, vec_len_dens.density, linewidth=3, subplot=1)
            scatter!(x_val_two_largest_modes_vl, y_val_two_largest_modes_vl, subplot=1)
            #
            plot!(h_norm_va, subplot=2, xlabel="saccade angles / rad")
            plot!(vec_angl_dens.x, vec_angl_dens.density, linewidth=3, subplot=2)
            scatter!(x_val_eight_largest_modes_va, y_val_eight_largest_modes_va, subplot=2)
            return p
        end

    end

end

#################################################################################
#################################################################################

function i_dt(
    xy_o::Array{Array{Float64,1},1},
    sf,
    fix_duration_threshold_sec,
    fix_dispersion_threshold,
    saccade_velocity_threshold,
    saccade_duration_threshold
    )

    ###########
    #
    # Implements I-DT algorithm to detect fixations, sacades
    # extended to identify and interpolate blinks
    #
    # INPUT
    # fix_duration_threshold_sec: minimum duration of a fixation to be counted in sec
    # fix_dispersion_threshold: can be left empty then it is inferred from central fixation period in beginning of trials
    # saccade_velocity_threshold: not normalized to sampling rate!
    # saccade_duration_threshold: in sec
    #
    # OUTPUT
    #   xy_: eye trace with blinks interpolated
    #   sac_idx:
    #
    #
    #
    ###########

    if size(xy_o[1],1) > 1000 # if number of samples is lower empty arrays are returned

        xy = deepcopy(xy_o)
        [deleteat!(xy[ii], findall(isnan.(xy[ii]))) for ii in 1:length(xy)] # remove NaNs
        xy_ = deepcopy(xy) # to be returned with blinks interpolated
        xy_idxs = [collect(1:length(xy[1])), collect(1:length(xy[2]))] # array with corresponding indices
        fix_window_length = Int(fix_duration_threshold_sec * sf)

        if fix_dispersion_threshold == []

            fix_dispersion_threshold = maximum(xy[1][1:499])-minimum(xy[1][1:499]) + maximum(xy[2][1:499])-minimum(xy[2][1:499]) # deg

        end

        #
        sac_idx = Array{Array{Int,1}, 1}(undef,0)#[]
        fix_idx = Array{Array{Int,1}, 1}(undef,0)#[]
        blink_idx = Array{Array{Int,1}, 1}(undef,0)#[]

        ### I-DT loop:
        while !isempty(xy_idxs[1])

            if length(xy_idxs[1]) > fix_window_length # check if length(xy) > fix_window_length (lower duration threshold of saccades)
                current_window_idxs = collect(1:fix_window_length)
            else
                # @show length(xy_idxs[1])
                current_window_idxs = collect(1:length(xy_idxs[1])) # if length(xy) > fix_window_length
            end

            # println(current_window_idxs)

            current_window_vals = [xy[ii][current_window_idxs] for ii in 1:2]
            wd = (maximum(current_window_vals[1])-minimum(current_window_vals[1])) + (maximum(current_window_vals[2])-minimum(current_window_vals[2]))

            ### if dispersion in current window is smaller than threshold
            if wd <= fix_dispersion_threshold

                # while this is the case and window does not exceed amount of data points add next point to window
                while wd<=fix_dispersion_threshold && length(current_window_idxs)<length(xy[1])

                    ### append to indices of current window
                    current_window_idxs = append!( current_window_idxs, maximum(current_window_idxs)+1 )
                    ### append to values of current widnow
                    current_window_vals[1] = xy[1][current_window_idxs]
                    current_window_vals[2] = xy[2][current_window_idxs]
                    wd = (maximum(current_window_vals[1])-minimum(current_window_vals[1])) + (maximum(current_window_vals[2])-minimum(current_window_vals[2]))

                end

                ### add fixation to array and delete window from data
                # push!( fix_idx, [xy_idxs[ii][current_window_idxs] for ii in 1:2] )
                push!( fix_idx, xy_idxs[1][current_window_idxs] )
                [deleteat!( xy[ii], current_window_idxs ) for ii in 1:length(xy)]
                [deleteat!( xy_idxs[ii], current_window_idxs ) for ii in 1:length(xy_idxs)]

            else # if dispersion in current window is greater than threshold delete first data point

                [deleteat!( xy[ii], 1 ) for ii in 1:length(xy)]
                [deleteat!( xy_idxs[ii], 1 ) for ii in 1:length(xy_idxs)]

            end

        end

        ### get saccade-like indices, i.e. those not identified as fixations
        for (idx,fixs) in enumerate(fix_idx[1:end-1])
            #if (fix_idx[idx+1][1][1]-fixs[1][end]) > blink_velocity_threshold # blink velocity threshold
                # push!( sac_idx, fixs[1][end]:fix_idx[idx+1][1][1] )
                push!( sac_idx, fixs[end]:fix_idx[idx+1][1] )

            #end
        end

        ### filter blinks: saccades with velocities > threshold
        saccade_velocities = [diff(xy_[1][sac]) for sac in sac_idx]
        blinks = findall(maximum.(saccade_velocities) .> saccade_velocity_threshold)

        for bl in blinks

            push!( blink_idx, sac_idx[bl] )

        end
        # delete blinks from saccades collection
        deleteat!( sac_idx, blinks )
        deleteat!(sac_idx, findall(length.(sac_idx) .< saccade_duration_threshold*sf)) # remove implausible short saccades

        ### inerpolate blinks (linear)
        interpolated_blinks = []

        for bl in 1:length(blink_idx)

            push!(interpolated_blinks, range( xy_[1][minimum(blink_idx[bl])], xy_[1][maximum(blink_idx[bl])], length=length(blink_idx[bl]) ))

        end

        [xy_[dim][blink_idx[bl]] .= interpolated_blinks[bl] for bl in 1:length(blink_idx), dim=1:2 ] # dim is x,y dimensions of eyedata -> bot needed

        return xy_, fix_idx, sac_idx, blink_idx

    else
        return [], [], [], []
    end

end # i_dt

end # end Behavior.jl


using CSV
using DataFrames
using DelimitedFiles
using YAML
using Query
using Base

"""
    function disagg!(d_shr::Dict, d_cal::Dict)
"""
function disagg!(d::Dict, set::Dict; save = true, overwrite = false)

    d_read = read_build("disagg"; save = save, overwrite = overwrite);
    if !isempty(d_read)
        [d[k] = v for (k,v) in d_read]
        return d
    end

    # d = merge(d, Dict(
    #     :r => fill_with((r = set[:r],), 1.0),
    #     (:yr,:r,:g) => fill_with((yr = set[:yr], r = set[:r], g = set[:g]), 1.0)))
    # d = merge(d, copy(d_shr), copy(d_cal))

    d[:ys0] = _disagg_ys0!(d)
    d[:id0] = _disagg_id0!(d)
    d[:ty0] = _disagg_ty0(d, set)
    d[:va0] = _disagg_va0!(d, set)
    d[:ld0] = _disagg_ld0!(d)
    d[:kd0] = _disagg_kd0(d)

    d[:fd0] = _disagg_fdcat!(d)       # not state model
    d[:g0] = _disagg_g0!(d)
    d[:i0] = _disagg_i0!(d)
    d[:cd0] = _disagg_cd0!(d)
    d[:c0] = _disagg_c0!(d)

    d[:yh0_temp] = _disagg_yh0!(d)
    d[:fe0] = _disagg_fe0!(d)         # not state model
    d[:x0_temp] = _disagg_x0!(d, set)
    d[:s0_temp] = _disagg_s0!(d)
    d[:a0] = _disagg_a0!(d)           # BE CAREFUL. DON'T REPLACE ARMINGTON SUPPLY
    d[:ta0] = _disagg_ta0!(d)
    d[:tm0] = _disagg_tm0!(d)
    d[:thetaa] = _disagg_thetaa!(d)   # not in state model

    d[:m0] = _disagg_m0!(d)
    d[:md0] = _disagg_md0!(d)
    d[:rx0_temp] = _disagg_rx0!(d)

    d[:pt0_temp] = _disagg_pt0!(d)    # not in state model
    d[:diff] = _disagg_diff!(d)       # not in state model

    d[:rx0] = _disagg_rx0!(d)
    d[:s0] = _disagg_s0!(d)
    d[:x0] = _disagg_x0!(d, set)
    d[:yh0] = _disagg_yh0!(d)
    d[:bopdef0] = _disagg_bopdef0!(d)

    d[:pt0] = _disagg_pt0!(d)
    d[:dc0] = _disagg_dc0!(d)

    d[:dd0max] = _disagg_dd0max(d)
    d[:dd0] = _disagg_dd0!(d)
    d[:nd0] = _disagg_nd0!(d)
    d[:dm0] = _disagg_dm0!(d)
    d[:nm0] = _disagg_nm0!(d)
    d[:xd0] = _disagg_xd0!(d)
    d[:xn0] = _disagg_xn0!(d)

    d[:hhadj] = _disagg_hhadj!(d)

    d_save = Dict()
    d_save[:a0]  = ensurenames(d[:a0], [:yr, :r, :g, :value])
    d_save[:bopdef0] = ensurenames(d[:bopdef0], [:yr, :r, :value])
    d_save[:c0]  = ensurenames(d[:c0], [:yr, :r, :value])
    d_save[:cd0] = ensurenames(d[:cd0], [:yr, :r, :s, :value])
    d_save[:dd0] = ensurenames(d[:dd0], [:yr, :r, :g, :value])
    d_save[:dm0] = ensurenames(d[:dm0], [:yr, :r, :g, :m, :value])
    d_save[:g0]  = ensurenames(d[:g0], [:yr, :r, :s, :value])
    d_save[:hhadj] = ensurenames(d[:hhadj], [:yr, :r, :value])
    d_save[:i0]  = ensurenames(d[:i0], [:yr, :r, :s, :value])
    d_save[:id0] = ensurenames(d[:id0], [:yr, :r, :g, :s, :value])
    d_save[:kd0] = ensurenames(d[:kd0], [:yr, :r, :s, :value])
    d_save[:ld0] = ensurenames(d[:ld0], [:yr, :r, :s, :value])
    d_save[:m0]  = ensurenames(d[:m0], [:yr, :r, :g, :value])
    d_save[:md0] = ensurenames(d[:md0], [:yr, :r, :m, :g, :value])
    d_save[:nd0] = ensurenames(d[:nd0], [:yr, :r, :g, :value])
    d_save[:nm0] = ensurenames(d[:nm0], [:yr, :r, :g, :m, :value])
    d_save[:rx0] = ensurenames(d[:rx0], [:yr, :r, :g, :value])
    d_save[:s0]  = ensurenames(d[:s0], [:yr, :r, :g, :value])
    d_save[:ta0] = ensurenames(d[:ta0], [:yr, :r, :g, :value])
    d_save[:tm0] = ensurenames(d[:tm0], [:yr, :r, :g, :value])
    d_save[:ty0] = ensurenames(d[:ty0], [:yr, :r, :s, :value])
    d_save[:x0]  = ensurenames(d[:x0], [:yr, :r, :g, :value])
    d_save[:xd0] = ensurenames(d[:xd0], [:yr, :r, :g, :value])
    d_save[:xn0] = ensurenames(d[:xn0], [:yr, :r, :g, :value])
    d_save[:yh0] = ensurenames(d[:yh0], [:yr, :r, :s, :value])
    d_save[:ys0] = ensurenames(d[:ys0], [:yr, :r, :s, :g, :value])

    write_build("disagg", d_save; save = save)
    return d_save
end

"`ys0(yr,r,s,g)`: Regional sectoral output"
function _disagg_ys0!(d::Dict)
    println("  Disaggregating ys0(yr,r,s,g), regional sectoral output")
    :r in propertynames(d[:ys0]) && (return d[:ys0])

    d[:ys0] = d[:region] * d[:ys0]
    return d[:ys0]
end

"`id0(yr,r,g,s)`: Regional intermediate demand"
function _disagg_id0!(d::Dict)
    println("  Disaggregating id0(yr,r,g,s), regional intermediate demand")
    :r in propertynames(d[:id0]) && (return d[:id0])

    d[:id0] = d[:region] * d[:id0]
    return d[:id0]
end

"`ty0_rev`: Production tax payments"
function _disagg_ty0_rev(d::Dict, set::Dict)
    if :va in propertynames(d[:va0])
        d[:va0] = edit_with(unstack(copy(d[:va0]), :va, :value),
            Replace.(Symbol.(set[:va]), missing, 0.0))
    end
    d[:region] * d[:va0][:,[:yr,:s,:othtax]]
end

"`ty0(yr,r,s)`: Production tax rate"
function _disagg_ty0(d::Dict, set::Dict)
    println("  Disaggregating ty0(yr,r,s), production tax rate")
    # !!!! test that returns error if va0 has already been edited.

    ty0_rev = _disagg_ty0_rev(d, set)
    df = dropnan(ty0_rev / combine_over(d[:ys0], :g))
    return df
end

"`va0`: Regional value added"
function _disagg_va0!(d::Dict, set::Dict)
    println("  Disaggregating va0, regional share of value added.")
    # If va0 has already been edited, don't edit it again.
    :r in propertynames(d[:va0]) && (return d[:va0])

    if :va in propertynames(d[:va0])
        d[:va0] = edit_with(unstack(copy(d[:va0]), :va, :value), Replace.(Symbol.(set[:va]), missing, 0.0))
    end

    # !(:compen in propertynames(d[:va0])) && return d[:va0]

    df = d[:va0][:,[:yr,:s,:compen]] + d[:va0][:,[:yr,:s,:surplus]]
    d[:va0] = d[:region] * df
    return d[:va0]
end

"`ld0`: Labor demand"
function _disagg_ld0!(d::Dict)
    println("  Disaggregating ld0(yr,r,s), labor demand")
    d[:ld0] = d[:labor] * d[:va0]
    return d[:ld0]
end

"`kd0`: Capital demand"
_disagg_kd0(d::Dict) = d[:va0] - d[:ld0]

function _disagg_fdcat!(d::Dict)
    # !!!! check if this has been edited.
    x = Map(joinpath("crosswalk","fd.csv"), [:fd], [:fdcat], [:fd], [:fdcat], :inner)

    # df = copy(d[:fd0])
    if !(:fdcat in propertynames(d[:fd0]))
        d[:fd0] = edit_with(d[:fd0], x)
        d[:fd0] = combine_over(d[:fd0], :fd)
    end
    return d[:fd0]
end

"`g0(yr,r,s)`: National government demand"
function _disagg_g0!(d::Dict)
    println("  Disaggregating g0(yr,r,s), national government demand")
    cols = [:yr,:r,:g,:value]

    df_fd0 = ensurenames(d[:fd0], [:yr,:fdcat,:g,:value])
    df = filter_with(df_fd0, (fdcat = "G",))[:,[:yr,:g,:value]]
    d[:g0] = ensurenames(d[:sgf], cols) * df
    # d[:g0] = edit_with(df, Rename(:s,:g))
    return d[:g0]
end

"`i0(yr,g)`: National investment demand"
function _disagg_i0!(d::Dict)
    println("  Disaggregating i0(yr,r,s), national investment demand")
    cols = [:yr,:r,:g,:value]
    df_fd0 = ensurenames(d[:fd0], [:yr,:fdcat,:g,:value])
    df = filter_with(df_fd0, (fdcat = "I",))[:,[:yr,:g,:value]]
    d[:i0] = ensurenames(d[:region], cols) * df
    # d[:i0] = edit_with(df, Rename(:s,:g))
    return d[:i0]
end

"`cd0(yr,r,s)`: National final consumption"
function _disagg_cd0!(d::Dict)
    println("  Disaggregating cd0(yr,r,s), national final consumption")
    cols = [:yr,:r,:g,:value]
    df_fd0 = ensurenames(d[:fd0], [:yr,:fdcat,:g,:value])
    df = filter_with(df_fd0, (fdcat = "C",))[:,[:yr,:g,:value]]
    d[:cd0] = ensurenames(d[:pce], cols) * df
    # d[:cd0] = edit_with(df, Rename(:s,:g))
    return d[:cd0]
end

"`c0(yr,r)`: Total final household consumption"
function _disagg_c0!(d::Dict)
    println("  Disaggregating c0(yr,r), total final household consumption")
    !(:cd0 in keys(d)) && _disagg_cd0!(d)
    d[:c0] = combine_over(d[:cd0], :g)
    return d[:c0]
end

############################################################################################
"`yh0(yr,r,s)`: Household production"
function _disagg_yh0!(d::Dict)
    println("  Disaggregating yh0(yr,r,s), household production")
    if !(:diff in keys(d))
        df_region = ensurenames(d[:region], [:yr,:r,:g,:value])
        df_fs0 = ensurenames(d[:fs0], [:yr,:g,:value])
        d[:yh0] = df_region * df_fs0
        # d[:yh0] = edit_with(df, Rename(:s,:g))
    else
        d[:yh0] = dropmissing(d[:yh0] + d[:diff])
    end
    return d[:yh0]
end

"`fe0(yr,r)`: Total factor supply"
function _disagg_fe0!(d::Dict)
    println("  Disaggregating fe0, total factor supply")
    d[:fe0] = combine_over(d[:va0], :s)
    return d[:fe0]
end

"`x0(yr,r,g)`: Foreign exports"
function _disagg_x0!(d::Dict, set::Dict)
    println("  Disaggregating x0(yr,r,g), foreign exports")
    if !(:diff in keys(d))
        cols = [:yr,:r,:g,:value]

        set[:notrd] = :s in propertynames(d[:utd]) ? setdiff(set[:s], d[:utd][:,:s]) :
            setdiff(set[:g], d[:utd][:,:g])

        df_exports = edit_with(filter_with(copy(d[:utd]), (t = "exports",)),
            [Drop.([:t,:units],"all","=="); Rename(:s,:g)])
        df_region = filter_with(ensurenames(d[:region], cols), (g = set[:notrd],))

        df_trd = dropmissing(df_exports * d[:x0])
        df_notrd = df_region * d[:x0]

        d[:x0] = [df_trd; df_notrd]
    else
        d[:x0] = dropmissing(d[:x0] + d[:diff])
    end

    return d[:x0]
end

"`s0(yr,r,g)`: Total supply"
function _disagg_s0!(d::Dict)
    println("  Disaggregating s0(yr,r,g), total supply")
    if !(:diff in keys(d))
        d[:s0] = combine_over(d[:ys0], :s) + d[:yh0]
    else
        d[:s0] = dropmissing(d[:s0] + d[:diff])
    end
    return d[:s0]
end

"`a0(yr,r,g)`: Domestic absorption"
function _disagg_a0!(d::Dict)
    cols = [:yr,:r,:g,:value]
    println("  Disaggregating a0(yr,r,g), domestic absorption")

    d[:a0] = dropmissing(d[:cd0] + d[:g0] + d[:i0] + combine_over(d[:id0], :s))
    return ensurenames!(d[:a0], cols)
end

"`ta0(yr,r,g)`: Absorption taxes"
function _disagg_ta0!(d::Dict)
    println("  Disaggregating ta0(yr,r,g), absorption taxes")
    d[:ta0] = ensurenames(d[:ta0] * d[:r], [:yr,:r,:g,:value])
    return d[:ta0]
end

"`tm0(yr,r,g)`: Import taxes"
function _disagg_tm0!(d::Dict)
    println("  Disaggregating tm0(yr,r,g), import taxes")
    d[:tm0] = ensurenames(d[:tm0] * d[:r], [:yr,:r,:g,:value])
    return d[:tm0]
end

"`thetaa(yr,r,g)`: Share of regional absorption"
function _disagg_thetaa!(d::Dict)
    println("  Disaggregating thetaa(yr,r,g), share of regional absorption")
    d[:thetaa] = dropnan(d[:a0] / transform_over(d[:a0], :r))
    return d[:thetaa]
end

"`m0(yr,r,g)`: Foreign Imports"
function _disagg_m0!(d::Dict)
    println("  Disaggregating m0(yr,r,g), foreign imports")
    d[:m0] = d[:thetaa] * d[:m0]
    return d[:m0]
end

"`md0(yr,r,m,g)`: Margin demand"
function _disagg_md0!(d::Dict)
    println("  Disaggregating md0(yr,r,m,g), margin demand")
    d[:md0] = dropmissing(d[:thetaa] * d[:md0])
    return d[:md0]
end

"`rx0(yr,r,g)`: re-exports"
function _disagg_rx0!(d::Dict)
    println("  Disaggregating rx0(yr,r,g), re-exports")
    if !(:diff in keys(d))
        d[:rx0] = d[:x0] - d[:s0]
        d[:rx0][round.(d[:rx0][:,:value]; digits = 10) .< 0, :value] .= 0.0
    else
        d[:rx0] = d[:rx0] + d[:diff]
    end

    return dropmissing(d[:rx0])
end

# ******************************************************************************************
"`pta0`"
_disagg_pta0(d::Dict) = dropmissing(((d[:yr,:r,:g] - d[:ta0]) * d[:a0]) + d[:rx0])

"`ptm0`"
_disagg_ptm0(d::Dict) = dropmissing(((d[:yr,:r,:g] + d[:tm0]) * d[:m0]) + combine_over(d[:md0], :m))

"`dc0`"
function _disagg_dc0!(d::Dict)
    d[:dc0] = dropmissing((d[:s0] - d[:x0] + d[:rx0]))
    d[:dc0][!,:value] .= round.(d[:dc0][:,:value]; digits = 10)
    return d[:dc0]
end

"`pt0`"
function _disagg_pt0!(d::Dict)
    df_pta0 = _disagg_pta0(d)
    df_ptm0 = _disagg_ptm0(d)
    d[:pt0] = df_pta0 - df_ptm0
    d[:pt0][!,:value] .= round.(d[:pt0][:,:value]; digits = 10)
    return d[:pt0]
end

"`diff`:"
function _disagg_diff!(d::Dict)
    df = copy(d[:pt0])
    df[!,:value] .= - min.(0, df[:,:value])
    d[:diff] = edit_with(df, Drop(:value,0.0,"=="))
end

"`bopdef0(yr,r)`: Balance of payments (closure parameter)"
function _disagg_bopdef0!(d::Dict)
    println("  Disaggregating bopdef0(yr,r), balance of payments (closure parameter)")
    cols = [:yr,:r,:value]
    d[:bopdef0] = combine_over((d[:m0] - d[:x0]), :g)
    return ensurenames!(d[:bopdef0], cols)
end

"`gm`: Commodities employed in margin supply"
function _disagg_gm!(d::Dict, set::Dict)
    cols = [:g,:yr]
    ms0_sum = ensurenames(combine_over(d[:ms0], [:yr,:m]), cols)
    md0_sum = ensurenames(combine_over(d[:md0], [:yr,:m]), cols)

    ms0_sum[!,:value] .= ms0_sum[:,:value] .!= 0.0
    md0_sum[!,:value] .= md0_sum[:,:value] .!= 0.0

    gm = ms0_sum + md0_sum
    set[:gm] = gm[gm[:,:value] .> 0, :g]
    return set[:gm]
end

"`dd0max(yr,r,g)`: Maximum regional demand from local market"
function _disagg_dd0max(d::Dict)
    println("  Disaggregating dd0max(yr,r,g), maximum regional demand from local market")
    cols = [:yr,:r,:g,:value]
    df = SLiDE._join_to_operate(d[:pt0],d[:dc0]; colnames = [:pt0,:dc0])
    df[!,:value] = min.(df[:,:pt0], df[:,:dc0])
    d[:dd0max] = copy(df)[:,cols]
    return d[:dd0max]
end

"`nd0max(yr,r,g)`: Maximum regional demand from national market"
function _disagg_nd0max(d::Dict)
    println("  Disaggregating nd0max(yr,r,g), maximum regional demand from national market")
    cols = [:yr,:r,:g,:value]
    df = SLiDE._join_to_operate(d[:pt0],d[:dc0]; colnames = [:pt0,:dc0])
    df[!,:value] = min.(df[:,:pt0], df[:,:dc0])
    d[:nd0max] = copy(df)[:,cols]
    return d[:nd0max]
end

"`dd0min(yr,r,g)`: Minimum regional demand from local market"
_disagg_dd0min(d::Dict) = d[:pt0] - _disagg_dd0max(d)

"`nd0min(yr,r,g)`: Minimum regional demand from national market"
_disagg_nd0min(d::Dict) = d[:pt0] - _disagg_nd0max(d)

"`dd0_(yr,r,g)`: Regional demand from local market"
function _disagg_dd0!(d::Dict)
    println("  Disaggregating dd0_(yr,r,g), regional demand from local market")
    df_dd0max = _disagg_dd0max(d)
    d[:dd0] = df_dd0max * d[:rpc]
    return d[:dd0]
end

"`nd0_(yr,r,g)`: Regional demand from national market"
function _disagg_nd0!(d::Dict)
    println("  Disaggregating nd0_(yr,r,g), regional demand from national market")
    d[:nd0] = d[:pt0] - d[:dd0]
    d[:nd0][!,:value] .= round.(d[:nd0][:,:value]; digits = 10)
    return d[:nd0]
end

# ##########################################################################################

"`mrgshr(yr,r,m)`: Share of margin demand by region"
function _disagg_mrgshr(d::Dict)
    df = combine_over(d[:md0], :g)
    df = df / transform_over(df, :r)
    return df
end

"`totmrgsupply(yr,r,m,g)`: Designate total supply of margins"
function _disagg_totmrgshr!(d::Dict)
    cols = [:yr,:r,:m,:g,:value]
    d[:totmrgshr] = _disagg_mrgshr(d) * d[:ms0]
    return ensurenames!(d[:totmrgshr], cols)
end

"`shrtrd(yr,r,m,g)`: Share of margin total by margin type"
function _disagg_shrtrd!(d::Dict)
    cols = [:yr,:r,:m,:g,:value]
    df = (:totmrgshr in keys(d)) ? d[:totmrgshr] : _disagg_totmrgshr!(d)
    d[:shrtrd] = dropnan(df / transform_over(df, :m))
    return ensurenames!(d[:shrtrd], cols)
end

"`dm0(yr,r,g,m)`: Margin supply from the local market"
function _disagg_dm0!(d::Dict)
    println("  Disaggregating dm0(yr,r,g,m), margin supply from the local market")
    cols = [:yr,:r,:m,:g,:value]
    !(:totmrgshr in keys(d)) && _disagg_totmrgshr!(d)
    !(:shrtrd in keys(d)) && _disagg_shrtrd!(d)

    dm1 = dropmissing(d[:totmrgshr] * d[:rpc])
    dm2 = dropmissing((d[:shrtrd] * d[:dc0]) - d[:dd0])

    df = SLiDE._join_to_operate(dm1, dm2; colnames = [:dm1,:dm2])
    df[!,:value] .= min.(df[:,:dm1], df[:,:dm2])
    d[:dm0] = df[:,cols]
    return d[:dm0]
end

"`nm0(yr,r,g,m)`: Margin demand from the national market"
function _disagg_nm0!(d::Dict)
    println("  Disaggregating nm0(yr,r,g,m), margin demand from the national market")
    d[:nm0] = d[:totmrgshr] - d[:dm0]
    return d[:nm0]
end

"`xd0(yr,r,g)`: Regional supply to local market"
function _disagg_xd0!(d::Dict)
    println("  Disaggregating xd0(yr,r,g), regional supply to local market")
    d[:xd0] = combine_over(d[:dm0], :m) + d[:dd0]
    return d[:xd0]
end

"`xn0(yr,r,g)`: Regional supply to national market"
function _disagg_xn0!(d::Dict)
    println("  Disaggregating xn0(yr,r,g), regional supply to national market")
    d[:xn0] = d[:s0] + d[:rx0] - d[:xd0] - d[:x0]
    return d[:xn0]
end

"`hhadj`: Household adjustment"
function _disagg_hhadj!(d::Dict)
    println("  Disaggregating hhadj, household adjustment")
    dh = Dict(k => edit_with(copy(d[k]), Rename(:g,:s))
        for k in [:c0,:ld0,:kd0,:yh0,:bopdef0,:ta0,:a0,:tm0,:ty0,:g0,:i0,:m0])
    dh[:ys0] = copy(d[:ys0])

    d[:hhadj] = dh[:c0] -
        combine_over(dh[:ld0] + dh[:kd0] + dh[:yh0], :s) - dh[:bopdef0] -
        combine_over(dh[:ta0]*dh[:a0] + dh[:tm0]*dh[:m0] + dh[:ty0]*combine_over(dh[:ys0],:g), :s) +
        combine_over(dh[:g0] + dh[:i0], :s)

    return d[:hhadj]
end


module SupervisedLearning

export train!

using Reexport
@reexport using MLBase

include("common.jl")
include("state.jl")
include("classencoding.jl")
include("datasource.jl")
include("traininghistory.jl")
include("solver.jl")
include("supervisedmodel.jl")
include("empiricalrisks.jl")

end # module


"""
The `assembleFreeEnergy` function accepts an `Algorithm` and populates
required fields for computing the variational free energy.
"""
function assembleFreeEnergy!(algo=currentInferenceAlgorithm())
    # Find counting numbers for energies and entropies
    pfz = algo.posterior_factorization
    assembleCountingNumbers!(pfz)

    # Convert energy counting numbers to energy inbounds
    average_energies_vect = Vector{Dict{Symbol, Any}}()
    entropies_vect = Vector{Dict{Symbol, Any}}()

    energy_counting_numbers_vect = [(node, cnt) for (node, cnt) in pfz.energy_counting_numbers] # Order the Dict in a Vector
    for (node, cnt) in sort(energy_counting_numbers_vect)
        if cnt != 0
            average_energy = Dict{Symbol, Any}(:counting_number => cnt,
                                               :node => typeof(node),
                                               :node_id => node.id,
                                               :inbounds => collectAverageEnergyInbounds(node))
            push!(average_energies_vect, average_energy)
        end
    end

    # Convert entropy counting numbers to entropy inbounds
    entropy_counting_numbers_vect = [(target, cnt) for (target, cnt) in pfz.entropy_counting_numbers] # Order the Dict in a Vector
    for (target, cnt) in sort(entropy_counting_numbers_vect)
        if cnt != 0
            entropy = Dict{Symbol, Any}(:counting_number => cnt,
                                        :target => target,
                                        :inbound => algo.target_to_marginal_entry[target])
            push!(entropies_vect, entropy)
        end
    end

    algo.average_energies = average_energies_vect
    algo.entropies = entropies_vect

    return algo
end

"""
The `assembleCountingNumbers` function accepts an `Algorithm` and
populates the counting numbers for the average energies and entropies.
"""
function assembleCountingNumbers!(pfz=currentPosteriorFactorization())
    pfz.free_energy_flag || error("Required quantities for free energy evaluation are not computed by the algorithm. Make sure to flag free_energy=true upon algorithm construction to schedule computation of required quantities.")

    energy_counting_numbers = Dict{FactorNode, Int64}()
    entropy_counting_numbers = Dict{Region, Int64}()

    # Collect regions
    internal_edges = Set{Edge}()
    for (_, pf) in pfz.posterior_factors
        union!(internal_edges, pf.internal_edges)
    end
    nodes_connected_to_internal_edges = nodes(internal_edges)

    # Iterate over large regions
    for node in nodes_connected_to_internal_edges
        target_regions = unique!(localStochasticRegions(node, pfz)) # Collect all unique stochastic regions around node
        if isDeltaConstraint(node)
            continue
        elseif !isa(node, DeltaFactor) # Node is stochastic
            increase!(energy_counting_numbers, node, 1) # Count average energy
            for target in target_regions
                increase!(entropy_counting_numbers, target, 1) # Count (joint) entropy
            end
        elseif isa(node, Equality)
            increase!(entropy_counting_numbers, node.i[1].edge.variable, 1) # Count univariate entropy
        elseif length(node.interfaces) >= 2 # Node is deterministic and not equality, requires the counting of the joint inbounds region
            outbound_region = region(node, node.interfaces[1].edge)
            for target in target_regions
                if target != outbound_region # Exclude outbound region
                    increase!(entropy_counting_numbers, target, 1) # Count (joint) entropy
                end
            end
        end
    end

    # Iterate over small regions
    for edge in internal_edges
        increase!(entropy_counting_numbers, edge.variable, -(degree(edge) - 1)) # Discount univariate entropy
        if edge in pfz.deterministic_edges
            entropy_counting_numbers[edge.variable] = 0 # However, ignore entropies of deterministic beliefs
        end
    end

    pfz.energy_counting_numbers = energy_counting_numbers
    pfz.entropy_counting_numbers = entropy_counting_numbers

    return pfz
end

"""
Increase (or decrease) a counting number
"""
function increase!(dict::Dict, key::Any, increase::Number)
    if haskey(dict, key)
        dict[key] += increase
    else
        dict[key] = increase
    end

    return dict
end