import QEDbase.mass

"""
    QEDModel <: AbstractPhysicsModel

Singleton definition for identification of the QED-Model.
"""
struct QEDModel <: AbstractPhysicsModel end

"""
    QEDParticle

Base type for all particles in the [`QEDModel`](@ref).

Its template parameter specifies the particle's direction.

The concrete types contain singletons of the types that they are, like `Photon` and `Electron` from QEDbase, and their state descriptions.
"""
abstract type QEDParticle{Direction <: ParticleDirection} <: AbstractParticle end

"""
    QEDProcessDescription <: AbstractProcessDescription

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

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

"""
    QEDProcessInput <: AbstractProcessInput

Input for a QED Process. Contains the [`QEDProcessDescription`](@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 QEDProcessInput <: AbstractProcessInput
    process::QEDProcessDescription
    inParticles::Vector{QEDParticle}
    outParticles::Vector{QEDParticle}
end

QEDParticleValue{ParticleType <: QEDParticle} =
    ParticleValue{ParticleType, <:Union{BiSpinor, AdjointBiSpinor, DiracMatrix, SLorentzVector{ComplexF64}, ComplexF64}}

"""
    PhotonStateful <: QEDParticle

A photon of the [`QEDModel`](@ref) with its state.
"""
struct PhotonStateful{Direction <: ParticleDirection} <: QEDParticle{Direction}
    momentum::SFourMomentum
    # this will maybe change to the full polarization vector? or do i need both
    polarization::AbstractPolarization
end

PhotonStateful{Direction}(mom::SFourMomentum) where {Direction <: ParticleDirection} =
    PhotonStateful{Direction}(mom, AllPolarization())

PhotonStateful{Dir1}(ph::PhotonStateful{Dir2}) where {Dir1 <: ParticleDirection, Dir2 <: ParticleDirection} =
    PhotonStateful{Dir1}(ph.momentum, ph.polarization)

"""
    FermionStateful <: QEDParticle

A fermion of the [`QEDModel`](@ref) with its state.
"""
struct FermionStateful{Direction <: ParticleDirection} <: QEDParticle{Direction}
    momentum::SFourMomentum
    spin::AbstractSpin
end

FermionStateful{Direction}(mom::SFourMomentum) where {Direction <: ParticleDirection} =
    FermionStateful{Direction}(mom, AllSpin())

FermionStateful{Dir1}(f::FermionStateful{Dir2}) where {Dir1 <: ParticleDirection, Dir2 <: ParticleDirection} =
    FermionStateful{Dir1}(f.momentum, f.spin)

"""
    AntiFermionStateful <: QEDParticle

An anti-fermion of the [`QEDModel`](@ref) with its state.
"""
struct AntiFermionStateful{Direction <: ParticleDirection} <: QEDParticle{Direction}
    momentum::SFourMomentum
    spin::AbstractSpin
end

AntiFermionStateful{Direction}(mom::SFourMomentum) where {Direction <: ParticleDirection} =
    AntiFermionStateful{Direction}(mom, AllSpin())

AntiFermionStateful{Dir1}(f::AntiFermionStateful{Dir2}) where {Dir1 <: ParticleDirection, Dir2 <: ParticleDirection} =
    AntiFermionStateful{Dir1}(f.momentum, f.spin)

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

For two given particle types that can interact, return the third.
"""
function interaction_result(t1::Type{T1}, t2::Type{T2}) where {T1 <: QEDParticle, T2 <: QEDParticle}
    @assert false "Invalid interaction between particles of types $t1 and $t2"
end

interaction_result(::Type{FermionStateful{Incoming}}, ::Type{FermionStateful{Outgoing}}) = PhotonStateful{Incoming}
interaction_result(::Type{FermionStateful{Incoming}}, ::Type{AntiFermionStateful{Incoming}}) = PhotonStateful{Incoming}
interaction_result(::Type{FermionStateful{Incoming}}, ::Type{<:PhotonStateful}) = FermionStateful{Outgoing}

interaction_result(::Type{FermionStateful{Outgoing}}, ::Type{FermionStateful{Incoming}}) = PhotonStateful{Incoming}
interaction_result(::Type{FermionStateful{Outgoing}}, ::Type{AntiFermionStateful{Outgoing}}) = PhotonStateful{Incoming}
interaction_result(::Type{FermionStateful{Outgoing}}, ::Type{<:PhotonStateful}) = FermionStateful{Incoming}

# antifermion mirror
interaction_result(::Type{AntiFermionStateful{Incoming}}, t2::Type{<:QEDParticle}) =
    interaction_result(FermionStateful{Outgoing}, t2)
interaction_result(::Type{AntiFermionStateful{Outgoing}}, t2::Type{<:QEDParticle}) =
    interaction_result(FermionStateful{Incoming}, t2)

# photon commutativity
interaction_result(t1::Type{<:PhotonStateful}, t2::Type{<:QEDParticle}) = interaction_result(t2, t1)

# but prevent stack overflow
function interaction_result(t1::Type{<:PhotonStateful}, t2::Type{<:PhotonStateful})
    @assert false "Invalid interaction between particles of types $t1 and $t2"
end

"""
    propagation_result(t1::Type{T}) where {T <: QEDParticle}

Return the type of the inverted direction. E.g.
"""
propagation_result(::Type{FermionStateful{Incoming}}) = FermionStateful{Outgoing}
propagation_result(::Type{FermionStateful{Outgoing}}) = FermionStateful{Incoming}
propagation_result(::Type{AntiFermionStateful{Incoming}}) = AntiFermionStateful{Outgoing}
propagation_result(::Type{AntiFermionStateful{Outgoing}}) = AntiFermionStateful{Incoming}
propagation_result(::Type{PhotonStateful{Incoming}}) = PhotonStateful{Outgoing}
propagation_result(::Type{PhotonStateful{Outgoing}}) = PhotonStateful{Incoming}

"""
    types(::QEDModel)

Return a Vector of the possible types of particle in the [`QEDModel`](@ref).
"""
function types(::QEDModel)
    return [PhotonStateful, FermionStateful, AntiFermionStateful]
end

# type piracy?
String(::Type{Incoming}) = "Incoming"
String(::Type{Outgoing}) = "Outgoing"

function String(::Type{PhotonStateful{Dir}}) where {Dir <: ParticleDirection}
    return "Photon" * String(Dir)
end
function String(::Type{FermionStateful{Dir}}) where {Dir <: ParticleDirection}
    return "Fermion" * String(Dir)
end
function String(::Type{AntiFermionStateful{Dir}}) where {Dir <: ParticleDirection}
    return "Anti-Fermion" * String(Dir)
end

@inline particle(::PhotonStateful) = Photon()
@inline particle(::FermionStateful) = Fermion()
@inline particle(::AntiFermionStateful) = AntiFermion()

@inline momentum(p::PhotonStateful)::SFourMomentum{ComplexF64} = p.momentum
@inline momentum(p::FermionStateful)::SFourMomentum{ComplexF64} = p.momentum
@inline momentum(p::AntiFermionStateful)::SFourMomentum{ComplexF64} = p.momentum

@inline spin_or_pol(p::PhotonStateful)::AbstractPolarization = p.polarization
@inline spin_or_pol(p::FermionStateful)::AbstractSpin = p.spin
@inline spin_or_pol(p::AntiFermionStateful)::AbstractSpin = p.spin

@inline direction(::PhotonStateful{Dir}) where {Dir <: ParticleDirection} = Dir()
@inline direction(::FermionStateful{Dir}) where {Dir <: ParticleDirection} = Dir()
@inline direction(::AntiFermionStateful{Dir}) where {Dir <: ParticleDirection} = Dir()

"""
    caninteract(T1::Type{<:QEDParticle}, T2::Type{<:QEDParticle})

For two given [`QEDParticle`](@ref) types, return whether they can interact at a vertex. This is equivalent to `!issame(T1, T2)`.

See also: [`issame`](@ref) and [`interaction_result`](@ref)
"""
function caninteract(T1::Type{<:QEDParticle}, T2::Type{<:QEDParticle})
    if (T1 == T2)
        return false
    end
    if (T1 <: PhotonStateful && T2 <: PhotonStateful)
        return false
    end

    for (P1, P2) in [(T1, T2), (T2, T1)]
        if (P1 == FermionStateful{Incoming} && P2 == AntiFermionStateful{Outgoing})
            return false
        end
        if (P1 == FermionStateful{Outgoing} && P2 == AntiFermionStateful{Incoming})
            return false
        end
    end

    return true
end

"""
    issame(T1::Type{<:QEDParticle}, T2::Type{<:QEDParticle})

For two given [`QEDParticle`](@ref) types, return whether they are equivalent for the purpose of a Feynman Diagram. That means e.g. an `Incoming` `AntiFermion` is the same as an `Outgoing` `Fermion`. This is equivalent to `!caninteract(T1, T2)`.

See also: [`caninteract`](@ref) and [`interaction_result`](@ref)
"""
function issame(T1::Type{<:QEDParticle}, T2::Type{<:QEDParticle})
    return !caninteract(T1, T2)
end

"""
    QED_vertex()

Return the factor of a vertex in a QED feynman diagram.
"""
@inline function QED_vertex()::SLorentzVector{DiracMatrix}
    return ComplexF64(0, 1) * gamma()
end

"""
    QED_preserve_momentum(p1::QEDParticle, p2::QEDParticle)

Calculate and return a new particle from two given interacting ones at a vertex.
"""
function QED_preserve_momentum(p1::QEDParticle, p2::QEDParticle)
    T3 = interaction_result(typeof(p1), typeof(p2))
    # TODO: probably also need to do something about the spin/pol
    p3 = T3(p1.momentum + p2.momentum)
    return p3
end
