#!/usr/bin/env python3
"""Temporal obligations benchmark for Paper 48.

This synthetic harness holds one event stream fixed and compares configured
representations on four future-question families. It tests information
preservation under explicit budgets. It does not benchmark trained language
models, neural attention quality, or hardware throughput.
"""

from __future__ import annotations

import argparse
from collections import Counter, OrderedDict, defaultdict
from dataclasses import dataclass
from random import Random
from typing import Any, Iterable, Protocol

STATUSES = ("OPEN", "HOLD", "CLOSED")
TAGS = tuple(f"T{i:02d}" for i in range(16))


@dataclass(frozen=True)
class Event:
    step: int
    event_id: str
    kind: str
    key: str
    value: str
    tag: str
    committed: bool = True


class Memory(Protocol):
    name: str

    def note(self, key: str) -> str | None: ...
    def adjacency(self, first: str, second: str) -> int | None: ...
    def count(self, tag: str) -> int | None: ...
    def state(self, key: str) -> str | None: ...


def build_stream(n: int = 4096, seed: int = 48) -> list[Event]:
    """Generate a deterministic mixture of signals, notes, and status records."""
    if n < 256:
        raise ValueError("n must be at least 256 so every query family is populated")

    rng = Random(seed)
    events: list[Event] = []
    note_number = 0

    for step in range(n):
        draw = rng.random()
        tag = rng.choice(TAGS)

        if draw < 0.48:
            kind = "signal"
            key = f"sensor-{rng.randrange(96):03d}"
            value = f"reading-{rng.randrange(1000):03d}"
            committed = True
        elif draw < 0.68:
            kind = "note"
            key = f"note-{note_number:04d}"
            note_number += 1
            value = f"phrase-{rng.randrange(10_000_000):07d}"
            committed = True
        else:
            kind = "status"
            key = f"asset-{rng.randrange(128):03d}"
            value = rng.choice(STATUSES)
            committed = rng.random() < 0.68

        events.append(
            Event(
                step=step,
                event_id=f"E{rng.randrange(10**8):08d}",
                kind=kind,
                key=key,
                value=value,
                tag=tag,
                committed=committed,
            )
        )

    return events


class Replay:
    """Exact full-history access used as an information-preservation oracle."""

    name = "Full-history replay oracle"

    def __init__(self, events: list[Event]) -> None:
        self.events = events

    def note(self, key: str) -> str | None:
        for event in self.events:
            if event.kind == "note" and event.key == key:
                return event.value
        return None

    def adjacency(self, first: str, second: str) -> int:
        return sum(
            1
            for left, right in zip(self.events, self.events[1:])
            if left.tag == first and right.tag == second
        )

    def count(self, tag: str) -> int:
        return sum(event.tag == tag for event in self.events)

    def state(self, key: str) -> str | None:
        for event in reversed(self.events):
            if event.kind == "status" and event.key == key and event.committed:
                return event.value
        return None


class SlidingAttention(Replay):
    """A bounded recent window standing in for cost-constrained attention."""

    def __init__(self, events: list[Event], width: int = 128) -> None:
        if width < 1:
            raise ValueError("attention width must be positive")
        self.width = width
        self.events = events[-width:]
        self.name = f"{width}-event attention"


class Retrieval:
    """An addressable archive whose model-facing interface exposes top-k items."""

    def __init__(self, events: list[Event], k: int = 8) -> None:
        if k < 1:
            raise ValueError("retrieval k must be positive")
        self.k = k
        self.name = f"Top-{k} retrieval"
        self.by_note: dict[str, Event] = {}
        self.by_tag: dict[str, list[Event]] = defaultdict(list)
        self.by_key: dict[str, list[Event]] = defaultdict(list)

        for event in events:
            if event.kind == "note":
                self.by_note[event.key] = event
            self.by_tag[event.tag].append(event)
            if event.kind == "status":
                self.by_key[event.key].append(event)

    def note(self, key: str) -> str | None:
        event = self.by_note.get(key)
        return event.value if event else None

    def adjacency(self, first: str, second: str) -> int:
        # The model sees only the most recent k matches for each tag.
        candidates = {
            event.step: event
            for event in self.by_tag[first][-self.k :]
            + self.by_tag[second][-self.k :]
        }
        ordered = [candidates[index] for index in sorted(candidates)]
        return sum(
            1
            for left, right in zip(ordered, ordered[1:])
            if right.step == left.step + 1
            and left.tag == first
            and right.tag == second
        )

    def count(self, tag: str) -> int:
        # The interface does not reveal posting-list cardinality.
        return len(self.by_tag[tag][-self.k :])

    def state(self, key: str) -> str | None:
        # Recency is used without commit authority, so proposals can win.
        matches = self.by_key.get(key, [])[-self.k :]
        return matches[-1].value if matches else None


class RecurrentSummary:
    """A fixed-size, task-shaped summary updated once per event."""

    name = "Fixed recurrent summary"

    def __init__(
        self,
        events: list[Event],
        note_slots: int = 32,
        state_slots: int = 24,
        pair_slots: int = 32,
    ) -> None:
        if min(note_slots, state_slots, pair_slots) < 0:
            raise ValueError("slot counts cannot be negative")

        all_pairs = [(left, right) for left in TAGS for right in TAGS]
        self.selected_pairs = set(all_pairs[:pair_slots])
        self.tag_counts: Counter[str] = Counter()
        self.pair_counts: Counter[tuple[str, str]] = Counter()
        self.notes: dict[str, str] = {}
        self.states: OrderedDict[str, str] = OrderedDict()
        self.note_slots = note_slots
        self.state_slots = state_slots

        previous: Event | None = None
        notes_seen = 0

        for event in events:
            self.tag_counts[event.tag] += 1

            if previous and (previous.tag, event.tag) in self.selected_pairs:
                self.pair_counts[(previous.tag, event.tag)] += 1
            previous = event

            if event.kind == "note":
                # Deterministic sparse retention: one note in every 25.
                if notes_seen % 25 == 0 and len(self.notes) < note_slots:
                    self.notes[event.key] = event.value
                notes_seen += 1

            if event.kind == "status" and event.committed:
                if event.key in self.states:
                    self.states.pop(event.key)
                self.states[event.key] = event.value
                while len(self.states) > state_slots:
                    self.states.popitem(last=False)

    def note(self, key: str) -> str | None:
        return self.notes.get(key)

    def adjacency(self, first: str, second: str) -> int:
        return self.pair_counts.get((first, second), 0)

    def count(self, tag: str) -> int:
        return self.tag_counts[tag]

    def state(self, key: str) -> str | None:
        return self.states.get(key)


class ExplicitState:
    """A typed current-state view that advances only on committed updates."""

    name = "Typed current state"

    def __init__(self, events: list[Event]) -> None:
        self.current: dict[str, tuple[int, str]] = {}
        for event in events:
            if event.kind == "status" and event.committed:
                self.current[event.key] = (event.step, event.value)

    def note(self, key: str) -> None:
        return None

    def adjacency(self, first: str, second: str) -> None:
        return None

    def count(self, tag: str) -> None:
        return None

    def state(self, key: str) -> str | None:
        item = self.current.get(key)
        return item[1] if item else None


def exact_accuracy(got: Iterable[Any], expected: Iterable[Any]) -> float:
    pairs = list(zip(got, expected))
    if not pairs:
        raise ValueError("at least one query is required")
    return 100.0 * sum(actual == target for actual, target in pairs) / len(pairs)


def run_benchmark(
    n: int = 4096,
    seed: int = 48,
    queries: int = 256,
    attention_width: int = 128,
    retrieval_k: int = 8,
) -> list[tuple[str, float, float, float, float]]:
    if queries < 1:
        raise ValueError("queries must be positive")

    rng = Random(seed + 1)
    events = build_stream(n=n, seed=seed)
    oracle = Replay(events)

    notes = [event.key for event in events if event.kind == "note"]
    assets = sorted(
        {
            event.key
            for event in events
            if event.kind == "status" and event.committed
        }
    )
    if not notes or not assets:
        raise RuntimeError("generated stream did not contain required event types")

    note_queries = [rng.choice(notes) for _ in range(queries)]
    pair_queries = [(rng.choice(TAGS), rng.choice(TAGS)) for _ in range(queries)]
    count_queries = [rng.choice(TAGS) for _ in range(queries)]
    state_queries = [rng.choice(assets) for _ in range(queries)]

    expected = {
        "note": [oracle.note(key) for key in note_queries],
        "relation": [oracle.adjacency(a, b) for a, b in pair_queries],
        "count": [oracle.count(tag) for tag in count_queries],
        "state": [oracle.state(key) for key in state_queries],
    }

    systems: list[Memory] = [
        oracle,
        SlidingAttention(events, width=attention_width),
        Retrieval(events, k=retrieval_k),
        RecurrentSummary(events),
        ExplicitState(events),
    ]

    rows: list[tuple[str, float, float, float, float]] = []
    for system in systems:
        rows.append(
            (
                system.name,
                exact_accuracy(
                    (system.note(key) for key in note_queries), expected["note"]
                ),
                exact_accuracy(
                    (system.adjacency(a, b) for a, b in pair_queries),
                    expected["relation"],
                ),
                exact_accuracy(
                    (system.count(tag) for tag in count_queries), expected["count"]
                ),
                exact_accuracy(
                    (system.state(key) for key in state_queries), expected["state"]
                ),
            )
        )
    return rows


def print_table(rows: list[tuple[str, float, float, float, float]]) -> None:
    print(
        f"{'Mechanism':30} {'Exact note':>11} {'Relation':>11} "
        f"{'Count':>11} {'State':>11}"
    )
    for name, note, relation, count, state in rows:
        print(f"{name:30} {note:10.1f}% {relation:10.1f}% {count:10.1f}% {state:10.1f}%")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compare memory representations on matched temporal obligations."
    )
    parser.add_argument("--events", type=int, default=4096, help="stream length")
    parser.add_argument("--queries", type=int, default=256, help="queries per family")
    parser.add_argument("--seed", type=int, default=48, help="random seed")
    parser.add_argument("--attention-width", type=int, default=128)
    parser.add_argument("--retrieval-k", type=int, default=8)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    rows = run_benchmark(
        n=args.events,
        seed=args.seed,
        queries=args.queries,
        attention_width=args.attention_width,
        retrieval_k=args.retrieval_k,
    )
    print_table(rows)


if __name__ == "__main__":
    main()