using MetagraphOptimization
using QEDbase
using AccurateArithmetic
using Random

import MetagraphOptimization.ABCParticle
import MetagraphOptimization.interaction_result

const RTOL = sqrt(eps(Float64))
RNG = Random.default_rng()

function check_particle_reverse_moment(p1::SFourMomentum, p2::SFourMomentum)
    @test isapprox(abs(p1.E), abs(p2.E))
    @test isapprox(p1.px, -p2.px)
    @test isapprox(p1.py, -p2.py)
    @test isapprox(p1.pz, -p2.pz)
    return nothing
end

function ground_truth_graph_result(input::ABCProcessInput)
    # formula for one diagram:
    # u_Bp * iλ * u_Ap * S_C * u_B * iλ * u_A
    # for the second diagram:
    # u_B * iλ * u_Ap * S_C * u_Bp * iλ * u_Ap
    # the "u"s are all 1, we ignore the i, λ is 1/137.

    constant = (1 / 137.0)^2

    # calculate particle C in diagram 1
    diagram1_C = ParticleC(input.inParticles[1].momentum + input.inParticles[2].momentum)
    diagram2_C = ParticleC(input.inParticles[1].momentum + input.outParticles[2].momentum)

    diagram1_Cp = ParticleC(input.outParticles[1].momentum + input.outParticles[2].momentum)
    diagram2_Cp = ParticleC(input.outParticles[1].momentum + input.inParticles[2].momentum)

    check_particle_reverse_moment(diagram1_Cp.momentum, diagram1_C.momentum)
    check_particle_reverse_moment(diagram2_Cp.momentum, diagram2_C.momentum)
    @test isapprox(getMass2(diagram1_C.momentum), getMass2(diagram1_Cp.momentum))
    @test isapprox(getMass2(diagram2_C.momentum), getMass2(diagram2_Cp.momentum))

    inner1 = MetagraphOptimization.ABC_inner_edge(diagram1_C)
    inner2 = MetagraphOptimization.ABC_inner_edge(diagram2_C)

    diagram1_result = inner1 * constant
    diagram2_result = inner2 * constant

    return sum_kbn([diagram1_result, diagram2_result])
end

machine = get_machine_info()

process_2_2 = ABCProcessDescription(
    Dict{Type, Int64}(ParticleA => 1, ParticleB => 1),
    Dict{Type, Int64}(ParticleA => 1, ParticleB => 1),
)

particles_2_2 = ABCProcessInput(
    process_2_2,
    ABCParticle[
        ParticleA(SFourMomentum(0.823648, 0.0, 0.0, 0.823648)),
        ParticleB(SFourMomentum(0.823648, 0.0, 0.0, -0.823648)),
    ],
    ABCParticle[
        ParticleA(SFourMomentum(0.823648, -0.835061, -0.474802, 0.277915)),
        ParticleB(SFourMomentum(0.823648, 0.835061, 0.474802, -0.277915)),
    ],
)
expected_result = ground_truth_graph_result(particles_2_2)

@testset "AB->AB no optimization" begin
    for _ in 1:10   # test in a loop because graph layout should not change the result
        graph = parse_dag(joinpath(@__DIR__, "..", "input", "AB->AB.txt"), ABCModel())
        @test isapprox(execute(graph, process_2_2, machine, particles_2_2), expected_result; rtol = RTOL)

        # graph should be fully scheduled after being executed
        @test is_scheduled(graph)

        func = get_compute_function(graph, process_2_2, machine)
        @test isapprox(func(particles_2_2), expected_result; rtol = RTOL)
    end
end

@testset "AB->AB after random walk" begin
    for i in 1:200
        graph = parse_dag(joinpath(@__DIR__, "..", "input", "AB->AB.txt"), ABCModel())
        optimize!(RandomWalkOptimizer(RNG), graph, 50)

        @test is_valid(graph)

        @test isapprox(execute(graph, process_2_2, machine, particles_2_2), expected_result; rtol = RTOL)

        # graph should be fully scheduled after being executed
        @test is_scheduled(graph)
    end
end

process_2_4 = ABCProcessDescription(
    Dict{Type, Int64}(ParticleA => 1, ParticleB => 1),
    Dict{Type, Int64}(ParticleA => 1, ParticleB => 3),
)
particles_2_4 = gen_process_input(process_2_4)
graph = parse_dag(joinpath(@__DIR__, "..", "input", "AB->ABBB.txt"), ABCModel())
expected_result = execute(graph, process_2_4, machine, particles_2_4)

@testset "AB->ABBB no optimization" begin
    for _ in 1:5   # test in a loop because graph layout should not change the result
        graph = parse_dag(joinpath(@__DIR__, "..", "input", "AB->ABBB.txt"), ABCModel())
        @test isapprox(execute(graph, process_2_4, machine, particles_2_4), expected_result; rtol = RTOL)

        func = get_compute_function(graph, process_2_4, machine)
        @test isapprox(func(particles_2_4), expected_result; rtol = RTOL)
    end
end

@testset "AB->ABBB after random walk" begin
    for i in 1:50
        graph = parse_dag(joinpath(@__DIR__, "..", "input", "AB->ABBB.txt"), ABCModel())
        optimize!(RandomWalkOptimizer(RNG), graph, 100)
        @test is_valid(graph)

        @test isapprox(execute(graph, process_2_4, machine, particles_2_4), expected_result; rtol = RTOL)
    end
end

@testset "AB->AB large sum fusion" begin
    for _ in 1:20
        graph = parse_dag(joinpath(@__DIR__, "..", "input", "AB->AB.txt"), ABCModel())

        # push a fusion with the sum node
        ops = get_operations(graph)
        for fusion in ops.nodeFusions
            if isa(fusion.input[3].task, ComputeTaskABC_Sum)
                push_operation!(graph, fusion)
                break
            end
        end

        # push two more fusions with the fused node
        for _ in 1:15
            ops = get_operations(graph)
            for fusion in ops.nodeFusions
                if isa(fusion.input[3].task, FusedComputeTask)
                    push_operation!(graph, fusion)
                    break
                end
            end
        end

        # try execute
        @test is_valid(graph)
        expected_result = ground_truth_graph_result(particles_2_2)
        @test isapprox(execute(graph, process_2_2, machine, particles_2_2), expected_result; rtol = RTOL)
    end
end


@testset "AB->AB large sum fusion" begin
    for _ in 1:20
        graph = parse_dag(joinpath(@__DIR__, "..", "input", "AB->AB.txt"), ABCModel())

        # push a fusion with the sum node
        ops = get_operations(graph)
        for fusion in ops.nodeFusions
            if isa(fusion.input[3].task, ComputeTaskABC_Sum)
                push_operation!(graph, fusion)
                break
            end
        end

        # push two more fusions with the fused node
        for _ in 1:15
            ops = get_operations(graph)
            for fusion in ops.nodeFusions
                if isa(fusion.input[3].task, FusedComputeTask)
                    push_operation!(graph, fusion)
                    break
                end
            end
        end

        # try execute
        @test is_valid(graph)
        expected_result = ground_truth_graph_result(particles_2_2)
        @test isapprox(execute(graph, process_2_2, machine, particles_2_2), expected_result; rtol = RTOL)
    end
end

@testset "AB->AB fusion edge case" begin
    for _ in 1:20
        graph = parse_dag(joinpath(@__DIR__, "..", "input", "AB->AB.txt"), ABCModel())

        # push two fusions with ComputeTaskABC_V
        for _ in 1:2
            ops = get_operations(graph)
            for fusion in ops.nodeFusions
                if isa(fusion.input[1].task, ComputeTaskABC_V)
                    push_operation!(graph, fusion)
                    break
                end
            end
        end

        # push fusions until the end
        cont = true
        while cont
            cont = false
            ops = get_operations(graph)
            for fusion in ops.nodeFusions
                if isa(fusion.input[1].task, FusedComputeTask)
                    push_operation!(graph, fusion)
                    cont = true
                    break
                end
            end
        end

        # try execute
        @test is_valid(graph)
        expected_result = ground_truth_graph_result(particles_2_2)
        @test isapprox(execute(graph, process_2_2, machine, particles_2_2), expected_result; rtol = RTOL)
    end
end
