# SPDX-License-Identifier: MIT

import argparse

import torch
import triton

from aiter.ops.triton.fla.fused_sigmoid_gating import (
    fused_sigmoid_gating_delta_rule_update,
)
from op_tests.triton_tests.utils.fused_sigmoid_gating_ref import (
    fused_sigmoid_gating_delta_rule_ref,
)


def estimate_bytes_from_inputs(inp: dict, args) -> int:
    # Estimate only bytes touched by this invocation using real tensor dtypes.
    num_tokens = args.num_reqs * args.seq_len
    num_v_heads = args.num_v_heads
    head_k_dim = args.head_k_dim
    head_v_dim = args.head_v_dim
    num_k_heads = args.num_k_heads

    qkv_elem = inp["q"].element_size()
    ab_elem = inp["a"].element_size()
    alog_elem = inp["A_log"].element_size()
    dt_elem = inp["dt_bias"].element_size()
    state_elem = inp["state"].element_size()
    out_elem = inp["v"].element_size()

    # q/k/v read once per token.
    qkv_elems = num_tokens * (2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim)
    bytes_qkv = qkv_elems * qkv_elem

    # a/b are loaded once per processed token/head (not full ab table).
    bytes_ab = 2 * num_tokens * num_v_heads * ab_elem
    bytes_A_log = num_v_heads * alog_elem
    bytes_dt_bias = num_v_heads * dt_elem

    # State read once per sequence, write once per processed token.
    state_elems_per_slot = num_v_heads * head_v_dim * head_k_dim
    num_seqs = int(inp["cu_seqlens"].numel() - 1) if inp["cu_seqlens"] is not None else args.num_reqs
    state_write_slots = int(inp["state_indices"].numel()) if inp["state_indices"] is not None else num_tokens
    bytes_state_rw = (num_seqs + state_write_slots) * state_elems_per_slot * state_elem

    bytes_out = num_tokens * num_v_heads * head_v_dim * out_elem
    return bytes_qkv + bytes_ab + bytes_A_log + bytes_dt_bias + bytes_state_rw + bytes_out


def gbps(num_bytes: int, ms: float) -> float:
    return num_bytes / (ms * 1e-3) / 1e9


def bytes_to_mb(num_bytes: int) -> float:
    return num_bytes / 1e6


def build_inputs(args):
    torch.manual_seed(args.seed)
    device = torch.device("cuda")

    num_tokens = args.num_reqs * args.seq_len
    total_entries = max(args.state_pool, num_tokens * 2)

    mixed_qkv_dim = args.num_k_heads * args.head_k_dim * 2 + args.num_v_heads * args.head_v_dim
    mixed_qkv = torch.randn(num_tokens, mixed_qkv_dim, device=device, dtype=args.dtype) * args.input_scale
    q, k, v = torch.split(
        mixed_qkv,
        [args.num_k_heads * args.head_k_dim, args.num_k_heads * args.head_k_dim, args.num_v_heads * args.head_v_dim],
        dim=-1,
    )
    q = q.view(1, num_tokens, args.num_k_heads, args.head_k_dim)
    k = k.view(1, num_tokens, args.num_k_heads, args.head_k_dim)
    v = v.view(1, num_tokens, args.num_v_heads, args.head_v_dim)

    A_log = torch.randn(args.num_v_heads, device=device, dtype=args.dtype) * args.input_scale
    dt_bias = torch.randn(args.num_v_heads, device=device, dtype=args.dtype) * args.input_scale
    ab_tokens = max(args.ab_tokens, num_tokens)
    a = torch.randn(ab_tokens, args.num_v_heads, device=device, dtype=args.dtype) * args.input_scale
    b = torch.randn(ab_tokens, args.num_v_heads, device=device, dtype=args.dtype) * args.input_scale

    state = torch.randn(total_entries, args.num_v_heads, args.head_v_dim, args.head_k_dim, device=device, dtype=args.dtype)
    if args.spec_tokens > 0:
        seq_width = args.spec_tokens + 1
        assert num_tokens % seq_width == 0
        num_reqs = num_tokens // seq_width
        state_indices = torch.randperm(total_entries, device=device, dtype=torch.int64)[:num_tokens]
        state_indices = state_indices.to(torch.int32).view(num_reqs, seq_width)
        cu_base = args.cu_seqlens_base
        if cu_base < 0:
            cu_base = 0
        cu_seqlens = torch.arange(
            cu_base,
            cu_base + num_tokens + 1,
            seq_width,
            device=device,
            dtype=torch.int32,
        )
        num_accepted_tokens = torch.randint(1, seq_width, (num_reqs,), device=device, dtype=torch.int32)
    else:
        state_indices = torch.randperm(total_entries, device=device, dtype=torch.int64)[:num_tokens].to(torch.int32)
        cu_seqlens = torch.arange(0, num_tokens + 1, args.seq_len, device=device, dtype=torch.int32)
        num_accepted_tokens = None

    return {
        "A_log": A_log,
        "a": a,
        "b": b,
        "dt_bias": dt_bias,
        "q": q,
        "k": k,
        "v": v,
        "state": state,
        "state_indices": state_indices,
        "cu_seqlens": cu_seqlens,
        "num_accepted_tokens": num_accepted_tokens,
    }


def main():
    p = argparse.ArgumentParser("benchmark fused_sigmoid_gating_delta_rule_update")
    p.add_argument("--num-reqs", type=int, default=2)
    p.add_argument("--seq-len", type=int, default=4)
    p.add_argument("--num-k-heads", type=int, default=4)
    p.add_argument("--num-v-heads", type=int, default=16)
    p.add_argument("--head-k-dim", type=int, default=128)
    p.add_argument("--head-v-dim", type=int, default=128)
    p.add_argument("--state-pool", type=int, default=1639)
    p.add_argument("--spec-tokens", type=int, default=3, help="0 means non-spec")
    p.add_argument("--ab-tokens", type=int, default=13628, help="Rows for a/b tables.")
    p.add_argument(
        "--cu-seqlens-base",
        type=int,
        default=-1,
        help="Base offset for cu_seqlens in spec mode; -1 uses 0 (relative window).",
    )
    p.add_argument("--dtype", type=str, default="bf16", choices=["fp16", "bf16", "fp32"])
    p.add_argument("--input-scale", type=float, default=0.2)
    p.add_argument("--warmup", type=int, default=25)
    p.add_argument("--rep", type=int, default=100)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--bv", type=int, default=0, help="Kernel override. 0 means use config/default.")
    args = p.parse_args()

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required")

    if args.dtype == "fp16":
        args.dtype = torch.float16
    elif args.dtype == "bf16":
        args.dtype = torch.bfloat16
    else:
        args.dtype = torch.float32

    inp = build_inputs(args)
    kernel_cfg = None
    if args.bv > 0:
        kernel_cfg = {}
        if args.bv > 0:
            kernel_cfg["BV"] = args.bv
    out_ref, _ = fused_sigmoid_gating_delta_rule_ref(
        A_log=inp["A_log"],
        a=inp["a"],
        b=inp["b"],
        dt_bias=inp["dt_bias"],
        q=inp["q"],
        k=inp["k"],
        v=inp["v"],
        initial_state=inp["state"].clone(),
        inplace_final_state=True,
        cu_seqlens=inp["cu_seqlens"],
        ssm_state_indices=inp["state_indices"],
        num_accepted_tokens=inp["num_accepted_tokens"],
        use_qk_l2norm_in_kernel=True,
    )
    out_cur, _ = fused_sigmoid_gating_delta_rule_update(
        A_log=inp["A_log"],
        a=inp["a"],
        b=inp["b"],
        dt_bias=inp["dt_bias"],
        q=inp["q"],
        k=inp["k"],
        v=inp["v"],
        initial_state=inp["state"].clone(),
        inplace_final_state=True,
        cu_seqlens=inp["cu_seqlens"],
        ssm_state_indices=inp["state_indices"],
        num_accepted_tokens=inp["num_accepted_tokens"],
        use_qk_l2norm_in_kernel=True,
        kernel_cfg=kernel_cfg,
    )
    torch.testing.assert_close(out_cur.float(), out_ref.float(), atol=2e-2, rtol=2e-2)

    fn = lambda: fused_sigmoid_gating_delta_rule_update(  # noqa: E731
        A_log=inp["A_log"],
        a=inp["a"],
        b=inp["b"],
        dt_bias=inp["dt_bias"],
        q=inp["q"],
        k=inp["k"],
        v=inp["v"],
        initial_state=inp["state"],
        inplace_final_state=True,
        cu_seqlens=inp["cu_seqlens"],
        ssm_state_indices=inp["state_indices"],
        num_accepted_tokens=inp["num_accepted_tokens"],
        use_qk_l2norm_in_kernel=True,
    )
    fn_ref = lambda: fused_sigmoid_gating_delta_rule_ref(  # noqa: E731
        A_log=inp["A_log"],
        a=inp["a"],
        b=inp["b"],
        dt_bias=inp["dt_bias"],
        q=inp["q"],
        k=inp["k"],
        v=inp["v"],
        initial_state=inp["state"],
        inplace_final_state=True,
        cu_seqlens=inp["cu_seqlens"],
        ssm_state_indices=inp["state_indices"],
        num_accepted_tokens=inp["num_accepted_tokens"],
        use_qk_l2norm_in_kernel=True,
    )

    ms = triton.testing.do_bench(fn, warmup=args.warmup, rep=args.rep)
    ms_ref = triton.testing.do_bench(fn_ref, warmup=args.warmup, rep=args.rep)
    num_tokens = args.num_reqs * args.seq_len
    bytes_est = estimate_bytes_from_inputs(inp, args)

    print("=" * 96)
    print(
        f"fused_sigmoid_gating_delta_rule_update | reqs={args.num_reqs} seq={args.seq_len} spec={args.spec_tokens} "
        f"H={args.num_k_heads} HV={args.num_v_heads} K={args.head_k_dim} V={args.head_v_dim} dtype={args.dtype} "
        f"kernel_cfg={kernel_cfg}"
    )
    print(f"bytes_est: {bytes_to_mb(bytes_est):.3f} MB")
    print(
        f"new={ms*1000:8.2f} us {gbps(bytes_est, ms):7.2f} GB/s | "
        f"ref={ms_ref*1000:8.2f} us {gbps(bytes_est, ms_ref):7.2f} GB/s | "
        f"speedup={ms_ref / ms:.3f}x"
    )
    print("=" * 96)


if __name__ == "__main__":
    main()
