#!/usr/bin/env python3
"""Empirical Muon/BBP campaign for the three-layer quadratic ResNet.

The experiment deliberately separates three objects:

1. the exact population trajectory in the finite teacher subspace;
2. fixed-sample empirical training, evaluated on exact Gauss--Hermite test risk;
3. a fresh-sample ambient Hessian lift.  Its orthogonal block is a block-Wishart
   matrix, so the variational edges of Montanari--Saeed apply conditionally.

The Muon powers 1/7 and 1/3 are transfer controls from the power-law
phase-retrieval project.  They are not asserted to be optimal for this model.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Callable

import numpy as np
import torch

os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib-codex-hierarchical3")
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt


WORKSPACE_ROOT = Path(__file__).resolve().parents[3]
if str(WORKSPACE_ROOT) not in sys.path:
    sys.path.insert(0, str(WORKSPACE_ROOT))

from experiments.brainstorm.resnetl3.resnet_order4_full_verify_v2 import (
    ReducedQuadraticResNet,
    build_exact_setup,
    eval_multiindices,
    exact_risk_from_coeff,
    order_errors,
    project_coeffs,
    reduced_alignment_rows,
    seed_all,
)
from experiments.muon.lowrank_block_wishart_spectral_experiments import (
    block_wishart,
    spectral_clip_blocks,
    variational_edges,
)


ROOT = Path(__file__).resolve().parent
DEFAULT_OUT = ROOT / "results/hierarchical3_muon_bbp_latest"


@dataclass
class Config:
    outdir: str = str(DEFAULT_OUT)
    seed: int = 314159
    k: int = 3
    p: int = 3
    q: int = 3
    nq: int = 7
    steps: int = 1200
    diag_every: int = 40
    repeats: int = 3
    train_n: int = 1024
    init_scale: float = 0.12
    readout_scale: float = 0.05
    gd_lr: float = 0.02
    muon_lr: float = 0.004
    vector_lr: float = 0.02
    spectral_eps: float = 1e-6
    tie_u_v: bool = False
    d: int = 32
    alpha: float = 5.0
    bbp_checkpoints: int = 4
    block_cap: float = 100.0
    edge_scan_points: int = 56
    mde_iterations: int = 240
    mde_damping: float = 0.35
    mde_tolerance: float = 3e-9
    density_eta: float = 0.05
    density_points: int = 81


@dataclass(frozen=True)
class OptimizerSpec:
    name: str
    power: float | None


OPTIMIZERS = (
    OptimizerSpec("gd", None),
    OptimizerSpec("muon_a0", 0.0),
    OptimizerSpec("muon_a1_7_transfer", 1.0 / 7.0),
    OptimizerSpec("muon_a1_3_normalized", 1.0 / 3.0),
)


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if not rows:
        path.write_text("", encoding="utf-8")
        return
    fields: list[str] = []
    for row in rows:
        for key in row:
            if key not in fields:
                fields.append(key)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def json_ready(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(k): json_ready(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [json_ready(v) for v in value]
    if isinstance(value, np.ndarray):
        return value.tolist()
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, torch.Tensor):
        return value.detach().cpu().tolist()
    if isinstance(value, float) and not math.isfinite(value):
        return None
    return value


def make_setup(cfg: Config):
    # The amplitudes decrease across teacher coordinates and mixed terms force
    # a genuinely hierarchical degree-2 -> degree-3 -> degree-4 problem.
    return build_exact_setup(
        k=cfg.k,
        nq=cfg.nq,
        stage="234",
        mu_max=1.0,
        mu_min=0.35,
        q2mix=0.18,
        c23=0.18,
        q3diag=0.04,
        q3111=0.04,
        q4diag=0.18,
        q422=0.12,
        q431=0.035,
        q4211=0.035,
    )


def new_model(cfg: Config, setup, seed: int) -> ReducedQuadraticResNet:
    seed_all(seed)
    return ReducedQuadraticResNet(
        setup.k,
        cfg.p,
        cfg.q,
        init_scale=cfg.init_scale,
        dtype=setup.z.dtype,
        tie_u_v=cfg.tie_u_v,
        alpha_scale=cfg.readout_scale,
        beta_scale=cfg.readout_scale,
        a0_scale=0.02,
    )


def clone_state(model: ReducedQuadraticResNet) -> dict[str, torch.Tensor]:
    return {name: value.detach().clone() for name, value in model.state_dict().items()}


def teacher_on(z: torch.Tensor, setup) -> torch.Tensor:
    return eval_multiindices(z, setup.basis) @ setup.teacher_coeff


def population_objective(model: ReducedQuadraticResNet, setup) -> torch.Tensor:
    coeff = project_coeffs(model(setup.z), setup)
    return exact_risk_from_coeff(coeff, setup)


def empirical_objective(
    model: ReducedQuadraticResNet, z: torch.Tensor, y: torch.Tensor
) -> torch.Tensor:
    return 0.5 * torch.mean((model(z) - y) ** 2)


def exact_svd_power(gradient: torch.Tensor, power: float, eps: float) -> torch.Tensor:
    """Muon-a update with the practical unit-RMS normalization.

    For power zero this is the exact polar/SignSVD update.  Dividing singular
    values by their mean makes the power dimensionless; the final RMS scaling
    matches the convention used by the repository's nanoGPT Muon runner.
    """

    work = gradient.float()
    transposed = False
    if work.shape[0] > work.shape[1]:
        work = work.T
        transposed = True
    u, singular, vh = torch.linalg.svd(work, full_matrices=False)
    scale = singular.mean().clamp_min(eps)
    weights = (singular / scale).clamp_min(eps).pow(power)
    update = (u * weights.unsqueeze(0)) @ vh
    if transposed:
        update = update.T
    update = update.to(dtype=gradient.dtype)
    return update / update.square().mean().sqrt().clamp_min(1e-12)


def proposed_updates(
    model: ReducedQuadraticResNet, spec: OptimizerSpec, cfg: Config
) -> dict[str, tuple[torch.Tensor, float]]:
    updates: dict[str, tuple[torch.Tensor, float]] = {}
    for name, parameter in model.named_parameters():
        if parameter.grad is None:
            continue
        gradient = parameter.grad.detach()
        if spec.power is not None and gradient.ndim == 2 and name in {"V", "U", "A"}:
            updates[name] = (
                exact_svd_power(gradient, spec.power, cfg.spectral_eps),
                cfg.muon_lr,
            )
        else:
            lr = cfg.gd_lr if spec.power is None else cfg.vector_lr
            updates[name] = (gradient.clone(), lr)
    return updates


def stable_step(
    model: ReducedQuadraticResNet,
    objective: Callable[[], torch.Tensor],
    spec: OptimizerSpec,
    cfg: Config,
    lr_scale: float,
) -> tuple[float, float, int]:
    model.zero_grad(set_to_none=True)
    old_loss_t = objective()
    old_loss = float(old_loss_t.detach())
    old_loss_t.backward()
    updates = proposed_updates(model, spec, cfg)
    state = clone_state(model)
    accepted_scale = lr_scale
    rejects = 0
    for _ in range(10):
        model.load_state_dict(state)
        with torch.no_grad():
            named = dict(model.named_parameters())
            for name, (update, base_lr) in updates.items():
                named[name].add_(update, alpha=-base_lr * accepted_scale)
        candidate = float(objective().detach())
        if math.isfinite(candidate) and candidate <= old_loss * (1.0 + 1e-8) + 1e-12:
            return candidate, min(1.0, accepted_scale * 1.01), rejects
        accepted_scale *= 0.5
        rejects += 1
    model.load_state_dict(state)
    return old_loss, max(accepted_scale, 2.0**-16), rejects


def spectral_stats(matrix: torch.Tensor) -> tuple[float, float, float]:
    singular = torch.linalg.svdvals(matrix.detach()).cpu().numpy()
    if singular.size == 0:
        return float("nan"), float("nan"), 0.0
    total = float(np.sum(singular))
    probs = singular / max(total, 1e-30)
    entropy = -float(np.sum(probs * np.log(np.maximum(probs, 1e-30))))
    return float(singular[0]), float(singular[-1]), float(np.exp(entropy))


def diagnostics(
    model: ReducedQuadraticResNet,
    setup,
    objective: Callable[[], torch.Tensor],
    regime: str,
    spec: OptimizerSpec,
    repeat: int,
    step: int,
    train_loss: float,
    lr_scale: float,
    rejects: int,
) -> dict[str, Any]:
    with torch.no_grad():
        coeff = project_coeffs(model(setup.z), setup)
        test_risk = float(exact_risk_from_coeff(coeff, setup))
        errors = order_errors(coeff, setup)
        align_v = reduced_alignment_rows(model.V)
        align_u = reduced_alignment_rows(model.second_weight())

    model.zero_grad(set_to_none=True)
    objective().backward()
    named = dict(model.named_parameters())
    row: dict[str, Any] = {
        "regime": regime,
        "optimizer": spec.name,
        "repeat": repeat,
        "step": step,
        "train_loss": train_loss,
        "test_population_risk": test_risk,
        "lr_scale": lr_scale,
        "backtracking_rejects": rejects,
    }
    row.update(errors)
    for coord in range(setup.k):
        row[f"align_V_{coord + 1}"] = float(np.max(align_v[:, coord]))
        row[f"align_U_{coord + 1}"] = float(np.max(align_u[:, coord]))
    for name in ("V", "U", "A"):
        if name not in named or named[name].grad is None:
            continue
        smax, smin, erank = spectral_stats(named[name].grad)
        row[f"grad_{name}_smax"] = smax
        row[f"grad_{name}_smin"] = smin
        row[f"grad_{name}_effective_rank"] = erank
    model.zero_grad(set_to_none=True)
    return row


def run_trajectory(
    cfg: Config,
    setup,
    spec: OptimizerSpec,
    regime: str,
    repeat: int,
    init_state: dict[str, torch.Tensor],
    z_train: torch.Tensor | None,
    y_train: torch.Tensor | None,
) -> tuple[list[dict[str, Any]], dict[int, dict[str, torch.Tensor]]]:
    model = new_model(cfg, setup, cfg.seed + 1000 * repeat)
    model.load_state_dict(init_state)
    if regime == "population":
        objective = lambda: population_objective(model, setup)
    else:
        assert z_train is not None and y_train is not None
        objective = lambda: empirical_objective(model, z_train, y_train)

    checkpoint_grid = set(
        int(v)
        for v in np.unique(
            np.rint(np.linspace(0, cfg.steps, cfg.bbp_checkpoints)).astype(int)
        )
    )
    rows: list[dict[str, Any]] = []
    snapshots: dict[int, dict[str, torch.Tensor]] = {}
    lr_scale = 1.0
    rejects_since_diag = 0
    train_loss = float(objective().detach())
    for step in range(cfg.steps + 1):
        record = step % cfg.diag_every == 0 or step == cfg.steps or step in checkpoint_grid
        if record:
            rows.append(
                diagnostics(
                    model,
                    setup,
                    objective,
                    regime,
                    spec,
                    repeat,
                    step,
                    train_loss,
                    lr_scale,
                    rejects_since_diag,
                )
            )
            rejects_since_diag = 0
        if step in checkpoint_grid:
            snapshots[step] = clone_state(model)
        if step == cfg.steps:
            break
        train_loss, lr_scale, rejects = stable_step(
            model, objective, spec, cfg, lr_scale
        )
        rejects_since_diag += rejects
    return rows, snapshots


def link_hessian_blocks(
    model: ReducedQuadraticResNet,
    z: np.ndarray,
    target: np.ndarray,
) -> np.ndarray:
    """Per-example Hessian in the direct preactivations (V z, U z)."""

    v = model.V.detach().cpu().numpy()
    u = model.second_weight().detach().cpu().numpy()
    a = model.A.detach().cpu().numpy()
    b = model.b.detach().cpu().numpy()
    c = model.c.detach().cpu().numpy()
    alpha = model.alpha.detach().cpu().numpy()
    beta = model.beta.detach().cpu().numpy()
    a0 = float(model.a0.detach())

    h1 = z @ v.T + b[None, :]
    qfeat = h1 * h1 - 1.0
    hu = z @ u.T
    t2 = hu + qfeat @ a.T + c[None, :]
    sfeat = t2 * t2 - 1.0
    pred = a0 + qfeat @ alpha + sfeat @ beta
    residual = pred - target

    cq = alpha[None, :] + 2.0 * (t2 * beta[None, :]) @ a
    grad1 = 2.0 * h1 * cq
    grad2 = 2.0 * t2 * beta[None, :]
    grad = np.concatenate([grad1, grad2], axis=1)

    n = len(z)
    r = model.p + model.q
    hfun = np.zeros((n, r, r), dtype=float)
    weighted_ata = a.T @ (beta[:, None] * a)
    for index in range(n):
        d1 = np.diag(h1[index])
        h11 = 2.0 * np.diag(cq[index]) + 8.0 * d1 @ weighted_ata @ d1
        h12 = 4.0 * h1[index, :, None] * (a.T * beta[None, :])
        h22 = 2.0 * np.diag(beta)
        hfun[index, : model.p, : model.p] = h11
        hfun[index, : model.p, model.p :] = h12
        hfun[index, model.p :, : model.p] = h12.T
        hfun[index, model.p :, model.p :] = h22
    return grad[:, :, None] * grad[:, None, :] + residual[:, None, None] * hfun


def bbp_state(
    cfg: Config,
    setup,
    spec: OptimizerSpec,
    step: int,
    state: dict[str, torch.Tensor],
) -> dict[str, Any]:
    model = new_model(cfg, setup, cfg.seed)
    model.load_state_dict(state)
    # Common random numbers across optimizers: at step zero all methods have the
    # same state and must therefore return the same spectral diagnostic.
    rng = np.random.default_rng(cfg.seed + 100_003 + step)
    d_perp = cfg.d - cfg.k
    if d_perp <= 0:
        raise ValueError("Ambient dimension d must exceed teacher dimension k.")
    n = int(round(cfg.alpha * d_perp))
    z = rng.standard_normal((n, cfg.k))
    xi = rng.standard_normal((n, d_perp))
    x = np.concatenate([z, xi], axis=1)
    z_t = torch.tensor(z, dtype=setup.z.dtype)
    target = teacher_on(z_t, setup).detach().cpu().numpy()
    blocks_raw = link_hessian_blocks(model, z, target)
    blocks, truncation = spectral_clip_blocks(blocks_raw, cfg.block_cap)

    full_raw = block_wishart(x, blocks_raw)
    bulk_raw = block_wishart(xi, blocks_raw)
    full = block_wishart(x, blocks)
    bulk = block_wishart(xi, blocks)
    eig_full_raw, vec_full_raw = np.linalg.eigh(full_raw)
    eig_bulk_raw = np.linalg.eigvalsh(bulk_raw)
    eig_full, vec_full = np.linalg.eigh(full)
    eig_bulk = np.linalg.eigvalsh(bulk)
    alpha_perp = n / d_perp
    edge_error = ""
    try:
        left, right, left_diag, right_diag = variational_edges(
            blocks, alpha_perp, eig_bulk, cfg
        )
    except Exception as exc:  # retain empirical evidence if edge solver is singular
        left, right = float(np.min(eig_bulk)), float(np.max(eig_bulk))
        left_diag = {"derivative": float("nan"), "residual": float("nan")}
        right_diag = {"derivative": float("nan"), "residual": float("nan")}
        edge_error = f"{type(exc).__name__}: {exc}"

    r = cfg.p + cfg.q
    reshaped = vec_full.reshape(r, cfg.d, -1)
    teacher_overlap = np.sum(reshaped[:, : cfg.k, :] ** 2, axis=(0, 1))
    reshaped_raw = vec_full_raw.reshape(r, cfg.d, -1)
    teacher_overlap_raw = np.sum(reshaped_raw[:, : cfg.k, :] ** 2, axis=(0, 1))
    span = max(1e-9, right - left)
    tolerance = 0.015 * max(1.0, span)
    outlier = (eig_full < left - tolerance) | (eig_full > right + tolerance)
    raw_left = float(eig_bulk_raw[0])
    raw_right = float(eig_bulk_raw[-1])
    raw_span = max(1e-9, raw_right - raw_left)
    raw_tolerance = 0.015 * max(1.0, raw_span)
    outlier_raw = (eig_full_raw < raw_left - raw_tolerance) | (
        eig_full_raw > raw_right + raw_tolerance
    )
    informative_threshold = max(0.20, 2.5 * cfg.k / cfg.d)
    informative = outlier & (teacher_overlap >= informative_threshold)
    informative_raw = outlier_raw & (teacher_overlap_raw >= informative_threshold)
    with torch.no_grad():
        pop_risk = float(population_objective(model, setup))
    return {
        "optimizer": spec.name,
        "step": step,
        "population_risk": pop_risk,
        "n": n,
        "d": cfg.d,
        "block_size": r,
        "alpha_perp": alpha_perp,
        "block_truncation_fraction": truncation,
        "truncation_operator_norm": float(np.linalg.norm(full_raw - full, ord=2)),
        "edge_left_variational": left,
        "edge_right_variational": right,
        "edge_left_empirical_bulk": float(np.min(eig_bulk)),
        "edge_right_empirical_bulk": float(np.max(eig_bulk)),
        "edge_left_abs_error": abs(left - float(np.min(eig_bulk))),
        "edge_right_abs_error": abs(right - float(np.max(eig_bulk))),
        "left_edge_derivative": float(left_diag.get("derivative", float("nan"))),
        "right_edge_derivative": float(right_diag.get("derivative", float("nan"))),
        "edge_solver_error": edge_error,
        "outlier_count": int(np.sum(outlier)),
        "informative_outlier_count": int(np.sum(informative)),
        "max_outlier_teacher_overlap": (
            float(np.max(teacher_overlap[outlier])) if np.any(outlier) else 0.0
        ),
        "max_teacher_overlap": float(np.max(teacher_overlap)),
        "min_full_eigenvalue": float(eig_full[0]),
        "max_full_eigenvalue": float(eig_full[-1]),
        "raw_edge_left_empirical_bulk": raw_left,
        "raw_edge_right_empirical_bulk": raw_right,
        "raw_outlier_count": int(np.sum(outlier_raw)),
        "raw_informative_outlier_count": int(np.sum(informative_raw)),
        "raw_max_outlier_teacher_overlap": (
            float(np.max(teacher_overlap_raw[outlier_raw]))
            if np.any(outlier_raw)
            else 0.0
        ),
        "raw_min_full_eigenvalue": float(eig_full_raw[0]),
        "raw_max_full_eigenvalue": float(eig_full_raw[-1]),
    }


def aggregate_training(rows: list[dict[str, Any]], cfg: Config) -> dict[str, Any]:
    empirical = [row for row in rows if row["regime"] == "empirical"]
    summary: dict[str, Any] = {}
    for spec in OPTIMIZERS:
        final = [
            row
            for row in empirical
            if row["optimizer"] == spec.name and row["step"] == max(r["step"] for r in empirical)
        ]
        values = np.asarray([row["test_population_risk"] for row in final], dtype=float)
        summary[spec.name] = {
            "final_test_risk_mean": float(np.mean(values)),
            "final_test_risk_std": float(np.std(values, ddof=1)) if len(values) > 1 else 0.0,
            "final_test_risk_sem": float(np.std(values, ddof=1) / math.sqrt(len(values))) if len(values) > 1 else 0.0,
            "repeats": len(values),
            "final_test_risk_median": float(np.median(values)),
            "final_test_risk_geometric_mean": float(
                math.exp(np.mean(np.log(np.maximum(values, 1e-300))))
            ),
            "final_test_risk_by_repeat": values.tolist(),
            "total_backtracking_rejects": int(
                sum(row["backtracking_rejects"] for row in empirical if row["optimizer"] == spec.name)
            ),
        }
        gd_final = {
            int(row["repeat"]): float(row["test_population_risk"])
            for row in empirical
            if row["optimizer"] == "gd"
            and row["step"] == max(r["step"] for r in empirical)
        }
        paired_ratios = np.asarray(
            [
                float(row["test_population_risk"]) / gd_final[int(row["repeat"])]
                for row in final
            ],
            dtype=float,
        )
        summary[spec.name]["paired_risk_ratio_vs_gd"] = paired_ratios.tolist()
        summary[spec.name]["wins_vs_gd"] = int(np.sum(paired_ratios < 1.0))
        base = [
            row
            for row in empirical
            if row["optimizer"] == spec.name and row["repeat"] == 0 and row["step"] == 0
        ][0]
        trajectory = sorted(
            (
                row
                for row in empirical
                if row["optimizer"] == spec.name and row["repeat"] == 0
            ),
            key=lambda row: row["step"],
        )
        for degree in (2, 3, 4):
            key = f"E_order_{degree}"
            threshold = 0.5 * float(base[key])
            crossing = next((row["step"] for row in trajectory if row[key] <= threshold), None)
            summary[spec.name][f"order_{degree}_half_error_step"] = crossing
        for coord in range(1, cfg.k + 1):
            key = f"align_V_{coord}"
            crossing = next((row["step"] for row in trajectory if row[key] >= 0.8), None)
            summary[spec.name][f"V_coord_{coord}_align80_step"] = crossing
    return summary


def plot_training(rows: list[dict[str, Any]], outdir: Path) -> None:
    empirical = [row for row in rows if row["regime"] == "empirical"]
    fig, axes = plt.subplots(1, 2, figsize=(12.0, 4.8), dpi=190)
    colors = dict(zip((s.name for s in OPTIMIZERS), plt.cm.viridis(np.linspace(0.05, 0.92, len(OPTIMIZERS)))))
    for spec in OPTIMIZERS:
        mode_rows = [row for row in empirical if row["optimizer"] == spec.name]
        steps = sorted({int(row["step"]) for row in mode_rows})
        medians, minima, maxima = [], [], []
        for step in steps:
            values = np.asarray(
                [row["test_population_risk"] for row in mode_rows if row["step"] == step]
            )
            medians.append(float(np.median(values)))
            minima.append(float(np.min(values)))
            maxima.append(float(np.max(values)))
        for repeat in sorted({int(row["repeat"]) for row in mode_rows}):
            individual = sorted(
                (row for row in mode_rows if int(row["repeat"]) == repeat),
                key=lambda row: int(row["step"]),
            )
            axes[0].plot(
                [int(row["step"]) for row in individual],
                [row["test_population_risk"] for row in individual],
                color=colors[spec.name],
                lw=0.7,
                alpha=0.22,
            )
        axes[0].plot(steps, medians, color=colors[spec.name], lw=2.0, label=spec.name)
        axes[0].fill_between(steps, minima, maxima, color=colors[spec.name], alpha=0.10)
        first = sorted(
            (row for row in mode_rows if row["repeat"] == 0),
            key=lambda row: row["step"],
        )
        for degree, style in zip((2, 3, 4), ("-", "--", ":")):
            axes[1].plot(
                [row["step"] for row in first],
                [row[f"E_order_{degree}"] for row in first],
                color=colors[spec.name],
                ls=style,
                lw=1.3,
                label=f"{spec.name}, order {degree}",
            )
    axes[0].set_title("Finite-sample training, exact population test risk")
    axes[0].set_ylabel("risk")
    axes[1].set_title("Sequential Hermite-sector recovery (repeat 0)")
    axes[1].set_ylabel("weighted coefficient error")
    for ax in axes:
        ax.set_xlabel("gradient steps")
        ax.set_yscale("log")
        ax.grid(alpha=0.20)
        ax.legend(fontsize=7, ncol=2)
    fig.tight_layout()
    fig.savefig(outdir / "training_and_sequential_recovery.png", bbox_inches="tight")
    plt.close(fig)


def plot_bbp(rows: list[dict[str, Any]], outdir: Path) -> None:
    fig, axes = plt.subplots(1, 2, figsize=(12.0, 4.8), dpi=190)
    colors = dict(zip((s.name for s in OPTIMIZERS), plt.cm.plasma(np.linspace(0.05, 0.90, len(OPTIMIZERS)))))
    for spec in OPTIMIZERS:
        selected = sorted((row for row in rows if row["optimizer"] == spec.name), key=lambda row: row["step"])
        steps = [row["step"] for row in selected]
        axes[0].plot(steps, [row["raw_informative_outlier_count"] for row in selected], "o-", color=colors[spec.name], label=spec.name)
        axes[1].plot(steps, [row["raw_max_outlier_teacher_overlap"] for row in selected], "o-", color=colors[spec.name], label=spec.name)
    axes[0].set_title("Dynamic BBP diagnostic")
    axes[0].set_ylabel("informative eigenvalues outside bulk")
    axes[1].set_title("Residue of detached modes")
    axes[1].set_ylabel("max teacher-subspace overlap")
    for ax in axes:
        ax.set_xlabel("gradient steps")
        ax.grid(alpha=0.20)
        ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(outdir / "dynamic_bbp_fresh_sample.png", bbox_inches="tight")
    plt.close(fig)


def write_report(
    cfg: Config,
    training_summary: dict[str, Any],
    bbp_rows: list[dict[str, Any]],
    outdir: Path,
) -> None:
    lines = [
        "# Hierarchical three-layer Muon/BBP empirical report\n\n",
        "The run separates exact population dynamics, fixed-sample training, and a fresh-sample ambient Hessian. ",
        "The orthogonal Hessian block is conditionally block-Wishart; reported variational edges therefore concern the fresh bulk only.\n\n",
        "## Finite-sample endpoint\n\n",
        "| optimizer | mean exact test risk | median | std | wins vs GD | order-2 half time | order-3 half time | order-4 half time |\n",
        "|---|---:|---:|---:|---:|---:|---:|---:|\n",
    ]
    for spec in OPTIMIZERS:
        item = training_summary[spec.name]
        lines.append(
            f"| {spec.name} | {item['final_test_risk_mean']:.6g} | {item['final_test_risk_median']:.6g} | {item['final_test_risk_std']:.3g} | {item['wins_vs_gd']}/{item['repeats']} | "
            f"{item['order_2_half_error_step']} | {item['order_3_half_error_step']} | {item['order_4_half_error_step']} |\n"
        )
    lines.extend(
        [
            "\n## Fresh-sample spectral audit\n\n",
            "| optimizer | checkpoint | risk | clipped variational bulk | clipped empirical bulk | clipped informative | raw empirical bulk | raw informative | raw max overlap | truncation |\n",
            "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|\n",
        ]
    )
    for row in bbp_rows:
        lines.append(
            f"| {row['optimizer']} | {row['step']} | {row['population_risk']:.3g} | "
            f"[{row['edge_left_variational']:.3g}, {row['edge_right_variational']:.3g}] | "
            f"[{row['edge_left_empirical_bulk']:.3g}, {row['edge_right_empirical_bulk']:.3g}] | "
            f"{row['informative_outlier_count']} | "
            f"[{row['raw_edge_left_empirical_bulk']:.3g}, {row['raw_edge_right_empirical_bulk']:.3g}] | "
            f"{row['raw_informative_outlier_count']} | {row['raw_max_outlier_teacher_overlap']:.3g} | "
            f"{row['block_truncation_fraction']:.3g} |\n"
        )
    lines.extend(
        [
            "\n## Scope\n\n",
            "- `muon_a0` is exact SignSVD/polar Muon on matrix blocks; vectors and biases use gradient descent.\n",
            "- `a=1/7` is a raw phase-retrieval transfer control and `a=1/3` a normalized-branch control, not L3 optima.\n",
            "- The training plot shows the median, individual trajectories, and the min--max envelope over paired repetitions.\n",
            "- Hessian samples are fresh and independent of training. Same-sample DMFT/BBP needs a leave-one-out theorem and is not claimed here.\n",
            "- The variational edge controls the bulk. A theorem for detached eigenvalues additionally needs an anisotropic local law and finite Schur determinant.\n",
            "- The polynomial link has unbounded Hessian weights. Variational edges are therefore reported for explicitly clipped blocks; raw empirical bulk/outlier columns are kept separate.\n",
        ]
    )
    (outdir / "REPORT.md").write_text("".join(lines), encoding="utf-8")


def parse_args() -> Config:
    parser = argparse.ArgumentParser(description=__doc__)
    defaults = Config()
    for name, value in asdict(defaults).items():
        flag = "--" + name.replace("_", "-")
        if isinstance(value, bool):
            parser.add_argument(flag, action=argparse.BooleanOptionalAction, default=value)
        else:
            parser.add_argument(flag, type=type(value), default=value)
    return Config(**vars(parser.parse_args()))


def main() -> None:
    cfg = parse_args()
    if cfg.k >= cfg.d:
        raise ValueError("Require k < d for a non-empty orthogonal bulk.")
    if cfg.p != cfg.q and cfg.tie_u_v:
        raise ValueError("Tied U=V requires p=q.")
    outdir = Path(cfg.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    setup = make_setup(cfg)

    all_rows: list[dict[str, Any]] = []
    empirical_snapshots: dict[str, dict[int, dict[str, torch.Tensor]]] = {}
    for repeat in range(cfg.repeats):
        model_seed = cfg.seed + 1000 * repeat
        base = new_model(cfg, setup, model_seed)
        init_state = clone_state(base)
        generator = torch.Generator().manual_seed(cfg.seed + 20_000 + repeat)
        z_train = torch.randn(cfg.train_n, cfg.k, generator=generator, dtype=setup.z.dtype)
        y_train = teacher_on(z_train, setup)
        for spec in OPTIMIZERS:
            if repeat == 0:
                population_rows, _ = run_trajectory(
                    cfg, setup, spec, "population", repeat, init_state, None, None
                )
                all_rows.extend(population_rows)
            empirical_rows, snapshots = run_trajectory(
                cfg,
                setup,
                spec,
                "empirical",
                repeat,
                init_state,
                z_train,
                y_train,
            )
            all_rows.extend(empirical_rows)
            if repeat == 0:
                empirical_snapshots[spec.name] = snapshots
            print(
                f"[{spec.name}] repeat {repeat + 1}/{cfg.repeats}: "
                f"test risk={empirical_rows[-1]['test_population_risk']:.6g}",
                flush=True,
            )

    bbp_rows: list[dict[str, Any]] = []
    for spec in OPTIMIZERS:
        for step, state in sorted(empirical_snapshots[spec.name].items()):
            row = bbp_state(cfg, setup, spec, step, state)
            bbp_rows.append(row)
            print(
                f"[BBP {spec.name}] step={step} outliers={row['outlier_count']} "
                f"informative={row['informative_outlier_count']}",
                flush=True,
            )

    training_summary = aggregate_training(all_rows, cfg)
    summary = {
        "config": asdict(cfg),
        "optimizer_semantics": {
            "gd": "ordinary gradient descent on every parameter",
            "muon_matrix_update": "U diag((sigma/mean(sigma))^a) V^T, normalized to unit RMS",
            "muon_vector_update": "ordinary gradient descent",
            "a_1_7_status": "phase-retrieval transfer control, not an L3 theorem",
            "a_1_3_status": "normalized phase-retrieval branch control, not an L3 theorem",
        },
        "training": training_summary,
        "bbp": bbp_rows,
    }
    write_csv(outdir / "training_trajectories.csv", all_rows)
    write_csv(outdir / "fresh_sample_bbp.csv", bbp_rows)
    (outdir / "summary.json").write_text(
        json.dumps(json_ready(summary), indent=2), encoding="utf-8"
    )
    plot_training(all_rows, outdir)
    plot_bbp(bbp_rows, outdir)
    write_report(cfg, training_summary, bbp_rows, outdir)
    print(f"Wrote {outdir}", flush=True)


if __name__ == "__main__":
    main()