"""Build the second optimization report from recorded timings and real PNGs."""

from __future__ import annotations

import json
import statistics
from pathlib import Path

from h3_mlx.compare import image_metrics, make_comparison

ROOT = Path("reports/optimization-v2")
OUTPUT = Path("outputs/optimization-v2")
GALLERIES = Path("comparisons/optimization-v2")
NAMES = {
    "quiet-baseline": "BF16 baseline (packing still active)",
    "quiet-native": "Exact optimized BF16, final quiet session",
    "quiet-packed": "Exact optimized packed BF16, final quiet session",
    "wired-baseline": "BF16 baseline",
    "wired-compiled": "Compiled BF16, retained refiners",
    "packed12": "Lossless packed, retained refiners",
    "prefetch": "Compiled BF16 + exact caches + prefetch",
    "compact": "Lossless packed + exact caches + prefetch",
    "fp16": "FP16 projection failure (non-finite)",
    "five-pass": "Five original passes (skip head 1)",
}


def latest(case, profile):
    files = sorted(
        (OUTPUT / case).glob(f"{profile}-*.png"),
        key=lambda p: int(p.stem.rsplit("-", 1)[1]),
    )
    return files[-1] if files else None


def main():
    from PIL import Image

    ROOT.mkdir(parents=True, exist_ok=True)
    ledger = [
        json.loads(line) for line in (OUTPUT / "trials.jsonl").read_text().splitlines()
    ]
    summary = []
    for case in ("product", "portrait", "native-4k"):
        reference = latest(case, "quiet-baseline") or latest(case, "wired-baseline")
        if reference is None:
            old = list((Path("outputs/suite") / case).glob("native-*.png"))
            reference = sorted(old)[-1] if old else None
        if reference is None:
            continue
        ref = Image.open(reference)
        for profile in NAMES:
            path = latest(case, profile)
            if path is None:
                continue
            group = [r for r in ledger if r["case"] == case and r["profile"] == profile]
            # If a suite was rerun with the same filenames, keep its latest row.
            group = list({r["output"]: r for r in group}.values())
            hot = [r for r in group if r["dit_was_resident"]]
            import numpy as np
            from safetensors import safe_open

            with safe_open(
                path.with_suffix(".trace.safetensors"), framework="numpy"
            ) as trace:
                finite = bool(np.isfinite(trace.get_tensor("latent")).all())
            row = {
                "valid_denoising": finite,
                "case": case,
                "profile": profile,
                "label": NAMES[profile],
                "image": str(path),
                "reference": str(reference),
                "resident_samples": len(hot),
                "load_included_s": [
                    r["timings"]["total_s"] for r in group if not r["dit_was_resident"]
                ],
                "resident_total_s": [r["timings"]["total_s"] for r in hot],
                "resident_dit_s": [
                    sum(v for k, v in r["timings"].items() if k.startswith("step_"))
                    for r in hot
                ],
                "peak_mlx_gib": max(r["memory"]["peak_gb"] for r in group),
                **image_metrics(ref, Image.open(path)),
            }
            row["resident_median_s"] = (
                statistics.median(row["resident_total_s"]) if hot else None
            )
            row["resident_dit_median_s"] = (
                statistics.median(row["resident_dit_s"]) if hot else None
            )
            summary.append(row)
        if case in ("product", "portrait"):
            candidates = []
            labels = ["Native six-pass BF16"]
            for profile in ("prefetch", "compact", "fp16", "five-pass"):
                selected = {"prefetch": "quiet-native", "compact": "quiet-packed"}.get(
                    profile, profile
                )
                path = latest(case, selected) or latest(case, profile)
                if path:
                    candidates.append(path)
                    labels.append(NAMES[profile])
            for profile, label in (
                ("q8", "Affine Q8 (earlier trial)"),
                ("q4", "Affine Q4 (earlier trial)"),
            ):
                paths = sorted((Path("outputs/suite") / case).glob(f"{profile}-*.png"))
                if paths:
                    candidates.append(paths[-1])
                    labels.append(label)
            make_comparison(reference, candidates, GALLERIES / case, labels=labels)
        elif latest(case, "compact"):
            make_comparison(
                reference,
                [latest(case, "compact")],
                GALLERIES / case,
                labels=["Native 4K six-pass BF16", "Optimized 4K lossless packed BF16"],
            )
    (ROOT / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    lines = [
        "| Case | Configuration | Warm samples | Warm median (s) | Peak MLX GiB | SSIM | Exact pixels |",
        "|---|---|---:|---:|---:|---:|---|",
    ]
    for r in summary:
        seconds = (
            "—" if r["resident_median_s"] is None else f"{r['resident_median_s']:.2f}"
        )
        if not r["valid_denoising"]:
            seconds = "invalid"
        similarity = f"{r['ssim_rgb']:.6f}" if r["valid_denoising"] else "invalid state"
        exact = (
            ("yes" if r["pixel_identical"] else "no")
            if r["valid_denoising"]
            else "non-finite"
        )
        lines.append(
            f"| {r['case']} | {r['label']} | {r['resident_samples']} | {seconds} | {r['peak_mlx_gib']:.2f} | {similarity} | {exact} |"
        )
    (ROOT / "table.md").write_text("\n".join(lines) + "\n")
    target = ROOT / "RESULTS.md"
    if target.exists():
        text = target.read_text()
        start = "<!-- RESULTS_TABLE -->"
        end = "<!-- END_RESULTS_TABLE -->"
        if start in text and end in text:
            text = (
                text.split(start)[0]
                + start
                + "\n\n"
                + "\n".join(lines)
                + "\n\n"
                + end
                + text.split(end)[1]
            )
            target.write_text(text)
    try:
        import matplotlib

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

        rows = [
            r
            for r in summary
            if r["case"] == "product"
            and r["resident_median_s"] is not None
            and r["profile"] in ("quiet-native", "quiet-packed")
        ]
        fig, ax = plt.subplots(figsize=(10, 6), layout="constrained")
        for r in rows:
            ax.scatter(
                r["resident_median_s"],
                r["peak_mlx_gib"],
                s=90,
                c=[r["ssim_rgb"]],
                vmin=0.6,
                vmax=1,
                cmap="viridis",
                edgecolor="black",
            )
            ax.annotate(
                r["profile"],
                (r["resident_median_s"], r["peak_mlx_gib"]),
                xytext=(5, 6),
                textcoords="offset points",
                fontsize=9,
            )
        ax.set(
            xlabel="Recorded warm end-to-end seconds (lower is faster)",
            ylabel="Peak MLX allocator GiB (lower is smaller)",
            title="M5 Pro 64 GB · 2048² product · exact and approximate options",
        )
        ax.grid(alpha=0.2)
        fig.savefig(ROOT / "tradeoffs.png", dpi=170)
        plt.close(fig)
    except ImportError:
        pass
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()