#!/usr/bin/env python3
"""Render the audited 27-cell Plain-CNN accuracy/time Pareto figure."""
import argparse
import hashlib
import json
import math
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.ticker import FuncFormatter, FixedLocator


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
METHODS = (
    "bp", "fa", "dfa", "pepita", "ff", "ep", "dualprop",
    "clean_kp", "sdil",
)
ARCHITECTURES = ("minicnn", "vgglike", "vgg16")
LABELS = {
    "bp": "BP",
    "fa": "FA",
    "dfa": "DFA",
    "pepita": "PEPITA",
    "ff": "Forward-Forward",
    "ep": "EP",
    "dualprop": "Dual Prop",
    "clean_kp": "clean KP",
    "sdil": "SDIL",
}
ARCH_LABELS = {
    "minicnn": "miniCNN · 3 trainable layers",
    "vgglike": "VGGlike · 5 trainable layers",
    "vgg16": "VGG16 · 16 trainable layers",
}
COLORS = {
    "bp": "#222222",
    "fa": "#009E73",
    "dfa": "#E69F00",
    "pepita": "#8C8C8C",
    "ff": "#56B4E9",
    "ep": "#CC79A7",
    "dualprop": "#7B61A8",
    "clean_kp": "#2E8B57",
    "sdil": "#D55E00",
}
MARKERS = {
    "bp": "*",
    "fa": "s",
    "dfa": "^",
    "pepita": "v",
    "ff": "P",
    "ep": "D",
    "dualprop": "h",
    "clean_kp": "X",
    "sdil": "o",
}
PDF_METADATA = {
    "Creator": "SDIL audited crossover figure pipeline",
    "Producer": "Matplotlib",
    "CreationDate": None,
    "ModDate": None,
}


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 read_json(path):
    with open(path, encoding="utf-8") as handle:
        return json.load(handle)


def finite(record):
    accuracy = record.get("best_validation_accuracy")
    wall = record.get("driver_wall_seconds")
    return (
        record.get("status") == "completed"
        and record.get("outcome") == "finite"
        and isinstance(accuracy, (int, float))
        and math.isfinite(float(accuracy))
        and isinstance(wall, (int, float))
        and math.isfinite(float(wall))
        and wall > 0
    )


def point(record):
    return (
        float(record["driver_wall_seconds"]) / 3600.0,
        float(record["best_validation_accuracy"]),
    )


def pareto(records):
    candidates = [record for record in records if finite(record)]
    frontier = []
    for record in candidates:
        wall, accuracy = point(record)
        dominated = any(
            other is not record
            and point(other)[0] <= wall
            and point(other)[1] >= accuracy
            and (
                point(other)[0] < wall
                or point(other)[1] > accuracy
            )
            for other in candidates
        )
        if not dominated:
            frontier.append(record)
    return sorted(frontier, key=lambda row: point(row)[0])


def validate(report):
    assert report["gate"] == "pass"
    assert report["stage"] == "plain_cnn_p2"
    assert report["num_expected_records"] == 27
    assert report["num_audited_records"] == 27
    assert report["missing_experiments"] == []
    assert report["test_policy"] == "none"
    records = report["records"]
    assert len(records) == 27
    actual = {
        (record["architecture"], record["method"]) for record in records
    }
    expected = {
        (architecture, method)
        for architecture in ARCHITECTURES for method in METHODS
    }
    assert actual == expected
    for record in records:
        wall = record.get("driver_wall_seconds")
        accuracy = record.get("best_validation_accuracy")
        assert isinstance(wall, (int, float)) and wall > 0
        assert (
            isinstance(accuracy, (int, float))
            and 0 <= accuracy <= 100
        )
    return records


def frontier_ids(records, include_bp):
    eligible = [
        record for record in records
        if include_bp or record["method"] != "bp"
    ]
    return [
        f"{record['architecture']}::{record['method']}"
        for record in pareto(eligible)
    ]


def render(records, outdir):
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 8.2,
        "axes.labelsize": 8.5,
        "axes.titlesize": 10.2,
        "legend.fontsize": 7.4,
        "xtick.labelsize": 7.4,
        "ytick.labelsize": 7.4,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "savefig.bbox": "tight",
    })
    figure, axes = plt.subplots(
        1, 3, figsize=(11.5, 4.15), sharey=True,
        gridspec_kw={"wspace": 0.12},
    )
    panels = {}
    for panel_index, (axis, architecture) in enumerate(
        zip(axes, ARCHITECTURES)
    ):
        subset = [
            record for record in records
            if record["architecture"] == architecture
        ]
        local_frontier = pareto([
            record for record in subset if record["method"] != "bp"
        ])
        overall_frontier = pareto(subset)
        panels[architecture] = {
            "local_frontier": [
                f"{architecture}::{record['method']}"
                for record in local_frontier
            ],
            "overall_frontier": [
                f"{architecture}::{record['method']}"
                for record in overall_frontier
            ],
        }
        axis.plot(
            [point(record)[0] for record in local_frontier],
            [point(record)[1] for record in local_frontier],
            color="#6B8E9B",
            linewidth=2.1,
            solid_capstyle="round",
            zorder=1,
        )
        axis.plot(
            [point(record)[0] for record in overall_frontier],
            [point(record)[1] for record in overall_frontier],
            color="#222222",
            linewidth=1.3,
            linestyle=(0, (3, 2)),
            zorder=2,
        )
        for record in subset:
            wall, accuracy = point(record)
            method = record["method"]
            is_finite = finite(record)
            marker = MARKERS[method] if is_finite else "x"
            size = 93 if method == "sdil" else (78 if method == "bp" else 56)
            linewidth = 1.8 if method == "sdil" else 1.0
            scatter_style = {
                "s": size,
                "marker": marker,
                "color": COLORS[method],
                "linewidth": linewidth,
                "alpha": 1.0 if is_finite else 0.82,
                "zorder": 5 if method == "sdil" else 4,
            }
            if marker != "x":
                scatter_style["edgecolor"] = "white"
            axis.scatter(
                wall,
                accuracy,
                **scatter_style,
            )
            if method == "sdil":
                axis.annotate(
                    "SDIL",
                    (wall, accuracy),
                    xytext=(5, -11 if architecture == "vgg16" else 6),
                    textcoords="offset points",
                    color=COLORS["sdil"],
                    fontweight="bold",
                    fontsize=7.5,
                    zorder=7,
                )
        local_ids = panels[architecture]["local_frontier"]
        sdil_is_frontier = f"{architecture}::sdil" in local_ids
        badge = (
            "SDIL on local frontier"
            if sdil_is_frontier
            else "clean KP dominates SDIL"
        )
        badge_color = "#E8F3F7" if sdil_is_frontier else "#F8E9E4"
        axis.text(
            0.035,
            0.96,
            badge,
            transform=axis.transAxes,
            ha="left",
            va="top",
            fontsize=7.2,
            fontweight="bold",
            color="#33444A" if sdil_is_frontier else "#8A3A22",
            bbox={
                "boxstyle": "round,pad=0.28",
                "facecolor": badge_color,
                "edgecolor": "none",
            },
        )
        axis.axhline(
            10,
            color="#A0A0A0",
            linewidth=0.75,
            linestyle=":",
            zorder=0,
        )
        if panel_index == 0:
            axis.text(
                0.035,
                0.125,
                "chance",
                transform=axis.transAxes,
                fontsize=6.8,
                color="#777777",
            )
        axis.set_xscale("log")
        axis.xaxis.set_major_locator(
            FixedLocator([0.01, 0.03, 0.1, 0.3, 1, 3, 10])
        )
        axis.xaxis.set_major_formatter(FuncFormatter(
            lambda value, _: f"{value:g}"
        ))
        axis.set_xlim(0.0075, 9.2)
        axis.set_ylim(0, 100)
        axis.set_xlabel("Measured training wall time (hours, log scale)")
        axis.set_title(
            f"{chr(97 + panel_index)}  {ARCH_LABELS[architecture]}",
            loc="left",
            fontweight="bold",
        )
        axis.grid(
            True,
            which="major",
            color="#D9D9D9",
            linewidth=0.55,
            alpha=0.68,
            zorder=0,
        )
        axis.grid(
            True,
            which="minor",
            axis="x",
            color="#EEEEEE",
            linewidth=0.4,
            alpha=0.55,
            zorder=0,
        )
    axes[0].set_ylabel("Best validation accuracy (%)")

    method_handles = [
        Line2D(
            [0], [0],
            marker=MARKERS[method],
            linestyle="none",
            markerfacecolor=COLORS[method],
            markeredgecolor="white",
            markeredgewidth=0.7,
            markersize=6.8 if method != "sdil" else 7.8,
            label=LABELS[method],
        )
        for method in METHODS
    ]
    line_handles = [
        Line2D(
            [0], [0],
            color="#6B8E9B",
            linewidth=2.1,
            label="local-method frontier (BP excluded)",
        ),
        Line2D(
            [0], [0],
            color="#222222",
            linewidth=1.3,
            linestyle=(0, (3, 2)),
            label="overall frontier (BP included)",
        ),
        Line2D(
            [0], [0],
            marker="x",
            color="#666666",
            linestyle="none",
            markersize=6,
            label="nonfinite trajectory",
        ),
    ]
    figure.legend(
        handles=method_handles + line_handles,
        loc="upper center",
        bbox_to_anchor=(0.5, 1.02),
        ncol=6,
        frameon=False,
        handlelength=2.2,
        columnspacing=1.25,
    )
    figure.suptitle(
        "Matched local learning: accuracy–time Pareto frontiers",
        x=0.5,
        y=1.15,
        fontsize=13.0,
        fontweight="bold",
    )
    figure.text(
        0.5,
        1.075,
        "Complete 27/27 validation panel · CIFAR-10 · seed 0 · "
        "single-GPU GTX 1080 timing",
        ha="center",
        fontsize=8.2,
        color="#555555",
    )
    os.makedirs(outdir, exist_ok=True)
    pdf_path = os.path.join(outdir, "figure7_plain_cnn_pareto.pdf")
    png_path = os.path.join(outdir, "figure7_plain_cnn_pareto.png")
    figure.savefig(pdf_path, metadata=PDF_METADATA)
    figure.savefig(png_path, dpi=320)
    plt.close(figure)
    return pdf_path, png_path, panels


def write_caption(path):
    caption = (
        "**Figure 7: Complete matched Plain-CNN accuracy–time crossover.** "
        "Best CIFAR-10 validation accuracy is plotted against measured "
        "single-GPU GTX-1080 training wall time for every registered method "
        "at miniCNN, VGGlike, and VGG16. Solid lines are empirical Pareto "
        "frontiers among non-backpropagation methods; dashed lines include "
        "BP as an optimization reference. Crosses retain nonfinite "
        "trajectories at their last finite validation metric. SDIL lies on "
        "the local-method frontier for miniCNN and VGGlike, while clean KP "
        "slightly dominates it at VGG16. Dual Propagation remains more "
        "accurate than SDIL at VGG16 but requires substantially more wall "
        "time. The figure supports local-method scaling and cost "
        "competitiveness, not global dominance over BP."
    )
    with open(path, "w", encoding="utf-8") as handle:
        handle.write(caption + "\n")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--audit",
        default=os.path.join(ROOT, "results", "plain_cnn_p2_audit.json"),
    )
    parser.add_argument(
        "--outdir",
        default=os.path.join(ROOT, "results", "figs"),
    )
    args = parser.parse_args()
    report = read_json(args.audit)
    records = validate(report)
    pdf_path, png_path, panels = render(records, args.outdir)

    expected_sdil_frontier = {
        "minicnn": True,
        "vgglike": True,
        "vgg16": False,
    }
    observed = {
        architecture:
            f"{architecture}::sdil"
            in panels[architecture]["local_frontier"]
        for architecture in ARCHITECTURES
    }
    assert observed == expected_sdil_frontier
    assert "vgg16::clean_kp" in panels["vgg16"]["local_frontier"]

    caption_path = os.path.join(
        args.outdir, "figure7_plain_cnn_pareto_caption.md"
    )
    write_caption(caption_path)
    script_path = os.path.abspath(__file__)
    manifest = {
        "audit_status": "passed",
        "figure": "figure7_plain_cnn_pareto",
        "source": {
            "audit_path": os.path.relpath(
                os.path.abspath(args.audit), ROOT
            ),
            "audit_sha256": sha256(args.audit),
            "script_path": os.path.relpath(script_path, ROOT),
            "script_sha256": sha256(script_path),
        },
        "num_expected_cells": 27,
        "num_audited_cells": len(records),
        "metric": "best_validation_accuracy",
        "cost": "driver_wall_seconds",
        "local_frontier_excludes": ["bp"],
        "panels": panels,
        "sdil_on_local_frontier": observed,
        "outputs": {
            os.path.basename(pdf_path): sha256(pdf_path),
            os.path.basename(png_path): sha256(png_path),
            os.path.basename(caption_path): sha256(caption_path),
        },
    }
    manifest_path = os.path.join(
        args.outdir, "figure7_plain_cnn_pareto_manifest.json"
    )
    with open(manifest_path, "w", encoding="utf-8") as handle:
        json.dump(manifest, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps({
        "audit_status": "passed",
        "num_audited_cells": len(records),
        "sdil_on_local_frontier": observed,
        "png": png_path,
        "pdf": pdf_path,
    }, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()