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 = r'''
#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 N, int C, int L_in, int L_out,
    int K, int stride, int padding, bool has_bias) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C * L_out;
    if (idx >= total) return;
    
    int l_out = idx % L_out;
    int c = (idx / L_out) % C;
    int n = idx / (L_out * C);
    
    float sum = 0.0f;
    for (int k = 0; k < K; ++k) {
        int l_in = l_out * stride + k - padding;
        if (l_in >= 0 && l_in < L_in) {
            sum += input[n * C * L_in + c * L_in + l_in] * weight[c * K + k];
        }
    }
    if (has_bias) sum += bias[c];
    output[n * C * L_out + c * L_out + l_out] = sum;
}

__global__ void pointwise_conv1d_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int N, int C_in, int C_out, int L,
    bool has_bias) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C_out * L;
    if (idx >= total) return;
    
    int l = idx % L;
    int c_out = (idx / L) % C_out;
    int n = idx / (L * C_out);
    
    float sum = 0.0f;
    for (int c_in = 0; c_in < C_in; ++c_in) {
        sum += input[n * C_in * L + c_in * L + l] * weight[c_out * C_in + c_in];
    }
    if (has_bias) sum += bias[c_out];
    output[n * C_out * L + c_out * L + l] = sum;
}

torch::Tensor depthwise_separable_conv_cuda(
    torch::Tensor x,
    torch::Tensor dw_weight,
    torch::Tensor pw_weight,
    torch::Tensor dw_bias,
    torch::Tensor pw_bias,
    int stride, int padding) {
    
    int N = x.size(0);
    int C_in = x.size(1);
    int L_in = x.size(2);
    int K = dw_weight.size(2);
    bool dw_has_bias = dw_bias.defined() && dw_bias.numel() > 0;
    bool pw_has_bias = pw_bias.defined() && pw_bias.numel() > 0;
    
    int L_out = (L_in + 2 * padding - K) / stride + 1;
    
    auto dw_out = torch::empty({N, C_in, L_out}, x.options());
    
    int total_dw = N * C_in * L_out;
    int threads = 256;
    int blocks = (total_dw + threads - 1) / threads;
    
    depthwise_conv1d_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        dw_weight.data_ptr<float>(),
        dw_has_bias ? dw_bias.data_ptr<float>() : nullptr,
        dw_out.data_ptr<float>(),
        N, C_in, L_in, L_out, K, stride, padding, dw_has_bias);
    
    int C_out = pw_weight.size(0);
    auto out = torch::empty({N, C_out, L_out}, x.options());
    
    int total_pw = N * C_out * L_out;
    blocks = (total_pw + threads - 1) / threads;
    
    pointwise_conv1d_kernel<<<blocks, threads>>>(
        dw_out.data_ptr<float>(),
        pw_weight.data_ptr<float>(),
        pw_has_bias ? pw_bias.data_ptr<float>() : nullptr,
        out.data_ptr<float>(),
        N, C_in, C_out, L_out, pw_has_bias);
    
    return out;
}
'''

cpp_source = r'''
torch::Tensor depthwise_separable_conv_cuda(
    torch::Tensor x,
    torch::Tensor dw_weight,
    torch::Tensor pw_weight,
    torch::Tensor dw_bias,
    torch::Tensor pw_bias,
    int stride, int padding);

'''

_ext = load_inline(
    name="depthwise_separable_conv",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["depthwise_separable_conv_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"]
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = 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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dw_w = self.depthwise.weight
        pw_w = self.pointwise.weight
        dw_b = self.depthwise.bias if self.depthwise.bias is not None else torch.empty(0, device=x.device)
        pw_b = self.pointwise.bias if self.pointwise.bias is not None else torch.empty(0, device=x.device)
        return _ext.depthwise_separable_conv_cuda(
            x, dw_w, pw_w, dw_b, pw_b, self.stride, self.padding
        )