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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor cross_scan_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
torch::Tensor cross_scan_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) { return cross_scan_forward_cuda(x, weight, bias); }
"""

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

__global__ void cross_scan_kernel3_32_64(const float* __restrict__ x,
                                   const float* __restrict__ weight,
                                   const float* __restrict__ bias,
                                   float* __restrict__ y,
                                   int total) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;
    int spat = tid & 1023;
    int w = tid & 31;
    int h = (tid >> 5) & 31;
    int c = (tid >> 10) & 63;
    int bc_base = tid - spat;
    const float* xp = x + bc_base;
    const float* wc = weight + c * 3;
    float w0 = wc[0], w1 = wc[1], w2 = wc[2];
    float xv = xp[spat];
    float acc = 4.0f * bias[c] + (4.0f * w1) * xv;
    float ws = w0 + w2;

    if ((unsigned)(h - 1) < 30u && (unsigned)(w - 1) < 30u) {
        acc = fmaf(ws, xp[spat - 1], acc);
        acc = fmaf(ws, xp[spat + 1], acc);
        acc = fmaf(ws, xp[spat - 32], acc);
        acc = fmaf(ws, xp[spat + 32], acc);
    } else {
        if (spat > 0)     acc = fmaf(w0, xp[spat - 1], acc);
        if (spat < 1023)  acc = fmaf(w2, xp[spat + 1], acc);
        if (w < 31)       acc = fmaf(w0, xp[spat + 1], acc);
        else if (h > 0)   acc = fmaf(w0, xp[(h - 1) * 32], acc);
        if (w > 0)        acc = fmaf(w2, xp[spat - 1], acc);
        else if (h < 31)  acc = fmaf(w2, xp[(h + 1) * 32 + 31], acc);
        if (h > 0)        acc = fmaf(w0, xp[spat - 32], acc);
        else if (w > 0)   acc = fmaf(w0, xp[992 + (w - 1)], acc);
        if (h < 31)       acc = fmaf(w2, xp[spat + 32], acc);
        else if (w < 31)  acc = fmaf(w2, xp[w + 1], acc);
        if (h < 31)       acc = fmaf(w0, xp[spat + 32], acc);
        else if (w > 0)   acc = fmaf(w0, xp[w - 1], acc);
        if (h > 0)        acc = fmaf(w2, xp[spat - 32], acc);
        else if (w < 31)  acc = fmaf(w2, xp[992 + (w + 1)], acc);
    }
    y[tid] = xv + acc;
}

__global__ void cross_scan_kernel3(const float* __restrict__ x,
                                   const float* __restrict__ weight,
                                   const float* __restrict__ bias,
                                   float* __restrict__ y,
                                   int C, int H, int W, int total) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;
    int s = tid;
    int w = s % W; s /= W;
    int h = s % H; s /= H;
    int c = s % C;
    int HW = H * W;
    int spat = h * W + w;
    int bc_base = tid - spat;
    const float* xp = x + bc_base;
    const float* wc = weight + c * 3;
    float w0 = wc[0], w1 = wc[1], w2 = wc[2];
    float xv = xp[spat];
    float acc = 4.0f * bias[c] + (4.0f * w1) * xv;
    if (spat > 0)      acc = fmaf(w0, xp[spat - 1], acc);
    if (spat < HW - 1) acc = fmaf(w2, xp[spat + 1], acc);
    if (w < W - 1)        acc = fmaf(w0, xp[spat + 1], acc);
    else if (h > 0)       acc = fmaf(w0, xp[(h - 1) * W], acc);
    if (w > 0)            acc = fmaf(w2, xp[spat - 1], acc);
    else if (h < H - 1)   acc = fmaf(w2, xp[(h + 1) * W + (W - 1)], acc);
    if (h > 0)            acc = fmaf(w0, xp[spat - W], acc);
    else if (w > 0)       acc = fmaf(w0, xp[(H - 1) * W + (w - 1)], acc);
    if (h < H - 1)        acc = fmaf(w2, xp[spat + W], acc);
    else if (w < W - 1)   acc = fmaf(w2, xp[w + 1], acc);
    if (h < H - 1)        acc = fmaf(w0, xp[spat + W], acc);
    else if (w > 0)       acc = fmaf(w0, xp[w - 1], acc);
    if (h > 0)            acc = fmaf(w2, xp[spat - W], acc);
    else if (w < W - 1)   acc = fmaf(w2, xp[(H - 1) * W + (w + 1)], acc);
    y[tid] = xv + acc;
}

__device__ __forceinline__ int spatial_from_rl(int n, int H, int W) { int row = n / W; int col = n - row * W; return row * W + (W - 1 - col); }
__device__ __forceinline__ int spatial_from_tb(int n, int H, int W) { int col = n / H; int row = n - col * H; return row * W + col; }
__device__ __forceinline__ int spatial_from_bt(int n, int H, int W) { int col = n / H; int rowf = n - col * H; return (H - 1 - rowf) * W + col; }

__global__ void cross_scan_kernel_generic(const float* __restrict__ x,
                                  const float* __restrict__ weight,
                                  const float* __restrict__ bias,
                                  float* __restrict__ y,
                                  int C, int H, int W, int K, int total) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;
    int s = tid; int w = s % W; s /= W; int h = s % H; s /= H; int c = s % C;
    int HW = H * W; int spat = h * W + w; int bc_base = tid - spat;
    int pos_lr = spat, pos_rl = h * W + (W - 1 - w), pos_tb = w * H + h, pos_bt = w * H + (H - 1 - h);
    int pad = K >> 1; const float* wx = weight + c * K; float acc = 4.0f * bias[c];
    for (int kk = 0; kk < K; ++kk) {
        int off = kk - pad; float wk = wx[kk]; int n = pos_lr + off;
        if ((unsigned)n < (unsigned)HW) acc = fmaf(wk, x[bc_base + n], acc);
        n = pos_rl + off; if ((unsigned)n < (unsigned)HW) acc = fmaf(wk, x[bc_base + spatial_from_rl(n, H, W)], acc);
        n = pos_tb + off; if ((unsigned)n < (unsigned)HW) acc = fmaf(wk, x[bc_base + spatial_from_tb(n, H, W)], acc);
        n = pos_bt + off; if ((unsigned)n < (unsigned)HW) acc = fmaf(wk, x[bc_base + spatial_from_bt(n, H, W)], acc);
    }
    y[tid] = x[tid] + acc;
}

torch::Tensor cross_scan_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    int B = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3), K = (int)weight.size(2);
    int total = B * C * H * W;
    auto y = torch::empty_like(x);
    int threads = 256, blocks = (total + threads - 1) / threads;
    if (K == 3 && C == 64 && H == 32 && W == 32) {
        cross_scan_kernel3_32_64<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), total);
    } else if (K == 3) {
        cross_scan_kernel3<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), C, H, W, total);
    } else {
        cross_scan_kernel_generic<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), C, H, W, K, total);
    }
    return y;
}
"""

_ext = load_inline(name="cross_scan_ext_final", cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=["cross_scan_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"], extra_cflags=["-O3"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, channels: int, kernel_size: int = 3):
        super().__init__()
        assert channels >= 1
        assert kernel_size % 2 == 1
        self.channels = channels
        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:
        return _ext.cross_scan_forward(x, self.dw_conv.weight, self.dw_conv.bias)