"""Benchmark full first DiT block with real saved conditioning/noise and weights."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import mlx.core as mx

from h3_mlx.benchmark import error, measure
from h3_mlx.config import SIGMAS
from h3_mlx.dit import DiT
from h3_mlx.layout import image_positions, rope_cache
from h3_mlx.runtime import configure_memory, environment


def main(a):
    configure_memory()
    data = mx.load(a.trace)
    model = DiT(a.model)
    model.prepare_modulation(SIGMAS)
    text = model.refine(data["conditioning"])
    x = mx.concatenate(
        (
            text,
            model._linear("video_patch_proj")(data["initial_rows"]).astype(mx.bfloat16),
        ),
        axis=0,
    )
    rope = rope_cache(
        image_positions(len(text), 128, 128), model._constant("rope.inv_freq")
    )
    indices = mx.concatenate((data["text_tags"], mx.zeros(4096, mx.int32)))
    params = tuple(model.modulation["block.0"][0, :, j] for j in range(6))
    block = model._block("blocks.0")
    dest = Path(a.output)
    dest.parent.mkdir(parents=True, exist_ok=True)
    report = {
        "environment": environment(),
        "trace": a.trace,
        "warmup": 2,
        "repeats": a.repeats,
        "trials": [],
    }
    ref = block(x, params=params, indices=indices, rope=rope, mlp_chunk=0)
    mx.eval(ref)
    for kernel in ("native", "metal", "metal-qk"):
        for chunk in (0, 512, 1024, 2048):
            for compiled in (False, True):
                if compiled and chunk:
                    continue  # Chunking intentionally synchronizes intermediates.

                def forward(x, params, indices, rope, chunk=chunk, kernel=kernel):
                    return block(
                        x,
                        params=params,
                        indices=indices,
                        rope=rope,
                        mlp_chunk=chunk,
                        kernel=kernel,
                    )

                fn = mx.compile(forward) if compiled else forward
                row = {"kernel": kernel, "mlp_chunk": chunk, "compiled": compiled}
                try:
                    row.update(
                        measure(
                            lambda fn=fn: fn(x, params, indices, rope),
                            repeats=a.repeats,
                        ),
                        **error(ref, fn(x, params, indices, rope)),
                    )
                except Exception as exc:
                    row["error"] = f"{type(exc).__name__}: {exc}"
                report["trials"].append(row)
                dest.write_text(json.dumps(report, indent=2) + "\n")
                print(json.dumps(row), flush=True)
    model.clear()


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument(
        "--trace", default="outputs/suite/product/native-3.trace.safetensors"
    )
    p.add_argument("--repeats", type=int, default=5)
    p.add_argument("--output", default="reports/optimization-v2/blocks.json")
    main(p.parse_args())