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

The renderer intentionally reads the ten untouched D4 records directly.  It
refuses incomplete seed panels, a failed gate, protocol drift, or disagreement
between the gate summary and the source records before writing any artifact.
"""
import argparse
import hashlib
import json
import math
import os
import statistics

import matplotlib

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


EXPECTED_SEEDS = tuple(range(10, 15))
PROTOCOL = "kp_dynamic_neutral_projection_confirmation_v1"
BLUE = "#0072B2"
LIGHT_BLUE = "#56B4E9"
GREEN = "#009E73"
GRAY = "#5A5A5A"
LIGHT_GRAY = "#B7B7B7"
PDF_METADATA = {
    "Creator": "SDIL audited D4 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 finite_leaves(value):
    if isinstance(value, bool) or value is None:
        return True
    if isinstance(value, (int, float)):
        return math.isfinite(float(value))
    if isinstance(value, dict):
        return all(finite_leaves(child) for child in value.values())
    if isinstance(value, (list, tuple)):
        return all(finite_leaves(child) for child in value)
    return True


def load_panel(input_dir, gate_path):
    gate = load_json(gate_path)
    require(gate.get("protocol") == PROTOCOL, "D4 protocol drift")
    require(gate.get("status") == "passed", "D4 gate is not passed")
    require(gate.get("review_score_after") == 7, "D4 score rule drift")
    require(
        gate.get("checks")
        and all(gate["checks"].values()),
        "D4 contains a failed check",
    )

    records = {}
    paths = []
    for seed in EXPECTED_SEEDS:
        for label, mode in (("clean_kp", "kp"), ("dynamic", "kp_traffic")):
            path = os.path.join(input_dir, f"seed{seed}_{label}.json")
            require(os.path.isfile(path), f"missing D4 record: {path}")
            row = load_json(path)
            paths.append(path)
            require(row.get("args", {}).get("seed") == seed, f"{path}: seed drift")
            require(row["args"].get("mode") == mode, f"{path}: mode drift")
            require(row["args"].get("depth") == 20, f"{path}: depth drift")
            require(row["args"].get("width") == 16, f"{path}: width drift")
            require(row["args"].get("epochs") == 200, f"{path}: epoch drift")
            require(row["args"].get("eval_split") == "test", f"{path}: split drift")
            require(
                row.get("provenance", {}).get("git_tracked_dirty") is False,
                f"{path}: tracked-dirty provenance",
            )
            require(
                row.get("final", {}).get("evaluation_split") == "test",
                f"{path}: final split drift",
            )
            require(row["final"].get("finite") is True, f"{path}: nonfinite final")
            require(
                len(row.get("epochs", [])) == 200,
                f"{path}: incomplete trajectory",
            )
            require(
                finite_leaves(
                    {
                        "final": row["final"],
                        "epochs": row["epochs"],
                        "diagnostics": row.get("diagnostics"),
                        "work": row.get("work"),
                    }
                ),
                f"{path}: nonfinite plotted source",
            )
            if mode == "kp_traffic":
                expected = {
                    "traffic_rule": "innovation",
                    "traffic_ratio": 4,
                    "neutral_projection": 1,
                    "predictor_mode": "closed_form",
                    "learn_P": 1,
                    "predictor_warmup_steps": 1,
                }
                for key, value in expected.items():
                    require(
                        row["args"].get(key) == value,
                        f"{path}: dynamic {key} drift",
                    )
                require(
                    row["work"].get("logical_batch_loss_queries") == 0,
                    f"{path}: nonzero task-loss queries",
                )
                require(
                    row["counters"].get("neutral_projection_examples")
                    == row["counters"].get("ordinary_examples"),
                    f"{path}: neutral observation count drift",
                )
            records[(seed, mode)] = row

    observed = sorted(
        name for name in os.listdir(input_dir) if name.endswith(".json")
    )
    expected_names = sorted(
        f"seed{seed}_{label}.json"
        for seed in EXPECTED_SEEDS
        for label in ("clean_kp", "dynamic")
    )
    require(observed == expected_names, "D4 directory contains an unexpected record set")

    clean = [records[(seed, "kp")]["final"]["accuracy"] for seed in EXPECTED_SEEDS]
    dynamic = [
        records[(seed, "kp_traffic")]["final"]["accuracy"]
        for seed in EXPECTED_SEEDS
    ]
    metrics = gate.get("metrics", {})
    require(
        np.allclose(metrics.get("accuracy_by_seed", {}).get("clean_kp"), clean),
        "D4 clean accuracy summary disagrees with records",
    )
    require(
        np.allclose(metrics.get("accuracy_by_seed", {}).get("dynamic"), dynamic),
        "D4 dynamic accuracy summary disagrees with records",
    )
    require(
        math.isclose(metrics.get("mean_accuracy", {}).get("clean_kp"), np.mean(clean)),
        "D4 clean mean disagrees with records",
    )
    require(
        math.isclose(metrics.get("mean_accuracy", {}).get("dynamic"), np.mean(dynamic)),
        "D4 dynamic mean disagrees with records",
    )
    return gate, records, paths


def mean_and_sd(values):
    array = np.asarray(values, dtype=float)
    return array.mean(axis=0), array.std(axis=0, ddof=1)


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=15,
        fontweight="bold",
        va="top",
    )


def plot_paired_accuracy(axis, records, gate):
    clean = np.asarray(
        [records[(seed, "kp")]["final"]["accuracy"] * 100 for seed in EXPECTED_SEEDS]
    )
    dynamic = np.asarray(
        [
            records[(seed, "kp_traffic")]["final"]["accuracy"] * 100
            for seed in EXPECTED_SEEDS
        ]
    )
    for index, seed in enumerate(EXPECTED_SEEDS):
        color = BLUE if dynamic[index] >= clean[index] else LIGHT_GRAY
        axis.plot([0, 1], [clean[index], dynamic[index]], color=color, alpha=0.55,
                  linewidth=1.4, zorder=1)
        axis.scatter([0], [clean[index]], s=34, facecolor="white", edgecolor=GRAY,
                     linewidth=1.0, zorder=2)
        axis.scatter([1], [dynamic[index]], s=38, facecolor=BLUE, edgecolor="white",
                     linewidth=0.8, zorder=3)
    means = [clean.mean(), dynamic.mean()]
    sems = [
        clean.std(ddof=1) / math.sqrt(len(clean)),
        dynamic.std(ddof=1) / math.sqrt(len(dynamic)),
    ]
    axis.errorbar(
        [0, 1],
        means,
        yerr=[1.96 * sem for sem in sems],
        fmt="D",
        markersize=7.5,
        color="#111111",
        markerfacecolor=[GREEN, BLUE][0],
        ecolor="#111111",
        elinewidth=1.5,
        capsize=4,
        zorder=5,
    )
    # Draw the second mean separately so its fill encodes the condition.
    axis.scatter([1], [means[1]], marker="D", s=58, facecolor=BLUE,
                 edgecolor="#111111", linewidth=0.8, zorder=6)
    axis.scatter([0], [means[0]], marker="D", s=58, facecolor=GREEN,
                 edgecolor="#111111", linewidth=0.8, zorder=6)
    delta = dynamic.mean() - clean.mean()
    upper_deficit = (
        gate["metrics"]["paired_deficit_one_sided_95pct_upper_bound"] * 100
    )
    axis.text(
        0.04,
        0.95,
        f"mean $\\Delta$ = {delta:+.3f} points\n"
        f"one-sided 95% deficit bound = {upper_deficit:.3f}",
        transform=axis.transAxes,
        va="top",
        fontsize=8.8,
        bbox={"boxstyle": "round,pad=0.35", "facecolor": "white",
              "edgecolor": "#D6D6D6"},
    )
    axis.set_xlim(-0.25, 1.25)
    axis.set_ylim(90.55, 92.02)
    axis.set_xticks([0, 1], ["Clean reciprocal\ncredit", "Innovation under\n4× traffic"])
    axis.set_ylabel("CIFAR-10 test accuracy (%)")
    axis.set_title("Untouched paired confirmation", loc="left", fontsize=11.5,
                   fontweight="bold", pad=10)
    style_axis(axis)
    panel_label(axis, "a")


def plot_layer_alignment(axis, records):
    raw = np.asarray(
        [
            records[(seed, "kp_traffic")]["diagnostics"][
                "raw_negative_gradient_cosine"
            ]
            for seed in EXPECTED_SEEDS
        ],
        dtype=float,
    )
    innovation = np.asarray(
        [
            records[(seed, "kp_traffic")]["diagnostics"][
                "used_negative_gradient_cosine"
            ]
            for seed in EXPECTED_SEEDS
        ],
        dtype=float,
    )
    require(raw.shape == innovation.shape == (5, 19), "D4 layer diagnostic drift")
    layers = np.arange(1, raw.shape[1] + 1)
    raw_mean, raw_sd = mean_and_sd(raw)
    innovation_mean, innovation_sd = mean_and_sd(innovation)
    axis.fill_between(layers, raw_mean - raw_sd, raw_mean + raw_sd,
                      color=LIGHT_GRAY, alpha=0.35, linewidth=0)
    axis.plot(layers, raw_mean, color=GRAY, linewidth=2.0, marker="o",
              markersize=3.2, label="Raw apical")
    axis.fill_between(
        layers,
        innovation_mean - innovation_sd,
        innovation_mean + innovation_sd,
        color=LIGHT_BLUE,
        alpha=0.25,
        linewidth=0,
    )
    axis.plot(layers, innovation_mean, color=BLUE, linewidth=2.2, marker="o",
              markersize=3.2, label="Innovation used")
    axis.axhline(0, color="#888888", linewidth=0.8, linestyle="--")
    axis.set_xlim(0.5, 19.5)
    axis.set_ylim(-0.06, 1.06)
    axis.set_xticks([1, 4, 7, 10, 13, 16, 19])
    axis.set_xlabel("Local credit layer")
    axis.set_ylabel(r"$\cos(\mathrm{signal},-\nabla_h\mathcal{L})$")
    axis.set_title("Residualization restores direction", loc="left",
                   fontsize=11.5, fontweight="bold", pad=10)
    axis.legend(frameon=False, loc="center left", bbox_to_anchor=(0.02, 0.56),
                fontsize=8.5)
    axis.text(
        0.04,
        0.78,
        "mean early-layer innovation cosine\n"
        f"= {innovation[:, :6].mean():.6f}",
        transform=axis.transAxes,
        fontsize=8.7,
        color=BLUE,
    )
    style_axis(axis)
    panel_label(axis, "b")


def plot_tracking(axis, records, gate):
    dynamic = np.asarray(
        [
            [
                epoch["feedback_tracking"]["mean_feedback_forward_cosine"]
                for epoch in records[(seed, "kp_traffic")]["epochs"]
            ]
            for seed in EXPECTED_SEEDS
        ],
        dtype=float,
    )
    clean = np.asarray(
        [
            [
                epoch["feedback_tracking"]["mean_feedback_forward_cosine"]
                for epoch in records[(seed, "kp")]["epochs"]
            ]
            for seed in EXPECTED_SEEDS
        ],
        dtype=float,
    )
    require(dynamic.shape == clean.shape == (5, 200), "D4 tracking trajectory drift")
    epochs = np.arange(1, 201)
    for values, color, label, alpha, linestyle, linewidth in (
        (dynamic, BLUE, "Innovation + traffic", 0.16, "-", 2.3),
        (clean, GREEN, "Clean reciprocal", 0.10, (0, (4, 2)), 1.8),
    ):
        mean, sd = mean_and_sd(values)
        axis.fill_between(epochs, mean - sd, mean + sd, color=color, alpha=alpha,
                          linewidth=0)
        axis.plot(epochs, mean, color=color, linewidth=linewidth, label=label,
                  linestyle=linestyle)
    axis.set_xlim(1, 200)
    axis.set_ylim(0.18, 1.02)
    axis.set_xticks([1, 50, 100, 150, 200])
    axis.set_xlabel("Training epoch")
    axis.set_ylabel("Feedback–forward cosine")
    axis.set_title("Credit tracking remains stable", loc="left",
                   fontsize=11.5, fontweight="bold", pad=10)
    axis.legend(frameon=False, loc="lower right", fontsize=8.4)
    metrics = gate["metrics"]
    dynamic_rows = [records[(seed, "kp_traffic")] for seed in EXPECTED_SEEDS]
    mac_ratio = statistics.mean(
        row["work"]["total_macs_estimate"] / metrics["bp_50k_mac_reference"]
        for row in dynamic_rows
    )
    peak_gib = statistics.mean(
        row["hardware"]["peak_memory_allocated_bytes"] / (1024 ** 3)
        for row in dynamic_rows
    )
    wall_ratio = statistics.mean(
        records[(seed, "kp_traffic")]["timing"]["total_timed_wall_s"]
        / records[(seed, "kp")]["timing"]["total_timed_wall_s"]
        for seed in EXPECTED_SEEDS
    )
    axis.text(
        0.04,
        0.50,
        f"{mac_ratio:.3f}× BP MAC estimate\n"
        "0 task-loss queries\n"
        "1 neutral observation / example\n"
        f"{peak_gib:.2f} GiB peak allocated\n"
        f"{wall_ratio:.2f}× clean-KP wall",
        transform=axis.transAxes,
        fontsize=8.6,
        va="top",
        bbox={"boxstyle": "round,pad=0.4", "facecolor": "white",
              "edgecolor": "#D6D6D6", "alpha": 0.96},
    )
    style_axis(axis)
    panel_label(axis, "c")


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

**Figure 4 | Somato-dendritic innovation survives a frozen standard-ResNet
confirmation.** All task endpoints are CIFAR-10 test results from the five
untouched seeds 10--14 on a standard ResNet-20; lines or bands show paired
seeds or mean ± sample standard deviation. **a,** Dynamic neutral-projection
innovation under four-times-RMS soma-predictable apical traffic matches clean
reciprocal Kolen--Pollack credit. Diamonds show means and 95% normal intervals.
Dynamic innovation gains 0.196 accuracy points on average; the predeclared
one-sided 95% upper bound on its deficit is 0.131 points. **b,** Final
layerwise cosine between the raw apical or innovation signal and the exact
negative hidden-state gradient.
Exact gradients are audit-only and never used for learning. Residualization
removes the directionally contaminating traffic while retaining an early-layer
cosine of 0.999687. **c,** Mean reciprocal-feedback tracking across all 200
epochs. Resource annotations report measured GTX 1080 wall time and peak
allocation, hardware-independent MAC estimates relative to the matched
50,000-example BP accounting reference, and logical task-loss queries. The
paired instruction-off neutral observation is explicitly counted and is not
claimed to be free or single-phase.
"""
    with open(path, "w") as handle:
        handle.write(caption)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_dir", default="results/kp_dynamic_projection_confirmation"
    )
    parser.add_argument(
        "--gate", default="results/kp_dynamic_projection_confirmation_gate.json"
    )
    parser.add_argument("--outdir", default="results/figs")
    args = parser.parse_args()

    gate, records, paths = load_panel(args.input_dir, args.gate)
    os.makedirs(args.outdir, exist_ok=True)

    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 9.5,
            "axes.linewidth": 1.0,
            "xtick.major.width": 1.0,
            "ytick.major.width": 1.0,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )
    figure, axes = plt.subplots(1, 3, figsize=(13.2, 4.05))
    plot_paired_accuracy(axes[0], records, gate)
    plot_layer_alignment(axes[1], records)
    plot_tracking(axes[2], records, gate)
    figure.subplots_adjust(left=0.068, right=0.985, bottom=0.20, top=0.84, wspace=0.37)

    pdf_path = os.path.join(args.outdir, "figure4_resnet_confirmation.pdf")
    png_path = os.path.join(args.outdir, "figure4_resnet_confirmation.png")
    caption_path = os.path.join(
        args.outdir, "figure4_resnet_confirmation_caption.md"
    )
    manifest_path = os.path.join(
        args.outdir, "figure4_resnet_confirmation_manifest.json"
    )
    figure.savefig(pdf_path, metadata=PDF_METADATA)
    figure.savefig(png_path, dpi=320)
    plt.close(figure)
    write_caption(caption_path)

    clean = np.asarray(
        [records[(seed, "kp")]["final"]["accuracy"] for seed in EXPECTED_SEEDS]
    )
    dynamic = np.asarray(
        [
            records[(seed, "kp_traffic")]["final"]["accuracy"]
            for seed in EXPECTED_SEEDS
        ]
    )
    manifest = {
        "strict": True,
        "protocol": PROTOCOL,
        "gate_status": gate["status"],
        "required_seeds": list(EXPECTED_SEEDS),
        "statistics": {
            "clean_kp_test_accuracy_mean": float(clean.mean()),
            "dynamic_test_accuracy_mean": float(dynamic.mean()),
            "dynamic_test_accuracy_min": float(dynamic.min()),
            "dynamic_minus_clean_points_mean": float(
                (dynamic - clean).mean() * 100
            ),
            "clean_minus_dynamic_one_sided_95pct_upper_points": float(
                gate["metrics"]["paired_deficit_one_sided_95pct_upper_bound"] * 100
            ),
            "dynamic_early_alignment_mean": float(
                gate["metrics"]["mean_dynamic_early_alignment"]
            ),
            "dynamic_mac_ratio_to_bp": float(
                statistics.mean(
                    records[(seed, "kp_traffic")]["work"]["total_macs_estimate"]
                    / gate["metrics"]["bp_50k_mac_reference"]
                    for seed in EXPECTED_SEEDS
                )
            ),
            "dynamic_peak_allocated_gib_mean": float(
                statistics.mean(
                    records[(seed, "kp_traffic")]["hardware"][
                        "peak_memory_allocated_bytes"
                    ]
                    / (1024 ** 3)
                    for seed in EXPECTED_SEEDS
                )
            ),
            "dynamic_wall_ratio_to_clean_kp_mean": float(
                statistics.mean(
                    records[(seed, "kp_traffic")]["timing"]["total_timed_wall_s"]
                    / records[(seed, "kp")]["timing"]["total_timed_wall_s"]
                    for seed in EXPECTED_SEEDS
                )
            ),
            "dynamic_logical_task_loss_queries": sorted(
                {
                    records[(seed, "kp_traffic")]["work"][
                        "logical_batch_loss_queries"
                    ]
                    for seed in EXPECTED_SEEDS
                }
            ),
            "dynamic_neutral_observations_per_ordinary_example": sorted(
                {
                    records[(seed, "kp_traffic")]["counters"][
                        "neutral_projection_examples"
                    ]
                    / records[(seed, "kp_traffic")]["counters"][
                        "ordinary_examples"
                    ]
                    for seed in EXPECTED_SEEDS
                }
            ),
        },
        "sources": [
            {"path": args.gate, "sha256": sha256(args.gate)}
        ]
        + [{"path": path, "sha256": sha256(path)} for path in sorted(paths)],
        "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()