import torch
from task import input_t, output_t

# Import the module that contains the CK-based blockscale GEMM
from aiter.ops import gemm_op_a8w8 as aiter_ck_gemm_ops


# Helper to ensure tensors are Row-Major (C-contiguous)
def ensure_row_major(tensor: torch.Tensor) -> torch.Tensor:
    if not tensor.is_contiguous(memory_format=torch.contiguous_format):
        return tensor.clone(memory_format=torch.contiguous_format)
    return tensor


def custom_kernel(data: input_t) -> output_t:
    A_cm: torch.Tensor = data[0]
    B_cm: torch.Tensor = data[1]
    alpha_cm: torch.Tensor = data[2]
    beta_cm: torch.Tensor = data[3]
    C_cm_orig: torch.Tensor = data[4]

    A_api_rm = ensure_row_major(A_cm)
    B_api_rm = ensure_row_major(B_cm)
    alpha_api_rm = ensure_row_major(alpha_cm)
    beta_api_rm = ensure_row_major(beta_cm)

    # C_cm_orig is already BF16 and column-major as per problem spec.
    # We are passing it directly.
    aiter_ck_gemm_ops.gemm_a8w8_blockscale(
        A_api_rm,
        B_api_rm,
        alpha_api_rm,
        beta_api_rm,
        C_cm_orig,  # Pass C_cm_orig (column-major) directly
    )

    return C_cm_orig


# k: 7168; m: 1024; n: 1536; seed: 8135
#  ⏱ 174 ± 0.2 µs
#  ⚡ 172 µs 🐌 179 µs

# k: 1536; m: 1024; n: 3072; seed: 6251
#  ⏱ 86.8 ± 0.09 µs
#  ⚡ 85.3 µs 🐌 91.4 µs

# k: 7168; m: 1024; n: 576; seed: 12346
#  ⏱ 109 ± 0.1 µs
#  ⚡ 107 µs 🐌 113 µs

# k: 256; m: 1024; n: 7168; seed: 5364
#  ⏱ 67.5 ± 0.17 µs
#  ⚡ 64.5 µs 🐌 77.0 µs

# k: 2048; m: 1024; n: 7168; seed: 6132
#  ⏱ 169 ± 0.2 µs
#  ⚡ 165 µs 🐌 175 µs

# k: 7168; m: 1024; n: 4608; seed: 7531
#  ⏱ 368 ± 0.7 µs
#  ⚡ 356 µs 🐌 386 µs

# k: 2304; m: 1024; n: 7168; seed: 12345
#  ⏱ 196 ± 0.2 µs
#  ⚡ 192 µs 🐌 202 µs

# k: 7168; m: 1024; n: 512; seed: 6563
#  ⏱ 109 ± 0.7 µs
#  ⚡ 105 µs 🐌 143 µs

# k: 512; m: 1024; n: 4096; seed: 17512
#  ⏱ 64.5 ± 0.15 µs
#  ⚡ 62.3 µs 🐌 72.3 µs

# k: 7168; m: 6144; n: 1536; seed: 6543
#  ⏱ 701 ± 2.4 µs
#  ⚡ 657 µs 🐌 742 µs

# k: 1536; m: 6144; n: 3072; seed: 234
#  ⏱ 230 ± 0.2 µs
#  ⚡ 228 µs 🐌 231 µs

# k: 7168; m: 6144; n: 576; seed: 9863
#  ⏱ 503 ± 2.6 µs
#  ⚡ 464 µs 🐌 548 µs

# k: 256; m: 6144; n: 7168; seed: 764243
#  ⏱ 120 ± 0.2 µs
#  ⚡ 118 µs 🐌 128 µs

# k: 2048; m: 6144; n: 7168; seed: 76547
#  ⏱ 502 ± 0.9 µs
#  ⚡ 482 µs 🐌 522 µs

# k: 7168; m: 6144; n: 4608; seed: 65436
#  ⏱ 1297 ± 2.6 µs
#  ⚡ 1238 µs 🐌 1361 µs

# k: 2304; m: 6144; n: 7168; seed: 452345
#  ⏱ 596 ± 1.0 µs
#  ⚡ 581 µs 🐌 623 µs

# k: 7168; m: 6144; n: 512; seed: 12341
#  ⏱ 496 ± 2.6 µs
#  ⚡ 450 µs 🐌 550 µs

# k: 512; m: 6144; n: 4096; seed: 45245
#  ⏱ 125 ± 0.1 µs
#  ⚡ 122 µs 🐌 131 µs

# Successful on MI300: 225 µs

python eval.py test /tmp/tmpwjc1cngd
[aiter] start build [module_gemm_a8w8_blockscale] under /home/runner/aiter/aiter/jit/build/module_gemm_a8w8_blockscale
[92mSuccessfully preprocessed all matching files.[0m
[aiter] finish build [module_gemm_a8w8_blockscale], cost 70.99105085s
[aiter] type hints mismatch, override to --> gemm_a8w8_blockscale(XQ: torch.Tensor, WQ: torch.Tensor, x_scale: torch.Tensor, w_scale: torch.Tensor, Out: torch.Tensor) -> torch.Tensor
