"""Matched compiled-block validation of the operator finalists."""

import argparse
import copy
import json
import types
from pathlib import Path

import mlx.core as mx

from experiments.r3.common import paired
from experiments.r3.fusions import TunedFC2, qk_native_reduce
from h3_mlx.benchmark import error
from h3_mlx.config import SIGMAS
from h3_mlx.dit import DiT
from h3_mlx.layout import image_positions, rope_cache
from h3_mlx.memory import FootprintMonitor
from h3_mlx.ops import attention, modulate, rms_norm, rotate
from h3_mlx.runtime import configure_memory, environment


def main(a):
    configure_memory(wired_gb=14)
    report = {"environment": environment(), "trace": a.trace, "trials": []}
    with FootprintMonitor() as 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")
        native_fc2 = block.fc2
        control = mx.compile(
            lambda x, params, indices, rope: block(
                x, params=params, indices=indices, rope=rope, mlp_chunk=0
            )
        )
        ref = control(x, params, indices, rope)
        mx.eval(ref)
        xx = modulate(x, block.n1, params[0], params[1], indices, block.config.norm_eps)
        qkv = block.qkv(xx).reshape(-1, 3, 56, 128)
        qref = tuple(
            rotate(rms_norm(qkv[None, :, j], w, block.config.qk_norm_eps), *rope)
            for j, w in ((0, block.qn), (1, block.kn))
        )
        mx.eval(qkv, qref)
        for head_major in (False, True):

            def fn(head_major=head_major):
                return qk_native_reduce(
                    qkv,
                    block.qn,
                    block.kn,
                    *rope,
                    block.config.qk_norm_eps,
                    head_major=head_major,
                )

            result = fn()
            if head_major:
                result = tuple(t.transpose(0, 2, 1, 3) for t in result)
            row = {
                "group": "qk_native_reduce",
                "head_major": head_major,
                **error(mx.stack(qref), mx.stack(result)),
            }

            def native():
                return tuple(
                    rotate(
                        rms_norm(qkv[None, :, j], w, block.config.qk_norm_eps), *rope
                    )
                    for j, w in ((0, block.qn), (1, block.kn))
                )

            row.update(paired(native, fn, a.repeats))
            report["trials"].append(row)
            print(json.dumps(row), flush=True)
        block.fc2 = TunedFC2(native_fc2)
        candidate = mx.compile(
            lambda x, params, indices, rope: block(
                x, params=params, indices=indices, rope=rope, mlp_chunk=0
            )
        )
        row = {
            "group": "fc2_compiled_block",
            **error(ref, candidate(x, params, indices, rope)),
        }
        row.update(
            paired(
                lambda: control(x, params, indices, rope),
                lambda: candidate(x, params, indices, rope),
                a.repeats,
            )
        )
        report["trials"].append(row)
        print(json.dumps(row), flush=True)
        # Separate block objects keep native compiled closures independent.
        for tuned_fc2 in (False, True):
            fused = copy.copy(block)
            fused.fc2 = TunedFC2(native_fc2) if tuned_fc2 else native_fc2

            def fused_attn(
                self, xx, rope, query_chunk, kernel="native", projection_weights=None
            ):
                qq = self.qkv(xx).reshape(-1, 3, 56, 128)
                q, k = qk_native_reduce(
                    qq, self.qn, self.kn, *rope, self.config.qk_norm_eps
                )
                v = qq[None, :, 2]
                return self.out(
                    attention(q, k, v, query_chunk=query_chunk).reshape(xx.shape[0], -1)
                )

            fused.attn = types.MethodType(fused_attn, fused)
            runner = mx.compile(
                lambda x, params, indices, rope, fused=fused: fused(
                    x, params=params, indices=indices, rope=rope, mlp_chunk=0
                )
            )
            row = {
                "group": "qk_and_fc2_compiled_block"
                if tuned_fc2
                else "qk_compiled_block",
                **error(ref, runner(x, params, indices, rope)),
            }
            row.update(
                paired(
                    lambda: control(x, params, indices, rope),
                    lambda runner=runner: runner(x, params, indices, rope),
                    a.repeats,
                )
            )
            report["trials"].append(row)
            print(json.dumps(row), flush=True)
        model.clear()
    report["memory"] = memory.result()
    Path(a.output).write_text(json.dumps(report, indent=2) + "\n")


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=10)
    p.add_argument("--output", default="reports/optimization-r3/blocks.json")
    main(p.parse_args())