"""Fresh end-to-end timings after the final decoder arithmetic audit."""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path

from h3_mlx.runtime import Pipeline


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument("--text-encoder", required=True)
    a = p.parse_args()
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S"
    )
    case = next(
        c
        for c in json.loads(Path("experiments/cases.json").read_text())
        if c["id"] == "product"
    )
    pipe = Pipeline(a.model, a.text_encoder)
    try:
        for name, trial, kernel, batch in [
            ("native", 2, "native", 1),
            ("native", 3, "native", 1),
            ("fused", 2, "metal-qk", 8),
        ]:
            pipe.dit_options.update(kernel=kernel)
            pipe.vae_options.update(tile_batch=batch)
            if pipe.model is not None:
                pipe.model.kernel = kernel
            if pipe.vae is not None:
                pipe.vae.tile_batch = batch
            pipe.generate(
                case["prompt"],
                width=2048,
                height=2048,
                seed=case["seed"],
                output=f"outputs/suite/product/{name}-{trial}.png",
                save_trace=True,
            )
    finally:
        pipe.clear()


if __name__ == "__main__":
    main()