# SPDX-License-Identifier: MIT

import argparse

import torch
import triton

from aiter.ops.triton.fla.sglang.chunk_delta_h import (
    _get_chunk_delta_h_config,
    chunk_gated_delta_rule_fwd_h,
    launch_chunk_gated_delta_rule_fwd_kernel_h_blockdim64,
    prepare_chunk_offsets,
    prepare_chunk_indices,
)
from aiter.ops.triton.fla.sglang.chunk_o import (
    _get_chunk_o_config,
    chunk_fwd_o,
    launch_chunk_fwd_kernel_o,
)
from op_tests.triton_tests.utils.chunk_delta_h_sglang_ref import (
    chunk_gated_delta_rule_fwd_h as chunk_gated_delta_rule_fwd_h_ref,
)
from op_tests.triton_tests.utils.chunk_o_sglang_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 gbps(num_bytes, ms):
    return num_bytes / (ms * 1e-3) / 1e9


def estimate_launch_h_bytes(
    *,
    batch_size,
    seqlen,
    num_heads,
    grouped_heads,
    k_dim,
    v_dim,
    chunk_size,
    cu_seqlens,
    use_g,
    save_new_value,
    use_initial_state,
):
    cfg = _get_chunk_delta_h_config(k_dim, v_dim, chunk_size, num_heads)
    bv = cfg["BV"]
    nv = triton.cdiv(v_dim, bv)

    seq_lens = (
        [int((cu_seqlens[i + 1] - cu_seqlens[i]).item()) for i in range(len(cu_seqlens) - 1)]
        if cu_seqlens is not None
        else [seqlen] * batch_size
    )
    total_nt = sum(triton.cdiv(s, chunk_size) for s in seq_lens)

    # Per (seq, head, v-tile, chunk) iteration traffic.
    # dtypes in this benchmark: k/w/u/h/v_new=f16 (2B), g=f32 (4B).
    read_w = chunk_size * k_dim * 2
    read_k = chunk_size * k_dim * 2
    read_v = chunk_size * bv * 2
    write_h = bv * k_dim * 2
    write_v_new = chunk_size * bv * 2 if save_new_value else 0
    read_g = (chunk_size + 1) * 4 if use_g else 0
    per_iter = read_w + read_k + read_v + write_h + write_v_new + read_g

    total_bytes = total_nt * num_heads * nv * per_iter

    # One-time per (seq, head, v-tile) initial/final state traffic (fp32 state).
    if use_initial_state:
        per_prog_state = bv * k_dim * 4 + bv * k_dim * 4
        total_bytes += len(seq_lens) * num_heads * nv * per_prog_state

    # k is grouped by Hg but loaded per H program (reuse via caches is hardware-dependent).
    _ = grouped_heads
    return total_bytes




def estimate_launch_o_bytes(
    *,
    batch_size,
    seqlen,
    num_heads,
    k_dim,
    v_dim,
    chunk_size,
    cu_seqlens,
    use_g,
):
    # Approximate memory traffic for launch_chunk_fwd_kernel_o.
    # q/k/v/h/o are fp16 (2B), g is fp32 (4B).
    cfg = _get_chunk_o_config(k_dim, v_dim, chunk_size)
    bv = cfg["BV"]
    nv = triton.cdiv(v_dim, bv)

    seq_lens = (
        [int((cu_seqlens[i + 1] - cu_seqlens[i]).item()) for i in range(len(cu_seqlens) - 1)]
        if cu_seqlens is not None
        else [seqlen] * batch_size
    )
    total_nt = sum(triton.cdiv(s, chunk_size) for s in seq_lens)

    # Per (seq, head, v-tile, chunk) iteration traffic.
    # q/k/g are reloaded for each v-tile program in this kernel.
    read_q = chunk_size * k_dim * 2
    read_k = chunk_size * k_dim * 2
    read_h = bv * k_dim * 2
    read_v = chunk_size * bv * 2
    write_o = chunk_size * bv * 2
    read_g = chunk_size * 4 if use_g else 0
    per_iter = read_q + read_k + read_h + read_v + write_o + read_g

    return total_nt * num_heads * nv * per_iter


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_bytes = estimate_launch_h_bytes(
        batch_size=args.batch,
        seqlen=args.seqlen,
        num_heads=args.heads,
        grouped_heads=args.grouped_heads,
        k_dim=args.k_dim,
        v_dim=args.v_dim,
        chunk_size=args.chunk_size,
        cu_seqlens=cu_seqlens,
        use_g=g is not None,
        save_new_value=True,
        use_initial_state=initial_state_cur is not None,
    )
    o_bytes = estimate_launch_o_bytes(
        batch_size=args.batch,
        seqlen=args.seqlen,
        num_heads=args.heads,
        k_dim=args.k_dim,
        v_dim=args.v_dim,
        chunk_size=args.chunk_size,
        cu_seqlens=cu_seqlens,
        use_g=g is not None,
    )

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

    N = args.batch if cu_seqlens is None else (len(cu_seqlens) - 1)
    NT = triton.cdiv(args.seqlen, args.chunk_size) if cu_seqlens is None else len(chunk_indices)
    H = args.heads
    Hg = args.grouped_heads
    K = args.k_dim
    V = args.v_dim
    BT = args.chunk_size
    h_buf = k.new_empty((args.batch, NT, H, V, K))
    v_new_buf = torch.empty_like(u)
    chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) if cu_seqlens is not None else None
    kernel_cfg_h = _get_chunk_delta_h_config(K, V, BT, H)
    o_buf = torch.empty_like(u)

    def fn_h_cur():
        launch_chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
            k=k,
            u=u,
            w=w,
            v_new=v_new_buf,
            g=g,
            gk=None,
            h=h_buf,
            initial_state=initial_state_cur,
            initial_state_indices=initial_state_indices_cur,
            cu_seqlens=cu_seqlens,
            chunk_offsets=chunk_offsets,
            N=N,
            T=args.seqlen,
            H=H,
            Hg=Hg,
            K=K,
            V=V,
            BT=BT,
            kernel_cfg=kernel_cfg_h,
        )
        return h_buf, v_new_buf
    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()

    def fn_o_cur():
        launch_chunk_fwd_kernel_o(
            q=q,
            k=k,
            v=v_new_cur_cached,
            h=h_cur_cached,
            g=g,
            o=o_buf,
            cu_seqlens=cu_seqlens,
            chunk_indices=chunk_indices,
            scale=args.k_dim ** -0.5,
            T=args.seqlen,
            H=H,
            Hg=Hg,
            K=K,
            V=V,
            BT=BT,
            NT=NT,
            B=args.batch,
            kernel_cfg=None,
        )
        return o_buf
    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_o_cur = triton.testing.do_bench(fn_o_cur, warmup=args.warmup, rep=args.rep)
    ms_e2e_cur = triton.testing.do_bench(fn_e2e_cur, warmup=args.warmup, rep=args.rep)
    # ms_e2e_fused = triton.testing.do_bench(fn_e2e_fused, 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 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, {gbps(h_bytes, ms_h_cur):.3f} GB/s)")
    print(
        f"o_only : cur {ms_o_cur:.3f} ms ({tflops(o_flops, ms_o_cur):.3f} TF, {gbps(o_bytes, ms_o_cur):.3f} GB/s)")
    print(
        f"e2e(cur two-kernel) : {ms_e2e_cur:.3f} ms ({tflops(h_flops + o_flops, ms_e2e_cur):.3f} TF)"
    )
    return
    
    print(
        f"h_only : cur {ms_h_cur:.3f} ms ({tflops(h_flops, ms_h_cur):.3f} TF, {gbps(h_bytes, ms_h_cur):.3f} GB/s) | "
        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()
