using DifferentialEquations

const ħc = 197.327 # ħc in MeVfm

# Values defined in C. J. Horowitz and J. Piekarewicz, Phys. Rev. Lett. 86, 5647 (2001)
# Values taken from Hartree.f (FSUGarnet) and
const m_ρ = 763.0
const m_ω = 782.5
const m_s = 496.939473213388
const m_γ = 0.000001000 # not defined in paper
const g2_s = 110.349189097820
const g2_v = 187.694676506801
const g2_ρ = 192.927428365698
const g2_g = 0.091701236 # not defined in paper
const κ = 3.260178893447
const λ = -0.003551486718 # LambdaSS
const ζ = 0.023499504053 # LambdaVV
const Λv = 0.043376933644 # LambdaVR

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

"Coupled radial meson equations that returns ddu=[ddΦ0, ddW0, ddB0, ddA0] in-place where
    du=[dΦ0, dW0, dB0, dA0] are the first derivatives of meson fields evaluated at r,
    u=[Φ0, W0, B0, A0] are the meson fields evaluated at r,
    κ is the generalized angular momentum,
    ρ_sp, ρ_vp are the scalar and vector proton densities as funtions of r in fm,
    ρ_sn, ρ_vn are the scalar and vector neutron densities as funtions of r in fm,
    Φ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 mesons!(ddu, du, u, (ρ_sp, ρ_vp, ρ_sn, ρ_vn), r)
    (dΦ0, dW0, dB0, dA0) = du
    (Φ0, W0, B0, A0) = u

    ddΦ0 = -(2/(r + r_reg)) * dΦ0 + m_s^2 * Φ0 / ħc + g2_s * ((κ/2) * Φ0^2 + (λ/6) * Φ0^3 - (ρ_sp(r) + ρ_sn(r))) / ħc
    ddW0 = -(2/(r + r_reg)) * dW0 + m_ω^2 * W0 / ħc + g2_v * ((ζ/6) * W0^3 + 2 * Λv * B0^2 * W0 - (ρ_vp(r) + ρ_vn(r))) / ħc
    ddB0 = -(2/(r + r_reg)) * dB0 + m_ρ^2 * B0 / ħc + g2_ρ * (2 * Λv * W0^2 * B0 - (ρ_vp(r) - ρ_vn(r)) / 2) / ħc
    ddA0 = -(2/(r + r_reg)) * dA0 - ρ_vp(r) / ħc # e goes here?

    ddu .= [ddΦ0, ddW0, ddB0, ddA0]
end

"Solve meson equations and return the wave functions u(r)=[Φ0(r), W0(r), B0(r), A0(r)] where
    divs is the number of mesh divisions so solution would be returned as a 4×(1+divs) matrix,
    the other parameters are the same from mesons!(...)."
function solveWfs(ρ_sp, ρ_vp, ρ_sn, ρ_vn, r_max, divs)
    prob = SecondOrderODEProblem(mesons!, BigFloat.([0, 0, 0, 0]), BigFloat.([0, 0, 0, 0]), (r_max, 0))
    sol = solve(prob, Feagin14(), p=(ρ_sp, ρ_vp, ρ_sn, ρ_vn), saveat=r_max/divs)
    wfs = hcat(sol.u...)

    return wfs
end