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,
    f1(r) =  M-Φ0(r)-W0(r)-(p-0.5)*2B0(r)-p*A0(r) is a function of r in MeV (see optimized_dirac_potentials()),
    f2(r) = -M+Φ0(r)-W0(r)-(p-0.5)*2B0(r)-p*A0(r) is a function of r in MeV (see optimized_dirac_potentials()),
    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, (κ, E, f1, f2), r)
    (g, f) = u
    @inbounds du[1] = -(κ/(r + r_reg)) * g + (E + f1(r)) * f / ħc
    @inbounds du[2] =  (κ/(r + r_reg)) * f - (E + f2(r)) * g / ħc
end

"Get the potentials f1 and f2 that goes into the Dirac equation, defined as
    f1(r) =  M-Φ0(r)-W0(r)-(p-0.5)*2B0(r)-p*A0(r),
    f2(r) = -M+Φ0(r)-W0(r)-(p-0.5)*2B0(r)-p*A0(r)."
function optimized_dirac_potentials(p, s::system)
    M = p ? M_p : M_n

    f1s = zero_array(s)
    f2s = zero_array(s)

    @. f1s =  M - s.Φ0 - s.W0 - (p - 0.5) * 2 * s.B0 - p * s.A0
    @. f2s = -M + s.Φ0 - s.W0 - (p - 0.5) * 2 * s.B0 - p * s.A0

    f1 = linear_interpolation(rs(s), f1s)
    f2 = linear_interpolation(rs(s), f2s)

    return (f1, f2)
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, s::system; shooting=true, normalize=true, algo=Tsit5())
    (f1, f2) = optimized_dirac_potentials(p, s)

    if shooting
        @assert s.divs % 2 == 0 "divs must be an even number when shooting=true"
        prob = ODEProblem(dirac!, [0, 1], (s.r_max, s.r_max / 2))
        sol = solve(prob, algo, p=(κ, E, f1, f2), saveat=Δr(s))
        wf_right = reverse(hcat(sol.u...); dims=2)
        next_r_max = s.r_max / 2 # for the next step
    else
        next_r_max = s.r_max
    end
    
    prob = ODEProblem(dirac!, [0, 1], (0, next_r_max))
    sol = solve(prob, algo, p=(κ, E, f1, f2), saveat=Δr(s))
    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(s) # integration by Reimann sum
        wf = wf ./ sqrt(norm)
    end

    return wf
end

"Returns a function that solves the Dirac equation and returns g(r=r_max) where
    r_max is the outer boundary in fm,
    the other parameters are the same from dirac!(...)."
function boundaryValueFunc(κ, p, s::system; dtype=Float64, algo=Tsit5())
    prob = ODEProblem(dirac!, convert.(dtype, [0, 1]), (0, s.r_max))
    (f1, f2) = optimized_dirac_potentials(p, s)
    func(E) = solve(prob, algo, p=(κ, E, f1, f2), saveat=[s.r_max], save_idxs=[1])[1, 1]
    return func
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, s::system, E_min=0, E_max=(p ? M_p : M_n), tol_digits=5)
    func = boundaryValueFunc(κ, p, s)
    Es = find_zeros(func, (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, s::system, 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, s, 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 Z_or_N lowest lying orbitals and return occupancy numbers"
function fillNucleons(Z_or_N::Int, κs, Es)
    sort_i = sortperm(Es)

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

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

    Z_or_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, s::system)
    ρr2 = zeros(2, s.divs + 1) # ρ×r² values
    
    for (κ, E, occ) in zip(κs, Es, occs)
        wf = solveNucleonWf(κ, p, E, s; shooting=true, normalize=true)
        wf2 = wf .* wf
        ρr2 += (occ / (4 * pi)) * wf2 # 2j+1 factor is accounted in the occupancy number
    end

    r2s = rs(s).^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(p, s::system, E_min=800, E_max=939)
    κs, Es = findAllOrbitals(p, s, E_min, E_max)
    occs = fillNucleons(Z_or_N(s, p), κs, Es)
    return calculateNucleonDensity(κs, Es, occs, p, s)
end