import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a depthwise-separable 1D convolution: a depthwise Conv1d (each channel independently)
    followed by a pointwise Conv1d (1x1, mixes channels). Used in Enformer and genomic sequence
    models for efficient spatial+channel factorized processing.

    Args:
        in_channels (int): Number of channels in the input.
        out_channels (int): Number of channels produced by the pointwise convolution.
        kernel_size (int): Kernel size of the depthwise convolution.
        stride (int, optional): Stride of the depthwise convolution. Defaults to 1.
        padding (int, optional): Padding for the depthwise convolution. Defaults to 0.
        bias (bool, optional): If True, adds a learnable bias. Defaults to False.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(Model, self).__init__()
        self.depthwise = nn.Conv1d(in_channels, in_channels, kernel_size, stride=stride, padding=padding, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

batch_size = 32
in_channels = 256
out_channels = 512
kernel_size = 9
seq_len = 4096
stride = 1
padding = 4

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

def get_init_inputs():
    return [in_channels, out_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(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    float* __restrict__ output,
    int N, int C, int L, int K, int stride, int padding) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C * L;
    if (idx >= total) return;
    
    int l_out = idx % L;
    int cin = (idx / L) % C;
    int n = idx / (L * C);
    
    float sum = 0.0f;
    #pragma unroll
    for (int k = 0; k < 9; ++k) {
        int l_in = l_out * stride + k - padding;
        if (l_in >= 0 && l_in < L) {
            sum += input[n * C * L + cin * L + l_in] * weight[cin * K + k];
        }
    }
    output[n * C * L + cin * L + l_out] = sum;
}

torch::Tensor depthwise_conv(torch::Tensor x, torch::Tensor weight, int stride, int padding) {
    int N = x.size(0);
    int C = x.size(1);
    int L = x.size(2);
    int K = weight.size(2);
    
    auto output = torch::empty({N, C, L}, x.options());
    
    int total = N * C * L;
    depthwise_conv1d<<<(total + 255)/256, 256>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        output.data_ptr<float>(),
        N, C, L, K, stride, padding
    );
    
    return output;
}
'''

cpp_source = '''
torch::Tensor depthwise_conv(torch::Tensor x, torch::Tensor weight, int stride, int padding);
'''

module = load_inline(
    name='dw_conv',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['depthwise_conv'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=False):
        super(ModelNew, self).__init__()
        self.depthwise = nn.Conv1d(in_channels, in_channels, kernel_size, stride=stride, padding=padding, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=bias)
        self.stride = stride
        self.padding = padding
        self.dw_conv = module.depthwise_conv
        
    def forward(self, x):
        x = self.dw_conv(x, self.depthwise.weight, self.stride, self.padding)
        x = self.pointwise(x)
        return x