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 depthwise_pointwise_forward(torch::Tensor x,
                                          torch::Tensor dw_w,
                                          torch::Tensor pw_w,
                                          torch::Tensor dw_b,
                                          torch::Tensor pw_b,
                                          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>

__global__ void depthwise_to_blc_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        const float* __restrict__ b,
                                        float* __restrict__ h,
                                        int B, int C, int L, int K,
                                        int Lout, int stride, int padding,
                                        bool has_bias) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * Lout * C;
    if (idx >= total) return;

    int c = idx % C;
    int t = (idx / C) % Lout;
    int n = idx / (C * Lout);

    float acc = has_bias ? b[c] : 0.0f;
    int base = t * stride - padding;
    const float* x_base = x + ((n * C + c) * L);
    const float* w_base = w + c * K;

    #pragma unroll
    for (int k = 0; k < 16; ++k) {
        if (k < K) {
            int xi = base + k;
            if ((unsigned)xi < (unsigned)L) {
                acc += x_base[xi] * w_base[k];
            }
        }
    }
    h[idx] = acc;
}

torch::Tensor depthwise_pointwise_forward(torch::Tensor x,
                                          torch::Tensor dw_w,
                                          torch::Tensor pw_w,
                                          torch::Tensor dw_b,
                                          torch::Tensor pw_b,
                                          int64_t stride64,
                                          int64_t padding64) {
    const int B = (int)x.size(0);
    const int C = (int)x.size(1);
    const int L = (int)x.size(2);
    const int K = (int)dw_w.size(2);
    const int O = (int)pw_w.size(0);
    const int stride = (int)stride64;
    const int padding = (int)padding64;
    const int Lout = (L + 2 * padding - K) / stride + 1;

    auto h = torch::empty({(int64_t)B * Lout, C}, x.options());

    int total = B * Lout * C;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    const float* db = dw_b.numel() ? dw_b.data_ptr<float>() : nullptr;
    depthwise_to_blc_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), dw_w.data_ptr<float>(), db, h.data_ptr<float>(),
        B, C, L, K, Lout, stride, padding, dw_b.numel() != 0);

    // Pointwise 1x1 convolution is exactly a large GEMM: [B*Lout, C] x [C, O].
    auto y = at::matmul(h, pw_w.t());
    if (pw_b.numel() != 0) {
        y.add_(pw_b);
    }
    return y.view({B, Lout, O}).permute({0, 2, 1});
}
'''

_ext = load_inline(
    name="dw_sep_conv1d_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["depthwise_pointwise_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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(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
        self.register_buffer("_empty", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dw_b = self.depthwise.bias if self.depthwise.bias is not None else self._empty
        pw_b = self.pointwise.bias if self.pointwise.bias is not None else self._empty
        return _ext.depthwise_pointwise_forward(
            x, self.depthwise.weight, self.pointwise.weight, dw_b, pw_b, self.stride, self.padding
        )