import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Scalar-gated delta-rule linear attention recurrence.

    A recurrent state matrix is decayed by a scalar gate per head, updated by a
    rank-1 key/value outer product, then read out by a query vector.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        gate: torch.Tensor,
        beta: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, seq_len, num_heads, head_dim = q.shape
        state = torch.zeros(batch_size, num_heads, head_dim, head_dim, dtype=q.dtype, device=q.device)
        outputs = []

        for t in range(seq_len):
            gate_t = gate[:, t].unsqueeze(-1).unsqueeze(-1)
            beta_t = beta[:, t].unsqueeze(-1).unsqueeze(-1)
            kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2)
            state = gate_t * state + beta_t * kv_outer
            y_t = (q[:, t].unsqueeze(-1) * state).sum(dim=-2)
            outputs.append(y_t)

        return torch.stack(outputs, dim=1)


batch_size = 2
seq_len = 128
num_heads = 4
head_dim = 24


def get_inputs():
    p = popcorn_pri
    q = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    k = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    v = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    gate = torch.sigmoid(torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), dtype=torch.float32))
    beta = torch.sigmoid(torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), dtype=torch.float32))
    return [q, k, v, gate, beta]


def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor delta_rule_forward_cuda(torch::Tensor q,
                                      torch::Tensor k,
                                      torch::Tensor v,
                                      torch::Tensor gate,
                                      torch::Tensor beta);

torch::Tensor delta_rule_forward(torch::Tensor q,
                                 torch::Tensor k,
                                 torch::Tensor v,
                                 torch::Tensor gate,
                                 torch::Tensor beta) {
    return delta_rule_forward_cuda(q, k, v, gate, beta);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAStream.h>

__global__ __launch_bounds__(32, 8)
void delta_rule_kernel(const float* __restrict__ q,
                       const float* __restrict__ k,
                       const float* __restrict__ v,
                       const float* __restrict__ gate,
                       const float* __restrict__ beta,
                       float* __restrict__ out,
                       int B, int T, int H, int D) {
    const int lane = threadIdx.x;  // one warp per (batch, head, output column)
    const int col = blockIdx.x % D;
    int tmp = blockIdx.x / D;
    const int h = tmp % H;
    const int b = tmp / H;

    float state_i = 0.0f;  // element S[i, col], owned by lane i

    const int bh_base = (b * T * H + h) * D;
    const int gh_base = b * T * H + h;

    #pragma unroll 1
    for (int t = 0; t < T; ++t) {
        const int vec_base = bh_base + t * H * D;
        const int scalar_idx = gh_base + t * H;
        const float g = gate[scalar_idx];
        const float be = beta[scalar_idx];
        const float v_col = v[vec_base + col];

        float acc = 0.0f;
        if (lane < D) {
            state_i = g * state_i + be * k[vec_base + lane] * v_col;
            acc = q[vec_base + lane] * state_i;
        }

        // Warp reduction over the 24 contributing rows.
        acc += __shfl_down_sync(0xffffffffu, acc, 16);
        acc += __shfl_down_sync(0xffffffffu, acc, 8);
        acc += __shfl_down_sync(0xffffffffu, acc, 4);
        acc += __shfl_down_sync(0xffffffffu, acc, 2);
        acc += __shfl_down_sync(0xffffffffu, acc, 1);

        if (lane == 0) {
            out[vec_base + col] = acc;
        }
    }
}

torch::Tensor delta_rule_forward_cuda(torch::Tensor q,
                                      torch::Tensor k,
                                      torch::Tensor v,
                                      torch::Tensor gate,
                                      torch::Tensor beta) {
    const int B = (int)q.size(0);
    const int T = (int)q.size(1);
    const int H = (int)q.size(2);
    const int D = (int)q.size(3);

    auto out = torch::empty_like(q);

    const int blocks = B * H * D;
    const dim3 grid(blocks);
    const dim3 threads(32);
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();

    delta_rule_kernel<<<grid, threads, 0, stream>>>(
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(),
        gate.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(),
        B, T, H, D);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="delta_rule_linear_attention_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["delta_rule_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                gate: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
        return _ext.delta_rule_forward(q, k, v, gate, beta)