#!/usr/bin/env python3
"""A straggler silently undercounts a fired window; a watermark recovers it. DDIA 2e Ch. 12.

Group a stream of events into 1-minute EVENT-TIME tumbling windows and count per
window. Window [37, 38) accumulates 5 events. Then time moves on -- events start
arriving with minute-38 and minute-39 timestamps -- so the aggregator decides
window 37 is "complete" and FIRES it, emitting count = 5.

Then a STRAGGLER arrives: an event whose event_time is minute 37, delayed by a
network hiccup so it shows up *after* the minute-39 events. The window it belongs
to has already closed.

The surprise is what the completeness policy does with it:

  Policy A -- drop late (allowed lateness = 0): the window closed the instant the
  watermark passed its end, so the straggler is silently DROPPED. No error.
  Window 37's published count stays 5, but the true count is 6 -- undercount of 1.

  Policy B -- watermark + allowed lateness (grace = 2): the window still fired
  on time at count 5, but is kept OPEN for corrections until the watermark passes
  its end + grace. The straggler arrives while it is still open, so it is
  accepted and the window emits a CORRECTION: published 5 -> 6.

Same stream, same straggler; the only difference is how long the window waits
before it stops accepting late events. Pure standard library, and the arrival
order is modeled explicitly, so the transcript is fully deterministic.

The watermark used here is a perfect arrival high-watermark: watermark =
(max event_time seen so far). An event is "late" if its window's end is already
<= the watermark -- i.e. the watermark had already declared that window complete.
"""

# The arrival stream: (event_time_minute, event_id), in PROCESSING order -- the
# order the aggregator actually sees the events. Most events arrive roughly in
# event-time order. One straggler, s1, carries a minute-37 timestamp but arrives
# late, after the minute-38 and minute-39 events (a network hiccup delayed it).
STREAM = [
    (37, "e1"),
    (37, "e2"),
    (37, "e3"),
    (37, "e4"),
    (37, "e5"),
    (38, "e6"),
    (39, "e7"),
    (37, "s1"),   # <-- STRAGGLER: minute-37 timestamp, arrives out of order
    (40, "e8"),
]

WINDOW = 37                       # the window we watch: event-time minute [37, 38)
TRUE_37 = sum(1 for t, _ in STREAM if t == WINDOW)   # events that truly belong to it


class WindowAggregator:
    """1-minute event-time tumbling windows, counting events per window.

    A window is [w, w+1) for the minute w. The watermark is the max event_time
    seen so far. A window FIRES (emits its count once) when the watermark reaches
    its end. `allowed_lateness` is how long past the end the window is kept open
    to accept late events and emit CORRECTIONs; after that it is CLOSED and any
    further late event for it is DROPPED.
    """

    def __init__(self, allowed_lateness):
        self.allowed_lateness = allowed_lateness
        self.counts = {}          # window_start -> current count
        self.published = {}       # window_start -> last emitted count
        self.closed = set()       # window_starts no longer accepting events
        self.watermark = -1       # max event_time seen (perfect arrival watermark)

    def process(self, event_time, event_id):
        w = event_time            # 1-minute tumbling: window start == the minute
        end = w + 1
        is_straggler = event_id.startswith("s")
        tag = "STRAGGLER" if is_straggler else "event"
        print(f"  arrive: {tag:9s} {event_id}  event_time=min {event_time}   "
              f"(watermark is {self._wm()})")

        # A late event is one whose window the watermark has already passed.
        if end <= self.watermark:
            if w in self.closed:
                print(f"      DROP {event_id}: window [{w},{end}) is CLOSED "
                      f"(watermark {self.watermark} >= {end}+{self.allowed_lateness}) "
                      f"-- late event ignored, published count stays "
                      f"{self.published[w]}")
                return
            # Still within allowed lateness: accept it and correct the window.
            self.counts[w] += 1
            old = self.published[w]
            self.published[w] = self.counts[w]
            print(f"      LATE {event_id} accepted into still-open window [{w},{end}) "
                  f"(watermark {self.watermark} < {end}+{self.allowed_lateness})")
            print(f"      CORRECT window [{w},{end}): published {old} -> "
                  f"{self.published[w]}")
            return

        # On-time event: count it, then let the watermark advance and fire/close.
        self.counts[w] = self.counts.get(w, 0) + 1
        if event_time > self.watermark:
            self.watermark = event_time
            print(f"      watermark advances to {self.watermark}")

        for ws in sorted(self.counts):
            if ws + 1 <= self.watermark and ws not in self.published:
                self.published[ws] = self.counts[ws]
                print(f"      FIRE window [{ws},{ws + 1}) -> count "
                      f"{self.published[ws]} (watermark {self.watermark} >= {ws + 1})")

        for ws in sorted(self.counts):
            if ws in self.published and ws not in self.closed \
                    and ws + 1 + self.allowed_lateness <= self.watermark:
                self.closed.add(ws)
                print(f"      CLOSE window [{ws},{ws + 1}) (watermark "
                      f"{self.watermark} >= {ws + 1}+{self.allowed_lateness}) "
                      f"-- stops accepting late events")

    def _wm(self):
        return "-" if self.watermark < 0 else str(self.watermark)


def run_policy(label, allowed_lateness):
    print(label)
    print(f"    (allowed lateness = {allowed_lateness} min)")
    agg = WindowAggregator(allowed_lateness)
    for event_time, event_id in STREAM:
        agg.process(event_time, event_id)
    return agg.published.get(WINDOW)


def main():
    print("A stream of events, grouped into 1-minute event-time tumbling windows,")
    print("counted per window. We watch window [37, 38).\n")
    print(f"    arrival order (processing time): "
          f"{', '.join(f'{i}@min{t}' for t, i in STREAM)}")
    print(f"    events that TRULY belong to window [37, 38): {TRUE_37} "
          f"(e1 e2 e3 e4 e5 s1)\n")
    print("If every event landed in its window, window 37 would count "
          f"{TRUE_37}.\n")

    print("-" * 72)
    a = run_policy("Policy A -- DROP LATE (the default: fire, then ignore late events)",
                   allowed_lateness=0)
    print(f"\n  window [37, 38) FINAL published count = {a}")
    print(f"  events that truly belong there              = {TRUE_37}")
    print(f"  >> UNDERCOUNT: published {a}, actual {TRUE_37} "
          f"-- the straggler was dropped, silently, with no error.\n")

    print("-" * 72)
    b = run_policy("Policy B -- WATERMARK + ALLOWED LATENESS (keep open, emit corrections)",
                   allowed_lateness=2)
    print(f"\n  window [37, 38) FINAL published count = {b}")
    print(f"  events that truly belong there              = {TRUE_37}")
    print(f"  >> RECOVERED: fired at 5, then CORRECTED to {b} when the straggler "
          f"arrived -- count now matches reality.\n")

    print("-" * 72)
    print("Same stream, same straggler:")
    print(f"    drop-late  published {a}  (undercount of {TRUE_37 - a})")
    print(f"    watermark  published {b}  (correct)")


if __name__ == "__main__":
    main()