"""Actual H3 projection dimensions, M5 layouts, and new quantization formats.

Run as a module from the repository root. Each result is persisted immediately.
The inputs are synthetic; weights come from the original checkpoint.
"""

from __future__ import annotations

import argparse
import gc
import json
from pathlib import Path

import mlx.core as mx

from h3_mlx.benchmark import error, measure
from h3_mlx.ops import swiglu
from h3_mlx.runtime import configure_memory, environment
from h3_mlx.weights import Linear, WeightStore


def run(args):
    configure_memory()
    mx.random.seed(731)
    dest = Path(args.output)
    dest.parent.mkdir(parents=True, exist_ok=True)
    report = {
        "environment": environment(),
        "warmup": 2,
        "repeats": args.repeats,
        "scope": "Original block-0 weights; synthetic BF16 inputs. Operator latency only, excludes conversion and loading.",
        "trials": [],
    }

    def record(op, variant, fn, ref, **extra):
        try:
            timing = measure(fn, repeats=args.repeats)
            row = {
                "operator": op,
                "variant": variant,
                **extra,
                **timing,
                **error(ref, fn()),
            }
        except Exception as exc:
            row = {
                "operator": op,
                "variant": variant,
                **extra,
                "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)

    store = WeightStore(Path(args.model) / "transformer")
    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, bias = layer.weight, layer.bias
        x = mx.random.normal((args.tokens, w.shape[1])).astype(mx.bfloat16)
        mx.eval(x, w)
        ref = layer(x)
        mx.eval(ref)
        shape = [args.tokens, w.shape[1], w.shape[0]]
        record(name, "bf16_nt", lambda layer=layer, x=x: layer(x), ref, shape=shape)
        wt = mx.contiguous(w.T)
        mx.eval(wt)
        fn = (
            (lambda x=x, wt=wt: x @ wt)
            if bias is None
            else (lambda bias=bias, x=x, wt=wt: mx.addmm(bias, x, wt))
        )
        record(name, "bf16_nn", fn, ref, shape=shape)
        del wt, fn
        for pad in (64, 128, 256):
            n = ((args.tokens + pad - 1) // pad) * pad
            padded = mx.pad(x, ((0, n - args.tokens), (0, 0)))
            mx.eval(padded)
            record(
                name,
                f"bf16_pad{pad}",
                lambda layer=layer, padded=padded: layer(padded)[: args.tokens],
                ref,
                shape=shape,
                padded_tokens=n,
            )
            del padded
        half, halfw = x.astype(mx.float16), w.astype(mx.float16)
        halfb = None if bias is None else bias.astype(mx.float16)
        mx.eval(half, halfw, halfb)
        fn = (
            (lambda half=half, halfw=halfw: (half @ halfw.T).astype(mx.bfloat16))
            if bias is None
            else (
                lambda halfb=halfb, half=half, halfw=halfw: mx.addmm(
                    halfb, half, halfw.T
                ).astype(mx.bfloat16)
            )
        )
        record(name, "fp16_round_bf16", fn, ref, shape=shape)
        del half, halfw, halfb, fn
        for mode, bits, group in (
            ("mxfp8", 8, 32),
            ("mxfp4", 4, 32),
            ("nvfp4", 4, 16),
            ("affine", 8, 32),
            ("affine", 6, 32),
            ("affine", 4, 32),
        ):
            try:
                packed = mx.quantize(w, mode=mode, bits=bits, group_size=group)
                mx.eval(packed)

                def quant(
                    x=x, packed=packed, mode=mode, bits=bits, group=group, bias=bias
                ):
                    y = mx.quantized_matmul(
                        x, *packed, mode=mode, bits=bits, group_size=group
                    )
                    return y if bias is None else y + bias

                record(
                    name,
                    f"{mode}{bits}_g{group}",
                    quant,
                    ref,
                    shape=shape,
                    weight_gib=sum(v.nbytes for v in packed) / 1024**3,
                )
                del packed, quant
            except Exception as exc:
                record(
                    name,
                    f"{mode}{bits}_g{group}",
                    lambda exc=exc: (_ for _ in ()).throw(exc),
                    ref,
                )
        del x, w, bias, layer, ref
        gc.collect()
        mx.clear_cache()
    a = Linear.load(store, "blocks.0.mlp.fc1")
    b = Linear.load(store, "blocks.0.mlp.fc2")
    x = mx.random.normal((args.tokens, 5376)).astype(mx.bfloat16)

    def mlp(x):
        return b(swiglu(a(x)))

    ref = mlp(x)
    mx.eval(ref)
    record("mlp", "eager", lambda: mlp(x), ref)
    compiled = mx.compile(mlp)
    record("mlp", "compiled", lambda: compiled(x), ref)


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=5)
    p.add_argument("--output", default="reports/optimization-v2/operators.json")
    run(p.parse_args())