module ArrayInitialization

using DocStringExtensions
using Unitful
using AxisArrays
using ArgCheck
using ..ConstantArrays: ConstantArray
using ..UnitfulHartree
using ..Traits: components, ColinearSpin, SpinDegenerate, SpinCategory, ColinearSpin,
                ColinearSpinFirst, ColinearSpinLast, is_spin_polarized, concretize_type,
                SpinAware
using ..Dispatch
export wrap

macro lintpragma(s) end
@lintpragma("Ignore unused args")
@lintpragma("Ignore unused T")

const DD = Dispatch.Dimensions

_size(T::Type{<: DD.Scalars.All}, ::SpinDegenerate,
      ::SpinDegenerate, object::DD.AxisArrays.All) = size(object)
_size(T::Type{<: DD.Scalars.All}, ::ColinearSpin,
      ::SpinDegenerate, object::DD.AxisArrays.All) =
    (size(object)..., length(components(object)))

add_spin_axis(::SpinDegenerate, tup::Tuple, ::Any) = tup
add_spin_axis(::ColinearSpinLast, tup::Tuple, axis::Any) = tup..., axis
add_spin_axis(::ColinearSpin{N}, tup::Tuple, axis::Any) where N = begin
    @lintpragma("Ignore use of undeclared variable N")
    @lintpragma("Ignore use of undeclared variable axis")
    @lintpragma("Ignore use of undeclared variable tup")
    top, down = Base.IteratorsMD.split(tup, Val{N})

    Base.front(top)..., axis, Base.last(top), down...
end

replace_spin_axis(T::Type{<: DD.Scalars.All}, ::SpinDegenerate,
                  sizes::Tuple, ax::Tuple) = begin
    @assert length(sizes.parameters) == length(ax.parameters)
    @assert findfirst(is_spin_polarized, ax.parameters) == 0
    sizes, ax
end

"""
Position for the spin axis for a given spin category

$(SIGNATURES)

`n` is the number of dimensions in the array, and `s` is the location of the current spin
axis, or 0 if there are none.
"""
spin_axis_position(::Type{ColinearSpin{T}}, ::Integer, ::Integer) where T = T
spin_axis_position(::Type{ColinearSpinLast}, n::Integer, s::Integer) =
    !(1 ≤ s ≤ n) ? n + 1: n
spin_axis_position(S::SpinCategory, n::Integer, s::Integer) =
    spin_axis_position(typeof(S), n, s)
spin_axis_position(array::DD.AxisArrays.All) = findfirst(is_spin_polarized, axes(array))
spin_axis_position(T::Type{DD.AxisArrays.All}) =
    findfirst(is_spin_polarized, T.parameters[end].parameters)

function _move_and_re(name::Symbol, expr::Expr, orig::Integer, final::Integer, n::Integer)
    result = Expr(:tuple)
    for i in 1:n
        i > orig && push!(result.args, :($name[$i]))
        i == final && push!(result.args, expr)
        i < orig && push!(result.args, :($name[$i]))
    end
    final == n + 1 && push!(result.args, expr)
    result
end

@generated replace_spin_axis(T::Type{<: DD.Scalars.All},
                             S::ColinearSpin, sizes::Tuple, ax::Tuple) = begin
    @lintpragma("Ignore use of undeclared variable comps")
    @assert length(sizes.parameters) == length(ax.parameters)
    current = findfirst(is_spin_polarized, ax.parameters)
    final = spin_axis_position(S, length(sizes.parameters), current)
    @assert final ≠ 0
    d = _move_and_re(:sizes, :(length(comps)), current, final, length(sizes.parameters))
    a = _move_and_re(:ax, :(Axis{:spin}(comps)), current, final, length(sizes.parameters))
    quote
        comps = components(T, S)
        ($d, $a)
    end
end

@generated remove_spin_axis(sizes::Tuple, ax::Tuple) = begin
    @lintpragma("Ignore use of undeclared variable comps")
    @assert length(sizes.parameters) == length(ax.parameters)
    current = findfirst(is_spin_polarized, ax.parameters)
    d = _move_and_re(:sizes, :(length(comps)), current, 0, length(sizes.parameters))
    a = _move_and_re(:ax, :(Axis{:spin}(comps)), current, 0, length(sizes.parameters))
    :(($d, $a))
end

for extension in [:zeros, :ones, :rand]
    @eval begin
        """
        Creates an array for the given DFT quantity

        The spin axis is automatically added if required. Neither `dims` nor `ax` should
        include the spin dimension.
        """
        Base.$extension(T::Type{<:DD.Scalars.All}, C::SpinCategory,
                        dims::Tuple, ax::Tuple) = begin
            @argcheck length(ax) <= (length(dims) + 1) "Too many axes"
            comps = components(T, C)
            data = $extension(T.parameters[1], add_spin_axis(C, dims, length(comps)))
            # we can use a dummy array here since the underlying array (and indices) will be
            # the standard one
            defaults = AxisArrays.default_axes(ConstantArray(0, dims), ax)
            AxisArray(reinterpret(T, data), add_spin_axis(C, defaults, Axis{:spin}(comps)))
        end
        Base.$extension(T::Type{<:DD.Scalars.All}, C::SpinCategory, dims::Tuple) = begin
            comps = components(T, C)
            data = $extension(T, add_spin_axis(C, dims, length(comps)))
            defaults = AxisArrays.default_axes(ConstantArray(0, dims))
            AxisArray(data, add_spin_axis(C, defaults, Axis{:spin}(comps)))
        end
        Base.$extension(T::Type{<:DD.Scalars.All}, C::SpinCategory,
                        ax::Tuple{Axis, Vararg{Axis}}) = begin
            @argcheck(SpinCategory(ax) isa SpinDegenerate,
                      "Do not include Spin axis explicitly when specifying spin category")
            $extension(T, C, length.(ax), ax)
        end
        Base.$extension(T::Type{<:DD.Scalars.All}, C::SpinCategory, ax::Axis...) =
            $extension(T, C, ax)
        Base.$extension(T::Type{<:DD.Scalars.All}, ax::Axis...) = $extension(T, ax)
        Base.$extension(T::Type{<:DD.Scalars.All}, ax::Tuple{Axis, Vararg{Axis}}) = begin
            const C = SpinCategory(ax)
            if C isa ColinearSpin
                index = findfirst(a -> a isa Axis{:spin}, ax)
                @argcheck(length(components(T, C)) == length(ax[index]),
                          "Incorrect spin components for given input type")
            end
            AxisArray(reinterpret(T, $extension(T.parameters[1], length.(ax))), ax)
        end

        """
        Creates an array for the given DFT quantity

        The spin components, if any, are added as the dimension specified by the spin trait.
        Note that the spin trait is the only way polarization should enter the input
        arguments, e.g. the spin-axis should not be given explicitly, neither in the axes
        nor in the dimensions.
        """
        Base.$extension(T::Type{<:DD.Scalars.All}, P::SpinCategory,
                        args::Vararg{Union{Integer, Axis}}; kwargs...) = begin
            @lintpragma("Ignore use of undeclared variable x")
            dims = (
                (x for x in args if typeof(x) <: Integer)...,
                (length(v) for (_, v) in kwargs)...,
            )
            ax = (
                (x for x in args if typeof(x) <: Axis)...,
                (Axis{k}(v) for (k, v) in kwargs)...,
            )
            $extension(T, P, dims, ax)
        end
        Base.$extension(T::Type{<:DD.Scalars.All}, polarized::Bool,
                        args::Vararg{Union{Integer, Axis}}; kwargs...) =
            $extension(T, polarized ? ColinearSpin(): SpinDegenerate(), args...; kwargs...)
    end
end

for extension in [:zeros, :ones, :similar]
    private = Symbol(:_, extension)
    @eval begin
        $private(T::Type{<:DD.Scalars.All}, ::SpinDegenerate, ::SpinDegenerate,
            array::DD.AxisArrays.All) = $extension(array, concretize_type(T, array))

        $private(T::Type{<:DD.Scalars.All}, C::ColinearSpin,
                 ::SpinDegenerate, array::DD.AxisArrays.All) = begin
            comps = components(T, C)
            dims = add_spin_axis(C, size(array), length(comps))
            data = $extension(array.data, concretize_type(T, array), dims)
            defaults = AxisArrays.default_axes(ConstantArray(0, size(array)), axes(array))
            AxisArray(data, add_spin_axis(C, defaults, Axis{:spin}(comps)))
        end

        $private(T::Type{<:DD.Scalars.All}, C::ColinearSpin,
                 ::ColinearSpin, array::DD.AxisArrays.All) = begin
            dims = replace_spin_axis(T, C, size(array), axes(array))
            AxisArray($extension(array.data, concretize_type(T, array), dims[1]), dims[2])
        end

        $private(T::Type{<:DD.Scalars.All}, ::SpinDegenerate,
                    ::ColinearSpin, array::DD.AxisArrays.All) = begin
            dims = remove_spin_axis(size(array), axes(array))
            AxisArray($extension(array.data, concretize_type(T, array), dims[1]), dims[2])
        end

        Base.$extension(array::DD.AxisArrays.All, T::Type{<:DD.Scalars.All}, ::SpinAware) =
            $private(T, SpinCategory(array), SpinCategory(array), array)
        Base.$extension(array::DD.AxisArrays.All, ::SpinAware) =
            $private(eltype(array), SpinCategory(array), SpinCategory(array), array)
        Base.$extension(array::DD.AxisArrays.All, C::SpinCategory) =
            $private(eltype(array), C, SpinCategory(array), array)
        Base.$extension(array::DD.AxisArrays.All, T::Type{<:DD.Scalars.All},
                        wanted::SpinCategory) =
            $private(T, wanted, SpinCategory(array), array)

        Base.$extension(array::DD.DenseArrays.All, T::Type{<:DD.Scalars.All},
                        ::SpinDegenerate) = similar(array, T)
        Base.$extension(array::DD.DenseArrays.All, T::Type{<:DD.Scalars.All},
                        C::ColinearSpin{N}) where N = begin
            @lintpragma("Ignore use of undeclared variable array")
            @lintpragma("Ignore use of undeclared variable C")
            @lintpragma("Ignore use of undeclared variable T")
            @argcheck 1 ≤ N ≤ ndims(array)
            left, right = Base.IteratorsMD.split(size(array), Val{N})
            n = Base.last(left)
            @argcheck n == length(components(eltype(array), C))

            n′ = length(components(T, C))
            similar(array, T, Base.front(left)..., n′, right...)
        end
        Base.$extension(array::DD.DenseArrays.All, T::Type{<:DD.Scalars.All},
                        C::ColinearSpinLast) = begin
            @argcheck size(array, ndims(array)) == length(components(eltype(array), C))
            n′ = length(components(T, C))
            similar(array, T, Base.front(size(array))..., n′)
        end
    end
end

Base.reinterpret(T::Type{<: DD.Scalars.All}, ::SpinDegenerate, array::DenseArray) =
    AxisArray(reinterpret(concretize_type(T, array), array), axes(array))
Base.reinterpret(T::Type{<: DD.Scalars.All}, ::ColinearSpinFirst, array::DenseArray) =
    AxisArray(reinterpret(concretize_type(T, array), array),
              Axis{:spin}(components(T, ColinearSpinFirst())), Base.tail(axes(array))...)
Base.reinterpret(T::Type{<: DD.Scalars.All}, C::ColinearSpin, array::DenseArray) =
    AxisArray(reinterpret(concretize_type(T, array), array),
              Base.front(axes(array))..., Axis{:spin}(components(T, C)))

"""
    wrap([T::Type{<: Dispatch.Scalars.All}], [P::SpinCategory], array)

Wraps an existing array as a DFT array. The default spin-category is `SpinDegenerate`.
If the type and units are explicitly specified (first argument)), then the array must be
dimensionless. If they are not, then the array must have an element type `T` such that
`T <: Dispatch.Scalars.All`.
"""
wrap(T::Type{<: DD.Scalars.All}, P::SpinCategory, array::DenseArray) = begin
    @argcheck dimension(eltype(array)) == NoDims "Input array should be dimensionless"
    reinterpret(T, P, array)
end
wrap(::SpinDegenerate, array::DenseArray{<: DD.Scalars.All}) = AxisArray(array)
wrap(P::ColinearSpinFirst, array::DenseArray{<: DD.Scalars.All}) = begin
    c = components(eltype(array), P)
    @argcheck length(c) == size(array, 1) "Size of spin-axis does not match expectations"
    AxisArray(array, Axis{:spin}(c))
end
wrap(P::ColinearSpin, array::DenseArray{<: DD.Scalars.All}) = begin
    c = components(eltype(array), P)
    @argcheck(length(c) == size(array, ndims(array)),
              "Size of spin-axis does not match expectations")
    AxisArray(array, Base.front(axes(array))..., Axis{:spin}(c))
end
wrap(array::DenseArray{<: DD.Scalars.All}) = wrap(SpinDegenerate(), array)
wrap(T::Type{<: DD.Scalars.All}, array::DenseArray) =
    wrap(concretize_type(T, array), SpinDegenerate(), array)

""" Helper functions to convert between arrays with different units and memory layout """
_convert(T::Type{<:DD.Scalars.All}, C::ColinearSpin,
         C′::ColinearSpin, array::DD.AxisArrays.All) = begin
    T′ = concretize_type(T, array)
    (T′ == T && C == C′) ? array: copy!(similar(array, T, C), array)
end

"""
Converts axis to the requisite spin-axis location

$(SIGNATURES)

If the spin-axis does not move, then a reference to the original array is returned.
Otherwise, a new array is returned.
"""
Base.convert(::SpinAware, array::DD.AxisArrays.All) = array
Base.convert(T::Type{<: DD.Scalars.All}, ::SpinAware, array::DD.AxisArrays.All) =
    _convert(concretize_type(T, array), SpinCategory(array), SpinCategory(array), array)
Base.convert(C::ColinearSpin, array::DD.AxisArrays.All) =
    _convert(eltype(array), C, SpinCategory(array), array)
Base.convert(T::Type{<: DD.Scalars.All}, C::SpinCategory, array::DD.AxisArrays.All) =
    _convert(T, C, SpinCategory(array), array)
Base.convert(T::Type{<: DD.Scalars.All}, array::DD.AxisArrays.All) =
    _convert(T, SpinCategory(array), SpinCategory(array), array)

Unitful.uconvert(u::Unitful.Units, array::AxisArray) = begin
    data = similar(array.data, typeof(uconvert(u, oneunit(eltype(array)))))
    data .= array.data
    AxisArray(data, axes(array))
end

end


@testset "smarts.smiles" begin

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

    methane = parse(SMILES, "C")
    @test nodeattr(methane, 1).symbol == :C
    @test edgecount(methane) == 0

    EtOH = parse(SMILES, "CCO")
    @test nodeattr(EtOH, 3).symbol == :O

    EtOH2 = parse(SMILES, "[CH3][CH2][OH]")
    @test nodecount(EtOH2) == 9
    @test edgecount(EtOH2) == 8

    nitro = parse(SMILES, "N#N")
    @test edgeattr(nitro, 1).order == 3

    vinylcl = parse(SMILES, "C=CCl")
    @test edgeattr(vinylcl, 1).order == 2
    @test nodeattr(vinylcl, 3).symbol == :Cl
end

@testset "branch" begin
    tBuOH = parse(SMILES, "CC(C)(C)O")
    @test nodeattr(tBuOH, 5).symbol == :O
    @test degree(tBuOH, 2) == 4

    TFA = parse(SMILES, "FC(F)(F)C(=O)O")
    @test edgeattr(TFA, 5).order == 2
    @test degree(TFA, 2) == 4

    EDTA = parse(SMILES, "C(CN(CC(=O)O)CC(=O)O)N(CC(=O)O)CC(=O)O")
    @test degree(EDTA, 5) == 3
    @test degree(EDTA, 12) == 3
    @test degree(EDTA, 14) == 3
end

@testset "ring" begin
    cychexe = parse(SMILES, "C1C=CCCC1")
    @test degree(cychexe, 1) == 2

    cychexe2 = parse(SMILES, "C=1CCCCC1")
    @test edgeattr(cychexe2, 6).order == 2

    cychexe3 = parse(SMILES, "C1CCCCC=1")
    @test edgeattr(cychexe3, 6).order == 2

    cychexe4 = parse(SMILES, "C=1CCCCC=1")
    @test edgeattr(cychexe4, 6).order == 2

    bicyclohexyl = parse(SMILES, "C1CCCCC1C1CCCCC1")
    @test edgeattr(bicyclohexyl, 6).order == 1
    @test edgeattr(bicyclohexyl, 13).order == 1

    bicyclohexyl2 = parse(SMILES, "C1CCCCC1C2CCCCC2")
    @test edgeattr(bicyclohexyl2, 6).order == 1
    @test edgeattr(bicyclohexyl2, 13).order == 1

    naphthalene = parse(SMILES, "C=1C=CC=C2C1C=CC=C2")
    @test edgeattr(naphthalene, 6).order == 2
    @test edgeattr(naphthalene, 11).order == 1

    naphthalene2 = parse(SMILES, "C1=CC=C2C(=C1)C=CC=C2")
    @test edgeattr(naphthalene2, 6).order == 1
    @test edgeattr(naphthalene2, 11).order == 1

    tetrahedrane = parse(SMILES, "C12C3C1C23")
    @test edgecount(tetrahedrane) == 6

    thiamin = parse(SMILES, "OCCc1c(C)[n+](=cs1)Cc2cnc(C)nc(N)2")
    @test edgecount(thiamin) == 19
end

@testset "component" begin
    NaCl = parse(SMILES, "[Na+].[Cl-]")
    @test nodeattr(NaCl, 1).symbol == :Na
    @test nodeattr(NaCl, 1).charge == 1
    @test nodeattr(NaCl, 2).symbol == :Cl
    @test nodeattr(NaCl, 2).charge == -1

    ferrous = parse(SMILES, "[Fe++]")
    @test nodeattr(ferrous, 1).symbol == :Fe
    @test nodeattr(ferrous, 1).charge == 2

    ferrous2 = parse(SMILES, "[Fe+2]")
    @test nodeattr(ferrous2, 1).symbol == :Fe
    @test nodeattr(ferrous2, 1).charge == 2

    CuSO4 = parse(SMILES, "[Cu+2].[O-]S(=O)(=O)[O-]")
    @test nodeattr(CuSO4, 1).charge == 2
    @test nodeattr(CuSO4, 6).charge == -1
    @test degree(CuSO4, 3) == 4

    PhONa = parse(SMILES, "c1cc([O-].[Na+])ccc1")
    @test degree(PhONa, 3) == 3
    @test degree(PhONa, 5) == 0
end

@testset "stereo" begin
    LAla = parse(SMILES, "N[C@H](C)C(=O)O")
    @test degree(LAla, 2) == 4
    @test nodeattr(LAla, 2).stereo == :anticlockwise
    DAla = parse(SMILES, "N[C@@H](C)C(=O)O")
    @test degree(DAla, 2) == 4
    @test nodeattr(DAla, 2).stereo == :clockwise
end

@testset "aromatic" begin
    benzene = parse(SMILES, "c1ccccc1")
    @test nodeattr(benzene, 1).isaromatic
    @test edgeattr(benzene, 6).order == 1
    pyrrole = parse(SMILES, "[nH]1cccc1")  # pyrrole H cannot be omitted
    @test nodeattr(pyrrole, 1).isaromatic
    @test degree(pyrrole, 1) == 3
end

end # smarts.smiles


module MDTest

using Test, PIMD

@testset "MD" begin
    phase = HamiltonianPhase([1., 2.], [3., 4.])
    @test phase.q ≈ [1., 2.]
    phase.p .= 0
    @test phase.p ≈ [0., 0.]
end

end


using Libtask

@testset "tapedfunction" begin
    # Test case 1: stack allocated objects are deep copied.
    @testset "Instruction{typeof(__new__)}" begin
        mutable struct S
            i::Int
            S(x, y) = new(x + y)
        end

        tf = Libtask.TapedFunction(S, 1, 2)
        s1 = tf(1, 2)
        @test s1.i == 3
        newins = findall(x -> isa(x, Libtask.Instruction{typeof(Libtask.__new__)}), tf.tape)
        @test length(newins) == 1
    end

    @testset "Compiled Tape" begin
        function g(x, y)
            if x>y
                r= string(sin(x))
            else
                r= sin(x) * cos(y)
            end
            return r
        end

        tf = Libtask.TapedFunction(g, 1., 2.)
        ctf = Libtask.compile(tf)
        r = ctf(1., 2.)

        @test typeof(r) === Float64
    end
end


# Function for plotting the performance of a specific kernel in order to illustrate
# performance difference between PMM and DRAM
using nGraph
using nGraph.Flux
using nGraph.NNlib

# This is the implementation of the convolution kernel for profiling.
# We do some addition tricks to make sure intermediate tensors exist.
function __kernel(x, w)
    x += ones(eltype(x), size(x))
    w += ones(eltype(x), size(w))
    z = NNlib.conv(x, w; pad = (1, 1, 1, 1))
    return log(z)
end

function profile_kernel()
    # Input data
    x = rand(Float32, 48, 48, 128, 128)
    # Input weights
    w = rand(Float32, 3, 3, 128, 128)

    # Create an Actualizer and do the profiling.
    f = () -> AutoTM.Actualizer(__kernel, x, w)
    kernel_profile(f)
end

# Plot the kernel performance
function plot_kernel_performance(;
        file = "plot.pdf",
        width = "15cm",
        height = "5cm",
    )

    # Step 1: Get the kernel from the cache
    cache = AutoTM.Profiler.CPUKernelCache(AutoTM.Experiments.CPU_CACHE)
    configs = (AutoTM.Utils.DRAM, AutoTM.Utils.PMEM)

    expected_input_sizes = (
        (48, 48, 128, 128),
        (3, 3, 128, 128),
    )

    times = Float64[]
    for i in Iterators.product(configs, configs)
        for o in configs
            config = AutoTM.Utils.IOConfig(wrap(i), wrap(o))
            for (k,v) in cache.cache
                params = k[1]
                io_config = k[2]
                if config == io_config &&
                        params.description == "Convolution" &&
                        params.input_sizes == expected_input_sizes
                    push!(times, v)
                end
            end
        end
    end

    # Normalize times to a multple of the "all dram" case.
    times .= times ./ first(times)

    # Now that we've found the kernel times, we generate the plot!!
    xticks = [
        "DRAM\\\\DRAM\\\\DRAM",
        "DRAM\\\\DRAM\\\\\\textbf{PMM}",
        "DRAM\\\\\\textbf{PMM}\\\\DRAM",
        "DRAM\\\\\\textbf{PMM}\\\\\\textbf{PMM}",

        "\\textbf{PMM}\\\\DRAM\\\\DRAM",
        "\\textbf{PMM}\\\\DRAM\\\\\\textbf{PMM}",
        "\\textbf{PMM}\\\\\\textbf{PMM}\\\\DRAM",
        "\\textbf{PMM}\\\\\\textbf{PMM}\\\\\\textbf{PMM}",
    ]

    plot = Plot(Coordinates((i, t) for (i,t) in enumerate(times)))

    plt = TikzDocument()
        push!(plt, """
    \\pgfplotsset{
        width=$width,
        height=$height
    }
    """)

    tikz = TikzPicture()
        axs = @pgf Axis(
        {
            ybar,
            enlarge_x_limits=0.10,
            nodes_near_coords_align={vertical},
            ylabel="Performance relative\\\\to all IO in DRAM",
            ymajorgrids,
            ymin=0,
            ylabel_style={
                align = "center",
            },

            xtick="data",
            xticklabels = xticks,
            xticklabel_style={
                align = "center",
            },
            yticklabel_style={
                "/pgf/number format/fixed",
                "/pgf/number format/precision=5",
            },
            bar_width="20pt",
            clip=false,
        },
        plot,
        """
        \\node[align=right] at (axis cs: -0.1, -0.78) {
        \\textbf{Data In:}\\\\
        \\textbf{Weight:}\\\\
        \\textbf{Data Out:}\\\\
        };
        """,
    )

    push!(tikz, axs)
    push!(plt, tikz)
    pgfsave(file, plt)
    return plt
end


@testset "Storage data misspecification" begin
    # See https://discourse.julialang.org/t/how-to-use-test-warn/15557/5 about testing for warning throwing
    warn_message = "The data doesn't include devices of type GenericBattery, consider changing the device models"
    model = DeviceModel(PSY.GenericBattery, PSI.BookKeeping)
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5)
    @test_logs (:warn, warn_message) construct_device!(op_problem, :Storage, model)
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys14)
    @test_logs (:warn, warn_message) construct_device!(op_problem, :Storage, model)
end

@testset "Storage Basic Storage With DC - PF" begin
    model = DeviceModel(PSY.GenericBattery, PSI.BookKeeping)
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_bat)
    construct_device!(op_problem, :Storage, model)
    moi_tests(op_problem, false, 72, 72, 0, 0, 24, false)
    psi_checkobjfun_test(op_problem, GAEVF)
end

@testset "Storage Basic Storage With AC - PF" begin
    model = DeviceModel(PSY.GenericBattery, PSI.BookKeeping)
    op_problem = OperationsProblem(TestOpProblem, ACPPowerModel, c_sys5_bat)
    construct_device!(op_problem, :Storage, model)
    moi_tests(op_problem, false, 96, 96, 0, 0, 24, false)
    psi_checkobjfun_test(op_problem, GAEVF)
end

@testset "Storage with Reservation DC - PF" begin
    model = DeviceModel(PSY.GenericBattery, PSI.BookKeepingwReservation)
    op_problem = OperationsProblem(TestOpProblem, DCPPowerModel, c_sys5_bat)
    construct_device!(op_problem, :Storage, model)
    moi_tests(op_problem, false, 96, 24, 48, 48, 24, true)
    psi_checkobjfun_test(op_problem, GAEVF)
end

@testset "Storage with Reservation With AC - PF" begin
    model = DeviceModel(PSY.GenericBattery, PSI.BookKeepingwReservation)
    op_problem = OperationsProblem(TestOpProblem, ACPPowerModel, c_sys5_bat)
    construct_device!(op_problem, :Storage, model)
    moi_tests(op_problem, false, 120, 48, 48, 48, 24, true)
    psi_checkobjfun_test(op_problem, GAEVF)
end


############################################################
# joPDEfull ################################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: January 2017

export joPDEfull, joPDEfullException, subsample

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

# Type for linear operator representing  Pr*A(m)^-1*Ps,
# i.e. it includes source and receiver projections

struct joPDEfull{DDT<:Number,RDT<:Number} <: joAbstractLinearOperator{DDT,RDT}
    name::String
    m::Integer
    n::Integer
	info::Info
	model::ModelAll
	srcGeometry::Geometry
	recGeometry::Geometry
	options::Options
	fop::Function              # forward
	fop_T::Nullable{Function}  # transpose
end

mutable struct joPDEfullException <: Exception
	msg :: String
end

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

function joModeling(info::Info,model::ModelAll, srcGeometry::Geometry, recGeometry::Geometry; options=Options(), DDT::DataType=Float32, RDT::DataType=DDT)
# JOLI wrapper for nonlinear forward modeling
	(DDT == Float32 && RDT == Float32) || throw(joPDEfullException("Domain and range types not supported"))

	# Determine dimensions
	if typeof(recGeometry) == GeometryOOC
		m = sum(recGeometry.nsamples)
	else
		m = 0
		for j=1:info.nsrc m += length(recGeometry.xloc[j])*recGeometry.nt[j] end

	end
	if typeof(srcGeometry) == GeometryOOC
		n = sum(srcGeometry.nsamples)
	else
		n = 0
		for j=1:info.nsrc n += length(srcGeometry.xloc[j])*srcGeometry.nt[j] end
	end

	if info.nsrc > 1
		srcnum = 1:info.nsrc
	else
		srcnum = 1
	end

	return F = joPDEfull{Float32,Float32}("Proj*F*Proj'", m, n, info, model, srcGeometry, recGeometry, options,
							  src -> time_modeling(model, srcGeometry, src.data, recGeometry, [], [], srcnum, 'F', 1, options),
							  rec -> time_modeling(model, srcGeometry, [], recGeometry, rec.data, [], srcnum, 'F', -1, options),
							  )
end

############################################################
## overloaded Base functions

# conj(joPDEfull)
conj{DDT,RDT}(A::joPDEfull{DDT,RDT}) =
	joPDEfull{DDT,RDT}("conj("*A.name*")",A.m,A.n,A.info,A.model,A.srcGeometry,A.recGeometry,A.options,
		get(A.fop),
		A.fop_T
		)

# transpose(joPDEfull)
transpose{DDT,RDT}(A::joPDEfull{DDT,RDT}) =
	joPDEfull{DDT,RDT}("Proj*F'*Proj'",A.n,A.m,A.info,A.model,A.srcGeometry,A.recGeometry,A.options,
		get(A.fop_T),
		A.fop
		)

# ctranspose(joPDEfull)
ctranspose{DDT,RDT}(A::joPDEfull{DDT,RDT}) =
	joPDEfull{DDT,RDT}("Proj*F'*Proj'",A.n,A.m,A.info,A.model,A.srcGeometry,A.recGeometry,A.options,
		get(A.fop_T),
		A.fop
		)

############################################################
## overloaded Base *(...joPDEfull...)

# *(joPDEfull,joData)
function *{ADDT,ARDT,vDT}(A::joPDEfull{ADDT,ARDT},v::joData{vDT})
	A.n == size(v,1) || throw(joPDEfullException("shape mismatch"))
	if compareGeometry(A.srcGeometry,v.geometry) == false && compareGeometry(A.recGeometry,v.geometry) == false
		throw(joPDEfullException("geometry mismatch"))
	end
	jo_check_type_match(ADDT,vDT,join(["DDT for *(joPDEfull,joData):",A.name,typeof(A),vDT]," / "))
	V = A.fop(v)
	jo_check_type_match(ARDT,eltype(V),join(["RDT from *(joPDEfull,joData):",A.name,typeof(A),eltype(V)]," / "))
	return V
end

# *(num,joPDEfull)
function *{ADDT,ARDT}(a::Number,A::joPDEfull{ADDT,ARDT})
	return joPDEfull{ADDT,ARDT}("(N*"*A.name*")",A.m,A.n,A.info,A.model,A.srcGeometry,A.recGeometry,A.options,
		v1 -> jo_convert(ARDT,a*A.fop(v1),false),
		v2 -> jo_convert(ADDT,a*A.fop_T(v2),false)
		)
end

############################################################
## overloaded Bases +(...joPDEfull...), -(...joPDEfull...)

# +(joPDEfull,num)
function +{ADDT,ARDT}(A::joPDEfull{ADDT,ARDT},b::Number)
	return joPDE{ADDT,ARDT}("("*A.name*"+N)",A.m,A.n,A.info,A.model,A.srcGeometry,A.recGeometry,A.options,
		v1 -> A.fop(v1)+joConstants(A.m,A.n,b;DDT=ADDT,RDT=ARDT)*v1,
		v2 -> get(A.fop_T)(v2)+joConstants(A.m,A.n,b;DDT=ADDT,RDT=ARDT)*v2
		)
end

# -(joPDEfull,num)
function -{ADDT,ARDT}(A::joPDEfull{ADDT,ARDT},b::Number)
	return joPDE{ADDT,ARDT}("("*A.name*"-N)",A.m,A.n,A.info,A.model,A.srcGeometry,A.recGeometry,A.options,
		v1 -> A.fop(v1)-joConstants(A.m,A.n,b;DDT=ADDT,RDT=ARDT)*v1,
		v2 -> get(A.fop_T)(v2)-joConstants(A.m,A.n,b;DDT=ADDT,RDT=ARDT)*v2
		)
end

# -(joPDEfull)
-{DDT,RDT}(A::joPDEfull{DDT,RDT}) =
	joPDEfull{DDT,RDT}("(-"*A.name*")",A.m,A.n,A.info,A.model,A.srcGeometry,A.recGeometry,A.options,
		v1 -> -A.fop(v1),
		v2 -> -get(A.fop_T)(v2)
		)

############################################################
## Additional overloaded functions

# Subsample Modeling operator
function subsample{ADDT,ARDT}(F::joPDEfull{ADDT,ARDT}, srcnum)

	srcGeometry = subsample(F.srcGeometry,srcnum)		# Geometry of subsampled data container
	recGeometry = subsample(F.recGeometry,srcnum)

	info = Info(F.info.n, length(srcnum), F.info.nt[srcnum])
	return joModeling(info, F.model, srcGeometry, recGeometry; options=F.options, DDT=ADDT, RDT=ARDT)
end

getindex(F::joPDEfull,a) = subsample(F,a)


using Pythia
cxx"""
#include "fastjet/PseudoJet.hh"
#include "fastjet/ClusterSequence.hh"
"""

jetdef = icxx"""
  double Rparam = 0.4;
  fastjet::Strategy               strategy = fastjet::Best;
  fastjet::RecombinationScheme    recombScheme = fastjet::E_scheme;
  new fastjet::JetDefinition(fastjet::kt_algorithm, Rparam,
                                      recombScheme, strategy);
"""
fjinputs = icxx"new std::vector <fastjet::PseudoJet>";

pythia = newPythia()

icxx"""$pythia->readString("Beams:idA = 2212");"""
icxx"""$pythia->readString("Beams:idB = 2212");"""
icxx"""$pythia->readString("Beams:eCM = 8000.");"""
icxx"""$pythia->readString("HardQCD:gg2qqbar = on");"""
icxx"""$pythia->readString("PhaseSpace:pTHatMin = 50.0");"""
icxx"""$pythia->init();"""

function cluster(jetDef,e)
    icxx"$fjInputs->resize(0);"
    for p in particles(e)
        # Final State only
        if !icxx"$p.isFinal();"
            continue
        end

        # No neutrinos
        idabs = abs(icxx"$p.id();")
        if idabs == 12 || idabs == 14 || idabs == 16
            continue
        end

        if abs(icxx"$p.eta();" > 3.6)
            continue
        end

        icxx"$fjInputs->push_back( fastjet::PseudoJet( $p.px(),
        $p.py(), $p.pz(), $p.e() ) );"
    end
    icxx"new fastjet::ClusterSequence(*$fjinputs, *$jetDef);"
end

function jets(jetDef,e)
    clusterSeq = cluster(jetDef,e)
    sortedJets = icxx"""
        vector <fastjet::PseudoJet> inclusiveJets;
        inclusiveJets = $clusterSeq->inclusive_jets(20.0);
        new vector <fastjet::PseudoJet> (sorted_by_pt(inclusiveJets));
    """
end

for e in take(events(pythia),10)
    icxx"$pythia->process->list()"
    sortedJets = jets(jetDef,e)
    icxx"""
        for(unsigned int j = 0; j < $sortedJets->size(); j++) {
        cout << "Jet #" << j << " (px,py,pz,E) = (" << (*$sortedJets)[j].px() << ", " << (*$sortedJets)[j].py() << ", " << (*$sortedJets)[j].pz() << ", "
        << (*$sortedJets)[j].e() << ")\n";
        }
    """
end


struct Element
    name::AbstractString
    refstate::AbstractString # reference state
    mass::Float64
    H298::Float64
    S298::Float64
end

abstract type AbstractGFunction end

mutable struct GFunction <: AbstractGFunction
    name::AbstractString
    temp::Tuple{<:Real,<:Real}
    temps::Vector{<:Real}
    exprs::Vector{AbstractString}
    funcname::AbstractString
    funcstr::AbstractString
end

function GFunction(name::AbstractString,
                   temp::Tuple{<:Real,<:Real},
                   funcstr::AbstractString)
    return GFunction(name, temp, Real[], AbstractString[], getfuncname(name), funcstr)
end

function GFunction(name::AbstractString,
                   temps::Vector{<:Real},
                   exprs::Vector{<:AbstractString})
    temp = (minimum(temps), maximum(temps))
    return GFunction(name, temp, temps, exprs, getfuncname(name), "")
end

const Sublattice = Vector{AbstractString}
const Constitution = Vector{Sublattice}

mutable struct Parameter <: AbstractGFunction
    name::AbstractString
    symbol::Char
    comb::Constitution
    order::Int
    temp::Tuple{<:Real,<:Real}
    temps::Vector{<:Real}
    exprs::Vector{AbstractString}
    funcname::AbstractString
    funcstr::AbstractString
end

function Parameter(name::AbstractString,
                   symbol::Char,
                   comb::Vector{<:Vector},
                   order::Int,
                   temps::Vector{<:Real},
                   exprs::Vector{<:AbstractString},
                   funcname::AbstractString)
    temp = (minimum(temps), maximum(temps))
    return Parameter(name, symbol, comb, order, temp, temps, exprs, funcname, "")
end

function Parameter(name::AbstractString,
                   symbol::Char,
                   comb::Vector{<:Vector},
                   order::Int,
                   temp::Tuple{<:Real,<:Real},
                   funcname::AbstractString,
                   funcstr::AbstractString)
    return Parameter(name, symbol, comb, order, temp,
                     Real[], AbstractString[], funcname, funcstr)
end

mutable struct Phase
    name::AbstractString
    state::Char
    type::AbstractString
    sites::Vector{<:Real}
    cons::Constitution
    params::Vector{Parameter}
end

function Phase(name::AbstractString,
               state::Char,
               sites::Vector{<:Real})
    return Phase(name, state, "%", sites, Sublattice[], Parameter[])
end

function Phase(name::AbstractString,
               state::Char,
               sites::Vector{<:Real},
               cons::Vector{<:Vector},
               params::Vector{Parameter})
    return Phase(name, state, "%", sites, cons, params)
end

function Phase(name::AbstractString,
               state::Char,
               type::AbstractString,
               sites::Vector{<:Real})
    phas = Phase(name, state, type, sites, Sublattice[], Parameter[])
end

struct TypeDefinition
    code::Char
    key::AbstractString # SEQ, GES, etc.
    args::AbstractString
end

struct Database
    elems::Vector{Element}
    funcs::Vector{GFunction}
    phass::Vector{Phase}
    types::Vector{TypeDefinition}
end

function Database(elems::Vector{Element},
                  funcs::Vector{GFunction},
                  phass::Vector{Phase})
    return Database(elems, funcs, phass, TypeDefinition[])
end

function Base.display(db::Database)
    println("Database Summary:")
    nelem = length(db.elems)
    println("\tElements: $nelem")
    nfunc = length(db.funcs)
    println("\tFunctions: $nfunc")
    nphas = length(db.phass)
    println("\tPhases: $nphas")
end

function Base.print(db::Database)
    println("Database:")

    nelem = length(db.elems)
    println("\tElements: $nelem")
    for elem in db.elems
        println("\t\t$(elem.name)")
    end

    nfunc = length(db.funcs)
    println("\tFunctions: $nfunc")
    for func in db.funcs
        println("\t\t$(func.name)")
    end

    nphas = length(db.phass)
    println("\tPhases: $nphas")
    for k in 1:nphas
        phas = db.phass[k]
        print("\t\t$k: $(phas.name); ")
        println(constitutionstring(phas))
        for param in phas.params
            println("\t\t\t" * param.symbol * '(' * localname(param) * ')')
        end
    end
end


"""
    Decryptor

A `Decryptor` can be used to decrypt a `Ciphertext` instance back into a `Plaintext` instance.

See also: [`Plaintext`](@ref), [`Ciphertext`](@ref)
"""
mutable struct Decryptor <: SEALObject
  handle::Ptr{Cvoid}

  function Decryptor(context::SEALContext, secret_key::SecretKey)
    handleref = Ref{Ptr{Cvoid}}(C_NULL)
    retval = ccall((:Decryptor_Create, libsealc), Clong,
                   (Ptr{Cvoid}, Ptr{Cvoid}, Ref{Ptr{Cvoid}}),
                   context, secret_key, handleref)
    @check_return_value retval
    return Decryptor(handleref[])
  end

  function Decryptor(handle::Ptr{Cvoid})
    object = new(handle)
    finalizer(destroy!, object)
    return object
  end
end

function destroy!(object::Decryptor)
  if isallocated(object)
    @check_return_value ccall((:Decryptor_Destroy, libsealc), Clong, (Ptr{Cvoid},), object)
    sethandle!(object, C_NULL)
  end

  return nothing
end

function decrypt!(destination::Plaintext, encrypted::Ciphertext, decryptor::Decryptor)
  retval = ccall((:Decryptor_Decrypt, libsealc), Clong,
                 (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}),
                 decryptor, encrypted, destination)
  @check_return_value retval
  return destination
end

function invariant_noise_budget(encrypted::Ciphertext, decryptor::Decryptor)
  budgetref = Ref{Cint}(0)
  retval = ccall((:Decryptor_InvariantNoiseBudget, libsealc), Clong,
                 (Ptr{Cvoid}, Ptr{Cvoid}, Ref{Cint}),
                 decryptor, encrypted, budgetref)
  @check_return_value retval
  return Int(budgetref[])
end


import Random: CloseOpen01_64, CloseOpen12_64, FloatInterval, MT_CACHE_F

mutable struct MemorizingRNG{S} <: AbstractRNG
    memory::Vector{Float64}
    start::Int
    finish::Int
    idx::Int
    source::S
end

MemorizingRNG(source) = MemorizingRNG(Float64[], 1, 0, 0, source)

# Low level API (copied from MersenneTwister)
@inline mr_avail(r::MemorizingRNG) = r.finish - r.idx
@inline mr_empty(r::MemorizingRNG) = r.idx == r.finish
@inline mr_pop!(r::MemorizingRNG) = @inbounds return r.memory[r.idx+=1]

@inline reserve_1(r::MemorizingRNG) = mr_empty(r) && gen_rand!(r, 1)
@inline reserve(r::MemorizingRNG, n::Integer) = mr_avail(r) < n && gen_rand!(r, n)

# precondition: !mr_empty(r)
@inline rand_inbounds(r::MemorizingRNG, ::Type{CloseOpen12_64}) = mr_pop!(r)
@inline rand_inbounds(r::MemorizingRNG, ::Type{CloseOpen01_64}) = rand_inbounds(r, CloseOpen12_64) - 1.0
@inline rand_inbounds(r::MemorizingRNG) = rand_inbounds(r, CloseOpen01_64)

Random.rng_native_52(rng::MemorizingRNG) = Random.rng_native_52(rng.source)

function rand(r::MemorizingRNG, ::Union{I, Type{I}}) where I <: FloatInterval
    reserve_1(r)
    rand_inbounds(r, I)
end

# rand(r::MemorizingRNG, ::Type{Float64}) = rand(r, CloseOpen01)

function gen_rand!(r::MemorizingRNG{MersenneTwister}, n::Integer)
    len = length(r.memory)
    if len < r.idx + n
        resize!(r.memory, len+MT_CACHE_F)
        Random.gen_rand(r.source) # could be faster to use dsfmt_fill_array_close1_open2
        r.memory[len+1:end] = r.source.vals
        Random.mt_setempty!(r.source)
    end
    r.finish = length(r.memory)
    return r
end


mutable struct Feature{T}
    fixation_time::Float64
    visible::Bool
	value::T
end

Feature(;fixation_time=0.0, visible=false, value) = Feature(fixation_time, visible, value)

function populate_features(features, values)
	vals = [Feature(value=v) for v in values]
	return NamedTuple{features}(vals)
end

mutable struct VisualObject{F,T} <:AbstractVisualObject
	features::F
    attended::Bool
    visible::Bool
	width::Float64
	angular_size::Float64
	activation::T
	bottomup_activation::Float64
	topdown_activation::Float64
	location::Vector{Float64}
	attend_time::Float64
	target::Bool
end

function VisualObject(;features, attended=false, visible=false, width=0.0, angular_size=0.0, location=[0.0,0.0],
		attend_time=0.0, target=false, activation=0.0)
	return VisualObject(features, attended, visible, width, 0.0, activation, 0.0, 0.0, location, attend_time,
		target)
end

"""
* `viewing_distance`: distance between model and screen in inches (30 inches)
* `current_time`: current model processing time during a trial
* `topdown_weight`: a weight for the influence of top-down activation (1.1)
* `bottomup_weight`: a weight for the influence of bottom-up activation (0.4)
* `noise`: noise added to visual activation
* `persistence`: time during which visible visual objects can stay in iconic memory
* `acuity`: NamedTuple of acuity parameters
* `n_finst`: number of visual objects in finst
* `finst_span`: the duration of finst
* `τₐ`: activation threshold for terminating search
*  `Δτ`: activation threshold increment following a fixation to a distractor
"""
mutable struct Parm{A,T,T1} <: AbstractParms
    viewing_distance::Float64
	topdown_weight::T1
	bottomup_weight::Float64
	noise::Bool
	rnd_time::Bool
	σ::Float64
	persistence::Float64
	acuity::A
	n_finst::Int64
	finst_span::Float64
	β₀exe::Float64
	Δexe::Float64
	τₐ::Float64
	Δτ::Float64
	misc::T
end

function Parm(;viewing_distance=30.0, topdown_weight=.66, bottomup_weight=1.1, noise=false, rnd_time=false,
	σ=.2*π/sqrt(3), persistence=4.0, a_color=.104, b_color=.85, a_shape=.142, b_shape=.96, n_finst=4,
	finst_span=3.0, β₀exe=.02, Δexe=.002, τₐ=0.0, Δτ=0.39, args...)
	acuity = (color = (a=a_color,b=b_color), shape = (a=a_shape,b=b_shape))
	return Parm(viewing_distance, topdown_weight, bottomup_weight,
		noise, rnd_time, σ, persistence, acuity, n_finst, finst_span, β₀exe, Δexe, τₐ, Δτ, NamedTuple(args))
end

"""
**ACTRV**

ACTR model object
- `declarative`: declarative memory module
- `imaginal`: imaginal memory module
- `visual`: visual module
- `goal`: goal module
- `visual_location`: visual location module
- `parms`: model parameters
- `time`: model time

Constructor
````julia
ACTRV(;declarative=Declarative(), imaginal=Imaginal(), goal = Goal(),
    scheduler=nothing, visual=nothing, visual_location=nothing, time=0.0, parms...)
````
"""
@concrete mutable struct ACTRV <: AbstractACTR
    declarative
    imaginal
    visual
    visual_location
    goal
	visicon
    parms
    time
end

function ACTRV(;declarative=Declarative(), imaginal=Imaginal(), goal = Goal(),
    scheduler=nothing, visual=nothing, visual_location=nothing, visicon=nothing, time=0.0, parms...)
    parms′ = Parm(;parms...)
    ACTRV(declarative, imaginal, visual, visual_location, goal, visicon, parms′, time)
end

mutable struct Fixation
	target_color::Symbol
	target_shape::Symbol
	attend_time::Float64
	idx::Int
	stop::Bool
end

Fixation(;color, shape, attend_time, idx, stop) = Fixation(color, shape, attend_time, idx, stop)

mutable struct Data
	target_present::Symbol
	target_color::Symbol
	target_shape::Symbol
	trial_type::Symbol
	response::Symbol
	rt::Float64
end

Data() = Data(fill(:_,5)..., 0.0)

"""
** Experiment **

* `array_width`: with of visual array in pixels
* `n_cells`: number of cells in visual array grid
* `n_trials`: number of trials in simulation
* `cell_width`: width of cell that contains ≤ 1 visual objects
* `object_width`: width of visual object in pixels
* `n_color_distractors`: number of distractors for color
* `n_shape_distractors`: number of distractors for shape
* `set_size`: number of elements in visual array
* `colors`: tuple of colors
* `shapes`: tuple of shapes
* `base_rate`: probability that target is present
* `data`: Array of Data for all trials
* `trial_data`: Data for current trial
* `fixations`: Array of fixations for all trials
* `trial_fixations`: fixations for current trial
* `trace`: displays trace if true
* `window`: GUI window
* `canvas`: GUI canvas
* `visible`: displays GUI if true
* `speed`: how quickly to simulate the model if trace is on
* `populate_visicon`: a function that populates the visicon
"""
mutable struct Experiment{T1,T2,F<:Function}
	array_width::Float64
	n_cells::Int64
	n_trials::Int64
	cell_width::Float64
	object_width::Float64
	n_color_distractors::Int64
	n_shape_distractors::Int64
	set_size::Int64
	colors::Vector{Symbol}
	shapes::Vector{Symbol}
	base_rate::Float64
	data::Vector{Data}
	trial_data::Data
	fixations::Vector{Vector{Fixation}}
	trial_fixations::Vector{Fixation}
	trace::Bool
	window::T1
	canvas::T2
	visible::Bool
	speed::Float64
	populate_visicon::F
end

function Experiment(;array_width=428., object_width=32.0, n_cells=8, n_trials=20,
	n_color_distractors=20, set_size=40, n_shape_distractors=20, shapes=[:p,:q], colors=[:red,:blue],
	base_rate=.50, data=Data[], trial_data=Data(), fixations=Vector{Vector{Fixation}}(), trial_fixations=Fixation[],
	trace=false, window=nothing, canvas=nothing, visible=false, speed=1.0, populate_visicon=conjunctive_ratio)
	cell_width = array_width/n_cells
	@argcheck  cell_width > object_width
	@argcheck n_color_distractors + n_shape_distractors + 1 <= n_cells^2
	visible ? ((canvas,window) = setup_window(array_width)) : nothing
    visible ? Gtk.showall(window) : nothing
	return Experiment(array_width, n_cells, n_trials, cell_width, object_width,
		n_color_distractors, n_shape_distractors, set_size, colors, shapes, base_rate,
		data, trial_data, fixations, trial_fixations, trace, window, canvas, visible, speed, populate_visicon)
end

function import_gui()
    path = @__DIR__
    include(path*"/draw.jl")
end


abstract type AbstractSDEKernel
end

struct SDEKernel{fType,gType,tType,pType,NRPType} <: AbstractSDEKernel
  f::fType
  g::gType
  trange::tType
  p::pType
  noise_rate_prototype::NRPType
  constant_diffusity::Bool
end

struct SDEKernel!{fType,gType,gstep!Type,tType,pType,NRPType,wsType} <: AbstractSDEKernel
  f::fType
  g::gType
  gstep!::gstep!Type
  trange::tType
  p::pType
  noise_rate_prototype::NRPType
  ws::wsType
  constant_diffusity::Bool
end
function make_gstep(g) #g has signature  g(du,u,p,t)
  function gstep!(dx, ws, x, p, t, dw, noise_rate_prototype)
    if noise_rate_prototype===nothing
      g(ws, x, p, t)
      dx .+= ws .* dw
    else
      g(ws, x, p,t)
      dx .+= ws * dw
    end
  end
end

function SDEKernel(f,g,trange,p=nothing,noise_rate_prototype=nothing,constant_diffusity=false)
  SDEKernel{typeof(f),typeof(g),typeof(trange),typeof(p),typeof(noise_rate_prototype)}(f,g,trange,p,noise_rate_prototype,constant_diffusity)
end

function SDEKernel!(f,g,gstep!,trange,p=nothing,noise_rate_prototype=nothing,constant_diffusity=false; ws = copy(noise_rate_prototype))
  SDEKernel!{typeof(f),typeof(g),typeof(gstep!),typeof(trange),typeof(p),typeof(noise_rate_prototype),typeof(ws)}(f,g,gstep!, trange,p,noise_rate_prototype, ws, constant_diffusity)
end
abstract type AbstractFilteringAlgorithm end

struct CovarianceFilter <: AbstractFilteringAlgorithm end
struct InformationFilter <: AbstractFilteringAlgorithm end
struct LyapunovFilter <: AbstractFilteringAlgorithm end

struct solLyapunov{kernelType,ΣType,νTType,PTType,CTType}
  ktilde::kernelType
  Σ::ΣType
  νT::νTType
  PT::PTType
  cT::CTType
end

struct Message{kernelType,solType,sol2Type,tType,filterType}
  ktilde::kernelType
  sol::solType
  soldis::sol2Type
  ts::tType
  filter::filterType
end

function Message(sol, sdekernel::SDEKernel, filter, apply_timechange=false)
  soldis = construct_discrete_sol(sol)
  if OrdinaryDiffEq.isadaptive(sol.alg) || apply_timechange
    ts = reverse(sol.t)
  else
    ts = sdekernel.trange
  end
  Message{typeof(sdekernel),typeof(sol),typeof(soldis),typeof(ts),typeof(filter)}(sdekernel,
    sol,soldis,ts,filter)
end

struct GuidingDriftCache{kernelType,messageType}
  k::kernelType
  message::messageType
end

struct GuidingDiffusionVectorCache{gType}
  g::gType
end

struct GuidingDiffusionCache{gType,sizeType}
  g::gType
  padded_size::sizeType
end

abstract type AbstractRegression{inplace} end

struct Regression{kernelType,pJType,ifuncType,pfType,θType,duType,inplace} <: AbstractRegression{inplace}
  k::kernelType
  fjac::pJType
  ϕ0func::ifuncType
  pf::pfType
  θ::θType
  dy::duType
  isscalar::Bool
end

function Regression{iip}(sdekernel::Union{SDEKernel,SDEProblem}; paramjac=nothing,
    intercept=nothing, θ=sdekernel.p, yprototype=nothing,
    dyprototype=nothing, m=nothing) where iip

  if paramjac === nothing
    if iip
      pf = ParamJacobianWrapper2(sdekernel.f,first(sdekernel.trange),yprototype)
    else
      pf = ParamJacobianWrapper(sdekernel.f,first(sdekernel.trange),yprototype)
    end
  else
    pf = nothing
  end

  if !iip
    dy = nothing
  else
    dyprototype === nothing && error("dyprototype needs to be known for using inplace functions.")
    if m===nothing
      # default, diagonal noise
      dy = similar(dyprototype)
      isscalar = false
    else
      if m==1
        # scalar noise case
        dy = similar(dyprototype, length(dyprototype))
        isscalar = true
      else
        # non-diagonal noise
        dy = similar(dyprototype, length(dyprototype,m))
        isscalar = false
      end
    end
  end

  if m===nothing
    isscalar = false
  else
    isscalar = (m==1)
  end

  Regression{typeof(sdekernel),typeof(paramjac),typeof(intercept),typeof(pf),
    typeof(θ),typeof(dy),iip}(sdekernel,
    paramjac,intercept,pf,θ,dy,isscalar)
end

function Regression(sdekernel;kwargs...)
  iip=DiffEqBase.isinplace(sdekernel.g,4)
  Regression{iip}(sdekernel;kwargs...)
end

mutable struct Regression!{kernelType,pJType,ifuncType,phiType,uType,pfType,θType,duType,inplace} <: AbstractRegression{inplace}
  k::kernelType
  fjac!::pJType
  ϕ0func!::ifuncType
  ϕ::phiType
  ϕ0::uType
  y::uType
  y2::uType
  pf::pfType
  θ::θType
  dy::duType
  isscalar::Bool
end

function Regression!{iip}(sdekernel::Union{SDEKernel,SDEProblem}, yprototype;
      paramjac_prototype=nothing, paramjac=nothing, intercept=nothing, isdiagonal=true,
      θ=sdekernel.p,
      dyprototype=yprototype, m=length(yprototype)) where iip

  y = similar(yprototype)
  y2 = similar(y)
  ϕ0 = similar(y)

  if paramjac_prototype !== nothing
    ϕ = similar(paramjac_prototype)
  else
    ϕ = zeros((length(yprototype),length(θ)))
  end

  if paramjac === nothing
    if iip
      pf = ParamJacobianWrapper2(sdekernel.f,first(sdekernel.trange),y)
    else
      pf = ParamJacobianWrapper(sdekernel.f,first(sdekernel.trange),y)
    end
  else
    pf = nothing
  end

  if !iip
    dy = nothing
  else
    dyprototype === nothing && error("dyprototype needs to be known for using inplace functions.")
    if m==length(yprototype)
      # default, diagonal noise
      dy = similar(dyprototype)
    else
      if m==1
        # scalar noise case
        dy = similar(dyprototype, length(dyprototype))
      else
        # non-diagonal noise
        dy = similar(dyprototype, length(dyprototype,m))
      end
    end
  end

  isscalar = (m==1 && length(yprototype)!=1)

  Regression!{typeof(sdekernel),typeof(paramjac),typeof(intercept),typeof(ϕ),
    typeof(y),typeof(pf),typeof(θ),typeof(dy),iip}(sdekernel,paramjac,intercept,
    ϕ,ϕ0,y,y2,pf,θ,dy,isscalar)
end

function Regression!(sdekernel,yprototype;kwargs...)
  iip=DiffEqBase.isinplace(sdekernel.g,4)
  Regression!{iip}(sdekernel,yprototype;kwargs...)
end

# internal solvers
abstract type AbstractInternalSolver end

struct EulerMaruyama! <: AbstractInternalSolver end

struct DefaultForwardGuidingP end

"""
    tangent!(du, u, dz, P)

!! May change `du`, but has to return it. !!

"""
function tangent!
end

struct GuidedSDE{kType, mType}
  kernel::kType
  message::mType
end

struct GuidedSDE!{kType, mType}
  kernel::kType
  message::mType
end


"""
Helper function that inserts a new decision node in the search tree,
corresponding to the new time window event that is observed.
success_prob_fn takes two arguments: (i) reference time (ii) interaction event and returns a single floating
point probability of success.
"""
function insert_decision_node!(tree::SearchTree, agent_name::String, ie::InteractionEvent,
                                parent_id::Int, ref_time::Float64, utilval::Float64, success_prob_fn::Function)

    # Start with end index to add nodes
    idx = length(tree.nodes)

    # If ref_time is less than the start of the event, set true ref time to event start
    true_ref_time = (ref_time < ie.timestamps[START]) ? ie.timestamps[START] : ref_time
    @assert true_ref_time <= ie.timestamps[FINISH] "ref_time is $(ref_time) and ie is $(ie.timestamps)"

    # N.B. Decision nodes should be initialized with 0 util
    # Decision node corresponding to attempt
    idx += 1
    attempt_idx = idx
    dnode1 = DecisionNode(agent_name=agent_name, task_name=ie.task_name, attempt=true, timeval=true_ref_time,
                          util=0.0, idx=attempt_idx, utilset=false)
    push!(tree.nodes, dnode1)
    tree.parent_id[attempt_idx] = parent_id

    # Decision node corresponding to not attempt
    # Add node idx to decision leaf
    idx += 1
    no_attempt_idx = idx
    dnode2 = DecisionNode(agent_name=agent_name, task_name=ie.task_name, attempt=false, timeval=true_ref_time,
                          util=0.0, idx=no_attempt_idx, utilset=false)
    push!(tree.nodes, dnode2)
    tree.parent_id[no_attempt_idx] = parent_id
    new_decision_leaf = no_attempt_idx

    # set children of parent id to decision node idxs
    tree.child_ids[parent_id] = [attempt_idx, no_attempt_idx]

    # Use function to get probability of success
    # Arguments of success_prob are the ref_time and the time window
    success_prob = success_prob_fn(ref_time, ie)

    # For now, set util to -Inf because leaf will set it
    idx += 1
    fail_idx = idx
    onode1 = OutcomeNode(agent_name=agent_name, task_name=ie.task_name, outcome=FINISH,
                         timeval=ie.timestamps[FINISH], probability=1.0 - success_prob, util=0.0, idx=fail_idx, utilset=false)
    push!(tree.nodes, onode1)
    tree.parent_id[fail_idx] = attempt_idx

    idx += 1
    succ_idx = idx
    onode2 = OutcomeNode(agent_name=agent_name, task_name=ie.task_name, outcome=SUCCESS,
                         timeval=ie.timestamps[SUCCESS], probability=success_prob, util=utilval, idx=succ_idx, utilset=false)
    push!(tree.nodes, onode2)
    tree.parent_id[succ_idx] = attempt_idx

    # Set child IDs of attempt_idx and new outcome leaves to succ and fail
    tree.child_ids[attempt_idx] = [succ_idx, fail_idx]
    new_outcome_leaves = [succ_idx, fail_idx]

    return new_outcome_leaves, new_decision_leaf
end

"""
Generate the allocation policy tree for a given agent and set of tasks to consider.
The util_val_fn returns the utility of completing an interaction event successfully.
"""
function generate_search_tree!(server::GenericAllocation{A,AP}, agent_name::String,
                               tasks_to_consider::Set{String}, success_prob_fn::Function, util_val_fn::Function,
                               downtime::Float64) where {A,AP}

    # Must be at least one task to consider
    if isempty(tasks_to_consider)
        empty!(server.agent_prop_set[agent_name].tree)
        return
    end

    # Assume that updated time windows have been generated
    # and that agent is currently not occupied with task already
    # Have outcome leaf idxs and outcome decision idxs (when NOT ATTEMPTED)
    # If a new event subsumes all outcome leaf idxs it should NOT be added to not-attempted OUTCOME DECISION IDX
    tree = server.agent_prop_set[agent_name].tree
    interaction_events = server.agent_prop_set[agent_name].interaction_events

    outcome_leaf_idxs = Set{Int}()
    decision_leaf_idxs = Set{Int}()

    first_ie_stamp = Inf

    considered_ie = [ie for ie in interaction_events if ie.task_name in tasks_to_consider]

    num_to_consider = (length(considered_ie) > server.max_tasks_to_consider) ? server.max_tasks_to_consider : length(considered_ie)

    # IMP - The first interaction event timestamp may be shared by multiple
    # events; in that case, should consider all at root
    for ie in considered_ie[1:num_to_consider]

        ie_stamp = ie.timestamps[START]

        if ie_stamp < first_ie_stamp

            utilval = util_val_fn(ie)
            (new_outcome_leaves, new_decision_leaf) = insert_decision_node!(tree, agent_name, ie,
                                                                            0, server.current_time,
                                                                            utilval, success_prob_fn)
            for idx in new_outcome_leaves
                push!(outcome_leaf_idxs, idx)
            end
            push!(decision_leaf_idxs, new_decision_leaf)

            first_ie_stamp = ie_stamp

        else

            # Now iterate through outcome leaf nodes and compare
            new_outcome_leaf_idxs = Set{Int}()
            new_decision_leaf_idxs = Set{Int}()
            outcome_leaves_to_rm = Set{Int}()
            decision_leaves_to_rm = Set{Int}()

            for ol_idx in outcome_leaf_idxs
                onode = tree.nodes[ol_idx]
                @assert typeof(onode) == OutcomeNode # sanity check

                if onode.timeval <= ie.timestamps[START]
                    # Dominated - continue
                    continue
                end

                if onode.timeval >= ie.timestamps[FINISH] - downtime/2.0
                    continue
                end

                # Last case - new IE begins before outcome and ends after outcome
                # now add decision node to try and its two outcomes
                utilval = util_val_fn(ie)
                (temp_outcome_idxs, temp_dec_idx) = insert_decision_node!(tree, agent_name, ie,
                                                                          ol_idx, onode.timeval,
                                                                          utilval, success_prob_fn)
                push!(outcome_leaves_to_rm, ol_idx)

                for to_idx in temp_outcome_idxs
                    push!(new_outcome_leaf_idxs, to_idx)
                end
                push!(new_decision_leaf_idxs, temp_dec_idx)
            end

            # IMP - Not attempting is DIFFERENT from failing immediately, because it affects the dominance principle
            # Current no-attempt decision nodes MUST have had their attempt counterparts added already
            # So if new event is dominated, then it should not follow no-attempt leaf nodes, EVEN if otherwise valid
            @assert isempty(new_decision_leaf_idxs) == isempty(new_outcome_leaf_idxs)

            # Otherwise, loop through decision_leaf_idxs and do the same thing
            for dec_idx in decision_leaf_idxs
                dnode = tree.nodes[dec_idx]

                @assert dnode.attempt == false

                # We know that event is non-dominated by outcomes
                # and we did not jump forward in time, so do not need to do earlier checks
                if dnode.timeval >= ie.timestamps[FINISH] - downtime/2.0
                    continue
                end

                utilval = util_val_fn(ie)
                (temp_outcome_idxs, temp_dec_idx) = insert_decision_node!(tree, agent_name, ie,
                                                                        dec_idx, dnode.timeval,
                                                                        utilval, success_prob_fn)
                push!(decision_leaves_to_rm, dec_idx)

                for to_idx in temp_outcome_idxs
                    push!(new_outcome_leaf_idxs, to_idx)
                end
                push!(new_decision_leaf_idxs, temp_dec_idx)
            end

            # Finally, remove old leaves and add new ones
            setdiff!(outcome_leaf_idxs, outcome_leaves_to_rm)
            union!(outcome_leaf_idxs, new_outcome_leaf_idxs)

            setdiff!(decision_leaf_idxs, decision_leaves_to_rm)
            union!(decision_leaf_idxs, new_decision_leaf_idxs)
        end
    end

    # Now go up from leaves, assigning utilities to parent nodes
    # First assign utilities to actual leaves based on outcomes
    # Then, add parents to the current fringe
    rev_fringe_idxs = Set{Int64}()

    for ol_idx in outcome_leaf_idxs

        # Add to util of parent node
        ol_node = tree.nodes[ol_idx]
        ol_node.utilset = true

        push!(rev_fringe_idxs, ol_idx)
    end

    # Do the same for decision leaves, but now
    # Since there are only positive or zero utilities
    # can just do max with self
    for dec_idx in decision_leaf_idxs

        # this must be a non-attempt decision so util zero
        dec_node = tree.nodes[dec_idx]
        @assert dec_node.attempt == false
        dec_node.util = 0.0 # For good measure
        dec_node.utilset = true
        push!(rev_fringe_idxs, dec_idx)
    end

    # IMP - REV FRINGE IDXS can have their own children in rev_fringe_idxs too!
    # So keep adding and removing until empty
    # N.B - This is for setting PARENT UTIL
    while !isempty(rev_fringe_idxs)

        to_rm = Set{Int64}()
        to_add = Set{Int64}()

        for rfi in rev_fringe_idxs

            # If child of root, should remove (unless still needed)
            if tree.parent_id[rfi] == 0
                push!(to_rm, rfi)
                continue
            end

            node = tree.nodes[rfi]
            @assert node.utilset == true "Node in rev fringe does NOT have util set",node

            # If any sibling not set, continue WHILE keeping in rev_fringe_idxs
            par_id = tree.parent_id[rfi]
            par_node = tree.nodes[par_id]

            siblingset = true
            for child in tree.child_ids[par_id]
                if ~(child in rev_fringe_idxs)
                    siblingset = false
                    break
                end
            end
            if siblingset == false
                continue
            end

            # Do what is necessary for parent and add to rfi
            # Parent is decision node
            if typeof(node) == OutcomeNode
                par_node.util = par_node.util + node.probability*node.util
            else
                # Node type is decision
                # If parent is decision, just compare child decisions and choose max
                # If parent is outcome, add to exp. util if outcome success
                temp_node_util = node.util
                if typeof(par_node) == OutcomeNode && par_node.outcome == SUCCESS
                    temp_node_util = 1.0 + node.util
                end
                par_node.util = (temp_node_util > par_node.util) ? temp_node_util : par_node.util

            end

            push!(to_add, par_id)
            push!(to_rm, rfi)
        end

        for par_id in to_add
            tree.nodes[par_id].utilset = true
        end

        # First remove and then add
        union!(rev_fringe_idxs, to_add)
        setdiff!(rev_fringe_idxs, to_rm)
    end
end

# When there are multiple simultaneous objects, then it is not enough
# that the 'first' object is not attempted
function get_next_attempt_idx(tree::SearchTree)

    if isempty(tree)
        return -1
    end

    # Find the first node at which ATTEMPTING is the decision
    # Until that point, put all objs in entries
    dec_fringe = Deque{Int64}()
    push!(dec_fringe, 0)

    subtree_root_dec = -1

    while !isempty(dec_fringe)

        dec_top = popfirst!(dec_fringe)

        top_children = tree.child_ids[dec_top]

        @assert length(top_children) == 2 # TODO : Check this is valid
        attempt_idx = (tree.nodes[top_children[1]].attempt == true) ? top_children[1] : top_children[2]
        no_attempt_idx = top_children[1] + top_children[2] - attempt_idx

        if tree.nodes[no_attempt_idx].util > tree.nodes[attempt_idx].util
            push!(dec_fringe, no_attempt_idx)
        else
            subtree_root_dec = attempt_idx
        end
    end

    return subtree_root_dec
end


# System tests.
module TestSystem

using Orchestra.Types
using Orchestra.System
using Orchestra.Transformers
importall Orchestra.Util

include("fixture_learners.jl")
using .FixtureLearners
nfcp = NumericFeatureClassification()
fcp = FeatureClassification()

function all_concrete_subtypes(a_type::Type)
  a_subtypes = Type[]
  for a_subtype in subtypes(a_type)
    if isleaftype(a_subtype)
      push!(a_subtypes, a_subtype)
    else
      append!(a_subtypes, all_concrete_subtypes(a_subtype))
    end
  end
  return a_subtypes
end

concrete_learner_types = setdiff(
  all_concrete_subtypes(Learner),
  all_concrete_subtypes(TestLearner)
)

using FactCheck

facts("Orchestra system") do
  context("All learners train and predict on fixture data.") do
    for concrete_learner_type in concrete_learner_types
      learner = concrete_learner_type()
      fit_and_transform!(learner, nfcp)
    end

    @fact 1 => 1
  end

  context("All learners train and predict on iris dataset.") do
    # Get data
    dataset = readcsv(joinpath(dirname(@__FILE__), "iris.csv"))
    instances = dataset[:,1:(end-1)]
    labels = dataset[:, end]
    (train_ind, test_ind) = holdout(size(instances, 1), 0.3)
    train_instances = instances[train_ind, :]
    test_instances = instances[test_ind, :]
    train_labels = labels[train_ind]
    test_labels = labels[test_ind]

    # Test all learners
    for concrete_learner_type in concrete_learner_types
      learner = concrete_learner_type()
      fit!(learner, train_instances, train_labels)
      transform!(learner, test_instances)
    end

    @fact 1 => 1
  end

  context("Ensemble with learners from different libraries work.") do
    learners = Learner[]
    push!(learners, RandomForest())
    push!(learners, StackEnsemble())
    if LIB_SKL_AVAILABLE
      push!(learners, SKLLearner())
    end
    if LIB_CRT_AVAILABLE
      push!(learners, CRTLearner())
    end
    ensemble = VoteEnsemble({:learners => learners})
    predictions = fit_and_transform!(ensemble, nfcp)

    @fact 1 => 1
  end

  context("Pipeline works with fixture data.") do
    transformers = [
      OneHotEncoder(),
      Imputer(),
      StandardScaler(),
      BestLearner()
    ]
    pipeline = Pipeline({:transformers => transformers})
    predictions = fit_and_transform!(pipeline, fcp)

    @fact 1 => 1
  end
end

end # module


export Network

function Network(g::SimpleGraph,
                 vertexf::Union{VertexFunction, Vector{<:VertexFunction}},
                 edgef::Union{EdgeFunction, Vector{<:EdgeFunction}};
                 accumulator=+, accdim=:auto,
                 execution=:seq,
                 verbose=false)
    verbose && println("Create dynamic network with $(nv(g)) vertices and $(ne(g)) edges:")
    @check execution ∈ EXECUTION_STYLES "Exectuion type $execution not supportet (choose from $(EXECUTION_STYLES))"

    im = IndexManager()

    vtypes, idxs = batch_identical(vertexf, collect(1:nv(g)))
    vertexbatches = collect(VertexBatch(im, i, t; verbose) for (i, t) in zip(idxs, vtypes))

    nl = NetworkLayer(im, g, edgef, accumulator, accdim; verbose)

    @assert isdense(im) "The index structure is not dense"
    return Network{execution, typeof(nl), typeof(vertexbatches)}(vertexbatches, nl, im)
end

function VertexBatch(im::IndexManager,
                     vertex_idxs::Vector{Int},
                     vertexf::VertexFunction;
                     verbose)
    dim = vertexf.dim
    pdim = vertexf.pdim
    (firstidx, pfirstidx) = register_vertices!(im, vertex_idxs, dim, pdim)

    verbose && println(" - VertexBatch: dim=$dim, pdim=$pdim, length=$(length(vertex_idxs))")

    VertexBatch(vertex_idxs,
                vertexf,
                dim,
                firstidx,
                pdim,
                pfirstidx)
end

function NetworkLayer(im::IndexManager, g::SimpleGraph,
                      edgef::Union{EdgeFunction, Vector{<:EdgeFunction}},
                      accumulator, accdim;
                      verbose)
    @check hasmethod(accumulator, NTuple{2, AbstractFloat}) "Accumulator needs `acc(AbstractFloat, AbstractFloat) method`"

    accdim_max = maxaccdim(edgef)
    if accdim === :auto
        accdim = accdim_max
        verbose && println(" - auto accumulation dimension = $accdim")
    end
    @check accdim<=accdim_max "For this system acc dim is limited to $accdim_max by the edgefunctions"

    edgecolors = color_edges_greedy(g)
    verbose && println(" - found $(length(unique(edgecolors))) edgecolors (optimum would be $(maximum(degree(g))))")

    colors = unique(edgecolors)

    idx_per_color = [findall(isequal(c), edgecolors) for c in colors]

    edgef_per_color = batch_by_idxs(edgef, idx_per_color)

    colorbatches = collect(ColorBatch(im, g, idx, ef; verbose)
                         for (idx, ef) in zip(idx_per_color, edgef_per_color))

    NetworkLayer(g, colorbatches, accumulator, accdim, CachePool())
end

maxaccdim(e::EdgeFunction) = accdim(e)
maxaccdim(e::AbstractVector{<:EdgeFunction}) = minimum(accdim.(e))

function ColorBatch(im::IndexManager, g::SimpleGraph,
                    edge_idxs::Vector{Int},
                    edgef::Union{EdgeFunction, Vector{<:EdgeFunction}};
                    verbose)
    verbose && println(" - ColorBatch with $(length(edge_idxs)) edges")

    eftypes, idxs = batch_identical(edgef, edge_idxs)
    edgebatches = collect(EdgeBatch(im, g, i, ef; verbose) for (i, ef) in zip(idxs, eftypes))

    return ColorBatch(edge_idxs, edgebatches)
end

function EdgeBatch(im::IndexManager, g::SimpleGraph,
                   edge_idxs::Vector{Int},
                   edgef::EdgeFunction;
                   verbose)
    dim = edgef.dim
    pdim = edgef.pdim

    if edgef isa StaticEdge
        (firstidx, pfirstidx) = register_edges!(im, edge_idxs, 0, pdim)
    else
        (firstidx, pfirstidx) = register_edges!(im, edge_idxs, dim, pdim)
    end

    vidx_array = Array{Int}(undef, 6*length(edge_idxs))
    empty!(vidx_array)
    alledges = collect(edges(g))
    for (i, eidx) in enumerate(edge_idxs)
        edge = alledges[eidx]
        src, dst = edge.src, edge.dst
        src_r = vertex_data_range(im, src)
        dst_r = vertex_data_range(im, dst)
        src_idx, dst_idx = src_r[begin], dst_r[begin]
        src_dim, dst_dim = length(src_r), length(dst_r)

        append!(vidx_array, src, src_idx, src_dim, dst, dst_idx, dst_dim)
    end
    @assert length(vidx_array) == 6*length(edge_idxs)

    verbose && println("   - EdgeBatch: dim=$dim, pdim=$pdim, length=$(length(edge_idxs))")

    EdgeBatch(edge_idxs,
              edgef,
              dim,
              firstidx,
              pdim,
              pfirstidx,
              vidx_array)
end

batch_by_idxs(v, idxs::Vector{Vector{Int}}) = [v for batch in idxs]
function batch_by_idxs(v::AbstractVector, batches::Vector{Vector{Int}})
    @assert length(v) == sum(length.(batches))
    [v[batch] for batch in batches]
end

batch_identical(el, idxs::Vector{Int}) = [el], [idxs]
function batch_identical(v::Vector{T}, indices::Vector{Int}) where {T}
    @assert length(v) == length(indices)
    idxs_per_type = Vector{Int}[]
    types = T[]
    for i in eachindex(v)
        found = false
        for j in eachindex(types)
            if v[i] === types[j]
                found = true
                push!(idxs_per_type[j], indices[i])
                break
            end
        end
        if !found
            push!(types, v[i])
            push!(idxs_per_type, [indices[i]])
        end
    end
    @assert length(types) == length(idxs_per_type)
    return types, idxs_per_type
end


### Single type of emitters
function get_measurement(
    num_samples,
    receiver::AbstractReceiver,
    emitters::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{1} = IdealManifold(),
    rng = Random.GLOBAL_RNG
) where T <: AbstractFloat
    signal = StructArray{Complex{T}}(undef, num_samples)
    get_measurement!(signal, receiver, emitters, manifold, rng)
end

function get_measurement(
    num_samples,
    receiver::AbstractReceiver,
    emitters::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{N},
    rng = Random.GLOBAL_RNG
) where {N, T <: AbstractFloat}
    signal = StructArray{Complex{T}}(undef, num_samples, N)
    get_measurement!(signal, receiver, emitters, manifold, rng)
end

function get_measurement!(
    signal::Union{AbstractMatrix{Complex{T}}, AbstractVector{Complex{T}}},
    receiver::AbstractReceiver,
    emitters::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{N} = IdealManifold(),
    rng = Random.GLOBAL_RNG
) where {N, T <: AbstractFloat}
    num_samples = get_num_samples(signal)
    existing_emitters = filteremitters(get_existence, emitters)
    emitter_signals = gen_signal!.(
        existing_emitters,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    steer_vecs = get_steer_vecs(T, existing_emitters, receiver, manifold)
    randn!(rng, signal)
    signal .*= T(get_noise_std(receiver))
    foreach(emitter_signals, steer_vecs) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    measurement = multiply_gpmc_with_signal!(
        signal,
        get_gain_phase_mism_crosstalk(receiver)
    )

    next_receiver = propagate(receiver, num_samples, rng)
    next_emitters = propagate.(
        emitters,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    signal, next_receiver, next_emitters
end

### Two types of emitters
function get_measurement(
    num_samples,
    receiver::AbstractReceiver,
    emitters1::Tuple{Vararg{AbstractEmitter{T}}},
    emitters2::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{1} = IdealManifold(),
    rng = Random.GLOBAL_RNG
) where T <: AbstractFloat
    signal = StructArray{Complex{T}}(undef, num_samples)
    get_measurement!(signal, receiver, emitters1, emitters2, manifold, rng)
end

function get_measurement(
    num_samples,
    receiver::AbstractReceiver,
    emitters1::Tuple{Vararg{AbstractEmitter{T}}},
    emitters2::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{N},
    rng = Random.GLOBAL_RNG
) where {N, T <: AbstractFloat}
    signal = StructArray{Complex{T}}(undef, num_samples, N)
    get_measurement!(signal, receiver, emitters1, emitters2, manifold, rng)
end

function get_measurement!(
    signal::Union{AbstractMatrix{Complex{T}}, AbstractVector{Complex{T}}},
    receiver::AbstractReceiver,
    emitters1::Tuple{Vararg{AbstractEmitter{T}}},
    emitters2::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{N} = IdealManifold(),
    rng = Random.GLOBAL_RNG
) where {N, T <: AbstractFloat}
    num_samples = get_num_samples(signal)
    existing_emitters1 = filteremitters(get_existence, emitters1)
    existing_emitters2 = filteremitters(get_existence, emitters2)
    emitter_signals1 = gen_signal!.(
        existing_emitters1,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    emitter_signals2 = gen_signal!.(
        existing_emitters2,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    steer_vecs1 = get_steer_vecs(T, existing_emitters1, receiver, manifold)
    steer_vecs2 = get_steer_vecs(T, existing_emitters2, receiver, manifold)
    randn!(rng, signal)
    signal .*= T(get_noise_std(receiver))
    foreach(emitter_signals1, steer_vecs1) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    foreach(emitter_signals2, steer_vecs2) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    measurement = multiply_gpmc_with_signal!(
        signal,
        get_gain_phase_mism_crosstalk(receiver)
    )

    next_receiver = propagate(receiver, num_samples, rng)
    next_emitters1 = propagate.(
        emitters1,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    next_emitters2 = propagate.(
        emitters2,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    measurement, next_receiver, next_emitters1, next_emitters2
end

### Three types of emitters
function get_measurement(
    num_samples,
    receiver::AbstractReceiver,
    emitters1::Tuple{Vararg{AbstractEmitter{T}}},
    emitters2::Tuple{Vararg{AbstractEmitter{T}}},
    emitters3::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{1} = IdealManifold(),
    rng = Random.GLOBAL_RNG
) where {T <: AbstractFloat}
    signal = StructArray{Complex{T}}(undef, num_samples)
    get_measurement!(signal, receiver, emitters1, emitters2, emitters3, manifold, rng)
end

function get_measurement(
    num_samples,
    receiver::AbstractReceiver,
    emitters1::Tuple{Vararg{AbstractEmitter{T}}},
    emitters2::Tuple{Vararg{AbstractEmitter{T}}},
    emitters3::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{N},
    rng = Random.GLOBAL_RNG
) where {N, T <: AbstractFloat}
    signal = StructArray{Complex{T}}(undef, num_samples, N)
    get_measurement!(signal, receiver, emitters1, emitters2, emitters3, manifold, rng)
end

function get_measurement!(
    signal::Union{AbstractMatrix{Complex{T}}, AbstractVector{Complex{T}}},
    receiver::AbstractReceiver,
    emitters1::Tuple{Vararg{AbstractEmitter{T}}},
    emitters2::Tuple{Vararg{AbstractEmitter{T}}},
    emitters3::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{N} = IdealManifold(),
    rng = Random.GLOBAL_RNG
) where {N, T <: AbstractFloat}
    num_samples = get_num_samples(signal)
    existing_emitters1 = filteremitters(get_existence, emitters1)
    existing_emitters2 = filteremitters(get_existence, emitters2)
    existing_emitters3 = filteremitters(get_existence, emitters3)
    emitter_signals1 = gen_signal!.(
        existing_emitters1,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    emitter_signals2 = gen_signal!.(
        existing_emitters2,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    emitter_signals3 = gen_signal!.(
        existing_emitters3,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    steer_vecs1 = get_steer_vecs(T, existing_emitters1, receiver, manifold)
    steer_vecs2 = get_steer_vecs(T, existing_emitters2, receiver, manifold)
    steer_vecs3 = get_steer_vecs(T, existing_emitters3, receiver, manifold)
    randn!(rng, signal)
    signal .*= T(get_noise_std(receiver))
    foreach(emitter_signals1, steer_vecs1) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    foreach(emitter_signals2, steer_vecs2) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    foreach(emitter_signals3, steer_vecs3) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    measurement = multiply_gpmc_with_signal!(
        signal,
        get_gain_phase_mism_crosstalk(receiver)
    )

    next_receiver = propagate(receiver, num_samples, rng)
    next_emitters1 = propagate.(
        emitters1,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    next_emitters2 = propagate.(
        emitters2,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    next_emitters3 = propagate.(
        emitters3,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    signal, next_receiver, next_emitters1, next_emitters2, next_emitters3
end

### Four types of emitter
function get_measurement(
    num_samples,
    receiver::AbstractReceiver,
    emitters1::Tuple{Vararg{AbstractEmitter{T}}},
    emitters2::Tuple{Vararg{AbstractEmitter{T}}},
    emitters3::Tuple{Vararg{AbstractEmitter{T}}},
    emitters4::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{N},
    rng = Random.GLOBAL_RNG
) where {N, T <: AbstractFloat}
    signal = StructArray{Complex{T}}(undef, num_samples, N)
    get_measurement!(signal, receiver, emitters1, emitters2, emitters3, emitters4, manifold, rng)
end

function get_measurement!(
    signal::Union{AbstractMatrix{Complex{T}}, AbstractVector{Complex{T}}},
    receiver::AbstractReceiver,
    emitters1::Tuple{Vararg{AbstractEmitter{T}}},
    emitters2::Tuple{Vararg{AbstractEmitter{T}}},
    emitters3::Tuple{Vararg{AbstractEmitter{T}}},
    emitters4::Tuple{Vararg{AbstractEmitter{T}}},
    manifold::AbstractManifold{N} = IdealManifold(),
    rng = Random.GLOBAL_RNG
) where {N, T <: AbstractFloat}
    num_samples = get_num_samples(signal)
    existing_emitters1 = filteremitters(get_existence, emitters1)
    existing_emitters2 = filteremitters(get_existence, emitters2)
    existing_emitters3 = filteremitters(get_existence, emitters3)
    existing_emitters4 = filteremitters(get_existence, emitters4)
    emitter_signals1 = gen_signal!.(
        existing_emitters1,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    emitter_signals2 = gen_signal!.(
        existing_emitters2,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    emitter_signals3 = gen_signal!.(
        existing_emitters3,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    emitter_signals4 = gen_signal!.(
        existing_emitters4,
        get_sampling_frequency(receiver),
        get_intermediate_frequency(receiver),
        get_noise_density(receiver),
        num_samples,
        Ref(rng)
    )
    steer_vecs1 = get_steer_vecs(T, existing_emitters1, receiver, manifold)
    steer_vecs2 = get_steer_vecs(T, existing_emitters2, receiver, manifold)
    steer_vecs3 = get_steer_vecs(T, existing_emitters3, receiver, manifold)
    steer_vecs4 = get_steer_vecs(T, existing_emitters4, receiver, manifold)
    randn!(rng, signal)
    signal .*= T(get_noise_std(receiver))
    foreach(emitter_signals1, steer_vecs1) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    foreach(emitter_signals2, steer_vecs2) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    foreach(emitter_signals3, steer_vecs3) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    foreach(emitter_signals4, steer_vecs4) do emitter_signal, steer_vec
        signal .+= emitter_signal .* transpose(steer_vec)
    end
    measurement = multiply_gpmc_with_signal!(
        signal,
        get_gain_phase_mism_crosstalk(receiver)
    )

    next_receiver = propagate(receiver, num_samples, rng)
    next_emitters1 = propagate.(
        emitters1,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    next_emitters2 = propagate.(
        emitters2,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    next_emitters3 = propagate.(
        emitters3,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    next_emitters4 = propagate.(
        emitters4,
        num_samples,
        get_intermediate_frequency(receiver),
        get_sampling_frequency(receiver),
        get_noise_density(receiver),
        Ref(rng)
    )
    signal, next_receiver, next_emitters1, next_emitters2, next_emitters3, next_emitters4
end

function convert_complex_or_real(::Type{T}, a::Complex) where T <: Union{Float32, Float64}
    Complex{T}(a)
end

function convert_complex_or_real(::Type{T}, a) where T <: Union{Float32, Float64}
    T(a)
end

function get_steer_vecs(::Type{T}, emitters, receiver, manifold) where T <: Union{Float32, Float64}
    map(emitters) do emitter
        convert_complex_or_real.(T,
            get_steer_vec(manifold, emitter, get_attitude(receiver))
        )
    end
end

@inline function get_num_samples(signal)
    size(signal, 1)
end

function multiply_gpmc_with_signal!(signal, gpmc::Number)
    signal .*= gpmc
end

function multiply_gpmc_with_signal!(signal::AbstractMatrix{T}, gpmc::SMatrix) where T
    A_mul_Bt!(signal, gpmc)
end

function A_mul_Bt!(A::AbstractMatrix, B::SMatrix{N, N, T}) where {N, T}
    size(A, 2) == N || throw(DimensionMismatch("2nd dimension of A must have same dimension as of B"))
    @inbounds for i = 1:size(A, 1)
        v = SVector{N, T}(view(A, i, :))
        A[i, :] .= B * v
    end
    A
end

function A_mul_Bt!(A::StructArray, B::SMatrix{N, N, T}) where {N, T}
    size(A, 2) == N || throw(DimensionMismatch("2nd dimension of A must have same dimension as of B"))
    B_re = real.(B)
    B_im = imag.(B)
    @inbounds for i = 1:size(A, 1)
        v_re = SVector{N, T}(view(A.re, i, :))
        v_im = SVector{N, T}(view(A.im, i, :))
        A.re[i, :] .= B_re * v_re - B_im * v_im
        A.im[i, :] .= B_re * v_im + B_im * v_re
    end
    A
end