"""
    AbstractMutation{T}

Abstract type defining a mutation operation on entities of type `T`.

Implementations are expected to be callable using an entity of type `T` as only input.

May also implement a callable accepting a `AbstractVector{<:T}` if it is useful to work on
all items to mutate at once.
"""
abstract type AbstractMutation{T} end

(m::AbstractMutation{T})(es::AbstractVector{<:T}) where T = m.(es)

"""
    AbstractCrossover{T}

Type alias for `AbstractMutation{Tuple{T,T}}` defining a crossover of two entities of type `T`.

Implementations are expected to be callable using a tuple of two type `T` as only input.
"""
const AbstractCrossover{T} = AbstractMutation{Tuple{T,T}}

"""
    DecoratingMutation{T}

Abstract type indicating that the type itself does not perform any mutation but wraps a type which might do.

Must either implement callable method for `AbstractVector{<:T}` or accept keyword arguments `next=wrapped(m)` and
`noop=identity` along with a single `T`.
"""
abstract type DecoratingMutation{T} <: AbstractMutation{T} end
wrapped(m::DecoratingMutation) = m.m

mutationleaves(m::DecoratingMutation) = (mutationleaves(wrapped(m))...,)
mutationleaves(tm::Tuple) = mapreduce(mutationleaves, (t1,t2) -> (t1...,t2...), tm)
mutationleaves(m) = tuple(m)

# Apart from being overengineered this helps protecting against forgetting to handle arrays in a DecoratingMutation
# The next/noop happens to work with most existing DecoratingMutations, but it is a bit arbitrary and in some cases
# one must implement both the per-element and the vector of elements versions.
function (m::DecoratingMutation{T})(es::AbstractVector{<:T}) where T
    cnt = Ref(1)
    fornext = Int[]
    next = function(e)
        push!(fornext, cnt[])
        cnt[] += 1
        e
    end
    noop = function(e)
        cnt[] += 1
        e
    end

    allres = m.(es; next, noop)
    mres = wrapped(m)(es[fornext])

    # Mutation might accidentally widen the type compared to allres and then we can't insert mres into allres.
    # Lets fix that if it happens
    RT = typejoin(eltype(allres), eltype(mres))
    res = RT === eltype(allres) ? allres :  convert(Vector{RT}, allres)

    res[fornext] = mres
    return res
end

"""
    MutationProbability{T} <: DecoratingMutation{T}
    MutationProbability(m::AbstractMutation{T}, p::Probability)
    MutationProbability(m::AbstractMutation{T}, p::Number)

Applies `m` with probability `p`.
"""
struct MutationProbability{T, P<:Probability} <: DecoratingMutation{T}
    m::AbstractMutation{T}
    p::P
end
MutationProbability(m::AbstractMutation{T}, p::Number) where T = MutationProbability(m, Probability(p))
(m::MutationProbability{T})(e::T; next=m.m, noop=identity) where T = apply(() -> next(e), m.p, () -> noop(e))

"""
    WeightedMutationProbability{T,F} <: DecoratingMutation{T}
    WeightedMutationProbability(m::AbstractMutation::T, pfun::F)

Applies `m` to an entity `e` with a probability `pfun(e)`.
"""
struct WeightedMutationProbability{T,F} <: DecoratingMutation{T}
    m::AbstractMutation{T}
    pfun::F
end
(m::WeightedMutationProbability{T})(e::T; next=m.m, noop=identity) where T = apply(() -> next(e), m.pfun(e), () -> noop(e))

"""
    HighUtilityMutationProbability(m::AbstractMutation{T}, pbase::Real, rng=rng_default; spread=0.5)

Return a `WeightedMutationProbability` which applies `m` to vertices with an (approximately) average probability of `pbase` and where high `neuronutility` compared to other vertices in same graph means higher probability.

Parameter `spread` can be used to control how much the difference in probability is between high and low utlity. High spread means high difference while low spread means low difference.
"""
HighUtilityMutationProbability(m::AbstractMutation{T}, pbase::Real, rng=rng_default;spread=0.5) where T <: AbstractVertex = WeightedMutationProbability(m, weighted_neuronutility_high(pbase, rng,spread=spread))

"""
    LowUtilityMutationProbability(m::AbstractMutation{T}, pbase::Real, rng=rng_default; spread=2)

Return a `WeightedMutationProbability` which applies `m` to vertices with an (approximately) average probability of `pbase` and where low `neuronutility` compared to other vertices in same graph means higher probability.

Parameter `spread` can be used to control how much the difference in probability is between high and low utlity. High spread means high difference while low spread means low difference.
"""
LowUtilityMutationProbability(m::AbstractMutation{T}, pbase::Real, rng=rng_default;spread=2) where T <: AbstractVertex = WeightedMutationProbability(m, weighted_neuronutility_low(pbase, rng, spread=spread))

weighted_neuronutility_high(pbase, rng=rng_default; spread=0.5) = function(v::AbstractVertex)
    ismissing(NaiveNASflux.neuronutility(v)) && return pbase
    return Probability(fixnan(pbase ^ normexp(v, spread), pbase), rng)
end

weighted_neuronutility_low(pbase, rng=rng_default;spread=2) = function(v::AbstractVertex)
    ismissing(NaiveNASflux.neuronutility(v)) && return pbase
    return Probability(fixnan(pbase ^ (1/normexp(v, 1/spread)), pbase), rng)
end

fixnan(x, rep) = isnan(x) ? rep : clamp(x, 0.0, 1.0)

# This is pretty hacky and arbitrary. Change to something better
function normexp(v::AbstractVertex, s)
    allvertices = filter(allow_mutation, all_in_graph(v))
    allvalues = map(vi -> NaiveNASflux.neuronutility(vi), allvertices)
    meanvalues = map(mean, skipmissing(allvalues))
    meanvalue = mean(meanvalues)
    maxvalue = maximum(meanvalues)
    utlity = mean(NaiveNASflux.neuronutility(v))
    # Basic idea: maxvalue - utlity means the (to be) exponent is <= 0 while the division seems to normalize so that
    # average of pbase ^ normexp across allvertices is near pbase (no proof!). The factor 2 is just to prevent
    # probability of vertex with maxvalue to be 1.
    return (2maxvalue^s - utlity^s) / (2maxvalue^s - meanvalue^s)
end

"""
    MutationChain{T} <: DecoratingMutation{T}
    MutationChain(m::AbstractMutation{T}...)

Chains multiple `AbstractMutation{T}`s after each other.

Input entities will be mutated by the first `AbstractMutation{T}` in the chain and the output will be fed into the next `AbstractMutation{T}` in the chain and so on. The output from the last `AbstractMutation{T}` is returned.
"""
struct MutationChain{T} <: DecoratingMutation{T}
    m::Tuple{Vararg{AbstractMutation{T}}}
end
MutationChain(m::AbstractMutation{T}...) where T = MutationChain(m)
# Identical, but can't use Union due to ambiguity
(m::MutationChain{T})(es::AbstractVector{<:T}) where T = foldl((ei, mi) -> mi(ei), m.m; init=es)
(m::MutationChain{T})(e::T) where T = foldl((ei, mi) -> mi(ei), m.m; init=e)

"""
    RecordMutation{T} <: DecoratingMutation{T}
    RecordMutation(m::AbstractMutation{T})

Records all mutated entities.

Intended use case is to be able to do parameter selection on mutated vertices.
"""
struct RecordMutation{T} <: DecoratingMutation{T}
    m::AbstractMutation{T}
    mutated::Vector{T}
end
RecordMutation(m::AbstractMutation{T}) where T = RecordMutation(m, T[])
function (m::RecordMutation{T})(e::T; next=m.m, noop=identity) where T
    em = next(e)
    push!(m.mutated, em)
    return em
end
function fetchmutated!(m::RecordMutation)
    mutated = copy(m.mutated)
    deleteat!(m.mutated, eachindex(m.mutated))
    return mutated
end

"""
    LogMutation{T} < :DecoratingMutation{T}
    LogMutation(strfun, m::AbstractMutation{T}; level = Logging.Info, nextlogfun=e -> PrefixLogger("   "))
    LogMutation(strfun, level::LogLevel, nextlogfun, m::AbstractMutation{T})

Logs all mutation operations.

Argument `strfun` maps the mutated entity to the logged string.

Calling `nextlogfun(e)` where `e` is the entity to mutate produces an `AbstractLogger` which will be used when applying `m(e)`.

By default, this is used to add a level of indentation to subsequent logging calls which makes logs of hierarchical mutations (e.g. mutate a CompGraph by applying mutations to some of its vertices) easier to read. Set `nextlogfun = e -> current_logger()` to remove this behaviour.
"""
struct LogMutation{T,F,L<:LogLevel,LF} <: DecoratingMutation{T}
    strfun::F
    level::L
    nextlogfun::LF
    m::AbstractMutation{T}
end
LogMutation(strfun, m::AbstractMutation{T}; level = Logging.Info, nextlogfun=e -> PrefixLogger("   ")) where T = LogMutation(strfun, level, nextlogfun, m)
function (m::LogMutation{T})(e::T; next=m.m, noop=identity) where T
    @logmsg m.level m.strfun(e)
    return with_logger(() -> next(e), m.nextlogfun(e))
end

"""
    MutationFilter{T} <: DecoratingMutation{T}
    MutationFilter(predicate, m)

Applies mutation `m` only for entities `e` for which `predicate(e)` returns true.
"""
struct MutationFilter{T,P} <: DecoratingMutation{T}
    predicate::P
    m::AbstractMutation{T}
end
function (m::MutationFilter{T})(e::T; next=m.m, noop=identity) where T
    m.predicate(e) && return next(e)
    return noop(e)
end

"""
    VertexMutation <: DecoratingMutation{CompGraph}
    VertexMutation(m::AbstractMutation{AbstractVertex}, s::AbstractVertexSelection)
    VertexMutation(m::AbstractMutation{AbstractVertex})

Applies a wrapped `AbstractMutation{AbstractVertex}` to each selected vertex in a `CompGraph`.

Vertices to select is determined by the configured `AbstractVertexSelection`.
"""
struct VertexMutation{S<:AbstractVertexSelection} <: DecoratingMutation{CompGraph}
    m::AbstractMutation{AbstractVertex}
    s::S
end
VertexMutation(m::AbstractMutation{AbstractVertex}) = VertexMutation(m, FilterMutationAllowed())
function (m::VertexMutation)(g::CompGraph)
    m.m(select(m.s, g, m))
    return g
end

"""
    NoutMutation <:AbstractMutation{AbstractVertex}
    NoutMutation(l1::Real,l2::Real, rng::AbstractRNG)
    NoutMutation(limit, rng::AbstractRNG=rng_default)
    NoutMutation(l1,l2)

Mutate the out size of a vertex or vector of vertices.

Size is changed by `x * nout(v)` rounded away from from zero where `x` is drawn from `U(minrel, maxrel)` where
`minrel` and `maxrel` are `l1` and `l2` if `l1 < l2` and `l2` and `l1` otherwise.
"""
struct NoutMutation{R<:Real, RNG<:AbstractRNG} <:AbstractMutation{AbstractVertex}
    minrel::R
    maxrel::R
    rng::RNG
    function NoutMutation(l1::R1, l2::R2, rng::RNG) where {R1, R2, RNG}
        R = promote_type(R1, R2)
        return l1 < l2 ? new{R, RNG}(promote(l1, l2)..., rng) : new{R, RNG}(promote(l2, l1)..., rng)
    end
end
NoutMutation(limit, rng::AbstractRNG=rng_default) = NoutMutation(0, limit, rng)
NoutMutation(l1,l2) = NoutMutation(l1,l2, rng_default)
(m::NoutMutation)(v::AbstractVertex) = first(m([v]))
function (m::NoutMutation)(vs::AbstractVector{<:AbstractVertex})

    Δs = Dict{AbstractVertex, Int}()
    shift = m.minrel
    scale = m.maxrel - m.minrel

    for v in vs
        terminputs = findterminating(v, inputs)

        # We are basically just searching for Immutable vertices here, allow_mutation(trait(v)) happens to do just that
        any(tv -> allow_mutation(trait(tv)), terminputs) || continue

        Δfloat = rand(m.rng) * scale + shift

        Δ = ceil(Int, abs(Δfloat) * nout(v)) *  sign(Δfloat)
        minsize = minimum(nout.(terminputs))
        # Or else we might increase the size despite Δ being negative which would be surprising to a user who has specified
        # strictly negative size changes
        minsize + Δ <= 0 && continue

        Δs[v] = Δ
    end

    if !isempty(Δs)
        failmsg = (args...) -> "Could not change nout of $(join(NaiveNASlib.nameorrepr.(keys(Δs)), ", ", " and ")) by $(join(values(Δs), ", ", " and ")). No change!"

        strategy = TimeOutAction(;base=ΔNoutRelaxed(Δs), fallback=LogΔSizeExec(failmsg, Logging.Warn, ΔSizeFailNoOp()))

        Δsize!(strategy)
    end
    return vs
end

"""
    AddVertexMutation <:AbstractMutation{AbstractVertex}
    AddVertexMutation(s::AbstractArchSpace, outselect::Function, WeightInit::AbstractWeightInit, rng::AbstractRNG)
    AddVertexMutation(s, outselect::Function=identity)
    AddVertexMutation(s, rng::AbstractRNG)
    AddVertexMutation(s, wi::AbstractWeightInit)

Insert a vertex from the wrapped `AbstractArchSpace` `s` after a given vertex `v`.

The function `outselect` takes an `AbstractVector{AbstractVertex}` representing the output of `v` and returns an `AbstractVector{AbstractVertex}` which shall be reconnected to the vertex `v'` returned by `s`. Defaults to `identity` meaning all outputs of `v` are reconnected to `v'`.
"""
struct AddVertexMutation{S<:AbstractArchSpace, F, WI<:AbstractWeightInit, RNG<:AbstractRNG} <:AbstractMutation{AbstractVertex}
    s::S
    outselect::F
    weightinit::WI
    rng::RNG
end
AddVertexMutation(s, outselect::Function=identity) = AddVertexMutation(s, outselect, IdentityWeightInit(), rng_default)
AddVertexMutation(s, rng::AbstractRNG) = AddVertexMutation(s, identity, IdentityWeightInit(), rng)
AddVertexMutation(s, wi::AbstractWeightInit) = AddVertexMutation(s, identity, wi, rng_default)

function (m::AddVertexMutation)(v::AbstractVertex)
    insert!(v, vi -> m.s(name(vi), vi, m.rng, outsize=nout(vi), wi=m.weightinit), m.outselect)
    return v
end

"""
    RemoveVertexMutation <:AbstractMutation{AbstractVertex}
    RemoveVertexMutation(s::RemoveStrategy)
    RemoveVertexMutation()

Remove the given vertex `v` using the configured `RemoveStrategy`.

Default size align strategy is `IncreaseSmaller -> DecreaseBigger -> AlignSizeBoth -> FailAlignSizeWarn -> FailAlignSizeRevert`.

Default reconnect strategy is `ConnectAll`.

Note: High likelyhood of large accuracy degradation after applying this mutation.
"""
struct RemoveVertexMutation{S<:RemoveStrategy} <:AbstractMutation{AbstractVertex}
    s::S
end
function RemoveVertexMutation()
    alignstrat = IncreaseSmaller(fallback=DecreaseBigger(fallback=AlignSizeBoth(fallback=FailAlignSizeWarn(msgfun = (vin,vout) -> "Can not remove vertex $(name(vin))! Could not align sizes of neighbours!"))))
    return RemoveVertexMutation(RemoveStrategy(CheckAligned(CheckNoSizeCycle(alignstrat, FailAlignSizeWarn(msgfun = (vin,vout) -> "Can not remove vertex $(name(vin))! Size cycle detected!")))))
end

function (m::RemoveVertexMutation)(v::AbstractVertex)
    remove!(v, m.s)
    return v
end

default_neuronselect(args...) = NaiveNASlib.defaultutility(args...)

"""
    AddEdgeMutation <: AbstractMutation{AbstractVertex}
    AddEdgeMutation(p; rng=rng_default, mergefun=default_mergefun(rng=rng), filtfun=no_shapechange, utilityfun=default_neuronselect)
    AddEdgeMutation(p::Probability; rng=rng_default, mergefun=default_mergefun(rng=rng), filtfun=no_shapechange, utilityfun=default_neuronselect)

Add an edge from a vertex `vi` to another vertex `vo` randomly selected from `vs = filtfun(vi)`.

Higher values of `p` will give more preference to earlier vertices of `vs`.

If `vo` is not capable of having multiple inputs (determined by `singleinput(v) == true`), `vm = mergefun(voi)` where `voi` is a randomly selected input to `vo` will be used instead of `vo` and `vo` will be added as the output of `vm`.

When selecting neurons/outputs after any eventual size change the output of `utilityfun(v)` will be used to determine the utlity of each output in vertex `v`. Note that `length(utilityfun(v)) == nout(v)` must hold.

Note: High likelyhood of large accuracy degradation after applying this mutation.
"""
struct AddEdgeMutation{F1, F2, F3, P<:Probability, RNG} <: AbstractMutation{AbstractVertex}
    mergefun::F1
    filtfun::F2
    utilityfun::F3
    p::P
    rng::RNG
end
AddEdgeMutation(p; rng=rng_default, mergefun=default_mergefun(rng=rng), filtfun=no_shapechange, utilityfun=default_neuronselect) = AddEdgeMutation(Probability(p, rng), rng=rng, mergefun=mergefun, filtfun=filtfun, utilityfun=utilityfun)
AddEdgeMutation(p::Probability; rng=rng_default, mergefun=default_mergefun(rng=rng), filtfun=no_shapechange, utilityfun=default_neuronselect) = AddEdgeMutation(mergefun, filtfun, utilityfun, p, rng)

default_mergefun(pconc = 0.5; rng=rng_default, traitfun = MutationShield ∘ RemoveIfSingleInput ∘ validated() ∘ default_logging(), layerfun = ActivationContribution) = function(vin)
    if rand(rng) > pconc
        return invariantvertex(layerfun(+), vin, traitdecoration=traitfun ∘ named(name(vin) * ".add"))
    end
    return concat(vin, traitfun = traitfun ∘ named(name(vin) * ".cat"), layerfun=layerfun)
end

function no_shapechange(vi)
    # all_in_graph is not sorted, and we want some kind of topoligical order here so that earlier indices are closer to vi
    allsorted = mapreduce(ancestors, vcat, filter(v -> isempty(outputs(v)), all_in_graph(vi))) |> unique

    # Vertices which have the same input as vi and are singleinput
    #   Reason is that this will cause a new vertex to be added between the target output vertex vo
    #   and the input vertex to vi (vii) and this is detected as a size cycle which causes
    #   try_add_edge to fail.
    inouts = filter(singleinput, mapreduce(outputs, vcat, inputs(vi); init=[]))
    # All vertices which are after vi in the topology
    vsafter = setdiff(allsorted, ancestors(vi), outputs(vi), inouts)

    vitrace = shapetrace(vi)
    viorder = allΔshapetypes(vitrace)
    viΔshape = squashshapes(vitrace; order=viorder)

    return filter(vsafter) do vafter
        all(inputs(vafter)) do v
        t = shapetrace(v)
        vΔshape = squashshapes(t; order=union(viorder, allΔshapetypes(t)))
        return viΔshape == vΔshape
        end
    end
end

function (m::AddEdgeMutation)(vi::AbstractVertex)
    # All vertices for which it is allowed to add vi as an input
    allverts = filter(allow_mutation, m.filtfun(vi))
    isempty(allverts) && return vi

    # Higher probability to select a vertex close to v is desired behaviour
    # One line less than a for loop => FP wins!!
    selfun(::Nothing, vc) = apply(m.p) ? vc : nothing
    selfun(vs, vd) = vs
    vo = foldl(selfun, allverts, init=nothing)
    vo = isnothing(vo) ? rand(m.rng, allverts) : vo

    try_add_edge(vi, vo, m.mergefun, m.utilityfun)
    return vi
end

function try_add_edge(vi, vo, mergefun, utilityfun=default_neuronselect)

    # Need to add a vertex which can handle multiple inputs if vo is single input only
    # For cleaning up added vertex if the whole operation fails
    cleanup_failed = () -> nothing
    if singleinput(vo)
        voi = inputs(vo)[1]
        # If the input to vo is capable of multi input we don't need to create a new vertex
        # We must also check that this input does not happen to be an input to vi as this would create a cycle in the graph
        if singleinput(voi) || voi in ancestors(vi)
            vm = mergefun(voi)
            # Insert vm between voi and vo, i.e voi -> vo turns into voi -> vm -> vo
            # vs -> [vo] means only add the new vertex between voi and vo as voi could have other outputs
            insert!(voi, vv -> vm, vs -> [vo])
            cleanup_failed = function()
                length(inputs(vm)) > 1 && return
                remove!(vm, RemoveStrategy(NoSizeChange()))
            end
            vo = vm # vm is the one we shall add an edge to
            @debug "Create new vertex for merging $(name(vo))"
        else
            vo = voi
        end
    end
    # This is mainly because FailAlignSizeRevert does not work when the same vertex is input more than once, but it also seems kinda redundant.
    vi in inputs(vo) && return
    @debug "Create edge between $(name(vi)) and $(name(vo))"
    create_edge!(vi, vo, strategy = create_edge_strat(vo, utilityfun))
    cleanup_failed()
end
# Need to override this one for strange types e.g. layers which support exactly 2 inputs or something.
singleinput(v) = isempty(inputs(v)) || length(inputs(v)) == 1

create_edge_strat(v::AbstractVertex, utilityfun) = create_edge_strat(trait(v), utilityfun)
create_edge_strat(d::DecoratingTrait, utilityfun) = create_edge_strat(base(d), utilityfun)
function create_edge_strat(::SizeInvariant, utilityfun)
    warnfailalign = FailAlignSizeWarn(msgfun = (vin,vout) -> "Could not align sizes of $(name(vin)) and $(name(vout))!")
    alignstrat = AlignSizeBoth(;mapstrat=WithUtilityFun(utilityfun), fallback = warnfailalign)
    # Tricky failure case: It is possible that CheckCreateEdgeNoSizeCycle does not detect any size cycle until after the edge has been created?
    sizecyclewarn = FailAlignSizeWarn(msgfun = (vin,vout) -> "Could not align sizes of $(name(vin)) and $(name(vout))! Size cycle detected!")

    return CheckCreateEdgeNoSizeCycle(ifok=alignstrat, ifnok=sizecyclewarn)
end
function create_edge_strat(::SizeStack, utilityfun)
    warnfailalign = FailAlignSizeWarn(msgfun = (vin,vout) -> "Could not align sizes of $(name(vin)) and $(name(vout))!")
    alignstrat = PostAlign(TruncateInIndsToValid(WithUtilityFun(utilityfun, AlignNinToNout(;fallback=ΔSizeFailNoOp()))), fallback=warnfailalign)

    sizecyclewarn = FailAlignSizeWarn(msgfun = (vin,vout) -> "Could not align sizes of $(name(vin)) and $(name(vout))! Size cycle detected!")
    return CheckCreateEdgeNoSizeCycle(ifok=alignstrat, ifnok=sizecyclewarn)
end

"""
    RemoveEdgeMutation <: AbstractMutation{AbstractVertex}
    RemoveEdgeMutation(;utilityfun=default_neuronselect, rng=rng_default)

Remove an edge from a vertex `vi` to another vertex `vo` randomly selected from `outputs(vi)`.

Vertex `vi` must have more than one output and vertex `vo` must have more than one output for the edge to be removed. Otherwise no change is made.

If there are multiple edges between `vi` and `vo` no change will be made due to NaiveNASlib not being able to revert a failed operation in this case..

When selecting neurons/outputs after any eventual size change the output of `utilityfun(v)` will be used to determine the utlity of each output in vertex `v`. Note that `length(utilityfun(v)) == nout(v)` must hold.

Note: High likelyhood of large accuracy degradation after applying this mutation.
"""
struct RemoveEdgeMutation{F, RNG<:AbstractRNG} <: AbstractMutation{AbstractVertex}
    utilityfun::F
    rng::RNG
end
RemoveEdgeMutation(;utilityfun=default_neuronselect, rng=rng_default) = RemoveEdgeMutation(utilityfun, rng)

function (m::RemoveEdgeMutation)(vi::AbstractVertex)
    length(outputs(vi)) < 2 && return vi

    allverts = filter(vo -> length(inputs(vo)) > 1, outputs(vi))

    isempty(allverts) && return vi

    vo = rand(m.rng, allverts)
    sum(inputs(vo) .== vi) > 1 && return vi# Not implemented in NaiveNASlib

    @debug "Remove edge between $(name(vi)) and $(name(vo))"
    remove_edge!(vi, vo, strategy=remove_edge_strat(vo, m.utilityfun))
    return vi
end

remove_edge_strat(v::AbstractVertex, utilityfun) = remove_edge_strat(trait(v), utilityfun)
remove_edge_strat(d::DecoratingTrait, utilityfun) = remove_edge_strat(base(d), utilityfun)
remove_edge_strat(::SizeInvariant, utilityfun) = NoSizeChange()
remove_edge_strat(t::SizeStack, utilityfun) = create_edge_strat(t, utilityfun)

"""
    KernelSizeMutation{N} <: AbstractMutation{AbstractVertex}
    KernelSizeMutation(Δsizespace::AbstractParSpace{N, Int}; maxsize, pad, rng)
    KernelSizeMutation2D(absΔ::Integer;maxsize, pad, rng)
    KernelSizeMutation(absΔ::Integer...;maxsize, pad, rng)

Mutate the size of filter kernels of convolutional layers.

Note: High likelyhood of large accuracy degradation after applying this mutation.

`KernelSizeMutation2D` is a convenience constructor for `KernelSizeMutation(absΔ, absΔ;...)`.
"""
struct KernelSizeMutation{N,F,P} <: AbstractMutation{AbstractVertex}
    Δsizespace::AbstractParSpace{N, Int}
    maxsize::F
    pad::P
    rng::AbstractRNG
end
KernelSizeMutation(Δsizespace::AbstractParSpace{N, Int}; maxsize = v -> ntuple(i->Inf,N), pad=SamePad(), rng=rng_default) where N = KernelSizeMutation(Δsizespace, maxsize, pad, rng)
KernelSizeMutation2D(absΔ::Integer;maxsize = v -> (Inf,Inf), pad=SamePad(), rng=rng_default) = KernelSizeMutation(absΔ, absΔ, maxsize = maxsize, pad=pad, rng=rng)
KernelSizeMutation(absΔ::Integer...;maxsize = v -> ntuple(i->Inf, length(absΔ)), pad=SamePad(), rng=rng_default) = KernelSizeMutation(ParSpace(UnitRange.(.-absΔ, absΔ));maxsize = maxsize, pad=pad, rng=rng)

function (m::KernelSizeMutation{N})(v::AbstractVertex) where N
    layertype(v) isa FluxConvolutional{N} || return
    l = layer(v)

    currsize = size(NaiveNASflux.weights(l))[1:N]
    Δsize = Int.(clamp.(m.Δsizespace(m.rng), 1 .- currsize, m.maxsize(v) .- currsize)) # ensure new size is > 0 and < maxsize
    # This will eventually boil down to Setfield doing its thing, and that won't be using any convenience constructors
    pad = Flux.calc_padding(typeof(l), m.pad, currsize .+ Δsize, dilation(l), stride(l))
    KernelSizeAligned(Δsize, pad)(v)
    return v
end
dilation(l) = l.dilation
stride(l) = l.stride

"""
    ActivationFunctionMutation{T,R} <: AbstractMutation{AbstractVertex} where {T <: AbstractParSpace{1}, R <: AbstractRNG}
    ActivationFunctionMutation(actspace::AbstractParSpace{1}, rng::AbstractRNG)
    ActivationFunctionMutation(acts...;rng=rng_default)
    ActivationFunctionMutation(acts::AbstractVector;rng=rng_default)

Mutate the activation function of layers which have an activation function.

Note: High likelyhood of large accuracy degradation after applying this mutation.
"""
struct ActivationFunctionMutation{T,RNG} <: AbstractMutation{AbstractVertex} where {T <: AbstractParSpace{1}, R <: AbstractRNG}
    actspace::T
    rng::RNG
end
ActivationFunctionMutation(acts...;rng=rng_default) = ActivationFunctionMutation(collect(acts), rng=rng)
ActivationFunctionMutation(acts::AbstractVector;rng=rng_default) = ActivationFunctionMutation(ParSpace(acts), rng)

function (m::ActivationFunctionMutation)(v::AbstractVertex)
    m(layertype(v), v)
    return v
end
function (m::ActivationFunctionMutation)(t, v) end
(m::ActivationFunctionMutation)(::Union{FluxDense, FluxConvolutional}, v) = NaiveNASflux.setlayer!(v, (σ = m.actspace(m.rng),))
(m::ActivationFunctionMutation)(::FluxParNorm, v) = NaiveNASflux.setlayer!(v, (λ = m.actspace(m.rng),))
function (m::ActivationFunctionMutation)(::FluxRnn, v)
    newcell = setproperties(layer(v).cell, (σ = m.actspace(m.rng),))
    NaiveNASflux.setlayer!(v, (cell = newcell,))
end

"""
    PostMutation{T} <: DecoratingMutation{T}
    PostMutation(actions, m::AbstractMutation{T})
    PostMutation(m::AbstractMutation{T}, actions...)

Performs a set of actions after a wrapped `AbstractMutation` is applied.

Actions will be invoked with arguments (m::PostMutation{T}, e::T) where m is the enclosing `PostMutation` and `e` is the mutated entity of type `T`.
"""
struct PostMutation{T,A} <: DecoratingMutation{T}
    actions::A
    m::AbstractMutation{T}
end
PostMutation(m::AbstractMutation{T}, actions...) where T = PostMutation(actions, m)
PostMutation(action::Function, m::AbstractMutation{T}) where T = PostMutation(m, action)
function (m::PostMutation{T})(e::T; next=m.m, noop=identity) where T
    eout = next(e)
    foreach(a -> a(m, eout), m.actions)
    return eout
end

"""
    RemoveZeroNout()
    RemoveZeroNout(fallback)

Search for vertices with zero output size and remove them and all of their input vertices if possible to do so witout removing an input or output vertex.

Removal is only possible if a vertex is inside a parallel path which will later be concatenated.
"""
struct RemoveZeroNout
    fallback
end
RemoveZeroNout() = RemoveZeroNout(IncreaseZeroNout())
struct IncreaseZeroNout end

(r::RemoveZeroNout)(m, e) = r(e)
function (r::RemoveZeroNout)(g::CompGraph)
    topoligical_order = vertices(g)
    for v in topoligical_order
        nout(v) == 0 || continue

        # Beware! This turned out to be a lot harder than I first thought.
        # I'm not sure the algorithm works (or realizes that it won't work) for all possible cases.

        # This obviously only works if v is inside a parallel path which will later be concatenated

        # To make sure it is, we look ahead in forward and backward direction.
        # If we don't see 1) an input vertex 2) a vertex without outputs (i.e an output vertex) we are good to go
        fseen, fpoints = findforkpoint(v, topoligical_order, inputs, outputs)
        isempty(fpoints) && continue

        bseen, bpoints = findforkpoint(v, topoligical_order, outputs, inputs)
        isempty(bpoints) && continue

        # Ok, fpoints are all vertices where forks with v in them join and bpoints are all vertices where forks with v in them begin

        # If we start removing all seen vertices which are input to an fpoint until we hit a bpoint we should have removed the fork with v in it, right?
        seen = union(fseen, bseen)
        to_rm = intersect(vcat(inputs.(fpoints)...), seen)
        foreach(fpoint -> remove_all_inputs(fpoint, seen, bpoints), to_rm)
    end
end

function findforkpoint(v::AbstractVertex, topoligical_order, f1=inputs, f2=outputs, seen = vcat(v, f1(v)), points = AbstractVertex[])

    if all(x -> x in seen, f1(v))

        # Check if we came across a vertex which previously was thought to be a forkpoint and remove it if so
        if v in points
            deleteat!(points, indexin([v], points))
        end
        # Always visit in reverse topoligical order to mitigate chasing down paths just because we haven't explored them yet
        nextverts = reverse(topoligical_order[unique(indexin(f2(v), topoligical_order))])
        push!(seen, filter(v2 -> !in(v2, seen), nextverts)...)
        foreach(v2 -> findforkpoint(v2, topoligical_order, f1, f2, seen, points), nextverts)
    elseif !(v in points)
        push!(points, v)
    end

    return seen, points
end

function remove_all_inputs(v, seen, stop)
    v in stop && return
    foreach(vrm -> remove_all_inputs(vrm, seen, stop), intersect(inputs(v), seen))
    remove!(v, RemoveStrategy(ConnectNone(), NoSizeChange()))
end

"""
    struct OptimizerMutation{F} <: AbstractMutation{FluxOptimizer}
    OptimizerMutation(optfun)
    OptimizerMutation(os::Union{Tuple, <:AbstractArray})

Mutatates optimizers not wrapped in `ShieldedOpt` through `optfun`.

Invoked recursively for `Flux.Optimiser`s.
"""
struct OptimizerMutation{F} <: AbstractMutation{FluxOptimizer}
    optfun::F
end
OptimizerMutation(os::Union{Tuple, <:AbstractArray}, rng=rng_default) = OptimizerMutation(o -> rand(rng, os)(learningrate(o)))

"""
    LearningRateMutation(rng=rng_default)

Return an `OptimizerMutation` which mutates the learning rate of optimizers.
"""
LearningRateMutation(rng=rng_default) = OptimizerMutation(o -> nudgelr(o, rng))

(m::OptimizerMutation)(opt::Flux.Optimiser) = Flux.Optimiser(m.(opt.os))
(m::OptimizerMutation)(o::ShieldedOpt) = o;
(m::OptimizerMutation)(o::FluxOptimizer) = m.optfun(o)

nudgelr(o, rng=rng_default) = sameopt(o, nudgelr(learningrate(o), rng))
nudgelr(lr::Number, rng=rng_default) = clamp(lr + (rand(rng) - 0.5) * lr * 0.3, 1e-6, 1.0)

learningrate(o::Flux.Optimiser) = prod(learningrate.(o.os))
learningrate(o::ShieldedOpt) = learningrate(o.opt)
learningrate(o) = o.eta

newlr(o, lrf = nudgelr) = sameopt(o, lrf(learningrate(o)))
sameopt(o, lr) = @set o.eta = lr

"""
    AddOptimizerMutation{F} <: AbstractMutation{FluxOptimizer}

Adds optimizer generated by `optgen(os)` to the set of optimizers where `os` is the existing set.

An attempt to merge optimizers of the same type is made using `mergeopts`.
"""
struct AddOptimizerMutation{F} <: AbstractMutation{FluxOptimizer}
    optgen::F
end
(m::AddOptimizerMutation)(o::ShieldedOpt) = o;
(m::AddOptimizerMutation)(o::FluxOptimizer) = m(Flux.Optimiser([o]))
function (m::AddOptimizerMutation)(opt::Flux.Optimiser)
    newopt = m.optgen(opt)
    return Flux.Optimiser(mergeopts(typeof(newopt), newopt, opt.os...))
end


module Inspire

using Codex.Questionnaires: Items, get_scores
using DataFrames

export mike2scores

# The scoring is determined by guessing by me.
# I expect to have made a few mistakes, but overall the scoring was reasonably after knowing
# the number of items per predictor; thanks to
# https://repository.tudelft.nl/view/tno/uuid:98d97d5c-ff8e-4f26-9fba-f993bc585d39.
# For coping, the last 4 items seem to be reversed and added later by someone.
# Those are ignored.

coping_flexibility = Items(1:36, [])

emotional_stability = Items(
    [41, 44, 49, 50, 52, 54, 56],
    [42, 43, 45, 46, 47, 48, 51, 53, 55]
)

optimism = Items(
    [61, 63, 69],
    [58, 59, 67]
)

social_competence = Items(
    [57, 60, 62, 64, 65, 66, 68, 70],
    []
)

self_efficacy = Items(71:83, [])

self_reflection = Items(
    [88, 90, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103],
    [84, 85, 86, 87, 91, 92, 89]
)

function mike2scores(df::DataFrame)
    new_df = DataFrame(df)
    if !("completed_at" in names(new_df))
        # Cleaning up the 2018 data.
        new_df[:, :completed_at] = repeat([""], nrow(new_df))
        select!(new_df, :id, :completed_at, Not([:id, :completed_at]))
        new_df = filter(:id => x -> !contains(x, "unknown"), new_df)
    end

    new_df[:, :coping_flexibility] = get_scores(new_df, coping_flexibility)
    new_df[:, :emotional_stability] = get_scores(new_df, emotional_stability)
    new_df[:, :optimism] = get_scores(new_df, optimism)
    new_df[:, :social_competence] = get_scores(new_df, social_competence)
    new_df[:, :self_efficacy] = get_scores(new_df, self_efficacy)
    new_df[:, :self_reflection] = get_scores(new_df, self_reflection)

    scored = select(new_df, :id, :completed_at,
        :coping_flexibility, :emotional_stability, :optimism,
        :social_competence, :self_efficacy, :self_reflection
    )
    # This could have been done above, but I mistakenly used `disallowmissing`
    # instead and went from there.
    dropmissing!(scored)
    scored
end

end # module


export FlowModel, CompiledFlowModel
export compile, nr_params
export getlayers

using TupleTools: flatten

import Base: length, eltype

@doc raw"""
The FlowModel structure is the most generic type of Flow model, in which the layers are not constrained to be of a specific type. The FlowModel structure contains the input dimensionality and a tuple of layers and can be constructed as `FlowModel( dim, (layer1, layer2, ...) )`.

Note: this model can be specialized by constraining the types of layers. This potentially allows for more efficient specialized methods that can deal with specifics of these layers, such as triangular jacobian matrices.
"""
struct FlowModel{N,T<:NTuple{N,AbstractLayer}} <: AbstractFlowModel
    dim     :: Int
    layers  :: T
end

@doc raw"""
Creates a flow model with uninitialized parameters.

Input arguments:
- `dim::Int` - input dimensionality.
- `layers<:NTuple{N,AbstractLayer}` - arbitrarily sized tuple containing abstractlayers.

Return arguments:
- `::Flowmodel` - model containing layers of which are appropriately sized according to the input dimensionality.
"""
function FlowModel(dim::Int, layers::T) where { T <: NTuple{N,AbstractLayerPlaceholder} where { N } }
    return FlowModel(dim, flatten(prepare.(dim, layers)))
end
function FlowModel(layers::T) where { T <: NTuple{N, Union{AbstractLayer, AbstractLayerPlaceholder}} where { N } }
    @assert typeof(first(layers)) <: InputLayer "The FlowModel requires an input dimension to be specified. This can be achieved, either by preceding the layers tuple with an integer as `FlowModel(dim, layers)`, or by starting the tuple of layers with an `InputLayer(dim)` as `FlowModel((InputLayer(dim), layers...))`."
    return FlowModel(getdim(first(layers)), flatten(prepare.(getdim(first(layers)), Base.tail(layers))))
end

# prepare function for setting correct sizes in the layers (without assigning the parameters yet!)
prepare(dim::Int, layers::T) where { T <: NTuple{N,Union{AbstractLayer, AbstractLayerPlaceholder}} where { N }} = _prepare(dim, layers)
Broadcast.broadcasted(::typeof(prepare), dim::Int, layers::T) where { T <: NTuple{N,Union{AbstractLayer, AbstractLayerPlaceholder}} where { N }} = broadcast(_prepare, Ref(dim), layers)

@doc raw"""
The CompiledFlowModel structure is the most generic type of compiled Flow model, in which the layers are not constrained to be of a specific type. The FlowModel structure contains the input dimension and a tuple of compiled layers. Do not manually create a CompiledFlowModel! Instead create a FlowModel first and compile it with `compile(model::FlowModel)`. This will make sure that all layers/mappings are configured with the proper dimensionality and with randomly sampled parameters. Alternatively, if you would like to pass your own parameters, call `compile(model::FlowModel, params::Vector)`.

Note: this model can be specialized by constraining the types of layers. This potentially allows for more efficient specialized methods that can deal with specifics of these layers, such as triangular jacobian matrices.
"""
struct CompiledFlowModel{N,T<:NTuple{N,AbstractLayer}} <: AbstractCompiledFlowModel
    dim     :: Int
    layers  :: T
end

@doc raw"""
`compile()` compiles a model by setting its parameters. It randomly sets parameter values in the layers and flows such that inference in the model can be obtained.

Input arguments
- `model::FlowModel` - a model of which the dimensionality of its layers/flows has been initialized, but its parameters have not been set.

Return arguments
- `::CompiledFlowModel` - a compiled model with set parameters, such that it can be used for processing data.
"""
function compile(model::FlowModel)
    # do not create parameters here for decentralized initializers (TODO)
    return CompiledFlowModel(model.dim, compile.(model.layers))
end

@doc raw"""
`compile(model::FlowModel, params::Vector)` lets you initialize a model `model` with a vector of parameters `params`.

Input arguments
- `model::FlowModel` - a model of which the dimensionality of its layers/flows has been initialized, but its parameters have not been set.
- `params::Vector`   - a vector of parameters with which the model should be compiled.

Return arguments
- `::CompiledFlowModel` - a compiled model with set parameters, such that it can be used for processing data.
"""
function compile(model::FlowModel, params::Vector)

    # assert whether the parameter vector complies with the model
    @assert nr_params(model) == length(params) "The number of parameters in the model does not match the passed number of parameters."

    # compile layers
    return CompiledFlowModel(getdim(model), compile(getlayers(model), params))

end
function compile(x::Tuple, params::Vector)

    # fetch number of parameters in layer
    params_in_layer = nr_params(first(x))

    if params_in_layer > 0
        # if layer has parameters, compile that layer and the other layers with parameters
        return (compile(first(x), params[1:params_in_layer]), compile(Base.tail(x), params[params_in_layer+1:end])...)
    else
        # if layer has no parameters, compile that layer on its own and the other layers with parameters
        return (compile(first(x)), compile(Base.tail(x), params)...)
    end

end
function compile(x::Tuple)
    return (compile(first(x)), compile(Base.tail(x))...)
end
# compile functions to stop the recursion
compile(x::Tuple{}, params::Vector) = ()
compile(x::Tuple{})         = ()

# get-functions for the FlowModel structure
getdim(model::FlowModel)                    = model.dim
getdim(model::CompiledFlowModel)            = model.dim
getlayers(model::FlowModel)                 = model.layers
getlayers(model::CompiledFlowModel)         = model.layers
getforward(model::CompiledFlowModel)        = (x) -> forward(model, x)
getbackward(model::CompiledFlowModel)       = (x) -> backward(model, x)
getjacobian(model::CompiledFlowModel)       = (x) -> jacobian(model, x)
getinv_jacobian(model::CompiledFlowModel)   = (x) -> inv_jacobian(model, x)

# custom Base function for the FlowModel structure
eltype(model::CompiledFlowModel) = promote_type(map(eltype, model.layers)...)
length(model::CompiledFlowModel) = length(model.layers)

# fetch number of parameters of model
nr_params(model::FlowModel)         = nr_params(getlayers(model))
nr_params(model::CompiledFlowModel) = nr_params(getlayers(model))
nr_params(layers::T) where { T <: NTuple{N,AbstractLayer} where N} = nr_params(first(layers)) + nr_params(Base.tail(layers))
nr_params(layers::Tuple{}) = return 0

# forward pass through the Flow model
function _forward(model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for result
    output = zeros(Ti, size(input))

    # perform calculations
    forward!(output, model, input)

    # return result
    return output

end

# when calling forward, redirect to _forward
forward(model::CompiledFlowModel, input::AbstractVector{ <: Real }) = _forward(model, input)

# for broadcasting over forward, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(forward), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_forward, Ref(model), input)

# inplace forward pass through the Flow model
function forward!(output::AbstractVector{ <: Real }, model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # fetch layers
    layers = getlayers(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # decouple changing input from actual input
    input_new = zeros(Ti, size(input))
    input_new .= input

    # perform forward pass over all layers
    forward!(output, layers, input_new)

end

# inplace forward pass through a tuple of layers
function forward!(output::AbstractVector{ <: Real }, layers::T, input::AbstractVector{ <: Real }) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass over first layer
    forward!(output, first(layers), input)

    # update intermediate input
    input .= output

    # perform pass over all remaining layers
    forward!(output, Base.tail(layers), input)

end

# when no layers are left, stop the inplace recursion
forward!(output::AbstractVector{ <: Real }, layers::Tuple{}, input::AbstractVector{ <: Real }) = return

# backward pass through the Flow model
function _backward(model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for result
    input = zeros(Ti, size(output))

    # perform calculations
    backward!(input, model, output)

    # return result
    return input

end

# when calling backward, redirect to _backward
backward(model::CompiledFlowModel, output::AbstractVector{ <: Real }) = _backward(model, output)

# for broadcasting over backward, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(backward), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_backward, Ref(model), output)

# inplace backward pass through the Flow model
function backward!(input::AbstractVector{ <: Real }, model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # fetch layers
    layers = getlayers(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # decouple changing output from actual output
    output_new = zeros(Ti, size(output))
    output_new .= output

    # perform backward pass over all layers
    backward!(input, layers, output_new)

end

# inplace backward pass through a tuple of layers
function backward!(input::AbstractVector{ <: Real }, layers::T, output::AbstractVector{ <: Real }) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass over last layer
    backward!(input, last(layers), output)

    # update intermediate output
    output .= input

    # perform pass over all remaining layers
    backward!(input, Base.front(layers), output)

end

# when no layers are left, stop the inplace recursion
backward!(input::AbstractVector{ <: Real }, layers::Tuple{}, output::AbstractVector{ <: Real }) = return

# joint forward-jacobian of the Flow model
function _forward_jacobian(model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # fetch layers
    dim    = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for output and jacobian
    output = zeros(Ti, dim)
    J = zeros(Ti, dim, dim)
    for k = 1:dim
        J[k,k] = one(Ti)
    end

    # calculate jacobian
    forward_jacobian!(output, J, model, input)

    # return result
    return output, J

end

# when calling forward_jacobian, redirect to _forward_jacobian
forward_jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real }) = _forward_jacobian(model, input)

# for broadcasting over forward_jacobian, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(forward_jacobian), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_jacobian, Ref(model), input)

# inplace forward_jacobian of the Flow model
function forward_jacobian!(output::AbstractVector{ <: Real }, J::AbstractMatrix{T}, model::CompiledFlowModel, input::AbstractVector{ <: Real }) where { T <: Real }

    # fetch layers
    layers = getlayers(model)
    dim = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for intermediate output and jacobian
    input_new = zeros(Ti, dim)
    input_new .= input
    J_new = zeros(Ti, dim, dim)
    J_old = zeros(Ti, dim, dim)
    for k = 1:dim
        J_old[k,k] = one(Ti)
    end

    # calculate forward_jacobian over all layers
    forward_jacobian!(J, J_new, J_old, output, input_new, layers)

end

# inplace forward_jacobian calculation for a tuple of layers
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layers::T) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass through first layers
    forward_jacobian!(J, J_new, J_old, output, input_new, first(layers))

    # update unless we are on the last layer
    if length(layers) > 1

        # update intermediate input
        input_new .= output

        # update old jacobian
        J_old .= J

    end

    # calculate forward_jacobian of remaining layers
    forward_jacobian!(J, J_new, J_old, output, input_new, Base.tail(layers))

end

# specialized methods
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layer::PermutationLayer)
    # calculate new output
    forward!(output, layer, input_new)

    # calculate total jacobian
    mul!(J, jacobian(layer, input_new), J_old)
end

# standard method
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layer::AbstractLayer)
    # calculate new jacobian
    forward_jacobian!(output, J_new, layer, input_new)

    # calculate total jacobian
    mul!(J, J_new, J_old)
end

# when no layers are left, stop the inplace recursion
forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layers::Tuple{}) = return

# joing backward inverse jacobian of the Flow model
function _backward_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # fetch layers
    dim    = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for jacobian
    J = zeros(Ti, dim, dim)
    input = zeros(Ti, dim)
    for k = 1:dim
        J[k,k] = one(Ti)
    end

    # calculate jacobian
    backward_inv_jacobian!(input, J, model, output)

    # return result
    return input, J

end

# when calling backward inverse jacobian, redirect to _backward_inv_jacobian
backward_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real}) = _backward_inv_jacobian(model, output)

# for broadcasting over backward inverse jacobian, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(backward_inv_jacobian), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_backward_inv_jacobian, Ref(model), output)

# inplace inverse backward jacobian of the Flow model
function backward_inv_jacobian!(input::AbstractVector{ <: Real }, J::AbstractMatrix{T}, model::CompiledFlowModel, output::AbstractVector{ <: Real }) where { T <: Real }

    # fetch layers
    layers = getlayers(model)
    dim = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for intermediate output and jacobian
    output_new = zeros(Ti, dim)
    output_new .= output
    J_new = zeros(Ti, dim, dim)
    J_old = zeros(Ti, dim, dim)
    for k = 1:dim
        J_old[k,k] = one(Ti)
    end

    # calculate backward inverse jacobian over all layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, layers)

end

# inplace backward inverse jacobian calculation for a tuple of layers
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layers::T) where { T <: NTuple{N,AbstractLayer} where N}

    # perform backward pass through layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, last(layers))

    # update unless we are on the last/first layer
    if length(layers) > 1

        # update intermediate output
        output_new .= input

        # update old jacobian
        J_old .= J

    end

    # calculate backward inv jacobian of remaining layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, Base.front(layers))

end

# specialized methods
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layer::PermutationLayer)
    # perform forward pass over last layer
    backward!(input, layer, output_new)

    # calculate total jacobian
    mul!(J, inv_jacobian(layer, output_new), J_old)
end

# standard method
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layer::AbstractLayer)
    # perform forward pass over last layer
    backward_inv_jacobian!(input, J_new, layer, output_new)

    # calculate total jacobian
    mul!(J, J_new, J_old)
end

# when no layers are left, stop the inplace recursion
backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layers::Tuple{}) = return

# specify jacobian and inv_jacobian functions based on joint functions
_jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real })      = forward_jacobian(model, input)[2]
jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real })       = _jacobian(model, input)
Broadcast.broadcasted(::typeof(jacobian), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_jacobian, Ref(model), input)
_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real }) = backward_inv_jacobian(model, output)[2]
inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real })  = _inv_jacobian(model, output)
Broadcast.broadcasted(::typeof(inv_jacobian), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_inv_jacobian, Ref(model), output)
function jacobian!(J::AbstractMatrix{T1}, model::CompiledFlowModel, input::AbstractVector{T2}) where { T1 <: Real, T2 <: Real }

    # fetch dimension
    dim    = getdim(model)

    # promote type for allocating output
    T = promote_type(eltype(model), T1, T2)

    # allocate space for output and jacobian
    output = zeros(T, dim)
    J .= zero(T1)
    for k = 1:dim
        J[k,k] = one(T1)
    end

    # calculate jacobian
    forward_jacobian!(output, J, model, input)

end
function inv_jacobian!(J::AbstractMatrix{T1}, model::CompiledFlowModel, output::AbstractVector{T2}) where { T1 <: Real, T2 <: Real }

    # fetch dimension
    dim    = getdim(model)

    # promote type for allocating output
    T = promote_type(eltype(model), T1, T2)

    # allocate space for input and jacobian
    input = zeros(T, dim)
    J .= zero(T1)
    for k = 1:dim
        J[k,k] = one(T1)
    end

    # calculate jacobian
    backward_inv_jacobian!(input, J, model, output)

end

# fallback joint functions over layers
function forward_jacobian!(output::AbstractVector{ <: Real }, J_new::AbstractMatrix{ <: Real }, layer::AbstractLayer, input::AbstractVector{ <: Real })
    forward!(output, layer, input)
    jacobian!(J_new, layer, input)
end
function backward_inv_jacobian!(input::AbstractVector{ <: Real }, J_new::AbstractMatrix{ <: Real }, layer::AbstractLayer, output::AbstractVector{ <: Real })
    backward!(input, layer, output)
    inv_jacobian!(J_new, layer, output)
end

# extra utility functions
det_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})           = det(jacobian(model, input))
absdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})        = abs(det_jacobian(model, input))
logdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})        = logdet(jacobian(model, input))
logabsdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})     = logabsdet(jacobian(model, input))

detinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})       = det(inv_jacobian(model, output))
absdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})    = abs(detinv_jacobian(model, output))
logdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})    = logdet(inv_jacobian(model, output))
logabsdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real}) = logabsdet(inv_jacobian(model, output))

# throw an error when the model has not yet been compiled
_forward(model::FlowModel, input::AbstractVector{ <: Real })                                                        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
forward(model::FlowModel, input::AbstractVector{ <: Real })                                                         = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(forward), model::FlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } })   = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
forward!(output::AbstractVector{ <: Real }, model::FlowModel, input::AbstractVector{ <: Real })                     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_backward(model::FlowModel, output::AbstractVector{ <: Real })                                                      = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
backward(model::FlowModel, output::AbstractVector{ <: Real })                                                       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(backward), model::FlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
backward!(input::AbstractVector{ <: Real }, model::FlowModel, output::AbstractVector{ <: Real })                    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
jacobian!(J_new::AbstractMatrix{ <: Real }, model::FlowModel, input::AbstractVector{ <: Real })                     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
jacobian(model::FlowModel, input::AbstractVector{ <: Real })                                                        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(jacobian), model::FlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } })  = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_jacobian(model::FlowModel, input::AbstractVector{ <: Real })                                                       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_inv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })                                                  = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
inv_jacobian!(J_new::AbstractMatrix{ <: Real }, model::FlowModel, output::AbstractVector{ <: Real })                = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))

det_jacobian(model::FlowModel, input::AbstractVector{ <: Real })           = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
absdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logabsdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))

detinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
absdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logabsdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real }) = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))


function count_pawn_advances(board::Board, targets::Bitboard)
    n = 0
    if board.turn == WHITE
        n += count_wpawn_advances(board, targets)
    else
        n += count_bpawn_advances(board, targets)
    end
    return n
end

function count_wpawn_advances(board::Board, targets::Bitboard)
    n = 0
    pwns = (pawns(board) & friendly(board)) & (~pinned(board) | file(kings(board) & friendly(board)))

    # Single pawn advances
    dests = (pwns << 8) & empty(board) & ~RANK_8
    n += count(dests & targets)

    # Double pawn advances
    doubledests = (dests << 8) & targets & RANK_4
    n += count(doubledests)
    return n
end

function count_bpawn_advances(board::Board, targets::Bitboard)
    n = 0
    pwns = (pawns(board) & friendly(board)) & (~pinned(board) | file(kings(board) & friendly(board)))

    # Single pawn advances
    dests = (pwns >> 8) & empty(board) & ~RANK_1
    n += count(dests & targets)

    # Double pawn advances
    doubledests = (dests >> 8) & targets & RANK_5
    n += count(doubledests)
    return n
end

function count_free_pawn_captures(board::Board, targets::Bitboard)
    n = 0
    if board.turn == WHITE
        n += count_free_wpawn_captures(board, targets)
    else
        n += count_free_bpawn_captures(board, targets)
    end
    return n
end

function count_free_wpawn_captures(board::Board, targets::Bitboard)
    n = 0
    pwns = pawns(board) & friendly(board) & ~pinned(board)

    # "Left" pawn captures
    dests = (pwns << 9) & targets & ~FILE_H
    n += count(dests & ~RANK_8)

    n += count(dests & RANK_8) * 4

    # "Right" pawn captures
    dests = (pwns << 7) & targets & ~FILE_A
    n += count(dests & ~RANK_8)

    n += count(dests & RANK_8) * 4
    return n
end

function count_free_bpawn_captures(board::Board, targets::Bitboard)
    n = 0
    pwns = pawns(board) & friendly(board) & ~pinned(board)

    # "Left" pawn captures
    dests = (pwns >> 7) & targets & ~FILE_H
    n += count(dests & ~RANK_1)

    n += count(dests & RANK_1) * 4

    # "Right" pawn captures

    dests = (pwns >> 9) & targets & ~FILE_A
    n += count(dests & ~RANK_1)

    n += count(dests & RANK_1) * 4
    return n
end

function count_enpass_moves(board::Board)
    n = 0
    ep_sqr = board.enpass
    if ep_sqr !== zero(UInt8)
        king = square(kings(board) & friendly(board))
        for pwn in (pawnAttacks(!board.turn, ep_sqr) & pawns(board) & friendly(board))
            occ = occupied(board)
            occ &= ~Bitboard(pwn) & ~Bitboard(ep_sqr - 24 + (board.turn.val << 4))
            occ |= Bitboard(ep_sqr)
            if isempty(bishopMoves(king, occ) & bishoplike(board) & enemy(board)) && isempty(rookMoves(king, occ) & rooklike(board) & enemy(board))
                n += 1
            end
        end
    end
    return n
end

function count_pawn_advance_promos(board::Board, targets::Bitboard)
    n = 0
    if board.turn == WHITE
        n += count_wpawn_advance_promos(board, targets)
    else
        n += count_bpawn_advance_promos(board, targets)
    end
    return n
end

function count_wpawn_advance_promos(board::Board, targets::Bitboard)
    n = 0
    pwns = pawns(board) & friendly(board) & (~pinned(board) | file(kings(board) & friendly(board)))
    dests = (pwns << 8) & targets & RANK_8
    n += count(dests) * 4
    return n
end

function count_bpawn_advance_promos(board::Board, targets::Bitboard)
    n = 0
    pwns = pawns(board) & friendly(board) & (~pinned(board) | file(kings(board) & friendly(board)))
    dests = (pwns >> 8) & targets & RANK_1
    n += count(dests) * 4
    return n
end

function count_pinned_pawn_captures_and_promos(board::Board, targets::Bitboard)
    n = 0
    pwns = pawns(board) & friendly(board) & pinned(board)
    king = square(kings(board) & friendly(board))
    for pwn in pwns
        for move_to in (pawnAttacks(board.turn, pwn) & enemy(board) & targets)
            if !isempty(Bitboard(pwn) & blockers(king, move_to))
                if !isempty(Bitboard(move_to) & RANK_18)
                    n += 4
                else
                    n += 1
                end
            end
        end
    end
    return n
end

function count_knight_moves(board::Board, targets::Bitboard)
    n = 0
    for knight in (knights(board) & friendly(board) & ~pinned(board))
        n += count(knightMoves(knight) & targets)
    end
    return n
end

function count_bishop_moves(board::Board, targets::Bitboard)
    n = 0
    n += count_free_bishop_moves(board, targets)
    n += count_pinned_bishop_moves(board, targets)
    return n
end

function count_free_bishop_moves(board::Board, targets::Bitboard)
    n = 0
    occ = occupied(board)
    for bishop in (bishoplike(board) & friendly(board) & ~pinned(board))
        n += count(bishopMoves(bishop, occ) & targets)
    end
    return n
end

function count_pinned_bishop_moves(board::Board, targets::Bitboard)
    n = 0
    king = square(kings(board) & friendly(board))
    occ = occupied(board)
    for bishop in (bishoplike(board) & friendly(board) & pinned(board))
        block_1 = blockers(king, bishop)
        for move_to in (bishopMoves(bishop, occ) & targets)
            if !isempty(Bitboard(move_to) & block_1) || !isempty(Bitboard(bishop) & blockers(king, move_to))
                n += 1
            end
        end
    end
    return n
end

function count_rook_moves(board::Board, targets::Bitboard)
    n = 0
    n += count_free_rook_moves(board, targets)
    n += count_pinned_rook_moves(board, targets)
    return n
end

function count_free_rook_moves(board::Board, targets::Bitboard)
    n = 0
    occ = occupied(board)
    for rook in (rooklike(board) & friendly(board) & ~pinned(board))
        n += count(rookMoves(rook, occ) & targets)
    end
    return n
end

function count_pinned_rook_moves(board::Board, targets::Bitboard)
    n = 0
    king = square(kings(board) & friendly(board))
    occ = occupied(board)
    for rook in (rooklike(board) & friendly(board) & pinned(board))
        block_1 = blockers(king, rook)
        for move_to in (rookMoves(rook, occ) & targets)
            if !isempty(Bitboard(move_to) & block_1) || !isempty(Bitboard(rook) & blockers(king, move_to))
                n += 1
            end
        end
    end
    return n
end

function count_king_moves(board::Board, targets::Bitboard)
    n = 0
    king = square(kings(board) & friendly(board))
    for move_to in (kingMoves(king) & targets)
        if !isattacked_through_king(board, move_to)
            n += 1
        end
    end
    return n
end

function count_castling(board::Board)
    n = 0
    if cancastlekingside(board)
        n += count_kingside_castling(board)
    end
    if cancastlequeenside(board)
        n += count_queenside_castling(board)
    end
    return n
end

function count_kingside_castling(board::Board)
    n = 0
    king = square(kings(board) & friendly(board))
    rook = king - 3
    if isempty(occupied(board) & blockers(king, rook)) && !isattacked(board, king - 1) && !isattacked(board, king - 2)
        n += 1
    end
    return n
end

function count_queenside_castling(board)
    n = 0
    king = square(kings(board) & friendly(board))
    rook = king + 4
    if isempty(occupied(board) & blockers(king, rook)) && !isattacked(board, king + 1) && !isattacked(board, king + 2)
        n += 1
    end
    return n
end

function count_quiet_moves(board::Board)
    n = 0
    empty_bb = empty(board)
    if isdoublecheck(board)
        n += count_king_moves(board, empty_bb)
        return n
    end
    if ischeck(board)
        if isempty(checkers(board) & knights(board))
            targets = blockers(square(kings(board) & friendly(board)), square(checkers(board)))
        else
            # if a knight gives check, only *quiet* move is for king to step away
            n += count_king_moves(board, empty_bb)
            return n
        end
    else
        n += count_castling(board)
        targets = empty_bb
    end
    n += count_pawn_advances(board, targets)
    n += count_king_moves(board, empty_bb)
    n += count_knight_moves(board, targets)
    n += count_bishop_moves(board, targets)
    n += count_rook_moves(board, targets)
    return n
end

function count_noisy_moves(board::Board)
    n = 0
    enemies = enemy(board)
    if isdoublecheck(board)
        n += count_king_moves(board, enemies)
        return n
    end
    if ischeck(board)
        targets = checkers(board)
        n += count_pawn_advance_promos(board, blockers(square(kings(board) & friendly(board)), square(targets)))
    else
        targets = enemies
        n += count_pawn_advance_promos(board, empty(board))
    end
    n += count_free_pawn_captures(board, targets)
    n += count_pinned_pawn_captures_and_promos(board, targets)
    n += count_enpass_moves(board)
    n += count_king_moves(board, enemies)
    n += count_knight_moves(board, targets)
    n += count_bishop_moves(board, targets)
    n += count_rook_moves(board, targets)
    return n
end

function count_moves(board::Board)
    n = 0
    n += count_quiet_moves(board)
    n += count_noisy_moves(board)
    return n
end


struct ReactionPath
    forward::Bool
    spcsinds::Array{Int64,1}
    rxninds::Array{Int64,1}
    spcind::Int64
    branchfracts::Array{Float64,1}
    branchfract::Array{Float64,1}
    branchind::Array{Int64,1}
end
export ReactionPath

"""
Take a single step following the fluxes in the network from a single species
checks whether the flux is moving the forward/reverse direction
based on production/loss < 1/steptol for forward and
production/loss > steptol for reverse direction and stops accordingly
"""
function stepnetwork(sim,statespcind,ropp,ropl,rts;forward=true,i=0,steptol=1e-2)
    @views ratio = sum(ropp[:,statespcind])/sum(ropl[:,statespcind])
    if forward
        if ratio > 1.0/steptol
            return (0,[0,])
        elseif i == 0
            @views rxnind = argmax(ropl[:,statespcind])
        else
            @views rxnind = reverse(sortperm(ropl[:,statespcind]))[i]
        end
        rt = rts[rxnind]
        if rt > 0.0
            spcinds = sim.reactions[rxnind].productinds
        else
            spcinds = sim.reactions[rxnind].reactantinds
        end
        return rxnind,spcinds
    else
        if ratio < steptol
            return (0,[0,])
        elseif i == 0
            @views rxnind = argmax(ropp[:,statespcind])
        else
            @views rxnind = reverse(sortperm(ropp[:,statespcind]))[i]
        end
        rt = rts[rxnind]
        if rt > 0.0
            spcinds = sim.reactions[rxnind].reactantinds
        else
            spcinds = sim.reactions[rxnind].productinds
        end
        return rxnind,spcinds
    end
end

"""
Starting from a given species indicated by spcind attempt to follow the fluxes
in the forward or reverse direction to find a reaction indicated by rxnind
it will follow in all directions until the flux tracked reaches branchtol fraction
of the original flux, the flux stops (based on steptol) or it finds the reaction
if it finds the reaction it will then follow the flux past that reaction using
a flux pairs algorithm until the flux stops
"""
function follow(sim,rxnind,spcind,ropp,ropl,rts,forward;steptol=1e-2,branchtol=5e-2)

    rps = [ReactionPath(forward,[spcind],Array{Int64,1}(),spcind,Array{Float64,1}(),
                    Array{Float64,1}([1.0]),Array{Int64,1}())]
    boo = true
    boo2 = true
    i = 0
    rpouts = Array{ReactionPath,1}()
    newrps = Array{ReactionPath,1}()
    while boo #extend off the center species
        newrps = extendpath(rps[1],sim,ropp,ropl,rts,i=i,steptol=steptol)
        if length(newrps) == 0
            return ReactionPath(forward,Array{Int64,1}(),Array{Int64,1}(),spcind,Array{Float64,1}(),
                    Array{Float64,1}(),Array{Int64,1}())
        elseif newrps[1].rxninds[1] == rxnind
            rpouts = newrps
            boo2 = false
            break
        end
        if newrps[1].branchfract[1] < branchtol
            boo = false
        else
            append!(rps,newrps)
            i += 1
        end
    end
    deleteat!(rps,1)

    while boo2
        ind = argmax([rp.branchfract[1] for rp in rps])
        rp = rps[ind]
        if rp.branchfract[1] < branchtol
            return ReactionPath(forward,Array{Int64,1}(),Array{Int64,1}(),spcind,Array{Float64,1}(),
                    Array{Float64,1}(),Array{Int64,1})
        end
        boo3 = true
        i = 0
        while boo3 #extend off the center species
            newrps = extendpath(rp,sim,ropp,ropl,rts,i=i,steptol=steptol)
            if length(newrps) == 0
                deleteat!(rps,ind)
                break
            end
            if newrps[1].rxninds[end] == rxnind
                rpouts = newrps
                boo2 = false
                break
            end
            if newrps[1].branchfract[1] < branchtol
                deleteat!(rps,ind)
                boo3 = false
            else
                deleteat!(rps,ind)
                append!(rps,newrps)
                i += 1
            end
        end
        if length(rps) == 0
            return ReactionPath(forward,Array{Int64,1}(),Array{Int64,1}(),spcind,Array{Float64,1}(),
                    Array{Float64,1}(),Array{Int64,1}())
        end
    end

    rpout = newrps[argmax([getsimilarity(sim.species[nrp.spcsinds[end-1]],sim.species[nrp.spcsinds[end]]) for nrp in rpouts])]
    while true
        newrps = extendpath(rpout,sim,ropp,ropl,rts,i=0,steptol=steptol)
        if length(newrps) == 0 || newrps[1].branchfract[1] < branchtol
            return rpout
        else
            rpout = newrps[argmax([getsimilarity(sim.species[rpout.spcsinds[end]],
                            sim.domain.phase.species[nrp.spcsinds[end]]) for nrp in newrps])]
        end
    end
end
export follow

"""
expand a ReactionPath object out a single step
"""
function extendpath(rp,sim,ropp,ropl,rts;i=0,steptol=1e-2)
    rxnind,spcinds = stepnetwork(sim,rp.spcsinds[end],ropp,ropl,rts,forward=rp.forward,i=i,steptol=steptol)
    if rxnind == 0
        sind = rp.spcsinds[end]
        name = sim.names[sind]
        forward = rp.forward
        return Array{ReactionPath,1}()
    else
        rps = Array{ReactionPath,1}()
        for spcind in spcinds
            rpnew = deepcopy(rp)
            push!(rpnew.rxninds,rxnind)
            push!(rpnew.spcsinds,spcind)
            push!(rpnew.branchind,i)
            if rpnew.forward
                push!(rpnew.branchfracts,ropl[rxnind,rp.spcsinds[end]]/sum(ropl[:,rp.spcsinds[end]]))
            else
                push!(rpnew.branchfracts,ropp[rxnind,rp.spcsinds[end]]/sum(ropp[:,rp.spcsinds[end]]))
            end
            rpnew.branchfract[1] = rp.branchfract[1]*rpnew.branchfracts[end]
            push!(rps,rpnew)
        end
        return rps
    end
end

struct Branching
    spcind::Int64
    rxninds::Array{Int64,1}
    branchingratios::Array{Float64,1}
end
export Branching

"""
Identifies calculates and stores the reactions and
branching ratios associated with the "reactants" for the input reaction
defined by rxnind as defined by the direction the reaction is occuring in
"""
function getbranching(sim,rxnind,ropl,rts)
    rt = rts[rxnind]
    if rt > 0
        spcs = sim.reactions[rxnind].reactants
    else
        spcs = sim.reactions[rxnind].products
    end
    spcsinds = [findfirst(isequal(spc.name),sim.names) for spc in spcs]
    boospcs = falses(length(spcsinds))
    branches = Array{Branching,1}()
    for (i,spc) in enumerate(spcs)
        ind = findfirst(isequal(spc.name),sim.names)
        @views rinds = reverse(sortperm(ropl[:,ind]))
        @views rvals = (ropl[:,ind]/sum(ropl[:,ind]))[rinds]
        indend = findfirst(isequal(0.0),rvals)
        @views rinds = rinds[1:indend]
        @views rvals = rvals[1:indend]
        push!(branches,Branching(ind,rinds,rvals))
    end
    return branches
end
export getbranching

struct ReactionAnalysis
    branchings::Array{Branching,1}
    paths::Array{ReactionPath,1}
    radprodlossfract::Float64
    spcind::Int64
    spcname::String
    rxnind::Int64
    sens::Float64
end
export ReactionAnalysis

"""
Calculate Transitory Sensitivities for analysis
"""
function gettransitoryadjoint(sim,t,spcname,spcind,transitorysensitivitymethod)
    if applicable(transitorysensitivitymethod,sim,t,spcname)
        return transitorysensitivitymethod(sim,t,spcname)
    else
        dSdt = transitorysensitivitymethod(sim,t)
        return dSdt[spcind+sim.domain.indexes[1]-1,:]
    end
end

"""
At a given timepoint t targeting a species spcname run transitory sensitivity
analysis to identify important reactions for that species and then run
reaction path, branching, radical rop and other analyses to provide the information
necessary to identify why the reaction is important and flesh out important
pathways in the mechanism
Returns a list of ReactionAnalysis objects associated with each reaction
"""
function analyzespc(sim,spcname,t;N=10,tol=1e-3,branchthreshold=0.9,
        pathbranchthreshold=0.2,branchtol=1e-2,steptol=1e-2,
        transitorysensitivitymethod=transitorysensitivitiesfulltrapezoidal,
        eliminate=true
        )

    spcind = findfirst(isequal(spcname),sim.names)
    @assert spcind != nothing "$spcname not found"
    dSdt = gettransitoryadjoint(sim,t,spcname,spcind,transitorysensitivitymethod)

    rop = rops(sim,t)
    ropp = zeros(size(rop))
    for i in eachindex(rop)
        if rop[i] > 0
            ropp[i] = rop[i]
        end
    end
    ropl = zeros(size(rop))
    for i in eachindex(rop)
        if rop[i] < 0
            ropl[i] = abs(rop[i])
        end
    end
    rts = rates(sim,t)

    @views dSdtspc = dSdt[length(sim.domain.phase.species)+sim.domain.parameterindexes[1]:sim.domain.parameterindexes[2]]

    #find sensitive reactions
    inds = reverse(sortperm(abs.(dSdtspc)))
    dSdtmax = maximum(abs.(dSdtspc))
    maxthresh = dSdtmax*tol

    if N == 0
        N = length(inds)
    elseif N > length(inds)
        N = length(inds)
    end
    inds = inds[1:N]
    mval = abs(dSdtspc[inds[1]])
    minval = mval*tol
    k = 1
    while k < length(inds) && abs(dSdtspc[inds[k]]) >= minval
        k += 1
    end
    sensinds = inds[1:k]

    #Run analyses
    rxnanalysis = Array{ReactionAnalysis,1}()
    for rxnind in sensinds
        branches,rps = getbranchpathinfo(sim,spcind,rxnind,ropp,ropl,rts;steptol=steptol,branchtol=branchtol)
        radprodlossfract = getradprodlossfract(sim,rxnind,rts)
        if eliminate
            branches,rps = eliminatereasons(spcind,rxnind,branches,rps,dSdtspc;branchthreshold=branchthreshold,pathbranchthreshold=pathbranchthreshold)
        end
        push!(rxnanalysis,ReactionAnalysis(branches,rps,radprodlossfract,spcind,spcname,rxnind,dSdtspc[rxnind]))
    end
    return rxnanalysis
end
export analyzespc

"""
Calculate the fraction of the radical production or loss
attributable to the given reaction
"""
function getradprodlossfract(sim,rxnind,rts)
    radrops = rts.*getfield.(sim.reactions,:radicalchange)
    radprod = sum(r for r in radrops if r>0.0)
    radloss = sum(r for r in radrops if r<0.0)
    radflux = radrops[rxnind]
    if radflux > 0.0
        return radflux/radprod
    else
        return radflux/radloss
    end
end
export getradprodlossfract

"""
Identify species coupled to the target species whose pathways could
be important to the target species even if their pathways don't overlap
If the species has non-unimolecular loss reactions with branching greater
than branchtol the other species involved are considered coupled
"""
function getcoupledspecies(sim,spcind,rxnind,ropp,ropl,rts;branchtol=0.2)
    spcsinds = Array{Int64,1}()
    @views totflux = sum(ropl[:,spcind])
    @views inds = reverse(sortperm(ropl[:,spcind]))
    i = 1
    ind = inds[i]
    for ind in inds
        if ropl[ind,spcind]/totflux < branchtol
            break
        elseif length(sim.reactions[ind].reactants) > 1
            for rind in sim.reactions[ind].reactantinds
                if rind != spcind && rind > 0 && !(rind in spcsinds)
                    push!(spcsinds,rind)
                end
            end
        end
    end
    return spcsinds
end

"""
Run branching and pathway analysis for a give species and reaction
"""
function getbranchpathinfo(sim,spcind,rxnind,ropp,ropl,rts;steptol=1e-2,branchtol=5e-2,couplingbranchtol=0.2)
    branches = getbranching(sim,rxnind,ropl,rts)
    targetinds = [spcind]
    append!(targetinds,getcoupledspecies(sim,spcind,rxnind,ropp,ropl,rts;branchtol=couplingbranchtol))
    rps = Array{ReactionPath,1}()
    for spind in targetinds
        for forward in [true,false]
            rp = follow(sim,rxnind,spcind,ropp,ropl,rts,forward;steptol=steptol,branchtol=branchtol)
            if rp.branchfract != 0 && length(rp.spcsinds) > 0
                push!(rps,rp)
            end
        end
    end
    return branches,rps
end

"""
Attempt to remove importance hypotheses that don't make sense
if the reaction constitutes more than branchthreshold fraction of the flux we
consider it to not be sensitive due to branching
If the sensitivity is >0 increasing the rate coefficient increases concentration
so it can't be rate limited downstream so forward reaction paths don't make sense
Same for sensitivity <0 for reverse paths
"""
function eliminatereasons(spcind,rxnind,branches,rps,dSdtspc;branchthreshold=0.9,pathbranchthreshold=0.2)
    branchesout = Array{Branching,1}()
    for branch in branches
        ind = findfirst(isequal(rxnind),branch.rxninds)
        if ind !== nothing
            if branch.branchingratios[ind] < branchthreshold
                push!(branchesout,branch)
            end
        end
    end
    rpouts = Array{ReactionPath,1}()
    for rp in rps
        if rp.forward && dSdtspc[rxnind] > 0 #increasing forward (loss) paths should decrease spc concentration
            continue
        elseif !rp.forward && dSdtspc[rxnind] < 0 #increasing reverse (production) paths should increase spc concentration
            continue
        else
            rind = findfirst(isequal(rxnind),rp.rxninds)
            if rp.branchfracts[rind] > pathbranchthreshold
                push!(rpouts,rp)
            end
        end
    end
    return branchesout,rpouts
end
export eliminatereasons

"""
Generate a string report from the analysis objects
"""
function getrxnanalysisstring(sim,ra;branchingcutoff=1e-2,radbranchfract=0.01)
    spcname = sim.names[ra.spcind]
    rstr = getrxnstr(sim.domain.phase.reactions[ra.rxnind])
    sens = round(ra.sens,sigdigits=6)
    s = "Analyzing $spcname sensitivity to $rstr at a value of $sens \n"
    s *= "\n"
    for branch in ra.branchings
        sname = sim.names[branch.spcind]
        s *= "Key branching for $sname \n"
        for i = 1:length(branch.rxninds)+1
            br = round(branch.branchingratios[i],sigdigits=6)
            if br < branchingcutoff
                break
            else
                rstr = getrxnstr(sim.reactions[branch.rxninds[i]])
                s *= "$rstr had a branching ratio of $br \n"
            end
        end
        s *= "\n"
    end

    for rp in ra.paths
        spname = sim.names[rp.spcind]
        forward = rp.forward
        if rp.forward
            s *= "Associated key reaction path in $spname loss direction \n"
            for i = 1:length(rp.rxninds)
                rstr = getrxnstr(sim.reactions[rp.rxninds[i]])
                br = round(rp.branchfracts[i],sigdigits=6)
                s *= "$rstr at path branching of $br \n"
            end
        else
            s *= "Associated key reaction path in $spname production direction \n"
            revinds = reverse(rp.rxninds)
            for i = 1:length(rp.rxninds)
                rstr = getrxnstr(sim.reactions[revinds[i]])
                br = round(rp.branchfracts[i],sigdigits=6)
                s *= "$rstr at path step branching of $br \n"
            end
        end
        s *= "\n"
    end
    if abs(ra.radprodlossfract) > radbranchfract
        radfract = round(abs(ra.radprodlossfract),sigdigits=6)
        if ra.radprodlossfract > 0
            s *= "Reaction accounts for $radfract of the net radical production \n"
        else
            s *= "Reaction accounts for $radfract of the net radical loss \n"
        end
    end
    s *= "\n"
    return s
end
export getrxnanalysisstring

"""
Print out a string report from the analysis objects
"""
function printrxnanalysis(sim,ra;branchingcutoff=1e-2,radbranchfract=0.01)
    return println(getrxnanalysisstring(sim,ra;branchingcutoff=branchingcutoff,radbranchfract=radbranchfract))
end
export printrxnanalysis