# /// script
# requires-python = ">=3.13,<3.14"
# dependencies = ["scipy==1.16.2", "numpy==2.3.3"]
# ///

import math

from scipy.stats import betabinom, gamma, nbinom, poisson

shape, rate = 2 + 4, 1 + 2
posterior = gamma(a=shape, scale=1 / rate)
future_exposure = 1
predictive = nbinom(n=shape, p=rate / (rate + future_exposure))

assert math.isclose(posterior.mean(), 2)
assert math.isclose(posterior.var(), 2 / 3)
assert math.isclose(predictive.pmf(0), (3 / 4) ** 6)
assert math.isclose(predictive.var(), 8 / 3)
print(f"Gamma posterior rate mean: {posterior.mean():.7f}")
print(f"Predictive P(zero): {predictive.pmf(0):.7f}")
print(f"Plug-in P(zero): {poisson(mu=2).pmf(0):.7f}")

coin_prediction = betabinom(n=3, a=2 + 4, b=2 + 2)
assert math.isclose(coin_prediction.pmf(0), 1 / 11)
assert math.isclose(coin_prediction.mean(), 1.8)
print(f"Beta-binomial P(zero): {coin_prediction.pmf(0):.7f}")

assert math.isclose(poisson.sf(2, mu=2), gamma.cdf(1, a=3, scale=1 / 2))
print(f"Third arrival within one hour: {poisson.sf(2, mu=2):.7f}")