"""Compare packet-optimized schedules on transient event discovery.

The experiment deliberately separates background packet rate (the easy online
proxy) from event arrival rate (the monitoring target). It evaluates the same
four schedules in two environments:

1. aligned: event arrival rates increase with packet traffic;
2. rare-channel: event arrivals are concentrated on quiet channels.

Run:
    uv run event_replay.py

Output:
    event_replay.png
"""

# /// script
# requires-python = "==3.12.*"
# dependencies = [
#   "matplotlib>=3.9",
#   "numpy>=2.0",
# ]
# ///

from __future__ import annotations

import sys
from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from figure_quality import configure_matplotlib, figure_output_path, validate_figure

configure_matplotlib()

SEED = 0xC0FFEE
CHANNELS = 14
STEPS = 4_000
RUNS = 200
MAX_EVENT_DURATION = 12
MAX_REVISIT = 2 * CHANNELS

POLICIES = ("round-robin", "traffic-proportional", "packet TS", "quota + packet TS")
COLORS = ("#264653", "#2a9d8f", "#e9c46a", "#e76f51")


@dataclass(frozen=True)
class Event:
    channel: int
    start: int
    stop: int


def make_events(rng: np.random.Generator, rates: np.ndarray) -> list[Event]:
    events: list[Event] = []
    for channel, rate in enumerate(rates):
        starts = np.flatnonzero(rng.random(STEPS) < rate)
        durations = np.minimum(rng.geometric(0.45, len(starts)), MAX_EVENT_DURATION)
        events.extend(
            Event(channel, int(start), min(STEPS, int(start + duration)))
            for start, duration in zip(starts, durations, strict=True)
        )
    return events


def packet_ts(
    rng: np.random.Generator,
    packet_counts: np.ndarray,
    *,
    quota: bool,
) -> np.ndarray:
    alpha = np.ones(CHANNELS)
    beta = np.ones(CHANNELS)
    last_visit = np.full(CHANNELS, -MAX_REVISIT, dtype=int)
    schedule = np.empty(STEPS, dtype=int)

    for step in range(STEPS):
        overdue = np.flatnonzero(step - last_visit >= MAX_REVISIT)
        if quota and len(overdue):
            channel = int(overdue[np.argmin(last_visit[overdue])])
        else:
            channel = int(np.argmax(rng.gamma(alpha, 1.0 / beta)))

        schedule[step] = channel
        last_visit[channel] = step
        alpha[channel] += packet_counts[step, channel]
        beta[channel] += 1.0

    return schedule


def schedules(
    rng: np.random.Generator,
    packet_rates: np.ndarray,
    packet_counts: np.ndarray,
) -> dict[str, np.ndarray]:
    weights = packet_rates / packet_rates.sum()
    return {
        "round-robin": np.arange(STEPS) % CHANNELS,
        "traffic-proportional": rng.choice(CHANNELS, size=STEPS, p=weights),
        "packet TS": packet_ts(rng, packet_counts, quota=False),
        "quota + packet TS": packet_ts(rng, packet_counts, quota=True),
    }


def detection_fraction(schedule: np.ndarray, events: list[Event]) -> float:
    if not events:
        return 1.0
    detected = sum(
        bool(np.any(schedule[event.start : event.stop] == event.channel))
        for event in events
    )
    return detected / len(events)


def max_blind_run(schedule: np.ndarray) -> int:
    """Return the most consecutive decision intervals spent off one channel."""
    largest = 0
    for channel in range(CHANNELS):
        visits = np.flatnonzero(schedule == channel)
        boundaries = np.concatenate(([-1], visits, [STEPS]))
        largest = max(largest, int((np.diff(boundaries) - 1).max()))
    return largest


def run() -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
    root = np.random.SeedSequence(SEED)
    packet_rates = np.linspace(3.0, 48.0, CHANNELS)
    aligned = 0.003 + 0.017 * packet_rates / packet_rates.max()
    rare_channel = 0.003 + 0.017 * packet_rates[::-1] / packet_rates.max()

    detection = {
        "aligned": np.empty((RUNS, len(POLICIES))),
        "rare-channel": np.empty((RUNS, len(POLICIES))),
    }
    gaps = {
        "aligned": np.empty((RUNS, len(POLICIES))),
        "rare-channel": np.empty((RUNS, len(POLICIES))),
    }

    for run_index, run_seed in enumerate(root.spawn(RUNS)):
        packet_seed, policy_seed, aligned_seed, rare_seed = run_seed.spawn(4)
        packet_rng = np.random.default_rng(packet_seed)
        policy_rng = np.random.default_rng(policy_seed)
        packet_counts = packet_rng.poisson(packet_rates, size=(STEPS, CHANNELS))
        policy_schedules = schedules(policy_rng, packet_rates, packet_counts)

        for scenario, event_rates, event_seed in (
            ("aligned", aligned, aligned_seed),
            ("rare-channel", rare_channel, rare_seed),
        ):
            events = make_events(np.random.default_rng(event_seed), event_rates)

            for policy_index, policy in enumerate(POLICIES):
                schedule = policy_schedules[policy]
                detection[scenario][run_index, policy_index] = detection_fraction(
                    schedule, events
                )
                gaps[scenario][run_index, policy_index] = max_blind_run(schedule)

    return detection, gaps


def mean_ci(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    mean = values.mean(axis=0)
    ci = 1.96 * values.std(axis=0, ddof=1) / np.sqrt(values.shape[0])
    return mean, ci


def plot(detection: dict[str, np.ndarray], gaps: dict[str, np.ndarray]) -> None:
    figure, axes = plt.subplots(1, 3, figsize=(13.5, 4.2))
    x = np.arange(len(POLICIES))

    for axis, scenario, title in (
        (axes[0], "aligned", "Event rates increase with traffic"),
        (axes[1], "rare-channel", "Events favor quiet channels"),
    ):
        mean, ci = mean_ci(detection[scenario])
        axis.bar(x, 100 * mean, yerr=100 * ci, color=COLORS, capsize=3)
        axis.set_title(title)
        axis.set_ylabel("events detected (%)")
        axis.set_ylim(0, 100)
        axis.grid(axis="y", alpha=0.2)

    gap_mean, gap_ci = mean_ci(gaps["rare-channel"])
    gap_bars = axes[2].bar(x, gap_mean, yerr=gap_ci, color=COLORS, capsize=3)
    axes[2].axhline(MAX_REVISIT - 1, color="#333333", linestyle="--", linewidth=1)
    axes[2].set_title("Longest blind run")
    axes[2].set_ylabel("off-channel intervals (log scale)")
    axes[2].set_yscale("log")
    axes[2].set_ylim(10, 7_000)
    axes[2].bar_label(
        gap_bars, labels=[f"{value:.0f}" for value in gap_mean], padding=3
    )
    axes[2].grid(axis="y", which="major", alpha=0.2)

    for axis in axes:
        axis.set_xticks(x, POLICIES, rotation=22, ha="right")
        axis.spines[["top", "right"]].set_visible(False)

    figure.suptitle("Packet-rate optimization and event discovery", y=1.02)
    figure.tight_layout()
    validate_figure(figure)
    output = figure_output_path(__file__, "event_replay.png")
    figure.savefig(output, dpi=300, bbox_inches="tight")


if __name__ == "__main__":
    detection_results, gap_results = run()
    for scenario in ("aligned", "rare-channel"):
        means, cis = mean_ci(detection_results[scenario])
        print(scenario)
        for policy, mean, ci in zip(POLICIES, means, cis, strict=True):
            print(f"  {policy:22s} {100 * mean:5.1f}% +/- {100 * ci:.1f} pp")
    plot(detection_results, gap_results)