#!/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]


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, n, label, strict):
    missing = []
    for key in expected:
        count = len(groups.get(key, []))
        if count != n:
            missing.append(f"{key}: expected {n}, found {count}")
    if missing and strict:
        raise RuntimeError(f"incomplete {label}:\n" + "\n".join(missing))
    return missing


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 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)
    sg = scaling_groups(scale)
    miss_scale = require_cells(
        sg, [(m, d) for m in ("BP", "FA", "DFA", "SDIL") for d in DEPTHS],
        5, "CIFAR scaling", strict)
    miss_ep = require_cells(
        ep, [(m, d) for m in ("BP", "DFA", "SDIL", "EP") for d in (1, 2)],
        5, "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)
    groups = scaling_groups(rows)
    missing = require_cells(
        groups, [(m, d) for m in ("BP", "FA", "DFA", "SDIL") for d in DEPTHS],
        5, "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,
        "missing_cells": missing1 + missing2,
        "sources": [{"path": p, "sha256": file_hash(p)} for p in source_paths],
        "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()