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(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAStream.h>

#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK((x).scalar_type() == at::kFloat, #x " must be float32")

__global__ void cross_scan_32_k3_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int BC,
    int C)
{
    const int pix = blockIdx.x * blockDim.x + threadIdx.x;  // 0..1023
    const int bc = blockIdx.y;
    if (pix >= 1024 || bc >= BC) return;

    const int c = bc % C;
    const int base = bc << 10;  // *1024
    const int off = base + pix;

    const float w0 = __ldg(weight + c * 3 + 0);
    const float w1 = __ldg(weight + c * 3 + 1);
    const float w2 = __ldg(weight + c * 3 + 2);
    const float b  = __ldg(bias + c);
    const float xc = __ldg(x + off);

    const int h = pix >> 5;
    const int w = pix & 31;

    float acc = 0.0f;

    // LR: original row-major sequence
    float v = b + w1 * xc;
    if (pix > 0)    v += w0 * __ldg(x + base + pix - 1);
    if (pix < 1023) v += w2 * __ldg(x + base + pix + 1);
    acc += v;

    // RL: row-major sequence after horizontal flip, then unflip
    const int nR = (h << 5) + (31 - w);
    v = b + w1 * xc;
    if (nR > 0) {
        const int s = nR - 1;
        const int orig = (s & ~31) + (31 - (s & 31));
        v += w0 * __ldg(x + base + orig);
    }
    if (nR < 1023) {
        const int s = nR + 1;
        const int orig = (s & ~31) + (31 - (s & 31));
        v += w2 * __ldg(x + base + orig);
    }
    acc += v;

    // TB: column-major sequence (W,H), mapped back
    const int nT = (w << 5) + h;
    v = b + w1 * xc;
    if (nT > 0) {
        const int s = nT - 1;
        const int orig = ((s & 31) << 5) + (s >> 5);
        v += w0 * __ldg(x + base + orig);
    }
    if (nT < 1023) {
        const int s = nT + 1;
        const int orig = ((s & 31) << 5) + (s >> 5);
        v += w2 * __ldg(x + base + orig);
    }
    acc += v;

    // BT: column-major sequence after vertical flip, then unflip
    const int nB = (w << 5) + (31 - h);
    v = b + w1 * xc;
    if (nB > 0) {
        const int s = nB - 1;
        const int orig = ((31 - (s & 31)) << 5) + (s >> 5);
        v += w0 * __ldg(x + base + orig);
    }
    if (nB < 1023) {
        const int s = nB + 1;
        const int orig = ((31 - (s & 31)) << 5) + (s >> 5);
        v += w2 * __ldg(x + base + orig);
    }
    acc += v;

    out[off] = xc + acc;
}

__global__ void cross_scan_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int total,
    int C,
    int H,
    int W,
    int K)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int N = H * W;
    const int pix = idx % N;
    const int bc = idx / N;
    const int c = bc % C;
    const int base = bc * N;
    const int h = pix / W;
    const int w = pix - h * W;
    const int pad = K >> 1;
    const float* __restrict__ wt = weight + c * K;
    const float b = __ldg(bias + c);
    const float xc = __ldg(x + idx);

    float acc = 0.0f;

    // LR
    int n = pix;
    float v = b;
    for (int j = 0; j < K; ++j) {
        int s = n + j - pad;
        if ((unsigned)s < (unsigned)N) v += __ldg(wt + j) * __ldg(x + base + s);
    }
    acc += v;

    // RL
    n = h * W + (W - 1 - w);
    v = b;
    for (int j = 0; j < K; ++j) {
        int s = n + j - pad;
        if ((unsigned)s < (unsigned)N) {
            int hh = s / W;
            int ww = s - hh * W;
            int orig = hh * W + (W - 1 - ww);
            v += __ldg(wt + j) * __ldg(x + base + orig);
        }
    }
    acc += v;

    // TB
    n = w * H + h;
    v = b;
    for (int j = 0; j < K; ++j) {
        int s = n + j - pad;
        if ((unsigned)s < (unsigned)N) {
            int ww = s / H;
            int hh = s - ww * H;
            int orig = hh * W + ww;
            v += __ldg(wt + j) * __ldg(x + base + orig);
        }
    }
    acc += v;

    // BT
    n = w * H + (H - 1 - h);
    v = b;
    for (int j = 0; j < K; ++j) {
        int s = n + j - pad;
        if ((unsigned)s < (unsigned)N) {
            int ww = s / H;
            int hf = s - ww * H;
            int orig = (H - 1 - hf) * W + ww;
            v += __ldg(wt + j) * __ldg(x + base + orig);
        }
    }
    acc += v;

    out[idx] = xc + acc;
}

torch::Tensor cross_scan_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(bias);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight); CHECK_CONTIGUOUS(bias);
    CHECK_FLOAT(x); CHECK_FLOAT(weight); CHECK_FLOAT(bias);

    TORCH_CHECK(x.dim() == 4, "x must be B,C,H,W");
    TORCH_CHECK(weight.dim() == 3 && weight.size(1) == 1, "weight must be C,1,K");
    TORCH_CHECK(bias.dim() == 1, "bias must be C");

    const int B = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int K = (int)weight.size(2);
    TORCH_CHECK(weight.size(0) == C && bias.size(0) == C, "channel mismatch");

    auto out = torch::empty_like(x);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (H == 32 && W == 32 && K == 3) {
        const int BC = B * C;
        dim3 block(256);
        dim3 grid(4, BC);
        cross_scan_32_k3_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
            out.data_ptr<float>(), BC, C);
    } else {
        const int total = B * C * H * W;
        const int block = 256;
        const int grid = (total + block - 1) / block;
        cross_scan_generic_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
            out.data_ptr<float>(), total, C, H, W, K);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(
    name="cross_scan_depthwise_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["cross_scan_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    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, "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 forward(self, x: torch.Tensor) -> torch.Tensor:
        # Benchmark inputs are contiguous NCHW fp32 CUDA tensors.
        return _ext.cross_scan_forward(x, self.dw_conv.weight, self.dw_conv.bias)