# SPDX-License-Identifier: MIT

import argparse

import torch
import triton

from aiter.ops.triton.fla.fused_recurrent import (
    fused_recurrent_gated_delta_rule_packed_decode,
)
from op_tests.triton_tests.utils.fused_recurrent_ref import (
    fused_recurrent_gated_delta_rule_packed_decode_ref,
)


def estimate_packed_decode_flops(batch_size: int, hv: int, k_dim: int, v_dim: int) -> int:
    return 6 * batch_size * hv * v_dim * k_dim


def estimate_packed_decode_bytes(batch_size: int, h: int, hv: int, k_dim: int, v_dim: int) -> int:
    # Estimated logical traffic per call (not accounting for cache reuse):
    # read mixed_qkv + a/b + A_log/dt_bias + indices + state(read/write) + out(write)
    elem_fp16 = 2
    elem_fp32 = 4
    bytes_mixed_qkv = batch_size * (2 * h * k_dim + hv * v_dim) * elem_fp16
    bytes_ab = 2 * batch_size * hv * elem_fp16
    bytes_A_log = batch_size * hv * elem_fp32
    bytes_dt_bias = batch_size * hv * elem_fp16
    bytes_indices = batch_size * 4
    bytes_state_rw = 2 * batch_size * hv * v_dim * k_dim * elem_fp32
    bytes_out = batch_size * hv * v_dim * elem_fp16
    return (
        bytes_mixed_qkv
        + bytes_ab
        + bytes_A_log
        + bytes_dt_bias
        + bytes_indices
        + bytes_state_rw
        + bytes_out
    )


def tflops(flops: int, ms: float) -> float:
    return flops / (ms * 1e-3) / 1e12


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


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

    bsz = batch_size
    h_dim = args.heads
    hv = args.value_heads
    k_dim = args.k_dim
    v_dim = args.v_dim
    pool = max(args.pool_size, bsz)

    qkv_dim = 2 * h_dim * k_dim + hv * v_dim
    mixed_qkv = torch.randn((bsz, qkv_dim), device=device, dtype=torch.float16) * args.input_scale
    a = torch.randn((bsz, hv), device=device, dtype=torch.float16) * args.input_scale
    b = torch.randn((bsz, hv), device=device, dtype=torch.float16) * args.input_scale
    A_log = torch.randn((hv,), device=device, dtype=torch.float32) * args.input_scale
    dt_bias = torch.randn((hv,), device=device, dtype=torch.float16) * args.input_scale
    if args.allow_duplicate_indices:
        ssm_state_indices = torch.randint(0, pool, (bsz,), device=device, dtype=torch.int32)
    else:
        ssm_state_indices = torch.randperm(pool, device=device, dtype=torch.int64)[:bsz].to(torch.int32)
    if args.pad_slots > 0:
        ssm_state_indices[: min(args.pad_slots, bsz)] = -1

    out = torch.empty((bsz, 1, hv, v_dim), device=device, dtype=torch.float16)
    state = torch.randn((pool, hv, v_dim, k_dim), device=device, dtype=torch.float32) * 0.05

    return {
        "mixed_qkv": mixed_qkv,
        "a": a,
        "b": b,
        "A_log": A_log,
        "dt_bias": dt_bias,
        "ssm_state_indices": ssm_state_indices,
        "out": out,
        "state": state,
        "scale": k_dim ** -0.5,
    }


def bench_one(batch_size: int, args):
    inp = build_inputs(args, batch_size=batch_size)
    flops = estimate_packed_decode_flops(batch_size, args.value_heads, args.k_dim, args.v_dim)
    bytes_est = estimate_packed_decode_bytes(batch_size, args.heads, args.value_heads, args.k_dim, args.v_dim)

    out_cur = inp["out"].clone()
    out_ref = inp["out"].clone()
    state_cur = inp["state"].clone()
    state_ref = inp["state"].clone()

    fused_recurrent_gated_delta_rule_packed_decode(
        mixed_qkv=inp["mixed_qkv"],
        a=inp["a"],
        b=inp["b"],
        A_log=inp["A_log"],
        dt_bias=inp["dt_bias"],
        scale=inp["scale"],
        initial_state=state_cur,
        out=out_cur,
        ssm_state_indices=inp["ssm_state_indices"],
        use_qk_l2norm_in_kernel=args.use_l2norm,
    )
    fused_recurrent_gated_delta_rule_packed_decode_ref(
        mixed_qkv=inp["mixed_qkv"],
        a=inp["a"],
        b=inp["b"],
        A_log=inp["A_log"],
        dt_bias=inp["dt_bias"],
        scale=inp["scale"],
        initial_state=state_ref,
        out=out_ref,
        ssm_state_indices=inp["ssm_state_indices"],
        use_qk_l2norm_in_kernel=args.use_l2norm,
    )
    torch.cuda.synchronize()

    fn_cur = lambda: fused_recurrent_gated_delta_rule_packed_decode(  # noqa: E731
        mixed_qkv=inp["mixed_qkv"],
        a=inp["a"],
        b=inp["b"],
        A_log=inp["A_log"],
        dt_bias=inp["dt_bias"],
        scale=inp["scale"],
        initial_state=state_cur,
        out=out_cur,
        ssm_state_indices=inp["ssm_state_indices"],
        use_qk_l2norm_in_kernel=args.use_l2norm,
    )
    fn_ref = lambda: fused_recurrent_gated_delta_rule_packed_decode_ref(  # noqa: E731
        mixed_qkv=inp["mixed_qkv"],
        a=inp["a"],
        b=inp["b"],
        A_log=inp["A_log"],
        dt_bias=inp["dt_bias"],
        scale=inp["scale"],
        initial_state=state_ref,
        out=out_ref,
        ssm_state_indices=inp["ssm_state_indices"],
        use_qk_l2norm_in_kernel=args.use_l2norm,
    )

    ms_cur = triton.testing.do_bench(fn_cur, warmup=args.warmup, rep=args.rep)
    ms_ref = triton.testing.do_bench(fn_ref, warmup=args.warmup, rep=args.rep)

    print(
        f"B={batch_size:3d} | cur={ms_cur*1000:8.2f} us {tflops(flops, ms_cur):7.3f} TF {gbps(bytes_est, ms_cur):7.2f} GB/s | "
        f"ref={ms_ref*1000:8.2f} us {tflops(flops, ms_ref):7.3f} TF {gbps(bytes_est, ms_ref):7.2f} GB/s | speedup={ms_ref / ms_cur:.3f}x"
    )


def main():
    parser = argparse.ArgumentParser("Benchmark fused_recurrent packed decode: current vs ref")
    parser.add_argument("--batch", type=int, default=64)
    parser.add_argument("--batch-list", type=str, default="")
    parser.add_argument("--heads", type=int, default=4)
    parser.add_argument("--value-heads", type=int, default=16)
    parser.add_argument("--k-dim", type=int, default=128)
    parser.add_argument("--v-dim", type=int, default=128)
    parser.add_argument("--pool-size", type=int, default=512)
    parser.add_argument("--pad-slots", type=int, default=0)
    parser.add_argument("--allow-duplicate-indices", action="store_true", default=False)
    parser.add_argument("--use-l2norm", action="store_true", default=False)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--input-scale", type=float, default=0.2)
    parser.add_argument("--warmup", type=int, default=25)
    parser.add_argument("--rep", type=int, default=100)
    args = parser.parse_args()

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

    batch_list = [int(x) for x in args.batch_list.split(",") if x] if args.batch_list else [args.batch]

    print("=" * 112)
    print(
        "packed decode benchmark (current vs ref) | "
        f"H={args.heads} HV={args.value_heads} K={args.k_dim} V={args.v_dim} "
        f"pool={args.pool_size} pad_slots={args.pad_slots} dup_idx={args.allow_duplicate_indices} l2norm={args.use_l2norm}"
    )
    print("-" * 112)
    for bsz in batch_list:
        bench_one(bsz, args)
    print("=" * 112)


if __name__ == "__main__":
    main()
