import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a depthwise 1D convolution where each input channel is convolved with its own filter.
    Fundamental primitive in Hyena, Mamba, and genomic sequence models (Enformer, HyenaDNA).

    Args:
        channels (int): Number of input/output channels.
        kernel_size (int): Size of the convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        bias (bool, optional): If True, adds a learnable bias. Defaults to False.
    """
    def __init__(self, channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(Model, self).__init__()
        self.conv1d = nn.Conv1d(channels, channels, kernel_size, stride=stride, padding=padding, groups=channels, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.conv1d(x)

batch_size = 32
channels = 512
kernel_size = 3
seq_len = 8192
stride = 1
padding = 1

def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), channels, popcorn_pri.jitter_int(seq_len))]

def get_init_inputs():
    return [channels, kernel_size, stride, padding]

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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor depthwise_conv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, bool has_bias, int64_t stride, int64_t padding);
"""

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

#define THREADS 512
#define TILE 2048

__global__ void dwconv3_p1_nobias_vec4(const float* __restrict__ x,
                                       const float* __restrict__ w,
                                       float* __restrict__ y,
                                       int C, int L) {
    int row = blockIdx.y;
    int c = row % C;
    int l = blockIdx.x * TILE + threadIdx.x * 4;
    const float* __restrict__ xr = x + (int64_t)row * L;
    float* __restrict__ yr = y + (int64_t)row * L;

    float w0 = w[c * 3 + 0];
    float w1 = w[c * 3 + 1];
    float w2 = w[c * 3 + 2];

    if (l + 3 < L) {
        float4 v = *reinterpret_cast<const float4*>(xr + l);
        float xm1 = (l > 0) ? xr[l - 1] : 0.0f;
        float xp4 = (l + 4 < L) ? xr[l + 4] : 0.0f;
        float4 o;
        o.x = xm1 * w0 + v.x * w1 + v.y * w2;
        o.y = v.x * w0 + v.y * w1 + v.z * w2;
        o.z = v.y * w0 + v.z * w1 + v.w * w2;
        o.w = v.z * w0 + v.w * w1 + xp4 * w2;
        *reinterpret_cast<float4*>(yr + l) = o;
    } else {
        for (int j = 0; j < 4 && l + j < L; ++j) {
            int oi = l + j;
            float a = (oi > 0) ? xr[oi - 1] : 0.0f;
            float b = xr[oi];
            float cval = (oi + 1 < L) ? xr[oi + 1] : 0.0f;
            yr[oi] = a * w0 + b * w1 + cval * w2;
        }
    }
}

__global__ void dwconv_generic(const float* __restrict__ x,
                               const float* __restrict__ w,
                               const float* __restrict__ bias,
                               float* __restrict__ y,
                               int B, int C, int L, int K, int Lout,
                               int stride, int padding, bool has_bias) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * C * Lout;
    for (; idx < total; idx += blockDim.x * gridDim.x) {
        int lo = idx % Lout;
        int row = idx / Lout;
        int c = row % C;
        int base = lo * stride - padding;
        const float* __restrict__ xr = x + (int64_t)row * L;
        const float* __restrict__ wc = w + (int64_t)c * K;
        float acc = has_bias ? bias[c] : 0.0f;
        for (int k = 0; k < K; ++k) {
            int xi = base + k;
            if ((unsigned)xi < (unsigned)L) acc += xr[xi] * wc[k];
        }
        y[idx] = acc;
    }
}

torch::Tensor depthwise_conv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                       bool has_bias, int64_t stride64, int64_t padding64) {
    int B = (int)x.size(0);
    int C = (int)x.size(1);
    int L = (int)x.size(2);
    int K = (int)weight.size(2);
    int stride = (int)stride64;
    int padding = (int)padding64;
    int Lout = (L + 2 * padding - K) / stride + 1;
    auto y = torch::empty({B, C, Lout}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (K == 3 && stride == 1 && padding == 1 && Lout == L && !has_bias) {
        dim3 block(THREADS);
        dim3 grid((L + TILE - 1) / TILE, B * C);
        dwconv3_p1_nobias_vec4<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), y.data_ptr<float>(), C, L);
    } else {
        int total = B * C * Lout;
        int block = 256;
        int grid = min((total + block - 1) / block, 65535);
        dwconv_generic<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), B, C, L, K, Lout, stride, padding, has_bias);
    }
    return y;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super().__init__()
        self.conv1d = nn.Conv1d(channels, channels, kernel_size, stride=stride, padding=padding, groups=channels, bias=bias)
        self.stride = stride
        self.padding = padding

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv1d.bias
        has_bias = b is not None
        if b is None:
            b = self.conv1d.weight  # dummy CUDA tensor; ignored by no-bias kernels
        return _ext.depthwise_conv1d_forward(x, self.conv1d.weight, b, has_bias, self.stride, self.padding)