import torch
import torch.nn as nn

# 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 **2D bidirectional cross-scan** block for robotics **SSM-hybrid** backbones
    (e.g. Vision Mamba / VMamba-style encoders on overhead maps, costmaps, or
    multi-sensor tiles).

    One forward fuses four fixed scan permutations of ``(H, W)`` into 1D sequences,
    applies a **shared depthwise Conv1d** (lightweight stand-in for a selective SSM
    along the scan), maps each result back to the grid, **sums** the four
    directional contributions, and adds a **residual**.

    Scan order (same spirit as 2D cross-scan / S2D):
      LR — row-major left→right, rows top→bottom  
      RL — row-major right→left per row  
      TB — column-wise top→bottom, columns left→right  
      BT — flip vertical then same as TB  
    """

    def __init__(self, channels: int, kernel_size: int = 3):
        super().__init__()
        assert channels >= 1
        assert kernel_size % 2 == 1, "use odd kernel_size for same-length conv1d"
        self.channels = channels
        pad = kernel_size // 2
        self.dw_conv = nn.Conv1d(
            channels,
            channels,
            kernel_size,
            padding=pad,
            groups=channels,
            bias=True,
        )

    def _apply_scan_conv_unscan(self, x: torch.Tensor) -> torch.Tensor:
        """Shared depthwise conv along flattened 2D, then map back to (B,C,H,W)."""
        B, C, H, W = x.shape
        seq = x.flatten(2)
        y = self.dw_conv(seq)
        return y.view(B, C, H, W)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, C, H, W) spatial features (map / image / fused tile).
        Returns:
            (B, C, H, W) after fused four-way cross-scan + residual.
        """
        B, C, H, W = x.shape
        assert C == self.channels

        acc = torch.zeros_like(x)

        # LR: row-major
        acc = acc + self._apply_scan_conv_unscan(x)

        # RL: flip W, scan, unflip
        xr = x.flip(dims=(-1,))
        acc = acc + self._apply_scan_conv_unscan(xr).flip(dims=(-1,))

        # TB: columns as sequences (permute to W,H, flatten)
        xt = x.permute(0, 1, 3, 2).contiguous()
        Bt, Ct, Wt, Ht = xt.shape
        seq_t = xt.flatten(2)
        yt = self.dw_conv(seq_t).view(Bt, Ct, Wt, Ht).permute(0, 1, 3, 2).contiguous()
        acc = acc + yt

        # BT: flip H then same TB path
        xb = x.flip(dims=(-2,))
        xbt = xb.permute(0, 1, 3, 2).contiguous()
        Bb, Cb, Wb, Hb = xbt.shape
        seq_b = xbt.flatten(2)
        yb = self.dw_conv(seq_b).view(Bb, Cb, Wb, Hb).permute(0, 1, 3, 2).contiguous().flip(
            dims=(-2,)
        )
        acc = acc + yb

        return x + acc


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

    def __init__(self, channels: int, kernel_size: int = 3):
        super().__init__()
        self._impl = Model(channels, kernel_size=kernel_size)

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


batch_size = 4
channels = 64
height = 32
width = 32
kernel_size = 3


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), channels, p.jitter_int(height), p.jitter_int(width))
    return [x]


def get_init_inputs():
    return [channels, kernel_size]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

template <typename scalar_t>
__global__ void cross_scan_conv_kernel(
    const scalar_t* __restrict__ x,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ out,
    const int B, const int C, const int H, const int W, const int K,
    const int pad) {
    
    const int L = H * W;
    
    // Use 1D thread indexing where consecutive threads access consecutive positions
    // in row-major order (better coalescing)
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = B * C * L;
    
    if (idx >= total) return;
    
    // Compute position in row-major layout
    const int spatial_idx = idx % L;
    const int c_idx = (idx / L) % C;
    const int b_idx = idx / (C * L);
    
    const int h_idx = spatial_idx / W;
    const int w_idx = spatial_idx % W;
    
    // Load weights into registers for this channel
    scalar_t w_regs[16];  // Assuming K <= 16
    for (int k = 0; k < K; k++) {
        w_regs[k] = weight[c_idx * K + k];
    }
    scalar_t b_val = bias[c_idx];
    
    // Base pointer for this (batch, channel)
    const scalar_t* x_base = x + b_idx * C * L + c_idx * L;
    scalar_t* out_ptr = out + idx;
    
    scalar_t acc = scalar_t(0);
    
    // ===== LR scan: row-major left-to-right =====
    // Direct row-major access - good coalescing
    #pragma unroll
    for (int k = 0; k < K; k++) {
        const int seq_off = spatial_idx - pad + k;
        if (seq_off >= 0 && seq_off < L) {
            acc += w_regs[k] * x_base[seq_off];
        }
    }
    
    // ===== RL scan: flip W =====
    // For position (h,w), scan pos = h*W + (W-1-w)
    // Then unflip output
    {
        const int seq_pos = h_idx * W + (W - 1 - w_idx);
        #pragma unroll
        for (int k = 0; k < K; k++) {
            const int seq_off = seq_pos - pad + k;
            if (seq_off >= 0 && seq_off < L) {
                const int h_off = seq_off / W;
                const int w_off = seq_off % W;
                const int x_idx = h_off * W + (W - 1 - w_off);
                acc += w_regs[k] * x_base[x_idx];
            }
        }
    }
    
    // ===== TB scan: columns as sequences =====
    {
        const int seq_pos = w_idx * H + h_idx;
        #pragma unroll
        for (int k = 0; k < K; k++) {
            const int seq_off = seq_pos - pad + k;
            if (seq_off >= 0 && seq_off < L) {
                const int w_off = seq_off / H;
                const int h_off = seq_off % H;
                const int x_idx = h_off * W + w_off;
                acc += w_regs[k] * x_base[x_idx];
            }
        }
    }
    
    // ===== BT scan: flip H, columns as sequences =====
    {
        const int seq_pos = w_idx * H + (H - 1 - h_idx);
        #pragma unroll
        for (int k = 0; k < K; k++) {
            const int seq_off = seq_pos - pad + k;
            if (seq_off >= 0 && seq_off < L) {
                const int w_off = seq_off / H;
                const int h_off_flipped = seq_off % H;
                const int h_off_original = H - 1 - h_off_flipped;
                const int x_idx = h_off_original * W + w_off;
                acc += w_regs[k] * x_base[x_idx];
            }
        }
    }
    
    // Bias is added 4 times (once per scan)
    out_ptr[0] = acc + (scalar_t)4.0 * b_val;
}

torch::Tensor cross_scan_conv_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t kernel_size) {
    
    const int B = x.size(0);
    const int C = x.size(1);
    const int H = x.size(2);
    const int W = x.size(3);
    const int K = kernel_size;
    
    auto out = torch::empty_like(x);
    
    const int total = B * C * H * W;
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "cross_scan_conv_cuda", ([&] {
        cross_scan_conv_kernel<scalar_t><<<blocks, threads>>>(
            x.data_ptr<scalar_t>(),
            weight.data_ptr<scalar_t>(),
            bias.data_ptr<scalar_t>(),
            out.data_ptr<scalar_t>(),
            B, C, H, W, K, K/2);
    }));
    
    return out;
}
"""

cpp_source = r"""
torch::Tensor cross_scan_conv_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t kernel_size);
"""

module = load_inline(
    name='cross_scan_conv_inline',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['cross_scan_conv_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self, channels: int, kernel_size: int = 3):
        super().__init__()
        assert channels >= 1
        assert kernel_size % 2 == 1, "use odd kernel_size for same-length conv1d"
        self.channels = channels
        self.kernel_size = kernel_size
        pad = kernel_size // 2
        self.dw_conv = nn.Conv1d(
            channels,
            channels,
            kernel_size,
            padding=pad,
            groups=channels,
            bias=True,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        assert C == self.channels
        
        # Extract weight and bias from Conv1d
        weight = self.dw_conv.weight.data.squeeze(1)  # (C, 1, K) -> (C, K)
        bias = self.dw_conv.bias.data  # (C,)
        
        x_cuda = x.contiguous()
        weight_cuda = weight.contiguous()
        bias_cuda = bias.contiguous()
        
        out = module.cross_scan_conv_cuda(x_cuda, weight_cuda, bias_cuda, self.kernel_size)
        
        return x + out