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

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

torch::Tensor dsconv_forward(torch::Tensor x, torch::Tensor dw, torch::Tensor pw,
                             torch::Tensor db, torch::Tensor pb,
                             int64_t stride, int64_t padding);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>

__global__ void depthwise_kernel(const float* __restrict__ x,
                                 const float* __restrict__ w,
                                 const float* __restrict__ b,
                                 float* __restrict__ tmp,
                                 int N, int C, int L, int K, int Lo,
                                 int stride, int pad, int has_bias) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)C * N * Lo;
    if (idx >= total) return;

    int l = idx % Lo;
    long long t = idx / Lo;
    int n = t % N;
    int c = t / N;
    int in0 = l * stride - pad;

    float acc = has_bias ? b[c] : 0.0f;
    const float* xbase = x + ((long long)n * C + c) * L;
    const float* wbase = w + (long long)c * K;
#pragma unroll
    for (int k = 0; k < 32; ++k) {
        if (k >= K) break;
        int xi = in0 + k;
        if ((unsigned)xi < (unsigned)L) acc += xbase[xi] * wbase[k];
    }
    // Layout C x (N*Lo), which is column-major (N*Lo) x C for cuBLAS.
    tmp[(long long)c * N * Lo + (long long)n * Lo + l] = acc;
}

__global__ void add_pw_bias_kernel(float* __restrict__ y,
                                   const float* __restrict__ b,
                                   int O, long long M) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)O * M;
    if (idx < total) y[idx] += b[idx / M];
}

torch::Tensor dsconv_forward(torch::Tensor x, torch::Tensor dw, torch::Tensor pw,
                             torch::Tensor db, torch::Tensor pb,
                             int64_t stride64, int64_t padding64) {
    auto xc = x.contiguous();
    auto dwc = dw.contiguous();
    auto pwc = pw.contiguous();

    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int L = (int)xc.size(2);
    int K = (int)dwc.size(2);
    int O = (int)pwc.size(0);
    int stride = (int)stride64;
    int pad = (int)padding64;
    int Lo = (L + 2 * pad - K) / stride + 1;
    int M = N * Lo;

    auto tmp = torch::empty({C, N, Lo}, xc.options());
    auto ybase = torch::empty({O, N, Lo}, xc.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;
    long long total = (long long)C * N * Lo;
    depthwise_kernel<<<(total + threads - 1) / threads, threads, 0, stream>>>(
        xc.data_ptr<float>(), dwc.data_ptr<float>(),
        db.numel() ? db.data_ptr<float>() : nullptr,
        tmp.data_ptr<float>(), N, C, L, K, Lo, stride, pad, db.numel() != 0);

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, stream);
    cublasMath_t old_mode;
    cublasGetMathMode(handle, &old_mode);
    cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH);

    float alpha = 1.0f, beta = 0.0f;
    // ybase as column-major M x O = tmp(M x C) @ pw(C x O).  pw [O,C,1]
    // is exactly column-major C x O in memory.
    cublasStatus_t st = cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                                    M, O, C,
                                    &alpha,
                                    tmp.data_ptr<float>(), M,
                                    pwc.data_ptr<float>(), C,
                                    &beta,
                                    ybase.data_ptr<float>(), M);
    cublasSetMathMode(handle, old_mode);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "cublasSgemm failed");

    if (pb.numel() != 0) {
        long long ytot = (long long)O * M;
        add_pw_bias_kernel<<<(ytot + threads - 1) / threads, threads, 0, stream>>>(
            ybase.data_ptr<float>(), pb.data_ptr<float>(), O, M);
    }

    // Return logical [N,O,Lo] without a costly physical transpose.
    return ybase.as_strided({N, O, Lo}, {Lo, (long long)N * Lo, 1});
}
'''

_ext = load_inline(name='dsconv_ext_submit',
                   cpp_sources=_cpp_src,
                   cuda_sources=_cuda_src,
                   functions=['dsconv_forward'],
                   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:
        empty = torch.empty(0, device=x.device, dtype=x.dtype)
        db = self.depthwise.bias if self.depthwise.bias is not None else empty
        pb = self.pointwise.bias if self.pointwise.bias is not None else empty
        return _ext.dsconv_forward(x, self.depthwise.weight, self.pointwise.weight,
                                   db, pb, self.stride, self.padding)