"""Compare native, compiled and Core ML decoders on the exact same final latent."""

from __future__ import annotations

import argparse
import gc
import hashlib
import json
from pathlib import Path

import mlx.core as mx

from h3_mlx.benchmark import measure
from h3_mlx.compare import image_metrics, make_comparison
from h3_mlx.coreml_vae import CoreMLVAE
from h3_mlx.runtime import configure_memory, environment
from h3_mlx.vae import VAE, to_image


def main(a):
    configure_memory(wired_gb=5)
    root = Path(a.output)
    root.mkdir(parents=True, exist_ok=True)
    data = mx.load(a.trace)
    latent = data["latent"]
    mx.eval(latent)
    del data
    source = json.loads(Path(a.record).read_text())
    width, height = source["size"]
    results = []
    images = []
    labels = []
    reference = None
    variants = [
        ("native-b1", VAE, {"tile_batch": 1}),
        ("native-b4", VAE, {"tile_batch": 4}),
        ("compiled-b4", VAE, {"tile_batch": 4, "compile_blocks": True}),
        ("native-b8", VAE, {"tile_batch": 8}),
        ("compiled-b8", VAE, {"tile_batch": 8, "compile_blocks": True}),
    ]
    if a.include_fp32:
        variants.insert(
            0,
            (
                "fp32-b4-compiled",
                VAE,
                {"tile_batch": 4, "compile_blocks": True, "dtype": "float32"},
            ),
        )
    if a.package:
        variants.extend(
            (
                "coreml-" + unit,
                CoreMLVAE,
                {
                    "tile_batch": a.coreml_batch,
                    "package": a.package,
                    "compute_units": unit,
                },
            )
            for unit in a.units
        )
    for name, cls, options in variants:
        model = cls(a.vae, **options)
        timer = measure(
            lambda model=model: model.decode(latent), warmup=1, repeats=a.repeats
        )
        pixels = model.decode(latent)
        image = to_image(pixels, width, height)
        if reference is None:
            reference = image
        path = root / f"{name}.png"
        image.save(path)
        record = {
            k: source[k]
            for k in (
                "prompt",
                "seed",
                "size",
                "canvas",
                "rng",
                "reference_images",
                "reference_sha256",
                "sigmas",
                "steps",
            )
            if k in source
        }
        record.update(
            vae=options,
            vae_backend="coreml" if cls is CoreMLVAE else "mlx",
            environment=environment(),
            validation_scope="Same final latent; decoder only",
            timing_scope="Warm decoder only; model loading excluded",
            timings={"decode_s": timer["median_ms"] / 1000},
            source_trace=a.trace,
            source_trace_sha256=hashlib.sha256(Path(a.trace).read_bytes()).hexdigest(),
            output=str(path),
            png_sha256=hashlib.sha256(path.read_bytes()).hexdigest(),
        )
        path.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n")
        row = {"variant": name, **timer, **image_metrics(reference, image)}
        results.append(row)
        images.append(path)
        labels.append(name)
        print(json.dumps(row), flush=True)
        (root / "trials.json").write_text(
            json.dumps(
                {
                    "source_trace": a.trace,
                    "scope": "Same final latent, decoder-only hot latency",
                    "trials": results,
                },
                indent=2,
            )
            + "\n"
        )
        model.clear()
        del model
        gc.collect()
        mx.clear_cache()
    make_comparison(images[0], images[1:], root / "comparison", labels=labels)


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--vae", required=True)
    p.add_argument("--trace", required=True)
    p.add_argument("--record", required=True)
    p.add_argument("--package")
    p.add_argument(
        "--include-fp32",
        action="store_true",
        help="Use the original FP32 decoder as the same-latent reference",
    )
    p.add_argument("--coreml-batch", type=int, default=8)
    p.add_argument("--units", nargs="+", default=["CPU_AND_GPU"])
    p.add_argument("--repeats", type=int, default=3)
    p.add_argument("--output", default="comparisons/optimization-v2/decoder-product")
    main(p.parse_args())