using QEDbase
using Random
using Roots
using ForwardDiff

ComputeTaskABC_Sum() = ComputeTaskABC_Sum(0)

"""
    gen_process_input(processDescription::ABCProcessDescription)

Return a ProcessInput of randomly generated [`ABCParticle`](@ref)s from a [`ABCProcessDescription`](@ref). The process description can be created manually or parsed from a string using [`parse_process`](@ref).

Note: This uses RAMBO to create a valid process with conservation of momentum and energy.
"""
function gen_process_input(processDescription::ABCProcessDescription)
    inParticleTypes = keys(processDescription.inParticles)
    outParticleTypes = keys(processDescription.outParticles)

    massSum = 0
    inputMasses = Vector{Float64}()
    for (particle, n) in processDescription.inParticles
        for _ in 1:n
            massSum += mass(particle)
            push!(inputMasses, mass(particle))
        end
    end
    outputMasses = Vector{Float64}()
    for (particle, n) in processDescription.outParticles
        for _ in 1:n
            massSum += mass(particle)
            push!(outputMasses, mass(particle))
        end
    end

    # add some extra random mass to allow for some momentum
    massSum += rand(rng[threadid()]) * (length(inputMasses) + length(outputMasses))


    inputParticles = Vector{ABCParticle}()
    initialMomenta = generate_initial_moms(massSum, inputMasses)
    index = 1
    for (particle, n) in processDescription.inParticles
        for _ in 1:n
            mom = initialMomenta[index]
            push!(inputParticles, particle(mom))
            index += 1
        end
    end

    outputParticles = Vector{ABCParticle}()
    final_momenta = generate_physical_massive_moms(rng[threadid()], massSum, outputMasses)
    index = 1
    for (particle, n) in processDescription.outParticles
        for _ in 1:n
            mom = final_momenta[index]
            push!(outputParticles, particle(SFourMomentum(-mom.E, mom.px, mom.py, mom.pz)))
            index += 1
        end
    end

    processInput = ABCProcessInput(processDescription, inputParticles, outputParticles)

    return return processInput
end

####################
# CODE FROM HERE BORROWED FROM SOURCE: https://codebase.helmholtz.cloud/qedsandbox/QEDphasespaces.jl/
# use qedphasespaces directly once released
#
# quick and dirty implementation of the RAMBO algorithm
#
# reference: 
# * https://cds.cern.ch/record/164736/files/198601282.pdf
# * https://www.sciencedirect.com/science/article/pii/0010465586901190
####################

function generate_initial_moms(ss, masses)
    E1 = (ss^2 + masses[1]^2 - masses[2]^2) / (2 * ss)
    E2 = (ss^2 + masses[2]^2 - masses[1]^2) / (2 * ss)

    rho1 = sqrt(E1^2 - masses[1]^2)
    rho2 = sqrt(E2^2 - masses[2]^2)

    return [SFourMomentum(E1, 0, 0, rho1), SFourMomentum(E2, 0, 0, -rho2)]
end


Random.rand(rng::AbstractRNG, ::Random.SamplerType{SFourMomentum}) = SFourMomentum(rand(rng, 4))
Random.rand(rng::AbstractRNG, ::Random.SamplerType{NTuple{N, Float64}}) where {N} = Tuple(rand(rng, N))


function _transform_uni_to_mom(u1, u2, u3, u4)
    cth = 2 * u1 - 1
    sth = sqrt(1 - cth^2)
    phi = 2 * pi * u2
    q0 = -log(u3 * u4)
    qx = q0 * sth * cos(phi)
    qy = q0 * sth * sin(phi)
    qz = q0 * cth

    return SFourMomentum(q0, qx, qy, qz)
end

function _transform_uni_to_mom!(uni_mom, dest)
    u1, u2, u3, u4 = Tuple(uni_mom)
    cth = 2 * u1 - 1
    sth = sqrt(1 - cth^2)
    phi = 2 * pi * u2
    q0 = -log(u3 * u4)
    qx = q0 * sth * cos(phi)
    qy = q0 * sth * sin(phi)
    qz = q0 * cth

    return dest = SFourMomentum(q0, qx, qy, qz)
end

_transform_uni_to_mom(u1234::Tuple) = _transform_uni_to_mom(u1234...)
_transform_uni_to_mom(u1234::SFourMomentum) = _transform_uni_to_mom(Tuple(u1234))

function generate_massless_moms(rng, n::Int)
    a = Vector{SFourMomentum}(undef, n)
    rand!(rng, a)
    return map(_transform_uni_to_mom, a)
end

function generate_physical_massless_moms(rng, ss, n)
    r_moms = generate_massless_moms(rng, n)
    Q = sum(r_moms)
    M = sqrt(Q * Q)
    fac = -1 / M
    Qx = getX(Q)
    Qy = getY(Q)
    Qz = getZ(Q)
    bx = fac * Qx
    by = fac * Qy
    bz = fac * Qz
    gamma = getT(Q) / M
    a = 1 / (1 + gamma)
    x = ss / M

    i = 1
    while i <= n
        mom = r_moms[i]
        mom0 = getT(mom)
        mom1 = getX(mom)
        mom2 = getY(mom)
        mom3 = getZ(mom)

        bq = bx * mom1 + by * mom2 + bz * mom3

        p0 = x * (gamma * mom0 + bq)
        px = x * (mom1 + bx * mom0 + a * bq * bx)
        py = x * (mom2 + by * mom0 + a * bq * by)
        pz = x * (mom3 + bz * mom0 + a * bq * bz)

        r_moms[i] = SFourMomentum(p0, px, py, pz)
        i += 1
    end
    return r_moms
end

function _to_be_solved(xi, masses, p0s, ss)
    sum = 0.0
    for (i, E) in enumerate(p0s)
        sum += sqrt(masses[i]^2 + xi^2 * E^2)
    end
    return sum - ss
end

function _build_massive_momenta(xi, masses, massless_moms)
    vec = SFourMomentum[]
    i = 1
    while i <= length(massless_moms)
        massless_mom = massless_moms[i]
        k0 = sqrt(getT(massless_mom)^2 * xi^2 + masses[i]^2)

        kx = xi * getX(massless_mom)
        ky = xi * getY(massless_mom)
        kz = xi * getZ(massless_mom)

        push!(vec, SFourMomentum(k0, kx, ky, kz))

        i += 1
    end
    return vec
end

first_derivative(func) = x -> ForwardDiff.derivative(func, float(x))


function generate_physical_massive_moms(rng, ss, masses; x0 = 0.1)
    n = length(masses)
    massless_moms = generate_physical_massless_moms(rng, ss, n)
    energies = getT.(massless_moms)
    f = x -> _to_be_solved(x, masses, energies, ss)
    xi = find_zero((f, first_derivative(f)), x0, Roots.Newton())
    return _build_massive_momenta(xi, masses, massless_moms)
end
