using DifferentialEquations

const ħc = 197.33 # MeVfm

# Values defined in C. J. Horowitz and J. Piekarewicz, Phys. Rev. Lett. 86, 5647 (2001)
# Values taken from Hartree.f (FSUGarnet)
const m_s = 496.939473213388 # MeV/c2
const m_ω = 782.5 # MeV/c2
const m_ρ = 763.0 # MeV/c2
const m_γ = 0.000001000 # MeV/c2 # should be 0?
const g2_s = 110.349189097820 # dimensionless
const g2_v = 187.694676506801 # dimensionless
const g2_ρ = 192.927428365698 # dimensionless
const g2_γ = 0.091701236 # dimensionless # equal to 4πα
const κ = 3.260178893447 # MeV
const λ = -0.003551486718 # dimensionless # LambdaSS
const ζ = 0.023499504053 # dimensionless # LambdaVV
const Λv = 0.043376933644 # dimensionless # LambdaVR

const r_reg = 1E-8 # fm # regulator for Green's functions

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

"Green's function for Poisson's equation in natural units"
greensFunctionP(r, rp) = -1 / (max(r, rp) + r_reg)

"Solve the Klein-Gordon equation (or Poisson's equation when m=0) and return an array in MeV for a source array given in fm⁻³ where
    m is the mass of the meson in MeV/c2,
    r_max is the r-cutoff in fm."
function solveKG(m, source, r_max)
    Δr = r_max / (length(source) - 1)
    rs = range(0, r_max; length=length(source))
    int_measure = ħc .* Δr .* rs .^ 2

    greensFunction = m == 0 ? greensFunctionP : (r, rp) -> greensFunctionKG(m / ħc, r, rp)
    greenMat = greensFunction.(rs, transpose(rs))

    return greenMat * (int_measure .* 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,
    An initial guess initial_sol=(Φ0, W0, B0, A0) can be provided to speed up convergence (permuting!).
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=10; initial_sol=(zeros(1 + divs) for _ in 1:4))
    (Φ0, W0, B0, A0) = initial_sol
    (src_Φ0, src_W0, src_B0, src_A0) = (zeros(1 + divs) for _ in 1:4)

    # A0 doesn't need iterations
    @. src_A0 = -g2_γ * ρ_vp
    A0 .= solveKG(m_γ, src_A0, r_max)

    for _ in 1:iterations
        @. src_Φ0 = g2_s * ((κ/ħc)/2 * (Φ0/ħc)^2 + (λ/6) * (Φ0/ħc)^3) - g2_s * (ρ_sp + ρ_sn)
        @. src_W0 = g2_v * ((ζ/6) * (W0/ħc)^3 + 2Λv * (2B0/ħc)^2 * (W0/ħc)) - g2_v * (ρ_vp + ρ_vn)
        @. src_B0 = (2Λv * g2_ρ * (W0/ħc)^2 * (2B0/ħc) - g2_ρ/2 * (ρ_vp - ρ_vn)) / 2
        Φ0 .= solveKG(m_s, src_Φ0, r_max)
        W0 .= (solveKG(m_ω, src_W0, r_max) .+ W0) ./ 2 # to supress oscillation
        B0 .= (solveKG(m_ρ, src_B0, r_max) .+ B0) ./ 2 # to supress oscillation
    end

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