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

# 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 κ_ss = 3.260178893447 # MeV
const λ = -0.003551486718 # dimensionless # LambdaSS
const ζ = 0.023499504053 # dimensionless # LambdaVV
const Λv = 0.043376933644 # dimensionless # LambdaVR

"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."
function solveKG(m, source, s::system)

    @assert s.divs % 2 == 0 "Number of mesh divisions must be even for Simpson's rule"
    simpsons_weights = (Δr(s)/3) .* [1; repeat([2,4], s.divs ÷ 2)[2:end]; 1]
    int_measure = ħc .* simpsons_weights .* rs(s) .^ 2

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

    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),
    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)
    (Φ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)
    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