"""Experimental Metal 4 online-softmax attention with configurable Q/K tiles.

Dense, noncausal, batch-one, equal-head attention only. No approximated sparsity.
Not selected by the runtime unless full-pipeline validation supports a speed win.
"""

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.ops import attention
from h3_mlx.runtime import configure_memory, environment


@lru_cache(maxsize=None)
def build(bq, bk, groups):
    return mx.fast.metal_kernel(
        name=f"h3_flash_{bq}_{bk}_{groups}",
        input_names=["q", "k", "v"],
        output_names=["out"],
        ensure_row_contiguous=False,
        header="""
        #include <metal_tensor>
        #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
        using namespace mpp::tensor_ops;
        """,
        source=r"""
        const int N = q_shape[0], H = q_shape[1], D = 128;
        int sg = thread_position_in_threadgroup.x / 32;
        int row = (threadgroup_position_in_grid.x * SG + sg) * BQ;
        int head = threadgroup_position_in_grid.y;
        // SIMD groups own independent query tiles; the last group may be empty.
        auto Q = tensor(const_cast<device T*>(q) + head*q_strides[1], dextents<int,2>(D,N), array<int,2>{1,int(q_strides[0])});
        auto K = tensor(const_cast<device T*>(k) + head*k_strides[1], dextents<int,2>(D,N), array<int,2>{1,int(k_strides[0])});
        auto V = tensor(const_cast<device T*>(v) + head*v_strides[1], dextents<int,2>(D,N), array<int,2>{1,int(v_strides[0])});
        auto O = tensor(out + head*D, dextents<int,2>(D,N), array<int,2>{1,H*D});
        if (row >= N) return;
        constexpr auto qkdesc = matmul2d_descriptor(BQ,BK,128,false,true,false);
        constexpr auto pvdesc = matmul2d_descriptor(BQ,128,BK,false,false,false,matmul2d_descriptor::mode::multiply_accumulate);
        matmul2d<qkdesc,execution_simdgroup> qkop;
        matmul2d<pvdesc,execution_simdgroup> pvop;
        auto qtile = Q.template slice<128,dynamic_extent>(0,row);
        auto ktile = K.template slice<128,dynamic_extent>(0,0);
        auto vt = V.template slice<128,dynamic_extent>(0,0);
        auto vtile = V.template slice<128,dynamic_extent>(0,0);
        auto qs = qkop.template get_left_input_cooperative_tensor<T,T,float>();
        qs.load(qtile);
        auto scores = qkop.template get_destination_cooperative_tensor<decltype(qtile),decltype(ktile),float>();
        auto maximum = qkop.template get_row_reduction_destination_cooperative_tensor<decltype(qtile),decltype(ktile),float>();
        auto denom = maximum;
        auto tile_max = maximum;
        auto tile_sum = maximum;
        auto alpha = maximum;
        auto prob = pvop.template get_left_input_cooperative_tensor<T,T,float>();
        auto p_shape = tensor(const_cast<device T*>(q), extents<int,BK,BQ>(), array<int,2>{1,BK});
        auto accum = pvop.template get_destination_cooperative_tensor<decltype(p_shape),decltype(vtile),float>();
        #pragma unroll
        for (int i=0;i<maximum.get_capacity();++i) { maximum[i]=-INFINITY; denom[i]=0.0f; }
        #pragma unroll
        for (int i=0;i<accum.get_capacity();++i) accum[i]=0.0f;
        threadgroup T scratch[SG*BQ*BK];
        auto P = tensor(scratch+sg*BQ*BK, extents<int,BK,BQ>(), array<int,2>{1,BK});
        for (int start=0;start<N;start+=BK) {
            auto kk = K.template slice<128,dynamic_extent>(0,start);
            qkop.run(qs,kk,scores);
            #pragma unroll
            for (auto it=scores.begin();it!=scores.end();++it) {
                auto coord = scores.get_multidimensional_index(it);
                *it = start+coord[0]<N ? *it*0.08838834764831845f : -INFINITY;
            }
            reduce_rows(scores,tile_max,reduction_operation::max,-INFINITY);
            #pragma unroll
            for (int i=0;i<maximum.get_capacity();++i) {
                float next=max(maximum[i],tile_max[i]);
                alpha[i]=exp(maximum[i]-next); maximum[i]=next;
            }
            #pragma unroll
            for (auto it=scores.begin();it!=scores.end();++it) {
                *it=exp(*it-*maximum.map_iterator(it));
            }
            reduce_rows(scores,tile_sum,reduction_operation::sum,0.0f);
            #pragma unroll
            for (int i=0;i<denom.get_capacity();++i) denom[i]=denom[i]*alpha[i]+tile_sum[i];
            #pragma unroll
            for (auto it=accum.begin();it!=accum.end();++it) *it *= *alpha.map_iterator(it);
            // FP32 and BF16 cooperative destinations may have different lane
            // layouts. Convert by logical coordinates, never by array index.
            #pragma unroll
            for (auto it=scores.begin();it!=scores.end();++it) {
                if (scores.is_valid_element(it)) {
                    auto coord=scores.get_multidimensional_index(it);
                    scratch[sg*BQ*BK+coord[1]*BK+coord[0]]=T(*it);
                }
            }
            simdgroup_barrier(mem_flags::mem_threadgroup);
            prob.load(P);
            auto vv = V.template slice<128,dynamic_extent>(0,start);
            pvop.run(prob,vv,accum);
        }
        #pragma unroll
        for (auto it=accum.begin();it!=accum.end();++it) *it /= *denom.map_iterator(it);
        #pragma unroll
        for (auto it=accum.begin();it!=accum.end();++it) {
            if (accum.is_valid_element(it)) {
                auto coord=accum.get_multidimensional_index(it);
                if (row+coord[1]<N) out[((row+coord[1])*H+head)*D+coord[0]]=T(*it);
            }
        }
        """,
    )


def flash(q, k, v, bq=16, bk=64, groups=4):
    if (
        q.ndim != 4
        or q.shape[0] != 1
        or q.shape[-1] != 128
        or q.shape != k.shape
        or q.shape != v.shape
    ):
        raise ValueError("Experimental attention requires equal [1,N,H,128] tensors")
    n, h, d = q.shape[1:]
    return build(bq, bk, groups)(
        inputs=[q[0], k[0], v[0]],
        template=[("T", q.dtype), ("BQ", bq), ("BK", bk), ("SG", groups)],
        grid=(((n + bq * groups - 1) // (bq * groups)) * groups * 32, h, 1),
        threadgroup=(groups * 32, 1, 1),
        output_shapes=[(1, n, h, d)],
        output_dtypes=[q.dtype],
    )[0]


def main(a):
    configure_memory()
    mx.random.seed(731)
    out = Path(a.output)
    out.parent.mkdir(parents=True, exist_ok=True)
    report = {
        "environment": environment(),
        "repeats": a.repeats,
        "warmup": 2,
        "scope": "Synthetic BF16 Q/K/V; dense online softmax. Different reduction/rounding may affect generated images.",
        "trials": [],
    }
    for n in a.tokens:
        q, k, v = (
            mx.random.normal((1, n, 56, 128)).astype(mx.bfloat16) for _ in range(3)
        )
        mx.eval(q, k, v)
        ref = attention(q, k, v)
        mx.eval(ref)
        variants = [(0, 0, 0)] + [
            (bq, bk, sg) for bq in (16, 32) for bk in (32, 64, 128) for sg in (2, 4)
        ]
        for bq, bk, sg in variants:
            row = {"tokens": n, "bq": bq, "bk": bk, "simdgroups": sg}
            try:
                fn = (
                    (lambda q=q, k=k, v=v: attention(q, k, v))
                    if bq == 0
                    else (
                        lambda q=q, k=k, v=v, bq=bq, bk=bk, sg=sg: flash(
                            q, k, v, bq, bk, sg
                        )
                    )
                )
                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)
            out.write_text(json.dumps(report, indent=2) + "\n")
            print(json.dumps(row), flush=True)
            if "error" in row:
                return
        del q, k, v, ref
        gc.collect()
        mx.clear_cache()


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--tokens", nargs="+", type=int, default=[4181, 16469])
    p.add_argument("--repeats", type=int, default=5)
    p.add_argument("--output", default="reports/optimization-v2/attention.json")
    main(p.parse_args())