using AccurateArithmetic
using StaticArrays

"""
    compute(::ComputeTaskABC_P, data::ABCParticleValue)

Return the particle and value as is. 

0 FLOP.
"""
function compute(::ComputeTaskABC_P, data::ABCParticleValue{P})::ABCParticleValue{P} where {P <: ABCParticle}
    return data
end

"""
    compute(::ComputeTaskABC_U, data::ABCParticleValue)

Compute an outer edge. Return the particle value with the same particle and the value multiplied by an ABC_outer_edge factor.

1 FLOP.
"""
function compute(::ComputeTaskABC_U, data::ABCParticleValue{P})::ABCParticleValue{P} where {P <: ABCParticle}
    return ABCParticleValue{P}(data.p, data.v * ABC_outer_edge(data.p))
end

"""
    compute(::ComputeTaskABC_V, data1::ABCParticleValue, data2::ABCParticleValue)

Compute a vertex. Preserve momentum and particle types (AB->C etc.) to create resulting particle, multiply values together and times a vertex factor.

6 FLOP.
"""
function compute(
    ::ComputeTaskABC_V,
    data1::ABCParticleValue{P1},
    data2::ABCParticleValue{P2},
)::ABCParticleValue where {P1 <: ABCParticle, P2 <: ABCParticle}
    p3 = ABC_conserve_momentum(data1.p, data2.p)
    dataOut = ABCParticleValue{typeof(p3)}(p3, data1.v * ABC_vertex() * data2.v)
    return dataOut
end

"""
    compute(::ComputeTaskABC_S2, data1::ABCParticleValue, data2::ABCParticleValue)

Compute a final inner edge (2 input particles, no output particle).

For valid inputs, both input particles should have the same momenta at this point.

12 FLOP.
"""
function compute(
    ::ComputeTaskABC_S2,
    data1::ParticleValue{P},
    data2::ParticleValue{P},
)::Float64 where {P <: ABCParticle}
    #=
    @assert isapprox(abs(data1.p.momentum.E), abs(data2.p.momentum.E), rtol = 0.001, atol = sqrt(eps())) "E: $(data1.p.momentum.E) vs. $(data2.p.momentum.E)"
    @assert isapprox(data1.p.momentum.px, -data2.p.momentum.px, rtol = 0.001, atol = sqrt(eps())) "px: $(data1.p.momentum.px) vs. $(data2.p.momentum.px)"
    @assert isapprox(data1.p.momentum.py, -data2.p.momentum.py, rtol = 0.001, atol = sqrt(eps())) "py: $(data1.p.momentum.py) vs. $(data2.p.momentum.py)"
    @assert isapprox(data1.p.momentum.pz, -data2.p.momentum.pz, rtol = 0.001, atol = sqrt(eps())) "pz: $(data1.p.momentum.pz) vs. $(data2.p.momentum.pz)"
    =#
    inner = ABC_inner_edge(data1.p)
    return data1.v * inner * data2.v
end

"""
    compute(::ComputeTaskABC_S1, data::ABCParticleValue)

Compute inner edge (1 input particle, 1 output particle).

11 FLOP.
"""
function compute(::ComputeTaskABC_S1, data::ABCParticleValue{P})::ABCParticleValue{P} where {P <: ABCParticle}
    return ABCParticleValue{P}(data.p, data.v * ABC_inner_edge(data.p))
end

"""
    compute(::ComputeTaskABC_Sum, data...)
    compute(::ComputeTaskABC_Sum, data::AbstractArray)

Compute a sum over the vector. Use an algorithm that accounts for accumulated errors in long sums with potentially large differences in magnitude of the summands.

Linearly many FLOP with growing data.
"""
function compute(::ComputeTaskABC_Sum, data...)::Float64
    s = 0.0im
    for d in data
        s += d
    end
    return s
end

function compute(::ComputeTaskABC_Sum, data::AbstractArray)::Float64
    s = 0.0im
    for d in data
        s += d
    end
    return s
end
