"""Isolate VAE tile batching from diffusion drift using one saved final latent."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import mlx.core as mx

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


def 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("--batches", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32])
    p.add_argument("--repeats", type=int, default=3)
    p.add_argument("--output", default="reports/vae")
    a = p.parse_args()
    configure_memory()
    root = Path(a.output)
    root.mkdir(parents=True, exist_ok=True)
    data = mx.load(a.trace)
    latent = data["latent"]
    mx.eval(latent)
    del data
    record = json.loads(Path(a.record).read_text())
    width, height = record["size"]
    model = VAE(a.vae)
    baseline = model.decode(latent)
    mx.eval(baseline)
    ref = to_image(baseline, width, height)
    trials = []
    for batch in a.batches:
        if batch < 1:
            p.error("tile batches must be positive")
        model.tile_batch = batch
        timer = measure(lambda: model.decode(latent), warmup=1, repeats=a.repeats)
        pixels = model.decode(latent)
        image = to_image(pixels, width, height)
        image.save(root / f"batch-{batch}.png")
        item = {
            "batch": batch,
            **timer,
            **image_metrics(ref, image),
            "max_float_error": mx.max(mx.abs(baseline - pixels)).item(),
        }
        trials.append(item)
        print(json.dumps(item), flush=True)
    model.clear()
    (root / "trials.json").write_text(
        json.dumps(
            {
                "environment": environment(),
                "source_trace": a.trace,
                "scope": "Same final latent; decoder only, loaded weights and one warmup excluded",
                "trials": trials,
            },
            indent=2,
        )
        + "\n"
    )


if __name__ == "__main__":
    main()