Complex{BigFloat} log_gamma

alt_log_gamma.jl

#=
     There was some push-pull with feedback from Perplexity and from Chatgpt 5.4.
     log_gamma [loggamma] below is a strongly guided and multi-refined result, AI involved rather than an AI generated. 
     The inline docs are AI generated. 
     The bernoulli functions were AI generated -- where duplicated, they may be taken from Timon's implementation.
=#


# using SpecialFunctions

#export loggamma

# ─────────────────────────────────────────────────────────────────────────────
# Section 0 – Public API: extend SpecialFunctions.loggamma
# ─────────────────────────────────────────────────────────────────────────────
# currently named log_gamma for comparative testing

"""
    loggamma(z::Complex{BigFloat}) -> Complex{BigFloat}

Compute the principal value of log Γ(z) for complex `BigFloat` arguments.

This method extends `SpecialFunctions.loggamma` to `Complex{BigFloat}`,
filling the gap left by `SpecialFunctions.jl` which only implements
`loggamma` for `Complex{Float32/Float64}` and real `BigFloat`.

## Accuracy

The result is accurate to within a few ULP of the input precision. Internal
arithmetic uses `precision(real(z)) + 64` guard bits, and the output is
rounded back to `precision(real(z))` bits.

## Branch cut

The principal branch is used: `log` throughout takes its principal value
(imaginary part in `(−π, π]`).  The result is the analytic continuation of
the real log Γ from the positive real axis.

## Special values

| Input                        | Result                 |
|------------------------------|------------------------|
| `z = 1 + 0im`                | `0 + 0im`              |
| `z = 1/2 + 0im`              | `0.5·log(π) + 0im`     |
| `z = n + 0im`, `n ∈ ℤ₊`     | `log((n-1)!) + 0im`    |
| `z ∈ {0, −1, −2, …} + 0im`  | `+Inf + 0im` (pole)    |
| non-finite `z`               | `NaN + NaN·i`          |

## Examples

```julia
julia> using BigComplexGamma, SpecialFunctions
julia> setprecision(BigFloat, 256)

# Should be 0
julia> loggamma(Complex(BigFloat(1), BigFloat(0)))

# Should be 0.5 * log(π) ≈ 0.5723649…
julia> loggamma(Complex(BigFloat(1//2), BigFloat(0)))

# Cross-check against Float64
julia> z64 = 2.5 + 1.5im
julia> abs(loggamma(Complex(BigFloat(real(z64)), BigFloat(imag(z64)))) - loggamma(z64))
# Should be < 1e-15

# Very high precision (512 bits)
julia> setprecision(BigFloat, 512)
julia> loggamma(Complex(BigFloat(1//3), BigFloat(1//4)))
```
"""
#function SpecialFunctions.loggamma(z::Complex{BigFloat})
function log_gamma(z::Complex{BigFloat})
    prec = precision(real(z))
    guard = 64   # Internal guard bits to absorb rounding accumulation

    # Run all internal arithmetic at elevated precision.
    # setprecision is thread-local in Julia ≥1.7 and safe to use here.
    result = setprecision(BigFloat, prec + guard) do
        zw = Complex(BigFloat(real(z)), BigFloat(imag(z)))
        _loggamma_cbf(zw)
    end

    # Round back to the caller's requested precision.
    re_out = setprecision(BigFloat, prec) do
        BigFloat(real(result))
    end
    im_out = setprecision(BigFloat, prec) do
        BigFloat(imag(result))
    end
    return Complex(re_out, im_out)
end


# ─────────────────────────────────────────────────────────────────────────────
# Section 1 – Exact Bernoulli numbers as Rational{BigInt}
# ─────────────────────────────────────────────────────────────────────────────

"""
    _compute_bernoulli_coefficients(nmax::Int) -> Vector{Rational{BigInt}}

Return the Stirling-series coefficients
    c[k] = B_{2k} / (2k · (2k−1)),  k = 1, …, nmax
as exact `Rational{BigInt}` values computed via the standard recurrence
    B_m = −(1/(m+1)) · Σ_{k=0}^{m-1} C(m+1,k) · B_k.

Odd Bernoulli numbers (index ≥ 3) are zero; only even-indexed ones appear
in the Stirling series, so this function extracts only those.

The binomial coefficients C(m+1,k) are accumulated iteratively as exact
`BigInt`, avoiding any floating-point error in the recurrence.
"""
function _compute_bernoulli_coefficients(nmax::Int)::Vector{Rational{BigInt}}
    n_total = 2 * nmax          # We need B_0, B_2, B_4, …, B_{2·nmax}
    B = Vector{Rational{BigInt}}(undef, n_total + 1)

    # Seed values
    B[1] = one(Rational{BigInt})            # B_0 = 1
    n_total >= 1 && (B[2] = Rational{BigInt}(-1, 2))   # B_1 = -1/2

    for m in 2:n_total
        # Compute B_m = -(1/(m+1)) * sum_{k=0}^{m-1} C(m+1,k)*B_k
        # Binomial C(m+1, k) accumulated iteratively from C(m+1,0)=1.
        s = zero(Rational{BigInt})
        binom = one(BigInt)     # C(m+1, 0) = 1
        for k in 0:(m-1)
            s += binom * B[k+1]
            # C(m+1, k+1) = C(m+1, k) * (m+1-k) / (k+1)
            # The division is always exact because C(n,k+1) ∈ ℤ.
            binom = binom * BigInt(m + 1 - k) ÷ BigInt(k + 1)
        end
        B[m+1] = -s // BigInt(m + 1)
    end

    # Extract c_k = B_{2k} / (2k*(2k-1)) for k = 1..nmax
    return Rational{BigInt}[B[2k+1] // (BigInt(2k) * BigInt(2k - 1))
                            for k in 1:nmax]
end

# Thread-safe, lazily-growing cache of Stirling coefficients.
const _BCOEFF_CACHE = Ref{Vector{Rational{BigInt}}}(Rational{BigInt}[])
const _BCOEFF_CACHE_LOCK = ReentrantLock()
const _BCOEFF_BIGFLOAT_CACHE = Dict{Int,Vector{BigFloat}}()

"""
    _bernoulli_stirling(k::Int) -> Rational{BigInt}

Return the k-th Stirling coefficient c_k = B_{2k}/(2k·(2k−1)) as an exact
rational.  Uses a thread-safe cache that grows on demand.
"""
function _bernoulli_stirling(k::Int)::Rational{BigInt}
    lock(_BCOEFF_CACHE_LOCK) do
        if k > length(_BCOEFF_CACHE[])
            newlen = max(k, 2 * length(_BCOEFF_CACHE[]) + 16)
            _BCOEFF_CACHE[] = _compute_bernoulli_coefficients(newlen)
        end
        return _BCOEFF_CACHE[][k]
    end
end

"""
    _bernoulli_stirling_bigfloat(kmax::Int, prec::Int) -> Vector{BigFloat}

Return the first `kmax` Stirling coefficients converted to `BigFloat` at the
requested precision. The cache is thread-safe and grows geometrically, so the
hot Stirling loop can reuse already-converted coefficients without per-term
locking or rational-to-BigFloat conversion.
"""
function _bernoulli_stirling_bigfloat(kmax::Int, prec::Int)::Vector{BigFloat}
    lock(_BCOEFF_CACHE_LOCK) do
        coeffs = get!(_BCOEFF_BIGFLOAT_CACHE, prec) do
            BigFloat[]
        end

        if kmax > length(coeffs)
            newlen = max(kmax, 2 * length(coeffs) + 16)
            if newlen > length(_BCOEFF_CACHE[])
                _BCOEFF_CACHE[] = _compute_bernoulli_coefficients(newlen)
            end

            setprecision(BigFloat, prec) do
                sizehint!(coeffs, newlen)
                for idx in (length(coeffs)+1):newlen
                    ck_rat = _BCOEFF_CACHE[][idx]
                    push!(coeffs, BigFloat(numerator(ck_rat)) / BigFloat(denominator(ck_rat)))
                end
            end
        end

        return coeffs
    end
end

# ─────────────────────────────────────────────────────────────────────────────
# Section 2 – Stirling asymptotic series
# ─────────────────────────────────────────────────────────────────────────────

"""
    _stirling_loggamma(z::Complex{BigFloat}, prec::Int) -> Complex{BigFloat}

Evaluate the Stirling asymptotic expansion
    log Γ(z) = (z−½)·log z − z + ½·log(2π)
               + Σ_{k=1}^{N} B_{2k}/(2k·(2k−1)·z^{2k−1})

The caller must ensure that `|z|` (specifically `Re(z)`) is large enough for
the series to achieve `prec`-bit accuracy; see `_shift_for_stirling`.

Termination uses a *two-criterion* rule:
  - term magnitude drops below `2^{-(prec+4)}`  (converged), OR
  - term magnitude starts growing again          (asymptotic divergence onset).
The second criterion ensures we back off before the asymptotic tail blows up.

All BigFloat arithmetic executes at the precision active in the calling scope
(assumed to be `prec + guard` bits set by the caller).
"""
function _stirling_loggamma(z::Complex{BigFloat}, prec::Int)::Complex{BigFloat}
    T = BigFloat
    half = T(1) / T(2)

    # Leading terms: (z - 1/2)*log(z) - z + log(2π)/2
    result = (z - half) * log(z) - z + log(T(2) * T(π)) * half

    # Asymptotic correction sum: Σ c_k / z^{2k-1}
    threshold = T(2)^(-(prec + 4))
    z2 = z * z
    zpow = z                  # z^{2k-1} starting at k=1: z^1
    prev_mag = T(Inf)
    coeffs = _bernoulli_stirling_bigfloat(16, prec)

    for k in 1:100_000             # practical upper bound; never reached
        if k > length(coeffs)
            coeffs = _bernoulli_stirling_bigfloat(2 * k, prec)
        end
        term = coeffs[k] / zpow
        mag = abs(real(term)) + abs(imag(term))

        result += term

        # Stop when converged
        mag < threshold && break

        # Stop before asymptotic divergence takes over (series starts growing)
        if mag > prev_mag
            result -= term         # undo the diverging term
            break
        end

        prev_mag = mag
        zpow *= z2              # advance: z^1 → z^3 → z^5 → …
    end

    return result
end

# ─────────────────────────────────────────────────────────────────────────────
# Section 3 – Minimum shift for Stirling convergence
# ─────────────────────────────────────────────────────────────────────────────

"""
    _safe_shift_from_downcast(diff::BigFloat, ::Type{T}) -> Union{Nothing, Int}

Try to recover `ceil(diff) + 1` from a floating-point downcast of `diff`
without changing the integer result.

This is used as a tiered fast path in `_stirling_shift`:
- first try `Float64`, which is cheapest and usually sufficient;
- if `Float64` is too coarse near an integer boundary or at large magnitude,
  try `Float128` from `Quadmath`, which is slower but still much cheaper than
  exact `BigFloat`/`BigInt` fallback;
- if neither downcast can determine the ceiling unambiguously, return `nothing`
  and let `_stirling_shift` use exact arithmetic.

The downcast is accepted only when a conservative one-ULP uncertainty interval
around the floating-point value yields the same `ceil` at both ends.
"""
function _safe_shift_from_downcast(diff::BigFloat, ::Type{T})::Union{Nothing,Int} where {T<:AbstractFloat}
    downcast = T(diff)
    isfinite(downcast) || return nothing

    ulp = max(downcast - prevfloat(downcast), nextfloat(downcast) - downcast)
    ulp < one(T) || return nothing

    lo = downcast - ulp
    hi = downcast + ulp
    ceil_lo = ceil(Int, lo)
    ceil_hi = ceil(Int, hi)
    ceil_lo == ceil_hi || return nothing

    return ceil_hi + 1
end


"""
    _stirling_shift(re_z::BigFloat, prec::Int) -> Int

Return the non-negative integer `m` such that shifting `z → z + m` places
`Re(z+m) ≥ target` where
    target = prec · ln(2) / (2π) + 2.

This target ensures that the Stirling series achieves `prec`-bit accuracy.
The derivation follows from the asymptotic bound
    |c_k / z^{2k-1}| ≲ (k/(π·|z|))^{2k}
which decreases geometrically when |z| > k/(2πe), and the precision needed
sets the optimal k ~ prec·ln(2)/(2π).
"""
function _stirling_shift(re_z::BigFloat, prec::Int)::Int
    T = BigFloat
    target = T(prec) * log(T(2)) / (T(2) * T(π)) + T(2)
    diff = target - re_z
    diff <= zero(T) && return 0

    int_limit = T(typemax(Int) - 1)
    diff > int_limit && throw(OverflowError("required Stirling shift exceeds Int range"))

    # This ordering was benchmarked: Float64 wins on the common path, Float128
    # recovers most of the guarded-path overhead when Float64 is too coarse.
    maybe_shift = _safe_shift_from_downcast(diff, Float64)
    maybe_shift !== nothing && return maybe_shift

    # maybe_shift = _safe_shift_from_downcast(diff, Float128)
    # maybe_shift !== nothing && return maybe_shift

    return Int(ceil(BigInt, diff) + 1)
end

"""
    _reflection_branch_correction(re_z::BigFloat, im_z::BigFloat) -> Complex{BigFloat}

Let
    F(z) = log(π) - log(sin(πz)) - loggamma(1-z)
with principal `log`. Then `F(z)` differs from the principal branch of
`loggamma(z)` by an integer multiple of `2πi`.

Using the two-step recurrence together with `sin(π(z+2)) = sin(πz)`, one finds
that every backward shift `z -> z-2` changes this difference by exactly
    2πi * sign(imag(z))
as long as `imag(z) != 0`.

We normalize the correction to zero on the base strip `-1/2 < Re(z) < 3/2`.
This helper counts how many two-step backward shifts separate `z` from that
base strip and returns the accumulated branch jump.
"""
function _reflection_branch_correction(re_z::BigFloat, im_z::BigFloat)::Complex{BigFloat}
    T = BigFloat
    iszero(im_z) && return zero(Complex{T})

    # Number of backward two-step shifts z -> z-2 needed to leave the
    # normalization strip -1/2 < Re(z) < 3/2 and reach the current point.
    backward_twostep_shifts = floor(Int, (T(3) / T(2) - re_z) / T(2))
    backward_twostep_shifts <= 0 && return zero(Complex{T})

    return Complex(zero(T), T(2) * T(π) * T(backward_twostep_shifts) * sign(im_z))
end

# ─────────────────────────────────────────────────────────────────────────────
# Section 4 – Core recursive algorithm
# ─────────────────────────────────────────────────────────────────────────────

"""
    _loggamma_cbf(z::Complex{BigFloat}) -> Complex{BigFloat}

Internal implementation of log Γ for `Complex{BigFloat}`.  Assumes the
BigFloat global precision has been set to `prec + guard` by the caller
(`loggamma` public entry point).

The algorithm is:
  1. Non-finite → NaN.
  2. Pole (z ∈ {0, −1, −2, …}) → +Inf.
  3. Re(z) < 1/2: reflection formula (recursive call on 1−z).
  4. Re(z) ≥ 1/2: upward recurrence then Stirling series.
"""
function _loggamma_cbf(z::Complex{BigFloat})::Complex{BigFloat}
    T = BigFloat
    re_z = real(z)
    im_z = imag(z)
    prec = precision(re_z)

    # ── 1. Guard: non-finite input ────────────────────────────────────────────
    if !isfinite(re_z) || !isfinite(im_z)
        return Complex(T(NaN), T(NaN))
    end

    # ── 2. Poles of Γ at the non-positive integers ────────────────────────────
    # Γ has simple poles at z = 0, -1, -2, …; log Γ → +∞ there.
    if im_z == zero(T) && re_z <= zero(T) && isinteger(re_z)
        return Complex(T(Inf), zero(T))
    end

    # ── 3. Reflection formula for Re(z) < 1/2 ────────────────────────────────
    # Identity: Γ(z)·Γ(1−z) = π / sin(πz)
    # ⟹ log Γ(z) = log π − log sin(πz) − log Γ(1−z), modulo 2πi.
    # The helper below supplies the correct multiple of 2πi by counting the
    # branch jumps induced by repeated two-step recurrence moves z -> z-2.
    # Reflection maps Re(z) < 1/2 → Re(1−z) > 1/2, entering the convergence
    # region for the recursive call.
    if re_z < T(1) / T(2)
        π_bf = T(π)
        sinπz = sin(π_bf * z)
        # Recursive call: 1−z has Re(1−z) > 1/2, so no infinite recursion.
        lg1mz = _loggamma_cbf(one(T) - z)
        return log(π_bf) - log(sinπz) - lg1mz - _reflection_branch_correction(re_z, im_z)
    end

    # ── 4. Re(z) ≥ 1/2: shift upward, then Stirling ──────────────────────────
    m = _stirling_shift(re_z, prec)
    log_prefix = zero(Complex{T})
    for k in 0:(m-1)
        log_prefix += log(z + T(k))
    end

    z_shifted = z + T(m)
    lg_shifted = _stirling_loggamma(z_shifted, prec)

    return lg_shifted - log_prefix
end


# ===========================
# benchmark both for accuracy
# ===========================

#=
   results from one run

log_gamma may average ~14 bits of additional accuracy over loggamma
log_gamma may average ~1.05% faster than loggamma
=#

using Statistics: mean
using Chairmarks

BFPrec = 256
BFHiPrec = BFPrec + 64

setprecision(BigFloat, BFHiPrec)
bigrands = rand(Complex{BigFloat}, 4096)

include("Timon/gamma.jl")

tlg = map(x -> loggamma(x), bigrands);

include("jas/log_gammas.jl")

jas = map(x -> log_gamma(x), bigrands);

setprecision(BigFloat, BFPrec)
rands = map(x -> BigFloat(x.re) + BigFloat(x.im) * im, bigrands);

tlg2 = map(x -> loggamma(x), rands);
jas2 = map(x -> log_gamma(x), rands);

setprecision(BigFloat, BFHiPrec)

dtlg = map(x -> Complex{Float64}(x), tlg .- tlg2);
rtlg = map(x -> Complex{Float64}(x), (tlg .- tlg2) ./ tlg);

djas = map(x -> Complex{Float64}(x), jas .- jas2);
rjas = map(x -> Complex{Float64}(x), (jas .- jas2) ./ jas);

delta_bits_accuracy = round(log2(abs(mean(dtlg)) / abs(mean(djas))), sigdigits=4)
delta_rel_bits_accuracy = round(log2(abs(mean(rtlg)) / abs(mean(rjas))), sigdigits=4)

extra_bit_accuracy = round(Int, (7 * delta_bits_accuracy + 1 * delta_rel_bits_accuracy) / 8)
println("log_gamma may average ~$(extra_bit_accuracy) bits of additional accuracy over loggamma")

# log_gamma is 10+ bits more accurate than loggamma


setprecision(BigFloat, BFPrec)

function time_loggamma(randvals)
    map(loggamma, randvals)
end

function time_log_gamma(randvals)
    map(loggamma, randvals)
end

loggamma_timer = @b time_loggamma($rands)
loggamma_time = loggamma_timer.time

log_gamma_timer = @b time_log_gamma($rands)
log_gamma_time = log_gamma_timer.time

# times are quite similar, log_gamma may be 1% faster.
speedup = round(((log_gamma_time / loggamma_time) - 1) * 100, sigdigits=3);
println("log_gamma may be ~$(speedup)% faster than loggamma")