using StaticArrays

"""
    compute(::ComputeTaskQED_P, data::QEDParticleValue)

Return the particle as is and initialize the Value.
"""
function compute(::ComputeTaskQED_P, data::QEDParticleValue{P}) where {P <: QEDParticle}
    # TODO do we actually need this for anything?
    return ParticleValue{P, DiracMatrix}(data.p, one(DiracMatrix))
end

"""
    compute(::ComputeTaskQED_U, data::QEDParticleValue)

Compute an outer edge. Return the particle value with the same particle and the value multiplied by an outer_edge factor.
"""
function compute(::ComputeTaskQED_U, data::PV) where {P <: QEDParticle, PV <: QEDParticleValue{P}}
    part::P = data.p
    state = base_state(particle(part), direction(part), momentum(part), spin_or_pol(part))
    return ParticleValue{P, typeof(state)}(
        data.p,
        state, # will return a SLorentzVector{ComplexF64}, BiSpinor or AdjointBiSpinor
    )
end

"""
    compute(::ComputeTaskQED_V, data1::QEDParticleValue, data2::QEDParticleValue)

Compute a vertex. Preserve momentum and particle types (e + gamma->p etc.) to create resulting particle, multiply values together and times a vertex factor.
"""
function compute(
    ::ComputeTaskQED_V,
    data1::PV1,
    data2::PV2,
) where {P1 <: QEDParticle, P2 <: QEDParticle, PV1 <: QEDParticleValue{P1}, PV2 <: QEDParticleValue{P2}}
    p3 = QED_conserve_momentum(data1.p, data2.p)
    P3 = interaction_result(P1, P2)
    state = QED_vertex()
    if (typeof(data1.v) <: AdjointBiSpinor)
        state = data1.v * state
    else
        state = state * data1.v
    end
    if (typeof(data2.v) <: AdjointBiSpinor)
        state = data2.v * state
    else
        state = state * data2.v
    end

    dataOut = ParticleValue{P3, typeof(state)}(P3(momentum(p3)), state)
    return dataOut
end

"""
    compute(::ComputeTaskQED_S2, data1::QEDParticleValue, data2::QEDParticleValue)

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(
    ::ComputeTaskQED_S2,
    data1::ParticleValue{P1},
    data2::ParticleValue{P2},
) where {P1 <: Union{AntiFermionStateful, FermionStateful}, P2 <: Union{AntiFermionStateful, FermionStateful}}
    #@assert isapprox(data1.p.momentum, data2.p.momentum, rtol = sqrt(eps()), atol = sqrt(eps())) "$(data1.p.momentum) vs. $(data2.p.momentum)"

    inner = QED_inner_edge(propagation_result(P1)(momentum(data1.p)))

    # inner edge is just a "scalar", data1 and data2 are bispinor/adjointbispinnor, need to keep correct order
    if typeof(data1.v) <: BiSpinor
        return (data2.v)::AdjointBiSpinor * inner * (data1.v)::BiSpinor
    else
        return (data1.v)::AdjointBiSpinor * inner * (data2.v)::BiSpinor
    end
end

function compute(
    ::ComputeTaskQED_S2,
    data1::ParticleValue{P1},
    data2::ParticleValue{P2},
) where {P1 <: PhotonStateful, P2 <: PhotonStateful}
    # TODO: assert that data1 and data2 are opposites
    inner = QED_inner_edge(data1.p)
    # inner edge is just a scalar, data1 and data2 are photon states that are just Complex numbers here
    return data1.v * inner * data2.v
end

"""
    compute(::ComputeTaskQED_S1, data::QEDParticleValue)

Compute inner edge (1 input particle, 1 output particle).
"""
function compute(::ComputeTaskQED_S1, data::QEDParticleValue{P}) where {P <: QEDParticle}
    newP = propagation_result(P)
    new_p = newP(momentum(data.p))
    # inner edge is just a scalar, can multiply from either side
    if typeof(data.v) <: BiSpinor
        return ParticleValue(new_p, QED_inner_edge(new_p) * data.v)
    else
        return ParticleValue(new_p, data.v * QED_inner_edge(new_p))
    end
end

"""
    compute(::ComputeTaskQED_Sum, data...)
    compute(::ComputeTaskQED_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(::ComputeTaskQED_Sum, data...)::ComplexF64
    # TODO: want to use sum_kbn here but it doesn't seem to support ComplexF64, do it element-wise?
    s = 0.0im
    for d in data
        s += d
    end
    return s
end

function compute(::ComputeTaskQED_Sum, data::AbstractArray)::ComplexF64
    # TODO: want to use sum_kbn here but it doesn't seem to support ComplexF64, do it element-wise?
    s = 0.0im
    for d in data
        s += d
    end
    return s
end
