"""
    iterates inside one partition
"""
mutable struct FeynmanPartitionIterator
    topology::TopologyPartition
    forest1_iterator::Union{ForestIterator,Missing}
    forest2_iterator::Union{ForestIterator,Missing}
    forest3_iterator::Union{ForestIterator,Missing}
end

mutable struct FeynmanDiagramTopologyIterator
    leaves::Int
    partitions_state::Vector{Int64}
    partition::FeynmanPartitionIterator
end

function FeynmanPartitionIterator(topology::TopologyPartition)
    return FeynmanPartitionIterator(topology, ForestIterator(topology.leaves1), ForestIterator(topology.leaves2), ForestIterator(topology.leaves3))
end

function FeynmanDiagramTopologyIterator(def::FeynmanDiagramDefinition)
    @assert def.n >= 4 "A Feynman diagram must have at least 4 legs"
    (p, _) = iterate(partitions(def.n))
    while length(p) != 3 || !is_valid(TopologyPartition(p[1], p[2], p[3]))
        (p, _) = iterate(partitions(def.n), p)
    end
    return FeynmanDiagramTopologyIterator(def.n, p, FeynmanPartitionIterator(TopologyPartition(p[1], p[2], p[3])))
end

function FeynmanDiagramTopology(iterator::FeynmanPartitionIterator)
    return FeynmanDiagramTopology(tree(iterator.forest1_iterator), tree(iterator.forest2_iterator), tree(iterator.forest3_iterator))
end

function FeynmanDiagramTopology(iterator::FeynmanDiagramTopologyIterator)
    return FeynmanDiagramTopology(iterator.partition)
end

function _iterate(iterator::FeynmanPartitionIterator)
    # todo: need symmetry breaking here
    if !is_end(iterator.forest3_iterator)
        _iterate(iterator.forest3_iterator)
    elseif !is_end(iterator.forest2_iterator)
        iterator.forest3_iterator = ForestIterator(iterator.topology.leaves3)
        _iterate(iterator.forest2_iterator)
    else
        iterator.forest3_iterator = ForestIterator(iterator.topology.leaves3)
        iterator.forest2_iterator = ForestIterator(iterator.topology.leaves2)
        _iterate(iterator.forest1_iterator)
    end
    return nothing
end

function _iterate(iterator::FeynmanDiagramTopologyIterator)
    if !is_end(iterator.partition)
        _iterate(iterator.partition)
    else
        while true
            ret = iterate(partitions(iterator.leaves), iterator.partitions_state)

            if (isnothing(ret))
                return false
            end

            (p, _) = ret
            iterator.partitions_state = p
            if length(p) == 3 && is_valid(TopologyPartition(p[1], p[2], p[3]))
                iterator.partition = FeynmanPartitionIterator(TopologyPartition(p[1], p[2], p[3]))
                break
            end
        end
    end
    return true
end

function is_end(iterator::FeynmanPartitionIterator)
    return is_end(iterator.forest1_iterator) && is_end(iterator.forest2_iterator) && is_end(iterator.forest3_iterator)
end

function iterate(def::FeynmanDiagramDefinition)
    state = FeynmanDiagramTopologyIterator(def)

    return (FeynmanDiagramTopology(state), state)
end

function iterate(def::FeynmanDiagramDefinition, state::FeynmanDiagramTopologyIterator)
    if !_iterate(state)
        return nothing
    end

    return (FeynmanDiagramTopology(state), state)
end