import QEDbase.mass

"""
    ABCModel <: AbstractPhysicsModel

Singleton definition for identification of the ABC-Model.
"""
struct ABCModel <: AbstractPhysicsModel end

"""
    ABCParticle

Base type for all particles in the [`ABCModel`](@ref).
"""
abstract type ABCParticle <: AbstractParticle end

"""
    ParticleA <: ABCParticle

An 'A' particle in the ABC Model.
"""
struct ParticleA <: ABCParticle
    momentum::SFourMomentum
end

"""
    ParticleB <: ABCParticle

A 'B' particle in the ABC Model.
"""
struct ParticleB <: ABCParticle
    momentum::SFourMomentum
end

"""
    ParticleC <: ABCParticle

A 'C' particle in the ABC Model.
"""
struct ParticleC <: ABCParticle
    momentum::SFourMomentum
end

"""
    ABCProcessDescription <: AbstractProcessDescription

A description of a process in the ABC-Model. Contains the input and output particles.

See also: [`in_particles`](@ref), [`out_particles`](@ref), [`parse_process`](@ref)
"""
struct ABCProcessDescription <: AbstractProcessDescription
    inParticles::Dict{Type, Int}
    outParticles::Dict{Type, Int}
end

"""
    ABCProcessInput <: AbstractProcessInput

Input for a ABC Process. Contains the [`ABCProcessDescription`](@ref) of the process it is an input for, and the values of the in and out particles.

See also: [`gen_process_input`](@ref)
"""
struct ABCProcessInput <: AbstractProcessInput
    process::ABCProcessDescription
    inParticles::Vector{ABCParticle}
    outParticles::Vector{ABCParticle}
end

ABCParticleValue{ParticleType <: ABCParticle} = ParticleValue{ParticleType, Float64}

"""
    PARTICLE_MASSES

A constant dictionary containing the masses of the different [`ABCParticle`](@ref)s.
"""
const PARTICLE_MASSES = Dict{Type, Float64}(ParticleA => 1.0, ParticleB => 1.0, ParticleC => 0.0)

"""
    mass(t::Type{T}) where {T <: ABCParticle}
    
Return the mass (at rest) of the given particle type.
"""
mass(t::Type{T}) where {T <: ABCParticle} = PARTICLE_MASSES[t]

"""
    interaction_result(t1::Type{T1}, t2::Type{T2}) where {T1 <: ABCParticle, T2 <: ABCParticle}

For 2 given (non-equal) particle types, return the third of ABC.
"""
function interaction_result(t1::Type{T1}, t2::Type{T2}) where {T1 <: ABCParticle, T2 <: ABCParticle}
    @assert t1 != t2
    if t1 != ParticleA && t2 != ParticleA
        return ParticleA
    elseif t1 != ParticleB && t2 != ParticleB
        return ParticleB
    else
        return ParticleC
    end
end

"""
    types(::ABCModel)

Return a Vector of the possible types of particle in the [`ABCModel`](@ref).
"""
function types(::ABCModel)
    return [ParticleA, ParticleB, ParticleC]
end

"""
    square(p::ABCParticle)

Return the square of the particle's momentum as a `Float` value.

Takes 7 effective FLOP.
"""
function square(p::ABCParticle)
    return getMass2(p.momentum)
end

"""
    ABC_inner_edge(p::ABCParticle)

Return the factor of the inner edge with the given (virtual) particle.

Takes 10 effective FLOP. (3 here + 7 in square(p))
"""
function ABC_inner_edge(p::ABCParticle)
    return 1.0 / (square(p) - mass(typeof(p)) * mass(typeof(p)))
end

"""
    ABC_outer_edge(p::ABCParticle)

Return the factor of the outer edge with the given (real) particle.

Takes 0 effective FLOP.
"""
function ABC_outer_edge(p::ABCParticle)
    return 1.0
end

"""
    ABC_vertex()

Return the factor of a vertex.

Takes 0 effective FLOP since it's constant.
"""
function ABC_vertex()
    i = 1.0
    lambda = 1.0 / 137.0
    return i * lambda
end

"""
    ABC_preserve_momentum(p1::ABCParticle, p2::ABCParticle)

Calculate and return a new particle from two given interacting ones at a vertex.

Takes 4 effective FLOP.
"""
function ABC_preserve_momentum(p1::ABCParticle, p2::ABCParticle)
    t3 = interaction_result(typeof(p1), typeof(p2))
    p3 = t3(p1.momentum + p2.momentum)
    return p3
end

"""
    type_from_name(name::String)

For a name of a particle, return the particle's [`Type`].
"""
function type_from_name(name::String)
    if startswith(name, "A")
        return ParticleA
    elseif startswith(name, "B")
        return ParticleB
    elseif startswith(name, "C")
        return ParticleC
    else
        throw("Invalid name for a particle in the ABC model")
    end
end

function String(::Type{ParticleA})
    return "A"
end
function String(::Type{ParticleB})
    return "B"
end
function String(::Type{ParticleC})
    return "C"
end

function in_particles(process::ABCProcessDescription)
    return process.inParticles
end

function in_particles(input::ABCProcessInput)
    return input.inParticles
end

function out_particles(process::ABCProcessDescription)
    return process.outParticles
end

function out_particles(input::ABCProcessInput)
    return input.outParticles
end
