# Independent check of the extinction probabilities and the Otter-Dwass sum, by bisection and exact log-pmf.
import math
def G(s, R, k): return (1 + R / k * (1 - s)) ** (-k)
def root(f):
    lo, hi = 0.0, 1 - 1e-12
    # f(q) = G(q) - q is positive at 0 and negative just below the nontrivial root's right side up to 1
    for _ in range(200):
        mid = (lo + hi) / 2
        if f(mid) > 0: lo = mid
        else: hi = mid
    return (lo + hi) / 2
R = 2.5
for k in (0.1, 1.0):
    # bisect on [0, q*) where q* is the minimum of G(q) - q (supercritical: one root below 1)
    q = root(lambda x: G(x, R, k) - x if x < 0.999 else -1)
    print('k', k, 'q1 %.8f' % q, 'j*', math.ceil(math.log(0.05) / math.log(q)))
qp = root(lambda x: math.exp(-R * (1 - x)) - x if x < 0.999 else -1)
print('Poisson q1 %.8f' % qp, 'j*', math.ceil(math.log(0.05) / math.log(qp)))
k = 0.1
q = root(lambda x: G(x, R, k) - x if x < 0.999 else -1)
def pmf(s):  # P(S = s) = (1/s) P(NB(mean sR, size sk) = s - 1)
    n, size = s - 1, s * k
    p = size / (size + s * R)
    return math.exp(math.lgamma(n + size) - math.lgamma(n + 1) - math.lgamma(size) + size * math.log(p) + n * math.log(1 - p)) / s
p1 = pmf(1); print('P(S=1) %.6f' % p1, 'share %.2f%%' % (100 * p1 / q))
s500 = math.fsum(pmf(s) for s in range(1, 501)); tail = math.fsum(pmf(s) for s in range(501, 200001))
print('sum to 500 %.10f, tail 501..2e5 %.3e, q - sum500 %.3e' % (s500, tail, q - s500))
print('pmf(500) %.3e pmf(100) %.3e' % (pmf(500), pmf(100)))
# MERS-like and plague-like rows
for R2, k2 in ((1.2, 0.16), (1.5, 0.5)):
    q2 = root(lambda x: G(x, R2, k2) - x if x < 0.9999 else -1)
    print(R2, k2, '%.6f' % q2, math.ceil(math.log(0.05) / math.log(q2)))