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, int64_t kernel_size);

torch::Tensor cross_scan_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t kernel_size) {
    return cross_scan_forward_cuda(x, weight, bias, kernel_size);
}
"""

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

__global__ void cross_scan_32_k3_shared_opt_kernel(const float* __restrict__ x,
                                                   const float* __restrict__ w,
                                                   const float* __restrict__ b,
                                                   float* __restrict__ out,
                                                   int C) {
    __shared__ float s[1024];
    __shared__ float sw[4];
    int tid = threadIdx.x;
    int tile = blockIdx.x;
    int c = tile % C;
    int base = tile << 10;
    s[tid] = x[base + tid];
    if (tid == 0) {
        sw[0] = w[c * 3 + 0];
        sw[1] = w[c * 3 + 1];
        sw[2] = w[c * 3 + 2];
        sw[3] = b[c];
    }
    __syncthreads();

    int ww = tid & 31;
    int hh = tid >> 5;
    float w0 = sw[0], w1 = sw[1], w2 = sw[2];
    float wh = w0 + w2;
    float v = 4.0f * sw[3] + (1.0f + 4.0f * w1) * s[tid];

    // Horizontal scans. Interior pixels receive the same two neighbors from LR/RL.
    if (ww > 0 && ww < 31) {
        v += wh * (s[tid - 1] + s[tid + 1]);
    } else {
        if (tid > 0) v += w0 * s[tid - 1];
        if (tid < 1023) v += w2 * s[tid + 1];
        int pf = (hh << 5) + (31 - ww);
        if (pf > 0) {
            int q = pf - 1;
            int hq = q >> 5;
            int wf = q & 31;
            v += w0 * s[(hq << 5) + (31 - wf)];
        }
        if (pf < 1023) {
            int q = pf + 1;
            int hq = q >> 5;
            int wf = q & 31;
            v += w2 * s[(hq << 5) + (31 - wf)];
        }
    }

    // Vertical scans. Interior pixels receive the same two neighbors from TB/BT.
    if (hh > 0 && hh < 31) {
        v += wh * (s[tid - 32] + s[tid + 32]);
    } else {
        int pt = (ww << 5) + hh;
        if (pt > 0) {
            int q = pt - 1;
            int wq = q >> 5;
            int hq = q & 31;
            v += w0 * s[(hq << 5) + wq];
        }
        if (pt < 1023) {
            int q = pt + 1;
            int wq = q >> 5;
            int hq = q & 31;
            v += w2 * s[(hq << 5) + wq];
        }
        int pb = (ww << 5) + (31 - hh);
        if (pb > 0) {
            int q = pb - 1;
            int wq = q >> 5;
            int hf = q & 31;
            v += w0 * s[((31 - hf) << 5) + wq];
        }
        if (pb < 1023) {
            int q = pb + 1;
            int wq = q >> 5;
            int hf = q & 31;
            v += w2 * s[((31 - hf) << 5) + wq];
        }
    }
    out[base + tid] = v;
}

__global__ void cross_scan_k3_kernel(const float* __restrict__ x,
                                     const float* __restrict__ w,
                                     const float* __restrict__ b,
                                     float* __restrict__ out,
                                     int B, int C, int H, int W) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int N = B * C * H * W;
    if (idx >= N) return;
    int tmp = idx;
    int ww = tmp % W; tmp /= W;
    int hh = tmp % H; tmp /= H;
    int cc = tmp % C;
    int bc = tmp / C;
    int L = H * W;
    int base_bc = (bc * C + cc) * L;
    float v = x[idx] + 4.0f * b[cc];
    const float w0 = w[cc * 3 + 0], w1 = w[cc * 3 + 1], w2 = w[cc * 3 + 2];
    int p = hh * W + ww;
    if (p > 0) v += w0 * x[base_bc + p - 1];
    v += w1 * x[base_bc + p];
    if (p + 1 < L) v += w2 * x[base_bc + p + 1];
    int pf = hh * W + (W - 1 - ww);
    if (pf > 0) { int q=pf-1; int hq=q/W; int wf=q-hq*W; v += w0*x[base_bc+hq*W+(W-1-wf)]; }
    v += w1 * x[idx];
    if (pf + 1 < L) { int q=pf+1; int hq=q/W; int wf=q-hq*W; v += w2*x[base_bc+hq*W+(W-1-wf)]; }
    int pt = ww * H + hh;
    if (pt > 0) { int q=pt-1; int wq=q/H; int hq=q-wq*H; v += w0*x[base_bc+hq*W+wq]; }
    v += w1 * x[idx];
    if (pt + 1 < L) { int q=pt+1; int wq=q/H; int hq=q-wq*H; v += w2*x[base_bc+hq*W+wq]; }
    int pb = ww * H + (H - 1 - hh);
    if (pb > 0) { int q=pb-1; int wq=q/H; int hf=q-wq*H; v += w0*x[base_bc+(H-1-hf)*W+wq]; }
    v += w1 * x[idx];
    if (pb + 1 < L) { int q=pb+1; int wq=q/H; int hf=q-wq*H; v += w2*x[base_bc+(H-1-hf)*W+wq]; }
    out[idx] = v;
}

__global__ void cross_scan_generic_kernel(const float* __restrict__ x,
                                          const float* __restrict__ w,
                                          const float* __restrict__ b,
                                          float* __restrict__ out,
                                          int B, int C, int H, int W, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int N = B * C * H * W;
    if (idx >= N) return;
    int tmp = idx;
    int ww = tmp % W; tmp /= W;
    int hh = tmp % H; tmp /= H;
    int cc = tmp % C;
    int bc = tmp / C;
    int L = H * W;
    int base_bc = (bc * C + cc) * L;
    int pad = K >> 1;
    float v = x[idx] + 4.0f * b[cc];
    int p_lr = hh * W + ww;
    int p_rl = hh * W + (W - 1 - ww);
    int p_tb = ww * H + hh;
    int p_bt = ww * H + (H - 1 - hh);
    const float* wc = w + cc * K;
    for (int j = 0; j < K; ++j) {
        int off = j - pad;
        float wj = wc[j];
        int q = p_lr + off;
        if ((unsigned)q < (unsigned)L) v += wj * x[base_bc + q];
        q = p_rl + off;
        if ((unsigned)q < (unsigned)L) {
            int hq = q / W;
            int wf = q - hq * W;
            v += wj * x[base_bc + hq * W + (W - 1 - wf)];
        }
        q = p_tb + off;
        if ((unsigned)q < (unsigned)L) {
            int wq = q / H;
            int hq = q - wq * H;
            v += wj * x[base_bc + hq * W + wq];
        }
        q = p_bt + off;
        if ((unsigned)q < (unsigned)L) {
            int wq = q / H;
            int hf = q - wq * H;
            v += wj * x[base_bc + (H - 1 - hf) * W + wq];
        }
    }
    out[idx] = v;
}

torch::Tensor cross_scan_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t kernel_size) {
    auto out = torch::empty_like(x);
    int B = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    if ((int)kernel_size == 3 && H == 32 && W == 32) {
        cross_scan_32_k3_shared_opt_kernel<<<B * C, 1024>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), C);
    } else {
        int N = B * C * H * W, threads = 256, blocks = (N + threads - 1) / threads;
        if ((int)kernel_size == 3)
            cross_scan_k3_kernel<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), B, C, H, W);
        else
            cross_scan_generic_kernel<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), B, C, H, W, (int)kernel_size);
    }
    return out;
}
"""

_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
        self.kernel_size = kernel_size
        self.dw_conv = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2, 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, self.kernel_size)