#!/usr/bin/env python3
"""Figure 1: system map of the experiment (donor → localization → map → graft → evaluation), with the models,
backends and hosts actually used (read from run manifests). Vector PDF; no data values plotted."""
from __future__ import annotations

import json
import sys
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from wp_ai_0001.visualization.style import CATEGORICAL, DOUBLE_WIDTH_IN, ENTITY_COLORS, GRID, TEXT_PRIMARY, TEXT_SECONDARY, apply_style  # noqa: E402

runs = [json.loads(l) for l in (ROOT / "experiments" / "registry" / "runs.jsonl").read_text().splitlines() if l.strip()]
n_completed = len({r["run_id"] for r in runs if r["status"] == "completed"})
n_invalid = len({r["run_id"] for r in runs if r["status"] == "invalidated"})
man = ROOT / "experiments" / "manifests"
hosts = set()
for f in man.glob("*.json"):
    m = json.loads(f.read_text())
    hosts.add(m["hardware"].get("chip", "?"))

apply_style()
fig, ax = plt.subplots(figsize=(DOUBLE_WIDTH_IN, 3.6))
ax.set_xlim(0, 100)
ax.set_ylim(0, 60)
ax.axis("off")


def box(x, y, w, h, title, lines, color, fc="#ffffff"):
    ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle="round,pad=0.6,rounding_size=1.5", linewidth=1.0, edgecolor=color, facecolor=fc))
    ax.text(x + 1.0, y + h - 1.4, title, fontsize=6.2, fontweight="bold", color=color, va="top")
    ax.text(x + 1.0, y + h - 5.0, "\n".join(lines), fontsize=4.6, color=TEXT_PRIMARY, va="top", linespacing=1.3)


def arrow(x0, y0, x1, y1, label=None, color=TEXT_SECONDARY):
    ax.add_patch(FancyArrowPatch((x0, y0), (x1, y1), arrowstyle="-|>", mutation_scale=9, linewidth=0.9, color=color))
    if label:
        ax.text((x0 + x1) / 2, (y0 + y1) / 2 + 2.2, label, fontsize=4.4, color=TEXT_SECONDARY, ha="center")


# Tier-0 chain (top row)
box(1, 36, 17.5, 22, "Donor D", ["toy transformer 4 L, d=128", "all 8 operators (incl. Ω)", "3 seeds; Tier 1–2: 1.5B–7B"], ENTITY_COLORS["donor"])
box(21.5, 36, 17.5, 22, "Causal localization", ["interchange patching", "zero-ablation, heads", "ablation → rescue", "organ: post-L0 row, read pos."], CATEGORICAL[6])
box(42, 36, 17.5, 22, "Alignment map A", ["Procrustes / ridge / CCA", "fitted on task-free rows", "sham: random orthogonal", "Prop. 1: regression kills Ω"], CATEGORICAL[3])
box(62.5, 36, 17.5, 22, "Graft into R", ["activation (donor in loop)", "subspace k, steering,", "adapter, weights, persistent", "controls: random/shuffled/unrelated"], ENTITY_COLORS["graft"])
box(83, 36, 16, 22, "Evaluation", ["held-out target / control", "Δ_T^adj vs sham", "cluster bootstrap", "preregistered freeze"], ENTITY_COLORS["recipient"])
arrow(18.5, 39, 21.5, 39, "activations")
arrow(39, 39, 42, 39, "component")
arrow(59.5, 39, 62.5, 39, "A(S_D)")
arrow(80, 39, 83, 39, "R′")
# Tier 0 vs Tier 1 lanes (bottom row)
box(1, 6, 46, 24, "Tier 0 — micro-world (PyTorch, MPS)", ["OPSEQ: invented operators on Z_11, scratchpad trace", "recipient trained without Ω (manufactured gap)", "17 models (matrix_v1); runs: pilot, v1, v2, confirmatory", "fresh-process persistence; supervised / distillation references"], ENTITY_COLORS["donor"], fc="#f7f9fc")
box(52, 6, 47, 24, "Tier 1–2 — open-weight models (MLX, wrapped layers)", ["Qwen2.5-1.5B/3B (laptop), Llama-3.2-3B, Qwen2.5-7B (M3U96b)", "capability: zero-shot task (antonym, capital, morph, add-three)", "organ: few-shot − zero-shot last-token state; task-free map", "dev split selects, confirmatory split tests (frozen rule, A1–A2)"], CATEGORICAL[1], fc="#fdf8f5")
ax.text(1, 2.0, f"Registry: {n_completed} completed runs, {n_invalid} invalidated; hosts: {', '.join(sorted(h for h in hosts if h != '?'))}. Hooks: bitwise no-op when idle (PyTorch and MLX).", fontsize=4.8, color=TEXT_SECONDARY)
fig.savefig(ROOT / "paper" / "figures" / "fig_system_map.pdf")
print("written fig_system_map.pdf")