"""One actual image for a Core ML device/batch, with NumPy bridge diagnostics."""

from __future__ import annotations

import argparse
import hashlib
import json
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from PIL import Image

from h3_mlx.compare import image_metrics, make_comparison
from h3_mlx.coreml_vae import CoreMLVAE
from h3_mlx.runtime import configure_memory
from h3_mlx.vae import to_image


def main(a):
    configure_memory(wired_gb=5)
    data = mx.load(a.trace)
    latent = data["latent"]
    mx.eval(latent)
    del data
    tick = time.perf_counter()
    model = CoreMLVAE(
        a.vae, package=a.package, tile_batch=a.batch, compute_units=a.units
    )
    load_s = time.perf_counter() - tick
    sample = (
        np.random.default_rng(731).normal(size=(a.batch, 24, 16, 16)).astype(np.float16)
    )
    raw = model.model.predict({"raw": sample})["pixels"]
    audit = {
        "output_shape": list(raw.shape),
        "output_strides": list(raw.strides),
        "output_c_contiguous": bool(raw.flags.c_contiguous),
        "original_mlx_numpy_import_exact": bool(
            np.array_equal(np.asarray(mx.array(raw)), raw, equal_nan=True)
        ),
    }
    source = json.loads(Path(a.reference).with_suffix(".json").read_text())
    tick = time.perf_counter()
    pixels = model.decode(latent)
    mx.eval(pixels)
    seconds = time.perf_counter() - tick
    image = to_image(pixels, *source["size"])
    root = Path(a.output)
    root.mkdir(parents=True, exist_ok=True)
    path = root / "image.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(
        validation_scope="Same final latent; decoder only",
        timing_scope="One complete decoder pass; model loading excluded",
        timings={"decode_s": seconds},
        load_s=load_s,
        vae_backend="coreml",
        vae={"batch": a.batch, "compute_units": a.units},
        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(),
        numpy_bridge=audit,
        **image_metrics(Image.open(a.reference), image),
    )
    path.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n")
    print(json.dumps(record), flush=True)
    make_comparison(
        a.reference,
        [path],
        root / "comparison",
        labels=[
            "Native BF16 decoder",
            f"Core ML {a.units}, batch {a.batch}, contiguous bridge",
        ],
    )
    model.clear()


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--vae", required=True)
    p.add_argument("--package", required=True)
    p.add_argument("--batch", type=int, default=8)
    p.add_argument("--units", default="CPU_AND_NE")
    p.add_argument(
        "--trace", default="outputs/suite/product/native-3.trace.safetensors"
    )
    p.add_argument("--reference", default="outputs/suite/product/native-3.png")
    p.add_argument("--output", required=True)
    main(p.parse_args())