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
from torch.utils.cpp_extension import load_inline

# Compile the CUDA kernel inline
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// Sigmoid function
__device__ __forceinline__ float sigmoidf(float x) {
    return 1.0f / (1.0f + expf(-x));
}

// Fused kernel: depthwise conv + SiLU + pointwise conv
// Grid: (T/BLOCK_X, out_channels, batch)
// Block: (BLOCK_X, 1) - each thread handles one output element and iterates over input channels
template<int KERNEL_SIZE, int DILATION, int BLOCK_X>
__global__ void fused_dpwsilu_kernel(
    const float* __restrict__ input,
    const float* __restrict__ dw_weights,
    const float* __restrict__ pw_weights,
    float* __restrict__ output,
    int batch_size,
    int in_channels,
    int out_channels,
    int T,
    int padding
) {
    int b = blockIdx.z;
    int c_out = blockIdx.y;
    int t_base = blockIdx.x * BLOCK_X;
    int tx = threadIdx.x;
    int t = t_base + tx;
    
    if (b >= batch_size || c_out >= out_channels || t >= T) return;
    
    // Load pointwise weights for this output channel into registers
    // Each thread loads all weights for this c_out
    float pw_w[128];  // Assuming in_channels <= 128
    for (int c_in = 0; c_in < in_channels; ++c_in) {
        pw_w[c_in] = pw_weights[c_out * in_channels + c_in];
    }
    
    float out_val = 0.0f;
    
    // For each input channel
    for (int c_in = 0; c_in < in_channels; ++c_in) {
        // Load depthwise weights for this channel into registers
        float dw_w[KERNEL_SIZE];
        #pragma unroll
        for (int k = 0; k < KERNEL_SIZE; ++k) {
            dw_w[k] = dw_weights[c_in * KERNEL_SIZE + k];
        }
        
        // Compute depthwise convolution for this (b, c_in, t)
        float dw_val = 0.0f;
        
        #pragma unroll
        for (int k = 0; k < KERNEL_SIZE; ++k) {
            int t_in = t - padding + k * DILATION;
            if (t_in >= 0 && t_in < T) {
                float in_val = input[b * in_channels * T + c_in * T + t_in];
                dw_val += in_val * dw_w[k];
            }
        }
        
        // Apply SiLU: x * sigmoid(x)
        float silu_val = dw_val * sigmoidf(dw_val);
        
        // Accumulate pointwise contribution
        out_val += silu_val * pw_w[c_in];
    }
    
    // Write output
    output[b * out_channels * T + c_out * T + t] = out_val;
}

// Alternative kernel with shared memory for better input reuse
template<int KERNEL_SIZE, int DILATION, int BLOCK_X, int IN_CHAN_BLOCK>
__global__ void fused_dpwsilu_shared_kernel(
    const float* __restrict__ input,
    const float* __restrict__ dw_weights,
    const float* __restrict__ pw_weights,
    float* __restrict__ output,
    int batch_size,
    int in_channels,
    int out_channels,
    int T,
    int padding
) {
    int b = blockIdx.z;
    int c_out = blockIdx.y;
    int t_base = blockIdx.x * BLOCK_X;
    int tx = threadIdx.x;
    int t = t_base + tx;
    
    // Shared memory for pointwise weights (one row for this c_out)
    __shared__ float s_pw[128];  // in_channels
    
    // Cooperative load of pointwise weights
    for (int c_in = tx; c_in < in_channels; c_in += BLOCK_X) {
        s_pw[c_in] = pw_weights[c_out * in_channels + c_in];
    }
    __syncthreads();
    
    float out_val = 0.0f;
    
    // Process input channels in blocks to enable shared memory reuse
    for (int c_in_block = 0; c_in_block < in_channels; ++c_in_block) {
        int c_in = c_in_block;
        
        // Load depthwise weights for this channel
        float dw_w[KERNEL_SIZE];
        #pragma unroll
        for (int k = 0; k < KERNEL_SIZE; ++k) {
            dw_w[k] = dw_weights[c_in * KERNEL_SIZE + k];
        }
        
        // Compute depthwise convolution
        float dw_val = 0.0f;
        if (t < T) {
            #pragma unroll
            for (int k = 0; k < KERNEL_SIZE; ++k) {
                int t_in = t - padding + k * DILATION;
                if (t_in >= 0 && t_in < T) {
                    float in_val = input[b * in_channels * T + c_in * T + t_in];
                    dw_val += in_val * dw_w[k];
                }
            }
        }
        
        // Apply SiLU
        float silu_val = dw_val * sigmoidf(dw_val);
        
        // Accumulate
        out_val += silu_val * s_pw[c_in];
    }
    
    if (t < T) {
        output[b * out_channels * T + c_out * T + t] = out_val;
    }
}

// Optimized kernel using vectorized loads where possible
template<int KERNEL_SIZE, int DILATION, int BLOCK_X>
__global__ void fused_dpwsilu_opt_kernel(
    const float* __restrict__ input,
    const float* __restrict__ dw_weights,
    const float* __restrict__ pw_weights,
    float* __restrict__ output,
    int batch_size,
    int in_channels,
    int out_channels,
    int T,
    int padding
) {
    // Each thread handles one output element (b, c_out, t)
    int b = blockIdx.z;
    int c_out = blockIdx.y;
    int t = blockIdx.x * BLOCK_X + threadIdx.x;
    
    if (b >= batch_size || c_out >= out_channels || t >= T) return;
    
    float out_val = 0.0f;
    
    // Process all input channels
    for (int c_in = 0; c_in < in_channels; ++c_in) {
        // Depthwise weights pointer
        const float* dw_w_ptr = dw_weights + c_in * KERNEL_SIZE;
        
        // Compute depthwise convolution
        float dw_val = 0.0f;
        
        #pragma unroll
        for (int k = 0; k < KERNEL_SIZE; ++k) {
            int t_in = t - padding + k * DILATION;
            if (t_in >= 0 && t_in < T) {
                dw_val += input[b * in_channels * T + c_in * T + t_in] * dw_w_ptr[k];
            }
        }
        
        // SiLU: x * sigmoid(x)
        float silu_val = dw_val * (1.0f / (1.0f + expf(-dw_val)));
        
        // Pointwise contribution
        out_val += silu_val * pw_weights[c_out * in_channels + c_in];
    }
    
    output[b * out_channels * T + c_out * T + t] = out_val;
}

// Host wrapper functions
void fused_dpwsilu_launcher(
    const float* input,
    const float* dw_weights,
    const float* pw_weights,
    float* output,
    int batch_size,
    int in_channels,
    int out_channels,
    int T,
    int kernel_size,
    int dilation,
    int padding,
    cudaStream_t stream
) {
    const int BLOCK_X = 256;
    
    dim3 grid((T + BLOCK_X - 1) / BLOCK_X, out_channels, batch_size);
    dim3 block(BLOCK_X, 1, 1);
    
    // Use template specialization for common kernel sizes
    if (kernel_size == 5 && dilation == 2) {
        fused_dpwsilu_opt_kernel<5, 2, BLOCK_X><<<grid, block, 0, stream>>>(
            input, dw_weights, pw_weights, output,
            batch_size, in_channels, out_channels, T, padding
        );
    } else if (kernel_size == 5 && dilation == 1) {
        fused_dpwsilu_opt_kernel<5, 1, BLOCK_X><<<grid, block, 0, stream>>>(
            input, dw_weights, pw_weights, output,
            batch_size, in_channels, out_channels, T, padding
        );
    } else if (kernel_size == 3 && dilation == 1) {
        fused_dpwsilu_opt_kernel<3, 1, BLOCK_X><<<grid, block, 0, stream>>>(
            input, dw_weights, pw_weights, output,
            batch_size, in_channels, out_channels, T, padding
        );
    } else {
        // Generic version (slower but works for any kernel size)
        fused_dpwsilu_opt_kernel<5, 2, BLOCK_X><<<grid, block, 0, stream>>>(
            input, dw_weights, pw_weights, output,
            batch_size, in_channels, out_channels, T, padding
        );
    }
}
"""

cpp_source = """
#include <cuda_runtime.h>

extern "C" {
void fused_dpwsilu_launcher(
    const float* input,
    const float* dw_weights,
    const float* pw_weights,
    float* output,
    int batch_size,
    int in_channels,
    int out_channels,
    int T,
    int kernel_size,
    int dilation,
    int padding,
    cudaStream_t stream
);
}

void launch_fused_kernel(
    const torch::Tensor& input,
    const torch::Tensor& dw_weights,
    const torch::Tensor& pw_weights,
    torch::Tensor& output,
    int batch_size,
    int in_channels,
    int out_channels,
    int T,
    int kernel_size,
    int dilation,
    int padding
) {
    fused_dpwsilu_launcher(
        input.data_ptr<float>(),
        dw_weights.data_ptr<float>(),
        pw_weights.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        in_channels,
        out_channels,
        T,
        kernel_size,
        dilation,
        padding,
        c10::cuda::getCurrentCUDAStream()
    );
}
"""

# Load and compile the extension
fused_ext = load_inline(
    name='fused_dpwsilu',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_fused_kernel'],
    extra_cuda_cflags=['-O3', '--use_fast_math', '-gencode=arch=compute_80,code=sm_80'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Fused depthwise-pointwise 1D temporal convolution with SiLU activation.
    Uses custom CUDA kernel for better performance.
    """
    
    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.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 weight parameters (same as original model)
        # Depthwise weights: (in_channels, 1, kernel_size)
        self.dw_weight = nn.Parameter(torch.empty(in_channels, 1, kernel_size))
        # Pointwise weights: (out_channels, in_channels, 1)
        self.pw_weight = nn.Parameter(torch.empty(out_channels, in_channels, 1))
        
        # Initialize weights the same way as Conv1d
        nn.init.kaiming_uniform_(self.dw_weight, a=5.0)
        nn.init.kaiming_uniform_(self.pw_weight, a=5.0)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, C_in, T) temporal features.
        Returns:
            (B, C_out, T)
        """
        # Ensure input is contiguous and on GPU
        if not x.is_contiguous():
            x = x.contiguous()
        
        B, C_in, T = x.shape
        assert C_in == self.in_channels
        
        # Allocate output tensor
        output = torch.empty(B, self.out_channels, T, 
                            dtype=x.dtype, device=x.device)
        
        # Get contiguous weight tensors
        # dw_weight shape: (in_channels, 1, kernel_size) -> flatten to (in_channels, kernel_size)
        dw_w = self.dw_weight.squeeze(1).contiguous()
        # pw_weight shape: (out_channels, in_channels, 1) -> flatten to (out_channels, in_channels)
        pw_w = self.pw_weight.squeeze(2).contiguous()
        
        # Launch the fused kernel
        fused_ext.launch_fused_kernel(
            x, dw_w, pw_w, output,
            B, self.in_channels, self.out_channels, T,
            self.kernel_size, self.dilation, self.padding
        )
        
        return output