#!/usr/bin/env python3
"""Render the audited standard-depth ResNet scaling figure.

The renderer reads all 60 frozen records directly, applies the original record
validator, verifies every source hash against the passed gate, and refuses an
incomplete panel or disagreement between raw records and gate statistics.
"""
import argparse
import glob
import hashlib
import json
import math
import os
import statistics
import subprocess

import matplotlib

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

from analyze_oral_a_dynamic_scaling import (
    DEPTHS,
    METHODS,
    SEEDS,
    validate_record,
)


PROTOCOL = "oral_a_dynamic_innovation_scaling_recovery_v2"
BLUE = "#0072B2"
LIGHT_BLUE = "#56B4E9"
GREEN = "#009E73"
ORANGE = "#D55E00"
GRAY = "#4D4D4D"
LIGHT_GRAY = "#B7B7B7"
PDF_METADATA = {
    "Creator": "SDIL audited standard-depth figure pipeline",
    "Producer": "Matplotlib",
    "CreationDate": None,
    "ModDate": None,
}


def load_json(path):
    with open(path) as handle:
        return json.load(handle)


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def require(condition, message):
    if not condition:
        raise RuntimeError(message)


def git_blob_sha256(commit, relative_path):
    content = subprocess.run(
        ["git", "show", f"{commit}:{relative_path}"],
        check=True,
        capture_output=True,
    ).stdout
    return hashlib.sha256(content).hexdigest()


def expected_paths(d4_dir, new_dir):
    paths = {}
    for depth in DEPTHS:
        for seed in SEEDS:
            for method in METHODS:
                if depth == 20 and method in ("clean_kp", "dynamic"):
                    label = "clean_kp" if method == "clean_kp" else "dynamic"
                    path = os.path.join(d4_dir, f"seed{seed}_{label}.json")
                else:
                    path = os.path.join(
                        new_dir, f"{method}_d{depth}_s{seed}.json"
                    )
                paths[(method, depth, seed)] = path
    return paths


def require_exact_directories(d4_dir, new_dir):
    expected_d4 = {
        f"seed{seed}_{condition}.json"
        for seed in SEEDS
        for condition in ("clean_kp", "dynamic")
    }
    expected_new = {
        f"{method}_d{depth}_s{seed}.json"
        for depth in DEPTHS
        for seed in SEEDS
        for method in METHODS
        if not (depth == 20 and method in ("clean_kp", "dynamic"))
    }
    observed_d4 = {
        os.path.basename(path)
        for path in glob.glob(os.path.join(d4_dir, "*.json"))
    }
    observed_new = {
        os.path.basename(path)
        for path in glob.glob(os.path.join(new_dir, "*.json"))
    }
    require(observed_d4 == expected_d4, "D4 source directory drift")
    require(observed_new == expected_new, "standard-depth source directory drift")


def load_panel(d4_dir, new_dir, gate_path, d4_gate_path, bci_v2_gate_path):
    gate = load_json(gate_path)
    require(gate.get("protocol") == PROTOCOL, "standard-depth protocol drift")
    require(gate.get("status") == "passed", "standard-depth gate is not passed")
    require(gate.get("complete_grid") is True, "standard-depth grid incomplete")
    require(
        gate.get("standard_depth_scaling_established") is True,
        "standard-depth claim was not established",
    )
    require(gate.get("review_score_after") == 9, "score-change rule drift")
    require(
        gate.get("checks") and all(gate["checks"].values()),
        "standard-depth gate contains a failed check",
    )
    require_exact_directories(d4_dir, new_dir)

    records = {}
    rows = {}
    paths = expected_paths(d4_dir, new_dir)
    require(len(paths) == 60, "standard-depth panel must contain 60 cells")
    for key, path in paths.items():
        method, depth, seed = key
        require(os.path.isfile(path), f"missing standard-depth record: {path}")
        gate_digest = gate.get("source_sha256", {}).get(path)
        require(gate_digest is not None, f"gate omits source hash: {path}")
        require(sha256(path) == gate_digest, f"source hash drift: {path}")
        record = load_json(path)
        records[key] = record
        rows[key] = validate_record(record, path, method, depth, seed)

    d4_commits = {
        rows[(method, 20, seed)]["source_commit"]
        for method in ("clean_kp", "dynamic")
        for seed in SEEDS
    }
    new_commits = {
        rows[(method, depth, seed)]["source_commit"]
        for method in METHODS
        for depth in DEPTHS
        for seed in SEEDS
        if not (depth == 20 and method in ("clean_kp", "dynamic"))
    }
    require(len(d4_commits) == 1, "D4 source revision drift")
    require(len(new_commits) == 1, "standard-depth source revision drift")
    d4_commit = next(iter(d4_commits))
    new_commit = next(iter(new_commits))
    require(
        d4_commit == gate.get("d4_source_commit"),
        "D4 source commit disagrees with gate",
    )
    require(
        new_commit == gate.get("new_source_commit"),
        "standard-depth source commit disagrees with gate",
    )
    historical_sources = {
        "protocol_sha256": "ORAL_A_RECOVERY_V2.md",
        "runner_sha256": "experiments/oral_a_dynamic_scaling_v2.py",
        "analyzer_sha256": "experiments/analyze_oral_a_dynamic_scaling_v2.py",
        "legacy_runner_sha256": "experiments/oral_a_dynamic_scaling.py",
        "legacy_analyzer_sha256": "experiments/analyze_oral_a_dynamic_scaling.py",
        "conv_run_sha256": "experiments/conv_run.py",
    }
    for gate_key, relative_path in historical_sources.items():
        require(
            git_blob_sha256(new_commit, relative_path) == gate.get(gate_key),
            f"historical source binding drift: {relative_path}",
        )
    require(
        sha256(d4_gate_path) == gate.get("d4_gate_sha256"),
        "D4 prerequisite gate hash drift",
    )
    require(
        sha256(bci_v2_gate_path) == gate.get("bci_v2_gate_sha256"),
        "BCI-v2 prerequisite gate hash drift",
    )

    metrics = gate["metrics"]
    for method in METHODS:
        for depth in DEPTHS:
            observed = [
                rows[(method, depth, seed)]["accuracy"] for seed in SEEDS
            ]
            expected = metrics["accuracy_by_method_depth_seed"][method][str(depth)]
            require(
                np.allclose(observed, expected, rtol=0.0, atol=0.0),
                f"{method} depth-{depth} accuracy summary drift",
            )
            require(
                math.isclose(
                    statistics.mean(observed),
                    metrics["mean_accuracy"][method][str(depth)],
                    rel_tol=0.0,
                    abs_tol=1e-15,
                ),
                f"{method} depth-{depth} mean summary drift",
            )
    observed_alignment = {
        str(depth): [
            rows[("dynamic", depth, seed)]["early_alignment"] for seed in SEEDS
        ]
        for depth in DEPTHS
    }
    for depth in DEPTHS:
        require(
            np.allclose(
                observed_alignment[str(depth)],
                metrics["dynamic_alignment"][str(depth)],
                rtol=0.0,
                atol=0.0,
            ),
            f"depth-{depth} alignment summary drift",
        )
    return gate, records, rows, paths


def style_axis(axis):
    axis.spines["top"].set_visible(False)
    axis.spines["right"].set_visible(False)
    axis.tick_params(width=1.0, length=4)
    axis.grid(axis="y", color="#E7E7E7", linewidth=0.8, zorder=0)


def panel_label(axis, label):
    axis.text(
        -0.13,
        1.08,
        label,
        transform=axis.transAxes,
        fontsize=14.5,
        fontweight="bold",
        va="top",
    )


def accuracy_array(rows, method):
    return np.asarray(
        [
            [rows[(method, depth, seed)]["accuracy"] * 100 for depth in DEPTHS]
            for seed in SEEDS
        ],
        dtype=float,
    )


def mean_ci(values):
    mean = values.mean(axis=0)
    ci = 1.96 * values.std(axis=0, ddof=1) / math.sqrt(values.shape[0])
    return mean, ci


def plot_near_bp(axis, rows):
    x = np.arange(len(DEPTHS))
    specs = (
        ("bp", "Backprop", GRAY, "o", (0, (3, 2))),
        ("clean_kp", "Clean reciprocal", GREEN, "s", (0, (5, 2))),
        ("dynamic", "SDIL + 4× traffic", BLUE, "D", "-"),
    )
    for method, label, color, marker, linestyle in specs:
        values = accuracy_array(rows, method)
        for seed_values in values:
            axis.plot(x, seed_values, color=color, linewidth=0.65, alpha=0.12)
        mean, ci = mean_ci(values)
        axis.errorbar(
            x,
            mean,
            yerr=ci,
            color=color,
            marker=marker,
            markersize=5.5,
            linewidth=2.0,
            linestyle=linestyle,
            capsize=3,
            label=label,
            zorder=3,
        )
    dynamic = accuracy_array(rows, "dynamic")
    gain = (dynamic[:, -1] - dynamic[:, 0]).mean()
    axis.text(
        0.045,
        0.08,
        f"SDIL depth gain = {gain:+.3f} points\n"
        "all 5 paired seeds improve",
        transform=axis.transAxes,
        fontsize=8.6,
        color=BLUE,
        bbox={
            "boxstyle": "round,pad=0.34",
            "facecolor": "white",
            "edgecolor": "#D6D6D6",
        },
    )
    axis.set_xticks(x, [f"ResNet-{depth}" for depth in DEPTHS])
    axis.set_ylim(90.4, 94.0)
    axis.set_ylabel("CIFAR-10 test accuracy (%)")
    axis.set_title(
        "Useful depth without a BP gap",
        loc="left",
        fontsize=11.3,
        fontweight="bold",
        pad=10,
    )
    axis.legend(frameon=False, fontsize=8.1, loc="upper left", ncol=1)
    style_axis(axis)
    panel_label(axis, "a")


def plot_local_contrast(axis, rows):
    x = np.arange(len(DEPTHS))
    specs = (
        ("dfa", "Fixed DFA", ORANGE, "o"),
        ("clean_kp", "Clean reciprocal", GREEN, "s"),
        ("dynamic", "SDIL + 4× traffic", BLUE, "D"),
    )
    for method, label, color, marker in specs:
        values = accuracy_array(rows, method)
        mean, ci = mean_ci(values)
        axis.errorbar(
            x,
            mean,
            yerr=ci,
            color=color,
            marker=marker,
            markersize=5.5,
            linewidth=2.0,
            capsize=3,
            label=label,
            zorder=3,
        )
    dynamic = accuracy_array(rows, "dynamic")
    dfa = accuracy_array(rows, "dfa")
    advantages = (dynamic - dfa).mean(axis=0)
    axis.text(
        0.05,
        0.08,
        "SDIL − DFA:\n"
        + " / ".join(f"{value:+.1f}" for value in advantages)
        + " points",
        transform=axis.transAxes,
        fontsize=8.6,
        bbox={
            "boxstyle": "round,pad=0.34",
            "facecolor": "white",
            "edgecolor": "#D6D6D6",
        },
    )
    axis.set_xticks(x, [f"R{depth}" for depth in DEPTHS])
    axis.set_ylim(25, 96)
    axis.set_ylabel("CIFAR-10 test accuracy (%)")
    axis.set_title(
        "Local-learning contrast",
        loc="left",
        fontsize=11.3,
        fontweight="bold",
        pad=10,
    )
    axis.legend(frameon=False, fontsize=8.1, loc="center right")
    style_axis(axis)
    panel_label(axis, "b")


def plot_paired_gain(axis, rows):
    methods = ("bp", "dynamic")
    labels = ("Backprop", "SDIL + 4× traffic")
    colors = (GRAY, BLUE)
    markers = ("o", "D")
    gains = []
    for method in methods:
        values = accuracy_array(rows, method)
        gains.append(values[:, -1] - values[:, 0])
    for index, values in enumerate(gains):
        jitter = np.linspace(-0.055, 0.055, len(SEEDS))
        axis.scatter(
            np.full(len(SEEDS), index) + jitter,
            values,
            color=colors[index],
            marker=markers[index],
            s=34,
            alpha=0.72,
            edgecolor="white",
            linewidth=0.6,
            zorder=3,
        )
        mean = values.mean()
        ci = 1.96 * values.std(ddof=1) / math.sqrt(len(values))
        axis.errorbar(
            [index],
            [mean],
            yerr=[ci],
            color="#111111",
            marker=markers[index],
            markerfacecolor=colors[index],
            markersize=7,
            capsize=4,
            linewidth=1.5,
            zorder=5,
        )
        axis.text(
            index,
            mean + ci + 0.11,
            f"{mean:+.3f}",
            ha="center",
            va="bottom",
            fontsize=8.5,
            color=colors[index],
        )
    axis.axhline(0, color="#7A7A7A", linewidth=0.9, linestyle="--")
    axis.set_xlim(-0.45, 1.45)
    axis.set_ylim(-0.1, 1.85)
    axis.set_xticks(np.arange(2), labels)
    axis.set_ylabel("ResNet-56 − ResNet-20\naccuracy (points)")
    axis.set_title(
        "Paired added-depth benefit",
        loc="left",
        fontsize=11.3,
        fontweight="bold",
        pad=10,
    )
    axis.text(
        0.04,
        0.08,
        "Each point is one untouched seed.\n"
        "Every SDIL and BP pair is positive.",
        transform=axis.transAxes,
        fontsize=8.6,
    )
    style_axis(axis)
    panel_label(axis, "c")


def plot_alignment_cost(axis, rows):
    x = np.arange(len(DEPTHS))
    alignment = np.asarray(
        [
            [rows[("dynamic", depth, seed)]["early_alignment"] for depth in DEPTHS]
            for seed in SEEDS
        ],
        dtype=float,
    )
    mean, ci = mean_ci(alignment)
    alignment_handle = axis.errorbar(
        x,
        mean,
        yerr=ci,
        color=BLUE,
        marker="D",
        markersize=5.5,
        linewidth=2.0,
        capsize=3,
        label="Early-third alignment",
        zorder=3,
    )
    axis.set_xticks(x, [f"R{depth}" for depth in DEPTHS])
    axis.set_ylim(0.9992, 0.99982)
    axis.set_ylabel(
        r"$\cos(\mathrm{signal},-\nabla_h\mathcal{L})$",
        color=BLUE,
    )
    axis.tick_params(axis="y", colors=BLUE)
    axis.spines["left"].set_color(BLUE)

    bp_macs = np.asarray(
        [
            [rows[("bp", depth, seed)]["total_macs"] for depth in DEPTHS]
            for seed in SEEDS
        ],
        dtype=float,
    )
    dynamic_macs = np.asarray(
        [
            [rows[("dynamic", depth, seed)]["total_macs"] for depth in DEPTHS]
            for seed in SEEDS
        ],
        dtype=float,
    )
    mac_ratio = dynamic_macs / bp_macs
    cost_axis = axis.twinx()
    cost_mean, _ = mean_ci(mac_ratio)
    cost_handle = cost_axis.plot(
        x,
        cost_mean,
        color=GREEN,
        marker="s",
        markersize=5,
        linewidth=1.8,
        linestyle=(0, (5, 2)),
        label="MAC estimate / BP",
    )[0]
    cost_axis.set_ylim(1.29, 1.345)
    cost_axis.set_ylabel("SDIL MAC estimate / BP", color=GREEN)
    cost_axis.tick_params(axis="y", colors=GREEN)
    cost_axis.spines["top"].set_visible(False)
    cost_axis.spines["right"].set_color(GREEN)
    cost_axis.grid(False)

    axis.legend(
        [alignment_handle, cost_handle],
        ["Early-third alignment", "MAC estimate / BP"],
        frameon=False,
        fontsize=8.1,
        loc="lower left",
    )
    axis.text(
        0.04,
        0.93,
        "0 task-loss queries\n1 neutral observation / example",
        transform=axis.transAxes,
        fontsize=8.5,
        va="top",
        bbox={
            "boxstyle": "round,pad=0.34",
            "facecolor": "white",
            "edgecolor": "#D6D6D6",
        },
    )
    axis.set_title(
        "Credit and cost remain bounded",
        loc="left",
        fontsize=11.3,
        fontweight="bold",
        pad=10,
    )
    style_axis(axis)
    panel_label(axis, "d")
    return alignment, mac_ratio


def write_caption(path):
    caption = """# Standard-depth ResNet scaling figure caption

**Figure 6 | Dynamic somato-dendritic innovation scales across standard
ResNet depth.** All points are frozen CIFAR-10 test endpoints from seeds
10--14 after 200 epochs; error bars show 95% normal intervals over paired
seeds. The renderer validates all 60 records and their source hashes against
the passed predeclared gate. **a,** BP, clean reciprocal Kolen--Pollack credit,
and dynamic SDIL under four-times-RMS soma-predictable apical traffic all gain
accuracy from ResNet-20 to ResNet-56. SDIL gains 1.176 points, and all five
paired seeds improve. **b,** On the full accuracy scale, fixed direct feedback
alignment remains far below both reciprocal local learners at every depth;
the complete nine-method crossover is reported separately rather than inferred
from this four-method panel. **c,** Paired ResNet-56 minus ResNet-20 gains for
BP and SDIL. **d,** SDIL retains a mean early-third teaching-signal cosine
above 0.9994 while its hardware-independent affine-MAC estimate remains below
1.34 times matched BP. Exact gradients are diagnostic-only and never used for
learning. The SDIL rule uses no task-loss queries but does use one explicitly
counted instruction-off neutral observation per ordinary example.
"""
    with open(path, "w") as handle:
        handle.write(caption)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--d4_dir", default="results/kp_dynamic_projection_confirmation"
    )
    parser.add_argument(
        "--new_dir", default="results/oral_a_dynamic_scaling_v2"
    )
    parser.add_argument(
        "--gate", default="results/oral_a_dynamic_scaling_v2_gate.json"
    )
    parser.add_argument(
        "--d4_gate",
        default="results/kp_dynamic_projection_confirmation_gate.json",
    )
    parser.add_argument(
        "--bci_v2_gate",
        default="results/bci_v2_calibrated_confirmation_gate.json",
    )
    parser.add_argument("--outdir", default="results/figs")
    args = parser.parse_args()

    gate, records, rows, paths = load_panel(
        args.d4_dir,
        args.new_dir,
        args.gate,
        args.d4_gate,
        args.bci_v2_gate,
    )
    del records
    os.makedirs(args.outdir, exist_ok=True)
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 9.3,
            "axes.linewidth": 1.0,
            "xtick.major.width": 1.0,
            "ytick.major.width": 1.0,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )
    figure, axes = plt.subplots(2, 2, figsize=(10.4, 7.45))
    plot_near_bp(axes[0, 0], rows)
    plot_local_contrast(axes[0, 1], rows)
    plot_paired_gain(axes[1, 0], rows)
    alignment, mac_ratio = plot_alignment_cost(axes[1, 1], rows)
    figure.subplots_adjust(
        left=0.085,
        right=0.91,
        bottom=0.09,
        top=0.93,
        hspace=0.48,
        wspace=0.38,
    )

    stem = "figure6_standard_depth_scaling"
    pdf_path = os.path.join(args.outdir, f"{stem}.pdf")
    png_path = os.path.join(args.outdir, f"{stem}.png")
    caption_path = os.path.join(args.outdir, f"{stem}_caption.md")
    manifest_path = os.path.join(args.outdir, f"{stem}_manifest.json")
    figure.savefig(pdf_path, metadata=PDF_METADATA)
    figure.savefig(png_path, dpi=320)
    plt.close(figure)
    write_caption(caption_path)

    dynamic = accuracy_array(rows, "dynamic")
    bp = accuracy_array(rows, "bp")
    clean_kp = accuracy_array(rows, "clean_kp")
    dfa = accuracy_array(rows, "dfa")
    manifest = {
        "strict": True,
        "protocol": PROTOCOL,
        "gate_status": gate["status"],
        "complete_grid": True,
        "required_methods": list(METHODS),
        "required_depths": list(DEPTHS),
        "required_seeds": list(SEEDS),
        "record_count": len(paths),
        "statistics": {
            "test_accuracy_mean_percent": {
                method: {
                    str(depth): float(
                        accuracy_array(rows, method)[:, index].mean()
                    )
                    for index, depth in enumerate(DEPTHS)
                }
                for method in METHODS
            },
            "dynamic_d20_to_d56_gain_points": [
                float(value) for value in dynamic[:, -1] - dynamic[:, 0]
            ],
            "dynamic_d20_to_d56_gain_points_mean": float(
                (dynamic[:, -1] - dynamic[:, 0]).mean()
            ),
            "bp_d20_to_d56_gain_points": [
                float(value) for value in bp[:, -1] - bp[:, 0]
            ],
            "dynamic_minus_bp_points_mean": {
                str(depth): float((dynamic[:, index] - bp[:, index]).mean())
                for index, depth in enumerate(DEPTHS)
            },
            "dynamic_minus_clean_kp_points_mean": {
                str(depth): float(
                    (dynamic[:, index] - clean_kp[:, index]).mean()
                )
                for index, depth in enumerate(DEPTHS)
            },
            "dynamic_minus_dfa_points_mean": {
                str(depth): float((dynamic[:, index] - dfa[:, index]).mean())
                for index, depth in enumerate(DEPTHS)
            },
            "dynamic_early_alignment_mean": {
                str(depth): float(alignment[:, index].mean())
                for index, depth in enumerate(DEPTHS)
            },
            "dynamic_mac_ratio_to_bp_mean": {
                str(depth): float(mac_ratio[:, index].mean())
                for index, depth in enumerate(DEPTHS)
            },
            "dynamic_logical_task_loss_queries": sorted(
                {
                    rows[("dynamic", depth, seed)]["logical_queries"]
                    for depth in DEPTHS
                    for seed in SEEDS
                }
            ),
            "dynamic_neutral_observations_per_ordinary_example": sorted(
                {
                    load_json(paths[("dynamic", depth, seed)])["counters"][
                        "neutral_projection_examples"
                    ]
                    / load_json(paths[("dynamic", depth, seed)])["counters"][
                        "ordinary_examples"
                    ]
                    for depth in DEPTHS
                    for seed in SEEDS
                }
            ),
        },
        "sources": [
            {"path": args.gate, "sha256": sha256(args.gate)},
            {"path": args.d4_gate, "sha256": sha256(args.d4_gate)},
            {"path": args.bci_v2_gate, "sha256": sha256(args.bci_v2_gate)},
        ]
        + [
            {"path": path, "sha256": sha256(path)}
            for path in sorted(paths.values())
        ],
        "outputs": [
            os.path.basename(pdf_path),
            os.path.basename(png_path),
            os.path.basename(caption_path),
        ],
    }
    with open(manifest_path, "w") as handle:
        json.dump(manifest, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps(manifest["statistics"], indent=2, sort_keys=True))


if __name__ == "__main__":
    main()