"""ZeroGPU demo: compile and run real custom CUDA kernels on Hugging Face's
free, dynamically-allocated GPU.

Two tabs:
  1. Sliding-window attention -- a from-scratch CUDA kernel (JIT-compiled
     with torch.utils.cpp_extension.load_inline) implementing Longformer-
     style local attention with online softmax, run against a dense-PyTorch
     reference for correctness + timing comparison.
  2. Kernel fusion compiler -- takes a small elementwise computation graph
     (y = gelu(x*w + b)), fuses it into ONE generated CUDA kernel, compiles
     it, and runs it against the naive 3-kernel-launch PyTorch equivalent.

All CUDA work happens inside functions decorated with @spaces.GPU, which is
required on ZeroGPU Spaces: the GPU is only attached to the process for the
duration of that call.
"""

import os
import sys
import time
import traceback

import gradio as gr
import spaces
import torch

sys.path.insert(0, os.path.dirname(__file__))

from fusion_compiler.graph import Graph
from fusion_compiler.fuser import fuse
from fusion_compiler.codegen import generate_cuda_source

# ---------------------------------------------------------------------------
# Kernel 1: sliding-window attention (naive-but-correct, single thread per
# query row, online softmax). See long-context-attention-kernels/csrc for
# the tiled/shared-memory production version this is a simplified sibling
# of -- this one is written to be trivially JIT-compilable in one file.
# ---------------------------------------------------------------------------

_ATTN_CPP_SRC = """
torch::Tensor sliding_window_attention_naive(torch::Tensor Q, torch::Tensor K, torch::Tensor V, int64_t window, double scale);
"""

_ATTN_CUDA_SRC = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#define MAX_HEAD_DIM 128

__global__ void sliding_window_attn_kernel(
    const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V,
    float* __restrict__ O, int batch_heads, int seq_len, int head_dim, int window, float scale)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_heads * seq_len;
    if (idx >= total) return;

    int bh = idx / seq_len;
    int q = idx % seq_len;

    const float* Qp = Q + ((long)bh * seq_len + q) * head_dim;
    const float* Kbase = K + (long)bh * seq_len * head_dim;
    const float* Vbase = V + (long)bh * seq_len * head_dim;
    float* Op = O + ((long)bh * seq_len + q) * head_dim;

    int k_lo = max(0, q - window);
    int k_hi = min(seq_len - 1, q + window);

    float m = -1e30f, l = 0.f;
    float acc[MAX_HEAD_DIM];
    for (int d = 0; d < head_dim; ++d) acc[d] = 0.f;

    for (int k = k_lo; k <= k_hi; ++k) {
        const float* Kp = Kbase + (long)k * head_dim;
        float score = 0.f;
        for (int d = 0; d < head_dim; ++d) score += Qp[d] * Kp[d];
        score *= scale;

        float m_new = fmaxf(m, score);
        float corr = expf(m - m_new);
        float p = expf(score - m_new);
        l = l * corr + p;

        const float* Vp = Vbase + (long)k * head_dim;
        for (int d = 0; d < head_dim; ++d) acc[d] = acc[d] * corr + p * Vp[d];
        m = m_new;
    }

    float inv_l = (l > 0.f) ? 1.f / l : 0.f;
    for (int d = 0; d < head_dim; ++d) Op[d] = acc[d] * inv_l;
}

torch::Tensor sliding_window_attention_naive(torch::Tensor Q, torch::Tensor K, torch::Tensor V, int64_t window, double scale) {
    TORCH_CHECK(Q.is_cuda() && K.is_cuda() && V.is_cuda(), "Q, K, V must be CUDA tensors");
    TORCH_CHECK(Q.scalar_type() == torch::kFloat32, "expected float32 tensors");
    auto Qc = Q.contiguous();
    auto Kc = K.contiguous();
    auto Vc = V.contiguous();

    int batch = Qc.size(0), heads = Qc.size(1), seq_len = Qc.size(2), head_dim = Qc.size(3);
    TORCH_CHECK(head_dim <= 128, "demo kernel supports head_dim <= 128");

    auto O = torch::empty_like(Qc);
    int batch_heads = batch * heads;
    int total = batch_heads * seq_len;
    int threads = 128;
    int blocks = (total + threads - 1) / threads;

    sliding_window_attn_kernel<<<blocks, threads>>>(
        Qc.data_ptr<float>(), Kc.data_ptr<float>(), Vc.data_ptr<float>(), O.data_ptr<float>(),
        batch_heads, seq_len, head_dim, (int)window, (float)scale);

    return O;
}
"""

_attn_module = None


def _get_attn_module():
    global _attn_module
    if _attn_module is None:
        from torch.utils.cpp_extension import load_inline
        _attn_module = load_inline(
            name="sliding_window_attention_demo",
            cpp_sources=[_ATTN_CPP_SRC],
            cuda_sources=[_ATTN_CUDA_SRC],
            functions=["sliding_window_attention_naive"],
            verbose=False,
        )
    return _attn_module


def _dense_reference(Q, K, V, window, scale):
    seq_len = Q.shape[2]
    idx = torch.arange(seq_len, device=Q.device)
    mask = (idx[:, None] - idx[None, :]).abs() <= window
    scores = torch.einsum("bhqd,bhkd->bhqk", Q, K) * scale
    scores = scores.masked_fill(~mask[None, None, :, :], float("-inf"))
    attn = torch.softmax(scores, dim=-1)
    return torch.einsum("bhqk,bhkd->bhqd", attn, V)


@spaces.GPU(duration=60)
def run_attention_kernel(seq_len: int, window: int, heads: int, head_dim: int):
    log = []
    try:
        log.append(f"Compiling CUDA kernel (first call only; cached after)...")
        mod = _get_attn_module()
        log.append("Compiled.")

        torch.manual_seed(0)
        device = "cuda"
        Q = torch.randn(1, heads, seq_len, head_dim, device=device, dtype=torch.float32)
        K = torch.randn(1, heads, seq_len, head_dim, device=device, dtype=torch.float32)
        V = torch.randn(1, heads, seq_len, head_dim, device=device, dtype=torch.float32)
        scale = 1.0 / (head_dim ** 0.5)

        # warmup + timed run of the custom kernel
        for _ in range(3):
            out_kernel = mod.sliding_window_attention_naive(Q, K, V, window, scale)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(10):
            out_kernel = mod.sliding_window_attention_naive(Q, K, V, window, scale)
        torch.cuda.synchronize()
        t_kernel = (time.perf_counter() - t0) / 10 * 1000

        # dense PyTorch reference, timed the same way
        for _ in range(3):
            out_ref = _dense_reference(Q, K, V, window, scale)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(10):
            out_ref = _dense_reference(Q, K, V, window, scale)
        torch.cuda.synchronize()
        t_ref = (time.perf_counter() - t0) / 10 * 1000

        max_err = (out_kernel - out_ref).abs().max().item()
        gpu_name = torch.cuda.get_device_name(0)

        log.append(f"GPU: {gpu_name}")
        log.append(f"seq_len={seq_len}  window=±{window}  heads={heads}  head_dim={head_dim}")
        log.append(f"custom CUDA kernel: {t_kernel:.3f} ms")
        log.append(f"dense PyTorch (masked full attention): {t_ref:.3f} ms")
        log.append(f"speedup: {t_ref / t_kernel:.2f}x")
        log.append(f"max abs error vs dense reference: {max_err:.2e}  ({'PASS' if max_err < 1e-3 else 'CHECK'})")
        return "\n".join(log)
    except Exception:
        return "\n".join(log) + "\n\nERROR:\n" + traceback.format_exc()


# ---------------------------------------------------------------------------
# Kernel 2: fusion compiler demo -- fuse mul+add+gelu into one kernel,
# compile it, run it, and compare against 3 separate PyTorch ops.
# ---------------------------------------------------------------------------

def _build_example_graph():
    g = Graph()
    g.add("mul", ["x", "w"], "t1")
    g.add("add", ["t1", "b"], "t2")
    g.add("gelu", ["t2"], "y")
    return g


@spaces.GPU(duration=60)
def run_fusion_kernel(n_elements: int):
    log = []
    try:
        from fusion_compiler.jit import compile_group

        graph = _build_example_graph()
        groups = fuse(graph)
        source = generate_cuda_source(groups[0], kernel_name="fused_demo_kernel")
        log.append(f"Graph: {len(graph)} ops -> fused into {len(groups)} kernel(s).")

        fn = compile_group(groups[0], kernel_name="fused_demo_kernel_live")
        log.append("Compiled fused kernel.")

        device = "cuda"
        torch.manual_seed(0)
        x = torch.randn(n_elements, device=device, dtype=torch.float32)
        w = torch.randn(n_elements, device=device, dtype=torch.float32)
        b = torch.randn(n_elements, device=device, dtype=torch.float32)

        for _ in range(3):
            out_fused = fn(x, w, b)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(20):
            out_fused = fn(x, w, b)
        torch.cuda.synchronize()
        t_fused = (time.perf_counter() - t0) / 20 * 1000

        def unfused(x, w, b):
            t1 = x * w
            t2 = t1 + b
            return torch.nn.functional.gelu(t2, approximate="tanh")

        for _ in range(3):
            out_ref = unfused(x, w, b)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(20):
            out_ref = unfused(x, w, b)
        torch.cuda.synchronize()
        t_unfused = (time.perf_counter() - t0) / 20 * 1000

        max_err = (out_fused - out_ref).abs().max().item()
        gpu_name = torch.cuda.get_device_name(0)

        log.append(f"GPU: {gpu_name}")
        log.append(f"n_elements={n_elements}")
        log.append(f"fused (1 kernel launch): {t_fused:.4f} ms")
        log.append(f"unfused (3 kernel launches, PyTorch eager): {t_unfused:.4f} ms")
        log.append(f"speedup: {t_unfused / t_fused:.2f}x")
        log.append(f"max abs error vs PyTorch reference: {max_err:.2e}  ({'PASS' if max_err < 1e-3 else 'CHECK'})")
        log.append("\n--- generated fused kernel source ---\n")
        log.append(source)
        return "\n".join(log)
    except Exception:
        return "\n".join(log) + "\n\nERROR:\n" + traceback.format_exc()


def preview_fused_source():
    """No GPU needed -- just show what the fuser/codegen produce."""
    graph = _build_example_graph()
    groups = fuse(graph)
    return generate_cuda_source(groups[0], kernel_name="fused_demo_kernel")


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="CUDA Kernels — live on ZeroGPU") as demo:
    gr.Markdown(
        """
        # CUDA Kernels, running live on a free GPU
        Both tabs JIT-compile real `.cu` source with `nvcc` and run it on a
        Hugging Face **ZeroGPU** allocation (an A100 slice, attached only
        for the duration of each click). First run per session compiles the
        kernel (a few seconds); later runs reuse the cached build.

        Source: [long-context-attention-kernels](https://github.com/data-geek-astronomy/long-context-attention-kernels) ·
        [cuda-fusion-compiler](https://github.com/data-geek-astronomy/cuda-fusion-compiler)
        """
    )

    with gr.Tab("Sliding-window attention"):
        gr.Markdown("Longformer-style local attention: a hand-written CUDA kernel vs. dense masked PyTorch attention.")
        with gr.Row():
            seq_len_in = gr.Slider(128, 4096, value=1024, step=128, label="seq_len")
            window_in = gr.Slider(8, 512, value=64, step=8, label="window (±)")
        with gr.Row():
            heads_in = gr.Slider(1, 16, value=4, step=1, label="heads")
            head_dim_in = gr.Slider(16, 128, value=64, step=16, label="head_dim")
        attn_btn = gr.Button("Run on GPU", variant="primary")
        attn_out = gr.Textbox(label="Result", lines=10)
        attn_btn.click(run_attention_kernel, inputs=[seq_len_in, window_in, heads_in, head_dim_in], outputs=[attn_out])

    with gr.Tab("Kernel fusion compiler"):
        gr.Markdown("`y = gelu(x*w + b)`: 3 elementwise ops auto-fused into 1 CUDA kernel by `fusion_compiler`.")
        preview_btn = gr.Button("Preview generated source (no GPU needed)")
        preview_out = gr.Code(label="Generated CUDA", language="cpp")
        preview_btn.click(preview_fused_source, outputs=[preview_out])

        n_elements_in = gr.Slider(1024, 1 << 22, value=1 << 20, step=1024, label="n_elements")
        fusion_btn = gr.Button("Compile + run on GPU", variant="primary")
        fusion_out = gr.Textbox(label="Result", lines=16)
        fusion_btn.click(run_fusion_kernel, inputs=[n_elements_in], outputs=[fusion_out])

if __name__ == "__main__":
    demo.launch()