using MetagraphOptimization
using QEDbase
using QEDprocesses
using StatsBase # for countmap
using Random
using UUIDs

import MetagraphOptimization.caninteract
import MetagraphOptimization.issame
import MetagraphOptimization.interaction_result
import MetagraphOptimization.propagation_result
import MetagraphOptimization.direction
import MetagraphOptimization.spin_or_pol
import MetagraphOptimization.QED_vertex

def_momentum = SFourMomentum(1.0, 0.0, 0.0, 0.0)

RNG = Random.default_rng()

testparticleTypes = [
    PhotonStateful{Incoming, PolX},
    PhotonStateful{Outgoing, PolX},
    FermionStateful{Incoming, SpinUp},
    FermionStateful{Outgoing, SpinUp},
    AntiFermionStateful{Incoming, SpinUp},
    AntiFermionStateful{Outgoing, SpinUp},
]

testparticleTypesPropagated = [
    PhotonStateful{Outgoing, PolX},
    PhotonStateful{Incoming, PolX},
    FermionStateful{Outgoing, SpinUp},
    FermionStateful{Incoming, SpinUp},
    AntiFermionStateful{Outgoing, SpinUp},
    AntiFermionStateful{Incoming, SpinUp},
]

function compton_groundtruth(input::QEDProcessInput)
    # p1k1 -> p2k2
    # formula: −(ie)^2 (u(p2) slashed(ε1) S(p2 − k1) slashed(ε2) u(p1) + u(p2) slashed(ε2) S(p1 + k1) slashed(ε1) u(p1))

    p1 = input.inFerms[1]
    p2 = input.outFerms[1]

    k1 = input.inPhotons[1]
    k2 = input.outPhotons[1]

    u_p1 = base_state(Electron(), Incoming(), p1.momentum, spin_or_pol(p1))
    u_p2 = base_state(Electron(), Outgoing(), p2.momentum, spin_or_pol(p2))

    eps_1 = base_state(Photon(), Incoming(), k1.momentum, spin_or_pol(k1))
    eps_2 = base_state(Photon(), Outgoing(), k2.momentum, spin_or_pol(k2))

    virt1_mom = p2.momentum - k1.momentum
    @test isapprox(p1.momentum - k2.momentum, virt1_mom)

    virt2_mom = p1.momentum + k1.momentum
    @test isapprox(p2.momentum + k2.momentum, virt2_mom)

    s_p2_k1 = propagator(Electron(), virt1_mom)
    s_p1_k1 = propagator(Electron(), virt2_mom)

    diagram1 = u_p2 * (eps_1 * QED_vertex()) * s_p2_k1 * (eps_2 * QED_vertex()) * u_p1
    diagram2 = u_p2 * (eps_2 * QED_vertex()) * s_p1_k1 * (eps_1 * QED_vertex()) * u_p1

    return diagram1 + diagram2
end


@testset "Interaction Result" begin
    import MetagraphOptimization.QED_conserve_momentum

    for p1 in testparticleTypes, p2 in testparticleTypes
        if !caninteract(p1, p2)
            @test_throws AssertionError interaction_result(p1, p2)
            continue
        end

        @test interaction_result(p1, p2) in setdiff(testparticleTypes, [p1, p2])
        @test issame(interaction_result(p1, p2), interaction_result(p2, p1))

        testParticle1 = p1(rand(RNG, SFourMomentum))
        testParticle2 = p2(rand(RNG, SFourMomentum))
        p3 = interaction_result(p1, p2)

        resultParticle = QED_conserve_momentum(testParticle1, testParticle2)

        @test issame(typeof(resultParticle), interaction_result(p1, p2))

        totalMom = zero(SFourMomentum)
        for (p, mom) in [(p1, testParticle1.momentum), (p2, testParticle2.momentum), (p3, resultParticle.momentum)]
            if (typeof(direction(p)) <: Incoming)
                totalMom += mom
            else
                totalMom -= mom
            end
        end

        @test isapprox(totalMom, zero(SFourMomentum); atol = sqrt(eps()))
    end
end

@testset "Propagation Result" begin
    for (p, propResult) in zip(testparticleTypes, testparticleTypesPropagated)
        @test issame(propagation_result(p), propResult)
        @test direction(propagation_result(p)(def_momentum)) != direction(p(def_momentum))
    end
end

@testset "Parse Process" begin
    @testset "Order invariance" begin
        @test parse_process("ke->ke", QEDModel()) == parse_process("ek->ke", QEDModel())
        @test parse_process("ke->ke", QEDModel()) == parse_process("ek->ek", QEDModel())
        @test parse_process("ke->ke", QEDModel()) == parse_process("ke->ek", QEDModel())

        @test parse_process("kkke->eep", QEDModel()) == parse_process("kkek->epe", QEDModel())
    end

    @testset "Known processes" begin
        compton_process = QEDProcessDescription(
            Dict{Type, Int}(PhotonStateful{Incoming, PolX} => 1, FermionStateful{Incoming, SpinUp} => 1),
            Dict{Type, Int}(PhotonStateful{Outgoing, PolX} => 1, FermionStateful{Outgoing, SpinUp} => 1),
        )

        @test parse_process("ke->ke", QEDModel()) == compton_process

        positron_compton_process = QEDProcessDescription(
            Dict{Type, Int}(PhotonStateful{Incoming, PolX} => 1, AntiFermionStateful{Incoming, SpinUp} => 1),
            Dict{Type, Int}(PhotonStateful{Outgoing, PolX} => 1, AntiFermionStateful{Outgoing, SpinUp} => 1),
        )

        @test parse_process("kp->kp", QEDModel()) == positron_compton_process

        trident_process = QEDProcessDescription(
            Dict{Type, Int}(PhotonStateful{Incoming, PolX} => 1, FermionStateful{Incoming, SpinUp} => 1),
            Dict{Type, Int}(FermionStateful{Outgoing, SpinUp} => 2, AntiFermionStateful{Outgoing, SpinUp} => 1),
        )

        @test parse_process("ke->eep", QEDModel()) == trident_process

        pair_production_process = QEDProcessDescription(
            Dict{Type, Int}(PhotonStateful{Incoming, PolX} => 2),
            Dict{Type, Int}(FermionStateful{Outgoing, SpinUp} => 1, AntiFermionStateful{Outgoing, SpinUp} => 1),
        )

        @test parse_process("kk->pe", QEDModel()) == pair_production_process

        pair_annihilation_process = QEDProcessDescription(
            Dict{Type, Int}(FermionStateful{Incoming, SpinUp} => 1, AntiFermionStateful{Incoming, SpinUp} => 1),
            Dict{Type, Int}(PhotonStateful{Outgoing, PolX} => 2),
        )

        @test parse_process("pe->kk", QEDModel()) == pair_annihilation_process
    end
end

@testset "Generate Process Inputs" begin
    @testset "Process $proc_str" for proc_str in ["ke->ke", "kp->kp", "kk->ep", "ep->kk"]
        # currently can only generate for 2->2 processes
        process = parse_process(proc_str, QEDModel())

        for i in 1:100
            input = gen_process_input(process)
            @test length(input.inFerms) == get(process.inParticles, FermionStateful{Incoming, SpinUp}, 0)
            @test length(input.inAntiferms) == get(process.inParticles, AntiFermionStateful{Incoming, SpinUp}, 0)
            @test length(input.inPhotons) == get(process.inParticles, PhotonStateful{Incoming, PolX}, 0)
            @test length(input.outFerms) == get(process.outParticles, FermionStateful{Outgoing, SpinUp}, 0)
            @test length(input.outAntiferms) == get(process.outParticles, AntiFermionStateful{Outgoing, SpinUp}, 0)
            @test length(input.outPhotons) == get(process.outParticles, PhotonStateful{Outgoing, PolX}, 0)

            @test isapprox(
                sum([
                    getfield.(input.inFerms, :momentum)...,
                    getfield.(input.inAntiferms, :momentum)...,
                    getfield.(input.inPhotons, :momentum)...,
                ]),
                sum([
                    getfield.(input.outFerms, :momentum)...,
                    getfield.(input.outAntiferms, :momentum)...,
                    getfield.(input.outPhotons, :momentum)...,
                ]);
                atol = sqrt(eps()),
            )
        end
    end
end

@testset "Compton" begin
    import MetagraphOptimization.insert_node!
    import MetagraphOptimization.insert_edge!
    import MetagraphOptimization.make_node

    model = QEDModel()
    process = parse_process("ke->ke", model)
    machine = Machine(
        [
            MetagraphOptimization.NumaNode(
                0,
                1,
                MetagraphOptimization.default_strategy(MetagraphOptimization.NumaNode),
                -1.0,
                UUIDs.uuid1(),
            ),
        ],
        [-1.0;;],
    )

    graph = MetagraphOptimization.DAG()

    # manually build a graph for compton
    graph = DAG()

    # s to output (exit node)
    d_exit = insert_node!(graph, make_node(DataTask(16)), track = false)

    sum_node = insert_node!(graph, make_node(ComputeTaskQED_Sum(2)), track = false)

    d_s0_sum = insert_node!(graph, make_node(DataTask(16)), track = false)
    d_s1_sum = insert_node!(graph, make_node(DataTask(16)), track = false)

    # final s compute
    s0 = insert_node!(graph, make_node(ComputeTaskQED_S2()), track = false)
    s1 = insert_node!(graph, make_node(ComputeTaskQED_S2()), track = false)

    # data from v0 and v1 to s0
    d_v0_s0 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_v1_s0 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_v2_s1 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_v3_s1 = insert_node!(graph, make_node(DataTask(96)), track = false)

    # v0 and v1 compute
    v0 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false)
    v1 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false)
    v2 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false)
    v3 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false)

    # data from uPhIn, uPhOut, uElIn, uElOut to v0 and v1
    d_uPhIn_v0 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_uElIn_v0 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_uPhOut_v1 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_uElOut_v1 = insert_node!(graph, make_node(DataTask(96)), track = false)

    # data from uPhIn, uPhOut, uElIn, uElOut to v2 and v3
    d_uPhOut_v2 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_uElIn_v2 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_uPhIn_v3 = insert_node!(graph, make_node(DataTask(96)), track = false)
    d_uElOut_v3 = insert_node!(graph, make_node(DataTask(96)), track = false)

    # uPhIn, uPhOut, uElIn and uElOut computes
    uPhIn = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false)
    uPhOut = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false)
    uElIn = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false)
    uElOut = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false)

    # data into U
    d_uPhIn = insert_node!(graph, make_node(DataTask(16), "ki1"), track = false)
    d_uPhOut = insert_node!(graph, make_node(DataTask(16), "ko1"), track = false)
    d_uElIn = insert_node!(graph, make_node(DataTask(16), "ei1"), track = false)
    d_uElOut = insert_node!(graph, make_node(DataTask(16), "eo1"), track = false)

    # now for all the edges
    insert_edge!(graph, d_uPhIn, uPhIn, track = false)
    insert_edge!(graph, d_uPhOut, uPhOut, track = false)
    insert_edge!(graph, d_uElIn, uElIn, track = false)
    insert_edge!(graph, d_uElOut, uElOut, track = false)

    insert_edge!(graph, uPhIn, d_uPhIn_v0, track = false)
    insert_edge!(graph, uPhOut, d_uPhOut_v1, track = false)
    insert_edge!(graph, uElIn, d_uElIn_v0, track = false)
    insert_edge!(graph, uElOut, d_uElOut_v1, track = false)

    insert_edge!(graph, uPhIn, d_uPhIn_v3, track = false)
    insert_edge!(graph, uPhOut, d_uPhOut_v2, track = false)
    insert_edge!(graph, uElIn, d_uElIn_v2, track = false)
    insert_edge!(graph, uElOut, d_uElOut_v3, track = false)

    insert_edge!(graph, d_uPhIn_v0, v0, track = false)
    insert_edge!(graph, d_uPhOut_v1, v1, track = false)
    insert_edge!(graph, d_uElIn_v0, v0, track = false)
    insert_edge!(graph, d_uElOut_v1, v1, track = false)

    insert_edge!(graph, d_uPhIn_v3, v3, track = false)
    insert_edge!(graph, d_uPhOut_v2, v2, track = false)
    insert_edge!(graph, d_uElIn_v2, v2, track = false)
    insert_edge!(graph, d_uElOut_v3, v3, track = false)

    insert_edge!(graph, v0, d_v0_s0, track = false)
    insert_edge!(graph, v1, d_v1_s0, track = false)
    insert_edge!(graph, v2, d_v2_s1, track = false)
    insert_edge!(graph, v3, d_v3_s1, track = false)

    insert_edge!(graph, d_v0_s0, s0, track = false)
    insert_edge!(graph, d_v1_s0, s0, track = false)

    insert_edge!(graph, d_v2_s1, s1, track = false)
    insert_edge!(graph, d_v3_s1, s1, track = false)

    insert_edge!(graph, s0, d_s0_sum, track = false)
    insert_edge!(graph, s1, d_s1_sum, track = false)

    insert_edge!(graph, d_s0_sum, sum_node, track = false)
    insert_edge!(graph, d_s1_sum, sum_node, track = false)

    insert_edge!(graph, sum_node, d_exit, track = false)

    input = [gen_process_input(process) for _ in 1:1000]

    compton_function = get_compute_function(graph, process, machine)
    @test isapprox(compton_function.(input), compton_groundtruth.(input))

    graph_generated = gen_graph(process)

    compton_function = get_compute_function(graph_generated, process, machine)
    @test isapprox(compton_function.(input), compton_groundtruth.(input))
end

@testset "Equal results after optimization" for optimizer in
                                                [ReductionOptimizer(), RandomWalkOptimizer(MersenneTwister(0))]
    @testset "Process $proc_str" for proc_str in ["ke->ke", "kp->kp", "kk->ep", "ep->kk", "ke->kke", "ke->kkke"]
        model = QEDModel()
        process = parse_process(proc_str, model)
        machine = Machine(
            [
                MetagraphOptimization.NumaNode(
                    0,
                    1,
                    MetagraphOptimization.default_strategy(MetagraphOptimization.NumaNode),
                    -1.0,
                    UUIDs.uuid1(),
                ),
            ],
            [-1.0;;],
        )
        graph = gen_graph(process)

        compute_function = get_compute_function(graph, process, machine)

        if (typeof(optimizer) <: RandomWalkOptimizer)
            optimize!(optimizer, graph, 100)
        elseif (typeof(optimizer) <: ReductionOptimizer)
            optimize_to_fixpoint!(optimizer, graph)
        end
        reduced_compute_function = get_compute_function(graph, process, machine)

        input = [gen_process_input(process) for _ in 1:100]

        @test isapprox(compute_function.(input), reduced_compute_function.(input))
    end
end
