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 = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

__device__ float gelu(float x) {
    return 0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x * x * x)));
}

__global__ void fused_dwsep_kernel(
    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__ proj_weight, const float* __restrict__ proj_bias,
    float* __restrict__ output,
    int batch, int channels, int expanded, int seq_len, int ksize, int pad
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch * channels * seq_len;
    if (idx >= total) return;

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

    float res = proj_bias[c];
    for (int ec = 0; ec < expanded; ++ec) {
        float dw_out = dw_bias[ec];
        for (int k = 0; k < ksize; ++k) {
            int is = s + k - pad;
            if (is >= 0 && is < seq_len) {
                float pw = pw_bias[ec];
                for (int ic = 0; ic < channels; ++ic) {
                    pw += input[(b * channels + ic) * seq_len + is] * pw_weight[ec * channels + ic];
                }
                dw_out += pw * dw_weight[ec * ksize + k];
            }
        }
        float act = gelu(dw_out);
        res += act * proj_weight[c * expanded + ec];
    }
    float residual = input[(b * channels + c) * seq_len + s];
    output[idx] = res + residual;
}

torch::Tensor fused_forward(
    torch::Tensor input,
    torch::Tensor pw_weight, torch::Tensor pw_bias,
    torch::Tensor dw_weight, torch::Tensor dw_bias,
    torch::Tensor proj_weight, torch::Tensor proj_bias
) {
    int batch = input.size(0);
    int channels = input.size(1);
    int seq_len = input.size(2);
    int expanded = pw_weight.size(0);
    int ksize = dw_weight.size(2);
    int pad = ksize / 2;

    auto output = torch::empty({batch, channels, seq_len}, input.options());

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

    fused_dwsep_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        pw_weight.data_ptr<float>(), pw_bias.data_ptr<float>(),
        dw_weight.data_ptr<float>(), dw_bias.data_ptr<float>(),
        proj_weight.data_ptr<float>(), proj_bias.data_ptr<float>(),
        output.data_ptr<float>(),
        batch, channels, expanded, seq_len, ksize, pad
    );

    return output;
}
'''

cpp_source = r'''
torch::Tensor fused_forward(
    torch::Tensor input,
    torch::Tensor pw_weight, torch::Tensor pw_bias,
    torch::Tensor dw_weight, torch::Tensor dw_bias,
    torch::Tensor proj_weight, torch::Tensor proj_bias
);
'''

module = load_inline(
    name='fused_dwsep',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['fused_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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)

    def forward(self, x):
        return module.fused_forward(
            x,
            self.pointwise.weight.squeeze(-1),
            self.pointwise.bias,
            self.depthwise.weight.squeeze(1),
            self.depthwise.bias,
            self.project_back.weight.squeeze(-1),
            self.project_back.bias
        )