#!/usr/bin/env python3
"""Dense, coupled fresh-probe spectroscopy for the hierarchical L3 model.

This script is a presentation experiment, not a same-data BBP theorem.  It
reproduces one paired training instance on a dense checkpoint grid and freezes
one Gaussian probe sample across all checkpoints and optimizers.  The coupling
removes irrelevant sample-to-sample jitter and makes spectral motion readable.

Outputs:
  dense_training_trajectories.csv
  exact_sector_geometry.csv
  coupled_fresh_spectrum.csv
  training_observables_comparison.png
  sector_recovery_gd_vs_muon.png
  spectral_time_stack_gd.png
  spectral_time_stack_muon_a0.png
  dmft_bbp_rmt_master_panel.png
  REPORT.md
"""

from __future__ import annotations

import argparse
import math
import os
from pathlib import Path
from typing import Any

import numpy as np
import torch

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

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize

from experiments.brainstorm.resnetl3.hierarchical3_muon_bbp import (
    Config,
    OPTIMIZERS,
    clone_state,
    link_hessian_blocks,
    make_setup,
    new_model,
    run_trajectory,
    teacher_on,
    write_csv,
)
from experiments.brainstorm.resnetl3.resnet_order4_full_verify_v2 import (
    exact_hessian_decomposition,
    make_weighted_gram,
    order_tangent_bases,
    projected_block_stats,
    schur_complements_by_order,
    weighted_parameter_jacobian_by_order,
)
from experiments.muon.lowrank_block_wishart_spectral_experiments import (
    block_wishart,
)


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

DISPLAY = {
    "gd": "Gradient descent",
    "muon_a0": r"Polar Muon ($a=0$)",
    "muon_a1_7_transfer": r"Muon ($a=1/7$)",
    "muon_a1_3_normalized": r"Muon ($a=1/3$)",
}
METHOD_COLORS = {
    "gd": "#355C9A",
    "muon_a0": "#E07A3F",
    "muon_a1_7_transfer": "#1B9E77",
    "muon_a1_3_normalized": "#8E5AA9",
}
DEGREE_COLORS = {2: "#277DA1", 3: "#43AA8B", 4: "#F8961E"}


def figure_style() -> None:
    plt.rcParams.update(
        {
            "figure.facecolor": "#FCFCFA",
            "axes.facecolor": "#FCFCFA",
            "axes.edgecolor": "#30343B",
            "axes.labelcolor": "#24272E",
            "axes.titleweight": "semibold",
            "axes.titlesize": 12,
            "font.size": 10,
            "grid.color": "#CDD3DA",
            "grid.alpha": 0.45,
            "legend.frameon": False,
            "savefig.facecolor": "#FCFCFA",
        }
    )


def coupled_fresh_probe(
    cfg: Config,
    setup: Any,
    state: dict[str, torch.Tensor],
    optimizer: str,
    step: int,
    z: np.ndarray,
    xi: np.ndarray,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Return the raw full spectrum and the empirical orthogonal bulk edges."""

    model = new_model(cfg, setup, cfg.seed)
    model.load_state_dict(state)
    target = teacher_on(
        torch.tensor(z, dtype=setup.z.dtype), setup
    ).detach().cpu().numpy()
    blocks = link_hessian_blocks(model, z, target)
    x = np.concatenate([z, xi], axis=1)
    full = block_wishart(x, blocks)
    orthogonal = block_wishart(xi, blocks)
    eigenvalues, eigenvectors = np.linalg.eigh(full)
    bulk_eigenvalues = np.linalg.eigvalsh(orthogonal)

    block_size = cfg.p + cfg.q
    reshaped = eigenvectors.reshape(block_size, cfg.d, -1)
    teacher_overlap = np.sum(reshaped[:, : cfg.k, :] ** 2, axis=(0, 1))
    bulk_left = float(bulk_eigenvalues[0])
    bulk_right = float(bulk_eigenvalues[-1])
    tolerance = 0.015 * max(1.0, bulk_right - bulk_left)
    outside = (eigenvalues < bulk_left - tolerance) | (
        eigenvalues > bulk_right + tolerance
    )
    informative_threshold = max(0.20, 2.5 * cfg.k / cfg.d)
    informative = outside & (teacher_overlap >= informative_threshold)

    rows = [
        {
            "optimizer": optimizer,
            "step": step,
            "eigen_index": index,
            "eigenvalue": float(value),
            "teacher_overlap": float(overlap),
            "outside_empirical_orthogonal_bulk": int(is_outside),
            "informative_outlier": int(is_informative),
        }
        for index, (value, overlap, is_outside, is_informative) in enumerate(
            zip(eigenvalues, teacher_overlap, outside, informative)
        )
    ]
    edge = {
        "optimizer": optimizer,
        "step": step,
        "bulk_left": bulk_left,
        "bulk_right": bulk_right,
        "informative_outlier_count": int(np.sum(informative)),
        "max_teacher_overlap": float(np.max(teacher_overlap)),
    }
    return rows, edge


def empirical_rows(rows: list[dict[str, Any]], optimizer: str) -> list[dict[str, Any]]:
    return sorted(
        (
            row
            for row in rows
            if row["regime"] == "empirical" and row["optimizer"] == optimizer
        ),
        key=lambda row: int(row["step"]),
    )


def mean_alignment(row: dict[str, Any], prefix: str, k: int) -> float:
    return float(np.mean([row[f"align_{prefix}_{j}"] for j in range(1, k + 1)]))


def half_error_times(
    trajectory: list[dict[str, Any]],
) -> dict[int, int | None]:
    result: dict[int, int | None] = {}
    for degree in (2, 3, 4):
        key = f"E_order_{degree}"
        threshold = 0.5 * float(trajectory[0][key])
        result[degree] = next(
            (
                int(row["step"])
                for row in trajectory
                if float(row[key]) <= threshold
            ),
            None,
        )
    return result


def exact_sector_geometry(
    cfg: Config,
    setup: Any,
    state: dict[str, torch.Tensor],
    optimizer: str,
    step: int,
) -> dict[str, Any]:
    """Evaluate the exact finite Hermite/Schur geometry at one state."""

    model = new_model(cfg, setup, cfg.seed)
    model.load_state_dict(state)
    hessian, gauss_newton, residual, coefficient, delta, jacobian = (
        exact_hessian_decomposition(model, setup)
    )
    _, coefficient_kernel, weighted_gram = make_weighted_gram(
        jacobian, setup.basis_fact.detach().cpu().numpy()
    )
    schur = schur_complements_by_order(weighted_gram, setup.orders)
    weighted_jacobians = weighted_parameter_jacobian_by_order(
        jacobian,
        setup.basis_fact_sqrt.detach().cpu().numpy(),
        setup.orders,
    )
    tangent_bases = order_tangent_bases(weighted_jacobians)
    projected = projected_block_stats(
        hessian, gauss_newton, residual, tangent_bases
    )

    row: dict[str, Any] = {
        "optimizer": optimizer,
        "step": step,
        "coefficient_error_norm": float(
            np.linalg.norm(
                np.sqrt(setup.basis_fact.detach().cpu().numpy()) * delta
            )
        ),
        "hessian_min": float(np.linalg.eigvalsh(hessian)[0]),
        "hessian_max": float(np.linalg.eigvalsh(hessian)[-1]),
    }
    all_orders = (2, 3, 4)
    for degree in all_orders:
        stats = schur[f"S{degree}"]
        positive = stats.eigs[stats.eigs > 1e-10]
        lambda_min_positive = (
            float(positive[0]) if positive.size else float("nan")
        )
        indices = setup.orders[degree]
        self_drift = -coefficient_kernel[np.ix_(indices, indices)] @ delta[
            indices
        ]
        cross_drift = np.zeros(len(indices), dtype=float)
        for other in all_orders:
            if other == degree:
                continue
            other_indices = setup.orders[other]
            cross_drift -= coefficient_kernel[
                np.ix_(indices, other_indices)
            ] @ delta[other_indices]
        rho = float(projected[f"rho_{degree}"])
        row.update(
            {
                f"schur_rank_{degree}": stats.rank,
                f"schur_trace_{degree}": stats.trace,
                f"schur_lambda_max_{degree}": stats.lam_max,
                f"schur_lambda_min_positive_{degree}": lambda_min_positive,
                f"residual_radius_{degree}": rho,
                f"access_margin_{degree}": lambda_min_positive - rho,
                f"self_drift_norm_{degree}": float(
                    np.linalg.norm(self_drift)
                ),
                f"cross_drift_norm_{degree}": float(
                    np.linalg.norm(cross_drift)
                ),
                f"projected_hessian_min_{degree}": float(
                    projected[f"lam_min_{degree}"]
                ),
            }
        )
    return row


def plot_comparison(rows: list[dict[str, Any]], cfg: Config, output: Path) -> None:
    figure_style()
    fig, axes = plt.subplots(1, 3, figsize=(14.2, 4.3), dpi=210)
    for spec in OPTIMIZERS:
        trajectory = empirical_rows(rows, spec.name)
        steps = [int(row["step"]) for row in trajectory]
        color = METHOD_COLORS[spec.name]
        axes[0].plot(
            steps,
            [row["test_population_risk"] for row in trajectory],
            color=color,
            lw=2.1,
            label=DISPLAY[spec.name],
        )
        axes[1].plot(
            steps,
            [mean_alignment(row, "V", cfg.k) for row in trajectory],
            color=color,
            lw=2.1,
        )
        axes[2].plot(
            steps,
            [mean_alignment(row, "U", cfg.k) for row in trajectory],
            color=color,
            lw=2.1,
        )
    axes[0].set_yscale("log")
    axes[0].set_title("Generalization")
    axes[0].set_ylabel("exact population risk")
    axes[1].set_title("First feature map")
    axes[1].set_ylabel(r"mean teacher overlap of $V$")
    axes[2].set_title("Second direct map")
    axes[2].set_ylabel(r"mean teacher overlap of $U$")
    for axis in axes:
        axis.set_xlabel("training step")
        axis.grid(True)
    axes[1].set_ylim(0.0, 1.02)
    axes[2].set_ylim(0.0, 1.02)
    axes[0].legend(loc="best", fontsize=8)
    fig.suptitle(
        "One shared instance: risk and teacher-subspace recovery",
        fontsize=14,
        fontweight="bold",
    )
    fig.tight_layout()
    fig.savefig(output, bbox_inches="tight")
    plt.close(fig)


def plot_sectors(rows: list[dict[str, Any]], output: Path) -> None:
    figure_style()
    methods = ("gd", "muon_a0")
    fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.3), dpi=210, sharey=True)
    for axis, optimizer in zip(axes, methods):
        trajectory = empirical_rows(rows, optimizer)
        steps = [int(row["step"]) for row in trajectory]
        for degree in (2, 3, 4):
            key = f"E_order_{degree}"
            initial = float(trajectory[0][key])
            axis.plot(
                steps,
                [float(row[key]) / initial for row in trajectory],
                color=DEGREE_COLORS[degree],
                lw=2.2,
                label=f"degree {degree}",
            )
        axis.axhline(0.5, color="#4D535B", ls="--", lw=1.2, label="half error")
        axis.set_yscale("log")
        axis.set_xlabel("training step")
        axis.set_title(DISPLAY[optimizer])
        axis.grid(True)
    axes[0].set_ylabel(r"sector error $E_r(t)/E_r(0)$")
    axes[1].legend(loc="best", fontsize=8)
    fig.suptitle(
        "Sequential recovery: lower curves mean a sector has been learned",
        fontsize=14,
        fontweight="bold",
    )
    fig.tight_layout()
    fig.savefig(output, bbox_inches="tight")
    plt.close(fig)


def plot_spectral_stack(
    training_rows: list[dict[str, Any]],
    spectrum_rows: list[dict[str, Any]],
    edge_rows: list[dict[str, Any]],
    cfg: Config,
    optimizer: str,
    output: Path,
) -> None:
    figure_style()
    trajectory = empirical_rows(training_rows, optimizer)
    spectrum = [row for row in spectrum_rows if row["optimizer"] == optimizer]
    edges = sorted(
        (row for row in edge_rows if row["optimizer"] == optimizer),
        key=lambda row: int(row["step"]),
    )
    steps = [int(row["step"]) for row in trajectory]
    event_times = half_error_times(trajectory)

    fig, axes = plt.subplots(
        3,
        1,
        figsize=(10.2, 11.0),
        dpi=220,
        sharex=True,
        gridspec_kw={"height_ratios": [1.0, 1.0, 2.15]},
    )
    color = METHOD_COLORS[optimizer]
    axes[0].plot(
        steps,
        [row["test_population_risk"] for row in trajectory],
        color=color,
        lw=2.4,
    )
    axes[0].set_yscale("log")
    axes[0].set_ylabel("exact population risk")
    axes[0].set_title("A. Generalization error")
    axes[0].grid(True)

    axes[1].plot(
        steps,
        [mean_alignment(row, "V", cfg.k) for row in trajectory],
        color="#277DA1",
        lw=2.2,
        label=r"first map $V$",
    )
    axes[1].plot(
        steps,
        [mean_alignment(row, "U", cfg.k) for row in trajectory],
        color="#F8961E",
        lw=2.2,
        label=r"second direct map $U$",
    )
    axes[1].set_ylim(0.0, 1.02)
    axes[1].set_ylabel("mean teacher overlap")
    axes[1].set_title("B. Recovery of the teacher subspace")
    axes[1].legend(loc="best")
    axes[1].grid(True)

    edge_steps = np.asarray([int(row["step"]) for row in edges])
    left = np.asarray([float(row["bulk_left"]) for row in edges])
    right = np.asarray([float(row["bulk_right"]) for row in edges])
    axes[2].fill_between(
        edge_steps,
        left,
        right,
        color="#83B5D1",
        alpha=0.28,
        label="fresh orthogonal Hessian bulk",
    )
    axes[2].plot(edge_steps, left, color="#4A86A8", lw=1.4)
    axes[2].plot(edge_steps, right, color="#4A86A8", lw=1.4)

    values = np.asarray([float(row["eigenvalue"]) for row in spectrum])
    overlaps = np.asarray([float(row["teacher_overlap"]) for row in spectrum])
    spectrum_steps = np.asarray([int(row["step"]) for row in spectrum])
    scatter = axes[2].scatter(
        spectrum_steps,
        values,
        c=overlaps,
        cmap="magma",
        norm=Normalize(0.0, 1.0),
        s=8,
        alpha=0.72,
        linewidths=0,
        rasterized=True,
    )
    informative = [
        row for row in spectrum if int(row["informative_outlier"]) != 0
    ]
    if informative:
        axes[2].scatter(
            [int(row["step"]) for row in informative],
            [float(row["eigenvalue"]) for row in informative],
            marker="o",
            s=34,
            facecolors="none",
            edgecolors="#00A087",
            linewidths=1.2,
            label="outside bulk + teacher aligned",
            zorder=5,
        )
    axes[2].axhline(0.0, color="#30343B", lw=0.8, alpha=0.55)
    axes[2].set_yscale("symlog", linthresh=0.2, linscale=0.8)
    axes[2].set_ylabel("raw full-Hessian eigenvalue")
    axes[2].set_xlabel("training step")
    axes[2].set_title(
        "C. Spectral motion — color is teacher-subspace weight"
    )
    axes[2].grid(True)
    axes[2].legend(loc="best", fontsize=8)
    colorbar = fig.colorbar(scatter, ax=axes[2], pad=0.012, aspect=28)
    colorbar.set_label("teacher-subspace weight")

    for degree, event in event_times.items():
        if event is None:
            continue
        for axis in axes:
            axis.axvline(
                event,
                color=DEGREE_COLORS[degree],
                ls=":",
                lw=1.0,
                alpha=0.72,
            )
        axes[0].text(
            event,
            0.98,
            f"$E_{degree}/2$",
            color=DEGREE_COLORS[degree],
            transform=axes[0].get_xaxis_transform(),
            ha="center",
            va="top",
            fontsize=8,
        )

    fig.suptitle(
        f"{DISPLAY[optimizer]}: training observables and fresh Hessian spectrum",
        fontsize=15,
        fontweight="bold",
    )
    fig.tight_layout()
    fig.savefig(output, bbox_inches="tight")
    plt.close(fig)


def plot_master_mechanism(
    training_rows: list[dict[str, Any]],
    geometry_rows: list[dict[str, Any]],
    spectrum_rows: list[dict[str, Any]],
    edge_rows: list[dict[str, Any]],
    output: Path,
) -> None:
    """Put deterministic access and fresh-sample spectroscopy on one clock."""

    figure_style()
    methods = ("gd", "muon_a0")
    fig, axes = plt.subplots(
        3,
        2,
        figsize=(14.2, 11.4),
        dpi=230,
        sharex="col",
        gridspec_kw={"height_ratios": [1.0, 1.0, 1.7]},
    )
    for column, optimizer in enumerate(methods):
        trajectory = empirical_rows(training_rows, optimizer)
        geometry = sorted(
            (
                row
                for row in geometry_rows
                if row["optimizer"] == optimizer
            ),
            key=lambda row: int(row["step"]),
        )
        spectrum = [
            row for row in spectrum_rows if row["optimizer"] == optimizer
        ]
        edges = sorted(
            (
                row for row in edge_rows if row["optimizer"] == optimizer
            ),
            key=lambda row: int(row["step"]),
        )
        steps = np.asarray([int(row["step"]) for row in trajectory])
        risk0 = float(trajectory[0]["test_population_risk"])
        axes[0, column].plot(
            steps,
            [
                float(row["test_population_risk"]) / risk0
                for row in trajectory
            ],
            color="#252A34",
            lw=2.5,
            label=r"risk $\mathcal{R}(t)/\mathcal{R}(0)$",
        )
        for degree in (2, 3, 4):
            initial = float(trajectory[0][f"E_order_{degree}"])
            axes[0, column].plot(
                steps,
                [
                    float(row[f"E_order_{degree}"]) / initial
                    for row in trajectory
                ],
                color=DEGREE_COLORS[degree],
                lw=2.0,
                label=rf"$E_{degree}(t)/E_{degree}(0)$",
            )
        axes[0, column].axhline(
            0.5, color="#737A84", lw=1.0, ls="--"
        )
        axes[0, column].set_yscale("log")
        axes[0, column].set_title(
            f"A{column + 1}. {DISPLAY[optimizer]} — population sectors"
        )
        axes[0, column].grid(True)

        geometry_steps = np.asarray([int(row["step"]) for row in geometry])
        for degree in (2, 3, 4):
            access_ratio = [
                float(row[f"schur_lambda_min_positive_{degree}"])
                / max(float(row[f"residual_radius_{degree}"]), 1e-15)
                for row in geometry
            ]
            axes[1, column].plot(
                geometry_steps,
                access_ratio,
                color=DEGREE_COLORS[degree],
                marker="o",
                ms=3.0,
                lw=1.8,
                label=rf"$\lambda^+_{{\min}}(S_{degree})/\rho_{degree}$",
            )
        axes[1, column].axhspan(
            1.0,
            1e2,
            color="#CDE9D8",
            alpha=0.28,
            zorder=-3,
        )
        axes[1, column].axhline(
            1.0, color="#30343B", lw=1.0, ls="--"
        )
        axes[1, column].set_yscale("log")
        axes[1, column].set_ylim(1e-11, 2.0)
        axes[1, column].set_title(
            f"B{column + 1}. Exact Schur access ratio (certificate above 1)"
        )
        axes[1, column].grid(True)

        edge_steps = np.asarray([int(row["step"]) for row in edges])
        left = np.asarray([float(row["bulk_left"]) for row in edges])
        right = np.asarray([float(row["bulk_right"]) for row in edges])
        axes[2, column].fill_between(
            edge_steps,
            left,
            right,
            color="#83B5D1",
            alpha=0.30,
            label="fresh orthogonal bulk",
        )
        axes[2, column].plot(edge_steps, left, color="#4A86A8", lw=1.3)
        axes[2, column].plot(edge_steps, right, color="#4A86A8", lw=1.3)
        scatter = axes[2, column].scatter(
            [int(row["step"]) for row in spectrum],
            [float(row["eigenvalue"]) for row in spectrum],
            c=[float(row["teacher_overlap"]) for row in spectrum],
            cmap="magma",
            norm=Normalize(0.0, 1.0),
            s=4,
            alpha=0.55,
            linewidths=0,
            rasterized=True,
        )
        informative = [
            row
            for row in spectrum
            if int(row["informative_outlier"]) != 0
        ]
        axes[2, column].scatter(
            [int(row["step"]) for row in informative],
            [float(row["eigenvalue"]) for row in informative],
            marker="D",
            s=36,
            facecolors="none",
            edgecolors="#00A087",
            linewidths=1.25,
            label="detached + teacher aligned",
            zorder=5,
        )
        axes[2, column].axhline(
            0.0, color="#30343B", lw=0.8, alpha=0.55
        )
        axes[2, column].set_yscale(
            "symlog", linthresh=0.2, linscale=0.8
        )
        axes[2, column].set_title(
            f"C{column + 1}. Fresh raw Hessian against orthogonal RMT bulk"
        )
        axes[2, column].set_xlabel("training step")
        axes[2, column].grid(True)

        for event in half_error_times(trajectory).values():
            if event is None:
                continue
            for row_index in range(3):
                axes[row_index, column].axvline(
                    event, color="#8A9099", ls=":", lw=0.8, alpha=0.55
                )

    axes[0, 0].set_ylabel("normalized error")
    axes[1, 0].set_ylabel(r"access ratio $\lambda_{\min}^+(S_r)/\rho_r$")
    axes[2, 0].set_ylabel("raw Hessian eigenvalue")
    axes[0, 1].legend(loc="best", fontsize=8, ncol=2)
    axes[1, 1].legend(loc="best", fontsize=8)
    axes[2, 1].legend(loc="best", fontsize=8)
    colorbar_axis = fig.add_axes([0.925, 0.075, 0.013, 0.27])
    colorbar = fig.colorbar(scatter, cax=colorbar_axis)
    colorbar.set_label("eigenvector weight in teacher coordinates")
    fig.suptitle(
        "Three-layer mechanism map: coefficient recovery → Schur access → "
        "fresh spectral separation",
        fontsize=15,
        fontweight="bold",
        y=0.995,
    )
    fig.subplots_adjust(
        left=0.075, right=0.91, bottom=0.065, top=0.955, hspace=0.25, wspace=0.16
    )
    fig.savefig(output, bbox_inches="tight")
    plt.close(fig)


def write_report(
    cfg: Config,
    training_rows: list[dict[str, Any]],
    geometry_rows: list[dict[str, Any]],
    edge_rows: list[dict[str, Any]],
    output: Path,
    step_zero_discrepancy: float,
) -> None:
    lines = [
        "# Dense spectral-time panel\n\n",
        "## Object\n\n",
        "One paired training instance is rerun with a dense checkpoint grid. "
        "A single Gaussian probe `(z, xi)` is drawn independently of training "
        "and then frozen across checkpoints and optimizers. This coupling makes "
        "spectral motion readable; it is not a same-data Hessian claim.\n\n",
        "## How to read the spectral panels\n\n",
        "- The blue band is the empirical spectrum of the orthogonal Hessian "
        "`sum T(z_mu) tensor xi_mu xi_mu^T / n`.\n",
        "- Every point is an eigenvalue of the raw full direct-weight Hessian.\n",
        "- Point color is the eigenvector mass in the teacher-coordinate "
        "subspace: dark means nearly orthogonal, yellow means teacher aligned.\n",
        "- A green ring marks a point that is both outside the orthogonal bulk "
        "and above the pre-specified teacher-overlap threshold.\n",
        "- Vertical dotted lines are the first recorded half-error events for "
        "Hermite degrees 2, 3, and 4.\n\n",
        "## How to read `dmft_bbp_rmt_master_panel.png`\n\n",
        "- Row A contains exact population quantities. The black curve is "
        "risk normalized by its initial value; the colored curves are exact "
        "Hermite-sector errors normalized in the same way.\n",
        "- Row B is deterministic finite algebra at the trained state. A "
        "colored curve is the ratio "
        "`lambda_min_positive(S_r) / rho_r`; values above one certify "
        "positive curvature on the residualized degree-r tangent space. "
        "Values below one mean only that this sufficient certificate fails.\n",
        "- Row C is a fresh-sample finite-size diagnostic. The band is an "
        "empirical orthogonal bulk, point color is teacher-coordinate mass, "
        "and green diamonds apply the declared detachment-and-overlap rule.\n",
        "- Columns use the same initialization, training observations, and "
        "fresh Gaussian probe. Dotted lines repeat the sector half-error "
        "events across all three rows; visual proximity is not a causality "
        "test.\n\n",
        "## Configuration\n\n",
        f"- seed: `{cfg.seed}`\n",
        f"- training steps: `{cfg.steps}`\n",
        f"- diagnostic interval: `{cfg.diag_every}`\n",
        f"- spectral checkpoints: `{cfg.bbp_checkpoints}`\n",
        f"- train sample size: `{cfg.train_n}`\n",
        f"- fresh probe size: `{int(round(cfg.alpha * (cfg.d - cfg.k)))}`\n",
        f"- ambient dimension: `{cfg.d}`; latent dimension: `{cfg.k}`\n",
        f"- maximum step-zero spectral discrepancy across optimizers: "
        f"`{step_zero_discrepancy:.3e}`\n\n",
        f"- exact geometry states evaluated: `{len(geometry_rows)}`\n\n",
        "## Half-error times on the displayed instance\n\n",
        "| optimizer | degree 2 | degree 3 | degree 4 |\n",
        "|---|---:|---:|---:|\n",
    ]
    for spec in OPTIMIZERS:
        times = half_error_times(empirical_rows(training_rows, spec.name))
        cells = ["not reached" if times[d] is None else str(times[d]) for d in (2, 3, 4)]
        lines.append(f"| {DISPLAY[spec.name]} | " + " | ".join(cells) + " |\n")
    lines.extend(
        [
            "\n## Scope\n\n",
            "The orthogonal bulk identity is exact conditional on the latent "
            "probe coordinates. The plotted edges are empirical finite-sample "
            "edges, not Montanari--Saeed variational limits. Colored detached "
            "points are diagnostics; a BBP theorem still requires an "
            "anisotropic local law and a controlled finite Schur limit.\n",
        ]
    )
    output.write_text("".join(lines), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--outdir", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--seed", type=int, default=424242)
    parser.add_argument("--steps", type=int, default=1200)
    parser.add_argument("--diag-every", type=int, default=20)
    parser.add_argument("--checkpoints", type=int, default=25)
    args = parser.parse_args()

    cfg = Config(
        outdir=str(args.outdir),
        seed=args.seed,
        repeats=1,
        steps=args.steps,
        diag_every=args.diag_every,
        bbp_checkpoints=args.checkpoints,
    )
    args.outdir.mkdir(parents=True, exist_ok=True)
    setup = make_setup(cfg)
    base = new_model(cfg, setup, cfg.seed)
    init_state = clone_state(base)
    generator = torch.Generator().manual_seed(cfg.seed + 20_000)
    z_train = torch.randn(
        cfg.train_n, cfg.k, generator=generator, dtype=setup.z.dtype
    )
    y_train = teacher_on(z_train, setup)

    training_rows: list[dict[str, Any]] = []
    snapshots: dict[str, dict[int, dict[str, torch.Tensor]]] = {}
    for spec in OPTIMIZERS:
        rows, states = run_trajectory(
            cfg,
            setup,
            spec,
            "empirical",
            0,
            init_state,
            z_train,
            y_train,
        )
        training_rows.extend(rows)
        snapshots[spec.name] = states
        print(
            f"[{spec.name}] endpoint risk={rows[-1]['test_population_risk']:.6g}",
            flush=True,
        )

    rng = np.random.default_rng(cfg.seed + 900_001)
    n_fresh = int(round(cfg.alpha * (cfg.d - cfg.k)))
    z_fresh = rng.standard_normal((n_fresh, cfg.k))
    xi_fresh = rng.standard_normal((n_fresh, cfg.d - cfg.k))
    spectrum_rows: list[dict[str, Any]] = []
    edge_rows: list[dict[str, Any]] = []
    geometry_rows: list[dict[str, Any]] = []
    for spec in OPTIMIZERS:
        for step, state in sorted(snapshots[spec.name].items()):
            geometry_rows.append(
                exact_sector_geometry(
                    cfg, setup, state, spec.name, step
                )
            )
            rows, edge = coupled_fresh_probe(
                cfg,
                setup,
                state,
                spec.name,
                step,
                z_fresh,
                xi_fresh,
            )
            spectrum_rows.extend(rows)
            edge_rows.append(edge)
        print(
            f"[{spec.name}] spectra={len(snapshots[spec.name])}",
            flush=True,
        )

    zero_spectra = {
        spec.name: np.asarray(
            [
                row["eigenvalue"]
                for row in spectrum_rows
                if row["optimizer"] == spec.name and row["step"] == 0
            ]
        )
        for spec in OPTIMIZERS
    }
    reference = zero_spectra["gd"]
    step_zero_discrepancy = max(
        float(np.max(np.abs(values - reference)))
        for values in zero_spectra.values()
    )
    if not math.isfinite(step_zero_discrepancy) or step_zero_discrepancy > 1e-10:
        raise RuntimeError(
            f"Common-state/common-probe check failed: {step_zero_discrepancy}"
        )

    write_csv(args.outdir / "dense_training_trajectories.csv", training_rows)
    write_csv(args.outdir / "exact_sector_geometry.csv", geometry_rows)
    write_csv(args.outdir / "coupled_fresh_spectrum.csv", spectrum_rows)
    write_csv(args.outdir / "coupled_fresh_bulk_edges.csv", edge_rows)
    plot_comparison(
        training_rows, cfg, args.outdir / "training_observables_comparison.png"
    )
    plot_sectors(
        training_rows, args.outdir / "sector_recovery_gd_vs_muon.png"
    )
    plot_spectral_stack(
        training_rows,
        spectrum_rows,
        edge_rows,
        cfg,
        "gd",
        args.outdir / "spectral_time_stack_gd.png",
    )
    plot_spectral_stack(
        training_rows,
        spectrum_rows,
        edge_rows,
        cfg,
        "muon_a0",
        args.outdir / "spectral_time_stack_muon_a0.png",
    )
    plot_master_mechanism(
        training_rows,
        geometry_rows,
        spectrum_rows,
        edge_rows,
        args.outdir / "dmft_bbp_rmt_master_panel.png",
    )
    write_report(
        cfg,
        training_rows,
        geometry_rows,
        edge_rows,
        args.outdir / "REPORT.md",
        step_zero_discrepancy,
    )
    print(f"Wrote {args.outdir}", flush=True)


if __name__ == "__main__":
    main()