# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Availability of unanimity against a majority.

Atomic commit needs every participant up; consensus needs a majority. With
each node up with probability 0.999: a^n against the binomial majority, for n
from 1 to 20; then 10 million simulated rounds for 10-node unanimity and a
3-of-5 majority, independent and with a correlated rack failure (probability
0.0005 per round of losing 5 nodes at once).
"""
import argparse, json
from math import comb
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--a", type=float, default=0.999, help="Node availability.")
p.add_argument("--rounds", type=int, default=10_000_000, help="Simulated rounds.")
p.add_argument("--rack-p", type=float, default=0.0005, help="Per-round probability that a rack of 5 fails together.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

a = args.a
def unanimity(n): return a ** n
def majority(n): return sum(comb(n, k) * a ** k * (1 - a) ** (n - k) for k in range(n // 2 + 1, n + 1))
formula = [{"n": n, "unanimity": unanimity(n), "majority": majority(n)} for n in (1, 3, 5, 10, 20)]
rng = np.random.default_rng(args.seed); sims = []
for correlated in (False, True):
    up10 = np.ones(args.rounds, dtype=np.int8); up5 = np.zeros(args.rounds, dtype=np.int16)
    for i in range(10): up10 &= (rng.random(args.rounds) < a)
    rack = rng.random(args.rounds) < args.rack_p if correlated else np.zeros(args.rounds, dtype=bool)
    all10 = up10.astype(bool) & ~rack                                        # the 10 participants span the failed rack
    for i in range(5): up5 += (rng.random(args.rounds) < a)
    maj = (up5 >= 3) & ~(rack & (rng.random(args.rounds) < 0.6))            # a rack failure takes 3 of the 5 replicas 60% of the time
    sims.append({"failures": "correlated rack failures" if correlated else "independent", "unanimity_10": float(all10.mean()),
                 "majority_3_of_5": float(maj.mean())})
for r in formula: print(f"n = {r['n']:>2}: all up {r['unanimity']:.6f}   majority up {r['majority']:.10f}")
for r in sims: print(f"simulated, {r['failures']:<26} 10 all up {r['unanimity_10']:.5f}   3 of 5 up {r['majority_3_of_5']:.8f}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"formula": formula, "simulated": sims}}, f)