"""Build matched galleries, raw record exports and charts from completed trials."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from pathlib import Path

from h3_mlx.compare import make_comparison

LABELS = {
    "native": "BF16 native / six passes",
    "tuned": "BF16 / fused AdaLN / tile batch 8",
    "fused": "BF16 / fused AdaLN + QK-RoPE / tile batch 8",
    "q8": "DiT Q8 / six passes",
    "q4": "DiT Q4 / six passes",
    "preview": "BF16 preview / steps 0,2,4",
    "q8-g32": "DiT Q8 / group 32 / six passes",
    "chunked": "BF16 / MLP chunks 1024",
}


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--runs", default="outputs/suite")
    p.add_argument("--output", default="comparisons")
    p.add_argument(
        "--cuda-reference-root",
        help="Optional saved single-bf16-stream CUDA run; quality=lossless",
    )
    a = p.parse_args()
    runs, output = Path(a.runs), Path(a.output)
    records = []
    for case in sorted(d for d in runs.iterdir() if d.is_dir()):
        files = sorted(case.glob("*.png"))
        native = sorted(case.glob("native-*.png"))
        if not native:
            continue
        reference = native[-1]
        candidates = []
        labels = [LABELS["native"]]
        for name, label in LABELS.items():
            if name == "native":
                continue
            selected = sorted(case.glob(f"{name}-*.png"))
            if selected:
                candidates.append(selected[-1])
                labels.append(label)
        if candidates:
            crop = (1024, 384, 1536, 896) if case.name == "portrait" else None
            make_comparison(
                reference, candidates, output / case.name, labels=labels, crop=crop
            )
        for file in files:
            record = json.loads(file.with_suffix(".json").read_text())
            record["case"] = case.name
            record["profile"] = file.stem.rsplit("-", 1)[0]
            record["trial"] = int(file.stem.rsplit("-", 1)[1])
            records.append(record)
            target = (
                Path("reports/inference") / case.name / file.with_suffix(".json").name
            )
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(file.with_suffix(".json"), target)
        if a.cuda_reference_root:
            source = Path(a.cuda_reference_root)
            image = source / f"{case.name}.png"
            if image.exists():
                rows = [
                    json.loads(line)
                    for line in (source / "results.jsonl").read_text().splitlines()
                ]
                original = next(
                    r for r in rows if r["case"] == case.name and r["trial"] == 1
                )
                if hashlib.sha256(image.read_bytes()).hexdigest() != original["sha256"]:
                    raise ValueError(f"CUDA source hash mismatch: {image}")
                native_record = json.loads(reference.with_suffix(".json").read_text())
                if tuple(map(int, original["size"].split("x"))) != tuple(
                    native_record["size"]
                ):
                    continue
                target = output / "cuda-reference" / f"{case.name}.png"
                target.parent.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(image, target)
                r = {
                    k: native_record[k]
                    for k in (
                        "prompt",
                        "seed",
                        "size",
                        "canvas",
                        "rng",
                        "sigmas",
                        "steps",
                    )
                }
                # The original request script explicitly used CPU random noise.
                r.update(
                    output=str(target),
                    timings={"total_s": original["wall_seconds"]},
                    environment={
                        "device": "RTX 5090",
                        "backend": "SGLang CUDA",
                        "quality": original["quality"],
                    },
                    source_record=original,
                    reference_images=[],
                    reference_sha256=[],
                )
                r["dit"] = {"precision": "BF16"}
                r["vae"] = {"quality": original["quality"]}
                target.with_suffix(".json").write_text(json.dumps(r, indent=2) + "\n")
                make_comparison(
                    target,
                    [reference],
                    output / f"{case.name}-cuda",
                    labels=[
                        "CUDA BF16 / upstream lossless quality",
                        "MLX BF16 / native tile recipe",
                    ],
                )
    Path("reports").mkdir(exist_ok=True)
    Path("reports/image-trials.json").write_text(json.dumps(records, indent=2) + "\n")
    plot(records, output)


def plot(records, output):
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from PIL import Image

    from h3_mlx.compare import image_metrics

    product = [
        r for r in records if r["case"] == "product" and r.get("dit_was_resident")
    ]
    if not product:
        return
    order = [n for n in LABELS if any(r["profile"] == n for r in product)]
    selected = [
        max((r for r in product if r["profile"] == n), key=lambda r: r["trial"])
        for n in order
    ]
    baseline = max(
        (r for r in records if r["case"] == "product" and r["profile"] == "native"),
        key=lambda r: r["trial"],
    )
    ref = Image.open(baseline["output"])
    scores = [image_metrics(ref, Image.open(r["output"])) for r in selected]
    summary = []
    for row, score in zip(selected, scores, strict=True):
        summary.append(
            {
                "profile": row["profile"],
                "passes": len(row["steps"]),
                "denoise_s": sum(
                    v for k, v in row["timings"].items() if k.startswith("step_")
                ),
                "fresh_end_to_end_s": row["timings"].get("total_s"),
                "peak_mlx_gib": row["memory"]["peak_gb"],
                **score,
            }
        )
    Path("reports/tradeoffs.json").write_text(json.dumps(summary, indent=2) + "\n")
    table = [
        "| Profile | Passes | DiT only | Fresh end-to-end | Recorded peak MLX | RGB SSIM | RGB MAE / 255 |",
        "|---|---:|---:|---:|---:|---:|---:|",
    ]
    for row in summary:
        e2e = (
            "—"
            if row["fresh_end_to_end_s"] is None
            else f"{row['fresh_end_to_end_s']:.2f} s"
        )
        table.append(
            f"| {row['profile']} | {row['passes']} | {row['denoise_s']:.2f} s | {e2e} | {row['peak_mlx_gib']:.2f} GiB | {row['ssim_rgb']:.4f} | {row['mae_8bit']:.3f} |"
        )
    result_path = Path("reports/RESULTS.md")
    if result_path.exists():
        content = result_path.read_text()
        begin, end = "<!-- PRODUCT_TABLE_START -->", "<!-- PRODUCT_TABLE_END -->"
        if begin in content and end in content:
            before, rest = content.split(begin, 1)
            _, after = rest.split(end, 1)
            result_path.write_text(
                before + begin + "\n\n" + "\n".join(table) + "\n\n" + end + after
            )
    fig, axes = plt.subplots(1, 3, figsize=(14, 5), layout="constrained")
    names = [r["profile"] for r in selected]
    values = [
        [
            sum(v for k, v in r["timings"].items() if k.startswith("step_"))
            for r in selected
        ],
        [r["memory"]["peak_gb"] for r in selected],
        [m["ssim_rgb"] for m in scores],
    ]
    colors = [
        "#456dba" if n in ("native", "tuned", "fused") else "#d28c43" for n in names
    ]
    for ax, vs, title in zip(
        axes,
        values,
        [
            "DiT denoising only: seconds",
            "Recorded peak MLX allocation: GiB",
            "RGB SSIM to MLX native",
        ],
        strict=True,
    ):
        bars = ax.bar(names, vs, color=colors)
        ax.set_title(title)
        ax.tick_params(axis="x", rotation=45)
        ax.spines[["top", "right"]].set_visible(False)
        ax.bar_label(bars, fmt="%.3g", padding=3, fontsize=9)
        ax.set_ylim(0, max(vs) * 1.2)
    fig.suptitle(
        "M5 Pro 64 GB · 2048×2048 product · recorded DiT stages, corrected decoder for every image",
        fontsize=12,
    )
    fig.savefig(Path("reports") / "tradeoffs.png", dpi=160)
    plt.close(fig)


if __name__ == "__main__":
    main()