
ComputeTaskQED_Sum() = ComputeTaskQED_Sum(0)

function _svector_from_type(processDescription::QEDProcessDescription, type, particles)
    if haskey(in_particles(processDescription), type)
        return SVector{in_particles(processDescription)[type], type}(filter(x -> typeof(x) <: type, particles))
    end
    if haskey(out_particles(processDescription), type)
        return SVector{out_particles(processDescription)[type], type}(filter(x -> typeof(x) <: type, particles))
    end
    return SVector{0, type}()
end

"""
    gen_process_input(processDescription::QEDProcessDescription)

Return a ProcessInput of randomly generated [`QEDParticle`](@ref)s from a [`QEDProcessDescription`](@ref). The process description can be created manually or parsed from a string using [`parse_process`](@ref).

Note: This uses RAMBO to create a valid process with conservation of momentum and energy.
"""
function gen_process_input(processDescription::QEDProcessDescription)
    massSum = 0
    inputMasses = Vector{Float64}()
    for (particle, n) in processDescription.inParticles
        for _ in 1:n
            massSum += mass(particle)
            push!(inputMasses, mass(particle))
        end
    end
    outputMasses = Vector{Float64}()
    for (particle, n) in processDescription.outParticles
        for _ in 1:n
            massSum += mass(particle)
            push!(outputMasses, mass(particle))
        end
    end

    # add some extra random mass to allow for some momentum
    massSum += rand(rng[threadid()]) * (length(inputMasses) + length(outputMasses))


    particles = Vector{QEDParticle}()
    initialMomenta = generate_initial_moms(massSum, inputMasses)
    index = 1
    for (particle, n) in processDescription.inParticles
        for _ in 1:n
            mom = initialMomenta[index]
            push!(particles, particle(mom))
            index += 1
        end
    end

    final_momenta = generate_physical_massive_moms(rng[threadid()], massSum, outputMasses)
    index = 1
    for (particle, n) in processDescription.outParticles
        for _ in 1:n
            push!(particles, particle(final_momenta[index]))
            index += 1
        end
    end

    inFerms = _svector_from_type(processDescription, FermionStateful{Incoming, SpinUp}, particles)
    outFerms = _svector_from_type(processDescription, FermionStateful{Outgoing, SpinUp}, particles)
    inAntiferms = _svector_from_type(processDescription, AntiFermionStateful{Incoming, SpinUp}, particles)
    outAntiferms = _svector_from_type(processDescription, AntiFermionStateful{Outgoing, SpinUp}, particles)
    inPhotons = _svector_from_type(processDescription, PhotonStateful{Incoming, PolX}, particles)
    outPhotons = _svector_from_type(processDescription, PhotonStateful{Outgoing, PolX}, particles)

    processInput =
        QEDProcessInput(processDescription, inFerms, outFerms, inAntiferms, outAntiferms, inPhotons, outPhotons)

    return processInput
end

"""
    gen_graph(process_description::QEDProcessDescription)

For a given [`QEDProcessDescription`](@ref), return the [`DAG`](@ref) that computes it.
"""
function gen_graph(process_description::QEDProcessDescription)
    initial_diagram = FeynmanDiagram(process_description)
    diagrams = gen_diagrams(initial_diagram)

    graph = DAG()

    COMPLEX_SIZE = sizeof(ComplexF64)
    PARTICLE_VALUE_SIZE = 96.0

    # TODO: Not all diagram outputs should always be summed at the end, if they differ by fermion exchange they need to be diffed
    # Should not matter for n-Photon Compton processes though
    sum_node = insert_node!(graph, make_node(ComputeTaskQED_Sum(0)), track = false, invalidate_cache = false)
    global_data_out = insert_node!(graph, make_node(DataTask(COMPLEX_SIZE)), track = false, invalidate_cache = false)
    insert_edge!(graph, sum_node, global_data_out, track = false, invalidate_cache = false)

    # remember the data out nodes for connection
    dataOutNodes = Dict()

    for particle in initial_diagram.particles
        # generate data in and U tasks
        data_in = insert_node!(
            graph,
            make_node(DataTask(PARTICLE_VALUE_SIZE), String(particle)),
            track = false,
            invalidate_cache = false,
        ) # read particle data node
        compute_u = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false, invalidate_cache = false) # compute U node
        data_out =
            insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)), track = false, invalidate_cache = false) # transfer data out from u (one ABCParticleValue object)

        insert_edge!(graph, data_in, compute_u, track = false, invalidate_cache = false)
        insert_edge!(graph, compute_u, data_out, track = false, invalidate_cache = false)

        # remember the data_out node for future edges
        dataOutNodes[String(particle)] = data_out
    end

    # TODO: this should be parallelizable somewhat easily
    for diagram in diagrams
        tie = diagram.tie[]

        # handle the vertices
        for vertices in diagram.vertices
            for vertex in vertices
                data_in1 = dataOutNodes[String(vertex.in1)]
                data_in2 = dataOutNodes[String(vertex.in2)]

                compute_V = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false, invalidate_cache = false) # compute vertex

                insert_edge!(graph, data_in1, compute_V, track = false, invalidate_cache = false)
                insert_edge!(graph, data_in2, compute_V, track = false, invalidate_cache = false)

                data_V_out = insert_node!(
                    graph,
                    make_node(DataTask(PARTICLE_VALUE_SIZE)),
                    track = false,
                    invalidate_cache = false,
                )

                insert_edge!(graph, compute_V, data_V_out, track = false, invalidate_cache = false)

                if (vertex.out == tie.in1 || vertex.out == tie.in2)
                    # out particle is part of the tie -> there will be an S2 task with it later, don't make S1 task
                    dataOutNodes[String(vertex.out)] = data_V_out
                    continue
                end

                # otherwise, add S1 task
                compute_S1 =
                    insert_node!(graph, make_node(ComputeTaskQED_S1()), track = false, invalidate_cache = false) # compute propagator

                insert_edge!(graph, data_V_out, compute_S1, track = false, invalidate_cache = false)

                data_S1_out = insert_node!(
                    graph,
                    make_node(DataTask(PARTICLE_VALUE_SIZE)),
                    track = false,
                    invalidate_cache = false,
                )

                insert_edge!(graph, compute_S1, data_S1_out, track = false, invalidate_cache = false)

                # overrides potentially different nodes from previous diagrams, which is intentional
                dataOutNodes[String(vertex.out)] = data_S1_out
            end
        end

        # handle the tie
        data_in1 = dataOutNodes[String(tie.in1)]
        data_in2 = dataOutNodes[String(tie.in2)]

        compute_S2 = insert_node!(graph, make_node(ComputeTaskQED_S2()), track = false, invalidate_cache = false)

        data_S2 = insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)), track = false, invalidate_cache = false)

        insert_edge!(graph, data_in1, compute_S2, track = false, invalidate_cache = false)
        insert_edge!(graph, data_in2, compute_S2, track = false, invalidate_cache = false)

        insert_edge!(graph, compute_S2, data_S2, track = false, invalidate_cache = false)

        insert_edge!(graph, data_S2, sum_node, track = false, invalidate_cache = false)
        add_child!(task(sum_node))
    end

    return graph
end
