using QEDprocesses
using StaticArrays
import QEDbase.mass

# TODO check
const e = sqrt(4π / 137)

"""
    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

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

"""
    PhotonStateful <: QEDParticle

A photon of the [`QEDModel`](@ref) with its state.
"""
struct PhotonStateful{Direction <: ParticleDirection, Pol <: AbstractDefinitePolarization} <: QEDParticle{Direction}
    momentum::SFourMomentum
end

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

PhotonStateful{Dir, Pol}(ph::PhotonStateful) where {Dir, Pol} = PhotonStateful{Dir, Pol}(ph.momentum)

"""
    FermionStateful <: QEDParticle

A fermion of the [`QEDModel`](@ref) with its state.
"""
struct FermionStateful{Direction <: ParticleDirection, Spin <: AbstractDefiniteSpin} <: QEDParticle{Direction}
    momentum::SFourMomentum
    # TODO: mass for electron/muon/tauon representation?
end

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

FermionStateful{Dir, Spin}(f::FermionStateful) where {Dir, Spin} = FermionStateful{Dir, Spin}(f.momentum)

"""
    AntiFermionStateful <: QEDParticle

An anti-fermion of the [`QEDModel`](@ref) with its state.
"""
struct AntiFermionStateful{Direction <: ParticleDirection, Spin <: AbstractDefiniteSpin} <: QEDParticle{Direction}
    momentum::SFourMomentum
    # TODO: mass for electron/muon/tauon representation?
end

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

AntiFermionStateful{Dir, Spin}(f::AntiFermionStateful) where {Dir, Spin} = AntiFermionStateful{Dir, Spin}(f.momentum)

"""
    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, Spin1}},
    ::Type{FermionStateful{Outgoing, Spin2}},
) where {Spin1, Spin2} = PhotonStateful{Incoming, PolX}
interaction_result(
    ::Type{FermionStateful{Incoming, Spin1}},
    ::Type{AntiFermionStateful{Incoming, Spin2}},
) where {Spin1, Spin2} = PhotonStateful{Incoming, PolX}
interaction_result(::Type{FermionStateful{Incoming, Spin1}}, ::Type{<:PhotonStateful}) where {Spin1} =
    FermionStateful{Outgoing, SpinUp}

interaction_result(
    ::Type{FermionStateful{Outgoing, Spin1}},
    ::Type{FermionStateful{Incoming, Spin2}},
) where {Spin1, Spin2} = PhotonStateful{Incoming, PolX}
interaction_result(
    ::Type{FermionStateful{Outgoing, Spin1}},
    ::Type{AntiFermionStateful{Outgoing, Spin2}},
) where {Spin1, Spin2} = PhotonStateful{Incoming, PolX}
interaction_result(::Type{FermionStateful{Outgoing, Spin1}}, ::Type{<:PhotonStateful}) where {Spin1} =
    FermionStateful{Incoming, SpinUp}

# antifermion mirror
interaction_result(::Type{AntiFermionStateful{Incoming, Spin}}, t2::Type{<:QEDParticle}) where {Spin} =
    interaction_result(FermionStateful{Outgoing, Spin}, t2)
interaction_result(::Type{AntiFermionStateful{Outgoing, Spin}}, t2::Type{<:QEDParticle}) where {Spin} =
    interaction_result(FermionStateful{Incoming, Spin}, 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, Spin}}) where {Spin <: AbstractDefiniteSpin} =
    FermionStateful{Outgoing, Spin}
propagation_result(::Type{FermionStateful{Outgoing, Spin}}) where {Spin <: AbstractDefiniteSpin} =
    FermionStateful{Incoming, Spin}
propagation_result(::Type{AntiFermionStateful{Incoming, Spin}}) where {Spin <: AbstractDefiniteSpin} =
    AntiFermionStateful{Outgoing, Spin}
propagation_result(::Type{AntiFermionStateful{Outgoing, Spin}}) where {Spin <: AbstractDefiniteSpin} =
    AntiFermionStateful{Incoming, Spin}
propagation_result(::Type{PhotonStateful{Incoming, Pol}}) where {Pol <: AbstractDefinitePolarization} =
    PhotonStateful{Outgoing, Pol}
propagation_result(::Type{PhotonStateful{Outgoing, Pol}}) where {Pol <: AbstractDefinitePolarization} =
    PhotonStateful{Incoming, Pol}

"""
    types(::QEDModel)

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

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

String(::Type{PolX}) = "polx"
String(::Type{PolY}) = "poly"

String(::Type{SpinUp}) = "spinup"
String(::Type{SpinDown}) = "spindown"

String(::Incoming) = "i"
String(::Outgoing) = "o"

function String(::Type{<:PhotonStateful})
    return "k"
end
function String(::Type{<:FermionStateful})
    return "e"
end
function String(::Type{<:AntiFermionStateful})
    return "p"
end

function unique_name(::Type{PhotonStateful{Dir, Pol}}) where {Dir, Pol}
    return String(PhotonStateful) * String(Dir) * String(Pol)
end
function unique_name(::Type{FermionStateful{Dir, Spin}}) where {Dir, Spin}
    return String(FermionStateful) * String(Dir) * String(Spin)
end
function unique_name(::Type{AntiFermionStateful{Dir, Spin}}) where {Dir, Spin}
    return String(AntiFermionStateful) * String(Dir) * String(Spin)
end

@inline particle(::PhotonStateful) = Photon()
@inline particle(::FermionStateful) = Electron()
@inline particle(::AntiFermionStateful) = Positron()

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

@inline spin_or_pol(p::PhotonStateful{Dir, Pol}) where {Dir, Pol <: AbstractDefinitePolarization} = Pol()
@inline spin_or_pol(p::FermionStateful{Dir, Spin}) where {Dir, Spin <: AbstractDefiniteSpin} = Spin()
@inline spin_or_pol(p::AntiFermionStateful{Dir, Spin}) where {Dir, Spin <: AbstractDefiniteSpin} = Spin()

@inline direction(
    ::Type{P},
) where {P <: Union{FermionStateful{Incoming}, AntiFermionStateful{Incoming}, PhotonStateful{Incoming}}} = Incoming()
@inline direction(
    ::Type{P},
) where {P <: Union{FermionStateful{Outgoing}, AntiFermionStateful{Outgoing}, PhotonStateful{Outgoing}}} = Outgoing()

@inline direction(
    ::P,
) where {P <: Union{FermionStateful{Incoming}, AntiFermionStateful{Incoming}, PhotonStateful{Incoming}}} = Incoming()
@inline direction(
    ::P,
) where {P <: Union{FermionStateful{Outgoing}, AntiFermionStateful{Outgoing}, PhotonStateful{Outgoing}}} = Outgoing()

@inline isincoming(::QEDParticle{Incoming}) = true
@inline isincoming(::QEDParticle{Outgoing}) = false
@inline isoutgoing(::QEDParticle{Incoming}) = false
@inline isoutgoing(::QEDParticle{Outgoing}) = true

@inline isincoming(::Type{<:QEDParticle{Incoming}}) = true
@inline isincoming(::Type{<:QEDParticle{Outgoing}}) = false
@inline isoutgoing(::Type{<:QEDParticle{Incoming}}) = false
@inline isoutgoing(::Type{<:QEDParticle{Outgoing}}) = true

@inline mass(::Type{<:FermionStateful}) = 1.0
@inline mass(::Type{<:AntiFermionStateful}) = 1.0
@inline mass(::Type{<:PhotonStateful}) = 0.0

@inline invert_momentum(p::FermionStateful{Dir, Spin}) where {Dir, Spin} =
    FermionStateful{Dir, Spin}(-p.momentum, p.spin)
@inline invert_momentum(p::AntiFermionStateful{Dir, Spin}) where {Dir, Spin} =
    AntiFermionStateful{Dir, Spin}(-p.momentum, p.spin)
@inline invert_momentum(k::PhotonStateful{Dir, Spin}) where {Dir, Spin} =
    PhotonStateful{Dir, Spin}(-k.momentum, k.polarization)


"""
    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

function type_index_from_name(::QEDModel, name::String)
    if startswith(name, "ki")
        return (PhotonStateful{Incoming, PolX}, parse(Int, name[3:end]))
    elseif startswith(name, "ko")
        return (PhotonStateful{Outgoing, PolX}, parse(Int, name[3:end]))
    elseif startswith(name, "ei")
        return (FermionStateful{Incoming, SpinUp}, parse(Int, name[3:end]))
    elseif startswith(name, "eo")
        return (FermionStateful{Outgoing, SpinUp}, parse(Int, name[3:end]))
    elseif startswith(name, "pi")
        return (AntiFermionStateful{Incoming, SpinUp}, parse(Int, name[3:end]))
    elseif startswith(name, "po")
        return (AntiFermionStateful{Outgoing, SpinUp}, parse(Int, name[3:end]))
    else
        throw("Invalid name for a particle in the QED model")
    end
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}
    # Peskin-Schroeder notation
    return -1im * e * gamma()
end

@inline function QED_inner_edge(p::QEDParticle)::DiracMatrix
    return propagator(particle(p), p.momentum)
end

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

Calculate and return a new particle from two given interacting ones at a vertex.
"""
function QED_conserve_momentum(
    p1::P1,
    p2::P2,
) where {
    Dir1 <: ParticleDirection,
    Dir2 <: ParticleDirection,
    SpinPol1 <: AbstractSpinOrPolarization,
    SpinPol2 <: AbstractSpinOrPolarization,
    P1 <: Union{FermionStateful{Dir1, SpinPol1}, AntiFermionStateful{Dir1, SpinPol1}, PhotonStateful{Dir1, SpinPol1}},
    P2 <: Union{FermionStateful{Dir2, SpinPol2}, AntiFermionStateful{Dir2, SpinPol2}, PhotonStateful{Dir2, SpinPol2}},
}
    P3 = interaction_result(P1, P2)
    p1_mom = p1.momentum
    if (Dir1 <: Outgoing)
        p1_mom *= -1
    end
    p2_mom = p2.momentum
    if (Dir2 <: Outgoing)
        p2_mom *= -1
    end

    p3_mom = p1_mom + p2_mom
    if (typeof(direction(P3)) <: Incoming)
        return P3(-p3_mom)
    end
    return P3(p3_mom)
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{N1, N2, N3, N4, N5, N6} <: AbstractProcessInput
    process::QEDProcessDescription
    inFerms::SVector{N1, FermionStateful{Incoming, SpinUp}}
    outFerms::SVector{N2, FermionStateful{Outgoing, SpinUp}}
    inAntiferms::SVector{N3, AntiFermionStateful{Incoming, SpinUp}}
    outAntiferms::SVector{N4, AntiFermionStateful{Outgoing, SpinUp}}
    inPhotons::SVector{N5, PhotonStateful{Incoming, PolX}}
    outPhotons::SVector{N6, PhotonStateful{Outgoing, PolX}}
end

"""
    model(::AbstractProcessDescription)

Return the model of this process description.
"""
model(::QEDProcessDescription) = QEDModel()
model(::QEDProcessInput) = QEDModel()

function copy(process::QEDProcessDescription)
    return QEDProcessDescription(copy(process.inParticles), copy(process.outParticles))
end

==(p1::QEDProcessDescription, p2::QEDProcessDescription) =
    p1.inParticles == p2.inParticles && p1.outParticles == p2.outParticles

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

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

function get_particle(input::QEDProcessInput, t::Type{Particle}, n::Int)::Particle where {Particle}
    if (t <: FermionStateful{Incoming})
        return input.inFerms[n]
    elseif (t <: FermionStateful{Outgoing})
        return input.outFerms[n]
    elseif (t <: AntiFermionStateful{Incoming})
        return input.inAntiferms[n]
    elseif (t <: AntiFermionStateful{Outgoing})
        return input.outAntiferms[n]
    elseif (t <: PhotonStateful{Incoming})
        return input.inPhotons[n]
    elseif (t <: PhotonStateful{Outgoing})
        return input.outPhotons[n]
    end
    @assert false "Invalid type given"
end
