# Copyright (C) 2021 Ilya Baran.  This program is distributed under the terms of the MIT license.

using OrdinaryDiffEq
using MAT

abstract type AbstractBlock end

struct BlockOutputDefinition
    outputs::Variables
    computeOutputs::Function # Takes variables, parameters, time range, solution (in Variables form), and outputs, and returns outputs
    includeRawSolution::Bool # If true, computeOutputs takes a raw solution before solution-as-Variables
end

mutable struct BlockExtraData
    timeRange::Union{AbstractRange, Nothing}
    discontinuities::Union{Function, Nothing} # Takes parameters, timeRange, returns vector of discontinuities
    outputDefinition::Union{BlockOutputDefinition, Nothing}
    BlockExtraData() = new(nothing, nothing, nothing)
end

struct Block <: AbstractBlock
    variables::Variables
    parameters::Variables
    reactions::Array{AbstractReaction}
    extraData::BlockExtraData
    Block(v, p, r) = new(deepcopy(v), deepcopy(p), deepcopy(r), BlockExtraData());
end

function getExtraData(block::AbstractBlock)::BlockExtraData block.extraData end

function getTimeRange(block::AbstractBlock)::AbstractRange
    range = getExtraData(block).timeRange;
    if range === nothing
        error("Expected time range but it was missing");
    end
    range;
end

function setTimeRange!(block::AbstractBlock, timeRange::AbstractRange) getExtraData(block).timeRange = timeRange; end

function getOutputDefinition(block::AbstractBlock)::BlockOutputDefinition
    outputs = getExtraData(block).outputDefinition;
    if outputs === nothing
        error("Expected output definition but it was missing");
    end
    outputs;
end

function setOutputDefinition!(block::AbstractBlock, outputs::Variables, computeOutputs::Function)
    getExtraData(block).outputDefinition = BlockOutputDefinition(outputs, computeOutputs, false);
end

function setOutputDefinition!(block::AbstractBlock, outputs::Variables, computeOutputs::Function, includeRawSolution::Bool)
    getExtraData(block).outputDefinition = BlockOutputDefinition(outputs, computeOutputs, includeRawSolution);
end

function renameOutputs!(block::AbstractBlock, renameFunction::Function)
    oldDef = getOutputDefinition(block);
    newOutputs = renameVariables(oldDef.outputs, renameFunction);
    if oldDef.includeRawSolution
        newComputeFunction = (v, p, t, r, s, o) -> begin
            oldOutputs = oldDef.computeOutputs(v, p, t, r, s, deepcopy(oldDef.outputs));
            o.values = oldOutputs.values;
            return o;
        end;
    else
        newComputeFunction = (v, p, t, s, o) -> begin
            oldOutputs = oldDef.computeOutputs(v, p, t, s, deepcopy(oldDef.outputs));
            o.values = oldOutputs.values;
            return o;
        end;
    end
    setOutputDefinition!(block, newOutputs, newComputeFunction, oldDef.includeRawSolution);
end

function getDiscontinuities(block::AbstractBlock)::Vector
    d = getExtraData(block).discontinuities;
    if d === nothing
        return []
    end
    return d(getParameters(block), getTimeRange(block));
end

function setDiscontinuities!(block::AbstractBlock, discontinuities::T) where T <: Function
    getExtraData(block).discontinuities = discontinuities;
end

function runBlock(block::AbstractBlock)
    timeRange = getTimeRange(block);
    floatInterval::Tuple{Float64, Float64} = convert(Tuple{Float64, Float64}, (minimum(timeRange), maximum(timeRange)));
    problem = ODEProblem((x, p, t) -> computeDerivatives(block, t, x), Vector{Float64}(getVariables(block).values), floatInterval);
    solve(problem, Rodas4P(), tstops = getDiscontinuities(block));
end

function setParameters!(block::AbstractBlock, nameToValue::AbstractDict{String, T}) where T
    for (name, value) in nameToValue
        setParameter!(block, name, value);
    end
end

function solutionToVariables(solution, block::AbstractBlock, timeRange::AbstractRange)
    resultArray = Array{Array{Float64, 1}, 1}(undef, length(getVariables(block).variables));
    for index = 1:length(resultArray)
        resultArray[index] = solution(timeRange, idxs=index).u;
    end
    return Variables(getVariables(block), resultArray);
end

function solutionToMatrix(solution, timeRange::AbstractRange)::Matrix
    return Matrix(Array(solution(timeRange)));
end

function solutionToMatlab(solution, timeRange::AbstractRange, filename::String)
    t = Vector(timeRange);
    y = solutionToMatrix(solution, timeRange);
    status = string(solution.retcode);
    file = matopen(filename, "w");
    write(file, "t", t);
    write(file, "y", y);
    write(file, "status", status);
    close(file);
end

function computeDerivatives(block::Block, t::Number, x::Vector)::Vector
    derivatives = Variables(block.variables, zero(t * x));
    variables = Variables(block.variables, x);
    extraVariables = Dict{String, Any}();
    for reaction = block.reactions
        apply!(reaction, derivatives, t, variables, block.parameters, extraVariables);
    end
    #println("Time = $t");#, x = $x, result = $(derivatives.values)");
    vType = typeof(t * x);
    return vType(derivatives.values);
end

function getVariables(block::Block)::Variables block.variables; end
function getParameters(block::Block)::Variables block.parameters; end
function setParameter!(block::Block, name::String, value)
    block.parameters[name] = value;
end
function setParameters!(block::Block, values::Vector)
    block.parameters.values = values;
end

function computeOutputs(blockWithOutputs::AbstractBlock)
    timeRange = getTimeRange(blockWithOutputs);
    solution = runBlock(blockWithOutputs);
    return computeOutputsFromSolution(blockWithOutputs, timeRange, solution);
end

# ============================================ BlockWithBindings =======================================
struct BlockWithBindings <: AbstractBlock
    subblock::AbstractBlock
    parameterToBinding::Dict{String, Function} # Function takes t, variables, and existing parameters and returns a value
    parameters::Variables
    extraData::BlockExtraData
end

function BlockWithBindings(subblock::AbstractBlock, parameterToBinding::AbstractDict{String, T}) where T
    subblock = deepcopy(subblock);
    subblockParameters::Variables = getParameters(subblock);
    blockParameterToBinding = Dict{String, Function}();
    for (parameter, binding) in parameterToBinding
        if !haskey(subblockParameters.nameToIndex, parameter)
            error("Cannot bind nonexistent parameter $(parameter)");
        end
        if isa(binding, Function)
            blockParameterToBinding[parameter] = binding;
        else
            setParameter!(subblock, parameter, binding);
        end
    end
    newParameters = variablesSubtract(subblockParameters, Set(keys(parameterToBinding)));
    extraData = BlockExtraData();
    extraData.timeRange = getExtraData(subblock).timeRange;
    return BlockWithBindings(subblock, blockParameterToBinding, newParameters, extraData);
end

# Keep the parameters that are provided, and bind the rest to their current values
function BlockWithBindings(subblock::AbstractBlock, parametersToKeep::AbstractArray{String})
    subblock = deepcopy(subblock);
    subblockParameters::Variables = getParameters(subblock);
    newParameters::Array{Variable} = [];
    newValues::Array = [];
    for parameter in parametersToKeep
        if !haskey(subblockParameters.nameToIndex, parameter)
            error("Cannot keep nonexistent parameter $(parameter)");
        end
        index = subblockParameters.nameToIndex[parameter];
        variable = subblockParameters.variables[index];
        push!(newParameters, variable);
        push!(newValues, subblockParameters.values[index]);
    end
    newParametersV = Variables(newParameters);
    newParametersV.values = newValues;
    extraData = BlockExtraData();
    extraData.timeRange = getExtraData(subblock).timeRange;
    return BlockWithBindings(subblock, Dict{String, Function}(), newParametersV, extraData);
end

function getVariables(block::BlockWithBindings)::Variables getVariables(block.subblock); end
function getParameters(block::BlockWithBindings)::Variables block.parameters; end
function setParameter!(block::BlockWithBindings, name::String, value)
    block.parameters[name] = value;
    setParameter!(block.subblock, name, value);
end
function setParameters!(block::BlockWithBindings, values::Vector)
    block.parameters.values = values;
    for (name, index) in block.parameters.nameToIndex
        setParameter!(block.subblock, name, block.parameters.values[index]);
    end
end

function computeDerivatives(block::BlockWithBindings, t::Number, x::Vector)::Vector
    blockVariables = Variables(getVariables(block), x);
    for (parameter, binding) in block.parameterToBinding
        setParameter!(block.subblock, parameter, binding(t, blockVariables, block.parameters));
    end
    return computeDerivatives(block.subblock, t, x);
end

function computeOutputsFromSolution(block::AbstractBlock, timeRange, solution)
    outputDefinition = getOutputDefinition(block)
    if outputDefinition.includeRawSolution
        return outputDefinition.computeOutputs(getVariables(block), getParameters(block), timeRange,
                                            solution, solutionToVariables(solution, block, timeRange),
                                            deepcopy(outputDefinition.outputs));
    else
        return outputDefinition.computeOutputs(getVariables(block), getParameters(block), timeRange,
                                            solutionToVariables(solution, block, timeRange),
                                            deepcopy(outputDefinition.outputs));
    end
end

function recursiveComputeOutputs(block::AbstractBlock, timeRange, solution)
    computeOutputsFromSolution(block, timeRange, solution);
end

function recursiveComputeOutputs(block::BlockWithBindings, timeRange, solution)
    outputDefinition = getExtraData(block).outputDefinition;
    if (outputDefinition === nothing)
        return recursiveComputeOutputs(block.subblock, timeRange, solution);
    end
    computeOutputsFromSolution(block, timeRange, solution);
end

function computeOutputs(blockWithOutputs::BlockWithBindings)
    timeRange = getTimeRange(blockWithOutputs);
    solution = runBlock(blockWithOutputs);
    recursiveComputeOutputs(blockWithOutputs, timeRange, solution);
end

function getDiscontinuities(block::BlockWithBindings)::Vector
    d = getExtraData(block).discontinuities;
    if d === nothing
        return getDiscontinuities(block.subblock);
    end
    return d(getParameters(block), getTimeRange(block));
end

# ============================================ BlockCombo =======================================

const EXTRA_NAME = "_extra_";

struct BlockCombo <: AbstractBlock
    subblocks::Vector{Tuple{String, AbstractBlock}}
    parameters::Variables
    parameterNameToBlocksAndNames::Dict{String, Vector{Tuple{String, String}}} # For setting parameters
    glueFunctions::Vector{Tuple{String, String, Function}} # Block name, parameter name, function that computes it
    # Glue functions take time, variables, parameters (of the entire block)
    variables::Variables
    extraData::BlockExtraData
end

function BlockCombo(subblocks::Vector{Tuple{String, T}},
                    glueFunctions::Vector{Tuple{String, String, F}},
                    extraParameters::Variables) where {T <: AbstractBlock, F <: Function}
    if (length(subblocks) < 1)
        error("Must have at least one subblock");
    end
    extraData = deepcopy(getExtraData(subblocks[1][2]));
    subblocks = deepcopy(subblocks);
    gluedParameters = Set((f[1], f[2]) for f in glueFunctions);

    # Gather parameters
    seenNames = Set{String}();
    parameterArray = Vector{Variable}();
    parameterValueArray = [];
    parameterNameToBlocksAndNames = Dict{String, Vector{Tuple{String, String}}}();

    namesAndParameterSets = [(b[1], getParameters(b[2])) for b in subblocks];
    push!(namesAndParameterSets, (EXTRA_NAME, extraParameters));

    for namesAndParameters in namesAndParameterSets
        parameters = namesAndParameters[2];
        for parameter::Variable in parameters.variables
            if (in((namesAndParameters[1], parameter.name), gluedParameters))
                continue;
            end
            if !in(parameter.name, seenNames)
                push!(seenNames, parameter.name);
                push!(parameterArray, parameter);
                push!(parameterValueArray, parameters.values[parameters.nameToIndex[parameter.name]]);
            end
            if (namesAndParameters[1] != EXTRA_NAME)
                push!(get!(parameterNameToBlocksAndNames, parameter.name, []), (namesAndParameters[1], parameter.name));
            end
        end
    end

    parameters = Variables(parameterArray);

    variableArray = Vector{Variable}();
    variableValueArray = [];

    for nameAndSubblock in subblocks
        variables = getVariables(nameAndSubblock[2]);
        for variable in variables.variables
            push!(variableArray, variable);
            push!(variableValueArray, variables.values[variables.nameToIndex[variable.name]]);
        end
    end

    variables = Variables(variableArray);
    variables.values = variableValueArray;

    result = BlockCombo(subblocks, parameters, parameterNameToBlocksAndNames, glueFunctions, variables, extraData);
    setParameters!(result, parameterValueArray); # Propagate to subblocks
    result;
end

function getSubblock(block::BlockCombo, name::String)
    for (subblockName, subblock) in block.subblocks
        if name == subblockName
            return subblock;
        end
    end
    error("Subblock $name not found");
end

function getVariables(block::BlockCombo)::Variables block.variables; end
function getParameters(block::BlockCombo)::Variables block.parameters; end
function setParameter!(block::BlockCombo, name::String, value)
    block.parameters[name] = value;
    if haskey(block.parameterNameToBlocksAndNames, name)
        blocksAndNames = block.parameterNameToBlocksAndNames[name];
        for blockAndName in blocksAndNames
            setParameter!(getSubblock(block, blockAndName[1]), blockAndName[2], value);
        end
    end
end

function setParameters!(block::BlockCombo, values::Vector)
    block.parameters.values = values;

    for parameter in block.parameters.variables
        blocksAndNames = get(block.parameterNameToBlocksAndNames, parameter.name, nothing);
        if blocksAndNames === nothing
            continue;
        end
        for blockAndName in blocksAndNames
            setParameter!(getSubblock(block, blockAndName[1]), blockAndName[2], values[block.parameters.nameToIndex[parameter.name]]);
        end
    end
end

function computeDerivatives(block::BlockCombo, t::Number, x::Vector)::Vector
    blockVariables = Variables(block.variables, x);
    for nameParameterAndGlue in block.glueFunctions
        value = nameParameterAndGlue[3](t, blockVariables, block.parameters);
        setParameter!(getSubblock(block, nameParameterAndGlue[1]), nameParameterAndGlue[2], value);
    end

    derivatives = [];
    variableIndex = 1;
    for (_, subblock) in block.subblocks
        nVariables = length(getVariables(subblock).variables);
        push!(derivatives, computeDerivatives(subblock, t, x[variableIndex:(variableIndex + nVariables - 1)]))
        variableIndex += nVariables;
    end

    vcat(derivatives...);
end


#!/usr/bin/env julia

using Luxor, Colors

function test_ellipse(fname)
  pagewidth, pageheight = 1200, 1400
  Drawing(pagewidth, pageheight, fname)
  origin() # move 0/0 to center
  background("ivory")
  setopacity(0.5)
  setline(0.3)

  pagetiles = Tiler(pagewidth, pageheight, 4, 5, margin=50)
  for (pos, n) in pagetiles
    randomhue()
    ellipse(pos, rand(50:pagetiles.tilewidth), rand(50:pagetiles.tileheight), :strokepreserve)
    clip()
    for i in 1:10
      ellipse(pos, rand(50:pagetiles.tilewidth), rand(50:pagetiles.tileheight), :fill)
      randomhue()
      ellipse(pos, rand(50:pagetiles.tilewidth), rand(50:pagetiles.tileheight), :stroke)
    end
    fill()
    sethue("black")
    ellipse(pos, 5, 5, :fill)
    clipreset()
  end

  finish()
  println("ellipse-test saved in $(fname)")
end

fname = "/tmp/ellipse-test1.pdf"
test_ellipse(fname)


import JuMP
@testset "Selection" begin

    # Helper methods
    nt(name) = named(name)
    tf(name) = t -> nt(name)(t)
    iv(size, name="in") = inputvertex(name, size)
    av(in, outsize, name) = absorbvertex(IndMem(MatMul(nout(in), outsize)), in; traitdecoration=tf(name))
    tv(in, name) = invariantvertex(IndMem(identity, in, nout(in)), in; traitdecoration=tf(name))

    concindmem(ins) = f ->  IndMem(f, ins, sum(nout, ins))
    cc(ins...; name) = conc(ins...; dims=1, traitdecoration=tf(name), outwrap=concindmem(ins))
    nc(name) = traitconf(nt(name))

    @testset "Split relaxed exact" begin
        import NaiveNASlib: split_exact_relaxed, Relaxed, Exact

        @test (:a => relaxed(3)) == (:a => 3 => Relaxed())

        @test split_exact_relaxed((:a=>2, :b=>3=>Relaxed(), :c=>4=>Exact()),) == (Dict(:a => 2, :c => 4), Dict(:b => 3))
        @test split_exact_relaxed((Dict(:a=>2, :b=>3=>Relaxed(), :c=>4=>Exact())),) == (Dict(:a => 2, :c => 4), Dict(:b => 3))
    end

    @testset "Parselect" begin
        import NaiveNASlib: parselect
        mat = reshape(collect(1:3*4), 3, 4)

        @test parselect(mat, 2 => [2,-1,-1,4]; newfun = (args...) -> 0) == [4 0 0 10;5 0 0 11;6 0 0 12]
        @test parselect(mat, 1 => [-1,1,3]; newfun = (args...) -> 0) == [0 0 0 0;1 4 7 10;3 6 9 12]

        @test parselect(mat, 1 => [2,-1,3,-1], 2 => [-1,1,-1,4]; newfun = (args...) -> 0) == [0 2 0 11;0 0 0 0;0 3 0 12;0 0 0 0]

        dfun = (T, d, s...) ->  d == 1 ? 10 : -10

        @test parselect(ones(2,2), 1 => [-1, 1], 2 => [1 , -1], newfun=dfun) == [10 -10; 1 -10]
        @test parselect(ones(2,2), 2 => [1 , -1], 1 => [-1, 1], newfun=dfun) == [10 10; 1 -10]
    end

    @testset "Make ΔNout" begin
        import NaiveNASlib: Exact, Relaxed, fallback
        @testset "All Exact" begin
            res = ΔNout(:a => 2, :b => 3 => Exact())
            @test res isa ΔNout{Exact}
            @test res.Δs == Dict(:a => 2, :b => 3)

            @test base(fallback(res)) isa ΔNout{Relaxed}
        end

        @testset "All Relaxed" begin
            import NaiveNASlib: Exact, Relaxed
            res = ΔNout(:a => 2 => Relaxed(), :b => relaxed(3))
            @test res isa ΔNout{Relaxed}
            @test res.Δs == Dict(:a => 2, :b => 3)

            @test base(fallback(res)) isa ΔSizeFailNoOp
        end

        @testset "Mix" begin
            import NaiveNASlib: ΔNoutMix
            res = ΔNout(:a => 2, :b => relaxed(3), :c => 4, :d => relaxed(5))
            @test res isa ΔNoutMix
            @test res.exact.Δs == Dict(:a => 2, :c => 4)
            @test res.relax.Δs == Dict(:b => 3, :d => 5)

            @test base(fallback(res)) isa ΔNout{Relaxed}
        end

        @testset "Fallback" begin
            @test fallback(ΔNout(:a => 2; fallback=ThrowΔSizeFailError())) isa ThrowΔSizeFailError
            @test fallback(ΔNout(:a => relaxed(2); fallback=ThrowΔSizeFailError())) isa ThrowΔSizeFailError
            @test fallback(ΔNout(:a => relaxed(2), :b => 3; fallback=ThrowΔSizeFailError())) isa ThrowΔSizeFailError
        end
    end

    @testset "Join extraced inds" begin
        import NaiveNASlib: join_extracted_inds

        @test join_extracted_inds([2,3,4,5], [0,0,0,0]) == [2,3,4,5]
        @test join_extracted_inds([2,3,4,5], [1,0,0,0,0,0,1]) == [2,-1,3,4,5,-1]
        @test join_extracted_inds([2,3,4,5], [1,2,3,4]) == [2,-1,3,-1,-1,4,-1,-1,-1,5,-1,-1,-1,-1,]
        @test join_extracted_inds([2], [2,4,0]) == [2,-1,-1,-1,-1,-1,-1]
        @test join_extracted_inds([], [3,0]) == [-1,-1,-1]
        @test join_extracted_inds([1,2,4], [0]) == [1,2,4]
        @test join_extracted_inds([1,2,3,4,8], [0,1,2]) == [1, 2, -1, 3, -1, -1, 4, 8]
    end

    @testset "Wrap and scale utility function" begin
        using NaiveNASlib: remaputilityfun
        inpt = iv(8)
        v1 = av(inpt, 9, "v1")
        v2 = av(v1, 7, "v2")
        v3 = av(v2, 8, "v3")

        @testset "Empty valid set" begin
            remapped = remaputilityfun(v -> 1e8, DefaultJuMPΔSizeStrategy(), [inpt], 1e-2 ,1e5)
            @test remapped(inpt) == 1e8
        end

        @testset "All zero" begin
            remapped = remaputilityfun(v -> 0, DefaultJuMPΔSizeStrategy(), [inpt, v1], 1e-2 ,1e5)
            @test remapped(inpt) == 0
            @test remapped(v1) == 0
        end

        @testset "No scaling" begin
            remapped = remaputilityfun(v -> v == v2 ? 0.0 : ones(nout(v)), DefaultJuMPΔSizeStrategy(), [inpt, v1, v2, v3], 1e-2, 1e5)

            @test remapped(inpt) == ones(nout(inpt))
            @test remapped(v1) == ones(nout(v1))
            @test remapped(v2) == 0.0
            @test remapped(v3) == ones(nout(v3))
        end

        @testset "Scale up" begin
            remapped = remaputilityfun(v -> 2.0 .^ -(1:nout(v)), DefaultJuMPΔSizeStrategy(), [inpt, v1, v2, v3], 1e-2, 1e5)

            @test remapped(inpt) ≈ 2.0 .^ -(1:nout(inpt))
            @test remapped(v1) ≈ [2.56, 1.28, 0.64, 0.32, 0.16, 0.08, 0.04, 0.02, 0.01]
            @test remapped(v2) ≈ [2.56, 1.28, 0.64, 0.32, 0.16, 0.08, 0.04]
            @test remapped(v3) ≈ [2.56, 1.28, 0.64, 0.32, 0.16, 0.08, 0.04, 0.02]
        end
        @testset "Scale down" begin
            remapped = remaputilityfun(v -> 10.0 .^ (1:nout(v)), DefaultJuMPΔSizeStrategy(), [inpt, v1, v2, v3], 1e-2, 1e5)

            @test remapped(inpt) ≈ 10.0 .^ (1:nout(inpt))
            @test remapped(v1) ≈ [0.01, 0.1, 1.0, 10.0, 100.0, 1000.0, 10000.0, 100000.0, 1.0e6]
            @test remapped(v2) ≈ [0.01, 0.1, 1.0, 10.0, 100.0, 1000.0, 10000.0]
            @test remapped(v3) ≈ [0.01, 0.1, 1.0, 10.0, 100.0, 1000.0, 10000.0, 100000.0]
        end

        @testset "Saturate upper" begin
            # We want to scale up alot, but we hit the upper bound
            remapped = remaputilityfun(DefaultJuMPΔSizeStrategy(), [inpt, v1, v2, v3], 1e-2, 1e5) do v
                v == inpt && return fill(100, nout(v))
                v == v1 && return 10.0 .^-(1:nout(v))
                v == v2 && return (1:nout(v)) .*  1e-3
                v == v3 && return fill(10, nout(v))
            end

            @test remapped(inpt) == fill(100, nout(inpt))
            @test remapped(v1) ≈ 10.0 .^ (3:-1:-5)
            @test remapped(v2) ≈ 10.0:10.0:70.0
            @test remapped(v3) ≈ fill(1e5, nout(v3))
        end

        @testset "Saturate lower" begin
            # We want to scale down alot, but we hit the upper bound
            remapped = remaputilityfun(DefaultJuMPΔSizeStrategy(), [inpt, v1, v2, v3], 1e-2, 1e5) do v
                v == inpt && return fill(100, nout(v))
                v == v1 && return 10.0 .^(1:nout(v))
                v == v2 && return 1:nout(v)
                v == v3 && return fill(10, nout(v))
            end

            @test remapped(inpt) ≈ fill(100, nout(inpt))
            @test remapped(v1) ≈ [0.1, 1.0, 10.0, 100.0, 1000.0, 10000.0, 100000.0, 1.0e6, 1.0e7]
            @test remapped(v2) ≈ 0.01:0.01:0.07
            @test remapped(v3) ≈ fill(0.1, nout(v3))
        end

        @testset "Saturate both" begin
            # We want to scale both up and down but we hit both the upper and lower bound
            remapped = remaputilityfun(DefaultJuMPΔSizeStrategy(), [inpt, v1, v2, v3], 1e-2, 1e5) do v
                v == inpt && return 13
                v == v1 && return 1e6
                v == v2 && return 1e-3
                v == v3 && return fill(10, nout(v))
            end

            @test remapped(inpt) == 13
            @test remapped(v1) ≈ 1e6
            @test remapped(v2) ≈ 1e-3
            @test remapped(v3) == fill(10, nout(v3))
        end

        @testset "Mixed signs" begin
            remapped = remaputilityfun(DefaultJuMPΔSizeStrategy(), [inpt, v1, v2, v3], 1e-2, 1e5) do v
                v == inpt && return 13
                v == v1 && return 1e3 .* repeat([1, -1], nout(v))
                v == v2 && return 1e-3 .* repeat([1, -1], nout(v))
                v == v3 && return fill(10, nout(v))
            end

            @test remapped(inpt) == 13
            @test remapped(v1) ≈ 1e4 .* repeat([1, -1], nout(v1))
            @test remapped(v2) ≈ 1e-2 .* repeat([1, -1], nout(v2))
            @test remapped(v3) == fill(100, nout(v3))
        end
    end

    @testset "sizeobjective" begin

        using NaiveNASlib: selectmodel, createselectvars!, sizeobjective!

        function solvesize(utilityfun, vs; case = NaiveNASlib.NeuronIndices(), strat= DefaultJuMPΔSizeStrategy())
            model = selectmodel(case, strat, vs, utilityfun)

            outselectvars, outinsertvars, noutdict = createselectvars!(case, strat, vs, (;model, utilityfun))
            JuMP.set_upper_bound.(Iterators.flatten(values(outinsertvars)), 20)

            objexpr = sizeobjective!(case, strat, vs, (;model, outselectvars, outinsertvars, noutdict, utilityfun))

            JuMP.@objective(model, Max, objexpr)
            JuMP.optimize!(model)
            return noutdict
        end

        @testset "all positive utility $label" for (label, utilfun) in (
            ("scalar", v -> 1.0),
            ("array", v -> fill(1.0, nout(v))),
        )
            inpt = iv(3)
            v1 = av(inpt, 5, "v1")
            v2 = av(v1, 6, "v2")

            vs = all_in_graph(inpt)
            noutdict = solvesize(utilfun, vs)

            @testset "Size of $(name(v))" for v in vs
                @test JuMP.value(noutdict[v]) == nout(v)
            end
        end

        @testset "all negative utility $label" for (label, utilfun) in (
            ("scalar", v -> -1.0),
            ("array", v -> fill(-1.0, nout(v))),
        )

            inpt = iv(3)
            v1 = av(inpt, 5, "v1")
            v2 = av(v1, 6, "v2")

            vs = all_in_graph(inpt)
            noutdict = solvesize(utilfun, vs)

            @testset "Size of $(name(v))" for v in vs
                @test JuMP.value(noutdict[v]) == 1
            end
        end

        @testset "Want half size" begin
            inpt = iv(4)
            v1 = av(inpt, 6, "v1")
            v2 = av(v1, 8, "v2")

            vs = all_in_graph(inpt)
            noutdict = solvesize(vs) do v
                repeat([1, -1], nout(v) ÷ 2)
            end

            @testset "Size of $(name(v))" for v in vs
                @test JuMP.value(noutdict[v]) == nout(v) ÷ 2
            end
        end
    end

    @testset "Absorb 2 Absorb" begin
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = av(v1, 4, "v2")

        g = CompGraph(inpt, v2)

        @test Δnout!(v -> 1:nout(v), v1, -2)

        @test lastouts(v1) == lastins(v2) == [3,4,5]
        @test size(g(ones(3))) == (nout(v2),)

        @test Δnout!(v1, 3)

        @test lastouts(v1) == lastins(v2) == [1,2,3,-1,-1,-1]
        @test size(g(ones(3))) == (nout(v2),)

        @test Δnin!(v -> 1:nout(v), v2, -2)
        @test lastouts(v1) == lastins(v2) == [3,4,5,6]
        @test size(g(ones(3))) == (nout(v2),)
    end

    @testset "Argument plumbing" begin
        function graphgen()
            inpt = iv(3)
            v0 = av(inpt, 3, "v0")
            v1 = av(v0, 5, "v1")
            v2 = av(v1, 4, "v2")
            return v1, v2
        end

        @testset "Δnout!$(use_fun ? " with utility function" : "")" for use_fun in (false, true)
            f = use_fun ? (v -> (1:nout(v)),) : ()
            @testset "Single Δnout!$(wrap === identity ? "" : " $wrap")" for wrap in (identity, relaxed)
                v1,v2 = graphgen()
                @test Δnout!(f..., v1, wrap(3))
                @test nout.((v1, v2)) == (8, 4)
            end

            @testset "Single Δnout! pair$(wrap === identity ? "" : " $wrap")" for wrap in (identity, relaxed)
                v1,v2 = graphgen()
                @test Δnout!(f..., v1=>wrap(3))
                @test nout.((v1, v2)) == (8, 4)
            end

            @testset "Single Δnout! Dict$(wrap === identity ? "" : " $wrap")" for wrap in (identity, relaxed)
                v1,v2 = graphgen()
                @test Δnout!(f..., Dict(v1=>wrap(3)))
                @test nout.((v1, v2)) == (8, 4)
            end

            @testset "Multi Δnout! pair$(wrap === identity ? "" : " $wrap")" for wrap in (identity, relaxed)
                v1,v2 = graphgen()
                @test Δnout!(f..., v1=>wrap(3), v2=>wrap(2))
                @test nout.((v1, v2)) == (8, 6)
            end

            @testset "Multi Δnout! Dict$(wrap === identity ? "" : " $wrap")" for wrap in (identity, relaxed)
                v1,v2 = graphgen()
                @test Δnout!(f..., Dict(v1=>wrap(3), v2=>wrap(2)))
                @test nout.((v1, v2)) == (8, 6)
            end

            @testset "Multi Δnout! pair mix" begin
                v1,v2 = graphgen()
                @test Δnout!(f..., v1=>3, v2=>relaxed(2))
                @test nout.((v1, v2)) == (8, 6)
            end

            @testset "Multi Δnout! Dict mix" begin
                v1,v2 = graphgen()
                @test Δnout!(f..., Dict(v1=>3, v2=>relaxed(2)))
                @test nout.((v1, v2)) == (8, 6)
            end
        end

        @testset "Δnin!$(use_fun ? " with utility function" : "")" for use_fun in (false, true)
            f = use_fun ? (v -> (1:nout(v)),) : ()
            @testset "Single Δnin!$(wrap === identity ? "" : " $wrap")" for wrap in (identity, relaxed)
                v1,v2 = graphgen()
                @test Δnin!(f..., v1, wrap(3))
                @test nin.((v1, v2)) == ([6], [5])
            end

            @testset "Single Δnin! pair$(wrap === identity ? "" : " $(repr(wrap))")" for wrap in
            (identity, tuple, relaxed, relaxed ∘ tuple, tuple ∘ relaxed)
                v1,v2 = graphgen()
                @test Δnin!(f..., v1=>wrap(3))
                @test nin.((v1, v2)) == ([6], [5])
            end

            @testset "Single Δnin! Dict$(wrap === identity ? "" : " $(repr(wrap))")" for wrap in
            (identity, tuple, relaxed, relaxed ∘ tuple, tuple ∘ relaxed)
                v1,v2 = graphgen()
                @test Δnin!(f..., Dict(v1=>wrap(3)))
                @test nin.((v1, v2)) == ([6], [5])
            end

            @testset "Multi Δnin! pair mix $(repr(wrap1)) and $(repr(wrap2)) " for wrap1 in
            (identity, tuple, relaxed, relaxed ∘ tuple, tuple ∘ relaxed), wrap2 in
            (identity, tuple, relaxed, relaxed ∘ tuple, tuple ∘ relaxed)
                v1,v2 = graphgen()
                @test Δnin!(f..., v1=>wrap1(3), v2=>wrap2(2))
                @test nin.((v1, v2)) == ([6], [7])
            end

            @testset "Multi Δnin! Dict $(repr(wrap1)) and $(repr(wrap2)) " for wrap1 in
                (identity, tuple, relaxed, relaxed ∘ tuple, tuple ∘ relaxed), wrap2 in
                (identity, tuple, relaxed, relaxed ∘ tuple, tuple ∘ relaxed)
                v1,v2 = graphgen()
                @test Δnin!(f..., Dict(v1=>wrap1(3), v2=>wrap2(2)))
                @test nin.((v1, v2)) == ([6], [7])
            end
        end
    end

    @testset "add_participants" begin
        using NaiveNASlib: add_participants!
        using Logging: Info

        inpt = iv(3)
        v1 = av(inpt, 4, "v1")
        v2 = av(v1, 4, "v2")
        v3 = "v3" >> v1 + v2

        @testset "ΔNout $wraplabel" for (wraplabel, wrap) in (
            (LogΔSizeExec, s -> LogΔSizeExec("", Info, s)),
            (AlignNinToNout, s -> AlignNinToNout(vstrat=s)),
            (TruncateInIndsToValid, TruncateInIndsToValid),
            (WithUtilityFun, s -> WithUtilityFun(identity, s)),
            (TimeLimitΔSizeStrategy, s -> TimeLimitΔSizeStrategy(10, s)),
            (TimeOutAction,  s -> TimeOutAction(base=s)),
        )
            @test add_participants!(wrap(ΔNout(v2 => 3))) == [v2, v3 ,v1]
        end

        v4 = av(inpt, 5, "v4")
        @testset "ΔNoutMix" begin
            @test add_participants!(ΔNout(v1 => 2, v4 => relaxed(3), v2 => relaxed(2))) == [v1, v2, v3, v4]
        end
    end

    @testset "Absorb 2 Absorb fail error" begin
        import NaiveNASlib: ΔSizeFailError
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = av(v1, 4, "v2")

        @test_throws ΔSizeFailError Δsize!(v -> 1:nout(v), ThrowΔSizeFailError("Success!?"), v1)
    end

    @testset "Absorb 2 Absorb fail NoOp" begin
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = av(v1, 4, "v2")

        @test Δsize!(v -> 1:nout(v), ΔSizeFailNoOp(), v1) == false
    end

    @testset "Fail impossible" begin
        v0 = iv(3)
        v1 = tv(v0, "v1")
        v2 = av(v0, 3, "v2")
        v3 = "v3" >> v2 + v1
        @test Δnout!(v0, 1) == false
        @test @test_logs(
                (:warn, "Could not change nout of v3 by 1! Relaxing constraints..."),
                (:warn, "Could not change nout of v3 by 1! Attempt aborted!"),
                Δnout!(v3, 1)) == false
        @test @test_logs(
                (:warn, "Could not change nin of v3 by 1 and 1! Relaxing constraints..."),
                (:warn, "Could not change nin of v3 by 1 and 1! Attempt aborted!"),
                 Δnin!(v3, 1, 1)) == false
    end

    @testset "Δ = $wrap(0)" for wrap in (identity, relaxed)
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = av(v1, 4, "v2")

        @test Δnout!(v1 => wrap(0)) == true
        @test [nout(v1)] == nin(v2) == [5]
    end

    @testset "SizeStack duplicate" begin
        inpt = iv(3)
        v1 = av(inpt, 7, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = cc(v2, v1, name="v3")
        v4 = cc(v3, v2, name="v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(3,2))) == (nout(v4), 2)

        @test Δnout!(v4, -4)

        @test nout(v1) == 7
        @test nout(v2) == 2

        @test size(g(ones(3))) == (nout(v4),)

        @test Δnout!(v4, 6)

        @test nout(v1) == 11
        @test nout(v2) == 3

        @test size(g(ones(3, 3))) == (nout(v4), 3)

        @test Δnin!(v4, relaxed(-2), 3)

        @test nout(v1) == 6
        @test nout(v2) == 6

        @test size(g(ones(3))) == (nout(v4),)

        @test Δnin!(v4, -2, missing)

        @test nout(v1) == 5
        @test nout(v2) == 5

        @test size(g(ones(3))) == (nout(v4),)
    end

    @testset "SizeInvariant duplicate" begin
        inpt = iv(3)
        v1 = av(inpt, 7, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(inpt, 3, "v3")

        v4 = cc(v1, v2, name="v4")
        v5 = cc(v2, v3, v2, name="v5")

        v6 = nc("v6") >> v4 + v5

        g = CompGraph(inpt, v6)
        @test size(g(ones(3))) == (nout(v6),)

        @test Δnout!(v->1:nout(v), v6, -4)

        @test nout(v1) == 5
        @test nout(v2) == 2
        @test nout(v3) == 3

        @test size(g(ones(3))) == (nout(v6),)

        @test Δnout!(v6, 6)

        @test nout(v1) == 9
        @test nout(v2) == 4
        @test nout(v3) == 5

        @test size(g(ones(3))) == (nout(v6),)

        @test Δnin!(v -> 1:nout(v), v6, -2, relaxed(-2))

        @test nout(v1) == 8
        @test nout(v2) == 3
        @test nout(v3) == 5

        @test size(g(ones(3))) == (nout(v6),)

        @test Δnin!(v -> 1:nout(v), v6, 3, missing)

        @test nout(v1) == 10
        @test nout(v2) == 4
        @test nout(v3) == 6

        @test size(g(ones(3))) == (nout(v6),)
    end

    @testset "SizeStack exact infeasible" begin
        inpt = iv(3)
        v1 = av(inpt, 4, "v1")
        v2a = tv(v1, "v2a")
        v2b = tv(v1, "v2b")
        v2c = tv(v1, "v2c")
        v3 = cc(v2a, v2b, v2c; name="v3")

        g = CompGraph(inpt, v3)
        @test size(g(ones(3))) == (nout(v3),)

        @test @test_logs (:warn, r"Could not change nout of") Δnout!(v3, 2)

        @test nout(v3) == 3nout(v1) == 15

        @test size(g(ones(3))) == (nout(v3),)
    end

    @testset "SizeInvariant exact Δnout! infeasible" begin
        inpt = iv(3)
        v1 = av(inpt, 4, "v1")
        v2 = av(v1, 4, "v2")
        v3 = "v3" >> v1 + v2

        @test @test_logs (:warn, r"Could not change nout of") Δnout!(v1 => -1, v2 =>0)
        @test nout(v3) == nout(v2) == nout(v1) == 3
    end

    @testset "SizeInvariant exact Δnin! infeasible" begin
        inpt = iv(3)
        v1 = av(inpt, 4, "v1")
        v2 = av(v1, 4, "v2")
        v3 = "v3" >> v1 + v2

        @test @test_logs (:warn, r"Could not change nin of") Δnin!(v3 => (-1, 0))
        @test nout(v3) == nout(v2) == nout(v1) == 3
    end

    @testset "SizeStack one immutable" begin
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = cc(inpt, v1, name="v2")

        g = CompGraph(inpt, v2)
        @test size(g(ones(3))) == (nout(v2),)

        @test Δnout!(v1, -3) do v
            # "Tempt" optimizer to not select inputs from inpt
            v === v2 && return (-nout(inpt):nout(v1)-1)
            return 1:nout(v)
        end

        @test nin(v2) == [nout(inpt), nout(v1)] == [3, 2]
        @test nout(v2) == 5

        @test size(g(ones(3))) == (nout(v2),)
    end

    @testset "Decrease SizeStack and SizeInvariant" begin
        inpt = iv(3)
        v1 = av(inpt, 10, "v1")
        v2 = av(inpt, 5, "v2")
        v3 = av(inpt, 10, "v3")
        v4 = av(inpt, 5, "v4")

        v5 = cc(v1, v2, v3, name="v5")
        v6 = cc(v2, v1, v2, v4, name="v6")

        v7 = nc("v7") >> v5 + v6

        g = CompGraph(inpt, v7)
        @test size(g(ones(3))) == (nout(v7),)

        @test Δnout!(v7, -7)

        @test nout(v1) == 6
        @test nout(v2) == 3
        @test nout(v3) == 9
        @test nout(v4) == 6

        @test size(g(ones(3))) == (nout(v7),)
    end

    @testset "Decorating strategy" begin
        import NaiveNASlib: DecoratingJuMPΔSizeStrategy

        struct DummyDecorator{S} <: DecoratingJuMPΔSizeStrategy
            s::S
        end
        NaiveNASlib.base(s::DummyDecorator) = s.s
        NaiveNASlib.add_participants!(s::DummyDecorator, vs=AbstractVertex[]) = NaiveNASlib.add_participants!(base(s), vs)

        function graphgen()
            inpt = iv(3)
            v1 = av(inpt, 10, "v1")
            v2 = av(inpt, 5, "v2")
            v3 = av(inpt, 10, "v3")
            v4 = av(inpt, 5, "v4")

            v5 = cc(v1, v2, v3, name="v5")
            v6 = cc(v2, v1, v2, v4, name="v6")

            v7 = nc("v7") >> v5 + v6

            CompGraph(inpt, v7)
        end

        genstrat(::Type{ΔNout{Exact}}, g) = ΔNoutExact(g.outputs[1], -7)
        genstrat(::Type{ΔNout{Relaxed}}, g) = ΔNoutRelaxed(g.outputs[1], -7)
        genstrat(::typeof(ΔNinExact), g) = ΔNinExact(g.outputs[1], (-7, missing))

        @testset "$basestrat" for basestrat in (
            ΔNout{Exact},
            ΔNout{Relaxed},
            ΔNinExact,
        )
            g_exp, g_act = graphgen(), graphgen()
            s_exp = genstrat(basestrat, g_exp)
            s_act = genstrat(basestrat, g_act)

            @test Δsize!(v -> 1:nout(v), s_exp)
            @test Δsize!(v -> 1:nout(v), DummyDecorator(s_act))

            @test nout.(vertices(g_exp)) == nout.(vertices(g_act))
        end
    end

    @testset "Increase SizeStack and SizeInvariant" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(inpt, 4, "v3")
        v4 = av(inpt, 5, "v4")

        v5 = cc(v1, v2, v4, name="v5")
        v6 = cc(v3, v4, v1, name="v6")

        v7 = nc("v7") >> v5 + v6

        g = CompGraph(inpt, v7)
        @test size(g(ones(3))) == (nout(v7),)

        oldg = deepcopy(g)

        Δnout!(v7, 6)

        @test nout(v1) == 5
        @test nout(v2) == 6
        @test nout(v3) == 6
        @test nout(v4) == 7

        @test size(g(ones(3))) == (nout(v7),)

        # Since we are just increasing we have the same mapping except a few zeros
        @test filter(>(0), g(1:3)) == oldg(1:3)
    end

    @testset "SizeStack increase decrease" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 10, "v2")
        v3 = cc(v1,v2, name="v3")
        v4 = av(v3, 4, "v4")

        g = CompGraph(inpt, v3)
        @test size(g(ones(3))) == (nout(v3),)

        @test Δnout!(v -> 1:nout(v), v3 => relaxed(-5), v1 => 2, v2 => relaxed(0))

        @test nin(v3) == [nout(v1), nout(v2)] == [5, 7]
        @test nout(v3) == sum(nin(v3)) == 12

        @test lastins(v3) == [lastouts(v1), lastouts(v2)] == [[1,2,3,-1,-1], [4,5,6,7,8,9,10]]
        @test lastouts(v3) == lastins(v4) == [1,2,3,-1,-1,7,8,9,10,11,12,13]

        @test size(g(ones(3))) == (nout(v3),)
    end

    @testset "Constrained by remote subtree" begin
        inpt = iv(3)
        v0 = av(inpt, 10, "v0")

        # Main branch, the one we want to change
        v1 = cc(v0, v0, name="v1")

        # Path to subtree
        # First hop
        v2 = av(inpt, 10, "v2")
        v3 = nc("v3") >> v2 + v0

        # Second hop (v4a and v4b should not be touhed)
        v4a = av(inpt, 3, "v4a")
        v4b = av(inpt, 5, "v4b")
        v5 = cc(v4a, v3, v4b, name="v5")

        # Subtree
        v6 = av(inpt, 4, "v6")
        v7 = av(inpt, nout(v5) - nout(v6), "v7")
        v8 = cc(v6,v7,name="v8")

        # Aaaand connect it to the path
        v9 = nc("v9") >> v8 + v5
        v10 = av(v9, 5, "v10")

        g = CompGraph(inpt, [v1, v9])
        @test size.(g(ones(3))) == ((nout(v1),), (nout(v9),))

        Δnout!(v2, -1)

        @test nout(v6) == 4
        @test nout(v7) == 13
        @test nout(v0) == 9

        @test size.(g(ones(3))) == ((nout(v1),), (nout(v9),))

        Δnout!(v2, 5)

        @test nout(v6) == 4
        @test nout(v7) == 18
        @test nout(v0) == 14

        @test size.(g(ones(3))) == ((nout(v1),), (nout(v9),))
    end

    @testset "Low utility in not-connected" begin
        inpt = iv(3)
        v1 = av(inpt, 4, "v1")
        v2 = av(v1, 5, "v2")
        v3 = av(v2, 3, "v3")
        v4 = av(v2, 5, "v4")

        @test Δnout!(v1 => 2) do v
            v == v2 ? (-2:nout(v2)-3) : 1
        end
        @test nout(v1) == 6
        @test nout(v2) == 5
    end

    @testset "Increase at vertex removal" begin
        inpt = iv(3)
        v1 = av(inpt, 2, "v1")
        v2 = tv(v1, "v2")
        v3 = av(v2, 4, "v3")
        v4 = av(v3, 3, "v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(3))) == (nout(v4),)

        @test remove!(v3, RemoveStrategy())

        @test lastins(v4)== [1,2,3,4]
        @test lastouts(v1) == lastouts(v2) == [1, 2, -1, -1]

        @test size(g(ones(3))) == (nout(v4),)
    end

    @testset "Decrease at vertex removal" begin
        inpt = iv(3)
        v1 = av(inpt, 2, "v1")
        v2 = tv(v1, "v2")
        v3 = av(v2, 4, "v3")
        v4 = av(v3, 3, "v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(3))) == (nout(v4),)

        @test remove!(v3, RemoveStrategy(DecreaseBigger()))

        # What happened now is that nin(v4) got decreased from 4 to 2.
        # However, there was absolutely no need at all to select anything from v2 and before as they have not changed.

        @test lastins(v4) == lastouts(v3) == [3, 4]
        @test lastouts(v1) == lastouts(v2) == [1, 2]

        @test size(g(ones(3))) == (nout(v4),)
    end

    @testset "Increase at vertex removal SizeStack" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(v1, 5, "v3")
        v4 = cc(v2, v3, name="v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(3))) == (nout(v4),)

        @test remove!(v3, RemoveStrategy())

        # v4 did not change
        @test lastins(v4) == [[1, 2, 3, 4], [1, 2, 3, 4, 5]]
        @test lastouts(v4) == 1:9

        @test [lastouts(v2), lastouts(v1)] == [[1,2,3,4], [1,2,3,-1,-1]]

        @test size(g(ones(3))) == (nout(v4),)
    end

    @testset "Decrease at vertex removal SizeStack" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(v1, 5, "v3")
        v4 = cc(v2, v3, name="v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(Float32, 3))) == (nout(v4),)

        @test remove!(v3, RemoveStrategy(DecreaseBigger(mapstrat=WithUtilityFun(v -> 1:nout(v)))))

        @test lastins(v4) == [[1, 2, 3, 4], [3, 4, 5]]
        @test lastouts(v4) == [1, 2, 3, 4, 7, 8, 9]

        # v2 and v1 did not change
        @test [lastouts(v2), lastouts(v1)] == [[1,2,3,4], [1,2,3]]

        @test size(g(ones(3))) == (nout(v4),)
    end

    @testset "TruncateInIndsToValid" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = tv(v1, "v2")
        v3 = av(v2, 2, "v3")

        g = CompGraph(inpt, v2)
        @test size(g(ones(3))) == (nout(v2),)

        vnew = av(inpt, 7, "vnew")
        create_edge!(vnew, v2;strategy=PostAlign(TruncateInIndsToValid(AlignNinToNout())))
        @test nin(v2) == [nout(v1), nout(vnew)] == [7, 7]

        @test lastouts(vnew) == 1:7
        @test lastins(v2) == [[1,2,3,-1,-1,-1,-1], [1,2,3,-1,-1,-1,-1]]
        @test lastouts(v2) == [1,2,3,-1,-1,-1,-1]
    end

    @testset "WithUtilityFun" begin
        inpt = iv(3)
        v1 = av(inpt, 6, "v1")

        @test Δsize!(WithUtilityFun(v -> v === v1 ? [-1, 1, -1, 1, -1, 1] : 1:nout(v), DefaultJuMPΔSizeStrategy()), v1)

        @test nout(v1) == 3
        @test lastouts(v1) == [2,4,6]
    end

    @testset "CompConstraint" begin
        struct CompConstraint end
        function NaiveNASlib.compconstraint!(::NaiveNASlib.NeuronIndices, ::AbstractJuMPΔSizeStrategy, ::CompConstraint, data)
            var = data.outselectvars[data.vertex];
            JuMP.@constraint(data.model, var[[1, 3]] .== 0)
        end
        NaiveNASlib.nout(::CompConstraint) = 4
        NaiveNASlib.nin(::CompConstraint) = 3
        ccouts = nothing
        NaiveNASlib.Δsize!(::CompConstraint, ins::AbstractVector, outs::AbstractVector) = ccouts=outs

        inpt = iv(3)
        v1 = absorbvertex(CompConstraint(), inpt, traitdecoration = tf("v1"))
        v2 = av(v1, 5, "v2")

        Δnout!(v1, 3)
        @test ccouts == lastins(v2) == [2, 4, -1, -1, -1, -1, -1]
    end

    @testset "Time limit" begin
        import NaiveNASlib: TimeLimitΔSizeStrategy, selectmodel, NeuronIndices
        @test JuMP.time_limit_sec(selectmodel(NeuronIndices(), TimeLimitΔSizeStrategy(123), AbstractVertex[], () -> 1)) == 123
    end

    @testset "Timeout action" begin
        import NaiveNASlib: TimeLimitΔSizeStrategy, TimeOutAction, ΔSizeFailNoOp

        v0 = av(iv(1), 1000, "v0")
        v1 = cc(v0, v0, v0; name="v1")
        v2 = cc(v1, v0, v1; name="v2")
        v3 = "v3" >> v2 + v2

        @test @test_logs (:warn, "Solver timed out") Δsize!(TimeLimitΔSizeStrategy(0.001, TimeOutAction(); fallback=ΔSizeFailNoOp()), v3) == false

        modelres = nothing
        action = m -> (modelres = m; return false)

        @test Δsize!(TimeLimitΔSizeStrategy(0.001, TimeOutAction(;action); fallback=ΔSizeFailNoOp()), v3) == false

        @test modelres isa JuMP.Model
    end

    @testset "All negative utility $label" for (label, utilfun) in (
        ("scalar", v -> -1.0),
        ("array", v -> fill(-1.0, nout(v))),
    )
        inpt = iv(3)
        v1 = av(inpt, 4, "v1")
        v2 = av(v1, 5, "v2")

        @test Δsize!(utilfun, CompGraph(inpt, v2)) == true

        @test [nout(v1)] == nin(v2) == [nout(v2)] == [1]
    end

    @testset "WithKwargs" begin
        mutable struct WantKwargs
            insize::Int
            outsize::Int
        end
        NaiveNASlib.nout(w::WantKwargs) = w.outsize
        NaiveNASlib.nin(w::WantKwargs) = [w.insize]

        seenkwargs = Dict()
        function NaiveNASlib.Δsize!(w::WantKwargs, ins::AbstractVector, outs::AbstractVector; kwargs...)
            if ins[] !== missing
                w.insize = length(ins[])
            end
            w.outsize = length(outs)
            merge!(seenkwargs, kwargs)
        end

        inpt = iv(3)
        v1 = absorbvertex(WantKwargs(3, 5), inpt; traitdecoration=tf("v1"))

        @test Δsize!(WithKwargs(ΔNout(v1 => 2); a=3, b=4)) == true

        @test nout(v1) == 7
        @test seenkwargs == Dict(:a => 3, :b => 4)

        @test Δsize!(validateafterΔsize(base=WithKwargs(ΔNout(v1 => relaxed(-2)); c = 5))) == true
        @test nout(v1) == 5
        @test seenkwargs == Dict(:a => 3, :b => 4, :c=>5)
    end

    @testset "UpperInsertBound" begin
        import Logging
        inpt = iv(3)
        v1 = av(inpt, 4, "v1")
        v2 = av(v1, 5, "v2")

        strategy = ΔNoutExact(v1=>2)
        fallback = LogΔSizeExec("Success!?", Logging.Warn)
        @test @test_logs (:warn, "Success!?") Δsize!(UpperInsertBound(0;strategy, fallback)) == false

        @test @test_logs (:warn, "Success!?") Δsize!(TimeOutAction(;base=UpperInsertBound(0;strategy), fallback)) == false

        @test Δsize!(UpperInsertBound([v2], 0;strategy, fallback)) == true

    end
end


using Knet
include("../helper.jl")
include("mpii.annotation.parser.jl")
include("variant.data.jl")

function get_mpii_single_person_batches_and_data_items(;
    batch_size = 1,
    use_global_scaling = false,
    train_size = 18432,
    test_size = 753,
    should_shuffle_dataset = false,
    fetch_only_test_batch = false,
)

    # TODO: throw error if train + test size is bigger than 19185
    dataset = read_cropped_mpii_annotations(; should_shuffle = false)
    dtrn = []
    step_size = 1024
    data_items = []

    single_person_dataset_length = train_size
    single_person_dataset_test_length = test_size
    total_single_person_length = train_size + test_size

    for (k, v) in dataset
        dataset[k] = v[1:total_single_person_length]
    end

    if should_shuffle_dataset
        shuffled_idxs = shuffle!(collect(1:total_single_person_length))
        for (k, v) in dataset
            dataset[k] = v[shuffled_idxs]
        end
    end

    if !fetch_only_test_batch
        if !use_global_scaling

            for i = 1:step_size:single_person_dataset_length
                println(i)
                train_dataset =
                    get_from_dataset(dataset, i, i + step_size - 1; should_use_pmap = true)
                if isempty(train_dataset)
                    println("continued")
                    continue
                end
                train_preprocessed = preprocess_dataset(
                    train_dataset;
                    use_global_scaling = use_global_scaling,
                )

                dtrn_part = get_batch(batch_size, train_preprocessed)
                if isempty(dtrn)
                    dtrn = dtrn_part
                    println("$(summary(dtrn))")
                else
                    append_to_data!(dtrn, dtrn_part)
                    println("$(summary(dtrn))")
                end

                push!(data_items, map(e -> e[1], train_preprocessed)...)

                train_dataset = 0
                train_preprocessed = 0
                GC.gc(true)
            end

        else
            train_dataset = get_from_dataset(
                dataset,
                1,
                single_person_dataset_length;
                should_use_pmap = true,
            )
            train_preprocessed =
                preprocess_dataset(train_dataset; use_global_scaling = use_global_scaling)
            dtrn = get_variant_batch(train_preprocessed)

            push!(data_items, map(e -> e[1], train_preprocessed)...)

            train_dataset = 0
            train_preprocessed = 0
            GC.gc(true)
        end
    end

    # Preparing Test Data
    test_dataset = get_from_dataset(
        dataset,
        single_person_dataset_length + 1,
        total_single_person_length;
        should_use_pmap = true,
    )

    if use_global_scaling
        test_preprocessed =
            preprocess_dataset(test_dataset; use_global_scaling = use_global_scaling)
        dtst = get_variant_batch(test_preprocessed; shuffle_in_minibatch = false)
    else
        test_preprocessed = preprocess_dataset(test_dataset)
        dtst = get_batch(batch_size, test_preprocessed; shuffle_in_minibatch = false)
    end

    push!(data_items, map(e -> e[1], test_preprocessed)...)

    test_dataset = 0
    test_preprocessed = 0
    GC.gc(true)

    return (dtrn, dtst, data_items)
end

function get_mpii_batches_and_data_items(
    batch_size;
    should_shuffle = false,
    should_use_pmap = true,
)
    dataset = read_cropped_mpii_annotations(; should_shuffle = should_shuffle)
    dtrn = []
    step_size = 1024

    for i = 1:step_size:(train_image_count+1-step_size)
        println(i)
        train_dataset = get_from_dataset(
            dataset,
            i,
            i + step_size - 1;
            should_use_pmap = should_use_pmap,
        )
        if isempty(train_dataset)
            println("continued")
            continue
        end
        train_preprocessed = preprocess_dataset(train_dataset)
        dtrn_part = get_batch(batch_size, train_preprocessed)
        if isempty(dtrn)
            dtrn = dtrn_part
            println("$(summary(dtrn))")
        else
            append_to_data!(dtrn, dtrn_part)
            println("$(summary(dtrn))")
        end
        train_dataset = 0
        train_preprocessed = 0
        GC.gc(true)
    end

    # TODO: Preparing Validation Data
    validation_dataset = get_from_dataset(
        dataset,
        train_image_count + 1,
        train_image_count + validation_image_count;
        should_use_pmap = should_use_pmap,
    )
    validation_preprocessed = preprocess_dataset(validation_dataset)
    dval = get_batch(batch_size, validation_preprocessed; shuffle_in_minibatch = false)

    validation_dataset = 0
    validation_preprocessed = 0
    GC.gc(true)

    # Preparing Test Data
    test_dataset = get_from_dataset(
        dataset,
        train_image_count + validation_image_count + 1,
        train_image_count + validation_image_count + test_image_count;
        should_use_pmap = should_use_pmap,
    )
    test_preprocessed = preprocess_dataset(test_dataset)
    dtst = get_batch(batch_size, test_preprocessed; shuffle_in_minibatch = false)

    test_dataset = 0
    test_preprocessed = 0
    GC.gc(true)

    data_items = get_from_dataset(
        dataset,
        1,
        train_image_count + validation_image_count + test_image_count;
        should_use_pmap = should_use_pmap,
    )
    data_items = map(d_i -> DataItem(d_i, read_image_h, read_image_w), data_items)

    return (dtrn, dval, dtst, data_items)
end

function get_batch(
    batch_size,
    preprocessed;
    shuffle_in_minibatch = true,
    include_scmap_weights = false,
)
    y_scmap = map(element -> element[3], preprocessed)
    y_scmap_weights = map(element -> add_dim(element[4]), preprocessed)
    y_locref_map = map(element -> element[5], preprocessed)
    y_locref_mask = map(element -> element[6], preprocessed)

    y_all = map(
        i -> merge_different_y(
            y_scmap[i],
            y_scmap_weights[i],
            y_locref_map[i],
            y_locref_mask[i];
            include_scmap_weights = include_scmap_weights,
        ),
        1:length(y_scmap),
    )
    x_all = map(element -> element[2], preprocessed)

    xall = cat(x_all..., dims = 4)
    yall = cat(y_all..., dims = 4)

    xsize = size(xall)[1:3]
    ysize = size(yall)[1:3]

    d = minibatch(
        xall,
        yall,
        batch_size;
        xsize = (xsize..., :),
        ysize = (ysize..., :),
        xtype = Knet.atype(),
        shuffle = shuffle_in_minibatch,
    )
    return d
end

function get_variant_batch(
    preprocessed;
    shuffle_in_minibatch = true,
    include_scmap_weights = false,
)
    total_element_count = length(preprocessed)
    x_all = map(element -> element[2], preprocessed)
    y_scmap = map(element -> element[3], preprocessed)
    y_scmap_weights = map(element -> add_dim(element[4]), preprocessed)
    y_locref_map = map(element -> element[5], preprocessed)
    y_locref_mask = map(element -> element[6], preprocessed)
    y_all = map(
        i -> merge_different_y(
            y_scmap[i],
            y_scmap_weights[i],
            y_locref_map[i],
            y_locref_mask[i];
            include_scmap_weights = include_scmap_weights,
        ),
        1:total_element_count,
    )
    return variant_minibatch(
        x_all,
        y_all;
        shuffle = shuffle_in_minibatch,
        xtype = Knet.atype(),
        ytype = Knet.atype(),
        add_random_mirroring = global_add_random_mirroring,
    )
end

function merge_different_y(
    scmap,
    scmap_w,
    locref_map,
    locref_mask;
    include_scmap_weights = false,
)
    if include_scmap_weights == true
        if use_locref_mask_weights
            return add_dim(cat(scmap, scmap_w, locref_map, locref_mask; dims = (3)))
        else
            return add_dim(cat(scmap, scmap_w, locref_map; dims = (3)))
        end
    else
        if use_locref_mask_weights
            return add_dim(cat(scmap, locref_map, locref_mask; dims = (3)))
        else
            return add_dim(cat(scmap, locref_map; dims = (3)))
        end
    end
end


using KernelAbstractions: NoneEvent
using OffsetArrays: OffsetArray

using Oceananigans.Operators: Δzᵃᵃᶜ
using Oceananigans.BoundaryConditions: left_gradient, right_gradient, linearly_extrapolate, FBC, VBC, GBC
using Oceananigans.BoundaryConditions: fill_bottom_halo!, fill_top_halo!, apply_z_bottom_bc!, apply_z_top_bc!
using Oceananigans.Grids: Flat, Bounded
using Oceananigans.Architectures: device_event

import Oceananigans.Utils: launch!
import Oceananigans.Grids: validate_size, validate_halo
import Oceananigans.BoundaryConditions: fill_halo_regions!
import Oceananigans.TurbulenceClosures: time_discretization
import Oceananigans.TurbulenceClosures: calculate_diffusivities!

#####
##### Implements a "single column model mode" for HydrostaticFreeSurfaceModel
#####

const SingleColumnGrid = AbstractGrid{<:AbstractFloat, <:Flat, <:Flat, <:Bounded}

#####
##### Model constructor utils
#####

PressureField(arch, ::SingleColumnGrid) = (pHY′ = nothing,)
FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, velocities, arch, ::SingleColumnGrid) = nothing
FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, arch, ::SingleColumnGrid) = nothing

validate_momentum_advection(momentum_advection, ::SingleColumnGrid) = nothing
validate_tracer_advection(tracer_advection::AbstractAdvectionScheme, ::SingleColumnGrid) = nothing, NamedTuple()
validate_tracer_advection(tracer_advection::Nothing, ::SingleColumnGrid) = nothing, NamedTuple()

#####
##### Time-step optimizations
#####

calculate_free_surface_tendency!(arch, ::SingleColumnGrid, args...) = NoneEvent()

# Fast state update and halo filling

function update_state!(model::HydrostaticFreeSurfaceModel, grid::SingleColumnGrid)

    fill_halo_regions!(prognostic_fields(model), model.architecture, model.clock, fields(model))

    compute_auxiliary_fields!(model.auxiliary_fields)

    # Calculate diffusivities
    calculate_diffusivities!(model.diffusivity_fields, model.closure, model)

    fill_halo_regions!(model.diffusivity_fields,
                       model.architecture,
                       model.clock,
                       fields(model))

    return nothing
end

import Oceananigans.TurbulenceClosures: ∂ⱼ_τ₁ⱼ, ∂ⱼ_τ₂ⱼ, ∂ⱼ_τ₃ⱼ, ∇_dot_qᶜ

using Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure

const ClosureArray = AbstractArray{<:AbstractTurbulenceClosure}

@inline function ∂ⱼ_τ₁ⱼ(i, j, k, grid::SingleColumnGrid, closure_array::ClosureArray, args...)
    @inbounds closure = closure_array[i, j]
    return ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, args...)
end

@inline function ∂ⱼ_τ₂ⱼ(i, j, k, grid::SingleColumnGrid, closure_array::ClosureArray, args...)
    @inbounds closure = closure_array[i, j]
    return ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, args...)
end

@inline function ∇_dot_qᶜ(i, j, k, grid::SingleColumnGrid, closure_array::ClosureArray, c, tracer_index, args...)
    @inbounds closure = closure_array[i, j]
    return ∇_dot_qᶜ(i, j, k, grid, closure, c, tracer_index, args...)
end

struct ColumnEnsembleSize{C<:Tuple{Int, Int}}
    ensemble :: C
    Nz :: Int
    Hz :: Int
end

ColumnEnsembleSize(; Nz, ensemble=(0, 0), Hz=1) = ColumnEnsembleSize(ensemble, Nz, Hz)

validate_size(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(e.ensemble[1], e.ensemble[2], e.Nz)
validate_halo(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(0, 0, e.Hz)

@inline time_discretization(::AbstractArray{<:AbstractTurbulenceClosure{TD}}) where TD = TD()

#####
##### TKEBasedVerticalDiffusivity helpers
#####

using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: CATKEVD, Kuᶜᶜᶜ, Kcᶜᶜᶜ, Keᶜᶜᶜ, _top_tke_flux, CATKEVDArray

import Oceananigans.TurbulenceClosures: calculate_diffusivities!
import Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: top_tke_flux

function calculate_diffusivities!(diffusivities, closure::CATKEVDArray, model)

    arch = model.architecture
    grid = model.grid
    velocities = model.velocities
    tracers = model.tracers
    buoyancy = model.buoyancy
    clock = model.clock
    top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers))

    event = launch!(arch, grid, :xyz,
                    calculate_CATKEArray_diffusivities!,
                    diffusivities, grid, closure, velocities, tracers, buoyancy, clock, top_tracer_bcs,
                    dependencies = device_event(arch))

    wait(device(arch), event)

    return nothing
end

@kernel function calculate_CATKEArray_diffusivities!(diffusivities, grid::SingleColumnGrid, closure_array::CATKEVDArray, args...)
    i, j, k, = @index(Global, NTuple)
    @inbounds begin
        closure = closure_array[i, j]
        diffusivities.Kᵘ[i, j, k] = Kuᶜᶜᶜ(i, j, k, grid, closure, args...)
        diffusivities.Kᶜ[i, j, k] = Kcᶜᶜᶜ(i, j, k, grid, closure, args...)
        diffusivities.Kᵉ[i, j, k] = Keᶜᶜᶜ(i, j, k, grid, closure, args...)
    end
end

@inline tracer_tendency_kernel_function(model::HydrostaticFreeSurfaceModel, closure::CATKEVDArray, ::Val{:e}) =
    hydrostatic_turbulent_kinetic_energy_tendency

""" Compute the flux of TKE through the surface / top boundary. """
@inline function top_tke_flux(i, j, grid::SingleColumnGrid, clock, fields, parameters, closure_array::CATKEVDArray, buoyancy)
    top_tracer_bcs = parameters.top_tracer_boundary_conditions
    top_velocity_bcs = parameters.top_velocity_boundary_conditions
    @inbounds closure = closure_array[i, j]

    return _top_tke_flux(i, j, grid, closure.surface_TKE_flux, closure,
                         buoyancy, fields, top_tracer_bcs, top_velocity_bcs, clock)
end

@inline function hydrostatic_turbulent_kinetic_energy_tendency(i, j, k, grid::SingleColumnGrid,
                                                               val_tracer_index::Val{tracer_index},
                                                               advection,
                                                               closure_array::CATKEVDArray, args...) where tracer_index

    @inbounds closure = closure_array[i, j]
    return hydrostatic_turbulent_kinetic_energy_tendency(i, j, k, grid, val_tracer_index, advection, closure, args...)
end

#####
##### Arrays of Coriolises
#####

import Oceananigans.Coriolis: x_f_cross_U, y_f_cross_U, z_f_cross_U

const FPlaneArray = AbstractArray{<:FPlane, 2}

@inline x_f_cross_U(i, j, k, grid::SingleColumnGrid, coriolis::FPlaneArray, U) = @inbounds - coriolis[i, j].f * ℑxyᶠᶜᵃ(i, j, k, grid, U[2])
@inline y_f_cross_U(i, j, k, grid::SingleColumnGrid, coriolis::FPlaneArray, U) = @inbounds   coriolis[i, j].f * ℑxyᶜᶠᵃ(i, j, k, grid, U[1])
@inline z_f_cross_U(i, j, k, grid::SingleColumnGrid, coriolis::FPlaneArray, U) = zero(eltype(grid))


################################################################################
#
# Routines for synthetic data generation
#
################################################################################

export gendata

### Wrapper for JUDI data types

function gendata(
	m::Model, wav::judiVector, rcv_geom::Geometry;
	opt = Options(),
	dt_comp::Union{Nothing, Array{Any, 1}} = nothing,
	save_wavefield::Bool = false)

	return gendata_rawinput(m, wav.geometry, rcv_geom, wav.data, 1:wav.nsrc; opt = opt, dt_comp = dt_comp, save_wavefield = save_wavefield)

end

### Objective wrapper for array input

function gendata(
	m::Array{R, 2}, n::NTuple{2, Int64}, d::NTuple{2, R}, o::NTuple{2, R},
	wav::judiVector, rcv_geom::Geometry;
	opt = Options(),
	dt_comp::Union{Nothing, Array{Any, 1}} = nothing,
	save_wavefield::Bool = false)

	return gendata(Model(n, d, o, m), wav, rcv_geom; opt = opt, dt_comp = dt_comp, save_wavefield = save_wavefield)

end

### Serial/parallel implementation

## Serial implementation

function gendata_rawinput(
	model_full::Model,
	src_geom::Geometry, rcv_geom::Geometry,
	src_data::Array{R, 2},;
	opt = Options(),
	dt_comp::Union{Nothing, R} = nothing,
	save_wavefield::Bool = false)

	# Setting pre-defined absorbing layer size
	model_full.nb = NB

	# Load full geometry for out-of-core geometry containers
	typeof(rcv_geom) == GeometryOOC && (rcv_geom = Geometry(rcv_geom))
	typeof(src_geom) == GeometryOOC && (src_geom = Geometry(src_geom))
	length(model_full.n) == 3 ? dims = [3, 2, 1] : dims = [2, 1] # model dimensions for Python are (z,y,x) and (z,x)

	# Limit model to area with sources/receivers
	if opt.limit_m == true
		model = deepcopy(model_full)
		model = limit_model_to_receiver_area(src_geom, rcv_geom, model, opt.buffer_size)
	else
		model = model_full
	end

	# Set up Python model structure
	model_py = devito_model(model, "F", 1, opt, 0)

	# Remove receivers outside the modeling domain (otherwise leads to segmentation faults)
	rcv_geom = remove_out_of_bounds_receivers(rcv_geom, model)

	# Call to objective with Julia/Devito interface function
	return gendata_jldevito(model_py, model.o, src_geom, rcv_geom, src_data; opt = opt, dt_comp = dt_comp, save_wavefield = save_wavefield)

end

## Parallel implementation

function gendata_rawinput(
	model::Model,
	src_geom::Geometry, rcv_geom::Geometry,
	src_data,
	src_idx::UnitRange{Int64};
	opt = Options(),
	dt_comp::Union{Nothing, Array{Any, 1}} = nothing,
	save_wavefield::Bool = false)

	# Setup parallelization
	p = default_worker_pool()
	time_modeling_par = remote(gendata_rawinput)
	time_modeling = retry(time_modeling_par)

	# Initialize output
	nsrc = length(src_idx)
	results = Array{Any}(undef, nsrc)

	# Process shots from source channel asynchronously
	@sync begin
		for j = 1:nsrc

			# Local geometry for current position
			opt_loc = subsample(opt, j)
			src_geom_loc = subsample(src_geom, j)
			rcv_geom_loc = subsample(rcv_geom, j)

			# Selecting variables for current shot index
			src_data_loc = src_data[j]
			if dt_comp == nothing
				dt_comp_loc = nothing
			else
				dt_comp_loc = dt_comp[j]
			end

			# Local result
			results[j] = @spawn gendata_rawinput(model, src_geom_loc, rcv_geom_loc, src_data_loc, opt = opt_loc; dt_comp = dt_comp_loc, save_wavefield = save_wavefield)

		end
	end

	# Aggregating results
	argout = Array{Array, 1}(undef, nsrc)
	for j = 1:nsrc
		argout[j] = fetch(results[j])
	end
	if save_wavefield
		return argout
	else
		return judiVector(rcv_geom, argout)
	end

end

### Julia/devito implementation

function gendata_jldevito(
	model_py::PyCall.PyObject,
	origin,
	src_geom::Geometry, rcv_geom::Geometry,
	src_data::Array{R, 2};
	opt = Options(),
	dt_comp::Union{Nothing, R} = nothing,
	save_wavefield::Bool = false)

	# Loading python modules for devito implementation of objTWRIdual
	devitomod = load_codegen_TWRIdual()

	# Interpolate input data to computational grid
	if dt_comp == nothing
		dt_comp = R(model_py.critical_dt)
	end
	q_in = time_resample(src_data, src_geom, dt_comp)[1]
	nt_comp = size(q_in, 1)
	nt_rcv = Int(trunc(rcv_geom.t[1]/dt_comp+1))

	# Set up coordinates with devito dimensions
	src_coords = setup_grid(src_geom, model_py.shape, origin)
	rcv_coords = setup_grid(rcv_geom, model_py.shape, origin)

	# Computing output
	if !save_wavefield
		argout = pycall(devitomod.forward,
						Tuple{Array{R, 2}, Nothing},
						model_py,
						PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(rcv_coords))),
						PyReverseDims(copy(transpose(q_in))),
						dt_comp, SPACE_ORDER, save_wavefield)
		return argout[1]
	else
		argout = pycall(devitomod.forward,
						Tuple{Array{R, 2}, PyObject},
						model_py,
						PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(rcv_coords))),
						PyReverseDims(copy(transpose(q_in))),
						dt_comp, SPACE_ORDER, save_wavefield)
		return argout[2].data[:, model_py.nbpml+1:end-model_py.nbpml, model_py.nbpml+1:end-model_py.nbpml]
	end

end


using StatsBase

"""
    function mutation!(pop::populations)

Mutate a base pair in a random species, chosen depending on subpopulation size.
If the species is extinct after mutation, it is removed. If mutation creates a new species,
it is added to the population, otherwise the subpopulation size of species the individual mutates to
is simply increased by one.
"""
function mutation!(pop::populations)
    # Choose random sequence depending on subpopulation size
    mutation_sequence = sample(collect(1:length(pop.seqs)), Weights(pop.freqs ./ pop.N))
    # Add sequence to mutate
    push!(pop.seqs, copy(pop.seqs[mutation_sequence]))
    push!(pop.freqs, 1)
    pop.freqs[mutation_sequence] -= 1
    # Choose random base to mutate
    mutation_base = rand(1:pop.l)
    # Exclude self mutations
    poss_bases = filter(x->x != pop.seqs[end][mutation_base], collect(1:pop.n))
    # Choose random new base
    pop.seqs[end][mutation_base] = rand(poss_bases)
    copied = findfirst(x->~any(.~(x.==pop.seqs[end])), pop.seqs[1:end-1])
    # Look if mutated sequence already exists
    if ~isnothing(copied)
        pop.freqs[copied] += 1
        deleteat!(pop.freqs, length(pop.seqs))
        deleteat!(pop.seqs, length(pop.seqs))
    end
    # Remove extinct species
    remove_empty!(pop.freqs, pop.seqs)
end

"""
    function mutation!(pop::driver_trailer_l)

Mutate a base pair in a random species, chosen depending on subpopulation size.
If the species is extinct after mutation, it is removed. If mutation creates a new species,
it is added to the population, otherwise the subpopulation size of species the individual mutates to
is simply increased by one.
"""
function mutation!(pop::driver_trailer_l)
    # Choose random sequence depending on subpopulation size
    mutation_sequence = sample(collect(1:length(pop.seqs)), Weights(pop.freqs ./ pop.N))
    # Add sequence to mutate
    push!(pop.seqs, deepcopy(pop.seqs[mutation_sequence]))
    push!(pop.freqs, 1)
    push!(pop.l, pop.l[mutation_sequence])
    pop.freqs[mutation_sequence] -= 1
    # Choose random base to mutate
    mutation_base = rand(1:pop.L)
    # Exclude self mutations
    poss_bases = filter(x->x != pop.seqs[end][mutation_base], collect(1:pop.n))
    # Choose random new base
    pop.seqs[end][mutation_base] = rand(poss_bases)
    copied = findfirst(x->~any(.~(x.==pop.seqs[end])), pop.seqs[1:end-1])
    # Look if mutated sequence already exists
    if ~isnothing(copied)
        pop.freqs[copied] += 1
        deleteat!(pop.freqs, length(pop.seqs))
        deleteat!(pop.seqs, length(pop.seqs))
        deleteat!(pop.l, length(pop.l))
    end
    # Remove extinct species
    remove_empty!(pop.freqs, pop.seqs, pop.l)
end

"""
    function driver_mutation!(pop::DT_population)

Mutate a random base in the driver sequence.
"""
function driver_mutation!(pop::DT_population)
    # Choose random base to mutate
    mutation_base = rand(1:length(pop.driver))
    # Exclude self mutations
    poss_bases = filter(x->x != pop.driver[mutation_base], collect(1:pop.m))
    # Choose random new base
    pop.driver[mutation_base] = rand(poss_bases)
end

"""
    function driver_mutation!(pop::driver_trailer)

Mutate the length of species and create a new species
"""
function length_mutation!(pop::driver_trailer_l)
    # Choose random sequence depending on subpopulation size
    mutation_sequence = sample( collect(1:length(pop.seqs)),
                                Weights(pop.freqs ./ pop.N))
    # Add sequence to mutate

    if rand() < 0.5 && pop.l[mutation_sequence] < pop.L
        var = 1
    elseif pop.l[mutation_sequence] > 2
        var = -1
    else
        return
    end

    push!(pop.seqs, copy(pop.seqs[mutation_sequence]))
    push!(pop.freqs, 1)
    push!(pop.l, pop.l[mutation_sequence])

    pop.freqs[mutation_sequence] -= 1
    pop.l[end] += var
    # Remove extinct species
    remove_empty!(pop.freqs, pop.seqs, pop.l)
end

"""
    bp_substitution!(pop::mono_pop, emat::Array{T, 2}, fitness_function::fitness_functions) where {T<:Real}

Attempt a base pair substitution. The population has to consist of one species
only. Mutation is accepted with probability given by the Kimura fixation
probability.
"""
function bp_substitution!(pop::mono_pop, emat::Array{T, 2}, fitness_function::fitness_functions) where {T<:Real}

    # Choose random base to mutate
    mutation_base = rand(eachindex(pop.seqs))
    # Exclude self mutations
    possible_bases = filter(x->x != pop.seqs[mutation_base], collect(1:pop.n))
    # Choose random new base
    temp_pop = deepcopy(pop)
    temp_pop.seqs[mutation_base] = rand(possible_bases)
    # Compute energies
    E = get_energy(pop, emat)
    E_mutant = get_energy(temp_pop, emat)
    # Compute fitness and selection coefficients
    F = fitness(E, length(pop.seqs), fitness_function)
    F_mutant = fitness(E_mutant, length(pop.seqs), fitness_function)
    s = F_mutant - F
    # Accept mutatant depending on Kimura probability
    if rand() < kimura_prob(s, pop.N)
        pop.seqs = temp_pop.seqs
    end
end

"""
    bp_substitution!(pop::populations, emat::Array{T, 2}, fitness_function::fitness_functions) where {T<:Real}

Attempt a base pair substitution. The population has to consist of one species
only. Mutation is accepted with probability given by the Kimura fixation
probability.
"""
function bp_substitution!(pop::populations, emat::Array{T, 2}, fitness_function::fitness_functions) where {T<:Real}
    if length(pop.seqs) > 1
        throw(ArgumentError("Input population has more than one species."))
    end
    # Choose random base to mutate
    mutation_base = rand(1:pop.L)
    # Exclude self mutations
    poss_bases = filter(x->x != pop.seqs[1][mutation_base], collect(1:pop.n))
    # Choose random new base
    temp_pop = deepcopy(pop)
    temp_pop.seqs[1][mutation_base] = rand(poss_bases)
    # Compute energies
    E = get_energy(pop, emat)
    E_mutant = get_energy(temp_pop, emat)
    # Compute fitness and selection coefficients
    F = fitness(E[1], pop.l[1], fitness_function)
    F_mutant = fitness(E_mutant[1], pop.l[1], fitness_function)
    s = F_mutant - F
    # Accept mutatant depending on Kimura probability
    if rand() < kimura_prob(s, pop.N)
        pop.seqs = temp_pop.seqs
    end
end

"""
    l_substitution!(
        pop::driver_trailer_l,
        emat::Array{T, 2},
        fitness_function::fitness_functions) where {T<:Real}

Attempt a substitution of a length mutation. Population has to be consisting
of one species only.
"""
function l_substitution!(
    pop::driver_trailer_l,
    emat::Array{T, 2},
    fitness_function::fitness_functions) where {T<:Real}

    if length(pop.seqs) > 1
        throw(ArgumentError("Input population has more than one species."))
    end

    if rand() < 0.5 && pop.l[1] < pop.L
        var = 1
    elseif pop.l[1] > 2
        var = -1
    else
        return
    end
    # Choose random new base
    temp_pop = deepcopy(pop)
    temp_pop.l[1] += var
    # Compute energies
    E = get_energy(pop, emat)
    E_mutant = get_energy(temp_pop, emat)
    # Compute fitness and selection coefficients
    F = fitness(E[1], pop.l[1], fitness_function)
    F_mutant = fitness(E_mutant[1], temp_pop.l[1], fitness_function)
    s = F_mutant - F
    # Accept mutatant depending on Kimura probability
    if rand() < kimura_prob(s, pop.N)
        pop.l = temp_pop.l
    end
end

"""
    l_substitution!(
        pop::mono_pop,
        emat::Array{T, 2},
        fitness_function::fitness_functions) where {T<:Real}

Attempt a substitution of a length mutation. Population has to be consisting
of one species only.
"""
function l_substitution!(
    pop::mono_pop,
    emat::Array{T, 2},
    fitness_function::fitness_functions) where {T<:Real}

    temp_pop = deepcopy(pop)
    if rand() < 0.5
        push!(temp_pop.seqs, rand(1:pop.n))
        push!(temp_pop.driver, rand(1:pop.m))
    elseif length(temp_pop.seqs) > 1
        pop!(temp_pop.seqs)
        pop!(temp_pop.driver)
    end
    # Compute energies
    E = get_energy(pop, emat)
    E_mutant = get_energy(temp_pop, emat)
    # Compute fitness and selection coefficients
    F = fitness(E, length(pop.seqs), fitness_function)
    F_mutant = fitness(E_mutant, length(temp_pop.seqs), fitness_function)
    s = F_mutant - F
    # Accept mutatant depending on Kimura probability
    if rand() < kimura_prob(s, pop.N)
        pop.seqs = temp_pop.seqs
        pop.driver = temp_pop.driver
    end
end

"""
    kimura_prob(s, N)

Probability of fixation of a mutation with selection coefficient
`s` in a monomorphic population of size `N`.
"""
function kimura_prob(s::Real, N::Int)
    if abs(s) <= 10^-8
        return 1 / N
    elseif s < -10^3
        return 0
    else
        return (1 - exp(-2s)) / (1 - exp(-2N * s))
    end
end


using LightGraphs
using LightGraphsIO

@testset "Types" begin
    dg = complete_digraph(5)
    g = complete_graph(5)

    # DOT
    @test typeof(DOTFormat()) <: AbstractGraphFormat
    @test load_graph(1, DOTFormat) == 1
    @test_broken write_graph(g, DOTFormat) == 1

    # EdgeList
    @test typeof(EdgeListFormat()) <: AbstractGraphFormat
    @test load_graph(1, EdgeListFormat) == 2
    @test write_graph(g, EdgeListFormat) == 2

    # GML
    @test typeof(GMLFormat()) <: AbstractGraphFormat
    @test load_graph(1, GMLFormat) == 3
    @test write_graph(g, GMLFormat) == 3

    # NET
    @test typeof(NETFormat()) <: AbstractGraphFormat
    @test load_graph(1, NETFormat) == 4
    @test write_graph(g, NETFormat) == 4

    # GEXF
    @test typeof(GEXFFormat()) <: AbstractGraphFormat
    @test load_graph(1, GEXFFormat) == 5
    @test write_graph(g, GEXFFormat) == 5

    # Graph6
    @test typeof(Graph6Format()) <: AbstractGraphFormat
    @test load_graph(1, Graph6Format) == 6
    @test write_graph(g, Graph6Format) == 6

    # GraphML
    @test typeof(GraphMLFormat()) <: AbstractGraphFormat
    @test load_graph(1, GraphMLFormat) == 7
    @test write_graph(g, GraphMLFormat) == 7

    # CDF
    @test typeof(CDFFormat()) <: AbstractGraphFormat
    @test load_graph(1, CDFFormat) == 8
    @test_broken write_graph(g, CDFFormat) == 8

    # LGCompressed
    @test typeof(LGCompressedFormat()) <: AbstractGraphFormat
    @test load_graph(1, LGCompressedFormat) == 9
    @test write_graph(g, LGCompressedFormat) == 9
end


"""
    symbol, idxnext = tryparseelement(s, idx, elements)

Extract the element symbol starting at position `idx` in `s`. It must be one of the items in `elements`.

`symbol` is the element symbol (a string), or `nothing` if no valid element can be found.
`idxnext` is the first position in `s` after the parsed element.
"""
function tryparseelement(s::AbstractString, idx::Int, elements)
    len = length(s)
    idx > len && return nothing, idx
    c = s[idx]
    isletter(c) || return nothing, idx
    idxnext = nextind(s, idx)
    cnext = idxnext <= len ? s[idxnext] : '1'
    if islowercase(cnext)
        # Try a 2-letter element
        element = c*cnext
        element ∈ elements && return element, nextind(s, idxnext)
    end
    # Try a 1-letter element
    element = string(c)
    element ∈ elements && return element, idxnext
    return nothing, idx
end

"""
    i, idxnext = tryparseint16(s, idx)

Parse an `Int16` starting at position `idx` in `s`.

`i` is the `Int16` value, or `nothing` if `s` does not contain an integer starting at `idx`.
`idxnext` is the first position in `s` after the parsed integer.
"""
function tryparseint16(s::AbstractString, idx::Int)
    len = length(s)
    lastidx = prevind(s, idx)
    while lastidx < len
        i = nextind(s, lastidx)
        isdigit(s[i]) || break
        lastidx = i
    end
    return lastidx >= idx ? (parse(Int16, SubString(s, idx, lastidx)), nextind(s, lastidx)) : (nothing, idx)
end

"""
    z, idxnext = tryparseint16(s, idx)

Parse charge state (as an `Int`) starting at position `idx` in `s`.

`z` is the charge, or `nothing` if `s` does not contain a charge starting at `idx`.
`idxnext` is the first position in `s` after the parsed charge.
"""
function tryparsecharge(s::AbstractString, idx::Int)
    len = length(s)
    idx > len && return nothing, idx
    c = s[idx]
    ispm(c) || return nothing, idx
    sgn = c == '+' ? 1 : -1
    idxnext = nextind(s, idx)
    count = 1
    while idxnext <= len
        cnext = s[idxnext]
        if cnext == c
            count += 1
            @warn("++ and -- are deprecated in the OpenSMILES specification, please use + or - followed by multiplicity.", maxlog=1)
            idxnext = nextind(s, idxnext)
        elseif ispm(cnext)
            error("+- and -+ are not allowed for charge in SMILES")
        elseif isdigit(cnext)
            count == 1 || error("cannot use ++ or -- in charge specification followed by multiplicity")
            break
        else
            return sgn*count, idxnext
        end
    end
    idxnext > len && return sgn*count, idxnext
    mult, idxnext = tryparseint16(s, idxnext)
    mult === nothing && return sgn, idxnext
    @assert count == 1
    return sgn*mult, idxnext
end

function tryparsechirality(s::AbstractString, idx::Int)
    m = match(r"(@TH[12]|@AL[12]|@SP[1-3]|@TB\d\d?|@OH\d\d?|@@?)", s, idx)
    (m === nothing || m.offset > idx) && return nothing, idx
    chirality = m.captures[1]::AbstractString
    return chirality, m.offset + lastindex(chirality)
end


export FactorGaussian

include(joinpath(@__DIR__, "../prob_operations.jl"))
include(joinpath(@__DIR__, "../graph_operations.jl"))
include(joinpath(@__DIR__, "../util.jl"))

mutable struct FactorGaussian
    """
    Composite Gaussian factor node.

        f(y, x, a, u, z) = 𝓝(y | a*x + u, z)

    This is a Gaussian node, but slightly more general: it contains a
    multiplication node with a transition coefficient and an addition node with
    a control coefficient. By default, these are clamped such that a standard
    Gaussian is employed.

    Examples
    ≡≡≡≡≡≡≡≡≡≡

    # States
    julia> x_1 = VarGaussian(:x_1);
    julia> x_2 = VarGaussian(:x_2);

    # Define Gaussian factor for state transition
    julia> f = FactorGaussian(:f, out=:x_2, mean=:x_1)

    """

    # Identifiers of edges/nodes in factor graph
    id::Symbol
    variables::Dict{String, Union{Float64, Symbol}}
    time::Union{Integer,Nothing}

    # Reaction parameters
    threshold::Float64
    fired::Bool

    # Additional properties
    verbose::Bool

    function FactorGaussian(id::Symbol;
                            out::Union{Float64, Symbol}=0.0,
                            mean::Union{Float64, Symbol}=0.0,
                            transition::Union{Float64, Symbol}=1.0,
                            control::Union{Float64, Symbol}=0.0,
                            precision::Union{Float64, Symbol}=1.0,
                            threshold::Float64=0.0,
                            time::Union{Integer,Nothing}=0,
                            fired::Bool=false,
                            verbose::Bool=false)

        # Keep track of recognition distributions
        variables = Dict{String, Union{Float64, Symbol}}("out" => out,
                                                         "mean" => mean,
                                                         "transition" => transition,
                                                         "control" => control,
                                                         "precision" => precision)

        # Construct instance
        return new(id, variables, time, threshold, fired, verbose)
    end
end

function message(graph::MetaGraph, node::FactorGaussian, var_id::Symbol)
    "Compute a message to an edge."

    # Extract means and variances of beliefs
    Ex, Vx = meanvar(graph, node.variables["out"])
    Em, Vm = meanvar(graph, node.variables["mean"])
    Ea, Va = meanvar(graph, node.variables["transition"])
    Eu, Vu = meanvar(graph, node.variables["control"])
    Eγ, Vγ = meanvar(graph, node.variables["precision"])

    # Find which variable belongs to var_id
    var_name = key_from_value(node.variables, var_id)

    if var_name == "out"

        message = Normal(Ea*Em, sqrt(inv(Eγ)))

    elseif var_name == "mean"

        message = Normal(Ex*Ea / (Va + Ea^2), sqrt(inv(Eγ*(Va + Ea^2))))

    elseif var_name == "transition"

        message = Normal(Ex*Em / (Vm + Em^2), sqrt(inv(Eγ*(Vm + Em^2))))

    elseif var_name == "control"

        error("Exception: not implemented yet.")

    elseif var_name == "precision"

        inv_scale = (Vx + Ex^2 - 2*Ex*Ea*Em + (Em^2 + Vm)*(Ea^2 + Va))/2
        message = Gamma(3/2., inv(inv_scale))

    else
        throw("Exception: variable unknown.")
    end

    return message
end

function act!(graph::MetaGraph, node::FactorGaussian, var_id::Symbol)
    "Send out message to connected edge."

    # Compute message
    message_f2v = message(graph, node, var_id)

    # Find edge id
    edge_id = Edge(graph[var_id, :id], graph[node.id, :id])

    # Pass message
    set_prop!(graph, edge_id, :message_factor2var, message_f2v)
end

function react!(graph::MetaGraph, node::FactorGaussian)
    "React to incoming messages from edges."

    # Keep track of whether node fired
    node.fired = false

    # Set of variable ids
    var_ids = Set{Symbol}([x for x in values(node.variables) if isa(x, Symbol)])

    # Loop over variables to check for incoming messages
    for var_id in var_ids

        # Find edge id
        edge = Edge(graph[var_id, :id], graph[node.id, :id])

        # If change in FE is large enough, record variable as updated
        if get_prop(graph, edge, :Δfree_energy) >= node.threshold

            # Loop over other variables
            for other_var in setdiff(var_ids, [var_id])

                # Pass message to other variable
                act!(graph, node, other_var)
            end

            # Record that node has fired
            node.fired = true
        end
    end
end


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#

relaxed_tols = (default_tol_relax = 1000,)
insts = OrderedDict()
insts["minimal"] = [
    ((2, 2, false),),
    ((4, 1, false),),
    ((4, 3, false),),
    ((4, 2, true), nothing, (default_tol_relax = 100,)),
    ]
insts["fast"] = [
    ((1000, 1, false),),
    ((200, 1, true), nothing, relaxed_tols),
    ((1000, 2, false),),
    ((200, 2, true), nothing, relaxed_tols),
    ((1000, 3, false), nothing, relaxed_tols),
    ]
insts["various"] = [
    ((1000, 1, false),),
    ((7000, 1, false),),
    ((500, 1, true), nothing, relaxed_tols),
    ((1000, 1, true), nothing, relaxed_tols),
    ((1000, 2, false),),
    ((500, 2, true), nothing, relaxed_tols),
    ((2000, 3, false), nothing, relaxed_tols),
    ((100, 3, true), nothing, relaxed_tols),
    ((400, 3, true), nothing, relaxed_tols),
    ((4000, 3, false), nothing, relaxed_tols),
    ]
return (NonparametricDistrJuMP, insts)


module BiobakeryUtils

export
    import_abundance_tables,
    import_abundance_table,
    clean_abundance_tables,
    taxfilter,
    taxfilter!,
    rm_strat!,
    permanova,
    humann2_regroup,
    humann2_rename,
    humann2_barplot,
    humann2_barplots,
    qvalue!

using DataFrames
using Statistics
using CSV
using RCall
using Microbiome

include("general.jl")
include("metaphlan.jl")
include("humann2.jl")

end