#!/usr/bin/env python3
"""Generate the two audited, publication-facing SDIL figures.

Strict mode refuses to render a final figure unless every required cell has
five clean-provenance seeds. Use ``--allow-incomplete`` only for layout
previews while long experiments are still running.
"""
import argparse
import glob
import hashlib
import json
import math
import os
import statistics

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter


COLORS = {
    "BP": "#4D4D4D",
    "FA": "#009E73",
    "DFA": "#D55E00",
    "SDIL": "#0072B2",
    "EP": "#CC79A7",
}
MARKERS = {"BP": "*", "FA": "s", "DFA": "^", "SDIL": "o", "EP": "D"}
DEPTHS = [5, 10, 20, 30, 60]
EXPECTED_SEEDS = set(range(5))


def read_many(patterns):
    paths = sorted({p for pattern in patterns for p in glob.glob(pattern)})
    rows = []
    for path in paths:
        with open(path) as f:
            row = json.load(f)
        row["_path"] = path
        rows.append(row)
    return rows


def method(row):
    args = row["args"]
    return (args.get("mode") or args.get("method")).upper()


def mean_sd(values):
    return (statistics.mean(values),
            statistics.stdev(values) if len(values) > 1 else 0.0)


def require_clean(rows):
    bad = [r["_path"] for r in rows
           if r.get("provenance", {}).get("git_dirty") is not False]
    if bad:
        raise RuntimeError("dirty or unknown provenance:\n" + "\n".join(bad))


def require_cells(groups, expected, label, strict):
    missing = []
    for key in expected:
        rows = groups.get(key, [])
        seeds = [r["args"].get("seed") for r in rows]
        if len(seeds) != len(set(seeds)):
            missing.append(f"{key}: duplicate seeds {sorted(seeds)}")
        elif set(seeds) != EXPECTED_SEEDS:
            missing.append(
                f"{key}: expected seeds {sorted(EXPECTED_SEEDS)}, found {sorted(seeds)}")
    if missing and strict:
        raise RuntimeError(f"incomplete {label}:\n" + "\n".join(missing))
    return missing


def require_value(errors, row, key, actual, expected):
    if actual != expected:
        errors.append(f"{row['_path']}: {key} expected {expected!r}, found {actual!r}")


def require_valid_finals(rows):
    errors = []
    for row in rows:
        final = row.get("final", {})
        acc, wall = final.get("test_acc"), final.get("wall_s")
        if not isinstance(acc, (int, float)) or not math.isfinite(acc) or not 0 <= acc <= 1:
            errors.append(f"{row['_path']}: invalid final test_acc {acc!r}")
        if not isinstance(wall, (int, float)) or not math.isfinite(wall) or wall <= 0:
            errors.append(f"{row['_path']}: invalid final wall_s {wall!r}")
    if errors:
        raise RuntimeError("invalid result values:\n" + "\n".join(errors))


def require_scaling_protocol(rows):
    errors = []
    for row in rows:
        args = row["args"]
        for key, expected in (("dataset", "cifar10"), ("width", 64),
                              ("epochs", 5), ("batch_size", 128),
                              ("act", "tanh"), ("residual", 1),
                              ("device", "cuda")):
            require_value(errors, row, key, args.get(key), expected)
        name = method(row)
        require_value(errors, row, "depth", args.get("depth"),
                      args.get("depth") if args.get("depth") in DEPTHS else "one of " + repr(DEPTHS))
        require_value(errors, row, "eta", args.get("eta"), 0.02 if name == "SDIL" else 0.05)
        if name == "SDIL":
            require_value(errors, row, "pert_mode", args.get("pert_mode"), "simultaneous")
            require_value(errors, row, "pert_ndirs", args.get("pert_ndirs"), 16)
    if errors:
        raise RuntimeError("mixed CIFAR scaling protocols:\n" + "\n".join(errors))


def require_ep_protocol(groups):
    errors = []
    dynamics = {
        1: {"epochs": 25, "beta": 0.5, "free_steps": 20,
            "nudge_steps": 4, "learning_rates": [0.1, 0.05]},
        2: {"epochs": 60, "beta": 1.0, "free_steps": 150,
            "nudge_steps": 6, "learning_rates": [0.4, 0.1, 0.01]},
    }
    for (name, depth), rows in groups.items():
        for row in rows:
            args = row["args"]
            for key, expected in (("dataset", "mnist"), ("width", 500),
                                  ("depth", depth), ("train_examples", 50000),
                                  ("device", "cuda"),
                                  ("epochs", dynamics[depth]["epochs"])):
                require_value(errors, row, key, args.get(key), expected)
            if name == "EP":
                require_value(errors, row, "batch_size", args.get("batch_size"), 20)
                require_value(errors, row, "ep_persistent", args.get("ep_persistent"), 1)
                resolved = row.get("resolved_protocol", {})
                for key in ("beta", "free_steps", "nudge_steps", "learning_rates"):
                    require_value(errors, row, f"resolved_protocol.{key}",
                                  resolved.get(key), dynamics[depth][key])
                require_value(errors, row, "persistent_particles",
                              resolved.get("persistent_particles"), True)
            else:
                require_value(errors, row, "batch_size", args.get("batch_size"), 128)
                require_value(errors, row, "act", args.get("act"), "tanh")
                require_value(errors, row, "residual", args.get("residual"), 0)
                require_value(errors, row, "eta", args.get("eta"), 0.05)
                if name == "SDIL":
                    require_value(errors, row, "pert_mode", args.get("pert_mode"), "simultaneous")
                    require_value(errors, row, "pert_ndirs", args.get("pert_ndirs"), 16)
    if errors:
        raise RuntimeError("mixed exact-EP protocols:\n" + "\n".join(errors))


def scaling_records(root):
    return read_many([
        os.path.join(root, "scale_v3_cifar10_*.json"),
        os.path.join(root, "fa_scale_v2_cifar10_*.json"),
    ])


def scaling_groups(rows):
    groups = {}
    for row in rows:
        groups.setdefault((method(row), row["args"]["depth"]), []).append(row)
    return groups


def ep_groups(root):
    rows = read_many([
        os.path.join(root, "ep_original_v2_*.json"),
        os.path.join(root, "ep_archmatch_v1_*.json"),
        os.path.join(root, "ep_depth2_v1_*.json"),
    ])
    groups = {}
    for row in rows:
        args = row["args"]
        # The d3 near-parameter runs are intentionally excluded: this panel is
        # exact forward-architecture comparison only.
        if args["depth"] in (1, 2) and args["width"] == 500:
            groups.setdefault((method(row), args["depth"]), []).append(row)
    return rows, groups


def nondominated(points):
    """Return increasing-time points that improve accuracy."""
    front, best = [], -math.inf
    for point in sorted(points, key=lambda p: (p[0], -p[1])):
        if point[1] > best:
            front.append(point)
            best = point[1]
    return front


def cell_summary(rows, include_alignment=False):
    if not rows:
        return None
    acc = [100 * r["final"]["test_acc"] for r in rows]
    wall = [r["final"]["wall_s"] for r in rows]
    acc_mean, acc_sd = mean_sd(acc)
    wall_mean, wall_sd = mean_sd(wall)
    summary = {
        "n": len(rows),
        "seeds": sorted(r["args"]["seed"] for r in rows),
        "accuracy_percent_mean": acc_mean,
        "accuracy_percent_sd": acc_sd,
        "wall_s_mean": wall_mean,
        "wall_s_sd": wall_sd,
    }
    if include_alignment:
        alignments = [v for r in rows if (v := early_alignment(r)) is not None]
        if alignments:
            summary["early_alignment_mean"], summary["early_alignment_sd"] = mean_sd(
                alignments)
    return summary


def audited_statistics(root):
    scale = scaling_records(root)
    sg = scaling_groups(scale)
    scaling = {}
    frontier_candidates = []
    for name in ("BP", "FA", "DFA", "SDIL"):
        for depth in DEPTHS:
            summary = cell_summary(sg.get((name, depth), []), include_alignment=name != "BP")
            if summary is None:
                continue
            scaling[f"{name}_d{depth}"] = summary
            if name != "BP":
                frontier_candidates.append((summary["wall_s_mean"],
                                            summary["accuracy_percent_mean"],
                                            name, depth))
    front = nondominated(frontier_candidates)

    _, eg = ep_groups(root)
    ep_exact = {}
    paired = {}
    for name in ("BP", "DFA", "SDIL", "EP"):
        for depth in (1, 2):
            summary = cell_summary(eg.get((name, depth), []))
            if summary is not None:
                ep_exact[f"{name}_d{depth}"] = summary
    for depth in (1, 2):
        sdil = {r["args"]["seed"]: r for r in eg.get(("SDIL", depth), [])}
        ep = {r["args"]["seed"]: r for r in eg.get(("EP", depth), [])}
        seeds = sorted(set(sdil) & set(ep))
        if seeds:
            gaps = [100 * (sdil[s]["final"]["test_acc"] - ep[s]["final"]["test_acc"])
                    for s in seeds]
            speedups = [ep[s]["final"]["wall_s"] / sdil[s]["final"]["wall_s"]
                        for s in seeds]
            paired[f"SDIL_minus_EP_d{depth}"] = {
                "seeds": seeds,
                "accuracy_point_gap_mean": mean_sd(gaps)[0],
                "accuracy_point_gap_sd": mean_sd(gaps)[1],
                "wall_speedup_geomean": math.exp(statistics.mean(math.log(v)
                                                                  for v in speedups)),
                "wall_speedups": speedups,
            }
    return {
        "cifar_scaling": scaling,
        "local_pareto_frontier": [
            {"method": p[2], "depth": p[3], "wall_s_mean": p[0],
             "accuracy_percent_mean": p[1]} for p in front
        ],
        "sdil_is_high_accuracy_frontier_endpoint": bool(front and front[-1][2] == "SDIL"),
        "mnist_exact_ep": ep_exact,
        "paired_ep_comparison": paired,
    }


def setup_style():
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 9,
        "axes.titlesize": 10.5,
        "axes.labelsize": 9.5,
        "axes.linewidth": 0.8,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "xtick.direction": "out",
        "ytick.direction": "out",
        "legend.frameon": False,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
        "savefig.bbox": "tight",
    })


def error_point(ax, x, y, xerr, yerr, name, marker=None, alpha=1.0,
                label=None, zorder=3):
    ax.errorbar(
        x, y, xerr=xerr, yerr=yerr, fmt=marker or MARKERS[name],
        color=COLORS[name], markerfacecolor=COLORS[name],
        markeredgecolor="white", markeredgewidth=0.65, markersize=7,
        elinewidth=1.0, capsize=2.2, alpha=alpha,
        label=label, zorder=zorder)


def plot_tradeoff(root, outdir, strict):
    scale = scaling_records(root)
    ep_rows, ep = ep_groups(root)
    require_clean(scale + ep_rows)
    require_valid_finals(scale + ep_rows)
    sg = scaling_groups(scale)
    require_scaling_protocol(scale)
    require_ep_protocol(ep)
    miss_scale = require_cells(
        sg, [(m, d) for m in ("BP", "FA", "DFA", "SDIL") for d in DEPTHS],
        "CIFAR scaling", strict)
    miss_ep = require_cells(
        ep, [(m, d) for m in ("BP", "DFA", "SDIL", "EP") for d in (1, 2)],
        "exact EP comparison", strict)

    fig, axes = plt.subplots(1, 2, figsize=(7.35, 3.15),
                             gridspec_kw={"wspace": 0.30})
    ax = axes[0]
    frontier_candidates = []
    for name in ("BP", "FA", "DFA", "SDIL"):
        first = True
        for depth in DEPTHS:
            rows = sg.get((name, depth), [])
            if not rows:
                continue
            acc, acc_sd = mean_sd([100 * r["final"]["test_acc"] for r in rows])
            wall, wall_sd = mean_sd([r["final"]["wall_s"] for r in rows])
            error_point(ax, wall, acc, wall_sd, acc_sd, name,
                        label=name if first else None,
                        alpha=0.55 if name == "BP" else 0.92,
                        zorder=2 if name == "BP" else 3)
            ax.annotate(f"d{depth}", (wall, acc), xytext=(3, 3),
                        textcoords="offset points", fontsize=6.7,
                        color=COLORS[name], alpha=0.85)
            if name != "BP":
                frontier_candidates.append((wall, acc, name, depth))
            first = False
    front = nondominated(frontier_candidates)
    if len(front) > 1:
        ax.plot([p[0] for p in front], [p[1] for p in front],
                color="#222222", linewidth=1.15, alpha=0.75, zorder=1)
        ax.fill_between([p[0] for p in front], [p[1] for p in front],
                        min(p[1] for p in front) - 0.5,
                        color="#BDBDBD", alpha=0.10, zorder=0)
    ax.set_xscale("log")
    ax.xaxis.set_major_formatter(FuncFormatter(lambda value, _: f"{value:g}"))
    ax.set_xlabel("Wall time (s, log scale)")
    ax.set_ylabel("CIFAR-10 test accuracy (%)")
    ax.set_title("a  Local-learning Pareto frontier", loc="left", fontweight="bold")
    ax.grid(True, which="major", color="#D9D9D9", linewidth=0.55, alpha=0.65)
    ax.legend(ncol=2, loc="lower right", handletextpad=0.4, columnspacing=0.8)

    ax = axes[1]
    for name in ("EP", "SDIL"):
        xs, ys = [], []
        for depth in (1, 2):
            rows = ep.get((name, depth), [])
            if not rows:
                continue
            acc, acc_sd = mean_sd([100 * r["final"]["test_acc"] for r in rows])
            wall, wall_sd = mean_sd([r["final"]["wall_s"] for r in rows])
            error_point(ax, wall, acc, wall_sd, acc_sd, name,
                        marker="D" if depth == 1 else "o",
                        label=f"{name}, d{depth}")
            xs.append(wall); ys.append(acc)
        if len(xs) == 2:
            ax.plot(xs, ys, color=COLORS[name], linewidth=1.0, alpha=0.55)
    ax.set_xscale("log")
    ax.xaxis.set_major_formatter(FuncFormatter(lambda value, _: f"{value:g}"))
    ax.set_xlabel("Wall time (s, log scale)")
    ax.set_ylabel("MNIST test accuracy (%)")
    ax.set_title("b  Against canonical EP (exact architecture)",
                 loc="left", fontweight="bold")
    ax.grid(True, which="major", color="#D9D9D9", linewidth=0.55, alpha=0.65)
    ax.legend(loc="best", fontsize=8)
    fig.savefig(os.path.join(outdir, "figure1_pareto.pdf"))
    fig.savefig(os.path.join(outdir, "figure1_pareto.png"), dpi=320)
    plt.close(fig)
    return scale + ep_rows, miss_scale + miss_ep


def early_alignment(row):
    values = (row.get("final", {}).get("cos_r_negg")
              or row.get("final", {}).get("cos_fa_negg"))
    if not values:
        return None
    n = max(1, len(values) // 3)
    finite = [v for v in values[:n] if math.isfinite(v)]
    return statistics.mean(finite) if finite else None


def plot_scaling(root, outdir, strict):
    rows = scaling_records(root)
    require_clean(rows)
    require_valid_finals(rows)
    require_scaling_protocol(rows)
    groups = scaling_groups(rows)
    missing = require_cells(
        groups, [(m, d) for m in ("BP", "FA", "DFA", "SDIL") for d in DEPTHS],
        "CIFAR scaling", strict)
    fig, axes = plt.subplots(1, 3, figsize=(7.35, 2.75),
                             gridspec_kw={"wspace": 0.38})

    ax = axes[0]
    for name in ("BP", "FA", "DFA", "SDIL"):
        xs, ys, es = [], [], []
        for depth in DEPTHS:
            rs = groups.get((name, depth), [])
            if not rs:
                continue
            mu, sd = mean_sd([100 * r["final"]["test_acc"] for r in rs])
            xs.append(depth); ys.append(mu); es.append(sd)
        ax.errorbar(xs, ys, yerr=es, color=COLORS[name], marker=MARKERS[name],
                    markersize=5.5, linewidth=1.4, capsize=2.2, label=name,
                    markeredgecolor="white", markeredgewidth=0.5)
    ax.set_xlabel("Hidden depth")
    ax.set_ylabel("Test accuracy (%)")
    ax.set_title("a  Accuracy", loc="left", fontweight="bold")
    ax.grid(True, color="#D9D9D9", linewidth=0.55, alpha=0.65)
    ax.legend(fontsize=7.5, ncol=2, handlelength=1.5,
              columnspacing=0.7, handletextpad=0.35)

    ax = axes[1]
    bp_by_seed = {(r["args"]["depth"], r["args"]["seed"]): r
                  for r in rows if method(r) == "BP"}
    for name in ("FA", "DFA", "SDIL"):
        xs, ys, es = [], [], []
        for depth in DEPTHS:
            gaps = []
            for r in groups.get((name, depth), []):
                bp = bp_by_seed.get((depth, r["args"]["seed"]))
                if bp:
                    gaps.append(100 * (bp["final"]["test_acc"]
                                       - r["final"]["test_acc"]))
            if gaps:
                mu, sd = mean_sd(gaps)
                xs.append(depth); ys.append(mu); es.append(sd)
        ax.errorbar(xs, ys, yerr=es, color=COLORS[name], marker=MARKERS[name],
                    markersize=5.5, linewidth=1.4, capsize=2.2, label=name,
                    markeredgecolor="white", markeredgewidth=0.5)
    ax.axhline(0, color="#777777", linewidth=0.8, linestyle="--")
    ax.set_xlabel("Hidden depth")
    ax.set_ylabel("Gap to paired BP (points)")
    ax.set_title("b  Optimization gap", loc="left", fontweight="bold")
    ax.grid(True, color="#D9D9D9", linewidth=0.55, alpha=0.65)

    ax = axes[2]
    for name in ("FA", "DFA", "SDIL"):
        xs, ys, es = [], [], []
        for depth in DEPTHS:
            vals = [v for r in groups.get((name, depth), [])
                    if (v := early_alignment(r)) is not None]
            if vals:
                mu, sd = mean_sd(vals)
                xs.append(depth); ys.append(mu); es.append(sd)
        ax.errorbar(xs, ys, yerr=es, color=COLORS[name], marker=MARKERS[name],
                    markersize=5.5, linewidth=1.4, capsize=2.2, label=name,
                    markeredgecolor="white", markeredgewidth=0.5)
    ax.axhline(0, color="#777777", linewidth=0.8, linestyle="--")
    ax.set_xlabel("Hidden depth")
    ax.set_ylabel("Early-third cos(signal, −gradient)")
    ax.set_title("c  Credit alignment", loc="left", fontweight="bold")
    ax.grid(True, color="#D9D9D9", linewidth=0.55, alpha=0.65)
    fig.savefig(os.path.join(outdir, "figure2_scaling.pdf"))
    fig.savefig(os.path.join(outdir, "figure2_scaling.png"), dpi=320)
    plt.close(fig)
    return rows, missing


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


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--results", default="results")
    parser.add_argument("--outdir", default="results/figs")
    parser.add_argument("--allow-incomplete", action="store_true")
    args = parser.parse_args()
    strict = not args.allow_incomplete
    os.makedirs(args.outdir, exist_ok=True)
    setup_style()
    rows1, missing1 = plot_tradeoff(args.results, args.outdir, strict)
    rows2, missing2 = plot_scaling(args.results, args.outdir, strict)
    source_paths = sorted({r["_path"] for r in rows1 + rows2})
    manifest = {
        "strict": strict,
        "required_seeds": sorted(EXPECTED_SEEDS),
        "protocols": {
            "cifar_scaling": "width64 residual tanh; 5 epochs; d5/10/20/30/60",
            "ep_comparison": "exact d1/d2 width500; canonical EP dynamics",
        },
        "missing_cells": missing1 + missing2,
        "sources": [{"path": p, "sha256": file_hash(p)} for p in source_paths],
        "statistics": audited_statistics(args.results),
        "outputs": ["figure1_pareto.pdf", "figure1_pareto.png",
                    "figure2_scaling.pdf", "figure2_scaling.png"],
    }
    with open(os.path.join(args.outdir, "main_figure_manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)
    print(f"rendered figures from {len(source_paths)} audited result files; "
          f"missing cells={len(manifest['missing_cells'])}")


if __name__ == "__main__":
    main()