module PoissonRandom

using Random: Random, AbstractRNG, randexp
using LogExpFunctions: log1pmx
using PrecompileTools: @compile_workload

export pois_rand, PassthroughRNG

"""
    PassthroughRNG()

An `AbstractRNG` sentinel for backend-specific random number generation.

`PassthroughRNG` forwards scalar `rand`, `randn`, and `randexp` calls to the
ambient random number implementation. It is useful when a backend, such as a GPU
kernel overlay, provides the actual random number generation but an
`AbstractRNG` argument is still needed by the Poisson sampler.

# Examples

```julia
using PoissonRandom

x = pois_rand(PassthroughRNG(), 3.0)
x isa Integer
```
"""
struct PassthroughRNG <: AbstractRNG end

Base.rand(rng::PassthroughRNG) = rand()
Random.randexp(rng::PassthroughRNG) = randexp()
Base.randn(rng::PassthroughRNG) = randn()

# When an overlay method table (e.g. CUDA.jl's `@device_override
# Random.randexp(::AbstractRNG)`) shadows the methods above, the overlay body
# runs with rng::PassthroughRNG and may call `Random.rand(rng, UInt52Raw())`
# or `Random.rand(rng, T)`. The stdlib Sampler chain bottoms out at
# `_rand52(r, rng_native_52(r))` → `rand(r, UInt64)`; provide those so the
# chain still reaches bare rand(T) and the device-side default_rng path.
Random.rng_native_52(::PassthroughRNG) = UInt64
Base.rand(rng::PassthroughRNG, ::Type{T}) where {T} = rand(T)

count_rand(λ::Real) = count_rand(Random.default_rng(), λ)
function count_rand(rng::AbstractRNG, λ::Real)
    n = 0
    c = randexp(rng)
    while c < λ
        n += 1
        c += randexp(rng)
    end
    return n
end

# Algorithm from:
#
#   J.H. Ahrens, U. Dieter (1982)
#   "Computer Generation of Poisson Deviates from Modified Normal Distributions"
#   ACM Transactions on Mathematical Software, 8(2):163-179
#
#   For μ sufficiently large, (i.e. >= 10.0)
#
ad_rand(λ::Real) = ad_rand(Random.default_rng(), λ)
function ad_rand(rng::AbstractRNG, λ::Real)
    s = sqrt(λ)
    d = 6 * λ^2
    L = floor(Int, λ - 1.1484)
    # Step N
    G = λ + s * randn(rng)

    if G >= 0
        K = floor(Int, G)
        # Step I
        if K >= L
            return K
        end

        # Step S
        U = rand(rng)
        if d * U >= (λ - K)^3
            return K
        end

        # Step P
        px, py, fx, fy = procf(λ, K, s)

        # Step Q
        if fy * (1 - U) <= py * exp(px - fx)
            return K
        end
    end

    while true
        # Step E
        E = randexp(rng)
        U = 2 * rand(rng) - 1
        T = 1.8 + copysign(E, U)
        if T <= -0.6744
            continue
        end

        K = floor(Int, λ + s * T)
        px, py, fx, fy = procf(λ, K, s)
        c = 0.1069 / λ

        # Step H
        @fastmath if c * abs(U) <= py * exp(px + E) - fy * exp(fx + E)
            return K
        end
    end
    return
end

# Procedure F
function procf(λ::Real, K::Int, s::Real)
    # can be pre-computed, but does not seem to affect performance
    INV_SQRT_2PI = inv(sqrt(2pi))
    ω = INV_SQRT_2PI / s
    b1 = inv(24) / λ
    b2 = 0.3 * b1 * b1
    c3 = inv(7) * b1 * b2
    c2 = b2 - 15 * c3
    c1 = b1 - 6 * b2 + 45 * c3
    c0 = 1 - b1 + 3 * b2 - 15 * c3

    if K < 10
        px = -float(λ)
        py = λ^K / prod(2:K)
    else
        δ = inv(12) / K
        δ -= 4.8 * δ^3
        V = (λ - K) / K
        px = K * log1pmx(V) - δ # avoids need for table
        py = INV_SQRT_2PI / sqrt(K)
    end
    X = (K - λ + 0.5) / s
    X2 = X^2
    fx = X2 / -2 # missing negation in pseudo-algorithm, but appears in fortran code.
    fy = ω * (((c3 * X2 + c2) * X2 + c1) * X2 + c0)
    return px, py, fx, fy
end

"""
    pois_rand(λ::Real)
    pois_rand(rng::AbstractRNG, λ::Real)

Draw one random variate from the Poisson distribution with rate `λ`.

For small `λ`, `pois_rand` uses an exponential-counting method. For larger `λ`,
it switches to the Ahrens-Dieter modified normal algorithm.

# Arguments

- `λ`: Poisson rate parameter. It should be nonnegative.
- `rng`: optional random number generator. Defaults to `Random.default_rng()`.

# Examples

```julia
using PoissonRandom
using Random

rng = MersenneTwister(1234)
x = pois_rand(rng, 4.0)
x isa Integer

y = pois_rand(PassthroughRNG(), 4.0)
y isa Integer
```
"""
pois_rand(λ::Real) = pois_rand(Random.default_rng(), λ)
pois_rand(rng::AbstractRNG, λ::Real) = λ < 6 ? count_rand(rng, λ) : ad_rand(rng, λ)

@compile_workload begin
    # Precompile the most common code paths
    # Small λ uses count_rand, large λ uses ad_rand
    pois_rand(3.0)   # count_rand path (λ < 6)
    pois_rand(50.0)  # ad_rand path (λ >= 6)
    # PassthroughRNG for GPU compatibility
    pois_rand(PassthroughRNG(), 3.0)
    pois_rand(PassthroughRNG(), 50.0)
end

end # module
