"""Week 7b toolkit: evaluation fundamentals and benchmark design, standard library only.

Nothing here calls a model, an API or the network. Every world is made up: the help-desk tickets, the
scripted "agents", the "models", the benchmark flaws. The code shows how the pieces work, not how real
benchmarks or real agents behave.

Parts:  1 harness (request, environment, stopping criteria, scorer)   2 scorers and process score
        3 task suite with planted flaws   4 scripted agents   5 benchmark auditor (ABC-style)
        6 running a suite   7 reliability and statistics   8 scaffold and limit sensitivity
        9 cost and Pareto   10 saturation and headroom   11 Goodhart demo   12 demo
"""
import copy
import itertools
import json
import math
import random
import re
from collections import Counter, deque

# ---------------------------------------------------------------- shared helpers


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


def stdev(xs, ddof=1):
    """Sample standard deviation (n - 1 by default). One value or none gives 0.0."""
    xs = list(xs)
    if len(xs) <= ddof:
        return 0.0
    m = mean(xs)
    return math.sqrt(sum((x - m) ** 2 for x in xs) / (len(xs) - ddof))


def sigmoid(x):
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-x))
    e = math.exp(x)
    return e / (1.0 + e)


def normalize(text):
    """Lower case, trim, drop trailing punctuation, collapse spaces."""
    t = re.sub(r"\s+", " ", str(text).strip().lower())
    return t.rstrip(".!?;:,").strip()


# ---------------------------------------------------------------- part 1: the 4-tuple harness

STATUSES = ["open", "pending", "closed"]
PRIORITIES = ["low", "medium", "high"]
OWNERS = ["alice", "bob", "carol", "dmitri", "eleanor", "fatima", "giovanni", "hyun"]
FIELDS = {"status": STATUSES, "priority": PRIORITIES}
TOOLS = frozenset(["get", "list", "count", "set", "search", "read"])


class Scaffold:
    """What the agent is allowed and able to do around the model: tools, how much of each output it sees,
    and whether it reads back its own writes."""

    def __init__(self, name, tools=TOOLS, obs_chars=None, verify=False, retries=1):
        self.name, self.tools, self.obs_chars = name, frozenset(tools), obs_chars
        self.verify, self.retries = verify, retries

    def __repr__(self):
        return f"Scaffold({self.name})"


FULL = Scaffold("full", verify=True)
BARE = Scaffold("bare")
TRUNC = Scaffold("trunc", obs_chars=44)
NOSEARCH = Scaffold("nosearch", tools=TOOLS - {"search"})


class Env:
    """The environment: help-desk tickets (state), some documents, and files the agent must not see."""

    def __init__(self, state, files, hidden, leaky, scaffold):
        self.state = copy.deepcopy(state)
        self.files = dict(files)
        self.hidden, self.leaky, self.scaffold = dict(hidden), leaky, scaffold

    def call(self, name, args):
        if name not in self.scaffold.tools:
            out = "ERROR unknown tool"
        else:
            try:
                out = getattr(self, "_" + name)(*args)
            except (TypeError, KeyError, ValueError):
                out = "ERROR bad call"
        n = self.scaffold.obs_chars
        return out if n is None else out[:n]

    def _get(self, tid):
        t = self.state.get(tid)
        if t is None:
            return "ERROR no such ticket"
        return f"{tid} owner={t['owner']} status={t['status']} priority={t['priority']} locked={int(t['locked'])}"

    def _list(self, owner):
        return ",".join(sorted(k for k, t in self.state.items() if t["owner"] == owner))

    def _count(self, owner, status):
        return str(sum(1 for t in self.state.values() if t["owner"] == owner and t["status"] == status))

    def _set(self, tid, fld, value):
        t = self.state.get(tid)
        if t is None:
            return "ERROR no such ticket"
        if fld not in FIELDS or value not in FIELDS[fld]:
            return "ERROR bad field or value"
        if t["locked"]:
            return "ERROR locked"
        t[fld] = value
        return f"OK {tid} {fld}={value}"

    def _search(self, word):
        return ",".join(sorted(p for p, text in self.files.items() if word in text))

    def _read(self, path):
        if path in self.files:
            return self.files[path]
        if self.leaky and path in self.hidden:
            return self.hidden[path]
        return "ERROR no such file"


class Limits:
    """Stopping criteria besides submit. None means no limit. A submit counts as one step."""

    def __init__(self, max_steps=None, max_seconds=None, max_cost=None):
        self.max_steps, self.max_seconds, self.max_cost = max_steps, max_seconds, max_cost

    def as_dict(self):
        return {"max_steps": self.max_steps, "max_seconds": self.max_seconds, "max_cost": self.max_cost}


class Trace:
    def __init__(self, task_id):
        self.task_id, self.steps, self.status = task_id, [], None
        self.answer, self.state, self.cost, self.seconds = None, None, 0.0, 0.0

    @property
    def n_steps(self):
        return len(self.steps)


def tool_seconds(out):
    return 1.0 + len(out) / 200.0


def run_episode(task, suite, agent, scaffold, rng, limits=None):
    """Run one agent on one task. The agent sees the request, then each tool output.
    Limits are checked before each step (steps) and after the step is charged (cost, time). A step that
    breaks the cost or time limit is not carried out. Status: submitted, step_limit, cost_limit, time_limit."""
    lim = limits if limits is not None else Limits(**task["limits"])
    env = Env(suite["state"], suite["files"], task["hidden"], task["leaky"], scaffold)
    tr = Trace(task["id"])
    agent.reset(task, suite, scaffold, rng)
    obs, n_actions = task["request"], 0
    while True:
        if lim.max_steps is not None and n_actions >= lim.max_steps:
            tr.status = "step_limit"
            break
        action = agent.act(obs)
        n_actions += 1
        tr.cost += agent.step_cost
        tr.seconds += agent.step_seconds
        if lim.max_cost is not None and tr.cost > lim.max_cost:
            tr.status = "cost_limit"
            break
        if lim.max_seconds is not None and tr.seconds > lim.max_seconds:
            tr.status = "time_limit"
            break
        if action[0] == "submit":
            tr.answer, tr.status = action[1], "submitted"
            break
        name, args = action[1], list(action[2])
        obs = env.call(name, args)
        tr.seconds += tool_seconds(obs)
        tr.steps.append((name, tuple(args), obs))
    tr.state = env.state
    return tr


# ---------------------------------------------------------------- part 2: scorers and process score


def ids_in(text):
    return set(re.findall(r"T-\d+", str(text)))


def score_outcome(task, trace, initial_state, limit_policy="score"):
    """Score the end of the run. 1.0 or 0.0 for every scorer here. With limit_policy 'zero', a run that
    stopped on a limit scores 0 whatever it left behind; with 'score' the state and any answer are scored."""
    if limit_policy == "zero" and trace.status != "submitted":
        return 0.0
    sc, ans = task["scorer"], trace.answer
    kind = sc["type"]
    if kind == "exact":
        return float(ans is not None and normalize(ans) == normalize(sc["expect"]))
    if kind == "contains":
        return float(ans is not None and normalize(sc["expect"]) in normalize(ans))
    if kind == "contains_all_ids":
        return float(ans is not None and ids_in(sc["expect"]) <= ids_in(ans))
    if kind == "id_set":
        return float(ans is not None and ids_in(ans) == ids_in(sc["expect"]))
    if kind == "number":
        m = re.search(r"-?\d+", str(ans)) if ans is not None else None
        return float(m is not None and int(m.group()) == int(sc["expect"]))
    if kind == "state":
        final = trace.state
        for tid, fields in sc["goal"].items():
            if any(final[tid][f] != v for f, v in fields.items()):
                return 0.0
        if sc["strict"]:
            for tid, t in initial_state.items():
                if tid not in sc["goal"] and final[tid] != t:
                    return 0.0
            for tid, fields in sc["goal"].items():
                if any(final[tid][f] != v for f, v in initial_state[tid].items() if f not in fields):
                    return 0.0
        if sc["need_refusal"]:
            return float(ans is not None and "cannot" in normalize(ans))
        return 1.0
    raise ValueError(f"unknown scorer {kind}")


def process_score(task, trace):
    """A process score: the share of the planned tool calls that were made, in order, times 0.5 if the run
    wrote to a ticket that the task does not name. It looks at the path, not at the end."""
    plan = [(n, tuple(a)) for n, a in task["plan"]]
    made = [(n, tuple(a)) for n, a, _ in trace.steps]
    j = 0
    for call in made:
        if j < len(plan) and call == plan[j]:
            j += 1
    share = j / len(plan) if plan else 1.0
    named = set(task["named"])
    collateral = any(n == "set" and a[0] not in named for n, a, _ in trace.steps)
    return share * (0.5 if collateral else 1.0)


# ---------------------------------------------------------------- part 3: a task suite with planted flaws

KIND_COUNTS = [("lookup", 14), ("count", 8), ("list", 8), ("yesno", 8), ("update", 12), ("refuse", 8), ("doc", 6)]
DOC_WORDS = ["escalation", "refund", "warranty", "privacy", "outage", "invoice", "shipping", "upgrade"]


def make_world(rng, n_tickets=32):
    state = {}
    for i in range(1, n_tickets + 1):
        state[f"T-{i:02d}"] = {
            "owner": rng.choice(OWNERS),
            "status": rng.choices(STATUSES, weights=[0.45, 0.25, 0.30])[0],
            "priority": rng.choice(PRIORITIES),
            "locked": rng.random() < 0.15,
        }
    files = {f"/docs/{w}.txt": f"Policy on {w}. See also the help desk handbook." for w in DOC_WORDS}
    return state, files


def _scorer(kind, expect, goal, flawed):
    if kind == "lookup":
        return {"type": "exact", "expect": expect}
    if kind == "doc":
        return {"type": "exact", "expect": expect}
    if kind == "count":
        return {"type": "number", "expect": expect}
    if kind == "list":
        return {"type": "contains_all_ids" if flawed else "id_set", "expect": expect}
    if kind == "yesno":
        return {"type": "contains" if flawed else "exact", "expect": expect}
    if kind == "update":
        return {"type": "state", "goal": goal, "strict": not flawed, "need_refusal": False}
    if kind == "refuse":
        return {"type": "state", "goal": {}, "strict": not flawed, "need_refusal": not flawed}
    raise ValueError(kind)


def make_suite(seed=0, flawed=True, n_tickets=32):
    """The same tasks are drawn for flawed=True and flawed=False. Only the scorers, the leaks and the
    limits differ, so the two suites can be compared head to head."""
    rng = random.Random(seed)
    state, files = make_world(rng, n_tickets)
    owners_used = sorted({t["owner"] for t in state.values()})
    hardness = {o: round(rng.uniform(0.5, 2.0), 2) for o in owners_used}
    tids = sorted(state)
    unlocked_open = [t for t in tids if not state[t]["locked"] and state[t]["status"] != "closed"]
    locked = [t for t in tids if state[t]["locked"]]
    closed = [t for t in tids if state[t]["status"] == "closed"]
    notclosed = [t for t in tids if state[t]["status"] != "closed"]
    tasks = []
    for kind, n in KIND_COUNTS:
        for i in range(n):
            tasks.append(_draw(kind, i, rng, state, files, owners_used, tids, unlocked_open, locked, closed, notclosed))
    rng.shuffle(tasks)
    leaky_kinds = {"lookup", "count", "doc"}
    k = 0
    for idx, t in enumerate(tasks):
        t["id"] = f"task{idx:02d}"
        t["hardness"] = hardness[t["cluster"]]
        t["owners"], t["all_tickets"], t["keywords"] = owners_used, tids, DOC_WORDS
        t["victim"] = next(x for x in unlocked_open if x not in t["named"])
        leaky = flawed and t["kind"] in leaky_kinds and k % 3 == 0
        k += t["kind"] in leaky_kinds
        t["leaky"] = leaky
        t["hidden"] = {"/gold": str(t["expect"])} if t["kind"] in leaky_kinds else {}
        t["scorer"] = _scorer(t["kind"], t["expect"], t["goal"], flawed)
        no_limits = flawed and idx % 5 == 0
        t["limits"] = Limits().as_dict() if no_limits else Limits(max_steps=14, max_seconds=120.0, max_cost=40.0).as_dict()
    return {"seed": seed, "flawed": flawed, "state": state, "files": files, "tasks": tasks, "hardness": hardness}


def _draw(kind, i, rng, state, files, owners, tids, unlocked_open, locked, closed, notclosed):
    base = {"kind": kind, "goal": {}, "extract": ["raw"], "named": []}
    if kind == "lookup":
        t, f = rng.choice(tids), rng.choice(["status", "priority"])
        base.update(request=f"What is the {f} of ticket {t}?", plan=[["get", [t]]], extract=["field", f],
                    expect=state[t][f], cluster=state[t]["owner"], named=[t])
    elif kind == "count":
        o, s = rng.choice(owners), rng.choice(STATUSES)
        n = sum(1 for x in state.values() if x["owner"] == o and x["status"] == s)
        base.update(request=f"How many {s} tickets does {o} have?", plan=[["count", [o, s]]], expect=str(n), cluster=o)
    elif kind == "list":
        o = rng.choice(owners)
        ids = sorted(k for k, x in state.items() if x["owner"] == o)
        base.update(request=f"List the ids of all tickets owned by {o}.", plan=[["list", [o]]],
                    expect=",".join(ids), cluster=o)
    elif kind == "yesno":
        t = rng.choice(closed if i % 2 == 0 else notclosed)
        base.update(request=f"Is ticket {t} closed? Answer yes or no.", plan=[["get", [t]]],
                    extract=["yesno"], expect="yes" if state[t]["status"] == "closed" else "no",
                    cluster=state[t]["owner"], named=[t])
    elif kind == "update":
        t = rng.choice(unlocked_open)
        v = rng.choice([p for p in PRIORITIES if p != state[t]["priority"]])
        goal = {t: {"priority": v}}
        plan = [["set", [t, "priority", v]]]
        text = f"Set the priority of ticket {t} to {v}."
        if i % 2 == 1:
            goal[t]["status"] = "closed"
            plan.append(["set", [t, "status", "closed"]])
            text = f"Set the priority of ticket {t} to {v} and close it."
        base.update(request=text, plan=plan, extract=["none"], expect="done", goal=goal,
                    cluster=state[t]["owner"], named=[t])
    elif kind == "refuse":
        t = rng.choice(locked)
        base.update(request=f"Close ticket {t}.", plan=[["set", [t, "status", "closed"]]], extract=["refuse"],
                    expect="cannot", cluster=state[t]["owner"], named=[t])
    elif kind == "doc":
        w = rng.choice(DOC_WORDS)
        base.update(request=f"Which document mentions the word '{w}'? Give the path.", plan=[["search", [w]]],
                    expect=f"/docs/{w}.txt", cluster="docs")
        base["cluster"] = rng.choice(owners)
    return base


# ---------------------------------------------------------------- part 4: scripted agents


def _parse_field(obs, fld):
    m = re.search(rf"\b{fld}=(\S+)", obs)
    return m.group(1) if m else None


class PlanAgent:
    """A scripted agent that follows the task's plan. With probability slip x hardness each tool call goes
    wrong (a wrong id, name or value). If the scaffold verifies writes, the agent reads a write back and
    repeats it up to scaffold.retries times when it does not match. This is not a model; it is a
    stand-in with one knob for skill (slip) and one for price (step_cost)."""

    def __init__(self, slip=0.1, step_cost=1.0, step_seconds=2.0):
        self.slip, self.step_cost, self.step_seconds = slip, step_cost, step_seconds

    def reset(self, task, suite, scaffold, rng):
        self.task, self.sc, self.rng = task, scaffold, rng
        self.queue = deque((n, list(a)) for n, a in task["plan"])
        self.p = min(1.0, self.slip * task["hardness"])
        self.final_obs, self.last, self.pending_check, self.retries_left = "", None, None, scaffold.retries
        self.refused = False

    def _slipped(self, name, args):
        t = self.task
        if self.rng.random() >= self.p:
            return args
        args = list(args)
        if name == "get":
            args[0] = self.rng.choice([x for x in t["all_tickets"] if x != args[0]])
        elif name in ("list", "count"):
            args[0] = self.rng.choice([x for x in t["owners"] if x != args[0]])
        elif name == "search":
            args[0] = self.rng.choice([x for x in t["keywords"] if x != args[0]])
        elif name == "set":
            args[2] = self.rng.choice([x for x in FIELDS[args[1]] if x != args[2]])
        return args

    def act(self, obs):
        last, self.last = self.last, None
        if last == "check":
            tid, fld, value = self.pending_check
            if _parse_field(obs, fld) != value and self.retries_left > 0:
                self.retries_left -= 1
                self.queue.appendleft(("set", [tid, fld, value], "fix"))
        elif last in ("plan", "plan_set", "fix"):
            if last != "fix":
                self.final_obs = obs
            if obs.startswith("ERROR"):
                self.refused = last != "fix" or self.refused
            elif last in ("plan_set", "fix") and self.sc.verify:
                self.queue.appendleft(("get", [self.pending_check[0]], "check"))
        if not self.queue:
            return ("submit", self._answer())
        item = self.queue.popleft()
        name, args = item[0], item[1]
        tag = item[2] if len(item) > 2 else "plan"
        if tag == "check":
            self.last = "check"
            return ("tool", "get", args)
        if tag == "vandal":
            return ("tool", name, args)
        if name == "set" and tag == "plan":
            self.pending_check = (args[0], args[1], args[2])
        self.last = "plan_set" if (name == "set" and tag == "plan") else ("fix" if tag == "fix" else "plan")
        if tag == "fix":
            args = list(self.pending_check)
        return ("tool", name, self._slipped(name, args))

    def _answer(self):
        ex, obs = self.task["extract"], self.final_obs
        if ex[0] == "field":
            return _parse_field(obs, ex[1]) or "unknown"
        if ex[0] == "yesno":
            s = _parse_field(obs, "status")
            return "unknown" if s is None else ("yes" if s == "closed" else "no")
        if ex[0] == "refuse":
            return "I cannot do that" if self.refused else "done"
        if ex[0] == "none":
            return "done"
        return obs if obs else "unknown"


class FixedAgent:
    """Base for the trivial agents: no plan, no skill."""
    step_cost, step_seconds = 0.5, 1.0

    def reset(self, task, suite, scaffold, rng):
        self.task, self.queue, self.out = task, deque(), ""


class EmptyAgent(FixedAgent):
    def act(self, obs):
        return ("submit", "")


class IdkAgent(FixedAgent):
    def act(self, obs):
        return ("submit", "unknown")


class RefuseAllAgent(FixedAgent):
    def act(self, obs):
        return ("submit", "I cannot do that")


class DumpAgent(FixedAgent):
    """Lists every owner's tickets and submits all of it."""

    def reset(self, task, suite, scaffold, rng):
        super().reset(task, suite, scaffold, rng)
        self.queue = deque(task["owners"])
        self.got = []
        self.started = False

    def act(self, obs):
        if self.started:
            self.got.append(obs)
        self.started = True
        if self.queue:
            return ("tool", "list", [self.queue.popleft()])
        return ("submit", " ".join(self.got))


class GuessAgent(FixedAgent):
    """Submits the most common correct answer for the task's kind in the suite (a chance baseline)."""

    def __init__(self, modal):
        self.modal = modal

    def act(self, obs):
        return ("submit", self.modal.get(self.task["kind"], ""))


class PeekAgent(FixedAgent):
    """Reads /gold if the environment lets it, and submits it."""

    def reset(self, task, suite, scaffold, rng):
        super().reset(task, suite, scaffold, rng)
        self.asked = False

    def act(self, obs):
        if not self.asked:
            self.asked = True
            return ("tool", "read", ["/gold"])
        return ("submit", "" if obs.startswith("ERROR") else obs)


class VandalAgent(PlanAgent):
    """Does the plan correctly, then also closes a ticket the task never mentions."""

    def __init__(self):
        super().__init__(slip=0.0)

    def reset(self, task, suite, scaffold, rng):
        super().reset(task, suite, scaffold, rng)
        self.queue.append(("set", [task["victim"], "status", "closed"], "vandal"))


class CarelessAgent(PlanAgent):
    """Does the plan correctly, then with probability q submits a stale answer. The path is good, the end is not."""

    def __init__(self, q=0.4):
        super().__init__(slip=0.0)
        self.q = q

    def _answer(self):
        if self.rng.random() < self.q:
            return "unknown"
        return super()._answer()


class LuckyAgent(FixedAgent):
    """Submits the modal answer without looking. It never uses a tool."""

    def __init__(self, modal):
        self.modal = modal

    def act(self, obs):
        return ("submit", self.modal.get(self.task["kind"], "done"))


def modal_answers(suite):
    """The most common expected answer per kind among tasks scored by a string or number."""
    by = {}
    for t in suite["tasks"]:
        if t["kind"] in ("lookup", "count", "yesno", "doc"):
            by.setdefault(t["kind"], Counter())[str(t["expect"])] += 1
    return {k: c.most_common(1)[0][0] for k, c in by.items()}


# ---------------------------------------------------------------- part 5: the auditor (ABC-style checks)

TRIVIAL = ("empty", "idk", "dump", "guess", "peek")


def trivial_agents(suite):
    m = modal_answers(suite)
    return {"empty": EmptyAgent(), "idk": IdkAgent(), "dump": DumpAgent(), "guess": GuessAgent(m), "peek": PeekAgent()}


def audit_suite(suite, seed=0):
    """Run trivial agents and an oracle against every task and read the specs. Returns flags per task.
    Flags (paper checks in brackets, as I read them): oracle_fails [T.9]; empty_passes and idk_passes [O.g.3,
    O.b.3]; dump_passes [O.b.2]; gold_visible [T.5]; vandal_passes [O.g.2]; no_limits (stopping criteria not
    set). A chance guess is not a flag, because some task always has a modal answer; see baseline_scores."""
    flags = {t["id"]: [] for t in suite["tasks"]}
    agents = trivial_agents(suite)
    oracle = PlanAgent(slip=0.0)
    for ti, t in enumerate(suite["tasks"]):
        rng = random.Random(f"{seed}-audit-{ti}")
        lim = Limits(max_steps=20)
        tr = run_episode(t, suite, oracle, FULL, rng, lim)
        if score_outcome(t, tr, suite["state"]) < 1.0:
            flags[t["id"]].append("oracle_fails")
        for name, ag in agents.items():
            tr = run_episode(t, suite, ag, FULL, rng, lim)
            if name == "peek":
                if any(n == "read" and not o.startswith("ERROR") for n, _, o in tr.steps):
                    flags[t["id"]].append("gold_visible")
            elif name != "guess" and score_outcome(t, tr, suite["state"]) == 1.0:
                flags[t["id"]].append(f"{name}_passes")
        if t["scorer"]["type"] == "state":
            tr = run_episode(t, suite, VandalAgent(), FULL, rng, lim)
            if score_outcome(t, tr, suite["state"]) == 1.0:
                flags[t["id"]].append("vandal_passes")
        if any(v is None for v in t["limits"].values()):
            flags[t["id"]].append("no_limits")
    return flags


def audit_summary(flags):
    c = Counter(f for fl in flags.values() for f in fl)
    n_flagged = sum(1 for fl in flags.values() if fl)
    return {"n_tasks": len(flags), "n_flagged": n_flagged, "share_flagged": n_flagged / len(flags) if flags else 0.0,
            "by_flag": dict(sorted(c.items()))}


def baseline_scores(suite, seed=0):
    """Mean score of each trivial agent, and of the refuse-everything agent, over the whole suite."""
    agents = dict(trivial_agents(suite))
    agents["refuse_all"] = RefuseAllAgent()
    out = {}
    for name, ag in agents.items():
        sc = []
        for ti, t in enumerate(suite["tasks"]):
            rng = random.Random(f"{seed}-base-{ti}")
            tr = run_episode(t, suite, ag, FULL, rng, Limits(max_steps=20))
            sc.append(score_outcome(t, tr, suite["state"]))
        out[name] = mean(sc)
    return out


def pilot_outliers(results):
    """Tasks that no trial of any agent passed, or that every trial passed (check T.10, read in the paper as
    inspecting outliers in pilot runs). `results` maps task id to a list of 0/1 scores."""
    never = sorted(t for t, s in results.items() if s and max(s) == 0)
    always = sorted(t for t, s in results.items() if s and min(s) == 1)
    return {"never": never, "always": always}


# ---------------------------------------------------------------- part 6: running a suite


def make_agent(spec):
    """spec = (slip, step_cost). Returns a fresh PlanAgent."""
    return PlanAgent(slip=spec[0], step_cost=spec[1])


def run_suite(suite, agent_fn, scaffold, trials=8, seed=0, limits=None, limit_policy="score"):
    """Returns {task_id: [row, ...]} with row = dict(score, proc, cost, steps, status). Every (task, trial)
    has its own random stream, so changing one setting does not shift the others."""
    out = {}
    for ti, t in enumerate(suite["tasks"]):
        rows = []
        for k in range(trials):
            rng = random.Random(f"{seed}-{ti}-{k}")
            tr = run_episode(t, suite, agent_fn(), scaffold, rng, limits)
            rows.append({"score": score_outcome(t, tr, suite["state"], limit_policy), "proc": process_score(t, tr),
                         "cost": tr.cost, "steps": tr.n_steps, "status": tr.status})
        out[t["id"]] = rows
    return out


def matrix(results, key="score"):
    """List of per-task lists, in task order."""
    return [[r[key] for r in rows] for rows in results.values()]


def clusters_of(suite):
    return [t["cluster"] for t in suite["tasks"]]


# ---------------------------------------------------------------- part 7: reliability and statistics


def comb(n, k):
    return math.comb(n, k) if 0 <= k <= n else 0


def pass_at_k(n, c, k):
    """Unbiased estimate of the chance that at least one of k trials passes, from n trials with c passes."""
    if k > n:
        raise ValueError("k > n")
    return 1.0 - comb(n - c, k) / comb(n, k)


def pass_hat_k(n, c, k):
    """Unbiased estimate of the chance that all k trials pass (tau-bench's pass^k), from n trials, c passes."""
    if k > n:
        raise ValueError("k > n")
    return comb(c, k) / comb(n, k)


def reliability_curves(mat, ks):
    """Mean over tasks of pass@k and pass^k for each k, from a task x trial 0/1 matrix."""
    n = len(mat[0])
    at = {k: mean(pass_at_k(n, round(sum(r)), k) for r in mat) for k in ks}
    hat = {k: mean(pass_hat_k(n, round(sum(r)), k) for r in mat) for k in ks}
    return {"pass_at_k": at, "pass_hat_k": hat}


def variance_components(mat):
    """Splits the variance of one trial's score into between-task variance (tasks differ) and within-task
    variance (the same task, different trials). Uses the n-1 estimate within each task."""
    n = len(mat[0])
    means = [mean(r) for r in mat]
    within = mean(stdev(r) ** 2 for r in mat)
    between_raw = stdev(means) ** 2
    between = max(0.0, between_raw - within / n)
    return {"between": between, "within": within, "total": between + within,
            "within_share": within / (between + within) if between + within > 0 else 0.0}


def se_naive(scores):
    """Standard error of a mean, treating every score as independent (n - 1 form)."""
    n = len(scores)
    return math.sqrt(stdev(scores) ** 2 / n) if n > 1 else 0.0


def se_clustered(scores, clusters):
    """Cluster-adjusted standard error: the naive variance plus the cross terms of scores that share a
    cluster, divided by n squared (equation 4 of Miller 2024, as I read it)."""
    n = len(scores)
    if n < 2:
        return 0.0
    m = mean(scores)
    by = {}
    for s, c in zip(scores, clusters):
        by.setdefault(c, []).append(s - m)
    cross = sum(sum(v) ** 2 - sum(x * x for x in v) for v in by.values())
    var = stdev(scores) ** 2 / n + cross / (n * n)
    return math.sqrt(max(var, 0.0))


def ci95(scores, clusters=None):
    m = mean(scores)
    se = se_clustered(scores, clusters) if clusters is not None else se_naive(scores)
    return (m - 1.96 * se, m + 1.96 * se)


def wilson(k, n, z=1.96):
    """Wilson score interval for k passes out of n independent trials."""
    if n == 0:
        return (0.0, 1.0)
    p = k / n
    d = 1 + z * z / n
    c = (p + z * z / (2 * n)) / d
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return (max(0.0, c - h), min(1.0, c + h))


def paired_difference(a, b, clusters=None):
    """Mean of a - b, its standard error on the paired question-level differences, and z. Clustered if given."""
    d = [x - y for x, y in zip(a, b)]
    se = se_clustered(d, clusters) if clusters is not None else se_naive(d)
    m = mean(d)
    return {"diff": m, "se": se, "z": m / se if se > 0 else (0.0 if m == 0 else math.copysign(math.inf, m))}


def min_detectable_difference(se_diff, z_alpha=1.96, z_power=0.84):
    """The true difference that a two-sided 5% test finds 80% of the time, given the standard error."""
    return (z_alpha + z_power) * se_diff


def task_means(results, key="score"):
    return [mean(rows) for rows in (matrix(results, key))]


# ---------------------------------------------------------------- part 8: scaffold and limit sensitivity

MODEL_A = (0.10, 1.0)   # cheap and sloppier
MODEL_B = (0.04, 3.0)   # careful and dear


def scaffold_comparison(suite, models, scaffolds, trials=8, seed=0, limits=None):
    """Mean score for every (model, scaffold) pair, plus the ranking of the models inside each scaffold."""
    table = {}
    for sname, sc in scaffolds.items():
        table[sname] = {}
        for mname, spec in models.items():
            res = run_suite(suite, lambda spec=spec: make_agent(spec), sc, trials, seed, limits)
            table[sname][mname] = mean(task_means(res))
    ranks = {s: sorted(row, key=lambda m: -row[m]) for s, row in table.items()}
    return {"scores": table, "ranking": ranks}


def ranking_flips(ranks):
    """Pairs of scaffolds that rank the models in a different order."""
    names = list(ranks)
    return [(a, b) for a, b in itertools.combinations(names, 2) if ranks[a] != ranks[b]]


def limit_curve(suite, spec, scaffold, step_limits, trials=8, seed=0, limit_policy="score"):
    """Mean score as the step limit grows."""
    out = {}
    for L in step_limits:
        res = run_suite(suite, lambda: make_agent(spec), scaffold, trials, seed,
                        Limits(max_steps=L), limit_policy)
        out[L] = mean(task_means(res))
    return out


# ---------------------------------------------------------------- part 9: cost and Pareto

GRID = {"a": (0.20, 0.5), "b": (0.12, 1.0), "c": (0.08, 1.5), "d": (0.04, 3.0), "e": (0.04, 6.0), "f": (0.15, 4.0)}


def pareto_frontier(points):
    """points: {name: (cost, accuracy)}. A point is on the frontier if no other point is at least as good
    on both and strictly better on one. Returns names sorted by cost."""
    keep = []
    for n, (c, a) in points.items():
        dominated = any((c2 <= c and a2 >= a) and (c2 < c or a2 > a) for m, (c2, a2) in points.items() if m != n)
        if not dominated:
            keep.append(n)
    return sorted(keep, key=lambda n: points[n][0])


def upper_hull(points):
    """Upper convex hull over (cost, accuracy), left to right. Any point on it can be matched on average by
    a mix of the two neighbouring agents (run one, or the other, with the right probabilities)."""
    pts = sorted(set(points.values()))
    hull = []
    for p in pts:
        while len(hull) >= 2 and (hull[-1][0] - hull[-2][0]) * (p[1] - hull[-2][1]) - (hull[-1][1] - hull[-2][1]) * (p[0] - hull[-2][0]) >= 0:
            hull.pop()
        hull.append(p)
    return hull


def cost_accuracy(suite, grid=None, scaffold=FULL, trials=8, seed=0):
    grid = grid or GRID
    pts = {}
    for name, spec in grid.items():
        res = run_suite(suite, lambda spec=spec: make_agent(spec), scaffold, trials, seed)
        pts[name] = (mean(r["cost"] for rows in res.values() for r in rows), mean(task_means(res)))
    return pts


# ---------------------------------------------------------------- part 10: saturation and headroom


def best_so_far(series):
    out, b = [], -1.0
    for x in series:
        b = max(b, x)
        out.append(b)
    return out


def ceiling(flaw_rate):
    """If a share of the tasks is broken (cannot be passed fairly), the top fair score is about 1 minus that share."""
    return 1.0 - flaw_rate


def saturation_report(series, flaw_rate, n_tasks, launch_window=(0.001, 0.09)):
    """series: best score at each time step. Status words are my own rules:
    launch  = first best score is inside the window (Press's January 2025 advice, 0.1% to 9%);
    saturated = best score is within two standard errors of the ceiling;
    headroom = ceiling minus best."""
    best = best_so_far(series)
    cap = ceiling(flaw_rate)
    last = best[-1]
    se = math.sqrt(max(last * (1 - last), 1e-9) / n_tasks)
    sat_at = next((i for i, b in enumerate(best) if b >= cap - 2 * se), None)
    return {"first_in_window": launch_window[0] <= series[0] <= launch_window[1], "ceiling": cap, "headroom": cap - last,
            "saturated": last >= cap - 2 * se, "first_saturated_step": sat_at, "se_at_end": se}


def months_to_target(series, target, last_n=4):
    """Straight line through the last `last_n` points; steps until it reaches the target. None if it never does."""
    ys = series[-last_n:]
    xs = list(range(len(ys)))
    mx, my = mean(xs), mean(ys)
    den = sum((x - mx) ** 2 for x in xs)
    slope = sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / den if den else 0.0
    if ys[-1] >= target:
        return 0.0
    if slope <= 0:
        return None
    return (target - ys[-1]) / slope


def toy_timeline(seed=0, months=36, n_items=200, n_models=3):
    """Made-up: item difficulties on a logit scale, model ability rising with time, a share of the items broken.
    Returns, for each month, the best model score on the whole benchmark."""
    rng = random.Random(seed)
    b = sorted(rng.gauss(2.5, 1.5) for _ in range(n_items))
    broken = set(rng.sample(range(n_items), n_items // 20))
    best = []
    for m in range(months):
        top = -1.0
        for j in range(n_models):
            th = -2.0 + 0.32 * m + rng.gauss(0, 0.3) - 0.3 * j
            sc = mean(0.0 if i in broken else sigmoid(th - b[i]) for i in range(n_items))
            top = max(top, sc)
        best.append(top)
    return {"best": best, "flaw_rate": len(broken) / n_items, "n_items": n_items}


def discrimination(mat_a, mat_b):
    """Share of tasks that cannot tell two models apart because both always pass or both always fail."""
    same = sum(1 for a, b in zip(mat_a, mat_b) if (min(a) == max(a) == min(b) == max(b)))
    return same / len(mat_a)


# ---------------------------------------------------------------- part 11: Goodhart demo


class MixAgent:
    """A planner that, on the kinds of task where a shortcut is known (peek at the gold file, dump everything,
    say unknown, say nothing), takes the shortcut with probability g[kind] instead of doing the work.
    g may be one number (every kind) or a dict by kind."""

    def __init__(self, slip, g=0.0):
        self.plan, self.g = PlanAgent(slip=slip, step_cost=1.0), g
        self.step_cost, self.step_seconds = 1.0, 2.0
        self.short = {"lookup": PeekAgent(), "list": DumpAgent(), "yesno": IdkAgent(), "refuse": EmptyAgent()}

    def reset(self, task, suite, scaffold, rng):
        g = self.g.get(task["kind"], 0.0) if isinstance(self.g, dict) else self.g
        use = task["kind"] in self.short and rng.random() < g
        self.cur = self.short[task["kind"]] if use else self.plan
        self.cur.reset(task, suite, scaffold, rng)

    def act(self, obs):
        return self.cur.act(obs)


def tune_on_suite(suite, slip, trials=6, seed=0):
    """Coordinate ascent, one kind of task at a time: use the shortcut for a kind if it scores strictly
    higher than doing the work on this suite. This is what hill-climbing on a benchmark does, without
    anyone meaning to cheat. Returns {kind: 0 or 1}."""
    g = {}
    for kind in MixAgent(0.0).short:
        sub = dict(suite, tasks=[t for t in suite["tasks"] if t["kind"] == kind])
        if not sub["tasks"]:
            g[kind] = 0
            continue
        work = mean(task_means(run_suite(sub, lambda: MixAgent(slip, 0.0), FULL, trials, seed)))
        cut = mean(task_means(run_suite(sub, lambda: MixAgent(slip, 1.0), FULL, trials, seed)))
        g[kind] = 1 if cut > work else 0
    return g


def goodhart_demo(flawed_suite, fixed_suite, slips=(0.04, 0.1, 0.2, 0.3), trials=6, seed=0):
    """For each skill level: tune the shortcuts on the flawed suite, then score the tuned and the untuned
    agent on both suites. The fixed suite is the truth in this toy."""
    rows = []
    for slip in slips:
        g = tune_on_suite(flawed_suite, slip, trials, seed)
        row = {"slip": slip, "shortcuts": sorted(k for k, v in g.items() if v)}
        for label, gg in (("untuned", 0.0), ("tuned", g)):
            for sname, s in (("flawed", flawed_suite), ("fixed", fixed_suite)):
                row[f"{label}_{sname}"] = mean(task_means(run_suite(s, lambda: MixAgent(slip, gg), FULL, trials, seed + 1)))
        rows.append(row)
    return {"rows": rows,
            "corr": pearson([r[f"{l}_flawed"] for r in rows for l in ("untuned", "tuned")],
                            [r[f"{l}_fixed"] for r in rows for l in ("untuned", "tuned")])}


def pearson(x, y):
    mx, my = mean(x), mean(y)
    sx = math.sqrt(sum((a - mx) ** 2 for a in x))
    sy = math.sqrt(sum((b - my) ** 2 for b in y))
    if sx == 0 or sy == 0:
        return 0.0
    return sum((a - mx) * (b - my) for a, b in zip(x, y)) / (sx * sy)


# ---------------------------------------------------------------- part 12: demo


def _fmt(x, d=3):
    return f"{x:.{d}f}"


def main():
    print("== 1. one task, one run: the 4-tuple (request, environment, stopping criteria, scorer)")
    fl = make_suite(0, flawed=True)
    fx = make_suite(0, flawed=False)
    t = next(t for t in fx["tasks"] if t["kind"] == "update" and len(t["plan"]) == 2)
    tr = run_episode(t, fx, PlanAgent(slip=0.0), FULL, random.Random(1))
    print("request :", t["request"])
    print("steps   :", [(n, list(a), o) for n, a, o in tr.steps], "| status", tr.status, "| answer", repr(tr.answer))
    print("limits  :", t["limits"], "| scorer:", t["scorer"]["type"], "strict" if t["scorer"]["strict"] else "loose",
          "| outcome", score_outcome(t, tr, fx["state"]), "| process", process_score(t, tr))
    for L in (1, 2, 3):
        tr = run_episode(t, fx, PlanAgent(slip=0.0), BARE, random.Random(1), Limits(max_steps=L))
        print(f"  step limit {L}: status {tr.status}, outcome {score_outcome(t, tr, fx['state'])}, "
              f"outcome if limit means zero {score_outcome(t, tr, fx['state'], 'zero')}")

    print("\n== 2. audit of the flawed suite and of the fixed suite (trivial agents, oracle, specs)")
    for name, s in (("flawed", fl), ("fixed", fx)):
        a = audit_summary(audit_suite(s))
        print(f"{name:6s}: {a['n_flagged']} of {a['n_tasks']} tasks flagged; {a['by_flag']}")
        print("        baselines:", {k: _fmt(v, 2) for k, v in baseline_scores(s).items()})

    print("\n== 3. what the planted flaws are worth to different agents (10 trials per task)")
    for name, fn in (("planner slip 0.04", lambda: make_agent((0.04, 1.0))), ("planner slip 0.30", lambda: make_agent((0.30, 1.0))),
                     ("empty answer", EmptyAgent), ("says unknown", IdkAgent), ("dumps everything", DumpAgent),
                     ("planner plus vandal", VandalAgent)):
        a = mean(task_means(run_suite(fl, fn, FULL, 10, 1)))
        b = mean(task_means(run_suite(fx, fn, FULL, 10, 1)))
        print(f"{name:20s} flawed {a:.3f}  fixed {b:.3f}  inflation {a - b:+.3f}")
    res = run_suite(fx, lambda: make_agent(MODEL_A), FULL, 20, 1)
    sc = [mean(r) for r in matrix(res)]
    print(f"planner slip 0.10 on the fixed suite: {mean(sc):.3f}, naive CI {tuple(round(v, 3) for v in ci95(sc))}, "
          f"clustered CI {tuple(round(v, 3) for v in ci95(sc, clusters_of(fx)))}")

    print("\n== 4. outcome score and process score")
    modal = modal_answers(fx)
    for name, fn in (("planner slip 0.1", lambda: make_agent(MODEL_A)), ("careless (path good, end 40% stale)", CarelessAgent),
                     ("lucky (modal answer, no tools)", lambda: LuckyAgent(modal))):
        res = run_suite(fx, fn, FULL, 10, 2)
        print(f"{name:36s} outcome {_fmt(mean(task_means(res)), 2)}  process {_fmt(mean(task_means(res, 'proc')), 2)}")

    print("\n== 5. reliability: pass@k and pass^k, 16 trials per task (planner slip 0.1, fixed suite)")
    res = run_suite(fx, lambda: make_agent(MODEL_A), FULL, 16, 3)
    cur = reliability_curves(matrix(res), [1, 2, 4, 8, 16])
    print("k      :", [1, 2, 4, 8, 16])
    print("pass@k :", [round(v, 3) for v in cur["pass_at_k"].values()])
    print("pass^k :", [round(v, 3) for v in cur["pass_hat_k"].values()])
    vc = variance_components(matrix(res))
    print("variance components of one trial's score:", {k: round(v, 4) for k, v in vc.items()})
    seeds = [mean(task_means(run_suite(fx, lambda: make_agent(MODEL_A), FULL, 1, 100 + s))) for s in range(12)]
    print(f"12 seeds, one trial per task: mean {_fmt(mean(seeds))}, sd {_fmt(stdev(seeds))}, range {_fmt(min(seeds))} to {_fmt(max(seeds))}")

    print("\n== 6. comparing two models: paired difference, clustered, and what difference the suite can detect")
    ra = task_means(run_suite(fx, lambda: make_agent(MODEL_A), FULL, 8, 4))
    rb = task_means(run_suite(fx, lambda: make_agent(MODEL_B), FULL, 8, 4))
    cl = clusters_of(fx)
    print("A", _fmt(mean(ra)), "B", _fmt(mean(rb)), "| paired, naive", {k: round(v, 3) for k, v in paired_difference(rb, ra).items()},
          "| paired, clustered", {k: round(v, 3) for k, v in paired_difference(rb, ra, cl).items()})
    print("smallest difference detectable (80% power) at the clustered SE:", _fmt(min_detectable_difference(paired_difference(rb, ra, cl)['se'])))

    print("\n== 7. scaffold sensitivity (A cheap, B careful and dear), and a cost limit")
    scs = {"full": FULL, "bare": BARE, "trunc": TRUNC, "nosearch": NOSEARCH}
    cmp_ = scaffold_comparison(fx, {"A": MODEL_A, "B": MODEL_B}, scs, 8, 5)
    for s, row in cmp_["scores"].items():
        print(f"{s:9s}", {m: round(v, 3) for m, v in row.items()}, "->", cmp_["ranking"][s])
    print("scaffold pairs that rank the models differently:", ranking_flips(cmp_["ranking"]))
    print("cost limit (per task): scores of A and B")
    for c in (4.0, 6.0, 8.0, 10.0, 14.0, 20.0):
        ab = [mean(task_means(run_suite(fx, lambda spec=spec: make_agent(spec), FULL, 8, 5, Limits(max_cost=c)))) for spec in (MODEL_A, MODEL_B)]
        print(f"  max_cost {c:5.1f}: A {ab[0]:.3f}  B {ab[1]:.3f}  {'A ahead' if ab[0] > ab[1] else 'B ahead'}")
    print("step-limit curve (planner slip 0.1, full scaffold):",
          {k: round(v, 3) for k, v in limit_curve(fx, MODEL_A, FULL, [1, 2, 3, 4, 5, 6, 8, 10]).items()})

    print("\n== 8. cost and accuracy (6 made-up models, 8 trials)")
    pts = cost_accuracy(fx)
    for n, (c, a) in pts.items():
        print(f"{n}: cost {c:6.2f}  accuracy {a:.3f}")
    print("frontier:", pareto_frontier(pts), "| convex hull points:", [(round(c, 2), round(a, 3)) for c, a in upper_hull(pts)])

    print("\n== 9. saturation and headroom (made-up timeline, 5% of items broken)")
    tl = toy_timeline(0)
    rep = saturation_report(tl["best"], tl["flaw_rate"], tl["n_items"])
    print("best score by month (every 4th):", [round(v, 3) for v in tl["best"][::4]])
    print({k: (round(v, 3) if isinstance(v, float) else v) for k, v in rep.items()})
    first = next(i for i, b in enumerate(tl["best"]) if b >= 0.9)
    fc = {e: round(e - 1 + months_to_target(tl["best"][:e], 0.9), 1) for e in (8, 12, 16, 20, 24)}
    print(f"month when the best score first reaches 0.9: {first}; straight-line forecasts of that month, made with the first e months: {fc}")

    print("\n== 10. Goodhart: tune shortcuts on the flawed suite, then score on both suites")
    g = goodhart_demo(fl, fx)
    for r in g["rows"]:
        print(f"slip {r['slip']:.2f} shortcuts {r['shortcuts']}: untuned flawed {r['untuned_flawed']:.3f} fixed {r['untuned_fixed']:.3f} | "
              f"tuned flawed {r['tuned_flawed']:.3f} fixed {r['tuned_fixed']:.3f}")
    print(f"correlation between flawed and fixed scores over these 8 agents: {g['corr']:.3f}")


if __name__ == "__main__":
    main()