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

These pin every number the commentary quotes -- the 35.50x reservoir spread,
the 3.59x precision ratio that falsifies the usual explanation, the histogram's
1.000x at the same memory budget, the load-bearing knob, the inert ones, the
boundary, and the bound that rules out the fleet failure mode people expect.
They also pin the equivalence that makes the sweeps affordable: Algorithm R's
final buffer is a uniform k-subset, so `uniform_subset` may stand in for it.
"""
import math
import random
import statistics

from aggregator import (DEF_BUCKETS, Histogram, Reservoir, beta_inv, beta_sd,
                        exact_pct, latency_trace, uniform_subset)

N = 1_000_000
K = 100
SEEDS = 2000
TRACE_SEED = 20260803

# One trace, built once and shared: this is what "only the sampler's seed
# changed" means, and rebuilding it per test would quietly break the claim.
TRACE = sorted(latency_trace(random.Random(TRACE_SEED), N))
TRUE_P50 = exact_pct(TRACE, 0.50)
TRUE_P99 = exact_pct(TRACE, 0.99)

def sweep(srt, q, k=K, seeds=SEEDS, base=1000):
    """The reported percentile over `seeds` sampler seeds, sorted."""
    return sorted(exact_pct(uniform_subset(random.Random(base + s), srt, k), q)
                  for s in range(seeds))


def r2(x):
    return round(x * 1000, 2)


def test_exact_pct_is_nearest_rank():
    """ceil(q*n)-th smallest, no interpolation, and it never runs off the end."""
    xs = list(range(1, 101))                      # 1..100
    assert exact_pct(xs, 0.50) == 50              # ceil(50) -> slot 50
    assert exact_pct(xs, 0.99) == 99              # ceil(99) -> slot 99
    assert exact_pct(xs, 1.00) == 100
    assert exact_pct(xs, 0.0) == 1                # the max(1, ...) floor
    assert exact_pct([7], 0.99) == 7
    assert math.isnan(exact_pct([], 0.5))
    # p99 of a 100-slot sample reads slot 99, NOT slot 100. That off-by-one is
    # the whole of section 5.4: the estimator is asked for 0.99 and reads 0.98.
    assert math.ceil(0.99 * 100) == 99


def test_true_percentiles_of_the_shared_trace():
    assert r2(TRUE_P50) == 20.25, r2(TRUE_P50)
    assert r2(TRUE_P99) == 200.22, r2(TRUE_P99)


def test_the_headline_spread_on_one_unchanging_trace():
    """2000 sampler seeds, one trace: the p99 spans 35.50x, the p50 1.54x."""
    p99 = sweep(TRACE, 0.99)
    p50 = sweep(TRACE, 0.50)
    assert r2(p99[0]) == 38.14, r2(p99[0])
    assert r2(p99[-1]) == 1353.92, r2(p99[-1])
    assert round(p99[-1] / p99[0], 2) == 35.50
    assert r2(p50[0]) == 16.53 and r2(p50[-1]) == 25.45
    assert round(p50[-1] / p50[0], 2) == 1.54


def test_the_median_window_reports_less_than_half_the_true_p99():
    """No error bar, no outlier: the MIDDLE reading is 0.487x the truth."""
    median_reading = exact_pct(sweep(TRACE, 0.99), 0.50)
    assert r2(median_reading) == 97.47, r2(median_reading)
    assert round(median_reading / TRUE_P99, 3) == 0.487
    # The p50's median reading, by contrast, is dead on.
    assert round(exact_pct(sweep(TRACE, 0.50), 0.50) / TRUE_P50, 3) == 0.995


def test_the_p99_slot_is_located_more_tightly_than_the_p50_slot():
    """Falsifies "the tail is sampled sloppily". The sampler is 3.59x MORE
    precise about where slot 99 sits than about where slot 50 sits."""
    sd50, sd99 = beta_sd(50, K), beta_sd(99, K)
    assert round(sd50, 6) == 0.049505, sd50
    assert round(sd99, 6) == 0.013795, sd99
    assert round(sd50 / sd99, 2) == 3.59
    # Closed form against the definition, sqrt(j(k+1-j)/((k+1)^2 (k+2))).
    assert abs(sd99 - math.sqrt(99 * 2 / (101 ** 2 * 102))) < 1e-15
    # And the p99 slot does not sit at u=0.99. It sits at 99/101.
    assert round(beta_inv(0.50, 99, K), 4) == 0.9833
    assert round(99 / 101, 6) == 0.980198


def test_the_damage_is_the_quantile_function_not_the_sampler():
    """The p50 slot roams 4x further and moves the answer 6.5x less."""
    def Q(u):
        return TRACE[max(1, min(N, math.ceil(u * N))) - 1]

    lo50, hi50 = beta_inv(0.05, 50, K), beta_inv(0.95, 50, K)
    lo99, hi99 = beta_inv(0.05, 99, K), beta_inv(0.95, 99, K)
    assert round((hi50 - lo50) * 100, 2) == 16.30      # percentile-points
    assert round((hi99 - lo99) * 100, 2) == 4.30
    assert round(Q(hi50) / Q(lo50), 2) == 1.23         # value swing
    assert round(Q(hi99) / Q(lo99), 2) == 8.02


def test_beta_pushed_through_Q_predicts_the_measurement():
    """The headline is derived, not just observed: predicted within 1.8%."""
    def Q(u):
        return TRACE[max(1, min(N, math.ceil(u * N))) - 1]

    for j, q in ((50, 0.50), (99, 0.99)):
        for p in (0.05, 0.50, 0.95):
            predicted = Q(beta_inv(p, j, K))
            measured = exact_pct(sweep(TRACE, q), p)
            assert abs(predicted / measured - 1) < 0.02, (j, p, predicted, measured)
    assert r2(Q(beta_inv(0.95, 99, K))) == 415.75      # the number on the page
    assert r2(exact_pct(sweep(TRACE, 0.99), 0.95)) == 421.12


def test_algorithm_r_buffer_is_a_uniform_k_subset():
    """The equivalence that makes the sweeps affordable. Same stream, 2000
    seeds each: Algorithm R and `uniform_subset` agree within noise."""
    stream = latency_trace(random.Random(TRACE_SEED), 5_000)
    srt = sorted(stream)
    real, short = [], []
    for s in range(2000):
        res = Reservoir(K, s)
        for x in stream:
            res.add(x)
        real.append(res.pct(0.99))
        short.append(exact_pct(uniform_subset(random.Random(90_000 + s), srt, K),
                               0.99))
    diff = statistics.mean(real) - statistics.mean(short)
    se = math.sqrt(statistics.variance(real) / 2000
                   + statistics.variance(short) / 2000)
    assert abs(diff) / se < 2.5, (diff, se, abs(diff) / se)
    # Structural checks that do not depend on sampling luck.
    res = Reservoir(K, 7)
    for x in stream:
        res.add(x)
    assert len(res.buf) == K and res.n == 5_000
    assert set(res.buf) <= set(stream)


def test_algorithm_r_lands_inside_the_swept_range():
    """Three real one-pass runs over all 1,000,000 samples, for honesty."""
    p99 = sweep(TRACE, 0.99)
    for seed in (0, 1, 2):
        res = Reservoir(K, seed)
        for x in TRACE:
            res.add(x)
        assert p99[0] <= res.pct(0.99) <= p99[-1], (seed, res.pct(0.99))


def test_the_same_memory_spent_on_buckets_is_right_and_deterministic():
    """100 counters instead of 100 slots: 1.000x true, and no seed anywhere."""
    bounds = [0.001 * 10 ** (i * 5 / (K - 1)) for i in range(K)]
    first = Histogram(bounds)
    second = Histogram(bounds)
    for x in TRACE:
        first.add(x)
    for x in reversed(TRACE):                 # order cannot matter to counters
        second.add(x)
    assert first.counts == second.counts
    assert r2(first.quantile(0.99)) == 200.26, r2(first.quantile(0.99))
    assert round(first.quantile(0.99) / TRUE_P99, 3) == 1.000
    assert r2(first.quantile(0.50)) == 20.26
    assert first.total == N and not first.in_inf_bucket(0.99)


def test_prometheus_defbuckets_gets_within_1_06x_on_a_seventh_of_the_memory():
    h = Histogram(DEF_BUCKETS)
    for x in TRACE:
        h.add(x)
    assert len(h.counts) == 14
    assert r2(h.quantile(0.99)) == 213.03, r2(h.quantile(0.99))
    assert round(h.quantile(0.99) / TRUE_P99, 2) == 1.06


def test_the_inf_bucket_clamps_to_the_last_finite_bound_and_says_nothing():
    """The predicted failure: a true p99 of 1,013,571ms reported as 10,000ms."""
    trace = latency_trace(random.Random(TRACE_SEED), 200_000, tail_median=1000.0)
    h = Histogram(DEF_BUCKETS)
    for x in trace:
        h.add(x)
    assert r2(exact_pct(sorted(trace), 0.99)) == 1013571.75
    assert r2(h.quantile(0.99)) == 10000.00
    assert h.quantile(0.99) == DEF_BUCKETS[-1]      # the highest finite bound
    assert h.in_inf_bucket(0.99)                    # knowable, but not reported


def test_tail_shape_is_the_load_bearing_knob():
    """k, n, the sampler and the true answer all held fixed. Only Q(u)'s slope
    past u=0.95 changes: the true p99 moves 2.6%, the spread moves 29.6x."""
    out = {}
    for ts in (0.1, 1.6):
        t = sorted(latency_trace(random.Random(TRACE_SEED), 200_000,
                                 tail_sigma=ts))
        v = sweep(t, 0.99)
        out[ts] = (exact_pct(t, 0.99), v[-1] / v[0])
    assert r2(out[0.1][0]) == 200.34 and r2(out[1.6][0]) == 205.48
    assert round(out[1.6][0] / out[0.1][0], 3) == 1.026        # true p99: +2.6%
    assert round(out[0.1][1], 2) == 6.39, out[0.1][1]
    assert round(out[1.6][1], 2) == 188.98, out[1.6][1]


def test_stream_length_is_inert():
    """The reader's first guess. 100x more data moves the spread 9.05x -> 11.24x
    -- and in the wrong direction."""
    out = {}
    for n in (5_000, 500_000):
        rng = random.Random(TRACE_SEED)
        t = sorted(rng.lognormvariate(math.log(0.020), 1.0) for _ in range(n))
        v = sweep(t, 0.99)
        out[n] = v[-1] / v[0]
    assert round(out[5_000], 2) == 9.05, out[5_000]
    assert round(out[500_000], 2) == 11.24, out[500_000]


def test_the_percentile_definition_is_decorative():
    """The other first guess: `ceil` is not what makes the tail jumpy."""
    def interp_pct(s, q):
        pos = q * (len(s) - 1)
        lo = math.floor(pos)
        hi = min(lo + 1, len(s) - 1)
        return s[lo] + (pos - lo) * (s[hi] - s[lo])

    t = sorted(latency_trace(random.Random(TRACE_SEED), 200_000))
    out = {}
    for name, f in (("nearest", exact_pct), ("interp", interp_pct)):
        v = [f(uniform_subset(random.Random(1000 + s), t, K), 0.99)
             for s in range(3000)]
        out[name] = max(v) / min(v)
    assert round(out["nearest"], 2) == 30.96, out["nearest"]
    assert round(out["interp"], 2) == 30.80, out["interp"]


def test_memory_buys_the_spread_back_slowly():
    """k IS partially load-bearing -- and 256x the memory only gets 9.27x to
    1.13x, on the very memory the technique exists to save."""
    rng = random.Random(TRACE_SEED)
    t = sorted(rng.lognormvariate(math.log(0.020), 1.0) for _ in range(200_000))
    small = sweep(t, 0.99, k=100)
    large = sweep(t, 0.99, k=25_600, seeds=400)
    assert round(small[-1] / small[0], 2) == 9.27, small[-1] / small[0]
    assert round(large[-1] / large[0], 2) == 1.13, large[-1] / large[0]


def test_the_boundary_a_flat_tail_kills_the_effect_entirely():
    """Where the result stops applying: no tail, no instability."""
    const = [0.020] * 200_000
    v = sweep(const, 0.99)
    assert v[0] == v[-1] == 0.020                # exactly 1.00x, all 2000 seeds
    rng = random.Random(TRACE_SEED)
    flat = sorted(rng.uniform(0.010, 0.030) for _ in range(200_000))
    w = sweep(flat, 0.99)
    assert round(w[-1] / w[0], 2) == 1.08, w[-1] / w[0]
    # And the ratio that predicts it: Q over the p99 slot's own 5-95 window.
    lo_u, hi_u = beta_inv(0.05, 99, K), beta_inv(0.95, 99, K)
    assert round(exact_pct(flat, hi_u) / exact_pct(flat, lo_u), 3) == 1.029
    assert exact_pct(const, hi_u) / exact_pct(const, lo_u) == 1.0


def test_pooled_p99_can_never_exceed_the_largest_per_host_p99():
    """The falsification. At most 0.01*n_i of each host's samples exceed that
    host's own p99, so at most 0.01*N pooled samples exceed max_i p99_i -- so
    the pooled p99 sits at or below it. 120 fleets, 0 violations."""
    tested = violations = 0
    for seed in range(40):
        for n_sick in (0, 1, 7):
            rng = random.Random(seed)
            traces = []
            for h in range(20):
                if h < n_sick:
                    traces.append([rng.lognormvariate(math.log(0.5), 0.6)
                                   for _ in range(500)])
                else:
                    traces.append(latency_trace(rng, 500))
            pooled = exact_pct(sorted(x for t in traces for x in t), 0.99)
            max_p99 = max(exact_pct(sorted(t), 0.99) for t in traces)
            tested += 1
            violations += pooled > max_p99 + 1e-12
    assert tested == 120 and violations == 0, (tested, violations)


def test_counters_merge_but_summaries_do_not():
    """Why the histogram is the one you can fan in: merge is exact by
    construction, and merging then reading is not the same as reading twice."""
    a, b = Histogram(DEF_BUCKETS), Histogram(DEF_BUCKETS)
    left = latency_trace(random.Random(1), 20_000)
    right = [x * 20 for x in latency_trace(random.Random(2), 20_000)]
    for x in left:
        a.add(x)
    for x in right:
        b.add(x)
    a.merge(b)
    whole = Histogram(DEF_BUCKETS)
    for x in left + right:
        whole.add(x)
    assert a.counts == whole.counts               # merging counts is exact
    assert a.total == 40_000


TESTS = [
    test_exact_pct_is_nearest_rank,
    test_true_percentiles_of_the_shared_trace,
    test_the_headline_spread_on_one_unchanging_trace,
    test_the_median_window_reports_less_than_half_the_true_p99,
    test_the_p99_slot_is_located_more_tightly_than_the_p50_slot,
    test_the_damage_is_the_quantile_function_not_the_sampler,
    test_beta_pushed_through_Q_predicts_the_measurement,
    test_algorithm_r_buffer_is_a_uniform_k_subset,
    test_algorithm_r_lands_inside_the_swept_range,
    test_the_same_memory_spent_on_buckets_is_right_and_deterministic,
    test_prometheus_defbuckets_gets_within_1_06x_on_a_seventh_of_the_memory,
    test_the_inf_bucket_clamps_to_the_last_finite_bound_and_says_nothing,
    test_tail_shape_is_the_load_bearing_knob,
    test_stream_length_is_inert,
    test_the_percentile_definition_is_decorative,
    test_memory_buys_the_spread_back_slowly,
    test_the_boundary_a_flat_tail_kills_the_effect_entirely,
    test_pooled_p99_can_never_exceed_the_largest_per_host_p99,
    test_counters_merge_but_summaries_do_not,
]


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")