Detecting Regional Distribution Drift

from __future__ import annotations

import math
import random
from dataclasses import dataclass
from typing import Mapping, Sequence


@dataclass(frozen=True)
class DriftSignal:
    region: str
    distance: float
    error_bar: float
    contributors: int
    breached: bool
    worst_bucket: int


def total_variation(p: Sequence[float], q: Sequence[float]) -> float:
    """TV distance between two normalised histograms; bounded in [0, 1]."""
    if len(p) != len(q):
        raise ValueError("histograms must share their bucket edges")
    sp, sq = sum(p), sum(q)
    if sp <= 0 or sq <= 0:
        return 0.0
    return 0.5 * sum(abs(a / sp - b / sq) for a, b in zip(p, q))


def jensen_shannon(p: Sequence[float], q: Sequence[float]) -> float:
    """JS distance — the square root of the JS divergence, also bounded in [0, 1]."""
    sp, sq = sum(p), sum(q)
    if sp <= 0 or sq <= 0:
        return 0.0
    pa = [a / sp for a in p]
    qa = [b / sq for b in q]

    def kl(x: Sequence[float], y: Sequence[float]) -> float:
        return sum(xi * math.log(xi / yi) for xi, yi in zip(x, y) if xi > 0 and yi > 0)

    m = [(a + b) / 2 for a, b in zip(pa, qa)]
    return math.sqrt(max(0.0, 0.5 * kl(pa, m) + 0.5 * kl(qa, m)) / math.log(2))


def detect_drift(
    region: str,
    current: Sequence[float],
    reference: Sequence[float],
    *,
    contributors: int,
    epsilon: float,
    threshold: float,
    seed: int = 0,
) -> DriftSignal:
    """Compare a region's noised histogram against a pinned reference.

    Both histograms are already DP aggregates, so no further noise is added here —
    adding it again would double-charge the budget for a statistic computed from
    released values. The error bar propagates the aggregate's own uncertainty into
    the distance, which is what makes the threshold comparison honest.
    """
    distance = total_variation(current, reference)
    # A histogram bucket's noise is Laplace(cap/eps); normalised, the distance
    # inherits roughly that over the contributor count, per bucket.
    per_bucket = math.sqrt(2.0) / (epsilon * max(1, contributors))
    error_bar = 1.96 * per_bucket * math.sqrt(len(current)) / 2.0
    diffs = [
        abs(a / max(1e-9, sum(current)) - b / max(1e-9, sum(reference)))
        for a, b in zip(current, reference)
    ]
    return DriftSignal(
        region=region,
        distance=distance,
        error_bar=error_bar,
        contributors=contributors,
        breached=distance > max(threshold, error_bar),
        worst_bucket=int(max(range(len(diffs)), key=lambda i: diffs[i])),
    )


def persistent_breaches(
    history: Mapping[str, Sequence[DriftSignal]], *, required: int = 2
) -> dict[str, int]:
    """Regions whose most recent `required` rounds all breached.

    Requiring persistence trades detection latency for a large reduction in false
    alarms — with independent rounds the false-alarm rate falls roughly as the
    per-round rate raised to `required`.
    """
    out: dict[str, int] = {}
    for region, signals in history.items():
        recent = list(signals)[-required:]
        if len(recent) == required and all(s.breached for s in recent):
            out[region] = recent[-1].worst_bucket
    return out

def _validate() -> None:
    rng = random.Random(31)
    reference = [1200.0, 3400.0, 2900.0, 1100.0, 400.0, 90.0]

    # 1. An unchanged distribution must not breach.
    same = [v * (1 + rng.gauss(0, 0.01)) for v in reference]
    quiet = detect_drift("stable", same, reference, contributors=800,
                         epsilon=1.0, threshold=0.05)
    assert not quiet.breached, quiet

    # 2. A real shape change must breach, and must name the moved bucket.
    shifted = [400.0, 1200.0, 2900.0, 3100.0, 1300.0, 200.0]
    loud = detect_drift("shifted", shifted, reference, contributors=800,
                        epsilon=1.0, threshold=0.05)
    assert loud.breached and loud.distance > 0.15, loud
    assert loud.worst_bucket in (0, 1, 3), loud.worst_bucket

    # 3. A tiny cohort must widen the bar enough to suppress a marginal signal.
    marginal = [1150.0, 3300.0, 3000.0, 1150.0, 410.0, 95.0]
    small = detect_drift("small", marginal, reference, contributors=12,
                         epsilon=1.0, threshold=0.01)
    assert small.error_bar > small.distance, small

    # 4. Both distance measures must agree on direction, if not on magnitude.
    assert jensen_shannon(shifted, reference) > jensen_shannon(same, reference)

    # 5. Persistence must require consecutive breaches, not merely frequent ones.
    hist = {
        "a": [loud, quiet, loud],           # alternating: not persistent
        "b": [quiet, loud, loud],           # last two breached: persistent
    }
    flagged = persistent_breaches(hist, required=2)
    assert "b" in flagged and "a" not in flagged

    # 6. Distances are bounded, so a degenerate region cannot produce a huge signal.
    assert 0.0 <= total_variation([0.0, 0.0, 1.0], reference) <= 1.0

    print("regional drift detection: all assertions passed")


_validate()