Computing tumbling and sliding windows over telemetry streams

import math
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta

EPOCH = datetime(1970, 1, 1)


@dataclass
class Welford:
    """Numerically stable mean and variance in constant space.

    A naive sum-of-squares accumulator loses precision badly once the values
    are large and similar — which is exactly the case for cumulative distance
    or UTM coordinates. Welford's method costs one extra multiply and does not.
    """
    n: int = 0
    mean: float = 0.0
    m2: float = 0.0
    vmin: float = math.inf
    vmax: float = -math.inf

    def add(self, x: float) -> None:
        self.n += 1
        d = x - self.mean
        self.mean += d / self.n
        self.m2 += d * (x - self.mean)
        self.vmin = min(self.vmin, x)
        self.vmax = max(self.vmax, x)

    @property
    def variance(self) -> float:
        return self.m2 / (self.n - 1) if self.n > 1 else 0.0


@dataclass
class SlidingWindower:
    """
    Event-time sliding windows with bounded state and epoch-anchored edges.

    Parameters
    ----------
    width, slide : timedelta
        width must be an exact multiple of slide, or windows do not tile and
        an event's membership depends on where the stream happened to start.
    allowed_lateness : timedelta
        How long a window stays open past its end before it is sealed.
    """
    width: timedelta
    slide: timedelta
    allowed_lateness: timedelta
    state: dict = field(default_factory=lambda: defaultdict(Welford))
    watermark: datetime | None = None

    def __post_init__(self):
        if self.slide <= timedelta(0) or self.width <= timedelta(0):
            raise ValueError("width and slide must be positive.")
        if self.width % self.slide != timedelta(0):
            raise ValueError(
                f"width {self.width} is not a multiple of slide {self.slide}; "
                "windows would not tile consistently."
            )

    def _windows_for(self, t: datetime) -> list[datetime]:
        """Every window start whose span contains t. Anchored to the epoch."""
        epoch = EPOCH.replace(tzinfo=t.tzinfo)
        k = (t - epoch) // self.slide
        n_open = self.width // self.slide
        return [epoch + (k - i) * self.slide for i in range(n_open)
                if epoch + (k - i) * self.slide + self.width > t]

    def push(self, entity: str, event_time: datetime, value: float) -> list[tuple]:
        if event_time.tzinfo is None:
            raise ValueError("event_time must be timezone-aware (UTC).")
        if self.watermark is None or event_time - self.allowed_lateness > self.watermark:
            self.watermark = event_time - self.allowed_lateness

        for w in self._windows_for(event_time):
            if w + self.width + self.allowed_lateness > self.watermark:
                self.state[(entity, w)].add(value)

        return self._emit_sealed()

    def _emit_sealed(self) -> list[tuple]:
        """Emit and REMOVE every window now behind the watermark."""
        if self.watermark is None:
            return []
        out, done = [], []
        for (ent, w), acc in self.state.items():
            if w + self.width <= self.watermark:
                out.append((ent, w, acc.n, acc.mean, acc.variance, acc.vmin, acc.vmax))
                done.append((ent, w))
        for k in done:
            del self.state[k]        # eviction is not optional
        return out

    def open_windows(self) -> int:
        """Expose the state size as a metric; alert on its growth."""
        return len(self.state)

def validate_windower(w: SlidingWindower, results: list, expected_per_entity: int) -> None:
    """Assert the three properties that distinguish a correct windower."""
    # 1. Every entity emitted the same number of windows over a full stream.
    from collections import Counter
    per_entity = Counter(r[0] for r in results)
    assert len(set(per_entity.values())) == 1, (
        f"uneven window counts {dict(per_entity)} — an entity's stream ended "
        "early, or edges were anchored to first-seen rather than the epoch"
    )
    # 2. Window starts are multiples of the slide from the epoch.
    for _, start, *_ in results:
        offset = (start - EPOCH.replace(tzinfo=start.tzinfo)) % w.slide
        assert offset == timedelta(0), f"window start {start} is not on the grid"
    # 3. State is bounded — nothing is retained past sealing.
    assert w.open_windows() <= expected_per_entity * len(per_entity), (
        f"{w.open_windows()} open windows — eviction is not running"
    )
    print(f"OK — {len(results)} results, {w.open_windows()} windows still open")