import time
import math
import torch
import torch.nn.functional as F

def naive_online_causal_attn(q, k, v, block_size=128, causal=True):
    """
    Block‑wise causal attention with streaming softmax.

    Args:
      q: (B, Hq, T, D)   — Query tensor
      k: (B, Hk, D, T)   — Key tensor
      v: (B, Hk, D, T)   — Value tensor
      block_size: int   — size of each sequence block
      causal: bool      — if True, masks out future keys

    Returns:
      out: (B, Hq, T, D) — Attention output
    """
    B, Hq, T, D = q.shape
    _, Hk, Dk, Tk = k.shape
    assert Dk == D, "Key/value head dim must match query head dim"
    assert T == Tk, "Sequence length of q and k/v must match"
    assert Hq % Hk == 0, "Number of q‑heads must be a multiple of k‑heads"
    div = Hq // Hk
    num_blocks = T // block_size
    scale = 1.0 / math.sqrt(D)

    out = torch.zeros_like(q)

    if causal:
        mask = torch.triu(
            torch.ones(block_size, block_size, device=q.device, dtype=torch.bool),
            diagonal=1
        ).view(1, 1, 1, block_size, block_size)

    for qb in range(num_blocks):
        q_blk = q[:, :, qb*block_size:(qb+1)*block_size, :].view(B, Hk, div, block_size, D)
        m_running = torch.full((B, Hk, div, block_size), -float('inf'), device=q.device)
        l_running = torch.zeros((B, Hk, div, block_size), device=q.device)
        S_running = torch.zeros((B, Hk, div, block_size, D), device=q.device)

        for kb in range(num_blocks):
            if causal and (kb > qb):
                break

            k_blk = k[:, :, :, kb*block_size:(kb+1)*block_size].transpose(-2, -1)
            v_blk = v[:, :, :, kb*block_size:(kb+1)*block_size].transpose(-2, -1)

            scores = torch.einsum('bhqld,bhkd->bhqlk', q_blk, k_blk) * scale
            if causal and kb == qb:
                scores = scores.masked_fill(mask, float('-inf'))

            max_k = scores.max(dim=-1).values
            prev_m = m_running
            m_new = torch.maximum(prev_m, max_k)

            exp_old = torch.exp(prev_m - m_new)
            exp_new = torch.exp(scores - m_new.unsqueeze(-1))

            l_running = l_running * exp_old + exp_new.sum(dim=-1)
            weighted_new = torch.einsum('bhqlk,bhkd->bhqld', exp_new, v_blk)
            S_running = S_running * exp_old.unsqueeze(-1) + weighted_new

            m_running = m_new

        out_block = S_running / l_running.unsqueeze(-1)
        out[:, :, qb*block_size:(qb+1)*block_size, :] = out_block.view(B, Hq, block_size, D)

    return out


def benchmark_attention(device, configs, num_iters=10, causal=True):
    results = []
    for cfg in configs:
        B, H, T, D, block_size = cfg
        # prepare inputs
        q = torch.randn(B, H, T, D, device=device)
        k = torch.randn(B, H, T, D, device=device)
        v = torch.randn(B, H, T, D, device=device)
        k_naive = k.transpose(2, 3)  # (B, H, D, T)
        v_naive = v.transpose(2, 3)

        # warm‑up
        _ = naive_online_causal_attn(q, k_naive, v_naive, block_size, causal)
        _ = F.scaled_dot_product_attention(q, k, v, attn_mask=None,
                                           dropout_p=0.0, is_causal=causal)
        if device.type == 'cuda':
            torch.cuda.synchronize()

        # benchmark naive
        torch.cuda.reset_peak_memory_stats(device)
        start = time.perf_counter()
        for _ in range(num_iters):
            out1 = naive_online_causal_attn(q, k_naive, v_naive, block_size, causal)
        if device.type == 'cuda':
            torch.cuda.synchronize()
        t1 = (time.perf_counter() - start) / num_iters
        mem1 = torch.cuda.max_memory_allocated(device) / (1024**2) if device.type == 'cuda' else -1

        # benchmark native
        torch.cuda.reset_peak_memory_stats(device)
        start = time.perf_counter()
        for _ in range(num_iters):
            out2 = F.scaled_dot_product_attention(q, k, v, attn_mask=None,
                                                  dropout_p=0.0, is_causal=causal)
        if device.type == 'cuda':
            torch.cuda.synchronize()
        t2 = (time.perf_counter() - start) / num_iters
        mem2 = torch.cuda.max_memory_allocated(device) / (1024**2) if device.type == 'cuda' else -1

        results.append({
            'B': B, 'H': H, 'T': T, 'D': D, 'blk': block_size,
            'naive_time_s': t1, 'native_time_s': t2,
            'naive_peak_mem_MB': mem1, 'native_peak_mem_MB': mem2
        })

    return results


if __name__ == '__main__':
    # device selection
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Running on {device}')

    # configurations: (batch, heads, seq_len, dim, block_size)
    configs = [
        (1, 4, 128, 384, 32),
        (2, 8, 256, 384, 64),
        (4, 16, 512, 768, 64),
        (4, 16, 512, 768, 128),
        (8, 32, 1024, 2048, 128),
        (1, 32, 4096, 4096, 128),
    ]

    # run benchmarks
    results = benchmark_attention(device, configs, num_iters=20, causal=True)

    # pretty‑print
    header = f"{'B':>2} {'H':>2} {'T':>4} {'D':>3} {'blk':>4} | {'naive (ms)':>10} {'native (ms)':>12} | {'naive MB':>8} {'native MB':>10}"
    print(header)
    print('-' * len(header))
    for r in results:
        print(f"{r['B']:>2} {r['H']:>2} {r['T']:>4} {r['D']:>3} {r['blk']:>4} | "
              f"{r['naive_time_s']*1e3:10.2f} {r['native_time_s']*1e3:12.2f} | "
              f"{r['naive_peak_mem_MB']:8.1f} {r['native_peak_mem_MB']:10.1f}")

: make this faster compared to F.scaled_dot_product_attention, do not suggest to use flash attention or existing attentions, this is a custom solution, just modify algorithmically etc, give full code modified then.