import pytest
import torch

from aiter.ops.tilelang import mhc_fused_tilelang


def generate_test_data(
    num_tokens: int,
    hidden_size: int,
    mhc_mult: int,
    device: str = "cuda",
) -> dict[str, torch.Tensor | float | int]:
    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=device)
    residual = torch.randn((num_tokens, mhc_mult, hidden_size), dtype=torch.bfloat16, device=device)
    post_layer_mix = torch.randn((num_tokens, mhc_mult, 1), dtype=torch.float32, device=device)
    comb_res_mix = torch.randn((num_tokens, mhc_mult, mhc_mult), dtype=torch.float32, device=device)

    mhc_mult2 = mhc_mult * mhc_mult
    mhc_mult3 = mhc_mult * 2 + mhc_mult2
    fn = (
        torch.randn((mhc_mult3, mhc_mult, hidden_size), dtype=torch.float32, device=device)
        * 1e-4
        * (1 + torch.arange(mhc_mult, device=device).mul(0.01).view(1, -1, 1))
    ).flatten(1, 2)
    mhc_scale = torch.randn((3,), dtype=torch.float32, device=device) * 0.1
    mhc_base = torch.randn((mhc_mult3,), dtype=torch.float32, device=device) * 0.1

    return {
        "x": x,
        "residual": residual,
        "post_layer_mix": post_layer_mix,
        "comb_res_mix": comb_res_mix,
        "fn": fn,
        "mhc_scale": mhc_scale,
        "mhc_base": mhc_base,
        "rms_eps": 1e-6,
        "mhc_pre_eps": 1e-6,
        "mhc_sinkhorn_eps": 1e-6,
        "mhc_post_mult_value": 1.0,
        "sinkhorn_repeat": 20,
    }


def _estimate_io_bytes(num_tokens: int, hidden_size: int, mhc_mult: int) -> int:
    # Rough fwd IO estimation for reporting consistency.
    mhc_mult2 = mhc_mult * mhc_mult
    mhc_mult3 = mhc_mult * (2 + mhc_mult)
    read_bytes = (
        num_tokens * hidden_size * 2
        + num_tokens * mhc_mult * hidden_size * 2
        + num_tokens * mhc_mult * 4
        + num_tokens * mhc_mult2 * 4
        + mhc_mult3 * mhc_mult * hidden_size * 4
        + 3 * 4
        + mhc_mult3 * 4
    )
    write_bytes = (
        num_tokens * mhc_mult * hidden_size * 2
        + num_tokens * mhc_mult * 4
        + num_tokens * mhc_mult2 * 4
        + num_tokens * hidden_size * 2
    )
    return read_bytes + write_bytes


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
@pytest.mark.benchmark
@pytest.mark.parametrize(
    "num_tokens,hidden_size,mhc_mult",
    [
        (1, 4096, 4),
        (4, 4096, 4),
        (8, 4096, 4),
        (32, 4096, 4),
        (32, 7168, 4),
        (64, 4096, 4),
        (64, 7168, 4),
        (128, 4096, 4),
        (128, 7168, 4),
        (1024, 4096, 4),
        (1024, 7168, 4),
    ],
)
def test_mhc_fused_tilelang_benchmark(
    num_tokens: int,
    hidden_size: int,
    mhc_mult: int,
    benchmark_timer,
    benchmark_record,
) -> None:
    td = generate_test_data(num_tokens=num_tokens, hidden_size=hidden_size, mhc_mult=mhc_mult)

    mhc_mult2 = mhc_mult * mhc_mult
    mhc_mult3 = mhc_mult * 2 + mhc_mult2

    # Keep split strategy aligned with correctness test and kernel constraints.
    if num_tokens <= 16:
        n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
        tile_n = 2 if num_tokens < 8 else 3
    else:
        n_splits = 2
        tile_n = 1

    gemm_out_mul = torch.empty((n_splits, num_tokens, mhc_mult3), dtype=torch.float32, device=td["x"].device)
    gemm_out_sqrsum = torch.empty((n_splits, num_tokens), dtype=torch.float32, device=td["x"].device)
    residual_out = torch.empty_like(td["residual"])

    def fn_fwd():
        return mhc_fused_tilelang(
            td["comb_res_mix"],
            td["residual"],
            td["post_layer_mix"].squeeze(-1),
            td["x"],
            td["fn"].view(mhc_mult3, mhc_mult, hidden_size),
            gemm_out_mul,
            gemm_out_sqrsum,
            residual_out,
            mhc_mult,
            hidden_size,
            mhc_mult3,
            tile_n=tile_n,
            split_k=n_splits,
        )

    fn_fwd()
    t_us = benchmark_timer(fn_fwd)
    io_bytes = _estimate_io_bytes(num_tokens, hidden_size, mhc_mult)
    bw_gbs = io_bytes / t_us / 1e3
    benchmark_record(
        kernel="mhc_fused_tilelang",
        operation="fwd",
        params={"num_tokens": num_tokens, "hidden_size": hidden_size, "mhc_mult": mhc_mult},
        time_us=t_us,
        bandwidth_gbs=bw_gbs,
        extras={"io_bytes": io_bytes},
    )
