# SPDX-License-Identifier: MIT

import argparse

import torch
import triton

from aiter.ops.triton.fla.vllm.chunk_delta_h import (
    chunk_gated_delta_rule_fwd_h,
    prepare_chunk_indices,
)
from aiter.ops.triton.fla.vllm.chunk_o import chunk_fwd_o
from op_tests.triton_tests.utils.chunk_delta_h_vllm_ref import (
    chunk_gated_delta_rule_fwd_h as chunk_gated_delta_rule_fwd_h_ref,
)
from op_tests.triton_tests.utils.chunk_o_vllm_ref import (
    chunk_fwd_o as chunk_fwd_o_ref,
)


def estimate_fwd_h_flops(batch_size, seqlen, num_heads, k_dim, v_dim, chunk_size):
    num_chunks = triton.cdiv(seqlen, chunk_size)
    total = 0
    for chunk_idx in range(num_chunks):
        chunk_len = min(chunk_size, seqlen - chunk_idx * chunk_size)
        total += 4 * batch_size * num_heads * chunk_len * k_dim * v_dim
    return total


def estimate_fwd_o_flops(batch_size, seqlen, num_heads, k_dim, v_dim, chunk_size):
    num_chunks = triton.cdiv(seqlen, chunk_size)
    total = 0
    for chunk_idx in range(num_chunks):
        chunk_len = min(chunk_size, seqlen - chunk_idx * chunk_size)
        total += 2 * batch_size * num_heads * chunk_len * k_dim * v_dim
        total += 2 * batch_size * num_heads * chunk_len * chunk_len * k_dim
        total += 2 * batch_size * num_heads * chunk_len * chunk_len * v_dim
    return total


def estimate_varlen_flops(cu_seqlens, num_heads, k_dim, v_dim, chunk_size, kind="h"):
    total = 0
    for i in range(len(cu_seqlens) - 1):
        seqlen = int((cu_seqlens[i + 1] - cu_seqlens[i]).item())
        if kind == "h":
            total += estimate_fwd_h_flops(1, seqlen, num_heads, k_dim, v_dim, chunk_size)
        else:
            total += estimate_fwd_o_flops(1, seqlen, num_heads, k_dim, v_dim, chunk_size)
    return total


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


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

    b = args.batch
    t = args.seqlen
    h = args.heads
    hg = args.grouped_heads
    k_dim = args.k_dim
    v_dim = args.v_dim

    q = torch.randn((b, t, hg, k_dim), device=device, dtype=torch.float16) * args.input_scale
    k = torch.randn((b, t, hg, k_dim), device=device, dtype=torch.float16) * args.input_scale
    w = torch.randn((b, t, h, k_dim), device=device, dtype=torch.float16) * args.input_scale
    u = torch.randn((b, t, h, v_dim), device=device, dtype=torch.float16) * args.input_scale
    g = torch.randn((b, t, h), device=device, dtype=torch.float32) * args.g_scale

    cu_seqlens = None
    chunk_indices = None
    n_seq = b

    if args.varlen:
        if b != 1:
            raise ValueError("Current varlen benchmark assumes batch=1 flattened token layout.")
        if args.varlen_splits:
            points = [int(x) for x in args.varlen_splits.split(",") if x.strip()]
            points = [x for x in points if 0 < x < t]
            pts = [0] + sorted(points)
            if pts[-1] != t:
                pts.append(t)
        else:
            pts = [0, t // 4, t // 2, (3 * t) // 4, t]
        cu_seqlens = torch.tensor(pts, device=device, dtype=torch.long)
        chunk_indices = prepare_chunk_indices(cu_seqlens, args.chunk_size)
        n_seq = len(cu_seqlens) - 1

    if args.state_index_mode == "none":
        state_rows = n_seq
        initial_state_cur = torch.randn((state_rows, h, v_dim, k_dim), device=device, dtype=torch.float32) * 0.02
        initial_state_ref = initial_state_cur.clone()
        initial_state_indices_cur = None
        initial_state_indices_ref = torch.arange(n_seq, device=device, dtype=torch.int32)
    elif args.state_index_mode == "identity":
        state_rows = n_seq
        initial_state_cur = torch.randn((state_rows, h, v_dim, k_dim), device=device, dtype=torch.float32) * 0.02
        initial_state_ref = initial_state_cur.clone()
        initial_state_indices_cur = torch.arange(n_seq, device=device, dtype=torch.int32)
        initial_state_indices_ref = initial_state_indices_cur
    elif args.state_index_mode == "reverse":
        state_rows = n_seq
        initial_state_cur = torch.randn((state_rows, h, v_dim, k_dim), device=device, dtype=torch.float32) * 0.02
        initial_state_ref = initial_state_cur.clone()
        initial_state_indices_cur = torch.arange(n_seq - 1, -1, -1, device=device, dtype=torch.int32)
        initial_state_indices_ref = initial_state_indices_cur
    elif args.state_index_mode == "random":
        state_rows = max(n_seq + 8, n_seq * 2)
        initial_state_cur = torch.randn((state_rows, h, v_dim, k_dim), device=device, dtype=torch.float32) * 0.02
        initial_state_ref = initial_state_cur.clone()
        initial_state_indices_cur = torch.randperm(state_rows, device=device, dtype=torch.int64)[:n_seq].to(torch.int32)
        initial_state_indices_ref = initial_state_indices_cur
    else:
        raise ValueError(f"Unsupported state_index_mode: {args.state_index_mode}")

    return (
        q,
        k,
        w,
        u,
        g,
        cu_seqlens,
        chunk_indices,
        initial_state_cur,
        initial_state_ref,
        initial_state_indices_cur,
        initial_state_indices_ref,
    )


def main():
    parser = argparse.ArgumentParser("Benchmark chunk gated kernels: current vs ref")
    parser.add_argument("--batch", type=int, default=1)
    parser.add_argument("--seqlen", type=int, default=13320)
    parser.add_argument("--heads", type=int, default=16)
    parser.add_argument("--grouped-heads", type=int, default=8)
    parser.add_argument("--k-dim", type=int, default=128)
    parser.add_argument("--v-dim", type=int, default=128)
    parser.add_argument("--chunk-size", type=int, default=64)
    parser.add_argument("--varlen", action="store_true", default=False)
    parser.add_argument("--varlen-splits", type=str, default="")
    parser.add_argument("--warmup", type=int, default=25)
    parser.add_argument("--rep", type=int, default=100)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--input-scale", type=float, default=0.2)
    parser.add_argument("--g-scale", type=float, default=0.05)
    parser.add_argument(
        "--state-index-mode",
        type=str,
        default="identity",
        choices=("none", "identity", "reverse", "random"),
    )
    args = parser.parse_args()

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

    (
        q,
        k,
        w,
        u,
        g,
        cu_seqlens,
        chunk_indices,
        initial_state_cur,
        initial_state_ref,
        initial_state_indices_cur,
        initial_state_indices_ref,
    ) = build_inputs(args)

    h_flops = (
        estimate_varlen_flops(cu_seqlens, args.heads, args.k_dim, args.v_dim, args.chunk_size, kind="h")
        if args.varlen
        else estimate_fwd_h_flops(args.batch, args.seqlen, args.heads, args.k_dim, args.v_dim, args.chunk_size)
    )
    o_flops = (
        estimate_varlen_flops(cu_seqlens, args.heads, args.k_dim, args.v_dim, args.chunk_size, kind="o")
        if args.varlen
        else estimate_fwd_o_flops(args.batch, args.seqlen, args.heads, args.k_dim, args.v_dim, args.chunk_size)
    )

    h_cur, v_new_cur, _ = chunk_gated_delta_rule_fwd_h(
        k=k,
        w=w,
        u=u,
        g=g,
        initial_state=initial_state_cur,
        initial_state_indices=initial_state_indices_cur,
        output_final_state=True,
        chunk_size=args.chunk_size,
        save_new_value=True,
        cu_seqlens=cu_seqlens,
        chunk_indices=chunk_indices,
        use_exp2=False,
        transpose_state_layout=True,
    )
    h_ref, v_new_ref = chunk_gated_delta_rule_fwd_h_ref(
        k=k,
        w=w,
        u=u,
        g=g,
        gk=None,
        initial_state=initial_state_ref,
        initial_state_indices=initial_state_indices_ref,
        save_new_value=True,
        cu_seqlens=cu_seqlens,
    )
    _ = chunk_fwd_o(
        q=q,
        k=k,
        v=v_new_cur,
        h=h_cur,
        g=g,
        scale=args.k_dim ** -0.5,
        cu_seqlens=cu_seqlens,
        chunk_size=args.chunk_size,
        chunk_indices=chunk_indices,
        use_exp2=False,
        transpose_state_layout=True,
    )
    _ = chunk_fwd_o_ref(
        q=q,
        k=k,
        v=v_new_ref,
        h=h_ref,
        g=g,
        scale=args.k_dim ** -0.5,
        cu_seqlens=cu_seqlens,
        chunk_size=args.chunk_size,
    )
    torch.cuda.synchronize()

    fn_h_cur = lambda: chunk_gated_delta_rule_fwd_h(  # noqa: E731
        k=k,
        w=w,
        u=u,
        g=g,
        initial_state=initial_state_cur,
        initial_state_indices=initial_state_indices_cur,
        output_final_state=True,
        chunk_size=args.chunk_size,
        save_new_value=True,
        cu_seqlens=cu_seqlens,
        chunk_indices=chunk_indices,
        use_exp2=False,
        transpose_state_layout=True,
    )
    fn_h_ref = lambda: chunk_gated_delta_rule_fwd_h_ref(  # noqa: E731
        k=k,
        w=w,
        u=u,
        g=g,
        gk=None,
        initial_state=initial_state_ref,
        initial_state_indices=initial_state_indices_ref,
        save_new_value=True,
        cu_seqlens=cu_seqlens,
    )

    h_cur_cached, v_new_cur_cached, _ = fn_h_cur()
    h_ref_cached, v_new_ref_cached = fn_h_ref()

    fn_o_cur = lambda: chunk_fwd_o(  # noqa: E731
        q=q,
        k=k,
        v=v_new_cur_cached,
        h=h_cur_cached,
        g=g,
        scale=args.k_dim ** -0.5,
        cu_seqlens=cu_seqlens,
        chunk_size=args.chunk_size,
        chunk_indices=chunk_indices,
        use_exp2=False,
        transpose_state_layout=True,
    )
    fn_o_ref = lambda: chunk_fwd_o_ref(  # noqa: E731
        q=q,
        k=k,
        v=v_new_ref_cached,
        h=h_ref_cached,
        g=g,
        scale=args.k_dim ** -0.5,
        cu_seqlens=cu_seqlens,
        chunk_size=args.chunk_size,
    )

    def fn_e2e_cur():
        h_e2e, v_new_e2e, _ = fn_h_cur()
        return chunk_fwd_o(
            q=q,
            k=k,
            v=v_new_e2e,
            h=h_e2e,
            g=g,
            scale=args.k_dim ** -0.5,
            cu_seqlens=cu_seqlens,
            chunk_size=args.chunk_size,
            chunk_indices=chunk_indices,
            use_exp2=False,
            transpose_state_layout=True,
        )

    def fn_e2e_ref():
        h_e2e, v_new_e2e = fn_h_ref()
        return chunk_fwd_o_ref(
            q=q,
            k=k,
            v=v_new_e2e,
            h=h_e2e,
            g=g,
            scale=args.k_dim ** -0.5,
            cu_seqlens=cu_seqlens,
            chunk_size=args.chunk_size,
        )

    ms_h_cur = triton.testing.do_bench(fn_h_cur, warmup=args.warmup, rep=args.rep)
    ms_h_ref = triton.testing.do_bench(fn_h_ref, warmup=args.warmup, rep=args.rep)
    ms_o_cur = triton.testing.do_bench(fn_o_cur, warmup=args.warmup, rep=args.rep)
    ms_o_ref = triton.testing.do_bench(fn_o_ref, warmup=args.warmup, rep=args.rep)
    ms_e2e_cur = triton.testing.do_bench(fn_e2e_cur, warmup=args.warmup, rep=args.rep)
    ms_e2e_ref = triton.testing.do_bench(fn_e2e_ref, warmup=args.warmup, rep=args.rep)

    print("=" * 112)
    print(
        f"chunk_gated benchmark (current vs ref) | varlen={args.varlen} | "
        f"B={args.batch} T={args.seqlen} H={args.heads} Hg={args.grouped_heads} K={args.k_dim} V={args.v_dim} "
        f"state_index_mode={args.state_index_mode}"
    )
    if cu_seqlens is not None:
        print(f"cu_seqlens={cu_seqlens.tolist()}")
    print("-" * 112)
    print(
        f"h_only : cur {ms_h_cur:.3f} ms ({tflops(h_flops, ms_h_cur):.3f} TF) | "
        f"ref {ms_h_ref:.3f} ms ({tflops(h_flops, ms_h_ref):.3f} TF) | speedup(cur/ref) {ms_h_ref / ms_h_cur:.3f}x"
    )
    print(
        f"o_only : cur {ms_o_cur:.3f} ms ({tflops(o_flops, ms_o_cur):.3f} TF) | "
        f"ref {ms_o_ref:.3f} ms ({tflops(o_flops, ms_o_ref):.3f} TF) | speedup(cur/ref) {ms_o_ref / ms_o_cur:.3f}x"
    )
    print(
        f"e2e    : cur {ms_e2e_cur:.3f} ms ({tflops(h_flops + o_flops, ms_e2e_cur):.3f} TF) | "
        f"ref {ms_e2e_ref:.3f} ms ({tflops(h_flops + o_flops, ms_e2e_ref):.3f} TF) | speedup(cur/ref) {ms_e2e_ref / ms_e2e_cur:.3f}x"
    )
    print("=" * 112)


if __name__ == "__main__":
    main()
