#!/usr/bin/env python3
"""Build ME1a model-comparison charts from recorded JSON artifacts.

Usage: .venv/bin/python reports/260929/build_me1a_comparison_charts.py
Outputs: reports/260929/assets/me1a_{parameters,cpu_forward,weak_loss}.png
"""
from __future__ import annotations
import json
from pathlib import Path
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parents[2]
OUT = Path(__file__).resolve().parent / "assets"
SMOKE = ROOT / "experiments/ME1a_DT/results/models/report.json"
FIT = ROOT / "experiments/ME1a_DT/results/models/cad_predictions/report.json"

def load(path: Path):
    return json.loads(path.read_text(encoding="utf-8"))

def bar(labels, values, ylabel, title, output, *, log=False, footnote=""):
    fig, ax = plt.subplots(figsize=(10.5, 4.7), constrained_layout=True)
    colors = ["#1769aa" if i % 2 == 0 else "#3597c8" for i in range(len(labels))]
    bars = ax.bar(range(len(labels)), values, color=colors)
    ax.set_xticks(range(len(labels)), labels, rotation=32, ha="right")
    ax.set_ylabel(ylabel); ax.set_title(title, loc="left", weight="bold")
    ax.grid(axis="y", alpha=.22); ax.set_axisbelow(True)
    if log: ax.set_yscale("log")
    for rect, value in zip(bars, values):
        text = f"{value:,.0f}" if value >= 100 else f"{value:.3g}"
        ax.text(rect.get_x()+rect.get_width()/2, rect.get_height(), text, ha="center", va="bottom", fontsize=8)
    if footnote: fig.text(.01, .01, footnote, fontsize=8, color="#526273")
    fig.savefig(output, dpi=180); plt.close(fig)

def main():
    OUT.mkdir(exist_ok=True)
    smoke = load(SMOKE)["models"]
    wanted = [
        ("2D CNN/U-Net", "2D_CNN_UNet", "images/cad_multiview"),
        ("FNO2d", "FNO2d", "images/cad_multiview"),
        ("Dense 3D U-Net", "3D_UNet", "voxels/cad_surface_dense"),
        ("Sparse 3D U-Net", "Sparse_UNet", "voxels/cad_surface_sparse"),
        ("PointNet", "PointNet_segmentation", "points/cad_surface_points"),
        ("DGCNN", "DGCNN_dynamic_feature_knn", "points/cad_surface_points"),
        ("Mesh GNN", "GNN_MeshGraphNet_style", "graphs/graph_12000"),
        ("Risk GNN", "GNN_CAD_risk_deformation_multitask", "graphs/graph_12000"),
        ("UV-Net", "UVNet_face_grid_adjacency", "brep/faces_uv"),
        ("Tabular MLP", "tabular_surrogate", "tabular/process"),
    ]
    rows = []
    for label, family, source_hint in wanted:
        row = next(x for x in smoke if x["family"] == family and source_hint in x["source"])
        rows.append((label, row["params"], row["forward_latency_ms_cpu"]))
    labels = [r[0] for r in rows]
    bar(labels, [r[1] for r in rows], "parameters (log scale)", "ME1a actual-input smoke: model parameters",
        OUT / "me1a_parameters.png", log=True,
        footnote="Recorded architecture parameter counts; this is not a trained-model comparison.")
    bar(labels, [r[2] for r in rows], "CPU forward latency (ms, log scale)", "ME1a actual-input smoke: CPU forward",
        OUT / "me1a_cpu_forward.png", log=True,
        footnote="One recorded forward/backward smoke execution on each model's native input; hardware and domains differ.")
    fit = load(FIT)["models"]
    labels = list(fit)
    initial = [fit[k]["loss_initial"] for k in labels]
    final = [fit[k]["loss_final"] for k in labels]
    fig, ax = plt.subplots(figsize=(10.5, 4.7), constrained_layout=True)
    x = list(range(len(labels))); width = .38
    ax.bar([i-width/2 for i in x], initial, width, label="initial", color="#a7c9e2")
    ax.bar([i+width/2 for i in x], final, width, label="final", color="#1769aa")
    ax.set_xticks(x, labels, rotation=32, ha="right"); ax.set_ylabel("recorded weak-label training loss")
    ax.set_title("ME1a 14h06 weak-label diagnostic: loss change", loc="left", weight="bold")
    ax.legend(); ax.grid(axis="y", alpha=.22); ax.set_axisbelow(True)
    fig.text(.01, .01, "Losses are recorded in family-specific same-run provisional weak-label fits; architectures, domains and schedules differ. They are not accuracy or a comparable loss benchmark.", fontsize=8, color="#526273")
    fig.savefig(OUT / "me1a_weak_loss.png", dpi=180); plt.close(fig)

if __name__ == "__main__": main()