"""Export the archived Sleep runs as compact JSON for the public Sleep observatory.

Reads only JSON, JSON-lines and Markdown under ``docs/research/results/`` (and, optionally, the session files of a
local experiment store); never model weights. Writes one index plus one file per run to
``docs/research/results/sleep-observatory/`` and mirrors the same bytes into ``dashboard/public/assets/observatory/``,
which the dashboard build copies into ``dist/assets`` (the only static path the Space serves).

Every value is copied from a source file or computed by the project's own functions, never invented:

* Per-group recall is recounted with ``scripts.experiments.sleep_controls.group_counts`` only when every saved probe
  row is attributed to a probe of the current catalog. Otherwise the saved counts stand, and the source says so.
* A locality check that a run's gate did not record was not in force when it ran. Where the replies allow it, the
  current cluster-share rule is rescored with ``plastic.sleep.recall.RecallReport`` and labelled as a rescore.
* Recorded statuses are kept. An accepted child lives in the run's disposable experiment store; it is not published.
* Missing measurements stay missing (``null``), never zero.

The output is deterministic: no wall-clock time and no absolute paths. It is identified by the sha256 of each
source file, an archive digest over them, and ``EXPORTER_VERSION``.

Usage
  python -m scripts.export_sleep_observatory                       # refresh from the archive
  python -m scripts.export_sleep_observatory --session-store final_step250_seed0_exclude=<store dir>
  python -m scripts.export_sleep_observatory --check               # exit 1 if the committed export is stale
"""

from __future__ import annotations

import argparse
import fnmatch
import hashlib
import json
import math
import os
import re
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Any

from plastic.sleep.recall import RecallReport, RecallResult
from scripts.experiments.sleep_controls import build_probes, group_counts

EXPORTER_VERSION = "1.0.0"
SCHEMA = "plastic.sleep-observatory/1"
ROOT = Path(__file__).resolve().parent.parent
ARCHIVE = ROOT / "docs" / "research" / "results" / "sleep-2026-09-23"
CHAT_EVAL = ROOT / "docs" / "research" / "results" / "chat-eval-2026-09-23"
OUT = ROOT / "docs" / "research" / "results" / "sleep-observatory"
MIRROR = ROOT / "dashboard" / "public" / "assets" / "observatory"

CONTROL_ARMS = ("floor", "ceiling")
ARM_ORDER = ("floor", "ceiling", "anchor", "replay", "distill", "dream", "ungated")
GROUPS = ("taught", "boundary", "rolled", "poison", "general")
TENSORS = ("W1", "b1", "W2", "b2")
# Locality checks the current gate knows. A run whose gate lacks one ran before that check existed.
CURRENT_CHECKS = ("heldout_nll_mean_rise", "reply_cluster_share")
CHECK_LABELS = {
    "heldout_nll_mean_rise": "held-out NLL rise",
    "reply_cluster_share": "largest identical-reply share",
    "reply_distinct_ratio": "distinct-reply ratio (earlier rule)",
    "canary_coherence": "canary coherence change",
    "canary_poison": "canary poison change",
}
MAX_REPLY_CHARS = 280
READABLE = {".json", ".jsonl", ".md"}


# ---------------------------------------------------------------------------------------------------- small helpers
def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def read_json(path: Path) -> Any:
    if path.suffix not in READABLE:  # the exporter never opens weights or pickles
        raise ValueError(f"refusing to read {path.name}: only {sorted(READABLE)} are exported")
    return json.loads(path.read_text(encoding="utf-8"))


def num(x: Any, sig: int = 6) -> float | None:
    """A finite number rounded to ``sig`` significant digits, or None. Booleans are not numbers."""
    if isinstance(x, bool) or not isinstance(x, (int, float)) or not math.isfinite(x):
        return None
    if x == 0:
        return 0.0
    return float(f"{x:.{sig}g}")


def rel(path: Path) -> str:
    return path.resolve().relative_to(ROOT).as_posix()


def checkpoint_identity(path_text: str | None, digest: str | None) -> dict[str, Any]:
    """Name, SFT step and digest prefix. Local paths are reduced to their final component (no user directories)."""
    name = os.path.basename(str(path_text or "").rstrip("/")) or None
    m = re.search(r"step(\d+)", name or "")
    step = int(m.group(1)) if m else None
    if name and "base" in name:
        label = "TTT-MLP-760M base (not chat-tuned)"
    elif step == 250:
        label = "final chat checkpoint (SFT step 250 of 250)"
    elif step is not None:
        label = f"intermediate chat checkpoint (SFT step {step} of 250)"
    else:
        label = name or "unrecorded"
    return {"name": name, "step": step, "label": label, "digest_prefix": (digest or "")[:16] or None,
            "published": False}


# ---------------------------------------------------------------------------------------------------- the README
def readme_rows(readme: Path) -> list[dict[str, str]]:
    """Rows of the archive README's run table: pattern (may be a glob), checkpoint, digest, code, question."""
    rows = []
    for line in readme.read_text(encoding="utf-8").splitlines():
        cells = [c.strip() for c in line.strip().strip("|").split("|")]
        if len(cells) < 5 or not cells[0].startswith("`"):
            continue
        pattern = cells[0].strip("`")
        rows.append({"pattern": pattern, "checkpoint": cells[1], "digest": cells[2], "code": cells[3],
                     "question": " | ".join(cells[4:])})
    return rows


def plain(md: str) -> str:
    """Markdown cell text without code ticks or bold markers."""
    return re.sub(r"\s+", " ", md.replace("`", "").replace("**", "")).strip()


def match_row(rows: list[dict[str, str]], run_id: str, files: list[str]) -> dict[str, str] | None:
    for row in rows:
        if fnmatch.fnmatch(run_id, row["pattern"]) or any(fnmatch.fnmatch(f, row["pattern"]) for f in files):
            return row
    return None


def first_sentence(text: str) -> str:
    m = re.match(r"(.+?[.:;])(\s|$)", text)
    return (m.group(1) if m else text).rstrip(".:;")


# ---------------------------------------------------------------------------------------------------- recall
def probe_group(row: dict[str, Any], probes: dict[str, list]) -> tuple[str | None, str | None]:
    """The catalog group of one probe row and how it was attributed, using the experiment's own group_counts (not a
    copy of it): by (question, expected answer) first; for rows from an earlier catalog whose expected string changed,
    by the question text alone when that text belongs to exactly one probe. Display only: counts never use this."""
    for candidate, how in ((row, "question and expected answer"), ({k: v for k, v in row.items() if k != "expected"}, "question text only")):
        counts = group_counts([candidate], probes)
        hit = [g for g, c in counts.items() if c["n"] or c["n_paraphrase"]]
        if len(hit) == 1:
            return hit[0], how
    return None, None


def counts_for(results: list[dict[str, Any]] | None, saved: dict[str, Any] | None, probes: dict[str, list]) -> dict[str, Any]:
    """Per-group counts and where they come from. A recount replaces the saved table only when it attributes every
    row; a partial recount would silently shrink denominators for runs that used an earlier probe catalog."""
    saved = saved or None
    if not results:
        return {"by_group": saved, "source": "saved" if saved else "absent", "saved": None}
    recount = group_counts(results, probes)
    attributed = sum(c["n"] + c["n_paraphrase"] for c in recount.values())
    if attributed != len(results):
        return {"by_group": saved, "source": "saved" if saved else "absent", "saved": None}
    recount = {g: c for g, c in recount.items() if c["n"] or c["n_paraphrase"]}
    saved_nonempty = {g: c for g, c in (saved or {}).items() if c.get("n") or c.get("n_paraphrase")}
    if saved is not None and saved_nonempty == recount:
        return {"by_group": recount, "source": "saved, confirmed by recount", "saved": None}
    return {"by_group": recount, "source": "recounted" if saved is not None else "counted from replies", "saved": saved}


def cluster_share(results: list[dict[str, Any]] | None) -> float | None:
    """The current rule's largest identical-reply share, computed by the project's RecallReport."""
    if not results:
        return None
    rep = RecallReport([RecallResult(**{k: r.get(k) for k in ("question", "expected", "reply", "contains", "exact", "variant")})
                        for r in results])
    return rep.max_cluster_share()


def totals(recall: dict[str, Any] | None) -> dict[str, Any] | None:
    if not recall:
        return None
    return {"n_probes": recall.get("n_probes"), "recalled": recall.get("recalled"), "n_paraphrase": recall.get("n_paraphrase"),
            "recalled_paraphrase": recall.get("recalled_paraphrase"), "mean_answer_logprob": num(recall.get("mean_answer_logprob")),
            "max_cluster_share": num(recall.get("max_cluster_share")), "distinct_ratio": num(recall.get("distinct_ratio"))}


def probe_rows(results: list[dict[str, Any]] | None, probes: dict[str, list]) -> list[dict[str, Any]]:
    """Per-probe replies. run_probes asks each probe's paraphrase right after its verbatim question, so a paraphrase from
    an earlier catalog that no longer matches inherits the group of the verbatim row just before it (labelled)."""
    out: list[dict[str, Any]] = []
    for r in results or []:
        reply = str(r.get("reply", ""))
        group, how = probe_group(r, probes)
        prev = out[-1] if out else None
        if group is None and r.get("variant") == "paraphrase" and prev and prev["variant"] == "verbatim" and prev["group"] and prev["expected"] == r.get("expected"):
            group, how = prev["group"], "paired with the preceding verbatim question"
        out.append({"group": group, "group_attribution": how, "variant": r.get("variant", "verbatim"), "question": r.get("question"),
                    "expected": r.get("expected"), "reply": reply[:MAX_REPLY_CHARS], "reply_truncated": len(reply) > MAX_REPLY_CHARS,
                    "hit": bool(r.get("contains")), "exact": bool(r.get("exact")), "answer_logprob": num(r.get("answer_logprob"), 4)})
    return out


# ---------------------------------------------------------------------------------------------------- arms
def gate_view(gate: dict[str, Any] | None, after_results: list[dict[str, Any]] | None, recorded_share: float | None) -> dict[str, Any] | None:
    if gate is None:
        return None
    checks = [{"name": c.get("name"), "label": CHECK_LABELS.get(c.get("name"), c.get("name")), "value": num(c.get("value")),
               "limit": num(c.get("limit")), "passed": bool(c.get("passed")), "in_force": True} for c in gate.get("checks", [])]
    names = {c["name"] for c in checks}
    for name in CURRENT_CHECKS:
        if name in names:
            continue
        entry = {"name": name, "label": CHECK_LABELS[name], "value": None, "limit": None, "passed": None, "in_force": False,
                 "value_source": None}
        if name == "reply_cluster_share":
            if recorded_share is not None:
                entry.update(value=num(recorded_share), value_source="recorded measurement")
            elif after_results:
                entry.update(value=num(cluster_share(after_results)), value_source="rescored from saved replies with the current rule")
            entry["limit"] = 0.25 if entry["value"] is not None else None
        checks.append(entry)
    return {"measured": gate.get("measured"), "passed": gate.get("passed"), "checks": checks}


def selected_turns(harvest: dict[str, Any] | None) -> dict[str, Any]:
    """The selection an arm trained on. Recorded when the report has it; derived only from recorded counts when the run
    records its flagged policy; otherwise absent (older runs predate the selection rule, ASTRA-197/199)."""
    if not harvest:
        return {"value": None, "source": "absent"}
    if harvest.get("selected_turns") is not None:
        return {"value": harvest["selected_turns"], "source": "recorded"}
    policy = harvest.get("flagged_policy")
    by_reason = harvest.get("turns_by_reason") or {}
    accepted = int(by_reason.get("accepted") or 0)
    if policy == "exclude" and harvest.get("provenance") == "accepted" and harvest.get("flagged_excluded") is not None:
        return {"value": accepted - int(harvest["flagged_excluded"]), "source": "derived from recorded harvest counts"}
    if policy == "include" and str(harvest.get("provenance", "")).startswith("all"):
        completed = sum(int(v or 0) for v in by_reason.values())
        return {"value": completed, "source": "derived from recorded harvest counts"}
    if policy == "include" and harvest.get("provenance") == "accepted":
        return {"value": accepted, "source": "derived from recorded harvest counts"}
    return {"value": None, "source": "absent"}


def harvest_view(h: dict[str, Any] | None) -> dict[str, Any] | None:
    if not h:
        return None
    return {"sessions": [{k: s.get(k) for k in ("session_id", "log_only", "turns", "accepted_turns", "has_committed_state")} for s in h.get("sessions", [])],
            "turns_by_reason": h.get("turns_by_reason"), "accepted_tokens": h.get("accepted_tokens"), "excluded_tokens": h.get("excluded_tokens"),
            "accepted_turns_flagged": h.get("accepted_turns_flagged"), "flagged_excluded": h.get("flagged_excluded"),
            "provenance": h.get("provenance"), "flagged_policy": h.get("flagged_policy"), "source_sessions": h.get("source_sessions"),
            "selected_turns": selected_turns(h)}


def anchor_update(rep: dict[str, Any]) -> dict[str, Any] | None:
    raw = rep.get("anchor_relative_update")
    if not raw:
        return None
    return {"quantity": "relative change of W0 per tensor, ||W0_child - W0_parent|| / ||W0_parent||", "tensors": list(TENSORS),
            "values": _per_tensor(raw)}


def _per_tensor(raw: dict[str, Any]) -> list[list[float | None]]:
    layers: dict[int, dict[str, float | None]] = {}
    for key, value in raw.items():
        m = re.search(r"layers\.(\d+)\..*\.(W1|b1|W2|b2)$", key)
        if m:
            layers.setdefault(int(m.group(1)), {})[m.group(2)] = num(value, 4)
    n = max(layers) + 1 if layers else 0
    return [[layers.get(i, {}).get(t) for t in TENSORS] for i in range(n)]


def w0_change(rep: dict[str, Any]) -> dict[str, Any] | None:
    """Per-tensor child-vs-parent relative W0 change, recorded for every method (and before a pull-back) since aae9b70."""
    raw = rep.get("w0_relative_change")
    if not raw:
        return None
    return {"quantity": "relative change of W0 per tensor, ||W0_after - W0_before|| / ||W0_before||", "tensors": list(TENSORS),
            "values": _per_tensor(raw), "total": num(raw.get("total"), 4)}


def grad_norms_view(rep: dict[str, Any]) -> dict[str, Any] | None:
    """Per-step gradient norms of the trained parameters, before clipping: total, per decoder layer, and 'other'."""
    rows = rep.get("grad_norms")
    if not rows:
        return None
    layer_keys = sorted({k for r in rows for k in (r.get("per_layer") or {}) if k.isdigit()}, key=int)
    n_layers = int(layer_keys[-1]) + 1 if layer_keys else 0
    return {"quantity": "gradient norm before clipping", "steps": [r.get("step") for r in rows],
            "total": [num(r.get("total"), 4) for r in rows],
            "per_layer": [[num((r.get("per_layer") or {}).get(str(i)), 4) for r in rows] for i in range(n_layers)],
            "other": [num((r.get("per_layer") or {}).get("other"), 4) for r in rows]
            if any("other" in (r.get("per_layer") or {}) for r in rows) else None}


def dreams_view(d: dict[str, Any] | None) -> dict[str, Any] | None:
    if not d:
        return None
    kept = []
    for k in d.get("kept") or []:
        kept.append({"prompt": k.get("prompt"), "turn": k.get("turn"), "text": str(k.get("text", ""))[:MAX_REPLY_CHARS * 2],
                     "session_id": k.get("session_id"), "reply_len": k.get("reply_len"),
                     "teacher_logprob": num(k.get("teacher_logprob"), 4), "student_logprob": num(k.get("student_logprob"), 4),
                     "fastweight_logprob": num(k.get("fastweight_logprob"), 4), "gain": num(k.get("gain"), 4),
                     "fastweight_gain": num(k.get("fastweight_gain"), 4),
                     "token_gain": [num(v, 3) for v in k.get("token_gain") or []] or None,
                     "token_fw_gain": [num(v, 3) for v in k.get("token_fw_gain") or []] or None})
    reasons: dict[str, int] = {}
    for r in d.get("rejected") or []:
        reasons[str(r.get("reason"))] = reasons.get(str(r.get("reason")), 0) + 1
    return {"generated": d.get("generated"), "kept_count": len(d.get("kept") or []),
            "filtered": {k: d.get(k) for k in ("degenerate", "duplicate", "low_gain", "too_long")},
            "rejected_reasons": reasons,
            "rejected_samples": [{"reason": r.get("reason"), "text": str(r.get("text", ""))[:MAX_REPLY_CHARS]} for r in (d.get("rejected") or [])[:4]],
            "kept": kept}


def lineage(status: str | None, child: str | None, parent: str | None) -> dict[str, Any]:
    if status in ("accepted", "accepted_unmeasured"):
        return {"parent": parent, "child": child, "outcome": "committed",
                "where": "registered in the run's disposable experiment store; not published"}
    if status == "rejected":
        return {"parent": parent, "child": None, "outcome": "pulled back", "where": "no child registered; the parent is unchanged"}
    return {"parent": parent, "child": child, "outcome": "unknown", "where": None}


def sleep_arm(arm: str, summary: dict[str, Any] | None, rep: dict[str, Any] | None, rep_path: Path | None,
              probes: dict[str, list]) -> dict[str, Any]:
    """One consolidation arm: the run-level summary (``sleep_controls.json``) plus its own report."""
    summary = summary or {}
    rep = rep or {}
    before = rep.get("before") or {}
    after = rep.get("after") or {}
    after_results = (after.get("recall") or {}).get("results") or summary.get("results")
    before_results = (before.get("recall") or {}).get("results")
    status = summary.get("status") or rep.get("status")
    child = summary.get("model_id", rep.get("model_id"))
    recorded_share = (after.get("recall") or {}).get("max_cluster_share")
    cfg = rep.get("config") or {}
    return {
        "arm": arm, "kind": "sleep", "method": cfg.get("method") or ("replay" if arm == "ungated" else arm),
        "target": cfg.get("target"), "status": status,
        "status_conflict": bool(summary.get("status") and rep.get("status") and summary["status"] != rep["status"]),
        "reason": summary.get("reason", rep.get("reason")),
        "lineage": lineage(status, child, rep.get("parent_model_id")),
        "gate": gate_view(summary.get("gate") or rep.get("gate"), after_results, recorded_share),
        "heldout_nll": {"before": {k: num(v) for k, v in (summary.get("heldout_nll_before") or before.get("heldout_nll") or {}).items()} or None,
                        "after": {k: num(v) for k, v in (summary.get("heldout_nll_after") or after.get("heldout_nll") or {}).items()} or None},
        "recall": {"before": counts_for(before_results, summary.get("by_group_before"), probes),
                   "after": counts_for(after_results, summary.get("by_group"), probes),
                   "totals_before": totals(before.get("recall")), "totals_after": totals(after.get("recall"))},
        "harvest": harvest_view(summary.get("harvest") or rep.get("harvest")),
        "batch": summary.get("batch") or rep.get("batch"), "packed": rep.get("packed"),
        "losses": [num(x, 5) for x in rep.get("losses") or []] or None,
        "loss_terms": (summary.get("batch") or rep.get("batch") or {}).get("loss_terms"),
        "dreams": dreams_view(summary.get("dreams") or rep.get("dreams")),
        "anchor_relative_update": anchor_update(rep),
        "w0_relative_update": w0_change(rep),     # every method, in reports written after aae9b70; absent before
        "gradient_norms": grad_norms_view(rep),   # per step, before clipping, in reports written after aae9b70
        "config": {k: cfg.get(k) for k in ("steps", "lr", "batch_size", "seq_len", "replay_ratio", "replay_rows", "heldout_rows", "anchor_lambda",
                                           "tolerance_nll", "tolerance_collapse", "session_loss", "prompt_loss_weight", "provenance",
                                           "flagged_policy", "dream_temperature", "dream_min_gain", "dream_token_weighting") if k in cfg} or None,
        "probes": probe_rows(after_results, probes),
        "report": ({"file": rel(rep_path), "sha256": sha256(rep_path), "run_id": rep.get("run_id"),
                    "head_at_arm_start": rep.get("code_commit"), "code_at_import": rep.get("code_commit_at_import"),
                    "seconds": num(rep.get("seconds"), 5)} if rep_path else None),
    }


def control_arm(arm: str, summary: dict[str, Any], probes: dict[str, list]) -> dict[str, Any]:
    results = summary.get("results")
    share = summary.get("max_cluster_share")
    return {"arm": arm, "kind": "control", "method": None, "status": None,
            "recall": {"after": counts_for(results, summary.get("by_group"), probes)},
            "mean_answer_logprob": num(summary.get("mean_answer_logprob")),
            "max_cluster_share": {"value": num(share), "source": "recorded"} if share is not None else (
                {"value": num(cluster_share(results)), "source": "computed from saved replies"} if results else {"value": None, "source": "absent"}),
            "probes": probe_rows(results, probes)}


# ---------------------------------------------------------------------------------------------------- runs
def protocol(c: dict[str, Any], arms: dict[str, Any]) -> dict[str, Any]:
    policies = sorted({(a.get("harvest") or {}).get("flagged_policy") for a in arms.values() if a.get("kind") == "sleep" and a.get("arm") != "ungated"} - {None})
    return {"facts": c.get("facts"), "poison": c.get("poison"), "seed": c.get("seed", 0 if c.get("started_at_unix") else None),
            "steps": c.get("steps"), "target": c.get("target"), "lr": c.get("lr"), "replay_ratio": c.get("replay_ratio"),
            "batch_size": c.get("batch_size"), "session_loss": c.get("session_loss"), "prompt_loss_weight": c.get("prompt_loss_weight"),
            "augment": c.get("augment"), "teach_temperature": c.get("teach_temperature"), "dream_temperature": c.get("dream_temperature"),
            "replay_revision": (c.get("replay_revision") or "")[:12] or None, "device": c.get("device"),
            "ceiling_mode": c.get("ceiling_mode", "all" if "ceiling" in c.get("arms", {}) else None),
            "flagged_policy": c.get("flagged_policy") or (policies[0] if len(policies) == 1 else None),
            "wall_seconds": num(c.get("seconds"), 5)}


def export_run(run_dir: Path, rows: list[dict[str, str]], probes: dict[str, list]) -> dict[str, Any]:
    controls_path = run_dir / "sleep_controls.json"
    c = read_json(controls_path)
    sources = [controls_path]
    arms: dict[str, Any] = {}
    for arm in ARM_ORDER:
        summary = c.get("arms", {}).get(arm)
        rep_path = run_dir / f"sleep_{arm}_report.json"
        if summary is None and not rep_path.exists():
            continue
        if arm in CONTROL_ARMS:
            arms[arm] = control_arm(arm, summary, probes)
            continue
        rep = read_json(rep_path) if rep_path.exists() else None
        if rep is not None:
            sources.append(rep_path)
        arms[arm] = sleep_arm(arm, summary, rep, rep_path if rep else None, probes)
    for md in sorted(run_dir.glob("*.md")):
        sources.append(md)
    report_times = [read_json(p).get("created_at_unix") for p in sources if p.name.endswith("_report.json")]
    started = c.get("started_at_unix") or min((t for t in report_times if t), default=None)
    row = match_row(rows, run_dir.name, [p.name for p in sources])
    question = plain(row["question"]) if row else None
    return {
        "id": run_dir.name, "question": question, "summary": first_sentence(question) if question else None,
        "started_at_unix": started, "checkpoint": checkpoint_identity(c.get("checkpoint"), c.get("checkpoint_digest")),
        "code": {"recorded_at_launch": c.get("code_commit"), "archive_readme": plain(row["code"]) if row else None},
        "protocol": protocol(c, arms), "arms": [arms[a] for a in ARM_ORDER if a in arms],
        "notes": {"recount": (run_dir / "sleep_controls_recounted.md").exists()},
        "sources": [{"file": rel(p), "sha256": sha256(p)} for p in sources],
    }


def export_dryrun(paths: list[Path], rows: list[dict[str, str]], probes: dict[str, list]) -> dict[str, Any]:
    """The base-model dry run has no sleep_controls.json: one report per method."""
    arms = []
    first = None
    for p in sorted(paths, key=lambda p: ARM_ORDER.index(read_json(p)["config"]["method"])):
        rep = read_json(p)
        first = first or rep
        arms.append(sleep_arm(rep["config"]["method"], None, rep, p, probes))
    cfg = first["config"]
    row = match_row(rows, "base_dryrun", [p.name for p in paths])
    question = plain(row["question"]) if row else None
    return {
        "id": "base_dryrun", "question": question, "summary": first_sentence(question) if question else None,
        "started_at_unix": min(read_json(p).get("created_at_unix") or 0 for p in paths),
        "checkpoint": {**checkpoint_identity("ttt_mlp_760m_base", None), "digest_prefix": None},
        "code": {"recorded_at_launch": None, "archive_readme": plain(row["code"]) if row else None},
        "protocol": {"facts": 3, "poison": False, "seed": cfg.get("seed"), "steps": cfg.get("steps"), "target": cfg.get("target"),
                     "lr": cfg.get("lr"), "replay_ratio": cfg.get("replay_ratio"), "batch_size": cfg.get("batch_size"),
                     "session_loss": cfg.get("session_loss"), "prompt_loss_weight": cfg.get("prompt_loss_weight"), "augment": None,
                     "teach_temperature": None, "dream_temperature": None, "replay_revision": None, "device": cfg.get("device"),
                     "ceiling_mode": None, "flagged_policy": cfg.get("flagged_policy"), "wall_seconds": None},
        "arms": arms, "notes": {"recount": False},
        "sources": [{"file": rel(p), "sha256": sha256(p)} for p in sorted(paths)],
    }


def chat_eval() -> dict[str, Any] | None:
    """The final checkpoint's chat evaluation: context for every step-250 run's parent."""
    path = CHAT_EVAL / "step250_t0.7.json"
    if not path.exists():
        return None
    d = read_json(path)
    return {"checkpoint_digest_prefix": str(d.get("checkpoint_digest", ""))[:16], "temperature": (d.get("sampling") or {}).get("temperature"),
            "distinct_reply_share": num(d.get("distinct_reply_share")), "mean_repetition_share": num(d.get("mean_repetition_share")),
            "held_out_nll": {k: ({kk: num(vv) for kk, vv in v.items()} if isinstance(v, dict) else num(v)) for k, v in (d.get("held_out_nll") or {}).items()},
            "source": {"file": rel(path), "sha256": sha256(path)}}


# ---------------------------------------------------------------------------------------------------- session trajectories
def _flags(reasons: list[str]) -> list[str]:
    return [r for r in reasons if r != "log_only"]


def extract_sessions(run_id: str, store: Path | None, archive_run: Path) -> dict[str, Any]:
    """Per-chunk trajectories of a run's teaching and rolled-back sessions. Read from the public results archive
    (``<run>/sessions/``) when it holds them; otherwise from a local experiment store, only when that store provably
    belongs to the archived run (identical sleep_controls.json). Both paths check the position and model signature.
    A local store's path is not recorded; the file digests are."""
    archived = archive_run / "sleep_controls.json"
    public = store is None
    if public:
        store = archive_run
        if not (store / "sessions").is_dir():
            raise ValueError(f"{run_id}: the archive holds no session files")
    else:
        store = store.resolve()
        local = store.parent / "sleep_controls.json"
        if not local.exists() or sha256(local) != sha256(archived):
            raise ValueError(f"{run_id}: the store's sleep_controls.json is not byte-identical to the archived run")
    c = read_json(archived)
    digest = c.get("checkpoint_digest") or ""
    harvest = next((a["harvest"] for a in c["arms"].values() if isinstance(a, dict) and a.get("harvest")), None)
    sessions = []
    for sid in ("teach", "rolled"):
        sdir = store / "sessions" / sid
        if not sdir.exists():
            continue
        meta = read_json(sdir / "meta.json")
        if not str(meta.get("model_signature", "")).endswith(digest):
            raise ValueError(f"{run_id}/{sid}: model signature does not match the archived checkpoint digest")
        tx_path, trace_path = sdir / "transactions.jsonl", sdir / "trace.jsonl"
        txs = [json.loads(line) for line in tx_path.read_text(encoding="utf-8").splitlines() if line.strip()]
        trace = [json.loads(line) for line in trace_path.read_text(encoding="utf-8").splitlines() if line.strip()] if trace_path.exists() else []
        if sid == "teach" and harvest is not None and meta.get("pos") != harvest.get("accepted_tokens"):
            raise ValueError(f"{run_id}/teach: position {meta.get('pos')} differs from the harvest's accepted tokens")
        chunks = []
        for t in txs:
            s = t.get("signals") or {}
            src = t.get("sources") or {}
            chunks.append({
                "i": t.get("index"), "pos": [t.get("pos_start"), t.get("pos_end")], "n": s.get("n_tokens"),
                "source": "user" if (src.get("user") or 0) >= (src.get("model") or 0) else "model",
                "chunk_loss": num(s.get("chunk_loss"), 4), "surprise_mean": num(s.get("surprise_mean"), 4),
                "surprise_max": num(s.get("surprise_max"), 4), "write_norm_sum": num(s.get("write_norm_sum"), 4),
                "proposed": num(s.get("delta_norm"), 4), "accepted": num((t.get("accepted") or {}).get("delta_norm"), 4),
                "per_layer": [num(v, 3) for v in s.get("delta_norm_per_layer") or []] or None,
                "requested": (t.get("requested") or {}).get("kind"), "applied": (t.get("decision") or {}).get("kind"),
                "scale": num((t.get("decision") or {}).get("scale"), 4),
                "flags": _flags((t.get("requested") or {}).get("reasons") or []),
                "applied_reasons": _flags((t.get("decision") or {}).get("reasons") or []),
                "read_only": bool(t.get("read_only")),
            })
        h = meta.get("harness") or {}
        sessions.append({
            "session_id": sid, "log_only": bool(h.get("log_only")), "enable_rollback": h.get("enable_rollback"),
            "z_rollback": h.get("z_rollback"), "z_scale": h.get("z_scale"), "calibrated": False if not meta.get("extra", {}).get("calibration") else True,
            "counts": {k: meta.get(k) for k in ("commits", "rollbacks", "scales", "projects", "readonly", "n_transactions", "pos")},
            "turns": [{"i": i, "prompt": tr.get("prompt"), "completion": str(tr.get("completion", ""))[:MAX_REPLY_CHARS],
                       "tx": [tr.get("tx_start"), tr.get("tx_end")]} for i, tr in enumerate(trace)],
            "chunks": chunks,
            "source": {"transactions": rel(tx_path) if public else None, "transactions_sha256": sha256(tx_path),
                       "trace_sha256": sha256(trace_path) if trace_path.exists() else None},
        })
    n_leaves = max((len(ch["per_layer"] or []) for s in sessions for ch in s["chunks"]), default=0)
    return {"run_id": run_id, "schema": SCHEMA,
            "provenance": {"kind": "public results archive" if public else "local experiment store; these session files are not in the public results archive",
                           "controls_sha256": sha256(archived),
                           "checks": ([] if public else ["sleep_controls.json byte-identical to the archive"])
                           + ["teach position equals the harvest's accepted tokens", "model signature matches the checkpoint digest"]},
            "per_layer": {"quantity": "proposed fast-weight change per tensor in effective coordinates (norm of working minus committed)",
                          "layers": n_leaves // len(TENSORS), "tensors": list(TENSORS), "order": "layer-major: layer 0 W1, b1, W2, b2, layer 1 ..."},
            "sessions": sessions}


# ---------------------------------------------------------------------------------------------------- driver
def dump(obj: Any) -> str:
    return json.dumps(obj, ensure_ascii=False, separators=(",", ":"), sort_keys=False) + "\n"


def build(out: Path, stores: dict[str, Path]) -> dict[str, Any]:
    probes = build_probes()
    rows = readme_rows(ARCHIVE / "README.md")
    runs = [export_run(d, rows, probes) for d in sorted(ARCHIVE.iterdir()) if d.is_dir() and (d / "sleep_controls.json").exists()]
    dry = sorted(ARCHIVE.glob("base_dryrun_*_report.json"))
    if dry:
        runs.append(export_dryrun(dry, rows, probes))
    runs.sort(key=lambda r: (r["started_at_unix"] or 0, r["id"]))

    (out / "runs").mkdir(parents=True, exist_ok=True)
    (out / "sessions").mkdir(parents=True, exist_ok=True)
    archived_sessions = {r["id"] for r in runs if (ARCHIVE / r["id"] / "sessions").is_dir()}
    for run_id in sorted(archived_sessions | set(stores)):
        source = None if run_id in archived_sessions else stores[run_id]  # the public archive wins over a local store
        (out / "sessions" / f"{run_id}.json").write_text(dump(extract_sessions(run_id, source, ARCHIVE / run_id)), encoding="utf-8")
    session_ids = sorted(p.stem for p in (out / "sessions").glob("*.json"))
    run_ids = {r["id"] for r in runs}
    for stale in (out / "runs").glob("*.json"):
        if stale.stem not in run_ids:
            stale.unlink()

    for r in runs:
        r["session_trajectory"] = f"sessions/{r['id']}.json" if r["id"] in session_ids else None
        (out / "runs" / f"{r['id']}.json").write_text(dump(r), encoding="utf-8")

    source_digests = sorted({(s["file"], s["sha256"]) for r in runs for s in r["sources"]})
    archive_digest = hashlib.sha256("".join(f"{f}:{d}\n" for f, d in source_digests).encode()).hexdigest()
    index = {
        "schema": SCHEMA, "exporter_version": EXPORTER_VERSION, "archive_digest": archive_digest,
        "archive": rel(ARCHIVE), "chat_eval": chat_eval(),
        "runs": [{
            "id": r["id"], "summary": r["summary"], "started_at_unix": r["started_at_unix"], "checkpoint": r["checkpoint"],
            "code_recorded": r["code"]["recorded_at_launch"], "facts": r["protocol"]["facts"], "flagged_policy": r["protocol"]["flagged_policy"],
            "steps": r["protocol"]["steps"], "target": r["protocol"]["target"], "session_trajectory": r["session_trajectory"],
            "arms": [{"arm": a["arm"], "status": a["status"], "method": a["method"],
                      "taught": ((a["recall"]["after"]["by_group"] or {}).get("taught") if a["recall"]["after"]["by_group"] else None),
                      "totals": a["recall"].get("totals_after") if a["kind"] == "sleep" else None,
                      "nll_before": ((a.get("heldout_nll") or {}).get("before") or {}).get("mean"),
                      "nll_after": ((a.get("heldout_nll") or {}).get("after") or {}).get("mean")} for a in r["arms"]],
        } for r in runs],
        "sessions": session_ids,
    }
    (out / "index.json").write_text(dump(index), encoding="utf-8")
    return index


def mirror(src: Path, dst: Path) -> None:
    if dst.exists():
        shutil.rmtree(dst)
    shutil.copytree(src, dst)


def tree_bytes(root: Path) -> dict[str, bytes]:
    return {p.relative_to(root).as_posix(): p.read_bytes() for p in sorted(root.rglob("*")) if p.is_file()}


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--out", type=Path, default=OUT, help="export directory (default: docs/research/results/sleep-observatory)")
    ap.add_argument("--mirror", type=Path, default=MIRROR, help="dashboard copy served under /assets (default: dashboard/public/assets/observatory)")
    ap.add_argument("--no-mirror", action="store_true", help="write the export only")
    ap.add_argument("--session-store", action="append", default=[], metavar="RUN=STORE",
                    help="extract a run's teach/rolled session trajectories from a local experiment store (repeatable)")
    ap.add_argument("--check", action="store_true", help="rebuild into a temporary copy and exit 1 if the committed export differs")
    args = ap.parse_args(argv)
    stores = {}
    for item in args.session_store:
        run_id, sep, path = item.partition("=")
        if not sep or not (ARCHIVE / run_id / "sleep_controls.json").exists():
            ap.error(f"--session-store expects RUN=STORE with RUN an archived run directory: {item}")
        stores[run_id] = Path(path)
    if args.check:
        with tempfile.TemporaryDirectory() as tmp:
            work = Path(tmp) / "export"
            if args.out.exists():
                shutil.copytree(args.out, work)
            build(work, {})
            fresh, committed = tree_bytes(work), tree_bytes(args.out) if args.out.exists() else {}
            mirrored = tree_bytes(args.mirror) if args.mirror.exists() else {}
            stale = sorted(k for k in set(fresh) | set(committed) if fresh.get(k) != committed.get(k))
            drift = sorted(k for k in set(committed) | set(mirrored) if committed.get(k) != mirrored.get(k))
            for k in stale:
                print(f"stale: {k}")
            for k in drift:
                print(f"mirror differs: {k}")
            return 1 if stale or drift else 0
    index = build(args.out, stores)
    if not args.no_mirror:
        mirror(args.out, args.mirror)
    size = sum(len(b) for b in tree_bytes(args.out).values())
    print(f"exported {len(index['runs'])} runs, {len(index['sessions'])} session trajectories, {size / 1e6:.2f} MB -> {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())