"""Stdlib-only tests (no pytest): plain asserts in functions called from a
__main__ block. Run: `python3 test_gossip.py`.

Every number asserted here also appears in commentary.html, so if the page
rots these fail. The suite takes about ten seconds because the headline
figures are means over 200 seeded runs of a 1024-node cluster — that
ensemble *is* the result, so it is worth the wall time. The 200 curves are
built once at import and shared by the tests that need them.
"""

import math
import statistics as st

from gossip import (MODES, band_costs, callers, first_round_at, schedule,
                    simulate)

N = 1024
SEEDS = 200
MAX_R = 40
BANDS = ((0, .25), (.25, .5), (.5, .75), (.75, .9), (.9, .99), (.99, 1.0))

CURVES = {m: [] for m in MODES}
for _seed in range(SEEDS):
    _sched = schedule(N, 1, MAX_R, _seed)
    for _mode in MODES:
        _res = simulate(N, _sched, mode=_mode)
        assert _res["converged"], (_mode, _seed)
        CURVES[_mode].append(_res["curve"])


def mean_rounds(mode):
    return st.mean(len(c) - 1 for c in CURVES[mode])


def test_the_schedule_is_reproducible_and_never_self_addressed():
    """No wall clock, one seed: this plan is fixed forever on any box. A
    node never draws itself, because randrange(n-1) skips over i."""
    plan = schedule(8, 2, 3, 0)
    assert plan == schedule(8, 2, 3, 0)
    assert plan != schedule(8, 2, 3, 1)
    assert plan[0][0] == (7, 4)
    for row in plan:
        for i, peers in enumerate(row):
            assert len(peers) == 2
            assert all(0 <= j < 8 and j != i for j in peers)


def test_all_three_modes_replay_one_identical_schedule():
    """The whole comparison rests on this: same seed, same arrows, only the
    direction of travel differs. Seed 0's residual traces, verbatim."""
    sched = schedule(N, 1, MAX_R, 0)
    left = {m: [N - c for c in simulate(N, sched, mode=m)["curve"]]
            for m in MODES}
    assert left["push"][-8:] == [47, 17, 5, 2, 1, 1, 1, 0]
    assert left["pull"][-4:] == [165, 22, 1, 0]
    assert left["pushpull"][-3:] == [165, 18, 0]
    assert (len(left["push"]), len(left["pull"]), len(left["pushpull"])) == (
        21, 14, 10)


def test_o_log_n_is_true_and_flat():
    """The claim a reader already predicts: r100/log2(n) is flat to within
    2% over a 128x range of n. 200 seeds per size."""
    ratios = {}
    for n in (8, 128, N):
        rounds = [simulate(n, schedule(n, 1, MAX_R, s), mode="push")["rounds"]
                  for s in range(SEEDS)]
        ratios[n] = st.mean(rounds) / math.log2(n)
    assert round(ratios[8], 3) == 1.820
    assert round(ratios[128], 3) == 1.843
    assert round(ratios[N], 3) == 1.820


def test_half_the_cluster_costs_less_than_the_last_node():
    """The aha, in rounds: 10.01 to reach 512 of 1024 nodes, 8.19 more for
    the last one — 45% of the wall clock after half the cluster knows."""
    r50 = st.mean(first_round_at(c, N, 0.5) for c in CURVES["push"])
    r100 = mean_rounds("push")
    assert round(r50, 2) == 10.01
    assert round(r100, 2) == 18.20
    assert round(r100 - r50, 2) == 8.19
    assert round(100 * (r100 - r50) / r100, 1) == 45.0


def test_the_endgame_burns_two_thirds_of_the_message_budget():
    """The aha, in messages: rounds beginning at >=90% coverage spend 63.6%
    of everything to inform 53.4 of 1023 nodes, at 396x the unit price of
    the first quarter."""
    calls, gained = band_costs(CURVES["push"], N, 1, BANDS)
    total = sum(calls.values())
    late = calls[(.9, .99)] + calls[(.99, 1.0)]
    late_nodes = (gained[(.9, .99)] + gained[(.99, 1.0)]) / SEEDS
    assert round(100 * late / total, 1) == 63.6
    assert round(late_nodes, 1) == 53.4
    first = calls[(0, .25)] / gained[(0, .25)]
    last = calls[(.99, 1.0)] / gained[(.99, 1.0)]
    assert (round(first, 1), round(last, 1)) == (1.2, 492.2)
    assert round(last / first) == 396


def test_band_costs_conserves_both_columns():
    """Every call and every newly-informed node lands in exactly one band —
    otherwise the 63.6% is an artefact of dropped rows."""
    calls, gained = band_costs(CURVES["push"], N, 1, BANDS)
    assert sum(gained.values()) == SEEDS * (N - 1)
    charged = sum(
        simulate(N, schedule(N, 1, MAX_R, s), mode="push", lazy=True)["calls"]
        for s in range(SEEDS))
    assert sum(calls.values()) == charged
    assert round(sum(calls.values()) / SEEDS, 1) == 8293.1


def test_push_divides_the_residual_by_e_and_pull_squares_it():
    """The mechanism behind the two cost curves, pooled over 200 seeds:
    push's s'/s sits near 1/e for small residuals; pull's s' tracks the
    predicted s*s/n."""
    def pooled(mode, lo, hi):
        return [(N - c[r], N - c[r + 1]) for c in CURVES[mode]
                for r in range(len(c) - 1) if lo <= N - c[r] < hi]

    push = pooled("push", 8, 16)
    ratio = st.mean(b for _, b in push) / st.mean(a for a, _ in push)
    assert round(ratio, 3) == 0.365
    assert abs(ratio - 1 / math.e) < 0.01
    for lo, hi in ((32, 64), (64, 128)):
        pull = pooled("pull", lo, hi)
        got = st.mean(b for _, b in pull)
        pred = st.mean(a for a, _ in pull) ** 2 / N
        assert abs(got - pred) / pred < 0.06, (lo, hi, got, pred)


def test_pull_clears_the_last_node_in_exactly_one_round():
    """Derived: with one node left, all 1023 pushes miss it with probability
    0.999022^1023 = 0.3677 (1/e = 0.3679), so push expects 1.582 rounds
    there. Pull's last node calls a knower with probability 1.000 — and no
    pull run in 200 ever spends more than one round at residual 1."""
    stuck = (1 - 1 / (N - 1)) ** (N - 1)
    assert round(stuck, 4) == 0.3677
    assert abs(stuck - 1 / math.e) < 0.0003
    assert round(1 / (1 - stuck), 3) == 1.582
    occ = {m: [sum(1 for v in c[:-1] if N - v == 1) for c in CURVES[m]]
           for m in MODES}
    push_hits = [o for o in occ["push"] if o]
    pull_hits = [o for o in occ["pull"] if o]
    assert (len(push_hits), len(pull_hits)) == (128, 39)
    assert round(st.mean(push_hits), 3) == 1.688
    assert max(pull_hits) == 1


def test_lazy_accounting_changes_the_bill_and_not_the_outcome():
    """The steelman for push is pure accounting: a node with nothing to say
    places no call, which could not have taught anyone anything anyway."""
    sched = schedule(N, 1, MAX_R, 0)
    naive = simulate(N, sched, mode="push")
    lazy = simulate(N, sched, mode="push", lazy=True)
    assert naive["curve"] == lazy["curve"]
    assert naive["calls"] == N * naive["rounds"] == 20480
    assert lazy["calls"] == sum(naive["curve"][:-1]) == 10196
    pp = simulate(N, sched, mode="pushpull", lazy=True)
    assert pp["calls"] == N * pp["rounds"]        # push-pull cannot be lazy


def test_callers_is_the_whole_charging_model():
    assert callers("push", 100, 7) == 100
    assert callers("push", 100, 7, lazy=True) == 7
    assert callers("pull", 100, 7, lazy=True) == 93
    assert callers("pushpull", 100, 7, lazy=True) == 100


def test_lazy_push_at_fanout_3_beats_push_pull_at_fanout_1():
    """The boundary that kills 'push-pull is free': under honest accounting
    push at fanout 3 wins on rounds AND on messages. 100 seeds."""
    seeds = 100
    rounds, calls = {}, {}
    for mode, fanout, cap, lazy in (("push", 3, 20, True),
                                    ("pushpull", 1, MAX_R, False)):
        runs = [simulate(N, schedule(N, fanout, cap, s), mode=mode, lazy=lazy)
                for s in range(seeds)]
        assert all(r["converged"] for r in runs)
        rounds[mode] = st.mean(r["rounds"] for r in runs)
        calls[mode] = st.mean(r["calls"] for r in runs)
    assert round(rounds["push"], 2) == 8.21
    assert round(rounds["pushpull"], 2) == 9.18
    assert round(calls["push"]) == 9285
    assert round(calls["pushpull"]) == 9400
    assert rounds["push"] < rounds["pushpull"]
    assert calls["push"] < calls["pushpull"]


def test_push_is_faster_than_pull_at_half_coverage():
    """The other boundary: the arrow direction only pays in the endgame.
    At 50% coverage push is ahead; it loses by 1.309x at 100%."""
    def ratio(frac):
        a = st.mean(first_round_at(c, N, frac) for c in CURVES["push"])
        b = st.mean(first_round_at(c, N, frac) for c in CURVES["pull"])
        return a / b

    assert round(ratio(0.5), 3) == 0.958
    assert round(ratio(0.9), 3) == 1.064
    assert round(ratio(1.0), 3) == 1.309


def test_the_snapshot_line_is_load_bearing_but_the_aha_survives_it():
    """Reading `inf` live instead of the start-of-round copy speeds every
    mode up and *widens* the push/push-pull gap, so the line sets the
    constants and not the result. 50 seeds here to keep the suite quick;
    the 200-seed means (18.20 -> 14.56, 9.18 -> 6.38) are in the
    commentary, from checks/cf1_counterfactuals.py."""
    seeds, snap, live = 50, {}, {}
    for mode in MODES:
        a, b = [], []
        for s in range(seeds):
            sched = schedule(N, 1, MAX_R, s)
            a.append(simulate(N, sched, mode=mode)["rounds"])
            b.append(simulate(N, sched, mode=mode, live=True)["rounds"])
        snap[mode], live[mode] = st.mean(a), st.mean(b)
    assert (round(snap["push"], 2), round(live["push"], 2)) == (17.98, 14.62)
    assert (round(snap["pushpull"], 2), round(live["pushpull"], 2)) == (
        9.20, 6.36)
    assert live["push"] / live["pushpull"] > snap["push"] / snap["pushpull"]
    assert round(snap["push"] / snap["pushpull"], 2) == 1.95
    assert round(live["push"] / live["pushpull"], 2) == 2.30


def test_a_run_that_hits_max_rounds_reports_non_convergence():
    """The cap is hard: a truncated run comes back converged=False with the
    coverage it managed, instead of looping until the epidemic finishes."""
    sched = schedule(N, 1, MAX_R, 0)
    short = simulate(N, sched, mode="push", max_rounds=5)
    assert short["converged"] is False
    assert short["rounds"] == 5
    assert short["known"] == 32 < N
    assert short["curve"] == [1, 2, 4, 8, 16, 32]
    assert simulate(N, sched, mode="push", max_rounds=999)["rounds"] == 20


def test_first_round_at_reports_the_crossing_and_the_miss():
    curve = [1, 2, 4, 250, 700, 1024]
    assert first_round_at(curve, N, 0.25) == 4
    assert first_round_at(curve, N, 1.0) == 5
    assert first_round_at(curve[:3], N, 0.5) is None


def test_any_start_node_works_and_a_bad_mode_is_refused():
    sched = schedule(N, 1, MAX_R, 7)
    assert simulate(N, sched, mode="push", start=511)["converged"]
    try:
        simulate(N, sched, mode="gossip")
    except ValueError as exc:
        assert "mode must be one of" in str(exc)
    else:
        raise AssertionError("simulate accepted an unknown mode")


TESTS = [
    test_the_schedule_is_reproducible_and_never_self_addressed,
    test_all_three_modes_replay_one_identical_schedule,
    test_o_log_n_is_true_and_flat,
    test_half_the_cluster_costs_less_than_the_last_node,
    test_the_endgame_burns_two_thirds_of_the_message_budget,
    test_band_costs_conserves_both_columns,
    test_push_divides_the_residual_by_e_and_pull_squares_it,
    test_pull_clears_the_last_node_in_exactly_one_round,
    test_lazy_accounting_changes_the_bill_and_not_the_outcome,
    test_callers_is_the_whole_charging_model,
    test_lazy_push_at_fanout_3_beats_push_pull_at_fanout_1,
    test_push_is_faster_than_pull_at_half_coverage,
    test_the_snapshot_line_is_load_bearing_but_the_aha_survives_it,
    test_a_run_that_hits_max_rounds_reports_non_convergence,
    test_first_round_at_reports_the_crossing_and_the_miss,
    test_any_start_node_works_and_a_bad_mode_is_refused,
]


if __name__ == "__main__":
    failed = 0
    for test in TESTS:
        try:
            test()
        except AssertionError as exc:
            failed += 1
            print(f"FAIL  {test.__name__}: {exc}")
        else:
            print(f"PASS  {test.__name__}")
    if failed:
        print(f"\n{failed} of {len(TESTS)} tests FAILED")
        raise SystemExit(1)
    print(f"\nAll {len(TESTS)} tests PASSED")