# SPDX-License-Identifier: MIT

import torch
import triton

from aiter.ops.triton.fla.fused_sigmoid_gating_recurrent import (
    fused_sigmoid_gating_delta_rule_update,
)
from aiter.ops.triton.fla.fused_sigmoid_gating_recurrent_ref import (
    fused_sigmoid_gating_delta_rule_update as fused_sigmoid_gating_delta_rule_update_ref,
)


def estimate_bytes_target_verify(inp) -> int:
    B = inp["indices"].numel()
    T = inp["cache_steps"]
    H = inp["q"].shape[2]
    HV = inp["v"].shape[2]
    K = inp["q"].shape[3]
    V = inp["v"].shape[3]
    tokens = B * T

    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()

    bytes_qkv = tokens * (2 * H * K + HV * V) * qkv_elem
    bytes_ab = 2 * tokens * HV * ab_elem
    bytes_A_log = HV * alog_elem
    bytes_dt_bias = HV * dt_elem
    bytes_out = tokens * HV * V * out_elem

    state_slot_elems = HV * K * V
    state_read = B * state_slot_elems * state_elem
    state_write = 0
    inter_rw = tokens * state_slot_elems * state_elem
    return bytes_qkv + bytes_ab + bytes_A_log + bytes_dt_bias + bytes_out + state_read + state_write + inter_rw


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


def make_inputs(B=48, T=4, H=4, HV=16, K=128, V=128, dtype=torch.bfloat16):
    A_log = torch.randn(HV, dtype=torch.float32, device="cuda")
    dt_bias = torch.randn(HV, dtype=dtype, device="cuda")
    a = torch.randn(1, B * T, HV, dtype=dtype, device="cuda")
    b = torch.randn(1, B * T, HV, dtype=dtype, device="cuda")
    q = torch.randn(1, B * T, H, K, dtype=dtype, device="cuda")
    k = torch.randn(1, B * T, H, K, dtype=dtype, device="cuda")
    v = torch.randn(1, B * T, HV, V, dtype=dtype, device="cuda")
    state = torch.randn(B + 1, HV, K, V, dtype=torch.float32, device="cuda")
    indices = torch.arange(B, dtype=torch.int32, device="cuda")
    cu_seqlens = torch.arange(0, B * T + 1, T, dtype=torch.int32, device="cuda")
    inter_states = torch.empty(B, T, HV, K, V, dtype=torch.float32, device="cuda")
    inter_indices = torch.arange(B, dtype=torch.int32, device="cuda")
    return {
        "A_log": A_log,
        "dt_bias": dt_bias,
        "a": a,
        "b": b,
        "q": q,
        "k": k,
        "v": v,
        "state": state,
        "indices": indices,
        "cu_seqlens": cu_seqlens,
        "inter_states": inter_states,
        "inter_indices": inter_indices,
        "cache_steps": T,
    }


def run_target_verify_new(inp):
    return fused_sigmoid_gating_delta_rule_update(
        A_log=inp["A_log"],
        dt_bias=inp["dt_bias"],
        q=inp["q"],
        k=inp["k"],
        v=inp["v"],
        a=inp["a"],
        b=inp["b"],
        initial_state_source=inp["state"],
        initial_state_indices=inp["indices"],
        cu_seqlens=inp["cu_seqlens"],
        use_qk_l2norm_in_kernel=True,
        softplus_beta=1.0,
        softplus_threshold=20.0,
        disable_state_update=True,
        intermediate_states_buffer=inp["inter_states"],
        intermediate_state_indices=inp["inter_indices"],
        cache_steps=inp["cache_steps"],
    )


def run_target_verify_orig(inp):
    return fused_sigmoid_gating_delta_rule_update_ref(
        A_log=inp["A_log"],
        dt_bias=inp["dt_bias"],
        q=inp["q"],
        k=inp["k"],
        v=inp["v"],
        a=inp["a"],
        b=inp["b"],
        initial_state_source=inp["state"],
        initial_state_indices=inp["indices"],
        cu_seqlens=inp["cu_seqlens"],
        use_qk_l2norm_in_kernel=True,
        softplus_beta=1.0,
        softplus_threshold=20.0,
        disable_state_update=True,
        intermediate_states_buffer=inp["inter_states"],
        intermediate_state_indices=inp["inter_indices"],
        cache_steps=inp["cache_steps"],
    )


def bench_one(B, T, H=4, HV=16):
    inp = make_inputs(B=B, T=T, H=H, HV=HV)

    for _ in range(10):
        run_target_verify_new(inp)
        run_target_verify_orig(inp)
    torch.cuda.synchronize()

    ms_new = triton.testing.do_bench(
        lambda: run_target_verify_new(inp), warmup=50, rep=200
    )
    ms_ref = triton.testing.do_bench(
        lambda: run_target_verify_orig(inp), warmup=50, rep=200
    )
    bytes_tv = estimate_bytes_target_verify(inp)
    print(
        f"B={B:2d} T={T:d} H={H:d} HV={HV:d} | new={ms_new*1000:8.2f} us {gbps(bytes_tv, ms_new):7.2f} GB/s | "
        f"ref={ms_ref*1000:8.2f} us {gbps(bytes_tv, ms_ref):7.2f} GB/s | speedup={ms_ref/ms_new:.3f}x"
    )


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Benchmark fused gating delta rule update")
    parser.add_argument("--H", type=int, default=4, help="number of heads (default: 4)")
    parser.add_argument("--HV", type=int, default=16, help="number of V heads (default: 16)")
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required.")
    for B in [1, 4, 8, 16, 32, 48]:
        T = 4
        bench_one(B, T, H=args.H, HV=args.HV)
