Staleness-Aware Weighting for Late Updates

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Sequence


@dataclass(frozen=True)
class Update:
    """One client update as it arrives at the aggregator."""

    client_id: str
    generated_at_round: int
    samples: int
    gnss_pdop: float
    cosine_to_aggregate: float
    l2_norm: float


@dataclass(frozen=True)
class WeightingPolicy:
    half_life: float = 5.0
    max_staleness: int = 15
    weight_floor: float = 0.05
    cosine_floor: float = -0.1
    clip_norm: float = 1.0

    def staleness_factor(self, lag: int) -> float:
        """Exponential decay in the round lag."""
        return 0.5 ** (lag / self.half_life)

    def quality_factor(self, update: Update) -> float:
        """Sample count and GNSS geometry, both bounded so neither dominates."""
        sample_term = min(1.0, math.log1p(update.samples) / math.log1p(500.0))
        pdop_term = min(1.0, 1.5 / max(1.0, update.gnss_pdop))
        return sample_term * pdop_term

    def weight(self, update: Update, current_round: int) -> float:
        """Final aggregation weight, or 0.0 if the update must be rejected.

        Rejection has two independent causes, and they are worth distinguishing in
        metrics: staleness rejection is an infrastructure signal, cosine rejection
        is a poisoning signal, and conflating them hides both.
        """
        lag = current_round - update.generated_at_round
        if lag < 0:
            raise ValueError("an update cannot be newer than the current round")
        if lag > self.max_staleness:
            return 0.0
        if update.cosine_to_aggregate < self.cosine_floor:
            return 0.0
        raw = self.staleness_factor(lag) * self.quality_factor(update)
        return max(self.weight_floor, raw)


def aggregate(
    updates: Sequence[Update], policy: WeightingPolicy, current_round: int
) -> tuple[float, dict[str, int]]:
    """Weighted mean of update norms, plus a rejection breakdown by cause.

    Returning the breakdown alongside the aggregate is deliberate: a change in the
    ratio between the two rejection causes is the earliest available signal that
    something has changed in the fleet.
    """
    total_w = 0.0
    total = 0.0
    reasons = {"stale": 0, "cosine": 0, "accepted": 0}
    for u in updates:
        lag = current_round - u.generated_at_round
        if lag > policy.max_staleness:
            reasons["stale"] += 1
            continue
        if u.cosine_to_aggregate < policy.cosine_floor:
            reasons["cosine"] += 1
            continue
        w = policy.weight(u, current_round)
        clipped = min(u.l2_norm, policy.clip_norm)
        total += w * clipped
        total_w += w
        reasons["accepted"] += 1
    return (total / total_w if total_w else 0.0), reasons


def effective_participation(
    updates: Sequence[Update], policy: WeightingPolicy, current_round: int
) -> float:
    """Sum of weights divided by the client count — how much of the fleet counted.

    A round with 800 returning clients and an effective participation of 0.2 is a
    round with the statistical power of 160 clients, and the variance of the
    aggregate should be read against that number rather than the headline one.
    """
    if not updates:
        return 0.0
    return sum(policy.weight(u, current_round) for u in updates) / len(updates)

def _validate() -> None:
    policy = WeightingPolicy(half_life=5.0, max_staleness=15)

    fresh = Update("u1", 100, samples=400, gnss_pdop=1.2,
                   cosine_to_aggregate=0.8, l2_norm=0.9)
    stale = Update("u2", 92, samples=400, gnss_pdop=1.2,
                   cosine_to_aggregate=0.8, l2_norm=0.9)
    ancient = Update("u3", 80, samples=400, gnss_pdop=1.2,
                     cosine_to_aggregate=0.8, l2_norm=0.9)
    hostile = Update("u4", 100, samples=400, gnss_pdop=1.2,
                     cosine_to_aggregate=-0.7, l2_norm=0.9)

    # 1. Weight must decay monotonically with lag.
    assert policy.weight(fresh, 100) > policy.weight(stale, 100)

    # 2. The half-life must be exactly that: half the weight after `half_life` rounds.
    five_late = Update("u5", 95, 400, 1.2, 0.8, 0.9)
    assert abs(policy.weight(five_late, 100) / policy.weight(fresh, 100) - 0.5) < 1e-9

    # 3. Beyond max_staleness the update is dropped, not merely down-weighted.
    assert policy.weight(ancient, 100) == 0.0

    # 4. An anti-aligned update is rejected regardless of freshness.
    assert policy.weight(hostile, 100) == 0.0

    # 5. The rejection breakdown must separate the two causes.
    _, reasons = aggregate([fresh, stale, ancient, hostile], policy, 100)
    assert reasons == {"stale": 1, "cosine": 1, "accepted": 2}, reasons

    # 6. Effective participation must fall when the fleet is mostly stale.
    slow_fleet = [Update(f"s{i}", 88, 400, 1.2, 0.8, 0.9) for i in range(20)]
    fast_fleet = [Update(f"f{i}", 100, 400, 1.2, 0.8, 0.9) for i in range(20)]
    assert effective_participation(slow_fleet, policy, 100) < 0.4
    assert effective_participation(fast_fleet, policy, 100) > 0.9

    # 7. The floor must keep slow-but-valid clients contributing something.
    barely = Update("u6", 87, 400, 1.2, 0.8, 0.9)
    assert policy.weight(barely, 100) >= policy.weight_floor

    print("staleness weighting: all assertions passed")


_validate()