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-9 # regulator for Green's functions in fm

"Green's function for Klein-Gordon equation"
greensFunctionKG(m, r, rp) = -(1 / m) * rp / (r + r_reg) * sinh(m * min(r, rp)) * exp(-m * max(r, rp))

"Green's function for Poisson's equation"
greensFunctionP(r, rp) = -rp^2 / (max(r, rp) + r_reg)

"Solve the Klein-Gordon equation (or Poisson's equation when m=0) for a given a source function (as an array) where
    the other parameters are the same from mesons!(...)."
function solveKG(m, source, r_max)
    greensFunction = m == 0 ? greensFunctionP : (r, rp) -> greensFunctionKG(m, r, rp)
    mesh = range(0, r_max; length=length(source))
    greenMat = greensFunction.(transpose(mesh), mesh)
    return greenMat * source
end

"Iteratively 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 the arrays are of length (1+divs),
    ρ_sp, ρ_vp are the scalar and vector proton densities as arrays,
    ρ_sn, ρ_vn are the scalar and vector neutron densities as arrays,
    Φ0, W0, B0, A0 are the mean-field potentials (couplings included) in MeV as as arrays,
    r is the radius in fm.
Reference: P. Giuliani, K. Godbey, E. Bonilla, F. Viens, and J. Piekarewicz, Frontiers in Physics 10, (2023)"
function solveMesonWfs(ρ_sp, ρ_vp, ρ_sn, ρ_vn, r_max, divs, iterations=3)
    (Φ0, W0, B0, A0) = (zeros(1 + divs) for _ in 1:4)
    (src_Φ0, src_W0, src_B0, src_A0) = (zeros(1 + divs) for _ in 1:4)

    for _ in iterations
        @. src_Φ0 = g2_s * ((κ/2) * Φ0^2 + (λ/6) * Φ0^3 - (ρ_sp + ρ_sn)) / ħc
        Φ0 .= solveKG(m_s / sqrt(ħc), src_Φ0, r_max)

        @. src_W0 = g2_v * ((ζ/6) * W0^3 + 2 * Λv * B0^2 * W0 - (ρ_vp + ρ_vn)) / ħc
        W0 .= solveKG(m_ω / sqrt(ħc), src_W0, r_max)

        @. src_B0 = g2_ρ * (2 * Λv * W0^2 * B0 - (ρ_vp - ρ_vn) / 2) / ħc
        W0 .= solveKG(m_ρ / sqrt(ħc), src_B0, r_max)

        @. src_A0 = ρ_vp / ħc
        A0 .= solveKG(0, src_A0, r_max) # this doesn't need iterations
    end

    return (Φ0, W0, B0, A0)
end