import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Channel-wise gated delta attention recurrence.

    The recurrent state decays independently for each output channel, which
    makes the update closer to channel-gated delta attention variants used in
    fast long-context baselines.
    """

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

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        gate: 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(-2)
            kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2)
            state = gate_t * state + 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), p.jitter_int(head_dim, align=8), dtype=torch.float32))
    return [q, k, v, gate]


def get_init_inputs():
    return []

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

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void gated_delta_attn_kernel(
    const float* __restrict__ q,
    const float* __restrict__ k,
    const float* __restrict__ v,
    const float* __restrict__ gate,
    float* __restrict__ output,
    int batch_size,
    int seq_len,
    int num_heads,
    int head_dim
) {
    int bh = blockIdx.x * blockDim.x + threadIdx.x;
    int total_bh = batch_size * num_heads;
    if (bh >= total_bh) return;
    int b = bh / num_heads;
    int h = bh % num_heads;
    
    float state[32*32];
    for (int i = 0; i < head_dim * head_dim; i++) state[i] = 0.0f;
    
    int base = b * seq_len * num_heads * head_dim + h * head_dim;
    int stride_t = num_heads * head_dim;
    
    for (int t = 0; t < seq_len; t++) {
        float gate_t[32];
        float k_t[32], v_t[32], q_t[32];
        int t_offset = base + t * stride_t;
        for (int d = 0; d < head_dim; d++) {
            gate_t[d] = gate[t_offset + d];
            k_t[d] = k[t_offset + d];
            v_t[d] = v[t_offset + d];
            q_t[d] = q[t_offset + d];
        }
        
        // update state[l,m] = gate[m] * state[l,m] + k[l]*v[m]
        for (int l = 0; l < head_dim; l++) {
            for (int m = 0; m < head_dim; m++) {
                float kv = k_t[l] * v_t[m];
                int sidx = l * head_dim + m;
                state[sidx] = gate_t[m] * state[sidx] + kv;
            }
        }
        
        // y[m] = sum_l q[l] * state[l,m]
        for (int m = 0; m < head_dim; m++) {
            float sum = 0.0f;
            for (int l = 0; l < head_dim; l++) {
                sum += q_t[l] * state[l * head_dim + m];
            }
            output[t_offset + m] = sum;
        }
    }
}

torch::Tensor gated_delta_attention_cuda(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    torch::Tensor gate
) {
    auto batch_size = q.size(0);
    auto seq_len = q.size(1);
    auto num_heads = q.size(2);
    auto head_dim = q.size(3);
    
    auto output = torch::zeros({batch_size, seq_len, num_heads, head_dim}, q.options());
    
    int total_bh = batch_size * num_heads;
    gated_delta_attn_kernel<<<(total_bh + 255)/256, 256>>>(
        q.data_ptr<float>(),
        k.data_ptr<float>(),
        v.data_ptr<float>(),
        gate.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        seq_len,
        num_heads,
        head_dim
    );
    
    return output;
}
'''

cpp_source = '''
torch::Tensor gated_delta_attention_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate);
'''

module = load_inline(
    name='gated_delta_attn',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['gated_delta_attention_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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

    def forward(self, q, k, v, gate):
        return self.module.gated_delta_attention_cuda(q, k, v, gate)