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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor forward_cuda(torch::Tensor x,
                           torch::Tensor pw_w, torch::Tensor pw_b,
                           torch::Tensor dw_w, torch::Tensor dw_b,
                           torch::Tensor pr_w, torch::Tensor pr_b);
"""

_cuda_src = r"""
#include <torch/types.h>
#include <ATen/ATen.h>
#include <ATen/Context.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <math.h>

#define CHECK_CUBLAS(expr) do { cublasStatus_t _s = (expr); TORCH_CHECK(_s == CUBLAS_STATUS_SUCCESS, "cuBLAS error: ", (int)_s); } while (0)

__global__ void depthwise_gelu_kernel(const float* __restrict__ p,
                                      const float* __restrict__ pw_b,
                                      const float* __restrict__ dw_w,
                                      const float* __restrict__ dw_b,
                                      float* __restrict__ g,
                                      int B, int E, int L, int K) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * E * L;
    if (idx >= total) return;

    int l = (int)(idx % L);
    long long t = idx / L;
    int e = (int)(t % E);
    int b = (int)(t / E);
    int pad = K >> 1;

    const float* pbe = p + ((long long)b * E + e) * L;
    const float* dwe = dw_w + (long long)e * K;
    float pb = pw_b[e];
    float acc = dw_b[e];

    #pragma unroll
    for (int j = 0; j < 16; ++j) {
        if (j >= K) break;
        int src = l + j - pad;
        if ((unsigned)src < (unsigned)L) {
            acc += dwe[j] * (pbe[src] + pb);
        }
    }

    const float inv_sqrt2 = 0.70710678118654752440f;
    float gelu = 0.5f * acc * (1.0f + erff(acc * inv_sqrt2));
    g[idx] = gelu;
}

__global__ void add_residual_bias_kernel(float* __restrict__ y,
                                         const float* __restrict__ x,
                                         const float* __restrict__ bias,
                                         int B, int C, int L) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * C * L;
    if (idx >= total) return;
    int l = (int)(idx % L);
    long long t = idx / L;
    int c = (int)(t % C);
    (void)l;
    y[idx] = y[idx] + x[idx] + bias[c];
}

torch::Tensor forward_cuda(torch::Tensor x,
                           torch::Tensor pw_w, torch::Tensor pw_b,
                           torch::Tensor dw_w, torch::Tensor dw_b,
                           torch::Tensor pr_w, torch::Tensor pr_b) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(pw_w.is_contiguous() && pw_b.is_contiguous() && dw_w.is_contiguous() && dw_b.is_contiguous() && pr_w.is_contiguous() && pr_b.is_contiguous(), "weights must be contiguous");

    const c10::cuda::CUDAGuard device_guard(x.device());

    int B = (int)x.size(0);
    int C = (int)x.size(1);
    int L = (int)x.size(2);
    int E = (int)pw_w.size(0);
    int K = (int)dw_w.size(2);

    auto p = at::empty({B, E, L}, x.options());
    auto g = at::empty({B, E, L}, x.options());
    auto y = at::empty({B, C, L}, x.options());

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    CHECK_CUBLAS(cublasSetStream(handle, stream));
#if defined(CUBLAS_PEDANTIC_MATH)
    if (at::globalContext().allowTF32CuDNN() || at::globalContext().allowTF32CuBLAS()) {
        CHECK_CUBLAS(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
    } else {
        CHECK_CUBLAS(cublasSetMathMode(handle, CUBLAS_PEDANTIC_MATH));
    }
#endif

    const float alpha = 1.0f;
    const float beta0 = 0.0f;

    // Pointwise 1x1 conv as B independent row-major GEMMs:
    // p[b](E x L) = pw_w(E x C) * x[b](C x L), stored contiguous row-major.
    CHECK_CUBLAS(cublasSgemmStridedBatched(
        handle, CUBLAS_OP_N, CUBLAS_OP_N,
        L, E, C,
        &alpha,
        x.data_ptr<float>(), L, (long long)C * L,
        pw_w.data_ptr<float>(), C, 0LL,
        &beta0,
        p.data_ptr<float>(), L, (long long)E * L,
        B));

    int threads = 256;
    long long total_ge = (long long)B * E * L;
    int blocks_ge = (int)((total_ge + threads - 1) / threads);
    depthwise_gelu_kernel<<<blocks_ge, threads, 0, stream>>>(
        p.data_ptr<float>(), pw_b.data_ptr<float>(), dw_w.data_ptr<float>(), dw_b.data_ptr<float>(),
        g.data_ptr<float>(), B, E, L, K);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    // Project back: y[b](C x L) = pr_w(C x E) * g[b](E x L).
    CHECK_CUBLAS(cublasSgemmStridedBatched(
        handle, CUBLAS_OP_N, CUBLAS_OP_N,
        L, C, E,
        &alpha,
        g.data_ptr<float>(), L, (long long)E * L,
        pr_w.data_ptr<float>(), E, 0LL,
        &beta0,
        y.data_ptr<float>(), L, (long long)C * L,
        B));

    long long total_y = (long long)B * C * L;
    int blocks_y = (int)((total_y + threads - 1) / threads);
    add_residual_bias_kernel<<<blocks_y, threads, 0, stream>>>(
        y.data_ptr<float>(), x.data_ptr<float>(), pr_b.data_ptr<float>(), B, C, L);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return y;
}
"""

_ext = load_inline(
    name="dsconv1d_fused_cublas_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["forward_cuda"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    with_cuda=True,
    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)

    def forward(self, x):
        return _ext.forward_cuda(x,
                                 self.pointwise.weight, self.pointwise.bias,
                                 self.depthwise.weight, self.depthwise.bias,
                                 self.project_back.weight, self.project_back.bias)