"""proactive_agents_toolkit.py: a small, stdlib-only lab for Week 11 (Proactive Agents).

It has seven parts:
  1. Events and triggers: a mini-cron, interval, one-shot, webhook and file-change triggers, a scheduler with
     active hours, an idle-only rule and a "NO_REPLY" convention, and a webhook gate.
  2. A user model in the style of a General User Model: confidence-weighted propositions with a staleness rate,
     a privacy audit, revision, a "confirm out of character updates" rule, and retrieval (BM25, recency, MMR).
  3. Decisions: expected utility of acting, asking or waiting, and the thresholds between them.
  4. Calibration, a rate limiter, quiet hours, a duplicate filter and a governor that combines them.
  5. A consent ledger with levels of permission and a hash-chained log.
  6. A simulated day of suggestions and a user who gets tired of them, to compare policies.
  7. A toy next action predictor with a confidence score and selective prediction, and metrics for logs.

Everything is deterministic. The agents and the user are programs, not language models or people. Every
number, cost, probability and threshold below is made up by me for teaching. Nothing here is a security
control. Run `python3 proactive_agents_toolkit.py` for the demo and
`python3 -m unittest test_proactive_agents_toolkit` for the tests.
"""
import hashlib
import hmac
import math
import random
import re
from collections import Counter
from dataclasses import dataclass, field

DAY = 1440
NO_REPLY = "NO_REPLY"


def rng_for(*keys):
    h = hashlib.sha256("|".join(str(k) for k in keys).encode()).hexdigest()
    return random.Random(int(h[:16], 16))


def mean(xs):
    xs = list(xs)
    return sum(xs) / len(xs) if xs else 0.0


def clip(x, lo=0.0, hi=1.0):
    return max(lo, min(hi, x))


def hhmm(t):
    m = int(t) % DAY
    return "%02d:%02d" % (m // 60, m % 60)


# ---------------------------------------------------------------------------------------------------
# 1. Events, triggers, a scheduler and a webhook gate
# ---------------------------------------------------------------------------------------------------
@dataclass(frozen=True)
class Event:
    t: int
    kind: str
    data: tuple = ()

    def get(self, key, default=None):
        for k, v in self.data:
            if k == key:
                return v
        return default


def ev(t, kind, **data):
    return Event(t, kind, tuple(sorted(data.items())))


def _field(spec, lo, hi):
    out = set()
    for part in spec.split(","):
        step = 1
        if "/" in part:
            part, s = part.split("/")
            step = int(s)
            if step < 1:
                raise ValueError("step must be at least 1")
            if part != "*" and "-" not in part:
                part = part + "-" + str(hi)
        if part == "*":
            a, b = lo, hi
        elif "-" in part:
            a, b = (int(x) for x in part.split("-"))
        else:
            a = b = int(part)
        if a < lo or b > hi or a > b:
            raise ValueError("field out of range: " + spec)
        out.update(range(a, b + 1, step))
    return out


def parse_cron(expr):
    """Three fields: minute hour weekday (0 is Monday). This is a small cousin of a real five-field cron."""
    f = expr.split()
    if len(f) != 3:
        raise ValueError("need three fields: minute hour weekday")
    return _field(f[0], 0, 59), _field(f[1], 0, 23), _field(f[2], 0, 6)


def cron_matches(expr, t):
    mins, hrs, dows = parse_cron(expr)
    m = t % DAY
    return (m % 60) in mins and (m // 60) in hrs and ((t // DAY) % 7) in dows


class Trigger:
    def fires(self, e):
        raise NotImplementedError


class Cron(Trigger):
    def __init__(self, expr):
        parse_cron(expr)
        self.expr = expr

    def fires(self, e):
        return e.kind == "tick" and cron_matches(self.expr, e.t)


class Every(Trigger):
    def __init__(self, minutes, start=0):
        if minutes < 1:
            raise ValueError("minutes must be at least 1")
        self.minutes, self.start = minutes, start

    def fires(self, e):
        return e.kind == "tick" and e.t >= self.start and (e.t - self.start) % self.minutes == 0


class At(Trigger):
    def __init__(self, t):
        self.t = t

    def fires(self, e):
        return e.kind == "tick" and e.t == self.t


class Webhook(Trigger):
    def __init__(self, source):
        self.source = source

    def fires(self, e):
        return e.kind == "webhook" and e.get("source") == self.source


class FileChange(Trigger):
    """Fires on a file event under a path prefix. With debounce, a second change within that many minutes of
    the last one that fired is ignored (the first change in a burst wins)."""

    def __init__(self, prefix, debounce=0):
        self.prefix, self.debounce, self.last = prefix, debounce, None

    def fires(self, e):
        if e.kind != "file" or not str(e.get("path", "")).startswith(self.prefix):
            return False
        if self.last is not None and e.t - self.last < self.debounce:
            return False
        self.last = e.t
        return True


def in_window(t, window):
    """window is None (always) or (start, end) in minutes of the day. Start is inclusive and end is exclusive.
    A window with start greater than end wraps past midnight. Start equal to end means the whole day (my choice)."""
    if window is None:
        return True
    s, e = window
    m = t % DAY
    if s == e:
        return True
    if s < e:
        return s <= m < e
    return m >= s or m < e


@dataclass
class Job:
    name: str
    trigger: Trigger
    handler: object = None
    active: object = None
    idle_only: bool = False
    min_gap: int = 0
    last: object = None
    pending: object = None


class Scheduler:
    """Each minute: find jobs whose trigger fires. Drop a firing that is outside the job's active hours or too
    soon after the job's last run (min_gap is my simplification of pacing). If the job is idle-only and the user
    is busy, keep the firing as pending and run it at the first idle minute inside the window. A handler that
    returns None or NO_REPLY runs silently: nothing is delivered."""

    def __init__(self, jobs):
        self.jobs, self.log = list(jobs), []

    def step(self, t, events=(), busy=False):
        evs = [ev(t, "tick")] + list(events)
        out = []
        for job in self.jobs:
            cand = None
            for e in evs:
                if job.trigger.fires(e):
                    cand = e
                    break
            if cand is None:
                cand = job.pending
            if cand is None:
                continue
            if not in_window(t, job.active):
                job.pending = None
                continue
            if job.last is not None and t - job.last < job.min_gap:
                job.pending = None
                continue
            if job.idle_only and busy:
                job.pending = cand
                continue
            job.pending = None
            reply = job.handler(cand, t) if job.handler else job.name
            job.last = t
            delivered = reply is not None and reply != NO_REPLY
            self.log.append((t, job.name, reply, delivered))
            if delivered:
                out.append((job.name, reply))
        return out

    def run(self, timeline, events_by_t=None, busy_fn=None):
        events_by_t = events_by_t or {}
        res = []
        for t in timeline:
            busy = bool(busy_fn(t)) if busy_fn else False
            for name, reply in self.step(t, events_by_t.get(t, ()), busy):
                res.append((t, name, reply))
        return res


class WebhookGate:
    """An inbound hook gate. Disabled by default. Needs the token. Identical wake text that is already the most
    recent pending item coalesces, even if the queue is full. A full queue refuses new items and evicts none.
    The docs I read give 200 for queued and coalesced and 503 for a full queue; the other codes are my choice."""

    def __init__(self, token, max_queue=8, enabled=False):
        self.token, self.max_queue, self.enabled, self.queue = token, max_queue, enabled, []

    def post(self, token, text):
        if not self.enabled:
            return 404, "disabled"
        if not hmac.compare_digest(str(token).encode(), self.token.encode()):
            return 401, "unauthorized"
        if self.queue and self.queue[-1] == text:
            return 200, "coalesced"
        if len(self.queue) >= self.max_queue:
            return 503, "full"
        self.queue.append(text)
        return 200, "queued"

    def drain(self):
        q, self.queue = self.queue, []
        return q


def route(e, trusted_sources=()):
    """Where an event goes. Short notifications from a source you control wake the main session. Anything that
    carries raw content from outside (mail, a notification body, a file) goes to an isolated reader with no tools."""
    if e.kind == "tick":
        return "main"
    if e.kind == "webhook" and e.get("source") in trusted_sources:
        return "main"
    return "isolated_reader"


# ---------------------------------------------------------------------------------------------------
# 2. A small user model: propositions with confidence and staleness, audit, revise, retrieve
# ---------------------------------------------------------------------------------------------------
STOP = {"the", "a", "an", "is", "are", "to", "of", "and", "in", "on", "for", "with", "user", "has", "this", "their"}


def tokens(text):
    return [w for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in STOP and len(w) > 1]


def jaccard(a, b):
    sa, sb = set(tokens(a)), set(tokens(b))
    if not sa and not sb:
        return 1.0
    return len(sa & sb) / len(sa | sb)


# (keywords, proposition, negates, base confidence, staleness rate per day)
RULES = [
    (("deadline", "due", "submit"), "User has an impending deadline", None, 0.6, 0.5),
    (("overleaf", "draft", "paper", "reviews"), "User is writing a paper", None, 0.6, 0.3),
    (("wedding", "invite", "rsvp"), "User is invited to a wedding", None, 0.7, 0.2),
    (("suit", "tuxedo", "formalwear"), "User needs formal wear", None, 0.4, 0.4),
    (("recipe", "icecream", "youtube"), "User is distracted by recipes while writing", None, 0.5, 0.6),
    (("billing", "invoice", "statement"), "User ignores billing notifications", None, 0.3, 0.2),
    (("vacation", "ooo", "leave"), "User is on vacation", "User has an impending deadline", 0.6, 0.5),
]
SENSITIVE = {
    "credentials": ("password", "passcode", "login", "pin"),
    "finance": ("bank", "routing", "card"),
    "health": ("diagnosis", "prescription", "symptom"),
}
CONFIRM_AT = 0.8


@dataclass
class Proposition:
    text: str
    conf: float
    alpha: float
    last_seen: float
    obs: list = field(default_factory=list)


class Audit:
    """A contextual-integrity style gate in miniature: should this observation be recorded at all? Sensitive
    categories are blocked unless the user has allowed that category."""

    def __init__(self, allowed=()):
        self.allowed, self.blocked = set(allowed), []

    def category(self, obs):
        words = set(re.findall(r"[a-z0-9]+", obs["text"].lower()))
        for cat, keys in SENSITIVE.items():
            if words & set(keys):
                return cat
        return None

    def allow(self, obs):
        cat = self.category(obs)
        if cat is None or cat in self.allowed:
            return True
        self.blocked.append(obs["id"])
        return False


class UserModel:
    def __init__(self, audit=None):
        self.props, self.pending, self.audit = [], [], audit

    def find(self, text):
        for p in self.props:
            if p.text == text:
                return p
        return None

    def propose(self, obs):
        words = set(re.findall(r"[a-z0-9]+", obs["text"].lower()))
        out = []
        for keys, text, neg, base, alpha in RULES:
            hits = len(words & set(keys))
            if hits:
                out.append((text, neg, min(1.0, base + 0.15 * (hits - 1)), alpha))
        return out

    def observe(self, obs):
        """Returns a dict: blocked (bool), added, revised, lowered, pending (lists of proposition texts)."""
        res = {"blocked": False, "added": [], "revised": [], "lowered": [], "pending": []}
        if self.audit is not None and not self.audit.allow(obs):
            res["blocked"] = True
            return res
        for text, neg, conf, alpha in self.propose(obs):
            old = self.find(text)
            if old is not None:
                old.conf = 1 - (1 - old.conf) * (1 - conf)
                old.alpha = min(old.alpha, alpha)
                old.last_seen = obs["t"]
                old.obs.append(obs["id"])
                res["revised"].append(text)
                continue
            target = self.find(neg) if neg else None
            if target is not None and target.conf >= CONFIRM_AT:
                self.pending.append((text, conf, alpha, obs))
                res["pending"].append(text)
                continue
            if target is not None:
                target.conf = target.conf * 0.5
                if target.conf < 0.05:
                    target.conf = 0.0
                res["lowered"].append(target.text)
            self.props.append(Proposition(text, conf, alpha, obs["t"], [obs["id"]]))
            res["added"].append(text)
        return res

    def confirm(self, text, yes):
        """The user answers a pending update. If yes, apply it and lower what it contradicts."""
        for i, (t, conf, alpha, obs) in enumerate(self.pending):
            if t == text:
                del self.pending[i]
                if yes:
                    for keys, ptext, neg, base, a in RULES:
                        if ptext == text and neg and self.find(neg):
                            self.find(neg).conf *= 0.5
                    self.props.append(Proposition(t, conf, alpha, obs["t"], [obs["id"]]))
                return True
        return False

    def scored(self, query, now, kdecay=2.0):
        """(proposition, score) for each proposition with confidence above 0: BM25 of the text against the query,
        times the recency factor exp(-alpha * kdecay * age in days)."""
        docs = [p for p in self.props if p.conf > 0]
        q = tokens(query)
        if not docs or not q:
            return []
        toks = [tokens(p.text) for p in docs]
        avg = mean(len(t) for t in toks) or 1.0
        n = len(docs)
        out = []
        for p, tk in zip(docs, toks):
            s = 0.0
            for w in set(q):
                df = sum(1 for t in toks if w in t)
                tf = tk.count(w)
                if tf == 0:
                    continue
                idf = math.log(1 + (n - df + 0.5) / (df + 0.5))
                s += idf * tf * 2.5 / (tf + 1.5 * (1 - 0.75 + 0.75 * len(tk) / avg))
            age = max(0.0, (now - p.last_seen) / DAY)
            out.append((p, s * math.exp(-p.alpha * kdecay * age)))
        return out

    def retrieve(self, query, now, k=3, lam=0.5, kdecay=2.0):
        """The top k by maximal marginal relevance over the scored propositions, with token Jaccard as similarity.
        Scores are scaled so the best is 1 (my choice). Propositions at confidence 0 are never returned."""
        sc = self.scored(query, now, kdecay)
        if not sc:
            return []
        top = max(s for p, s in sc)
        if top <= 0:
            return []
        docs = [p for p, s in sc]
        rel = [s / top for p, s in sc]
        idx = mmr_select(rel, lambda i, j: jaccard(docs[i].text, docs[j].text), k, lam)
        return [docs[i] for i in idx]


def mmr_select(rel, sim, k, lam=0.5):
    """Maximal marginal relevance: repeatedly take the item with the best lam * relevance minus (1 - lam) times its
    largest similarity to the items already taken. Items with zero relevance are never taken. Ties go to the
    lower index."""
    chosen, left = [], [i for i in range(len(rel)) if rel[i] > 0]
    while left and len(chosen) < k:
        def score(i):
            div = max((sim(i, j) for j in chosen), default=0.0)
            return lam * rel[i] - (1 - lam) * div
        best = max(left, key=lambda i: (score(i), -i))
        chosen.append(best)
        left.remove(best)
    return chosen


# ---------------------------------------------------------------------------------------------------
# 3. Decisions: act, ask or wait
# ---------------------------------------------------------------------------------------------------
@dataclass(frozen=True)
class Utilities:
    """Horvitz's outcomes. A means act, D means ask (a dialog), notA means wait. G is the user's goal being true."""
    a_g: float
    a_ng: float
    na_g: float
    na_ng: float
    d_g: float
    d_ng: float


def from_costs(benefit, c_fp, c_fn, c_ask=None, ask_gain=0.6):
    """Act on a wanted goal gains benefit. Act on an unwanted one costs c_fp. Wait on a wanted one costs c_fn.
    Asking costs c_ask whatever the goal (default half of c_fp), and when the goal is wanted the user keeps only
    ask_gain of the benefit, because they still have to answer and wait. Both defaults are made up."""
    if c_ask is None:
        c_ask = 0.5 * c_fp
    return Utilities(benefit, -c_fp, -c_fn, 0.0, ask_gain * benefit - c_ask, -c_ask)


def eu_act(p, u):
    return p * u.a_g + (1 - p) * u.a_ng


def eu_wait(p, u):
    return p * u.na_g + (1 - p) * u.na_ng


def eu_ask(p, u):
    return p * u.d_g + (1 - p) * u.d_ng


def _cross(y0a, y1a, y0b, y1b):
    """Where two lines in p cross. Line a is y0a at p=0 and y1a at p=1."""
    den = (y1a - y0a) - (y1b - y0b)
    if den == 0:
        raise ValueError("lines are parallel")
    return (y0b - y0a) / den


def decide(p, u):
    """The option with the highest expected utility. Ties go to the less intrusive option: wait, then ask, then act."""
    opts = [("wait", eu_wait(p, u)), ("ask", eu_ask(p, u)), ("act", eu_act(p, u))]
    best = opts[0]
    for o in opts[1:]:
        if o[1] > best[1] + 1e-12:
            best = o
    return best[0]


def thresholds(u):
    """(p_wait_to_ask, p_ask_to_act). Below the first, wait. Between them, ask. Above the second, act. If asking
    is never the best option the two numbers are equal and are the wait-to-act threshold."""
    p_wa = _cross(u.na_ng, u.na_g, u.a_ng, u.a_g)
    p_wd = _cross(u.na_ng, u.na_g, u.d_ng, u.d_g)
    p_da = _cross(u.d_ng, u.d_g, u.a_ng, u.a_g)
    if p_wd < p_da:
        return p_wd, p_da
    return p_wa, p_wa


def gumbo_eu(p, benefit, c_fp, c_fn):
    """The two expected utilities used by the interruption rule in the paper: interrupt or not."""
    return p * benefit - (1 - p) * c_fp, -p * c_fn


def gumbo_threshold(benefit, c_fp, c_fn):
    """Interrupt exactly when p is above c_fp / (benefit + c_fn + c_fp)."""
    return c_fp / (benefit + c_fn + c_fp)


def attention_cost(base, focus, k=3.0):
    """Made up: an unwanted interruption costs more the deeper the focus. focus is in [0, 1]."""
    return base * (1 + k * focus)


# ---------------------------------------------------------------------------------------------------
# 4. Calibration, rate limits, quiet hours, duplicates
# ---------------------------------------------------------------------------------------------------
def brier(ps, ys):
    return mean((p - (1.0 if y else 0.0)) ** 2 for p, y in zip(ps, ys))


def reliability(ps, ys, bins=5):
    """Rows of (low, high, n, mean predicted, observed frequency) for equal-width bins. Empty bins are left out."""
    rows = []
    for b in range(bins):
        lo, hi = b / bins, (b + 1) / bins
        idx = [i for i, p in enumerate(ps) if (lo <= p < hi) or (b == bins - 1 and p == 1.0)]
        if idx:
            rows.append((lo, hi, len(idx), mean(ps[i] for i in idx), mean(1.0 if ys[i] else 0.0 for i in idx)))
    return rows


def ece(ps, ys, bins=5):
    n = len(ps)
    return sum(r[2] / n * abs(r[3] - r[4]) for r in reliability(ps, ys, bins)) if n else 0.0


def sharpen(p, gamma):
    """Push a probability away from one half (gamma above 1) or toward it (below 1). Models an overconfident scorer."""
    a, b = p ** gamma, (1 - p) ** gamma
    return a / (a + b) if a + b > 0 else 0.5


class Calibrator:
    """Piecewise-constant recalibration fitted on a log of (score, outcome)."""

    def __init__(self, bins=10):
        self.bins, self.rate = bins, []

    def fit(self, ps, ys):
        self.rate = []
        for b in range(self.bins):
            lo, hi = b / self.bins, (b + 1) / self.bins
            idx = [i for i, p in enumerate(ps) if (lo <= p < hi) or (b == self.bins - 1 and p == 1.0)]
            if idx:
                pos = sum(1 for i in idx if ys[i])
                self.rate.append((pos + 1) / (len(idx) + 2))
            else:
                self.rate.append((lo + hi) / 2)
        return self

    def apply(self, p):
        b = min(self.bins - 1, int(p * self.bins))
        return self.rate[b]


class TokenBucket:
    def __init__(self, per_min, capacity):
        self.per_min, self.capacity = per_min, capacity
        self.level, self.t = float(capacity), None

    def _refill(self, t):
        if self.t is not None and t > self.t:
            self.level = min(self.capacity, self.level + (t - self.t) * self.per_min)
        self.t = t if self.t is None else max(self.t, t)

    def allow(self, t):
        self._refill(t)
        if self.level >= 1.0:
            self.level -= 1.0
            return True
        return False


class Governor:
    """Checks, in this order: quiet hours, daily cap, duplicate text, rate limit. Only a suggestion that passes
    all four is recorded and spends a token. check returns (ok, reason) with reason one of ok, quiet, cap,
    duplicate, rate."""

    def __init__(self, quiet=None, per_min=None, capacity=1, daily_cap=None, dup_thr=0.8, dup_window=240):
        self.quiet, self.daily_cap, self.dup_thr, self.dup_window = quiet, daily_cap, dup_thr, dup_window
        self.bucket = TokenBucket(per_min, capacity) if per_min else None
        self.sent, self.counts = [], Counter()

    def check(self, t, text):
        if self.quiet is not None and in_window(t, self.quiet):
            self.counts["quiet"] += 1
            return False, "quiet"
        day = t // DAY
        if self.daily_cap is not None and sum(1 for s in self.sent if s[0] // DAY == day) >= self.daily_cap:
            self.counts["cap"] += 1
            return False, "cap"
        for st, stext in self.sent:
            if t - st < self.dup_window and jaccard(text, stext) >= self.dup_thr:
                self.counts["duplicate"] += 1
                return False, "duplicate"
        if self.bucket is not None and not self.bucket.allow(t):
            self.counts["rate"] += 1
            return False, "rate"
        self.sent.append((t, text))
        self.counts["ok"] += 1
        return True, "ok"


# ---------------------------------------------------------------------------------------------------
# 5. A consent ledger
# ---------------------------------------------------------------------------------------------------
LEVELS = ["observe", "suggest", "draft", "act_reversible", "act_irreversible"]


class Ledger:
    """Standing grants per scope, each with a level, a start and an optional expiry, and an append-only log whose
    entries are chained by hash. A standing grant never covers an irreversible action: that needs a yes each time."""

    def __init__(self):
        self.grants, self.log = [], []

    def _append(self, t, what, outcome):
        prev = self.log[-1][4] if self.log else "start"
        h = hashlib.sha256(f"{prev}|{t}|{what}|{outcome}".encode()).hexdigest()
        self.log.append((t, what, outcome, prev, h))

    def grant(self, scope, level, t, ttl=None):
        if level not in LEVELS:
            raise ValueError("unknown level")
        self.grants.append({"scope": scope, "level": level, "start": t, "end": None if ttl is None else t + ttl, "revoked": None})
        self._append(t, f"grant {scope} {level}", "recorded")

    def revoke(self, scope, t):
        n = 0
        for g in self.grants:
            if g["scope"] == scope and g["revoked"] is None:
                g["revoked"] = t
                n += 1
        self._append(t, f"revoke {scope}", f"{n} grants")
        return n

    def max_level(self, scope, t):
        best = -1
        for g in self.grants:
            live = g["start"] <= t and (g["end"] is None or t < g["end"]) and (g["revoked"] is None or t < g["revoked"])
            if g["scope"] == scope and live:
                best = max(best, LEVELS.index(g["level"]))
        return best

    def authorize(self, scope, level, t, confirm=False):
        """Returns (decision, reason) with decision allow, ask or deny, and logs it."""
        need = LEVELS.index(level)
        have = self.max_level(scope, t)
        if have < 0:
            res = ("deny", "no live grant")
        elif level == "act_irreversible":
            if have < need:
                res = ("deny", "grant is below act_irreversible")
            elif confirm:
                res = ("allow", "confirmed for this action")
            else:
                res = ("ask", "irreversible: needs a yes each time")
        elif have >= need:
            res = ("allow", "covered by a grant")
        else:
            res = ("deny", "grant is below " + level)
        self._append(t, f"authorize {scope} {level}", res[0])
        return res

    def verify(self):
        prev = "start"
        for t, what, outcome, p, h in self.log:
            if p != prev or h != hashlib.sha256(f"{p}|{t}|{what}|{outcome}".encode()).hexdigest():
                return False
            prev = h
        return True


# ---------------------------------------------------------------------------------------------------
# 6. A simulated day of suggestions, and a user who tires of them
# ---------------------------------------------------------------------------------------------------
@dataclass(frozen=True)
class Block:
    kind: str
    start: int
    end: int
    focus: float


@dataclass(frozen=True)
class Opp:
    t: int
    useful: bool
    signal: float
    benefit: float
    text: str
    focus: float


@dataclass(frozen=True)
class Costs:
    c_fp0: float = 0.15
    k_focus: float = 3.0
    miss: float = 0.5
    ignored_frac: float = 0.3
    accept: float = 0.9
    fatigue: float = 0.85
    window: int = 90


BLOCKS = [("deep", 0.9, (45, 100), 0.35), ("meeting", 0.7, (25, 60), 0.25), ("admin", 0.3, (15, 40), 0.25), ("break", 0.1, (10, 25), 0.15)]
TOPICS = ["flights", "slides", "budget", "reviewers", "lunch", "invoice", "refactor", "reading", "taxes", "gift", "rsvp", "backup"]


def make_day(seed, start=420, end=1380):
    r = rng_for("day", seed)
    out, t = [], start
    while t < end:
        x, acc = r.random(), 0.0
        pick = BLOCKS[-1]
        for b in BLOCKS:
            acc += b[3]
            if x < acc:
                pick = b
                break
        dur = min(r.randint(*pick[2]), end - t)
        out.append(Block(pick[0], t, t + dur, pick[1]))
        t += dur
    return out


def focus_at(blocks, t):
    for b in blocks:
        if b.start <= t < b.end:
            return b.focus
    return 0.0


def make_opps(seed, blocks, per_hour=8.5, base=0.3):
    r = rng_for("opps", seed)
    out = []
    for t in range(blocks[0].start, blocks[-1].end):
        if r.random() < per_hour / 60:
            useful = r.random() < base
            sig = clip(r.gauss(0.62 if useful else 0.38, 0.2))
            out.append(Opp(t, useful, sig, r.uniform(0.5, 1.5), "suggest %s v%d" % (r.choice(TOPICS), r.randint(1, 3)), focus_at(blocks, t)))
    return out


def prob(o, gamma=1.0, cal=None):
    p = sharpen(clip(o.signal, 1e-6, 1 - 1e-6), gamma)
    return cal.apply(p) if cal is not None else p


def c_fp(costs, focus):
    return attention_cost(costs.c_fp0, focus, costs.k_focus)


def never():
    return lambda o, n=0: False


def always():
    return lambda o, n=0: True


def fixed(thr, gamma=1.0, cal=None):
    return lambda o, n=0: prob(o, gamma, cal) >= thr


def eu_policy(costs, gamma=1.0, cal=None, benefit=1.0, aware=False):
    """Interrupt when the calibrated chance is above the threshold from the expected utilities. If aware, the benefit
    is shrunk by fatigue ** n, where n is the number of suggestions the user ignored in the last window (my addition)."""
    def d(o, n=0):
        b = benefit * (costs.fatigue ** n if aware else 1.0)
        return prob(o, gamma, cal) > gumbo_threshold(b, c_fp(costs, o.focus), costs.miss * benefit)
    return d


@dataclass
class Row:
    opp: Opp
    shown: bool
    accepted: bool
    util: float
    reason: str


class Result:
    def __init__(self, rows):
        self.rows = rows

    def net(self):
        return sum(r.util for r in self.rows)

    def shown(self):
        return sum(1 for r in self.rows if r.shown)

    def accepted(self):
        return sum(1 for r in self.rows if r.accepted)

    def useful(self):
        return sum(1 for r in self.rows if r.opp.useful)

    def precision(self):
        s = self.shown()
        return self.accepted() / s if s else 0.0

    def recall(self):
        u = self.useful()
        return self.accepted() / u if u else 0.0

    def per_hour(self):
        if not self.rows:
            return 0.0
        span = (self.rows[-1].opp.t - self.rows[0].opp.t) / 60
        return self.shown() / span if span > 0 else 0.0


def run_policy(opps, decide_fn, seed, costs=None, governor=None):
    costs = costs or Costs()
    r = rng_for("user", seed)
    ignored, rows = [], []
    for o in opps:
        n = sum(1 for x in ignored if o.t - x < costs.window)
        want, reason = decide_fn(o, n), "policy"
        if want and governor is not None:
            want, reason = governor.check(o.t, o.text)
        shown = accepted = False
        if want:
            shown = True
            draw = r.random()
            accepted = o.useful and draw < costs.accept * costs.fatigue ** n
            if not accepted:
                ignored.append(o.t)
        cost = c_fp(costs, o.focus)
        if shown and accepted:
            u = o.benefit
        elif shown and o.useful:
            u = -costs.ignored_frac * cost
        elif shown:
            u = -cost
        elif o.useful:
            u = -costs.miss * o.benefit
        else:
            u = 0.0
        rows.append(Row(o, shown, accepted, u, reason))
    return Result(rows)


def study(policy_factory, seeds, costs=None, governor_factory=None):
    """Run a policy over several simulated days. Returns a list of Result objects."""
    out = []
    for s in seeds:
        opps = make_opps(s, make_day(s))
        g = governor_factory() if governor_factory else None
        out.append(run_policy(opps, policy_factory(), s, costs, g))
    return out


def summarize(results):
    return {
        "net": mean(r.net() for r in results),
        "shown": mean(r.shown() for r in results),
        "accepted": mean(r.accepted() for r in results),
        "precision": mean(r.precision() for r in results),
        "recall": mean(r.recall() for r in results),
        "per_hour": mean(r.per_hour() for r in results),
    }


def fit_calibrator(seeds, gamma=1.0, bins=10):
    ps, ys = [], []
    for s in seeds:
        for o in make_opps(s, make_day(s)):
            ps.append(prob(o, gamma))
            ys.append(o.useful)
    return Calibrator(bins).fit(ps, ys)


# ---------------------------------------------------------------------------------------------------
# 7. A toy next action predictor, selective prediction, and metrics for logs
# ---------------------------------------------------------------------------------------------------
def habit_world(seed, n=3000, n_actions=40, strength=0.2):
    """A sequence of actions where each action has a favorite next action. How strongly it is followed differs
    by action. Returns the sequence and the favorite map. Made-up."""
    r = rng_for("habit", seed)
    acts = ["act%02d" % i for i in range(n_actions)]
    fav = {a: r.choice(acts) for a in acts}
    strong = {a: clip(strength * r.uniform(0.2, 2.2), 0.0, 0.95) for a in acts}
    seq = [r.choice(acts)]
    for _ in range(n - 1):
        prev = seq[-1]
        seq.append(fav[prev] if r.random() < strong[prev] else r.choice(acts))
    return seq, fav


class NextActionModel:
    def __init__(self, actions, smooth=0.5):
        self.actions, self.smooth = list(actions), smooth
        self.counts = {a: Counter() for a in actions}

    def fit(self, seq):
        for a, b in zip(seq, seq[1:]):
            self.counts[a][b] += 1
        return self

    def sample(self, prev, r):
        w = [self.counts[prev][a] + self.smooth for a in self.actions]
        return r.choices(self.actions, weights=w)[0]

    def predict(self, prev, k, r):
        return [self.sample(prev, r) for _ in range(k)]


def agreement(samples):
    """The share of samples that agree with the most common one. A stand-in for the spread of sampled predictions."""
    return max(Counter(samples).values()) / len(samples)


def nap_eval(seq, seed=0, split=0.7, k=20, top=0.1):
    """Train on the first part and test on the rest. pass1 is the mean share of samples equal to the truth, passk
    is the share of cases where any of k samples equals it, and sel1 is pass1 among the top share of cases by
    agreement. Exact match replaces the paper's similarity judge. Returns a dict."""
    cut = int(len(seq) * split)
    model = NextActionModel(sorted(set(seq))).fit(seq[:cut])
    r = rng_for("nap", seed)
    rows = []
    for prev, truth in zip(seq[cut - 1:-1], seq[cut:]):
        s = model.predict(prev, k, r)
        rows.append((agreement(s), s.count(truth) / k, truth in s))
    rows.sort(key=lambda x: -x[0])
    m = max(1, int(len(rows) * top))
    return {"n": len(rows), "pass1": mean(x[1] for x in rows), "passk": mean(1.0 if x[2] else 0.0 for x in rows),
            "sel1": mean(x[1] for x in rows[:m]), "all_conf": mean(x[0] for x in rows)}


def risk_coverage(rows, cuts=(1.0, 0.5, 0.25, 0.1)):
    """rows are (confidence, correct). Returns (coverage, accuracy among the most confident share) per cut."""
    rows = sorted(rows, key=lambda x: -x[0])
    out = []
    for c in cuts:
        m = max(1, int(round(len(rows) * c)))
        out.append((c, mean(x[1] for x in rows[:m])))
    return out


def precision_from_log(records):
    """records are (shown, accepted). Precision is the only quantity a log of shown suggestions can give."""
    shown = [a for s, a in records if s]
    return sum(1 for a in shown if a) / len(shown) if shown else 0.0


def estimate_recall(accepted, suppressed_total, audited_n, audited_useful):
    """Recall needs labels for what was NOT shown. Audit a random sample of suppressed suggestions, estimate how many
    suppressed ones were useful, and compare with the accepted ones."""
    if audited_n <= 0:
        raise ValueError("need audited suppressed suggestions")
    missed = suppressed_total * audited_useful / audited_n
    return accepted / (accepted + missed) if accepted + missed > 0 else 0.0


def f1(p, r):
    return 2 * p * r / (p + r) if p + r > 0 else 0.0


def lu_accept(task_shown, need, user_accepts):
    """The paper's acceptance rule: accepted if a task was shown and accepted, or none was shown and none was needed."""
    return (task_shown and user_accepts) or ((not task_shown) and (not need))


def wilson(k, n, z=1.96):
    if n == 0:
        return 0.0, 1.0
    p = k / n
    d = 1 + z * z / n
    c = p + z * z / (2 * n)
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))
    return (c - h) / d, (c + h) / d


# ---------------------------------------------------------------------------------------------------
# Demo
# ---------------------------------------------------------------------------------------------------
def _h(n, title):
    print("\n== %d. %s" % (n, title))


def demo():
    C = Costs()
    test_seeds = range(10, 40)
    cal = fit_calibrator(range(0, 10))

    _h(1, "Triggers over one Monday, with active hours, an idle-only job and NO_REPLY (made up)")
    new_items = {t: 1 for t in (545, 553, 700, 702, 1010)}
    seen = {"n": 0}

    def inbox_check(e, t):
        n = sum(v for k, v in new_items.items() if k <= t)
        d, seen["n"] = n - seen["n"], n
        return NO_REPLY if d == 0 else "%d new item(s)" % d

    jobs = [
        Job("daily 09:00 brief", Cron("0 9 0-4"), lambda e, t: "brief at " + hhmm(t)),
        Job("inbox check", Every(30, start=480), inbox_check, active=(480, 1080), idle_only=True),
        Job("on file change", FileChange("src/", debounce=5), lambda e, t: "changed " + e.get("path")),
        Job("on ci hook", Webhook("ci"), lambda e, t: "ci says " + e.get("text")),
    ]
    events = {540: [ev(540, "file", path="src/a.py")], 542: [ev(542, "file", path="src/b.py")],
              548: [ev(548, "file", path="src/a.py")], 560: [ev(560, "file", path="docs/x.md")],
              900: [ev(900, "webhook", source="ci", text="build red")],
              901: [ev(901, "webhook", source="other", text="ignored")]}
    sch = Scheduler(jobs)
    out = sch.run(range(0, DAY), events, busy_fn=lambda t: 600 <= t < 660)
    for t, name, reply in out:
        print("  %s  %-18s %s" % (hhmm(t), name, reply))
    runs, quiet = Counter(l[1] for l in sch.log), sum(1 for l in sch.log if not l[3])
    print("  runs per job:", dict(runs), "| silent runs (NO_REPLY):", quiet)
    print("  routing: ci hook from a trusted source ->", route(ev(900, "webhook", source="ci"), ("ci",)),
          "| an email event ->", route(ev(900, "email", body="raw text")))

    _h(2, "A webhook gate (made-up token)")
    g = WebhookGate("s3cret", max_queue=2)
    for label, tok, txt, en in [("disabled", "s3cret", "a", False), ("wrong token", "nope", "a", True), ("first", "s3cret", "a", True),
                                ("same text again", "s3cret", "a", True), ("second", "s3cret", "b", True),
                                ("third (queue is full)", "s3cret", "c", True), ("same as last, full queue", "s3cret", "b", True)]:
        g.enabled = en
        print("  %-26s -> %s" % (label, g.post(tok, txt)))
    print("  drained:", g.drain())

    _h(3, "A small user model: propose, audit, revise, ask, retrieve (made-up observations)")
    m = UserModel(Audit())
    obs = [dict(id="o1", t=0, source="screen", text="Overleaf draft of the paper, reviews open, deadline Friday"),
           dict(id="o2", t=60, source="mail", text="Wedding invite from Sam, please RSVP"),
           dict(id="o3", t=120, source="screen", text="Bank site, password field, login"),
           dict(id="o4", t=200, source="mail", text="Reminder: submit by the deadline"),
           dict(id="o5", t=300, source="mail", text="Spam: you won a vacation, leave now, ooo")]
    for o in obs:
        r = m.observe(o)
        print("  %s %-6s blocked=%-5s added=%s revised=%s pending=%s" % (o["id"], o["source"], r["blocked"], len(r["added"]), len(r["revised"]), r["pending"]))
    print("  the user answers no to:", m.pending[0][0], "->", m.confirm(m.pending[0][0], False))
    now = 300 + 2 * DAY
    print("  retrieve 'paper deadline wedding' two days later:")
    for p in m.retrieve("paper deadline wedding", now, k=3):
        print("    conf %.2f  staleness %.1f/day  %s" % (p.conf, p.alpha, p.text))

    _h(4, "Act, ask or wait: thresholds from the costs (made up)")
    for name, b, fp, fn, ca in [("break (focus 0.1)", 1.0, attention_cost(0.15, 0.1), 0.5, None), ("admin (focus 0.3)", 1.0, attention_cost(0.15, 0.3), 0.5, None),
                                ("deep work (focus 0.9)", 1.0, attention_cost(0.15, 0.9), 0.5, None), ("send an email (c_fp 2.0)", 1.0, 2.0, 0.5, 0.2)]:
        u = from_costs(b, fp, fn, ca)
        lo, hi = thresholds(u)
        print("  %-24s c_fp %.3f  wait below %.3f, ask up to %.3f, act above | at p=0.25: %s, p=0.45: %s, p=0.70: %s, p=0.90: %s"
              % (name, fp, lo, hi, decide(0.25, u), decide(0.45, u), decide(0.70, u), decide(0.90, u)))
    print("  same numbers as an interrupt rule: eu(interrupt)=%.3f, eu(not)=%.3f at p=0.3" % gumbo_eu(0.3, 1.0, 0.555, 0.5))

    _h(5, "Policies over %d simulated days (made up; costs, fatigue and the model's scores are all invented)" % len(test_seeds))
    gov = lambda: Governor(quiet=(22 * 60, 8 * 60), per_min=1 / 10, capacity=3, daily_cap=40)
    rows = [("never", study(never, test_seeds, C)), ("always", study(always, test_seeds, C)),
            ("fixed 0.5 on raw scores", study(lambda: fixed(0.5), test_seeds, C)),
            ("expected utility, raw scores", study(lambda: eu_policy(C), test_seeds, C)),
            ("expected utility, calibrated", study(lambda: eu_policy(C, cal=cal), test_seeds, C)),
            ("  + rate limit and quiet hours", study(lambda: eu_policy(C, cal=cal), test_seeds, C, gov)),
            ("  fatigue-aware", study(lambda: eu_policy(C, cal=cal, aware=True), test_seeds, C)),
            ("fixed 0.5 on calibrated", study(lambda: fixed(0.5, cal=cal), test_seeds, C))]
    print("  %-30s %8s %7s %9s %10s %7s %6s" % ("policy", "net/day", "shown", "accepted", "precision", "recall", "per h"))
    for name, res in rows:
        s = summarize(res)
        print("  %-30s %8.2f %7.1f %9.1f %10.2f %7.2f %6.1f" % (name, s["net"], s["shown"], s["accepted"], s["precision"], s["recall"], s["per_hour"]))

    _h(6, "A threshold sweep on calibrated scores, and how well calibrated the raw scores are")
    for th in (0.2, 0.3, 0.4, 0.5, 0.6, 0.7):
        s = summarize(study(lambda: fixed(th, cal=cal), test_seeds, C))
        print("  threshold %.1f  net/day %6.2f  shown %5.1f  precision %.2f  recall %.2f" % (th, s["net"], s["shown"], s["precision"], s["recall"]))
    ps, ys = [], []
    for sd in test_seeds:
        for o in make_opps(sd, make_day(sd)):
            ps.append(prob(o)); ys.append(o.useful)
    print("  raw scores:        Brier %.3f  ECE %.3f   base rate of useful %.2f" % (brier(ps, ys), ece(ps, ys), mean(1.0 if y else 0.0 for y in ys)))
    for lo, hi, n, mp, fr in reliability(ps, ys):
        print("    scores %.1f to %.1f: n %4d, mean score %.2f, useful %.2f" % (lo, hi, n, mp, fr))
    cps = [cal.apply(p) for p in ps]
    print("  after recalibration: Brier %.3f  ECE %.3f" % (brier(cps, ys), ece(cps, ys)))
    print("  (the recalibrator was fitted on days 0 to 9 and is scored here on days 10 to 39)")

    _h(7, "What the governor refused (one run of 30 days, calibrated expected-utility policy)")
    tot = Counter()
    for sd in test_seeds:
        gv = gov()
        run_policy(make_opps(sd, make_day(sd)), eu_policy(C, cal=cal), sd, C, gv)
        tot.update(gv.counts)
    print("  checked %d suggestions that the policy wanted to show: %s" % (sum(tot.values()), dict(sorted(tot.items()))))

    _h(8, "Logs: precision is observable, recall needs audited labels (made-up logs from section 5)")
    res = study(lambda: eu_policy(C, cal=cal), test_seeds, C)
    recs = [(r.shown, r.accepted) for x in res for r in x.rows]
    acc = sum(1 for s, a in recs if a)
    supp = [r for x in res for r in x.rows if not r.shown]
    ar = rng_for("audit", 0)
    aud = [r for r in supp if ar.random() < 0.1]
    est = estimate_recall(acc, len(supp), len(aud), sum(1 for r in aud if r.opp.useful))
    p = precision_from_log(recs)
    true_rec = acc / (acc + sum(1 for r in supp if r.opp.useful))
    print("  precision from the log %.3f | recall (accepted over accepted plus useful but suppressed): true %.3f, estimated from %d audited %.3f | F1 %.3f"
          % (p, true_rec, len(aud), est, f1(p, est)))
    lo, hi = wilson(acc, sum(1 for s, a in recs if s))
    print("  acceptance rate %.3f, 95%% Wilson interval %.3f to %.3f (n=%d shown)" % (p, lo, hi, sum(1 for s, a in recs if s)))
    t3 = {"GPT-4o": (0.9811, 0.4815, 0.6460), "Qwen2-7B-Proactive": (1.0000, 0.4978, 0.6647)}
    for k, (rc, pr, ff) in t3.items():
        print("  check of a published row (%s): F1 from recall %.4f and precision %.4f is %.4f; the paper's table says %.4f" % (k, rc, pr, f1(pr, rc), ff))

    _h(9, "Next action prediction with a confidence score (made-up habits, exact match instead of a judge)")
    tot = Counter(); n = 0
    for sd in range(5):
        seq, _ = habit_world(sd)
        r = nap_eval(seq, seed=sd)
        for k in ("pass1", "passk", "sel1", "all_conf"):
            tot[k] += r[k] / 5
        n = r["n"]
    print("  5 users, %d test steps each, 20 samples per step: pass@1 %.3f, pass@20 %.3f, pass@1 among the top 10%% by agreement %.3f, mean agreement %.2f"
          % (n, tot["pass1"], tot["passk"], tot["sel1"], tot["all_conf"]))

    _h(10, "A consent ledger (made-up scopes)")
    L = Ledger()
    L.grant("screen", "observe", 0)
    L.grant("calendar", "draft", 0, ttl=120)
    L.grant("shop", "act_reversible", 0)
    L.grant("bank", "act_irreversible", 0)
    for scope, lvl, t, cf in [("screen", "observe", 10, False), ("screen", "suggest", 10, False), ("calendar", "draft", 100, False),
                              ("calendar", "draft", 130, False), ("shop", "act_reversible", 50, False), ("shop", "act_irreversible", 50, False),
                              ("bank", "act_irreversible", 50, False), ("bank", "act_irreversible", 50, True), ("mail", "suggest", 50, False)]:
        print("  t=%3d %-9s %-17s confirm=%-5s -> %s (%s)" % (t, scope, lvl, cf, *L.authorize(scope, lvl, t, cf)))
    L.revoke("screen", 200)
    print("  after revoke at t=200:", L.authorize("screen", "observe", 210)[0])
    ok1 = L.verify()
    L.log[1] = (L.log[1][0], "grant screen act_irreversible", L.log[1][2], L.log[1][3], L.log[1][4])
    print("  log entries %d, chain valid %s; after editing one entry, chain valid %s" % (len(L.log), ok1, L.verify()))


if __name__ == "__main__":
    demo()