include("common.jl")
include("system.jl")

"Solve the Klein-Gordon equation (or Poisson's equation when m=0) using the Green's function
method with cubic-spline interpolation of the source and adaptive Simpson's integration for
each grid point (matching the approach in Hartree.f). Returns an array of field values in MeV.
    m is the mass of the meson in MeV/c²,
    source is the source density array in fm⁻³."
function solveKG(m, source, s::system)
    N = s.divs
    dr = Δr(s)
    r_max = s.r_max
    m_nat = m / ħc    # mass in fm⁻¹
    r_grid = range(0, r_max, length=N+1)

    # Cubic-spline the source for evaluation at arbitrary quadrature points
    src_spline = cubic_spline_interpolation(r_grid, source)

    result = zeros(N + 1)
    ε = 1e-10

    if m_nat > ε    # Massive field (Klein-Gordon equation)
        for i in 0:N
            x = max(i * dr, ε)  # regularize at r=0

            # ∫₀ˣ r' × sinh(m×r') × exp(-m×x) × S(r') dr'
            I1 = x > ε ? adaptive_simps(0.0, x) do r
                r * sinh(m_nat * r) * exp(-m_nat * x) * src_spline(r)
            end : 0.0

            # ∫ₓ^r_max r' × sinh(m×x) × exp(-m×r') × S(r') dr'
            I2 = adaptive_simps(x, r_max) do r
                r * sinh(m_nat * x) * exp(-m_nat * r) * src_spline(r)
            end

            result[i+1] = -ħc * (I1 + I2) / (m_nat * x)
        end
    else    # Massless field (Poisson/Coulomb equation)
        # In the m→0 limit: sinh(m r')/m → r' and exp(-m r) → 1,
        # so field(x) = (1/x) ∫₀ˣ r'² S(r') dr' + ∫ₓ^∞ r' S(r') dr'
        for i in 0:N
            x = max(i * dr, ε)

            I1 = x > ε ? adaptive_simps(0.0, x) do r
                r^2 * src_spline(r)
            end : 0.0

            I2 = adaptive_simps(x, r_max) do r
                r * src_spline(r)
            end

            result[i+1] = -ħc * (I1 / x + I2)
        end
    end

    return result
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),
    r is the radius in fm,
    the inital solutions are read from s and the final solutions are saved in-place.
Reference: P. Giuliani, K. Godbey, E. Bonilla, F. Viens, and J. Piekarewicz, Frontiers in Physics 10, (2023)"
function solveMesonFields!(s::system, iterations=15, oscillation_control_parameter=0.3)
    (m_s, m_ω, m_ρ, m_γ, g2_s, g2_v, g2_ρ, g2_γ, κ_ss, λ, ζ, Λv) = (s.param.m_s, s.param.m_ω, s.param.m_ρ, s.param.m_γ, s.param.g2_s, s.param.g2_v, s.param.g2_ρ, s.param.g2_γ, s.param.κ_ss, s.param.λ, s.param.ζ, s.param.Λv)
    (Φ0, W0, B0, A0) = (s.Φ0, s.W0, s.B0, s.A0)
    (ρ_sp, ρ_vp, ρ_sn, ρ_vn) = (s.ρ_sp, s.ρ_vp, s.ρ_sn, s.ρ_vn)

    (src_Φ0, src_W0, src_B0, src_A0) = (zero_array(s) for _ in 1:4)

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

    for _ in 1:iterations
        @. src_Φ0 = g2_s * ((κ_ss/ħ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, s)

        # W0 and B0 keep a fraction of their previous solutions to suppress oscillations
        W0 .*= (1 - oscillation_control_parameter)
        B0 .*= (1 - oscillation_control_parameter)

        W0 .+= solveKG(m_ω, src_W0, s) .* oscillation_control_parameter
        B0 .+= solveKG(m_ρ, src_B0, s) .* oscillation_control_parameter
    end

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

"Calculate the total energy associated with meson fields"
function meson_E(s::system)
    (κ_ss, λ, ζ, Λv) = (s.param.κ_ss, s.param.λ, s.param.ζ, s.param.Λv)

    E_densities = map(zip(s.Φ0, s.W0, s.B0, s.A0, s.ρ_sp, s.ρ_vp, s.ρ_sn, s.ρ_vn)) do (Φ0, W0, B0, A0, ρ_sp, ρ_vp, ρ_sn, ρ_vn)
        E_σ = (1/2) * (Φ0/ħc) * (ρ_sp + ρ_sn) - ((κ_ss/ħc)/12 * (Φ0/ħc)^3 + (λ/24) * (Φ0/ħc)^4)
        E_ω = -(1/2) * (W0/ħc) * (ρ_vp + ρ_vn) + (ζ/24) * (W0/ħc)^4
        E_ρ = -(1/4) * (2B0/ħc) * (ρ_vp - ρ_vn)
        E_γ = -(1/2) * (A0/ħc) * ρ_vp
        E_ωρ = Λv * (W0/ħc)^2 * (2B0/ħc)^2
        E_σ + E_ω + E_ρ + E_γ + E_ωρ
    end
    
    return simpsons_integrate(E_densities .* rs(s).^2, Δr(s); coefficient=4π * ħc)
end