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(
    const float* __restrict__ input,
    const float* __restrict__ dw_weight,
    float* __restrict__ dw_out,
    int batch_size, int channels, int seq_len, int kernel_size, int stride, int padding
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * channels * seq_len;
    if (idx >= total) return;
    int b = idx / (channels * seq_len);
    int c = (idx / seq_len) % channels;
    int pos = idx % seq_len;
    float sum = 0.0f;
    for (int k = 0; k < kernel_size; ++k) {
        int in_pos = pos * stride + k - padding;
        if (in_pos >= 0 && in_pos < seq_len) {
            sum += input[b*channels*seq_len + c*seq_len + in_pos] * dw_weight[c*kernel_size + k];
        }
    }
    dw_out[b*channels*seq_len + c*seq_len + pos] = sum;
}

__global__ void pointwise_conv1d(
    const float* __restrict__ dw_out,
    const float* __restrict__ pw_weight,
    float* __restrict__ output,
    int batch_size, int in_channels, int out_channels, int seq_len
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * out_channels * seq_len;
    if (idx >= total) return;
    int b = idx / (out_channels * seq_len);
    int oc = (idx / seq_len) % out_channels;
    int pos = idx % seq_len;
    double sum = 0.0;
    for (int ic = 0; ic < in_channels; ++ic) {
        sum += (double)dw_out[b*in_channels*seq_len + ic*seq_len + pos] * (double)pw_weight[oc*in_channels + ic];
    }
    output[b*out_channels*seq_len + oc*seq_len + pos] = (float)sum;
}

torch::Tensor fused_sepconv_cuda(
    torch::Tensor input,
    torch::Tensor dw_weight,
    torch::Tensor pw_weight,
    int stride, int padding, int kernel_size, int in_channels, int out_channels
) {
    int batch_size = input.size(0);
    int seq_len = input.size(2);
    int out_seq_len = (seq_len + 2*padding - kernel_size)/stride + 1;
    auto dw_out = torch::empty({batch_size, in_channels, out_seq_len}, input.options());
    auto output = torch::empty({batch_size, out_channels, out_seq_len}, input.options());

    int total_dw = batch_size * in_channels * out_seq_len;
    int block = 256;
    depthwise_conv1d<<<(total_dw + block-1)/block, block>>>(
        input.data_ptr<float>(), dw_weight.data_ptr<float>(), dw_out.data_ptr<float>(),
        batch_size, in_channels, seq_len, kernel_size, stride, padding
    );

    int total_pw = batch_size * out_channels * out_seq_len;
    pointwise_conv1d<<<(total_pw + block-1)/block, block>>>(
        dw_out.data_ptr<float>(), pw_weight.data_ptr<float>(), output.data_ptr<float>(),
        batch_size, in_channels, out_channels, out_seq_len
    );
    return output;
}
'''

cpp_source = '''
torch::Tensor fused_sepconv_cuda(torch::Tensor, torch::Tensor, torch::Tensor, int, int, int, int, int);
'''

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

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().__init__()
        self.dw = nn.Conv1d(in_channels, in_channels, kernel_size, stride=stride, padding=padding, groups=in_channels, bias=bias)
        self.pw = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
        self.stride = stride
        self.padding = padding
        self.kernel_size = kernel_size
        self.in_c = in_channels
        self.out_c = out_channels
        self.mod = module
    def forward(self, x):
        return self.mod.fused_sepconv_cuda(x, self.dw.weight, self.pw.weight, self.stride, self.padding, self.kernel_size, self.in_c, self.out_c)