"""agent_data_toolkit.py - Week 6b lab: data for agents, in one readable file.

Standard library only (tested on Python 3.13). Nothing here calls a model, a network or a
subprocess. The "agent" is a mock with a made-up skill table, the "users" are a simulator, and
every task, rate, probability and price is MADE UP. The code shows the *shape* of the data
work (traces, demonstrations, preference pairs, eval cases, splits, flywheels, synthetic tasks,
implicit feedback, privacy, data cards). It is not a method from any paper, and no number it
prints is a result about a real model.

Parts
  1  Trace schema, logger, validation, JSON Lines, a span tree
  2  A mock world: six task types, a mock agent, a user simulator
  3  Trace conversion: SFT demonstrations, preference pairs, eval cases
  4  Splits and leakage checks: shingles, MinHash and LSH, group split, time split, report
  5  Flywheel simulation (deploy, log, label, improve) and a model-collapse toy
  6  Synthetic task generation with filters: valid, verified, dedupe, difficulty
  7  Implicit feedback signals, weak labels, a Bradley-Terry fit for pairwise votes
  8  PII scrubbing, consent flags, retention and deletion
  9  A data card generator
"""
from __future__ import annotations

import difflib
import hashlib
import json
import math
import random
import re
import sys
from dataclasses import dataclass, field, replace


# --------------------------------------------------------------------------- 0. shared bits
def stable_hash(*parts, bits: int = 64) -> int:
    """A hash that does not change between runs (Python's built-in hash() does)."""
    h = hashlib.sha256("\x1f".join(str(p) for p in parts).encode()).hexdigest()
    return int(h[: bits // 4], 16)


def rng_for(*parts) -> random.Random:
    return random.Random(stable_hash(*parts))


_WORD = re.compile(r"[a-z0-9']+")


def words(text) -> list[str]:
    return _WORD.findall(str(text).lower())


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


def entropy(counts) -> float:
    """Shannon entropy in bits of a list or dict of counts."""
    vals = [c for c in (counts.values() if isinstance(counts, dict) else counts) if c > 0]
    total = sum(vals)
    return -sum(c / total * math.log2(c / total) for c in vals) if total else 0.0


def count_by(items, key) -> dict:
    out: dict = {}
    for it in items:
        out[it[key]] = out.get(it[key], 0) + 1
    return out


# --------------------------------------------------------------------------- 1. traces
EVENT_KINDS = ("system", "user", "assistant", "tool_call", "tool_result", "feedback")
FEEDBACK_TYPES = ("thumbs_up", "thumbs_down", "regenerate", "edit", "copy", "abandon")
DEFAULT_CONSENT = {"log": True, "train": False, "share": False}


class TraceError(ValueError):
    pass


@dataclass
class Event:
    kind: str
    text: str
    ts: float
    meta: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {"kind": self.kind, "text": self.text, "ts": self.ts, "meta": dict(self.meta)}

    @staticmethod
    def from_dict(d: dict) -> "Event":
        return Event(d["kind"], d["text"], d["ts"], dict(d.get("meta", {})))


@dataclass
class Trace:
    trace_id: str
    session_id: str
    user_id: str
    task_type: str
    harness: str
    model: str
    started: float
    events: list = field(default_factory=list)
    reward: float | None = None
    verifier: dict | None = None
    consent: dict = field(default_factory=lambda: dict(DEFAULT_CONSENT))
    scrubbed: bool = False

    def prompt(self) -> str:
        for e in self.events:
            if e.kind == "user":
                return e.text
        return ""

    def final_text(self) -> str:
        for e in reversed(self.events):
            if e.kind == "assistant":
                return e.text
        return ""

    def feedback(self) -> list:
        return [e for e in self.events if e.kind == "feedback"]

    def to_dict(self) -> dict:
        return {"trace_id": self.trace_id, "session_id": self.session_id, "user_id": self.user_id,
                "task_type": self.task_type, "harness": self.harness, "model": self.model,
                "started": self.started, "events": [e.to_dict() for e in self.events],
                "reward": self.reward, "verifier": self.verifier, "consent": dict(self.consent),
                "scrubbed": self.scrubbed}

    @staticmethod
    def from_dict(d: dict) -> "Trace":
        return Trace(d["trace_id"], d["session_id"], d["user_id"], d["task_type"], d["harness"],
                     d["model"], d["started"], [Event.from_dict(e) for e in d["events"]],
                     d.get("reward"), d.get("verifier"), dict(d.get("consent", DEFAULT_CONSENT)),
                     bool(d.get("scrubbed", False)))


class TraceLogger:
    """Collects traces in memory. start() opens one, log() appends, end() closes it."""

    def __init__(self, prefix: str = "t"):
        self.prefix = prefix
        self._traces: dict[str, Trace] = {}
        self._closed: set[str] = set()
        self._n = 0

    def start(self, session_id, user_id, task_type, harness, model, ts, consent=None) -> str:
        self._n += 1
        tid = f"{self.prefix}{self._n:05d}"
        c = dict(DEFAULT_CONSENT)
        c.update(consent or {})
        self._traces[tid] = Trace(tid, session_id, user_id, task_type, harness, model, float(ts), consent=c)
        return tid

    def log(self, trace_id, kind, text, ts, **meta) -> None:
        if trace_id not in self._traces:
            raise TraceError(f"unknown trace {trace_id}")
        if trace_id in self._closed:
            raise TraceError(f"trace {trace_id} is closed")
        if kind not in EVENT_KINDS:
            raise TraceError(f"unknown event kind {kind!r}")
        tr = self._traces[trace_id]
        last = tr.events[-1].ts if tr.events else tr.started
        if ts < last:
            raise TraceError("timestamps must not go backwards")
        tr.events.append(Event(kind, str(text), float(ts), dict(meta)))

    def end(self, trace_id, reward=None, verifier=None) -> Trace:
        if trace_id not in self._traces:
            raise TraceError(f"unknown trace {trace_id}")
        if trace_id in self._closed:
            raise TraceError(f"trace {trace_id} is already closed")
        if reward is not None and not 0.0 <= reward <= 1.0:
            raise TraceError("reward must be between 0 and 1")
        tr = self._traces[trace_id]
        tr.reward, tr.verifier = reward, verifier
        self._closed.add(trace_id)
        return tr

    def all(self) -> list:
        return list(self._traces.values())

    def open_ids(self) -> list:
        return [t for t in self._traces if t not in self._closed]


def validate_trace(tr: Trace) -> list[str]:
    """Return a list of problems. An empty list means the trace is well formed."""
    problems = []
    if not tr.events:
        return ["no events"]
    pending = 0
    last_ts = tr.started
    seen_first = False
    for i, e in enumerate(tr.events):
        if e.kind not in EVENT_KINDS:
            problems.append(f"event {i}: unknown kind {e.kind!r}")
            continue
        if e.ts < last_ts:
            problems.append(f"event {i}: timestamps go backwards")
        last_ts = max(last_ts, e.ts)
        if e.kind != "system" and not seen_first:
            seen_first = True
            if e.kind != "user":
                problems.append("first message is not from the user")
        if e.kind == "tool_call":
            pending += 1
        elif e.kind == "tool_result":
            if pending == 0:
                problems.append(f"event {i}: tool_result without tool_call")
            else:
                pending -= 1
        elif e.kind == "feedback" and e.meta.get("type") not in FEEDBACK_TYPES:
            problems.append(f"event {i}: bad feedback type {e.meta.get('type')!r}")
    if pending:
        problems.append("ends with an unanswered tool_call")
    if tr.reward is not None and not 0.0 <= tr.reward <= 1.0:
        problems.append("reward outside 0..1")
    return problems


def dumps_jsonl(traces) -> str:
    return "".join(json.dumps(t.to_dict(), sort_keys=True) + "\n" for t in traces)


def loads_jsonl(text: str) -> list:
    return [Trace.from_dict(json.loads(line)) for line in text.splitlines() if line.strip()]


def to_span_tree(tr: Trace) -> dict:
    """A loose, OpenTelemetry-flavoured view: one root span, one child per model or tool step.
    The names are this lab's own, not the official semantic-convention names."""
    kids = []
    for i, e in enumerate(tr.events):
        if e.kind not in ("assistant", "tool_call"):
            continue
        nxt = tr.events[i + 1].ts if i + 1 < len(tr.events) else e.ts
        kids.append({"name": "llm_call" if e.kind == "assistant" else "tool_call",
                     "start": e.ts, "end": nxt, "attributes": {"chars": len(e.text)}})
    end = tr.events[-1].ts if tr.events else tr.started
    return {"name": "session", "trace_id": tr.trace_id, "start": tr.started, "end": end,
            "attributes": {"harness": tr.harness, "model": tr.model, "reward": tr.reward},
            "children": kids}


def span_summary(root: dict) -> dict:
    out = {"llm_call": 0, "tool_call": 0, "seconds": round(root["end"] - root["start"], 3)}
    for k in root["children"]:
        out[k["name"]] += 1
    return out


# --------------------------------------------------------------------------- 2. a mock world
TYPES = ("sum", "reverse", "upper", "vowels", "sort", "domain")
HARNESSES = ("react", "json")
HARNESS_TRANSFER = 0.6      # MADE UP: skill kept when the harness differs from every training harness
MEMORY_SKILL = 0.98         # MADE UP: success rate on a prompt the mock agent has been trained on
INIT_SKILL = {"sum": 0.80, "reverse": 0.50, "upper": 0.90, "vowels": 0.35, "sort": 0.55, "domain": 0.70}
_NAMES = ["ada", "grace", "linus", "alan", "edsger", "barbara", "ken", "margaret"]
_DOMAINS = ["example.org", "mail.test", "corp.example", "uni.example"]
_WORDLIST = ["banana", "queue", "rhythm", "education", "sky", "avocado", "strength", "aeiou",
             "kitchen", "orchestra", "pencil", "umbrella"]
_COLORS = ["red", "green", "blue", "amber", "violet", "teal", "pink", "grey"]


@dataclass(frozen=True)
class WorldTask:
    task_id: str
    type: str
    prompt: str
    answer: str


def make_prompt(kind: str, rng: random.Random) -> str:
    if kind == "sum":
        return "Add the numbers: " + ", ".join(str(rng.randint(1, 20)) for _ in range(rng.randint(2, 4)))
    if kind == "reverse":
        return "Reverse the order of these words: " + " ".join(rng.sample(_COLORS, rng.randint(3, 4)))
    if kind == "upper":
        return "Write this text in uppercase: " + " ".join(rng.sample(_COLORS, 2))
    if kind == "vowels":
        return "Count the vowels in the word: " + rng.choice(_WORDLIST)
    if kind == "sort":
        return "Sort these numbers from smallest to largest: " + ", ".join(str(rng.randint(1, 30)) for _ in range(rng.randint(3, 5)))
    if kind == "domain":
        return f"What is the domain part of this email address: {rng.choice(_NAMES)}@{rng.choice(_DOMAINS)}"
    raise ValueError(f"unknown task type {kind!r}")


def solve(kind: str, prompt: str) -> str:
    """The reference solver: it reads the payload after the first ': ' and computes the answer."""
    if ": " not in prompt:
        raise ValueError("no payload")
    payload = prompt.split(": ", 1)[1].strip()
    if kind == "sum":
        return str(sum(int(x) for x in payload.split(",")))
    if kind == "reverse":
        return " ".join(reversed(payload.split()))
    if kind == "upper":
        return payload.upper()
    if kind == "vowels":
        return str(sum(ch in "aeiou" for ch in payload.lower()))
    if kind == "sort":
        return ", ".join(str(n) for n in sorted(int(x) for x in payload.split(",")))
    if kind == "domain":
        return payload.split("@", 1)[1]
    raise ValueError(f"unknown task type {kind!r}")


def make_task(kind: str, rng: random.Random, task_id: str) -> WorldTask:
    p = make_prompt(kind, rng)
    return WorldTask(task_id, kind, p, solve(kind, p))


def verifier_spec(task: WorldTask) -> dict:
    return {"type": "exact", "expected": task.answer}


def run_verifier(spec: dict, answer: str) -> float:
    if spec.get("type") != "exact":
        raise ValueError("unknown verifier type")
    return 1.0 if str(answer).strip() == str(spec["expected"]).strip() else 0.0


def corrupt(kind: str, answer: str, rng: random.Random) -> str:
    """A plausible wrong answer."""
    if answer.lstrip("-").isdigit():
        return str(int(answer) + rng.choice([-1, 1, 2]))
    if len(answer) > 1:
        return answer[:-1]
    return answer + "?"


def prompt_key(prompt: str) -> str:
    return hashlib.sha256(prompt.encode()).hexdigest()[:12]


class MockAgent:
    """A mock policy. skill[type] is the chance that one attempt is right. Not a model."""

    def __init__(self, skill=None, harness="react", model="mock-v0", trained_harnesses=None,
                 memorized=None, seed=0):
        self.skill = dict(skill or INIT_SKILL)
        self.harness = harness
        self.model = model
        self.trained_harnesses = set(trained_harnesses) if trained_harnesses else None
        self.memorized = set(memorized or ())
        self.seed = seed

    def effective_skill(self, task: WorldTask) -> float:
        p = self.skill[task.type]
        if self.trained_harnesses is not None and self.harness not in self.trained_harnesses:
            p *= HARNESS_TRANSFER
        if prompt_key(task.prompt) in self.memorized:
            p = max(p, MEMORY_SKILL)
        return p

    def attempts(self, task: WorldTask) -> list[str]:
        """One or two attempts. The second runs only after a wrong first one, with half the skill."""
        rng = rng_for("agent", self.seed, self.model, task.task_id)
        p = self.effective_skill(task)
        out = [task.answer if rng.random() < p else corrupt(task.type, task.answer, rng)]
        if out[0] != task.answer:
            out.append(task.answer if rng.random() < p * 0.5 else corrupt(task.type, task.answer, rng))
        return out

    def _action_text(self, task: WorldTask) -> str:
        if self.harness == "json":
            return json.dumps({"name": "solve", "arguments": {"input": task.prompt.split(": ", 1)[1]}})
        return f"solve[{task.prompt.split(': ', 1)[1]}]"

    def _thought_text(self, task: WorldTask, n: int) -> str:
        t = f"This is a {task.type} task." if n == 0 else "That looked wrong. I will try again."
        return json.dumps({"thought": t}) if self.harness == "json" else f"Thought: {t}"

    def run(self, task: WorldTask, logger: TraceLogger, user_id: str, session_id: str, ts: float,
            consent=None, user_sim=None) -> Trace:
        tid = logger.start(session_id, user_id, task.type, self.harness, self.model, ts, consent)
        t = float(ts)
        logger.log(tid, "system", f"You are a helpful assistant. Harness: {self.harness}.", t)
        t += 0.5
        logger.log(tid, "user", task.prompt, t, task_id=task.task_id)
        answer = ""
        for n, ans in enumerate(self.attempts(task)):
            t += 1.0
            logger.log(tid, "assistant", self._thought_text(task, n), t)
            t += 0.5
            logger.log(tid, "tool_call", self._action_text(task), t)
            t += 1.0
            logger.log(tid, "tool_result", ans, t)
            answer = ans
        t += 1.0
        logger.log(tid, "assistant", f"Answer: {answer}", t)
        reward = run_verifier(verifier_spec(task), answer)
        if user_sim is not None:
            user_sim.react(task, tid, logger, t, reward, answer)
        return logger.end(tid, reward, verifier_spec(task))


class UserSim:
    """Simulates how a user reacts. All probabilities are MADE UP."""
    OK = (("copy", 0.45), ("thumbs_up", 0.20), ("thumbs_down", 0.05))
    BAD = (("regenerate", 0.25), ("edit", 0.25), ("thumbs_down", 0.15), ("abandon", 0.10),
           ("rephrase", 0.10), ("copy", 0.05))

    def __init__(self, seed=0):
        self.seed = seed

    def react(self, task, trace_id, logger, ts, reward, answer) -> None:
        rng = rng_for("user", self.seed, task.task_id)
        r = rng.random()
        acc = 0.0
        pick = None
        for name, p in (self.OK if reward >= 1.0 else self.BAD):
            acc += p
            if r < acc:
                pick = name
                break
        t = ts + 2.0
        if pick is None:
            return
        if pick == "rephrase":
            logger.log(trace_id, "user", "again: " + task.prompt, t)
        elif pick == "edit":
            logger.log(trace_id, "feedback", f"Answer: {task.answer}", t, type="edit")
        else:
            logger.log(trace_id, "feedback", "", t, type=pick)


def simulate_traffic(n, seed=0, weights=None, agent=None, user_sim=None, start=0.0, n_users=12,
                     logger=None, session_prefix="s", gap=30.0) -> list:
    """Run n mock sessions. Tasks are drawn by weight, users by id, consent by user."""
    agent = agent or MockAgent(seed=seed)
    logger = logger or TraceLogger()
    rng = rng_for("traffic", seed)
    w = [weights.get(k, 0.0) for k in TYPES] if weights else [1.0] * len(TYPES)
    out = []
    for i in range(n):
        kind = rng.choices(TYPES, weights=w)[0]
        task = make_task(kind, rng, f"{session_prefix}{seed}-{i:04d}")
        uid = f"u{rng.randrange(n_users):02d}"
        out.append(agent.run(task, logger, uid, f"{session_prefix}{seed}-{i:04d}", start + i * gap,
                             consent=user_consent(uid), user_sim=user_sim))
    return out


def user_consent(user_id: str) -> dict:
    """MADE UP consent: about three users in four allow training, one in four allows sharing."""
    h = stable_hash("consent", user_id)
    return {"log": True, "train": h % 4 != 0, "share": h % 4 == 1}


# --------------------------------------------------------------------------- 3. conversion
def traces_to_sft(traces, min_reward: float = 1.0, require_consent: bool = True,
                  max_events: int | None = None) -> list:
    """Demonstrations. Assistant steps are learned (train=True). The prompt and tool output
    are context only. Failed, invalid, unconsented and too-long traces are dropped."""
    out = []
    for tr in traces:
        if tr.reward is None or tr.reward < min_reward or validate_trace(tr):
            continue
        if require_consent and not tr.consent.get("train"):
            continue
        evs = [e for e in tr.events if e.kind != "feedback"]
        if max_events is not None and len(evs) > max_events:
            continue
        msgs = []
        for e in evs:
            role = {"system": "system", "user": "user", "assistant": "assistant",
                    "tool_call": "assistant", "tool_result": "tool"}[e.kind]
            msgs.append({"role": role, "content": e.text, "train": e.kind in ("assistant", "tool_call")})
        out.append({"trace_id": tr.trace_id, "task_type": tr.task_type, "harness": tr.harness,
                    "reward": tr.reward, "started": tr.started, "user_id": tr.user_id,
                    "prompt": tr.prompt(), "messages": msgs})
    return out


def traces_to_preferences(traces, require_consent: bool = True) -> list:
    """Outcome pairs: for one prompt, the best trace is chosen and the worst is rejected.
    Needs a reward gap. Ties in reward go to the shorter trace."""
    groups: dict[str, list] = {}
    for tr in traces:
        if tr.reward is None or validate_trace(tr):
            continue
        if require_consent and not tr.consent.get("train"):
            continue
        groups.setdefault(tr.prompt(), []).append(tr)
    pairs = []
    for prompt in sorted(groups):
        g = groups[prompt]
        best = max(g, key=lambda t: (t.reward, -len(t.events), t.trace_id))
        worst = min(g, key=lambda t: (t.reward, -len(t.events), t.trace_id))
        if best.reward > worst.reward:
            pairs.append({"prompt": prompt, "chosen": best.final_text(), "rejected": worst.final_text(),
                          "source": "outcome", "chosen_trace": best.trace_id, "rejected_trace": worst.trace_id})
    return pairs


def edit_pairs(traces, require_consent: bool = True) -> list:
    """Edit pairs: the user's edit is chosen, the assistant's text it replaced is rejected."""
    pairs = []
    for tr in traces:
        if require_consent and not tr.consent.get("train"):
            continue
        for e in tr.feedback():
            if e.meta.get("type") == "edit" and e.text.strip() and e.text.strip() != tr.final_text().strip():
                pairs.append({"prompt": tr.prompt(), "chosen": e.text, "rejected": tr.final_text(),
                              "source": "edit", "chosen_trace": tr.trace_id, "rejected_trace": tr.trace_id})
                break
    return pairs


def traces_to_eval_cases(traces, only_failed: bool = True, max_per_type: int | None = None) -> list:
    """Eval cases: the prompt and its verifier, kept once per prompt. By default only failures."""
    seen, out, per_type = set(), [], {}
    for tr in traces:
        if tr.reward is None or tr.verifier is None or validate_trace(tr):
            continue
        if only_failed and tr.reward >= 1.0:
            continue
        p = tr.prompt()
        if p in seen:
            continue
        if max_per_type is not None and per_type.get(tr.task_type, 0) >= max_per_type:
            continue
        seen.add(p)
        per_type[tr.task_type] = per_type.get(tr.task_type, 0) + 1
        out.append({"case_id": "c-" + prompt_key(p), "task_type": tr.task_type, "prompt": p,
                    "verifier": dict(tr.verifier), "from_trace": tr.trace_id, "started": tr.started,
                    "user_id": tr.user_id})
    return out


def loss_tokens(example: dict) -> tuple[int, int]:
    """(learned words, context words) for one SFT example. A rough count, not tokens."""
    learned = sum(len(words(m["content"])) for m in example["messages"] if m["train"])
    context = sum(len(words(m["content"])) for m in example["messages"] if not m["train"])
    return learned, context


def harness_mix(examples) -> dict:
    out: dict[str, int] = {}
    for ex in examples:
        out[ex["harness"]] = out.get(ex["harness"], 0) + 1
    return dict(sorted(out.items()))


# --------------------------------------------------------------------------- 4. splits and leakage
def shingles(text: str, k: int = 3) -> set:
    w = words(text)
    if len(w) < k:
        return {" ".join(w)} if w else set()
    return {" ".join(w[i:i + k]) for i in range(len(w) - k + 1)}


def jaccard(a: set, b: set) -> float:
    if not a and not b:
        return 1.0
    return len(a & b) / len(a | b)


def minhash(sh: set, num_perm: int = 64, seed: int = 0) -> tuple:
    if not sh:
        return tuple([2 ** 64 - 1] * num_perm)
    return tuple(min(stable_hash("mh", seed, i, s) for s in sh) for i in range(num_perm))


def minhash_similarity(a: tuple, b: tuple) -> float:
    if len(a) != len(b):
        raise ValueError("signatures differ in length")
    return sum(x == y for x, y in zip(a, b)) / len(a)


def exact_key(text: str) -> str:
    """Hash of the text after lowercasing and collapsing spaces, for exact-duplicate checks."""
    return hashlib.sha256(" ".join(words(text)).encode()).hexdigest()[:16]


def lsh_candidates(texts: dict, k: int = 3, num_perm: int = 64, bands: int = 16, seed: int = 0) -> list:
    """Candidate pairs: two texts are candidates when their MinHash signatures agree on every
    row of at least one band. Cheap, and it can include pairs that are not really similar."""
    if num_perm % bands:
        raise ValueError("bands must divide num_perm")
    rows = num_perm // bands
    sig = {i: minhash(shingles(t, k), num_perm, seed) for i, t in texts.items()}
    buckets: dict[tuple, list] = {}
    for i in sorted(texts):
        for b in range(bands):
            buckets.setdefault((b, sig[i][b * rows:(b + 1) * rows]), []).append(i)
    cands = set()
    for ids in buckets.values():
        for x in range(len(ids)):
            for y in range(x + 1, len(ids)):
                cands.add((ids[x], ids[y]))
    return sorted(cands)


def near_duplicates(texts: dict, threshold: float = 0.8, k: int = 3, num_perm: int = 64,
                    bands: int = 16, seed: int = 0) -> list:
    """Find pairs whose shingle Jaccard is at least threshold. MinHash with banding picks the
    candidates, and the exact Jaccard decides. Returns sorted (id_a, id_b, jaccard) tuples."""
    sh = {i: shingles(t, k) for i, t in texts.items()}
    out = []
    for a, b in lsh_candidates(texts, k, num_perm, bands, seed):
        j = jaccard(sh[a], sh[b])
        if j >= threshold:
            out.append((a, b, round(j, 4)))
    return out


def brute_force_near_duplicates(texts: dict, threshold: float = 0.8, k: int = 3) -> list:
    ids = sorted(texts)
    sh = {i: shingles(texts[i], k) for i in ids}
    out = []
    for x in range(len(ids)):
        for y in range(x + 1, len(ids)):
            j = jaccard(sh[ids[x]], sh[ids[y]])
            if j >= threshold:
                out.append((ids[x], ids[y], round(j, 4)))
    return out


def split_by_group(items, group_key, fractions=(0.8, 0.1, 0.1), names=("train", "dev", "test"),
                   salt: str = "split") -> dict:
    """All items with the same group (a user, a task, a prompt) land in the same split."""
    if abs(sum(fractions) - 1.0) > 1e-9 or len(fractions) != len(names):
        raise ValueError("fractions must sum to 1 and match names")
    out = {n: [] for n in names}
    for it in items:
        u = (stable_hash(salt, group_key(it)) % 10 ** 6) / 10 ** 6
        acc = 0.0
        for n, f in zip(names, fractions):
            acc += f
            if u < acc or n == names[-1]:
                out[n].append(it)
                break
    return out


def split_random(items, fractions=(0.8, 0.1, 0.1), names=("train", "dev", "test"), seed: int = 0) -> dict:
    """A plain shuffle split. It ignores groups and time, so it can leak."""
    if abs(sum(fractions) - 1.0) > 1e-9 or len(fractions) != len(names):
        raise ValueError("fractions must sum to 1 and match names")
    idx = list(range(len(items)))
    random.Random(seed).shuffle(idx)
    out, start = {}, 0
    for n, f in zip(names, fractions):
        end = len(idx) if n == names[-1] else start + int(round(f * len(items)))
        out[n] = [items[i] for i in sorted(idx[start:end])]
        start = end
    return out


def time_split(items, ts_key, train_end: float, dev_end: float) -> dict:
    """Train before train_end, dev before dev_end, test from dev_end on."""
    if not train_end <= dev_end:
        raise ValueError("train_end must not exceed dev_end")
    out = {"train": [], "dev": [], "test": []}
    for it in items:
        t = ts_key(it)
        out["train" if t < train_end else "dev" if t < dev_end else "test"].append(it)
    return out


def leakage_report(splits: dict, text_key, group_key=None, ts_key=None, threshold: float = 0.8) -> dict:
    """Look for leaks from train into dev and test: exact copies, near copies, shared groups,
    and test items older than the newest training item."""
    rep = {"exact": [], "near": [], "groups": [], "time_violations": 0}
    train = splits.get("train", [])
    for other in [n for n in splits if n != "train"]:
        a = {f"train:{i}": text_key(x) for i, x in enumerate(train)}
        b = {f"{other}:{i}": text_key(x) for i, x in enumerate(splits[other])}
        ka = {}
        for i, t in a.items():
            ka.setdefault(exact_key(t), []).append(i)
        for i, t in b.items():
            for j in ka.get(exact_key(t), []):
                rep["exact"].append((j, i))
        near = near_duplicates({**a, **b}, threshold)
        exact = {tuple(sorted(p)) for p in rep["exact"]}
        for x, y, jv in near:
            tr_id, ot_id = (x, y) if x.startswith("train:") else (y, x)
            if tr_id.startswith("train:") and ot_id.startswith(other + ":") and tuple(sorted((x, y))) not in exact:
                rep["near"].append((tr_id, ot_id, jv))
        if group_key is not None:
            tg = {group_key(x) for x in train}
            for g in sorted({group_key(x) for x in splits[other]} & tg):
                rep["groups"].append((other, g))
        if ts_key is not None and train and splits[other]:
            newest = max(ts_key(x) for x in train)
            rep["time_violations"] += sum(1 for x in splits[other] if ts_key(x) <= newest)
    rep["ok"] = not (rep["exact"] or rep["near"] or rep["groups"] or rep["time_violations"])
    return rep


def decontaminate(train, test, text_key, threshold: float = 0.8) -> tuple:
    """Drop training items that are exact or near copies of a test item. Returns (kept, dropped)."""
    a = {f"a{i}": text_key(x) for i, x in enumerate(train)}
    b = {f"b{i}": text_key(x) for i, x in enumerate(test)}
    drop = set()
    kb = {exact_key(t) for t in b.values()}
    for i, t in a.items():
        if exact_key(t) in kb:
            drop.add(i)
    for x, y, _ in near_duplicates({**a, **b}, threshold):
        if x[0] != y[0]:
            drop.add(x if x[0] == "a" else y)
    kept = [x for i, x in enumerate(train) if f"a{i}" not in drop]
    dropped = [x for i, x in enumerate(train) if f"a{i}" in drop]
    return kept, dropped


# --------------------------------------------------------------------------- 5. the flywheel
def train_on(skill: dict, sft, lr: float = 0.5, scale: float = 15.0) -> dict:
    """The mock 'training' step (MADE UP rule). For each type, the skill moves toward 1 in
    proportion to how many examples there are (up to `scale`), and moves AWAY if most of the
    examples are wrong. Noisy data can make the mock agent worse."""
    new = dict(skill)
    for kind in skill:
        ex = [e for e in sft if e["task_type"] == kind]
        if not ex:
            continue
        q = mean(e["reward"] for e in ex)
        step = lr * min(1.0, len(ex) / scale) * (2 * q - 1)
        room = (1 - skill[kind]) if step > 0 else skill[kind]
        new[kind] = min(0.99, max(0.01, skill[kind] + step * room))
    return new


def make_eval_set(seed: int = 99, per_type: int = 10) -> list:
    rng = rng_for("evalset", seed)
    return [make_task(k, rng, f"e{seed}-{k}-{i:03d}") for k in TYPES for i in range(per_type)]


def expected_score(agent: MockAgent, cases, exclude_prompts=None) -> float:
    """Exact expected pass rate of one attempt (no sampling), optionally skipping prompts."""
    ex = exclude_prompts or set()
    xs = [agent.effective_skill(c) for c in cases if prompt_key(c.prompt) not in ex]
    return mean(xs)


@dataclass
class FlywheelConfig:
    rounds: int = 6
    sessions: int = 80
    label_budget: int = 50
    label_noise: float = 0.0
    lr: float = 0.5
    scale: float = 15.0
    harnesses: tuple = ("react",)
    deploy_harness: str = "react"
    filter_failures: bool = True
    use_consent: bool = True
    seed: int = 0


def label_traces(traces, budget: int, noise: float, rng: random.Random) -> dict:
    """A labeling budget: pick `budget` traces at random and label them with the verifier
    reward, flipped with probability `noise`. Returns {trace_id: label}."""
    ids = [t.trace_id for t in traces if t.reward is not None]
    rng.shuffle(ids)
    by = {t.trace_id: t for t in traces}
    out = {}
    for tid in ids[:max(0, budget)]:
        lab = by[tid].reward
        out[tid] = 1.0 - lab if rng.random() < noise else lab
    return out


def run_flywheel(cfg: FlywheelConfig, init_skill=None) -> dict:
    """Deploy -> log -> label -> improve, round after round, on the mock world."""
    skill = dict(init_skill or INIT_SKILL)
    memorized: set = set()
    trained_h: set = set()
    evalset = make_eval_set()
    hist, all_traces = [], []
    for r in range(cfg.rounds):
        h = cfg.harnesses[r % len(cfg.harnesses)]
        agent = MockAgent(skill, h, f"mock-r{r}", trained_h or None, memorized,
                          seed=stable_hash("fw", cfg.seed, r) % 10 ** 6)
        logger = TraceLogger(prefix=f"r{r}-")
        traces = simulate_traffic(cfg.sessions, seed=cfg.seed * 1000 + r, agent=agent, user_sim=UserSim(r),
                                  start=r * 10000.0, logger=logger, session_prefix=f"r{r}-")
        all_traces += traces
        labels = label_traces(traces, cfg.label_budget, cfg.label_noise, rng_for("label", cfg.seed, r))
        labelled = [replace(t, reward=labels[t.trace_id]) for t in traces if t.trace_id in labels]
        sft = traces_to_sft(labelled, min_reward=1.0 if cfg.filter_failures else 0.0, require_consent=cfg.use_consent)
        skill = train_on(skill, sft, cfg.lr, cfg.scale)
        memorized |= {prompt_key(e["prompt"]) for e in sft}
        trained_h |= {e["harness"] for e in sft}
        dep = MockAgent(skill, cfg.deploy_harness, "deploy", trained_h or None, memorized)
        hist.append({"round": r, "online": round(mean(t.reward for t in traces), 3), "logged": len(traces),
                     "labelled": len(labelled), "sft": len(sft),
                     "eval_all": round(expected_score(dep, evalset), 3),
                     "eval_clean": round(expected_score(dep, evalset, memorized), 3),
                     "eval_overlap": sum(prompt_key(c.prompt) in memorized for c in evalset)})
    return {"history": hist, "skill": skill, "memorized": memorized, "traces": all_traces}


def collapse_demo(probs: dict, n: int, generations: int, seed: int = 0, real_mix: float = 0.0) -> list:
    """A toy of recursive training. Each generation samples n items from the current
    distribution and the next distribution is their frequencies, mixed with the original at
    weight real_mix. Rare items vanish when n is small and real_mix is 0. Not a language model."""
    total = sum(probs.values())
    orig = {k: v / total for k, v in probs.items()}
    cur = dict(orig)
    rng = random.Random(seed)
    keys = sorted(orig)
    rows = []
    for g in range(generations + 1):
        rows.append({"generation": g, "support": sum(1 for k in keys if cur[k] > 0),
                     "entropy": round(entropy([cur[k] for k in keys]), 3), "p_rarest": round(cur[keys[-1]], 4)})
        if g == generations:
            break
        draw = rng.choices(keys, weights=[cur[k] for k in keys], k=n)
        freq = {k: draw.count(k) / n for k in keys}
        cur = {k: (1 - real_mix) * freq[k] + real_mix * orig[k] for k in keys}
    return rows


# --------------------------------------------------------------------------- 6. synthetic tasks
_PARAPHRASE = {"Add the numbers": "Sum these numbers",
               "Reverse the order of these words": "Put these words in reverse order",
               "Write this text in uppercase": "Convert this text to capital letters",
               "Count the vowels in the word": "How many vowels are in the word",
               "Sort these numbers from smallest to largest": "Order these numbers, smallest first",
               "What is the domain part of this email address": "Give the domain of this email address"}


def paraphrase_light(prompt: str) -> str:
    head, payload = prompt.split(": ", 1)
    return _PARAPHRASE.get(head, head) + ": " + payload


def propose_candidates(n: int, seed: int = 0, p_bad_answer: float = 0.10, p_invalid: float = 0.08,
                       p_near: float = 0.12, p_dup: float = 0.05) -> list:
    """Stand-in for 'ask a model for new tasks'. Most candidates are fine; some are broken on
    purpose so the filters have work. The `origin` field is the hidden truth, for scoring filters."""
    rng = rng_for("propose", seed)
    out = []
    for i in range(n):
        r = rng.random()
        kind = rng.choice(TYPES)
        t = make_task(kind, rng, f"g{seed}-{i:04d}")
        c = {"id": t.task_id, "type": kind, "prompt": t.prompt, "answer": t.answer, "origin": "clean"}
        if r < p_bad_answer:
            c["answer"], c["origin"] = corrupt(kind, t.answer, rng), "bad_answer"
        elif r < p_bad_answer + p_invalid:
            c["prompt"], c["origin"] = rng.choice(["", "Do it", "TODO {fill in}"]), "invalid"
        elif r < p_bad_answer + p_invalid + p_near and any(o["origin"] == "clean" for o in out):
            src = rng.choice([o for o in out if o["origin"] == "clean"])
            c.update(type=src["type"], prompt=paraphrase_light(src["prompt"]), answer=src["answer"], origin="near_dup")
        elif r < p_bad_answer + p_invalid + p_near + p_dup and any(o["origin"] == "clean" for o in out):
            src = rng.choice([o for o in out if o["origin"] == "clean"])
            c.update(type=src["type"], prompt=src["prompt"], answer=src["answer"], origin="dup")
        out.append(c)
    return out


def payload_words(prompt: str) -> set:
    return set(words(prompt.split(": ", 1)[1])) if ": " in prompt else set(words(prompt))


def filter_valid(c: dict) -> str | None:
    p = c["prompt"]
    if not p.strip():
        return "empty"
    if len(p) < 15 or len(p) > 200:
        return "length"
    if "{" in p or "TODO" in p:
        return "placeholder"
    return None


def filter_verified(c: dict) -> str | None:
    try:
        ok = solve(c["type"], c["prompt"]) == str(c["answer"]).strip()
    except (ValueError, IndexError):
        return "unparseable"
    return None if ok else "wrong_answer"


def filter_dedupe(cands, existing=(), threshold: float = 0.8) -> tuple:
    """Greedy: keep a candidate unless an earlier kept one, or an existing task, has the same
    type and payload overlap (Jaccard) at or above the threshold. Compares payloads, not
    templates, because the templates are shared by design."""
    kept, rejected = [], []
    pool = [(e["type"], payload_words(e["prompt"])) for e in existing]
    for c in cands:
        pw = payload_words(c["prompt"])
        if any(t == c["type"] and jaccard(pw, w) >= threshold for t, w in pool):
            rejected.append(c)
        else:
            kept.append(c)
            pool.append((c["type"], pw))
    return kept, rejected


def pass_count(c: dict, agent: MockAgent, k: int, seed: int = 0) -> int:
    """How many of k independent mock attempts get the claimed answer right."""
    n = 0
    for j in range(k):
        rng = rng_for("diff", seed, c["id"], j)
        n += rng.random() < agent.skill[c["type"]]
    return n


def filter_difficulty(cands, agent: MockAgent, k: int = 8, seed: int = 0) -> tuple:
    """Keep tasks that the agent sometimes solves and sometimes misses (0 < passes < k)."""
    kept, rejected = [], []
    for c in cands:
        n = pass_count(c, agent, k, seed)
        (kept if 0 < n < k else rejected).append(dict(c, passes=n))
    return kept, rejected


def generate_pipeline(n: int, agent: MockAgent, seed: int = 0, existing=(), k: int = 8,
                      dedupe_threshold: float = 0.8) -> dict:
    cands = propose_candidates(n, seed)
    funnel = [("proposed", len(cands))]
    rejects: dict[str, int] = {}
    ok = []
    for c in cands:
        why = filter_valid(c) or filter_verified(c)
        if why:
            rejects[why] = rejects.get(why, 0) + 1
        else:
            ok.append(c)
    funnel.append(("valid and verified", len(ok)))
    ok, dup = filter_dedupe(ok, existing, dedupe_threshold)
    rejects["duplicate"] = len(dup)
    funnel.append(("after dedupe", len(ok)))
    ok, hard = filter_difficulty(ok, agent, k, seed)
    rejects["too easy or too hard"] = len(hard)
    funnel.append(("after difficulty", len(ok)))
    return {"kept": ok, "funnel": funnel, "rejects": dict(sorted(rejects.items())),
            "origin_kept": dict(sorted(count_by(ok, "origin").items())),
            "origin_proposed": dict(sorted(count_by(cands, "origin").items()))}


def diversity(tasks) -> dict:
    """Distinct types, entropy of the type mix, and distinct instruction openings (first 3 words)."""
    types = count_by(tasks, "type")
    opens = {" ".join(words(t["prompt"])[:3]) for t in tasks}
    return {"n": len(tasks), "types": len(types), "type_entropy": round(entropy(types), 3), "openings": len(opens)}


# --------------------------------------------------------------------------- 7. implicit feedback
def extract_signals(tr: Trace) -> dict:
    """Pull implicit-feedback signals out of one trace."""
    counts = {t: 0 for t in FEEDBACK_TYPES}
    edit_ratio = 0.0
    for e in tr.feedback():
        t = e.meta.get("type")
        if t in counts:
            counts[t] += 1
        if t == "edit":
            edit_ratio = max(edit_ratio, 1.0 - difflib.SequenceMatcher(None, tr.final_text().strip(), e.text.strip()).ratio())
    users = [e for e in tr.events if e.kind == "user"]
    rephrase = any(jaccard(set(words(users[i].text)), set(words(users[j].text))) >= 0.6
                   for i in range(len(users)) for j in range(i + 1, len(users)))
    return {**counts, "edit_ratio": round(edit_ratio, 3), "rephrase": rephrase,
            "n_feedback": sum(counts.values())}


SIGNAL_WEIGHTS = {"copy": 0.6, "thumbs_up": 1.0, "thumbs_down": -1.0, "regenerate": -0.5,
                  "abandon": -0.8, "rephrase": -0.6, "edit": -0.8}


def signals_to_score(sig: dict) -> float:
    """Combine signals into a score in [-1, 1]. The weights are my own, not from any paper.
    Any edit counts for half of its weight, and a big rewrite counts in full."""
    w = SIGNAL_WEIGHTS
    s = w["copy"] * min(sig["copy"], 1) + w["thumbs_up"] * min(sig["thumbs_up"], 1)
    s += w["thumbs_down"] * min(sig["thumbs_down"], 1)
    s += w["regenerate"] * min(sig["regenerate"], 2)
    s += w["abandon"] * min(sig["abandon"], 1)
    s += w["rephrase"] * (1 if sig["rephrase"] else 0)
    if sig["edit"]:
        s += w["edit"] * (0.5 + 0.5 * sig["edit_ratio"])
    return max(-1.0, min(1.0, s))


def weak_label(score: float, hi: float = 0.4, lo: float = -0.4) -> int | None:
    """1 for good, 0 for bad, None when the signal is too weak to use."""
    if score >= hi:
        return 1
    if score <= lo:
        return 0
    return None


def signal_quality(traces, hi: float = 0.4, lo: float = -0.4) -> dict:
    """Compare weak labels to the verifier reward (known in the mock). Reports coverage,
    precision and recall of 'good', and accuracy of the labels that were produced."""
    tp = fp = fn = tn = 0
    for tr in traces:
        wl = weak_label(signals_to_score(extract_signals(tr)), hi, lo)
        if wl is None:
            continue
        truth = 1 if tr.reward >= 1.0 else 0
        if wl == 1 and truth == 1:
            tp += 1
        elif wl == 1:
            fp += 1
        elif truth == 1:
            fn += 1
        else:
            tn += 1
    n = len(traces)
    done = tp + fp + fn + tn
    return {"n": n, "coverage": round(done / n, 3) if n else 0.0,
            "precision_good": round(tp / (tp + fp), 3) if tp + fp else 0.0,
            "recall_good": round(tp / (tp + fn), 3) if tp + fn else 0.0,
            "accuracy": round((tp + tn) / done, 3) if done else 0.0, "tp": tp, "fp": fp, "fn": fn, "tn": tn}


def bradley_terry(votes, iters: int = 500) -> dict:
    """Strengths from pairwise votes [(winner, loser), ...] by the standard minorize-maximize
    update. Scaled so the geometric mean is 1. Needs every model to win at least once to be
    finite. P(a beats b) = s_a / (s_a + s_b)."""
    models = sorted({m for v in votes for m in v})
    wins = {m: 0 for m in models}
    n: dict = {}
    for w, l in votes:
        wins[w] += 1
        n[(w, l)] = n.get((w, l), 0) + 1
    p = {m: 1.0 for m in models}
    for _ in range(iters):
        new = {}
        for i in models:
            den = sum((n.get((i, j), 0) + n.get((j, i), 0)) / (p[i] + p[j]) for j in models if j != i)
            new[i] = wins[i] / den if den else p[i]
        g = math.exp(mean(math.log(max(v, 1e-12)) for v in new.values()))
        p = {m: v / g for m, v in new.items()}
    return p


# --------------------------------------------------------------------------- 8. privacy
def luhn_ok(digits: str) -> bool:
    ds = [int(c) for c in digits if c.isdigit()]
    if len(ds) < 13 or len(ds) > 19:
        return False
    total = 0
    for i, d in enumerate(reversed(ds)):
        if i % 2:
            d *= 2
            if d > 9:
                d -= 9
        total += d
    return total % 10 == 0


_PII = [
    ("EMAIL", re.compile(r"\b[\w.+-]+@[\w-]+(?:\.[\w-]+)+\b")),
    ("KEY", re.compile(r"\b(?:sk|pk|ghp|AKIA)[-_A-Za-z0-9]{16,}\b")),
    ("SSN", re.compile(r"\b\d{3}-\d{2}-\d{4}\b")),
    ("CARD", re.compile(r"\b(?:\d[ -]?){13,19}\b")),
    ("PHONE", re.compile(r"(?<![\w-])(?:\+?1[ .-]?)?\(?\d{3}\)?[ .-]\d{3}[ .-]\d{4}\b")),
    ("IP", re.compile(r"\b(?:(?:25[0-5]|2[0-4]\d|1?\d?\d)\.){3}(?:25[0-5]|2[0-4]\d|1?\d?\d)\b")),
]


class Scrubber:
    """Replaces PII with stable placeholders such as <EMAIL_1>. The same value always gets the
    same placeholder, so a story stays readable, but the value itself is gone. It only finds
    what the patterns and the name list describe. It will miss other names, street
    addresses and anything unusual."""

    def __init__(self, names=()):
        self.names = [n for n in names if n]
        self.map: dict[tuple, str] = {}
        self.counts: dict[str, int] = {}

    def _ph(self, kind: str, value: str) -> str:
        key = (kind, value.lower())
        if key not in self.map:
            n = sum(1 for k in self.map if k[0] == kind) + 1
            self.map[key] = f"<{kind}_{n}>"
        self.counts[kind] = self.counts.get(kind, 0) + 1
        return self.map[key]

    def scrub(self, text: str) -> str:
        out = str(text)
        for kind, pat in _PII:
            def sub(m, kind=kind):
                v = m.group(0)
                if kind == "CARD" and not luhn_ok(v):
                    return v
                return self._ph(kind, v)
            out = pat.sub(sub, out)
        for name in self.names:
            out = re.sub(r"\b" + re.escape(name) + r"\b", lambda m: self._ph("NAME", m.group(0)), out, flags=re.I)
        return out


def pseudonymize_user(user_id: str, salt: str) -> str:
    return "p-" + hashlib.sha256(f"{salt}|{user_id}".encode()).hexdigest()[:10]


def scrub_trace(tr: Trace, salt: str, names=()) -> Trace:
    """A scrubbed copy: one Scrubber per trace, user id replaced by a salted hash."""
    sc = Scrubber(names)
    evs = [Event(e.kind, sc.scrub(e.text), e.ts, {k: (sc.scrub(v) if isinstance(v, str) else v) for k, v in e.meta.items()})
           for e in tr.events]
    ver = None
    if tr.verifier is not None:
        ver = {k: (sc.scrub(v) if isinstance(v, str) else v) for k, v in tr.verifier.items()}
    return replace(tr, events=evs, verifier=ver, user_id=pseudonymize_user(tr.user_id, salt), scrubbed=True)


def residual_pii(text: str) -> list:
    """Kinds of PII that the patterns still find. Run it after scrubbing as a gate."""
    found = []
    for kind, pat in _PII:
        for m in pat.finditer(str(text)):
            if kind != "CARD" or luhn_ok(m.group(0)):
                found.append(kind)
    return found


def trace_residual_pii(tr: Trace) -> list:
    out = []
    for e in tr.events:
        out += residual_pii(e.text)
        out += [k for v in e.meta.values() if isinstance(v, str) for k in residual_pii(v)]
    return sorted(set(out))


def consent_filter(traces, purpose: str) -> tuple:
    """Keep traces whose owner allowed `purpose` ('log', 'train' or 'share'). For 'train' and
    'share' the trace must also be scrubbed. Returns (kept, dropped_reasons)."""
    if purpose not in DEFAULT_CONSENT:
        raise ValueError(f"unknown purpose {purpose!r}")
    kept, why = [], {"no_consent": 0, "not_scrubbed": 0}
    for t in traces:
        if not t.consent.get(purpose):
            why["no_consent"] += 1
        elif purpose in ("train", "share") and not t.scrubbed:
            why["not_scrubbed"] += 1
        else:
            kept.append(t)
    return kept, why


def retention_filter(traces, now: float, max_age_days: float) -> tuple:
    """Drop traces started more than max_age_days before `now` (timestamps are seconds)."""
    limit = now - max_age_days * 86400.0
    kept = [t for t in traces if t.started >= limit]
    return kept, len(traces) - len(kept)


def delete_user(traces, user_id: str, salt: str | None = None) -> tuple:
    """Remove every trace of one user. If traces were pseudonymized, pass the salt as well."""
    ids = {user_id} | ({pseudonymize_user(user_id, salt)} if salt else set())
    kept = [t for t in traces if t.user_id not in ids]
    return kept, len(traces) - len(kept)


# --------------------------------------------------------------------------- 9. data cards
CARD_FIELDS = ("name", "version", "purpose", "sources", "counts", "labels", "time_range", "harnesses",
               "privacy", "splits", "filters", "known_limits", "not_for")


def make_data_card(name: str, version: str, traces, purpose: str, sources, filters, known_limits,
                   not_for, splits: dict | None = None, leakage: dict | None = None,
                   synthetic_note: str = "") -> dict:
    """Facts that can be counted are counted from the data. Everything else is passed in."""
    traces = list(traces)
    rewards = [t.reward for t in traces if t.reward is not None]
    ts = [t.started for t in traces]
    card = {
        "name": name, "version": version, "purpose": purpose, "sources": list(sources),
        "counts": {"traces": len(traces),
                   "by_task_type": dict(sorted(count_by([{"k": t.task_type} for t in traces], "k").items())),
                   "users": len({t.user_id for t in traces})},
        "labels": {"labelled": len(rewards), "mean_reward": round(mean(rewards), 3) if rewards else None},
        "time_range": [min(ts), max(ts)] if ts else None,
        "harnesses": dict(sorted(count_by([{"k": t.harness} for t in traces], "k").items())),
        "privacy": {"scrubbed": sum(t.scrubbed for t in traces), "unscrubbed": sum(not t.scrubbed for t in traces),
                    "train_consent": sum(bool(t.consent.get("train")) for t in traces),
                    "share_consent": sum(bool(t.consent.get("share")) for t in traces),
                    "residual_pii_traces": sum(1 for t in traces if trace_residual_pii(t))},
        "splits": {k: len(v) for k, v in (splits or {}).items()},
        "filters": list(filters), "known_limits": list(known_limits), "not_for": list(not_for),
        "synthetic_note": synthetic_note,
    }
    if leakage is not None:
        card["leakage"] = {"ok": leakage["ok"], "exact": len(leakage["exact"]), "near": len(leakage["near"]),
                           "shared_groups": len(leakage["groups"]), "time_violations": leakage["time_violations"]}
    return card


def validate_card(card: dict) -> list:
    """Missing or empty required fields, and obvious contradictions."""
    problems = [f"missing {f}" for f in CARD_FIELDS if f not in card]
    for f in ("purpose", "sources", "filters", "known_limits", "not_for"):
        if f in card and not card[f]:
            problems.append(f"empty {f}")
    if card.get("splits") and card.get("counts") and sum(card["splits"].values()) > card["counts"]["traces"]:
        problems.append("splits hold more items than the card counts")
    if "privacy" in card and card["privacy"]["unscrubbed"] and card["privacy"]["share_consent"]:
        problems.append("unscrubbed traces in a card that lists sharing")
    if "leakage" in card and not card["leakage"]["ok"]:
        problems.append("leakage report is not clean")
    return problems


def render_card(card: dict) -> str:
    lines = [f"# Data card: {card['name']} (version {card['version']})", "", f"**Purpose.** {card['purpose']}", ""]
    c = card["counts"]
    lines += ["## Contents", f"- Traces: {c['traces']} from {c['users']} users",
              "- By task type: " + ", ".join(f"{k} {v}" for k, v in c["by_task_type"].items()),
              "- Harnesses: " + ", ".join(f"{k} {v}" for k, v in card["harnesses"].items()),
              f"- Labelled: {card['labels']['labelled']}, mean reward {card['labels']['mean_reward']}"]
    if card["time_range"]:
        lines.append(f"- Time range (seconds from start): {card['time_range'][0]:.0f} to {card['time_range'][1]:.0f}")
    p = card["privacy"]
    lines += ["", "## Privacy and consent",
              f"- Scrubbed: {p['scrubbed']}, not scrubbed: {p['unscrubbed']}",
              f"- Owners who allow training: {p['train_consent']}, sharing: {p['share_consent']}",
              f"- Traces with residual PII found by the patterns: {p['residual_pii_traces']}"]
    if card["splits"]:
        lines += ["", "## Splits", "- " + ", ".join(f"{k} {v}" for k, v in card["splits"].items())]
    if "leakage" in card:
        lg = card["leakage"]
        lines.append(f"- Leakage checks: {'clean' if lg['ok'] else 'NOT clean'} (exact {lg['exact']}, near {lg['near']}, "
                     f"shared groups {lg['shared_groups']}, time violations {lg['time_violations']})")
    lines += ["", "## Sources and filters"] + [f"- Source: {s}" for s in card["sources"]] + [f"- Filter: {f}" for f in card["filters"]]
    if card.get("synthetic_note"):
        lines += ["", f"**Synthetic data.** {card['synthetic_note']}"]
    lines += ["", "## Known limits"] + [f"- {x}" for x in card["known_limits"]]
    lines += ["", "## Not for"] + [f"- {x}" for x in card["not_for"]]
    return "\n".join(lines) + "\n"


# --------------------------------------------------------------------------- demo
def demo() -> None:
    out = print
    out("== 1. Trace schema, logging, validation ==")
    traces = simulate_traffic(120, seed=1, user_sim=UserSim(1), n_users=12)
    tr = traces[0]
    out(f"{len(traces)} traces logged; traces with validation problems: {sum(bool(validate_trace(t)) for t in traces)}")
    out(f"first trace {tr.trace_id}: type {tr.task_type}, harness {tr.harness}, reward {tr.reward}, {len(tr.events)} events")
    for e in tr.events:
        out(f"  {e.ts:6.1f}  {e.kind:<11} {e.text[:60]!r}" + (f"  meta={e.meta}" if e.meta else ""))
    rt = loads_jsonl(dumps_jsonl(traces))
    out(f"JSON Lines round trip equal: {rt == traces}; size {len(dumps_jsonl(traces))} bytes")
    out("span summary of the first trace: " + str(span_summary(to_span_tree(tr))))
    bad = Trace("x", "s", "u", "sum", "react", "m", 0.0, [Event("tool_result", "7", 1.0), Event("assistant", "hi", 0.5)])
    out("a broken trace: " + str(validate_trace(bad)))
    out("")
    out("== 2. Converting traces ==")
    sft = traces_to_sft(traces)
    prefs = traces_to_preferences(traces)
    edits = edit_pairs(traces)
    evals = traces_to_eval_cases(traces)
    ok = sum(1 for t in traces if t.reward == 1.0)
    out(f"traces {len(traces)}: passed {ok}, failed {len(traces) - ok}; owners allowing training: {sum(t.consent['train'] for t in traces)}")
    out(f"SFT demonstrations: {len(sft)} (passed AND consented); harness mix {harness_mix(sft)}")
    ex = sft[0]
    out(f"  example {ex['trace_id']}: learned and context words {loss_tokens(ex)}; learned roles: " + str([m['role'] for m in ex['messages'] if m['train']]))
    out(f"preference pairs: {len(prefs)} by outcome, {len(edits)} by edit")
    if prefs:
        out("  e.g. " + str({k: prefs[0][k] for k in ('prompt', 'chosen', 'rejected')}))
    out(f"eval cases from failures: {len(evals)}; per type: {count_by(evals, 'task_type')}")
    out("")
    out("== 3. Splits and leakage ==")
    items = [{"id": t.trace_id, "text": t.prompt(), "group": t.user_id, "ts": t.started} for t in traces]
    rnd = split_random(items, seed=3)
    rr = leakage_report(rnd, lambda x: x["text"], lambda x: x["group"], lambda x: x["ts"])
    out(f"random split {({k: len(v) for k, v in rnd.items()})}: exact copies {len(rr['exact'])}, near copies {len(rr['near'])}, shared users {len(rr['groups'])}, time violations {rr['time_violations']}")
    grp = split_by_group(items, lambda x: x["text"])
    gr = leakage_report(grp, lambda x: x["text"], lambda x: x["text"])
    out(f"split by prompt group {({k: len(v) for k, v in grp.items()})}: exact {len(gr['exact'])}, near {len(gr['near'])}, shared prompts {len(gr['groups'])}")
    sneaky = {**grp, "test": grp["test"] + [{"text": "again: " + x["text"]} for x in grp["train"][:5]]}
    sr = leakage_report(sneaky, lambda x: x["text"], lambda x: x["text"])
    out(f"same split plus 5 reworded copies of training prompts in test: exact {len(sr['exact'])}, near {len(sr['near'])}, shared prompts {len(sr['groups'])}, ok {sr['ok']}")
    ts = time_split(items, lambda x: x["ts"], 0.7 * 120 * 30.0, 0.85 * 120 * 30.0)
    tl = leakage_report(ts, lambda x: x["text"], None, lambda x: x["ts"])
    kept, dropped = decontaminate(ts["train"], ts["test"], lambda x: x["text"])
    out(f"time split {({k: len(v) for k, v in ts.items()})}: time violations {tl['time_violations']}, exact copies train-to-test {len(tl['exact'])}; decontaminate drops {len(dropped)} of {len(ts['train'])} training items")
    texts = {t.trace_id: t.prompt() for t in traces[:40]}
    fast, slow = near_duplicates(texts, 0.8), brute_force_near_duplicates(texts, 0.8)
    out(f"MinHash found {len(fast)} near-duplicate pairs among 40 prompts; brute force found {len(slow)}; same: {fast == slow}")
    out("")
    out("== 4. Flywheel on the mock ==")
    for label, cfg in (("base: filter failures, 50 labels per round", FlywheelConfig()),
                       ("no filter: train on failures too", FlywheelConfig(filter_failures=False)),
                       ("label noise 30 percent", FlywheelConfig(label_noise=0.3)),
                       ("no labels at all", FlywheelConfig(label_budget=0))):
        res = run_flywheel(cfg)
        out(f"{label}")
        out("  round: online/eval_all/eval_clean/overlap/sft")
        out("  " + " ".join(f"{h['round']}:{h['online']:.2f}/{h['eval_all']:.2f}/{h['eval_clean']:.2f}/{h['eval_overlap']}/{h['sft']}" for h in res['history']))
    one = run_flywheel(FlywheelConfig(harnesses=("react",), deploy_harness="json"))["history"][-1]
    mix = run_flywheel(FlywheelConfig(harnesses=("react", "json"), deploy_harness="json"))["history"][-1]
    out(f"deployed in the json harness after 6 rounds: trained on react only eval_clean {one['eval_clean']:.3f}; trained on react and json {mix['eval_clean']:.3f}")
    out("")
    out("== 5. Model collapse toy (not a language model) ==")
    probs = {"a": 0.50, "b": 0.25, "c": 0.15, "d": 0.06, "e": 0.03, "f": 0.01}
    for mix_w in (0.0, 0.2):
        rows = collapse_demo(probs, n=40, generations=12, seed=4, real_mix=mix_w)
        out(f"real_mix {mix_w}: support by generation " + str([r['support'] for r in rows]) + f"; entropy {rows[0]['entropy']} -> {rows[-1]['entropy']}")
    out("")
    out("== 6. Synthetic tasks with filters ==")
    agent = MockAgent(seed=5)
    res = generate_pipeline(400, agent, seed=2)
    out("funnel: " + " -> ".join(f"{n} {c}" for n, c in res["funnel"]))
    out(f"rejects: {res['rejects']}")
    out(f"proposed by hidden origin: {res['origin_proposed']}")
    out(f"kept by hidden origin: {res['origin_kept']}")
    out(f"diversity of kept: {diversity(res['kept'])}; of all proposed: {diversity(propose_candidates(400, 2))}")
    out("")
    out("== 7. Implicit feedback ==")
    sig = extract_signals(next(t for t in traces if t.feedback()))
    out(f"signals of one trace: {sig}")
    out(f"weak labels from signals vs the verifier: {signal_quality(traces)}")
    votes = [("a", "b")] * 6 + [("b", "a")] * 2 + [("b", "c")] * 5 + [("c", "b")] * 3 + [("a", "c")] * 7 + [("c", "a")] * 1
    bt = bradley_terry(votes)
    out("Bradley-Terry strengths from 24 made-up votes: " + str({k: round(v, 3) for k, v in bt.items()}))
    out("")
    out("== 8. PII scrubbing and consent ==")
    fake_key = "sk-" + "demo" * 5
    raw = ("Hi, I am Ada Lovelace. Mail ada@example.org or call (415) 555-0132. Card 4111 1111 1111 1111, "
           f"SSN 123-45-6789, key {fake_key}, host 10.0.0.12. Ada again: ada@example.org. Order 4111 1111 1111 1112.")
    sc = Scrubber(names=["Ada Lovelace"])
    clean = sc.scrub(raw)
    out("before: " + raw)
    out("after:  " + clean)
    out(f"counts {sc.counts}; residual PII after: {residual_pii(clean)}; a street address passes through: {Scrubber().scrub('I live at 12 Oak Street, Springfield')!r}")
    sc_traces = [scrub_trace(t, "demo-salt") for t in traces]
    dom = [t for t in traces if t.task_type == "domain"][0]
    out(f"a domain-task prompt before: {dom.prompt()!r}; after: {scrub_trace(dom, 'demo-salt').prompt()!r}")
    out(f"traces with residual PII, raw: {sum(bool(trace_residual_pii(t)) for t in traces)}; scrubbed: {sum(bool(trace_residual_pii(t)) for t in sc_traces)}")
    keep_tr, why = consent_filter(sc_traces, "train")
    keep_sh, why_sh = consent_filter(sc_traces, "share")
    raw_keep, raw_why = consent_filter(traces, "train")
    out(f"scrubbed, for training: kept {len(keep_tr)}, dropped {why}; for sharing: kept {len(keep_sh)}, dropped {why_sh}")
    out(f"raw (not scrubbed), for training: kept {len(raw_keep)}, dropped {raw_why}")
    ret, gone = retention_filter(sc_traces, now=3600.0, max_age_days=0.5 / 24)
    out(f"retention of 30 minutes, checked at second 3600: kept {len(ret)}, dropped {gone}")
    left, n_del = delete_user(sc_traces, "u03", "demo-salt")
    out(f"delete user u03 (pseudonymized): removed {n_del}, left {len(left)}")
    out("")
    out("== 9. Data card ==")
    cards_in = [{"t": t, "text": t.prompt()} for t in keep_tr]
    csplit = split_by_group(cards_in, lambda x: x["text"])
    crep = leakage_report(csplit, lambda x: x["text"], lambda x: x["text"])
    card = make_data_card("mock-agent-traces", "0.1", keep_tr, "Teach a mock agent; show the workflow.",
                          ["mock traffic from a simulator (made up)"],
                          ["owner allowed training", "scrubbed by regex", "schema validated", "no reward filter: failures kept"],
                          ["synthetic users and tasks", "regex scrubbing misses names and addresses", "six task types only"],
                          ["training a real model", "any claim about real users"],
                          splits=csplit, leakage=crep,
                          synthetic_note="All traces come from a mock agent and a user simulator.")
    out(render_card(card).rstrip())
    out(f"card problems: {validate_card(card)}")


if __name__ == "__main__":
    demo()
    sys.exit(0)