"""Fused gather-GEMM-accumulate Metal kernel for submanifold sparse conv.

The stock MlxSparseConv3d materializes a (K, N, Ci) gather — ~77GB at
decoder scale — then chunks it 1-2 offsets at a time with an mx.eval sync
between chunks. This kernel never materializes the gather: each thread
accumulates out[n, co] = Σ_k Σ_ci feats[nbr[k,n], ci] · w[k, ci, co]
straight from the source buffers. Correctness-first implementation
(scalar loops, fp32 accumulation); simdgroup-matrix tiling is a later
optimization if profiling ever says the decoders matter for *time*
(they are ~10s/gen — this kernel's prize is PEAK MEMORY, which gates the
M1 Ultra).

Env gate: TRELLIS2_METAL_SPCONV=1 routes MlxSparseConv3d through this.
"""
import mlx.core as mx

_SRC = """
    uint n = thread_position_in_grid.x;      // voxel index
    uint co = thread_position_in_grid.y;     // output channel
    uint N = (uint)shape_info[0];
    uint K = (uint)shape_info[1];
    uint Ci = (uint)shape_info[2];
    uint Co = (uint)shape_info[3];
    if (n >= N || co >= Co) return;

    float acc = 0.0f;
    for (uint k = 0; k < K; ++k) {
        int nbr = nmap[k * N + n];           // N (==pad row) when invalid
        if (nbr >= (int)N) continue;          // pad row is zeros anyway; skip
        const device T* frow = feats + (size_t)nbr * Ci;
        const device T* wrow = w + ((size_t)k * Ci) * Co + co;
        for (uint ci = 0; ci < Ci; ++ci) {
            acc += (float)frow[ci] * (float)wrow[(size_t)ci * Co];
        }
    }
    out[(size_t)n * Co + co] = (T)acc;
"""

_kernel = None


def _get_kernel():
    global _kernel
    if _kernel is None:
        _kernel = mx.fast.metal_kernel(
            name="submconv3d_gather_gemm",
            input_names=["feats", "nmap", "w", "shape_info"],
            output_names=["out"],
            source=_SRC,
        )
    return _kernel


def sparse_conv_metal(feats_padded: mx.array, neighbor_map: mx.array,
                      w_kcico: mx.array, bias: mx.array) -> mx.array:
    """feats_padded: (N+1, Ci) — row N is zeros (pad; skipped anyway).
    neighbor_map: (K, N) int32 with N as the invalid sentinel.
    w_kcico: (K, Ci, Co). bias: (Co,). Returns (N, Co) in w's dtype."""
    K, N = neighbor_map.shape
    Ci = feats_padded.shape[1]
    Co = w_kcico.shape[2]
    shape_info = mx.array([N, K, Ci, Co], dtype=mx.int32)
    kern = _get_kernel()
    (out,) = kern(
        inputs=[feats_padded.astype(w_kcico.dtype), neighbor_map,
                w_kcico, shape_info],
        template=[("T", w_kcico.dtype)],
        grid=(N, Co, 1),
        threadgroup=(256, 1, 1),
        output_shapes=[(N, Co)],
        output_dtypes=[w_kcico.dtype],
    )
    return out + bias