"""Synchronized operator experiments on real H3 dimensions, with numerical error."""

from __future__ import annotations

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

import mlx.core as mx

from .kernels import norm_modulate, qk_norm_rope
from .ops import attention, modulate, rms_norm, rotate
from .runtime import configure_memory, environment


def measure(fn, *, warmup=2, repeats=7):
    for _ in range(warmup):
        mx.eval(fn())
    samples = []
    for _ in range(repeats):
        mx.synchronize()
        tick = time.perf_counter()
        mx.eval(fn())
        mx.synchronize()
        samples.append((time.perf_counter() - tick) * 1000)
    return {
        "median_ms": statistics.median(samples),
        "min_ms": min(samples),
        "max_ms": max(samples),
        "samples_ms": samples,
    }


def error(reference, candidate):
    a, b = reference.astype(mx.float32), candidate.astype(mx.float32)
    diff = a - b
    return {
        "max_abs": mx.max(mx.abs(diff)).item(),
        "relative_rmse": mx.sqrt(
            mx.mean(diff * diff) / (mx.mean(a * a) + 1e-30)
        ).item(),
        "equal_fraction": mx.mean((a == b).astype(mx.float32)).item(),
    }


def run(output, *, tokens=(256, 4096), repeats=7):
    configure_memory()
    mx.random.seed(731)
    result = {
        "environment": environment(),
        "warmup": 2,
        "repeats": repeats,
        "scope": "Synthetic operator microbenchmarks; excludes loading/quantization. Not image latency.",
        "trials": [],
    }

    def record(group, variant, shape, fn, reference):
        data = {
            "operator": group,
            "variant": variant,
            "shape": shape,
            **measure(fn, repeats=repeats),
            **error(reference, fn()),
        }
        result["trials"].append(data)
        print(json.dumps(data), flush=True)
        # Persist each completed trial so interrupted sweeps retain evidence.
        Path(output).parent.mkdir(parents=True, exist_ok=True)
        Path(output).write_text(json.dumps(result, indent=2) + "\n")

    for n in tokens:
        x = mx.random.normal((n, 5376)).astype(mx.bfloat16)
        w = mx.random.normal((5376,)).astype(mx.bfloat16)
        shift = (mx.random.normal((3, 5376)) * 0.1).astype(mx.bfloat16)
        scale = (mx.random.normal((3, 5376)) * 0.1).astype(mx.bfloat16)
        ids = mx.arange(n, dtype=mx.int32) % 3
        mx.eval(x, w, shift, scale, ids)
        ref = modulate(x, w, shift, scale, ids)
        mx.eval(ref)
        record(
            "rms_adaln",
            "native_rounding",
            [n, 5376],
            lambda x=x, w=w, shift=shift, scale=scale, ids=ids: modulate(
                x, w, shift, scale, ids
            ),
            ref,
        )
        for threads in (128, 256, 512):
            record(
                "rms_adaln",
                f"metal_tg{threads}",
                [n, 5376],
                lambda x=x, w=w, shift=shift, scale=scale, ids=ids, threads=threads: (
                    norm_modulate(x, w, shift, scale, ids, threads=threads)
                ),
                ref,
            )
        qkv = mx.random.normal((n, 3, 56, 128)).astype(mx.bfloat16)
        qw = mx.random.normal((128,)).astype(mx.bfloat16)
        angles = mx.random.normal((1, n, 1, 48))
        angles = mx.concatenate((angles, angles), axis=-1)
        cos, sin = (
            mx.cos(angles).astype(mx.bfloat16),
            mx.sin(angles).astype(mx.bfloat16),
        )

        def native_qk(qkv=qkv, qw=qw, cos=cos, sin=sin):
            return tuple(
                rotate(rms_norm(qkv[None, :, i], qw), cos, sin) for i in (0, 1)
            )

        def fused_qk(qkv=qkv, qw=qw, cos=cos, sin=sin):
            return qk_norm_rope(qkv, qw, qw, cos, sin)

        qr, kr = native_qk()
        both = mx.concatenate((qr, kr), axis=0)
        mx.eval(both)
        record(
            "qk_norm_rope",
            "native",
            [n, 56, 128],
            lambda fn=native_qk: mx.concatenate(fn(), axis=0),
            both,
        )
        record(
            "qk_norm_rope",
            "metal",
            [n, 56, 128],
            lambda fn=fused_qk: mx.concatenate(fn(), axis=0),
            both,
        )
        del qkv, qw, angles, cos, sin, qr, kr, both, native_qk, fused_qk
        mx.clear_cache()
        for k, out in ((5376, 21504), (5376, 28672), (14336, 5376)):
            a = mx.random.normal((n, k)).astype(mx.bfloat16)
            weight = (mx.random.normal((out, k)) * 0.02).astype(mx.bfloat16)
            mx.eval(a, weight)
            baseline = a @ weight.T
            mx.eval(baseline)
            record(
                "gemm",
                "bf16",
                [n, k, out],
                lambda a=a, weight=weight: a @ weight.T,
                baseline,
            )
            for bits in (8, 4):
                for size in (32, 64, 128):
                    q, s, z = mx.quantize(weight, bits=bits, group_size=size)
                    mx.eval(q, s, z)
                    record(
                        "gemm",
                        f"q{bits}_g{size}",
                        [n, k, out],
                        lambda a=a, q=q, s=s, z=z, size=size, bits=bits: (
                            mx.quantized_matmul(
                                a, q, s, z, transpose=True, group_size=size, bits=bits
                            )
                        ),
                        baseline,
                    )
            del a, weight, baseline, q, s, z
            gc.collect()
            mx.clear_cache()
        q = mx.random.normal((1, n, 56, 128)).astype(mx.bfloat16)
        k = mx.random.normal(q.shape).astype(mx.bfloat16)
        v = mx.random.normal(q.shape).astype(mx.bfloat16)
        mx.eval(q, k, v)
        baseline = attention(q, k, v)
        mx.eval(baseline)
        for chunk in (0, 128, 512, 1024):
            record(
                "sdpa",
                f"query_chunk_{chunk}",
                list(q.shape),
                lambda q=q, k=k, v=v, chunk=chunk: attention(
                    q, k, v, query_chunk=chunk
                ),
                baseline,
            )
        del q, k, v, baseline, x, w, shift, scale, ids, ref
        gc.collect()
        mx.clear_cache()
    return result


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--output", default="reports/operators.json")
    p.add_argument("--tokens", type=int, nargs="+", default=[256, 4096])
    p.add_argument("--repeats", type=int, default=7)
    a = p.parse_args()
    if a.repeats < 1 or any(n < 1 for n in a.tokens):
        p.error("tokens and repeats must be positive")
    run(a.output, tokens=a.tokens, repeats=a.repeats)


if __name__ == "__main__":
    main()