using DifferentialEquations, Roots

const ħc = 197.33 # MeVfm
const M_n = 939.0 # MeV/c2
const M_p = 939.0 # MeV/c2

const r_reg = 1E-8 # fm # regulator for the centrifugal term

"The spherical Dirac equation that returns du=[dg, df] in-place where
    u=[g, f] are the reduced radial components evaluated at r,
    κ is the generalized angular momentum,
    p is true for proton and false for neutron,
    E in the energy in MeV,
    Φ0, W0, B0, A0 are the mean-field potentials (couplings included) in MeV as functions of r in fm,
    r is the radius in fm.
Reference: P. Giuliani, K. Godbey, E. Bonilla, F. Viens, and J. Piekarewicz, Frontiers in Physics 10, (2023)"
function dirac!(du, u, (κ, p, E, Φ0, W0, B0, A0), r)
    M = p ? M_p : M_n
    common1 = E - W0(r) - (p - 0.5) * 2B0(r) - p * A0(r)
    common2 = M - Φ0(r)
    (g, f) = u
    du[1] = -(κ/(r + r_reg)) * g + (common1 + common2) * f / ħc
    du[2] =  (κ/(r + r_reg)) * f - (common1 - common2) * g / ħc
end

"Solve the Dirac equation and return the wave function u(r)=[g(r), f(r)] where
    divs is the number of mesh divisions so solution would be returned as a 2×(1+divs) matrix,
    shooting method divides the interval into two partitions at r_max/2, ensuring convergence at both r=0 and r=r_max,
    the other parameters are the same from dirac!(...)."
function solveNucleonWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs; shooting=true, normalize=true)
    Δr = r_max / divs

    if shooting
        @assert divs % 2 == 0 "divs must be an even number when shooting=true"
        prob = ODEProblem(dirac!, [0, 1], (r_max, r_max / 2))
        sol = solve(prob, RK4(), p=(κ, p, E, Φ0, W0, B0, A0), saveat=Δr)
        wf_right = reverse(hcat(sol.u...); dims=2)
        r_max = r_max / 2 # for the next step
    end
    
    prob = ODEProblem(dirac!, [0, 1], (0, r_max))
    sol = solve(prob, RK4(), p=(κ, p, E, Φ0, W0, B0, A0), saveat=Δr)
    wf = hcat(sol.u...)

    if shooting # join two segments
        rescale_factor_g = wf[1, end] / wf_right[1, 1]
        rescale_factor_f = wf[2, end] / wf_right[2, 1]
        @assert isfinite(rescale_factor_g) && isfinite(rescale_factor_f) "Cannot rescale the right partition"
        isapprox(rescale_factor_g, rescale_factor_f; rtol=0.03) || @warn "Discontinuity between two partitions"
        wf_right_rescaled = wf_right .* ((rescale_factor_g + rescale_factor_f) / 2)
        wf = hcat(wf[:, 1:(end - 1)], wf_right_rescaled)
    end

    if normalize
        norm = sum(wf .* wf) * Δr # integration by Reimann sum
        wf = wf ./ sqrt(norm)
    end

    return wf
end

"Solve the Dirac equation and return g(r=r_max) where
    r_max is the outer boundary in fm,
    the other parameters are the same from dirac!(...)."
function boundaryValue(κ, p, E, Φ0, W0, B0, A0, r_max; dtype=Float64, algo=RK4())
    prob = ODEProblem(dirac!, convert.(dtype, [0, 1]), (0, r_max))
    sol = solve(prob, algo, p=(κ, p, E, Φ0, W0, B0, A0), saveat=[r_max], save_idxs=[1])
    return sol[1, 1]
end

"Find all bound energies between E_min (=0) and E_max (=mass) where
    tol_digits determines the precision for root finding and the threshold for identifying duplicates,
    the other parameters are the same from dirac!(...)."
function findEs(κ, p, Φ0, W0, B0, A0, r_max, E_min=0, E_max=(p ? M_p : M_n), tol_digits=5)
    f(E) = boundaryValue(κ, p, E, Φ0, W0, B0, A0, r_max)
    Es = find_zeros(f, (E_min, E_max); xatol=1/10^tol_digits)
    return unique(E -> round(E; digits=tol_digits), Es)
end

"Find all orbitals and return two lists containing κ values and corresponding energies for a single species where
    the other parameters are defined above"
function findAllOrbitals(p, Φ0, W0, B0, A0, r_max, E_min=0, E_max=(p ? M_p : M_n))
    κs = Int[]
    Es = Float64[]
    # start from κ=-1 and go both up and down
    for direction in [-1, 1]
        for κ in direction * (1:100) # cutoff is 100
            new_Es = findEs(κ, p, Φ0, W0, B0, A0, r_max, E_min, E_max)
            if isempty(new_Es); break; end
            append!(Es, new_Es)
            append!(κs, fill(κ, length(new_Es)))
        end
    end
    return (κs, Es)
end

"For a given list of κ values with corresponding energies, attempt to fill N lowest lying orbitals and return occupancy numbers"
function fillNucleons(N::Int, κs, Es)
    sort_i = sortperm(Es)

    occ = zeros(Int, length(κs))

    for i in sort_i
        if N ≤ 0; break; end;
        max_occ = 2 * j_κ(κs[i]) + 1
        occ[i] = min(max_occ, N)
        N -= occ[i]
    end

    N == 0 || @warn "All orbitals could not be filled"
    return occ
end

"Total angular momentum j for a given κ value"
j_κ(κ::Int) = abs(κ) - 1/2

"Orbital angular momentum l for a given κ value"
l_κ(κ::Int) = abs(κ) - (κ < 0) # since true = 1 and false = 0

"Calculate scalar and vector densities of a nucleon species on [0,r_max] divided into (divs+1) points and returns them as vectors (ρ_s, ρ_v) where
    the arrays κs, Es, occs tabulate the energies and occupation numbers corresponding to each κ,
    the other parameters are defined above"
function calculateNucleonDensity(κs, Es, occs, p, Φ0, W0, B0, A0, r_max, divs, E_min=0, E_max=(p ? M_p : M_n))
    ρr2 = zeros(2, divs + 1) # ρ×r² values
    
    for (κ, E, occ) in zip(κs, Es, occs)
        wf = solveNucleonWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs; shooting=true, normalize=true)
        wf2 = wf .* wf
        ρr2 += (occ / (4 * pi)) * wf2 # 2j+1 factor is accounted in the occupancy number
    end

    r2s = (collect ∘ range)(0, r_max, length=divs+1).^2
    ρ = ρr2 ./ transpose(r2s)
    ρ[:, 1] .= ρ[:, 2] # dirty fix for NaN at r=0
    
    ρ_s = ρ[1, :] - ρ[2, :] # g^2 - f^2
    ρ_v = ρ[1, :] + ρ[2, :] # g^2 + f^2
    
    return (ρ_s, ρ_v)
end

"Solve the Dirac equation and calculate scalar and vector densities of a nucleon species where
    the other parameters are defined above"
function solveNucleonDensity(N, p, Φ0, W0, B0, A0, r_max, divs, E_min=0, E_max=(p ? M_p : M_n))
    κs, Es = findAllOrbitals(p, Φ0, W0, B0, A0, r_max, E_min, E_max)
    occs = fillNucleons(N, κs, Es)
    return calculateNucleonDensity(κs, Es, occs, p, Φ0, W0, B0, A0, r_max, divs, E_min, E_max)
end