"""A metrics aggregator: percentiles out of bounded state, two ways.

Two estimators run over one stream at an identical memory budget:

  Reservoir  -- k slots holding a uniform random sample of the stream
                (Vitter's Algorithm R). Randomised O(1).
  Histogram  -- k fixed counters with Prometheus `le` semantics.
                Counted O(1).

Plus the exact reference (`exact_pct`), the stream source (`latency_trace`),
and the order-statistic arithmetic (`beta_inv`, `beta_sd`) that lets the
commentary *derive* the reservoir's error instead of merely observing it.

Determinism: every source of randomness is an explicit `random.Random(seed)`
handed in by the caller. No wall clock, no builtin `hash()`, no module-level
RNG -- so a named number on the commentary page reproduces byte for byte, and
"only the sampler's seed changed" is a statement you can actually make.
"""
import bisect
import math
import random


# --------------------------------------------------------------- the truth

def exact_pct(sorted_xs, q):
    """Nearest-rank percentile: the ceil(q*n)-th smallest value.

    No interpolation, deliberately. The obvious worry is that `ceil` is what
    makes the tail jumpy; the commentary measures the linear-interpolation
    variant and gets 30.80x against 30.96x, so the choice is decorative.
    """
    n = len(sorted_xs)
    if n == 0:
        return float("nan")
    r = max(1, math.ceil(q * n))
    return sorted_xs[min(r, n) - 1]


# -------------------------------------------------------------- the stream

def latency_trace(rng, n, median=0.020, sigma=0.5,
                  p_tail=0.02, tail_median=0.200, tail_sigma=0.8):
    """A lognormal body plus a lognormal slow path. Seconds.

    `tail_sigma` is the load-bearing knob of the whole toy: it changes the
    SLOPE of the quantile function past u=0.95 while barely moving the true
    p99, and the reservoir's reported spread moves 30x with it.
    """
    mu_body, mu_tail = math.log(median), math.log(tail_median)
    out = []
    for _ in range(n):
        if rng.random() < p_tail:
            out.append(rng.lognormvariate(mu_tail, tail_sigma))
        else:
            out.append(rng.lognormvariate(mu_body, sigma))
    return out


# ------------------------------------------------------- estimator 1: sample

class Reservoir:
    """Vitter's Algorithm R. k slots, one pass, O(k) memory, and the buffer is
    a uniform random k-subset of everything seen so far."""

    def __init__(self, k, seed):
        self.k = k
        self.rng = random.Random(seed)
        self.buf = []
        self.n = 0

    def add(self, x):
        if len(self.buf) < self.k:
            self.buf.append(x)
        else:
            # Keep x with probability k/(n+1) -- the invariant that makes the
            # buffer uniform. The +1 is what counts x itself as a candidate;
            # `randrange(self.n)` would bias the sample towards early arrivals.
            j = self.rng.randrange(self.n + 1)
            if j < self.k:
                self.buf[j] = x
        self.n += 1

    def pct(self, q):
        return exact_pct(sorted(self.buf), q)


def uniform_subset(rng, sorted_xs, k):
    """Draw what `Reservoir`'s final buffer IS -- a uniform k-subset -- in
    O(k) rather than O(n). Returns it sorted.

    This is not an approximation of Algorithm R; it is the theorem Algorithm R
    exists to satisfy, sampled directly. A 2000-seed sweep over a million-
    sample trace is 2e9 `add` calls the long way round and 2e5 the short way.
    `test_aggregator.py` pins the two against each other.
    """
    n = len(sorted_xs)
    if k >= n:
        return list(sorted_xs)
    return [sorted_xs[i] for i in sorted(rng.sample(range(n), k))]


# ----------------------------------------------------- estimator 2: counters

class Histogram:
    """A Prometheus-style cumulative histogram: finite `le` upper bounds plus
    a final +Inf bucket. `quantile` reproduces promql's histogram_quantile --
    linear interpolation inside the bucket the rank lands in, and, the
    load-bearing bit, the HIGHEST FINITE BOUND when that bucket is +Inf, with
    no indication anywhere that it clamped."""

    def __init__(self, bounds):
        self.bounds = list(bounds)             # ascending finite upper bounds
        self.counts = [0] * (len(bounds) + 1)  # last slot is the +Inf bucket

    def add(self, x):
        # `le` semantics: bucket i counts x <= bounds[i], so bisect_left.
        self.counts[bisect.bisect_left(self.bounds, x)] += 1

    def merge(self, other):
        # Counters compose; summaries do not. Three lines, because merging
        # counts is genuinely trivial -- see commentary section 7.3.
        assert self.bounds == other.bounds
        for i, c in enumerate(other.counts):
            self.counts[i] += c

    @property
    def total(self):
        return sum(self.counts)

    def quantile(self, q):
        total = self.total
        if total == 0:
            return float("nan")
        rank = q * total
        cum = 0
        for i, c in enumerate(self.counts):
            cum += c
            if cum >= rank:
                break
        if i == len(self.bounds):              # the rank fell in +Inf
            return self.bounds[-1]             # <-- the silent clamp
        lo = 0.0 if i == 0 else self.bounds[i - 1]
        prev = cum - self.counts[i]
        frac = (rank - prev) / self.counts[i] if self.counts[i] else 0.0
        return lo + (self.bounds[i] - lo) * frac

    def in_inf_bucket(self, q):
        """Whether `quantile(q)` clamped. Cheap to know and never reported --
        Prometheus has no equivalent, which is the point of section 7.5."""
        rank = q * self.total
        cum = 0
        for i, c in enumerate(self.counts):
            cum += c
            if cum >= rank:
                return i == len(self.bounds)
        return True


# Prometheus client_golang's DefBuckets, in seconds. 13 bounds + 1 for +Inf.
DEF_BUCKETS = [.005, .01, .025, .05, .075, .1, .25, .5, .75, 1, 2.5, 5, 10]


# ------------------------------------------- where a sampled slot really sits

def beta_cdf(x, j, k):
    """P(U_(j) <= x) for the j-th smallest of k independent uniforms: the
    chance that at least j of the k land below x."""
    return sum(math.comb(k, i) * x ** i * (1 - x) ** (k - i)
               for i in range(j, k + 1))


def beta_inv(p, j, k):
    """The p-quantile of slot j's POSITION, by bisection on `beta_cdf`.

    Nearest-rank pQ of a k-slot sample reads slot j = ceil(Q*k), and that slot
    sits not at quantile Q of the true distribution but at a RANDOM quantile
    U_(j) ~ Beta(j, k+1-j). Feed these positions to the trace's own quantile
    function and you predict the reported percentile before measuring it.
    """
    lo, hi = 0.0, 1.0
    for _ in range(200):
        mid = (lo + hi) / 2
        if beta_cdf(mid, j, k) < p:
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2


def beta_sd(j, k):
    """sd of U_(j) ~ Beta(j, k+1-j), in units of quantile.

    The one number that falsifies "the tail slot is sampled sloppily": at
    k=100 this is 0.0138 for slot 99 against 0.0495 for slot 50.
    """
    return math.sqrt(j * (k + 1 - j) / ((k + 1) ** 2 * (k + 2)))


def ms(x):
    """Seconds -> a fixed-width millisecond column."""
    return f"{x * 1000:8.2f}ms"