"""Shape-specific Metal 4 TensorOps GEMM tile, K barrier and walk-order search."""

from __future__ import annotations

import argparse
import gc
import json
from functools import lru_cache
from pathlib import Path

import mlx.core as mx

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


@lru_cache(maxsize=None)
def kernel(bm, bn, bk, groups, walk):
    header = """
    #include <metal_tensor>
    #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
    using namespace mpp::tensor_ops;
    """
    source = r"""
    const int M = x_shape[0], K = x_shape[1], N = w_shape[0];
    int tiles_m = (M + BM - 1) / BM, tiles_n = N / BN;
    int tile = threadgroup_position_in_grid.x;
    int ti, tj;
    if (WALK == 0) {
        ti = tile / tiles_n; tj = tile % tiles_n;
    } else {
        // Group adjacent M rows, interleaving N tiles for L2 reuse.
        int group = tile / (WALK * tiles_n);
        int first_m = group * WALK;
        int size_m = min(tiles_m - first_m, WALK);
        int local = tile - group * WALK * tiles_n;
        ti = first_m + local % size_m;
        tj = local / size_m;
    }
    auto A = tensor(const_cast<device T*>(x), dextents<int, 2>(K, M), array<int, 2>{1, K});
    auto B = tensor(const_cast<device T*>(w), dextents<int, 2>(K, N), array<int, 2>{1, K});
    auto D = tensor(out, dextents<int, 2>(N, M), array<int, 2>{1, N});
    constexpr auto mode = matmul2d_descriptor::mode::multiply_accumulate;
    constexpr auto desc = matmul2d_descriptor(BM, BN, BK, false, true, false, mode);
    matmul2d<desc, execution_simdgroups<SG>> op;
    auto a0 = A.template slice<BK, dynamic_extent>(0, ti * BM);
    auto b0 = B.template slice<BK, BN>(0, tj * BN);
    auto accum = op.template get_destination_cooperative_tensor<decltype(a0), decltype(b0), float>();
    for (int e = 0; e < accum.get_capacity(); ++e) accum[e] = 0.0f;
    for (int k = 0; k < K; k += BK) {
        threadgroup_barrier(mem_flags::mem_none);
        auto a = A.template slice<BK, dynamic_extent>(k, ti * BM);
        auto b = B.template slice<BK, BN>(k, tj * BN);
        op.run(a, b, accum);
    }
    auto bias_tensor = tensor(const_cast<device T*>(bias), dextents<int, 2>(N, M), array<int, 2>{1, 0});
    auto bias_tile = bias_tensor.template slice<BN, dynamic_extent>(tj * BN, ti * BM);
    auto bias_regs = op.template get_destination_cooperative_tensor<decltype(a0), decltype(b0), T>();
    bias_regs.load(bias_tile);
    auto result = op.template get_destination_cooperative_tensor<decltype(a0), decltype(b0), T>();
    for (int e = 0; e < accum.get_capacity(); ++e) result[e] = T(accum[e] + float(bias_regs[e]));
    result.store(D.template slice<BN, dynamic_extent>(tj * BN, ti * BM));
    """
    return mx.fast.metal_kernel(
        name=f"h3_mpp_{bm}_{bn}_{bk}_{groups}_{walk}",
        input_names=["x", "w", "bias"],
        output_names=["out"],
        source=source,
        header=header,
    )


def gemm(x, w, bias, bm, bn, bk, groups, walk):
    m, k = x.shape
    n = w.shape[0]
    if k % bk or n % bn:
        raise ValueError("This experimental kernel requires aligned K and N")
    return kernel(bm, bn, bk, groups, walk)(
        inputs=[x, w, bias],
        template=[
            ("T", x.dtype),
            ("BM", bm),
            ("BN", bn),
            ("BK", bk),
            ("SG", groups),
            ("WALK", walk),
        ],
        grid=(((m + bm - 1) // bm) * (n // bn) * groups * 32, 1, 1),
        threadgroup=(groups * 32, 1, 1),
        output_shapes=[(m, n)],
        output_dtypes=[x.dtype],
    )[0]


def main(a):
    configure_memory()
    mx.random.seed(731)
    dest = Path(a.output)
    dest.parent.mkdir(parents=True, exist_ok=True)
    report = {
        "environment": environment(),
        "warmup": 2,
        "repeats": a.repeats,
        "scope": "Original projection weights, synthetic BF16 inputs. Custom Metal 4 TensorOps with FP32 accumulation and BF16 bias epilogue.",
        "trials": [],
    }
    store = WeightStore(Path(a.model) / "transformer")
    for name in a.operators:
        layer = Linear.load(
            store,
            "blocks.0." + name,
            grouped_heads=(56, 128) if name == "attn.qkv_proj" else None,
        )
        w = layer.weight
        b = layer.bias if layer.bias is not None else mx.zeros(w.shape[0], w.dtype)
        x = mx.random.normal((a.tokens, w.shape[1])).astype(w.dtype)
        ref = layer(x)
        mx.eval(x, w, b, ref)
        variants = [(0, 0, 0, 0, 0)]
        variants += [
            (bm, bn, bk, sg, walk)
            for bm, bn, sg in (
                (32, 64, 2),
                (64, 64, 4),
                (64, 128, 4),
                (128, 64, 4),
                (128, 128, 4),
                (128, 128, 8),
                (64, 128, 8),
            )
            for bk in (128, 256, 512)
            if w.shape[1] % bk == 0
            for walk in (0, 8)
        ]
        for bm, bn, bk, sg, walk in variants:
            row = {
                "operator": name,
                "tokens": a.tokens,
                "tile": [bm, bn, bk],
                "simdgroups": sg,
                "walk_m": walk,
            }
            try:
                fn = (
                    (lambda layer=layer, x=x: layer(x))
                    if bm == 0
                    else (
                        lambda x=x, w=w, b=b, bm=bm, bn=bn, bk=bk, sg=sg, walk=walk: (
                            gemm(x, w, b, bm, bn, bk, sg, walk)
                        )
                    )
                )
                row.update(measure(fn, repeats=a.repeats), **error(ref, fn()))
            except Exception as exc:
                row["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)
            if "error" in row and bm == 32:
                # A syntax or SDK incompatibility must be fixed before a sweep.
                return
        del layer, w, b, 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(
        "--operators",
        nargs="+",
        default=["attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"],
    )
    p.add_argument("--repeats", type=int, default=5)
    p.add_argument("--output", default="reports/optimization-v2/mpp.json")
    main(p.parse_args())