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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor gated_delta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define T_CONST 128
#define H_CONST 4
#define D_CONST 24

__global__ __launch_bounds__(128) void scan_kernel(
    const float* __restrict__ q,
    const float* __restrict__ k,
    const float* __restrict__ v,
    const float* __restrict__ gate,
    float* __restrict__ out)
{
    int j = blockIdx.x;
    int h = blockIdx.y;
    int b = blockIdx.z;
    int t = threadIdx.x;

    __shared__ float Gs[T_CONST];
    __shared__ float As[T_CONST];
    __shared__ float gates[T_CONST];
    __shared__ float vals[T_CONST];

    int base = ((b * T_CONST + t) * H_CONST + h) * D_CONST;
    gates[t] = gate[base + j];
    vals[t] = v[base + j];

    float acc = 0.0f;
    __syncthreads();

    #pragma unroll
    for (int row = 0; row < D_CONST; ++row) {
        float G = gates[t];
        float A = k[base + row] * vals[t];
        Gs[t] = G;
        As[t] = A;
        __syncthreads();

        #pragma unroll
        for (int off = 1; off < T_CONST; off <<= 1) {
            float gp = 1.0f;
            float ap = 0.0f;
            if (t >= off) {
                gp = Gs[t - off];
                ap = As[t - off];
            }
            __syncthreads();
            if (t >= off) {
                A = fmaf(G, ap, A);
                G = G * gp;
            }
            Gs[t] = G;
            As[t] = A;
            __syncthreads();
        }
        acc = fmaf(q[base + row], A, acc);
        __syncthreads();
    }
    out[base + j] = acc;
}

torch::Tensor gated_delta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate) {
    auto out = torch::empty({q.size(0), q.size(1), q.size(2), q.size(3)}, q.options());
    dim3 grid(D_CONST, H_CONST, (unsigned)q.size(0));
    scan_kernel<<<grid, 128, 0, at::cuda::getCurrentCUDAStream()>>>(
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), gate.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
"""

_ext = load_inline(
    name="gated_delta_scan_submit",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["gated_delta_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_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) -> torch.Tensor:
        return _ext.gated_delta_forward(q, k, v, gate)