"""Generate model-free GDN boundary fixtures with the installed MTPLX/MLX.

Run with MTPLX's Python environment; this never loads or downloads a model.
Large-row math follows GatedDeltaNet.__call__ in the pinned qwen4_exp.py;
small rows call MTPLX's actual fused kernel. All files are diagnostic outputs.
"""

import argparse
import hashlib
import json
import statistics
import time
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import numpy as np
from mtplx.kernels.gdn_conv_norm import fused_gdn_conv_norm_rows
from mtplx.models import qwen4_exp


def delta_reference(output):
    from mlx_lm.models import gated_delta

    rng = np.random.default_rng(12345)
    initial = mx.array(rng.normal(0, 0.01, (1, 48, 128, 128)).astype(np.float32))
    np.array(initial).tofile(output / "delta-initial.f32")
    source = Path(gated_delta.__file__).resolve()
    receipt = {"mlx_version": mx.__version__, "reference_source": str(source),
               "reference_sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
               "cases": []}
    for rows in (1, 2, 3, 4, 5, 6, 7, 32, 2048):
        def normalized(scale):
            x = rng.normal(size=(1, rows, 16, 128)).astype(np.float32)
            return mx.array(x / np.linalg.norm(x, axis=-1, keepdims=True) * scale).astype(mx.bfloat16)

        q, k = normalized(128 ** -0.5), normalized(1)
        v = mx.array(rng.normal(0, 0.2, (1, rows, 48, 128)).astype(np.float32)).astype(mx.bfloat16)
        g = mx.array(rng.uniform(0.8, 0.999, (1, rows, 48)).astype(np.float32))
        beta = mx.array(rng.uniform(0.1, 0.9, (1, rows, 48)).astype(np.float32)).astype(mx.bfloat16)
        mx.eval(q, k, v, g, beta, initial)

        def forward():
            return gated_delta.gated_delta_kernel(q, k, v, g, beta, initial)

        mx.eval(*forward())
        elapsed = []
        for _ in range(10):
            start = time.perf_counter()
            result = forward()
            mx.eval(*result)
            elapsed.append((time.perf_counter() - start) * 1000)
        hashes = {}
        for name, array in zip(("q", "k", "v", "g", "beta", "out", "state"),
                               (q, k, v, g, beta, *result)):
            raw = np.array(array.astype(mx.float32)).astype("<f4").tobytes()
            (output / f"delta-{rows}-{name}.f32").write_bytes(raw)
            hashes[name] = hashlib.sha256(raw).hexdigest()
        case = {"rows": rows, "wall_ms": elapsed,
                "median_ms": statistics.median(elapsed), "sha256": hashes}
        receipt["cases"].append(case)
        print(json.dumps(case), flush=True)
    (output / "delta-reference.json").write_text(json.dumps(receipt, indent=2) + "\n")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("output", type=Path)
    parser.add_argument("--delta", action="store_true", help="Check the active scalar-gated recurrence boundary instead of Conv/Norm")
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    if args.delta:
        delta_reference(args.output)
        return
    width = 10240
    qk = 2048
    # Binary fractions keep fixture inputs exactly representable in BF16.
    values = lambda size, modulus: (
        (np.arange(size, dtype=np.int32) % modulus - modulus // 2).astype(np.float32)
        / 32
    )
    state_np = values(3 * width, 17).reshape(3, width)
    weight_np = values(width * 4, 13).reshape(width, 4)
    state = mx.array(state_np).astype(mx.bfloat16)
    conv = nn.Conv1d(width, width, 4, groups=width, bias=False)
    conv.weight = mx.array(weight_np[:, :, None]).astype(mx.bfloat16)
    state_np.tofile(args.output / "state.f32")
    weight_np.tofile(args.output / "weight.f32")
    receipt = {
        "mlx_version": mx.__version__,
        "reference_source": str(Path(qwen4_exp.__file__).resolve()),
        "reference_sha256": hashlib.sha256(Path(qwen4_exp.__file__).read_bytes()).hexdigest(),
        "cases": [],
    }
    for rows in (1, 2, 3, 4, 5, 6, 7, 32, 2048):
        x_np = values(rows * width, 23).reshape(rows, width)
        x = mx.array(x_np).astype(mx.bfloat16)
        mx.eval(x, state, conv.weight)

        def forward():
            if rows <= 6:
                return fused_gdn_conv_norm_rows(x, state, conv.weight)
            stream = mx.concatenate([state, x], axis=0)[None]
            activated = nn.silu(conv(stream))
            q, k, v = mx.split(activated, [qk, 2 * qk], axis=-1)

            def l2norm(a):
                a = a.reshape(1, rows, 16, 128)
                f = a.astype(mx.float32)
                return (f * mx.rsqrt((f * f).sum(-1, keepdims=True) + 1e-6)).astype(a.dtype)

            q = (128 ** -0.5) * l2norm(q)
            k = l2norm(k)
            return q, k, v, mx.contiguous(stream[0, -3:])

        mx.eval(*forward())  # Compile/warmup excluded from steady timing.
        elapsed = []
        for _ in range(10):
            start = time.perf_counter()
            result = forward()
            mx.eval(*result)
            elapsed.append((time.perf_counter() - start) * 1000)
        x_np.tofile(args.output / f"{rows}-input.f32")
        hashes = {}
        for name, array in zip(("q", "k", "v", "state"), result):
            raw = np.array(array.astype(mx.float32)).astype("<f4").tobytes()
            (args.output / f"{rows}-{name}.f32").write_bytes(raw)
            hashes[name] = hashlib.sha256(raw).hexdigest()
        case = {"rows": rows, "path": "fused" if rows <= 6 else "conv1d_silu_l2",
                "wall_ms": elapsed, "median_ms": statistics.median(elapsed), "sha256": hashes}
        receipt["cases"].append(case)
        print(json.dumps(case), flush=True)
    (args.output / "reference.json").write_text(json.dumps(receipt, indent=2) + "\n")


if __name__ == "__main__":
    main()