"""Measure exact palette12 packing versus quantized and decoded-weight GEMM."""

from __future__ import annotations

import argparse
import gc
import json
import time
from pathlib import Path

import mlx.core as mx

from h3_mlx.benchmark import error, measure
from h3_mlx.packed import PackedLinear
from h3_mlx.runtime import configure_memory, environment
from h3_mlx.weights import Linear, WeightStore


def main(a):
    configure_memory(wired_gb=8)
    dest = Path(a.output)
    dest.parent.mkdir(parents=True, exist_ok=True)
    report = {
        "environment": environment(),
        "warmup": 2,
        "repeats": a.repeats,
        "scope": "Original block-0 weights; synthetic BF16 activations. Includes per-call GPU weight reconstruction, excludes initial packing.",
        "trials": [],
    }
    store = WeightStore(Path(a.model) / "transformer")
    mx.random.seed(731)
    for name in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"):
        layer = Linear.load(
            store,
            "blocks.0." + name,
            grouped_heads=(56, 128) if name == "attn.qkv_proj" else None,
        )
        w = layer.weight
        x = mx.random.normal((a.tokens, w.shape[1])).astype(mx.bfloat16)
        ref = layer(x)
        mx.eval(ref, x)
        tick = time.perf_counter()
        packed = PackedLinear.from_linear(layer)
        pack_s = time.perf_counter() - tick
        exact = mx.all(
            packed.weight.unpack().view(mx.uint16) == w.view(mx.uint16)
        ).item()
        if not exact:
            raise AssertionError("Lossless weight reconstruction failed")
        functions = [
            ("bf16", lambda layer=layer, x=x: layer(x), {}),
            (
                "packed12",
                lambda packed=packed, x=x: packed(x),
                {
                    "pack_s": pack_s,
                    "weight_bytes": packed.nbytes,
                    "ratio": packed.nbytes / layer.nbytes,
                    "weight_bits_identical": exact,
                    "escape_fraction": packed.weight.escape_indices.size / w.size,
                },
            ),
        ]
        # Explicit state inputs prevent compile from retaining a dense constant.
        def packed_forward(x, state, packed=packed):
            return packed(x, state=state)
        fn = mx.compile(packed_forward)
        state = packed.weight.state()
        functions.append(("packed12_compiled", lambda fn=fn, x=x, state=state: fn(x, state), {}))
        for variant, call, extra in functions:
            row = {
                "operator": name,
                "variant": variant,
                "shape": [a.tokens, w.shape[1], w.shape[0]],
                **extra,
                **measure(call, repeats=a.repeats),
                **error(ref, call()),
            }
            report["trials"].append(row)
            dest.write_text(json.dumps(report, indent=2) + "\n")
            print(json.dumps(row), flush=True)
        del functions, fn, packed
        for bits in (8, 6, 4):
            q, s, z = mx.quantize(w, bits=bits, group_size=32)
            mx.eval(q, s, z)

            def decoded(q=q, s=s, z=z, bits=bits, x=x, layer=layer):
                d = mx.dequantize(q, s, z, bits=bits, group_size=32)
                return x @ d.T if layer.bias is None else mx.addmm(layer.bias, x, d.T)

            def fused(q=q, s=s, z=z, bits=bits, x=x, layer=layer):
                y = mx.quantized_matmul(x, q, s, z, bits=bits, group_size=32)
                return y if layer.bias is None else y + layer.bias

            for variant, call in (
                (
                    f"q{bits}_fused",
                    fused,
                ),
                (f"q{bits}_dequant_gemm", decoded),
            ):
                row = {
                    "operator": name,
                    "variant": variant,
                    "shape": [a.tokens, w.shape[1], w.shape[0]],
                    **measure(call, repeats=a.repeats),
                    **error(ref, call()),
                }
                report["trials"].append(row)
                dest.write_text(json.dumps(report, indent=2) + "\n")
                print(json.dumps(row), flush=True)
            del q, s, z, decoded
        del layer, w, x, ref
        gc.collect()
        mx.clear_cache()


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument("--tokens", type=int, default=4181)
    p.add_argument("--repeats", type=int, default=7)
    p.add_argument("--output", default="reports/optimization-v2/storage.json")
    main(p.parse_args())