Exponentially weighted metrics for fleet telemetry

import math
from dataclasses import dataclass
from datetime import datetime


@dataclass
class TimeAwareEWM:
    """
    Time-aware exponentially weighted mean with bias correction.

    State is three floats per entity regardless of history length, which is
    what makes it usable for a live metric across a large fleet.

    Parameters
    ----------
    half_life_s : float
        Elapsed time after which an observation's weight halves. Prefer this
        to alpha: "the last five minutes dominate" is a statement anybody can
        check, and alpha is not.
    """
    half_life_s: float
    value: float = 0.0          # weighted sum
    weight: float = 0.0         # sum of weights, for bias correction
    last_t: datetime | None = None

    def __post_init__(self):
        if self.half_life_s <= 0:
            raise ValueError("half_life_s must be positive.")

    def update(self, t: datetime, x: float) -> float:
        """Ingest one observation and return the current corrected mean."""
        if t.tzinfo is None:
            raise ValueError("Timestamps must be timezone-aware (UTC).")
        if not math.isfinite(x):
            return self.current()          # skip, do not poison the accumulator

        if self.last_t is None:
            decay = 1.0
        else:
            dt = (t - self.last_t).total_seconds()
            if dt < 0:
                # A backwards timestamp means a clock step. Decaying by a
                # negative interval AMPLIFIES old state; reset instead.
                self.value, self.weight = 0.0, 0.0
                decay = 1.0
            else:
                decay = 0.5 ** (dt / self.half_life_s)

        self.value = self.value * decay + x
        self.weight = self.weight * decay + 1.0
        self.last_t = t
        return self.current()

    def current(self, at: datetime | None = None) -> float:
        """
        The bias-corrected mean, optionally decayed to a later time.

        Dividing by the accumulated WEIGHT rather than assuming it has
        converged to 1/(1-decay) is the bias correction: without it the first
        few observations are pulled toward the zero the accumulator started at.
        """
        if self.weight <= 0:
            return float("nan")
        if at is None or self.last_t is None:
            return self.value / self.weight
        dt = max((at - self.last_t).total_seconds(), 0.0)
        d = 0.5 ** (dt / self.half_life_s)
        return (self.value * d) / (self.weight * d)     # decay cancels; kept for clarity

    def staleness_s(self, now: datetime) -> float:
        """How old the newest observation is. A metric with no freshness
        signal cannot distinguish 'steady' from 'stopped reporting'."""
        return float("inf") if self.last_t is None else (now - self.last_t).total_seconds()

def validate_ewm(half_life_s: float = 300.0) -> None:
    """Property tests that catch the three classic implementation errors."""
    from datetime import timedelta, timezone

    t0 = datetime(2026, 5, 14, 9, 0, tzinfo=timezone.utc)
    e = TimeAwareEWM(half_life_s)

    # 1. Bias correction: a single observation must return itself, not half
    #    of it. Failing this means the accumulator starts at zero uncorrected.
    assert abs(e.update(t0, 10.0) - 10.0) < 1e-9, "no bias correction on the first value"

    # 2. Constant input converges to that constant regardless of spacing.
    for i in range(1, 40):
        e.update(t0 + timedelta(seconds=i * 37), 10.0)
    assert abs(e.current() - 10.0) < 1e-6, "constant input does not converge"

    # 3. Time-awareness: a step change decays by half over exactly one
    #    half-life, whether that is one observation or fifty.
    a, b = TimeAwareEWM(half_life_s), TimeAwareEWM(half_life_s)
    a.update(t0, 0.0); b.update(t0, 0.0)
    a.update(t0 + timedelta(seconds=half_life_s), 1.0)
    for i in range(1, 51):
        b.update(t0 + timedelta(seconds=half_life_s * i / 50), 1.0)
    # Both saw the same elapsed time; the count-based bug makes them differ.
    assert abs(a.current() - 0.667) < 0.02, f"single-step value {a.current():.3f} unexpected"
    assert b.current() > a.current(), "denser sampling should weight recent values more"

    # 4. A backwards clock resets rather than amplifying.
    c = TimeAwareEWM(half_life_s)
    c.update(t0, 50.0)
    c.update(t0 - timedelta(seconds=60), 1.0)
    assert abs(c.current() - 1.0) < 1e-9, "backwards timestamp did not reset the state"
    print("OK — bias-corrected, time-aware, clock-safe")