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

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

__global__ void depthwise_conv1d_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int batch_size,
    int channels,
    int seq_len,
    int kernel_size,
    int padding,
    int stride,
    bool has_bias
) {
    // Each thread handles one output element
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    
    int total_outputs = batch_size * channels * seq_len;
    if (tid >= total_outputs) return;
    
    int out_pos = tid % seq_len;
    int c = (tid / seq_len) % channels;
    int b = tid / (channels * seq_len);
    
    // Weight index: kernel for channel c is at index c
    // Weight layout: [channels, 1, kernel_size] -> indexed as [c * kernel_size + k]
    
    float sum = 0.0f;
    
    // For each kernel position
    for (int k = 0; k < kernel_size; k++) {
        int in_pos = out_pos * stride + k - padding;
        if (in_pos >= 0 && in_pos < seq_len) {
            int in_idx = b * channels * seq_len + c * seq_len + in_pos;
            int w_idx = c * kernel_size + k;
            sum += input[in_idx] * weight[w_idx];
        }
    }
    
    if (has_bias) {
        sum += bias[c];
    }
    
    int out_idx = b * channels * seq_len + c * seq_len + out_pos;
    output[out_idx] = sum;
}

torch::Tensor depthwise_conv1d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t stride,
    int64_t padding,
    int64_t groups
) {
    int batch_size = input.size(0);
    int channels = input.size(1);
    int seq_len = input.size(2);
    int kernel_size = weight.size(2);
    bool has_bias = bias.has_value();
    
    auto output = torch::empty({batch_size, channels, seq_len}, input.options());
    
    int total_outputs = batch_size * channels * seq_len;
    int threads = 256;
    int blocks = (total_outputs + threads - 1) / threads;
    
    depthwise_conv1d_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        has_bias ? bias.value().data_ptr<float>() : nullptr,
        output.data_ptr<float>(),
        batch_size,
        channels,
        seq_len,
        kernel_size,
        (int)padding,
        (int)stride,
        has_bias
    );
    
    return output;
}
"""

cpp_source = """
torch::Tensor depthwise_conv1d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t stride,
    int64_t padding,
    int64_t groups
);
"""

depthwise_conv1d_ext = load_inline(
    name='depthwise_conv1d_ext',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['depthwise_conv1d_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(ModelNew, self).__init__()
        self.channels = channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        
        # Create weights same as Conv1d for depthwise
        # Depthwise: weight shape is [channels, 1, kernel_size]
        self.weight = nn.Parameter(torch.empty(channels, 1, kernel_size))
        nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5)
        
        if bias:
            self.bias = nn.Parameter(torch.empty(channels))
            fan_in = channels * kernel_size
            bound = 1 / (fan_in ** 0.5) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)
        else:
            self.register_parameter('bias', None)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return depthwise_conv1d_ext.depthwise_conv1d_cuda(
            x, self.weight, self.bias, self.stride, self.padding, self.channels
        )