import torch
import torch.nn as nn
import torch.nn.functional as F

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).


class Model(nn.Module):
    """
    **Fused depthwise–pointwise 1D temporal convolution** on ``(B, C, T)`` sequences.

    Fuses a **temporal depthwise** ``k``-tap conv (per-channel along time) with
    **SiLU** and a **pointwise** ``1×1`` mix—standard in temporal encoders for
    proprio, audio, or token-wise time series in robotics / VLA backbones.

    Padding uses dilated effective radius so length is unchanged for odd
    ``kernel_size``.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        dilation: int = 1,
    ):
        super().__init__()
        assert kernel_size % 2 == 1, "odd kernel_size keeps T unchanged with symmetric padding"
        assert dilation >= 1
        self.padding = ((kernel_size - 1) * dilation) // 2
        self.dw = nn.Conv1d(
            in_channels,
            in_channels,
            kernel_size,
            padding=self.padding,
            dilation=dilation,
            groups=in_channels,
            bias=False,
        )
        self.pw = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, C_in, T) temporal features.
        Returns:
            (B, C_out, T)
        """
        x = self.dw(x)
        x = F.silu(x)
        return self.pw(x)


class ModelNew(nn.Module):
    """KernelBench candidate entry point; default delegates to ``Model``."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        dilation: int = 1,
    ):
        super().__init__()
        self._impl = Model(in_channels, out_channels, kernel_size, dilation=dilation)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._impl(x)


batch_size = 16
in_channels = 128
out_channels = 128
kernel_size = 5
dilation = 2
seq_len = 2048


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), in_channels, p.jitter_int(seq_len))
    return [x]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, dilation]

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.cpp_extension as cpp_extension

class ModelNew(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        dilation: int = 1,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.dilation = dilation
        self.padding = ((kernel_size - 1) * dilation) // 2
        
        # Create internal Model to get exact weight initialization
        internal_model = nn.Module()
        internal_model.dw = nn.Conv1d(
            in_channels,
            in_channels,
            kernel_size,
            padding=self.padding,
            dilation=dilation,
            groups=in_channels,
            bias=False,
        )
        internal_model.pw = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)
        
        # Extract weights from internal model
        self.dw_weight = nn.Parameter(internal_model.dw.weight.data.clone())
        self.pw_weight = nn.Parameter(internal_model.pw.weight.data.clone())
        
        # Compile the CUDA kernel
        self._kernel = cpp_extension.load_inline(
            name='fused_depthwise_pointwise_final',
            cpp_sources=self._get_cpp_source(),
            cuda_sources=self._get_cuda_source(),
            functions=['fused_dwpw_cuda'],
            extra_cuda_cflags=['-O3'],
            verbose=False
        )
    
    def _get_cpp_source(self):
        return '''
torch::Tensor fused_dwpw_cuda(
    torch::Tensor input,
    torch::Tensor dw_weight,
    torch::Tensor pw_weight,
    int padding,
    int dilation
);
'''
    
    def _get_cuda_source(self):
        return '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void fused_dwpw_kernel(
    const float* __restrict__ input,
    const float* __restrict__ dw_weight,
    const float* __restrict__ pw_weight,
    float* __restrict__ output,
    int batch_size,
    int in_channels,
    int out_channels,
    int seq_len,
    int kernel_size,
    int dilation,
    int padding
) {
    int b = blockIdx.z;
    int oc = blockIdx.y;
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (t >= seq_len) return;
    
    float sum = 0.0f;
    
    for (int ic = 0; ic < in_channels; ++ic) {
        float dw_sum = 0.0f;
        
        for (int k = 0; k < kernel_size; ++k) {
            int t_input = t - padding + k * dilation;
            if (t_input >= 0 && t_input < seq_len) {
                float val = input[(b * in_channels + ic) * seq_len + t_input];
                float w = dw_weight[ic * kernel_size + k];
                dw_sum += val * w;
            }
        }
        
        // SiLU: x * sigmoid(x)
        float sig = 1.0f / (1.0f + expf(-dw_sum));
        float activated = dw_sum * sig;
        float pw_w = pw_weight[oc * in_channels + ic];
        sum += activated * pw_w;
    }
    
    output[(b * out_channels + oc) * seq_len + t] = sum;
}

torch::Tensor fused_dwpw_cuda(
    torch::Tensor input,
    torch::Tensor dw_weight,
    torch::Tensor pw_weight,
    int padding,
    int dilation
) {
    int batch_size = input.size(0);
    int in_channels = input.size(1);
    int seq_len = input.size(2);
    int kernel_size = dw_weight.size(2);
    int out_channels = pw_weight.size(0);
    
    auto output = torch::empty({batch_size, out_channels, seq_len}, input.options());
    
    int threads = 256;
    int blocks = (seq_len + threads - 1) / threads;
    
    dim3 grid(blocks, out_channels, batch_size);
    dim3 block(threads);
    
    fused_dwpw_kernel<<<grid, block>>>(
        input.data_ptr<float>(),
        dw_weight.data_ptr<float>(),
        pw_weight.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        in_channels,
        out_channels,
        seq_len,
        kernel_size,
        dilation,
        padding
    );
    
    return output;
}
'''
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._kernel.fused_dwpw_cuda(
            x, self.dw_weight.contiguous(), self.pw_weight.contiguous(), 
            self.padding, self.dilation
        )