import torch
import torch.nn as nn

class Model(nn.Module):
    """
    1D depthwise separable convolution block: pointwise Conv1d (channel mixing), depthwise Conv1d (spatial),
    GELU activation, and residual addition. Used in genomic foundation models (Enformer) for efficient
    sequence processing of DNA/RNA.
    """
    def __init__(self, channels, expanded_channels, kernel_size):
        super(Model, self).__init__()
        self.pointwise = nn.Conv1d(channels, expanded_channels, 1)
        self.depthwise = nn.Conv1d(expanded_channels, expanded_channels, kernel_size,
                                   padding=kernel_size // 2, groups=expanded_channels)
        self.activation = nn.GELU()
        self.project_back = nn.Conv1d(expanded_channels, channels, 1)

    def forward(self, x):
        residual = x
        x = self.pointwise(x)
        x = self.depthwise(x)
        x = self.activation(x)
        x = self.project_back(x)
        x = x + residual
        return x

batch_size = 32
channels = 256
expanded_channels = 512
seq_len = 4096
kernel_size = 9

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, expanded_channels, kernel_size]

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>
#include <math.h>

__global__ void fused_depthwise_separable(
    const float* __restrict__ input,
    const float* __restrict__ pw_weight,
    const float* __restrict__ pw_bias,
    const float* __restrict__ dw_weight,
    const float* __restrict__ dw_bias,
    const float* __restrict__ pb_weight,
    const float* __restrict__ pb_bias,
    float* __restrict__ output,
    int batch_size,
    int channels,
    int expanded_channels,
    int seq_len,
    int kernel_size,
    int pad) {

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * channels * seq_len;
    if (idx >= total) return;

    int c_out = idx / (seq_len);
    int b = c_out / channels;
    c_out = c_out % channels;
    int s = idx % seq_len;

    float sum = 0.0f;

    // pointwise -> depthwise -> gelu -> project_back
    for (int ec = 0; ec < expanded_channels; ec++) {
        // pointwise
        float pw_val = 0.0f;
        for (int ic = 0; ic < channels; ic++) {
            pw_val += input[((b * channels + ic) * seq_len) + s] * 
                      pw_weight[ec * channels + ic];
        }
        if (pw_bias) pw_val += pw_bias[ec];

        // depthwise
        float dw_val = 0.0f;
        for (int k = 0; k < kernel_size; k++) {
            int in_s = s + k - pad;
            if (in_s >= 0 && in_s < seq_len) {
                dw_val += pw_val * dw_weight[ec * kernel_size + k];
            }
        }
        if (dw_bias) dw_val += dw_bias[ec];

        // GELU
        float gelu_val = 0.5f * dw_val * (1.0f + tanhf(0.79788456f * (dw_val + 0.044715f * dw_val * dw_val * dw_val)));

        // project back
        sum += gelu_val * pb_weight[c_out * expanded_channels + ec];
    }
    if (pb_bias) sum += pb_bias[c_out];

    // residual
    sum += input[((b * channels + c_out) * seq_len) + s];

    output[((b * channels + c_out) * seq_len) + s] = sum;
}

torch::Tensor fused_forward(
    torch::Tensor x,
    torch::Tensor pw_weight,
    torch::Tensor pw_bias,
    torch::Tensor dw_weight,
    torch::Tensor dw_bias,
    torch::Tensor pb_weight,
    torch::Tensor pb_bias,
    int kernel_size) {

    auto batch_size = x.size(0);
    auto channels = x.size(1);
    auto seq_len = x.size(2);
    auto expanded_channels = pw_weight.size(0);
    int pad = kernel_size / 2;

    auto output = torch::zeros_like(x);

    int total = batch_size * channels * seq_len;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;

    fused_depthwise_separable<<<blocks, threads>>>(
        x.data_ptr<float>(),
        pw_weight.data_ptr<float>(),
        pw_bias.defined() ? pw_bias.data_ptr<float>() : nullptr,
        dw_weight.data_ptr<float>(),
        dw_bias.defined() ? dw_bias.data_ptr<float>() : nullptr,
        pb_weight.data_ptr<float>(),
        pb_bias.defined() ? pb_bias.data_ptr<float>() : nullptr,
        output.data_ptr<float>(),
        batch_size,
        channels,
        expanded_channels,
        seq_len,
        kernel_size,
        pad
    );

    return output;
}
'''

cpp_source = '''
torch::Tensor fused_forward(
    torch::Tensor x,
    torch::Tensor pw_weight,
    torch::Tensor pw_bias,
    torch::Tensor dw_weight,
    torch::Tensor dw_bias,
    torch::Tensor pb_weight,
    torch::Tensor pb_bias,
    int kernel_size);
'''

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

class ModelNew(nn.Module):
    def __init__(self, channels, expanded_channels, kernel_size):
        super(ModelNew, self).__init__()
        self.pointwise = nn.Conv1d(channels, expanded_channels, 1)
        self.depthwise = nn.Conv1d(expanded_channels, expanded_channels, kernel_size,
                                   padding=kernel_size // 2, groups=expanded_channels)
        self.activation = nn.GELU()
        self.project_back = nn.Conv1d(expanded_channels, channels, 1)
        self.kernel_size = kernel_size

    def forward(self, x):
        return module.fused_forward(
            x,
            self.pointwise.weight,
            self.pointwise.bias,
            self.depthwise.weight,
            self.depthwise.bias,
            self.project_back.weight,
            self.project_back.bias,
            self.kernel_size
        )
```