using QEDbase
using Random
using Roots
using ForwardDiff

ComputeTaskABC_Sum() = ComputeTaskABC_Sum(0)

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

function _svector_from_type_out(processDescription::ABCProcessDescription, type, particles)
    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::ABCProcessDescription)

Return a ProcessInput of randomly generated [`ABCParticle`](@ref)s from a [`ABCProcessDescription`](@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::ABCProcessDescription)
    inParticleTypes = keys(processDescription.inParticles)
    outParticleTypes = keys(processDescription.outParticles)

    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))


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

    outputParticles = Vector{ABCParticle}()
    final_momenta = generate_physical_massive_moms(rng[threadid()], massSum, outputMasses)
    index = 1
    for (particle, n) in processDescription.outParticles
        for _ in 1:n
            mom = final_momenta[index]
            push!(outputParticles, particle(SFourMomentum(-mom.E, mom.px, mom.py, mom.pz)))
            index += 1
        end
    end

    inA = _svector_from_type_in(processDescription, ParticleA, inputParticles)
    inB = _svector_from_type_in(processDescription, ParticleB, inputParticles)
    inC = _svector_from_type_in(processDescription, ParticleC, inputParticles)

    outA = _svector_from_type_out(processDescription, ParticleA, outputParticles)
    outB = _svector_from_type_out(processDescription, ParticleB, outputParticles)
    outC = _svector_from_type_out(processDescription, ParticleC, outputParticles)

    processInput = ABCProcessInput(processDescription, inA, inB, inC, outA, outB, outC)

    return return processInput
end
