"""Autogenerated Triton kernel for tiled matrix multiplication."""

from __future__ import annotations

import torch
import warnings

try:
    import triton
    import triton.language as tl
except ImportError:  # pragma: no cover - optional dependency
    triton = None
    tl = None


BLOCK_M = 64
BLOCK_N = 64
BLOCK_K = 32
GROUP_M = 4


if triton is not None:

    @triton.jit
    def matmul_kernel(
        a_ptr,
        b_ptr,
        c_ptr,
        M,
        N,
        K,
        stride_am,
        stride_ak,
        stride_bk,
        stride_bn,
        stride_cm,
        stride_cn,
        *,
        BLOCK_M: tl.constexpr,
        BLOCK_N: tl.constexpr,
        BLOCK_K: tl.constexpr,
        GROUP_M: tl.constexpr,
    ):
        pid = tl.program_id(axis=0)
        grid_m = tl.cdiv(M, BLOCK_M)
        grid_n = tl.cdiv(N, BLOCK_N)
        group_id = pid // (GROUP_M * grid_n)
        group_count = tl.minimum(grid_m - group_id * GROUP_M, GROUP_M)
        pid_m = group_id * GROUP_M + pid % group_count
        pid_n = (pid // group_count) % grid_n

        offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
        offs_k = tl.arange(0, BLOCK_K)

        a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
        b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

        accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
        for k in range(0, K, BLOCK_K):
            a_block = tl.load(a_ptrs, mask=(offs_am[:, None] < M) & (offs_k[None, :] + k < K), other=0.0)
            b_block = tl.load(b_ptrs, mask=(offs_bn[None, :] < N) & (offs_k[:, None] + k < K), other=0.0)
            accumulator += tl.dot(a_block, b_block)
            a_ptrs += BLOCK_K * stride_ak
            b_ptrs += BLOCK_K * stride_bk

        c_ptrs = c_ptr + (offs_am[:, None] * stride_cm + offs_bn[None, :] * stride_cn)
        tl.store(c_ptrs, accumulator, mask=(offs_am[:, None] < M) & (offs_bn[None, :] < N))


def run_kernel(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """Multiply matrices using Triton when available else Torch fallback."""

    if a.dim() != 2 or b.dim() != 2:
        raise ValueError("Inputs must be 2D matrices")
    if a.size(1) != b.size(0):
        raise ValueError("Incompatible matrix dimensions")
    a = a.to(dtype=torch.float32)
    b = b.to(dtype=torch.float32)

    fallback = a @ b
    if triton is None:
        return fallback

    M, K = a.shape
    _, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
    try:
        matmul_kernel[
            grid
        ](
            a,
            b,
            c,
            M,
            N,
            K,
            a.stride(0),
            a.stride(1),
            b.stride(0),
            b.stride(1),
            c.stride(0),
            c.stride(1),
            BLOCK_M=BLOCK_M,
            BLOCK_N=BLOCK_N,
            BLOCK_K=BLOCK_K,
            GROUP_M=GROUP_M,
        )
    except Exception as exc:  # pragma: no cover - triton runtime issues
        warnings.warn(
            f"Triton matmul kernel failed; falling back to torch: {exc}",
            RuntimeWarning,
        )
        return fallback
    return c