using MetagraphOptimization

import MetagraphOptimization.gen_diagrams
import MetagraphOptimization.isincoming
import MetagraphOptimization.types


model = QEDModel()
compton = ("Compton Scattering", parse_process("ke->ke", model), 2)
compton_3 = ("3-Photon Compton Scattering", parse_process("kkke->ke", QEDModel()), 24)
compton_4 = ("4-Photon Compton Scattering", parse_process("kkkke->ke", QEDModel()), 120)
bhabha = ("Bhabha Scattering", parse_process("ep->ep", model), 2)
moller = ("Møller Scattering", parse_process("ee->ee", model), 2)
pair_production = ("Pair production", parse_process("kk->ep", model), 2)
pair_annihilation = ("Pair annihilation", parse_process("ep->kk", model), 2)
trident = ("Trident", parse_process("ke->epe", model), 8)

@testset "Known Processes" begin
    @testset "$name" for (name, process, n) in
                         [compton, bhabha, moller, pair_production, pair_annihilation, trident, compton_3, compton_4]
        initial_diagram = FeynmanDiagram(process)

        n_particles = 0
        for type in types(model)
            if (isincoming(type))
                n_particles += get(process.inParticles, type, 0)
            else
                n_particles += get(process.outParticles, type, 0)
            end
        end
        @test n_particles == length(initial_diagram.particles)
        @test ismissing(initial_diagram.tie[])
        @test isempty(initial_diagram.vertices)

        result_diagrams = gen_diagrams(initial_diagram)
        @test length(result_diagrams) == n

        for d in result_diagrams
            n_vertices = 0
            for vs in d.vertices
                n_vertices += length(vs)
            end
            @test n_vertices == n_particles - 2
            @test !ismissing(d.tie[])
        end
    end
end
