using LinearAlgebra, Random, Statistics

# Bayesian linear regression with a Normal-Inverse-Gamma prior:
#   σ² ~ InvGamma(a0, b0),  β | σ² ~ Normal(0, σ² V0)
# The posterior is the same family, so it has a closed form.
function nig_posterior(y, cols...; v0 = 100.0, a0 = 1.0, b0 = 0.01, intercept = true)
    Z  = intercept ? hcat(ones(length(y)), cols...) : hcat(cols...)
    Λn = I / v0 + Z' * Z                      # posterior precision (per σ²)
    mn = Λn \ (Z' * y)
    an = a0 + length(y) / 2
    bn = b0 + (y' * y - mn' * Λn * mn) / 2
    (m = mn, L = cholesky(Symmetric(inv(Λn))).L, a = an, b = bn)
end

# Gamma(k, 1) for k a multiple of 1/2, with the standard library only
gamma_draw(k) = sum(randexp(floor(Int, k))) + (isinteger(k) ? 0.0 : randn()^2 / 2)

function draw(post)
    σ² = post.b / gamma_draw(post.a)
    post.m .+ sqrt(σ²) .* (post.L * randn(length(post.m)))
end