import math

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):
    """
    Graph-Fused VLA (GF-VLA) style **scene graph relation update** in one forward.

    Objects are nodes with vision-language aligned embeddings; directed pairs
    (i → j) carry pairwise conditioning (geometry, language grounding, or fused
    VLA features). This reference fuses:

      LayerNorm(nodes) → multi-head attention over objects with **pairwise
      relation MLP bias** on edges → residual write-back.

    Intended as a fusion target for navigation / manipulation stacks that maintain
    explicit scene graphs and update them from multimodal observations.
    """

    def __init__(self, dim: int, pair_dim: int, num_heads: int):
        super().__init__()
        assert dim % num_heads == 0
        self.dim = dim
        self.pair_dim = pair_dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)

        self.norm = nn.LayerNorm(dim)
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.rel_mlp = nn.Sequential(
            nn.Linear(pair_dim, pair_dim),
            nn.SiLU(),
            nn.Linear(pair_dim, num_heads),
        )
        self.out_proj = nn.Linear(dim, dim, bias=False)

    def forward(
        self,
        nodes: torch.Tensor,
        pair_feats: torch.Tensor,
        edge_mask: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            nodes: (B, N, D) object / region embeddings.
            pair_feats: (B, N, N, P) pairwise VLA or geometric conditioning.
            edge_mask: (B, N, N) float in {0, 1}; 1 keeps attention over column j.
        Returns:
            (B, N, D) updated node embeddings (residual around attention block).
        """
        B, N, D = nodes.shape
        _, N2, N3, P = pair_feats.shape
        assert N == N2 == N3 and P == self.pair_dim
        h = self.norm(nodes)

        q = self.q_proj(h).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(h).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(h).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)

        logits = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        rel_bias = self.rel_mlp(pair_feats).permute(0, 3, 1, 2)
        logits = logits + rel_bias

        logits = logits.masked_fill(edge_mask.unsqueeze(1) < 0.5, float("-inf"))
        attn = F.softmax(logits, dim=-1)
        ctx = torch.matmul(attn, v).transpose(1, 2).reshape(B, N, D)
        return nodes + self.out_proj(ctx)


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

    def __init__(self, dim: int, pair_dim: int, num_heads: int):
        super().__init__()
        self._impl = Model(dim, pair_dim, num_heads)

    def forward(
        self,
        nodes: torch.Tensor,
        pair_feats: torch.Tensor,
        edge_mask: torch.Tensor,
    ) -> torch.Tensor:
        return self._impl(nodes, pair_feats, edge_mask)


batch_size = 4
num_objects = 24
dim = 96
pair_dim = 16
num_heads = 4


def get_inputs():
    p = popcorn_pri
    nodes = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_objects), dim)
    pair_feats = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_objects), p.jitter_int(num_objects), pair_dim)
    edge_mask = torch.ones(p.jitter_int(batch_size), p.jitter_int(num_objects), p.jitter_int(num_objects))
    return [nodes, pair_feats, edge_mask]


def get_init_inputs():
    return [dim, pair_dim, num_heads]

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

_cpp_src = r"""
#include <torch/extension.h>

at::Tensor gf_forward(
    at::Tensor nodes,
    at::Tensor pair_feats,
    at::Tensor edge_mask,
    at::Tensor norm_w,
    at::Tensor norm_b,
    at::Tensor q_w,
    at::Tensor k_w,
    at::Tensor v_w,
    at::Tensor rel1_w,
    at::Tensor rel1_b,
    at::Tensor rel2_w,
    at::Tensor rel2_b,
    at::Tensor out_w,
    double eps,
    double scale);
"""

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

#define D 96
#define N 24
#define H 4
#define HD 24
#define P 16

__global__ void ln_qkv_kernel(
    const float* __restrict__ nodes,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ qw,
    const float* __restrict__ kw,
    const float* __restrict__ vw,
    float* __restrict__ q,
    float* __restrict__ k,
    float* __restrict__ v,
    int B,
    float eps) {

    int idx = blockIdx.x; // B*N
    int b = idx / N;
    int n = idx - b * N;
    int tid = threadIdx.x;
    __shared__ float hbuf[D];
    __shared__ float mean_s;
    __shared__ float invstd_s;

    const float* x = nodes + (b * N + n) * D;
    if (tid == 0) {
        float m = 0.0f;
        #pragma unroll
        for (int i = 0; i < D; ++i) m += x[i];
        m *= (1.0f / (float)D);
        float var = 0.0f;
        #pragma unroll
        for (int i = 0; i < D; ++i) {
            float d = x[i] - m;
            var += d * d;
        }
        var *= (1.0f / (float)D);
        mean_s = m;
        invstd_s = rsqrtf(var + eps);
    }
    __syncthreads();

    if (tid < D) {
        hbuf[tid] = (x[tid] - mean_s) * invstd_s * gamma[tid] + beta[tid];
    }
    __syncthreads();

    if (tid < D) {
        float sq = 0.0f, sk = 0.0f, sv = 0.0f;
        const float* wq = qw + tid * D;
        const float* wk = kw + tid * D;
        const float* wv = vw + tid * D;
        #pragma unroll
        for (int e = 0; e < D; ++e) {
            float hv = hbuf[e];
            sq += hv * wq[e];
            sk += hv * wk[e];
            sv += hv * wv[e];
        }
        int hh = tid / HD;
        int dd = tid - hh * HD;
        int o = ((b * H + hh) * N + n) * HD + dd;
        q[o] = sq;
        k[o] = sk;
        v[o] = sv;
    }
}

__global__ void attn_kernel(
    const float* __restrict__ q,
    const float* __restrict__ k,
    const float* __restrict__ v,
    const float* __restrict__ pair_feats,
    const float* __restrict__ edge_mask,
    const float* __restrict__ r1w,
    const float* __restrict__ r1b,
    const float* __restrict__ r2w,
    const float* __restrict__ r2b,
    float* __restrict__ ctx,
    int B,
    float scale) {

    int b = blockIdx.x;
    int h = blockIdx.y;
    int i = blockIdx.z;
    int tid = threadIdx.x;

    __shared__ float logits[N];
    __shared__ float denom_s;

    if (tid < N) {
        int j = tid;
        float logit;
        float mask = edge_mask[(b * N + i) * N + j];
        if (mask < 0.5f) {
            logit = -INFINITY;
        } else {
            const float* qv = q + ((b * H + h) * N + i) * HD;
            const float* kv = k + ((b * H + h) * N + j) * HD;
            float dot = 0.0f;
            #pragma unroll
            for (int d = 0; d < HD; ++d) dot += qv[d] * kv[d];

            const float* pf = pair_feats + ((b * N + i) * N + j) * P;
            float hidden[P];
            #pragma unroll
            for (int r = 0; r < P; ++r) {
                float z = r1b[r];
                const float* wr = r1w + r * P;
                #pragma unroll
                for (int p = 0; p < P; ++p) z += pf[p] * wr[p];
                hidden[r] = z / (1.0f + expf(-z));
            }
            float rb = r2b[h];
            const float* w2 = r2w + h * P;
            #pragma unroll
            for (int r = 0; r < P; ++r) rb += hidden[r] * w2[r];
            logit = dot * scale + rb;
        }
        logits[j] = logit;
    }
    __syncthreads();

    if (tid == 0) {
        float m = logits[0];
        #pragma unroll
        for (int j = 1; j < N; ++j) m = fmaxf(m, logits[j]);
        float s = 0.0f;
        #pragma unroll
        for (int j = 0; j < N; ++j) {
            float a = expf(logits[j] - m);
            logits[j] = a;
            s += a;
        }
        denom_s = 1.0f / s;
    }
    __syncthreads();

    if (tid < HD) {
        float val = 0.0f;
        #pragma unroll
        for (int j = 0; j < N; ++j) {
            float a = logits[j] * denom_s;
            val += a * v[((b * H + h) * N + j) * HD + tid];
        }
        ctx[(b * N + i) * D + h * HD + tid] = val;
    }
}

__global__ void out_kernel(
    const float* __restrict__ nodes,
    const float* __restrict__ ctx,
    const float* __restrict__ ow,
    float* __restrict__ out,
    int B) {
    int idx = blockIdx.x;
    int b = idx / N;
    int n = idx - b * N;
    int tid = threadIdx.x;
    if (tid < D) {
        const float* c = ctx + (b * N + n) * D;
        const float* w = ow + tid * D;
        float sum = 0.0f;
        #pragma unroll
        for (int e = 0; e < D; ++e) sum += c[e] * w[e];
        int o = (b * N + n) * D + tid;
        out[o] = nodes[o] + sum;
    }
}

at::Tensor gf_forward(
    at::Tensor nodes,
    at::Tensor pair_feats,
    at::Tensor edge_mask,
    at::Tensor norm_w,
    at::Tensor norm_b,
    at::Tensor q_w,
    at::Tensor k_w,
    at::Tensor v_w,
    at::Tensor rel1_w,
    at::Tensor rel1_b,
    at::Tensor rel2_w,
    at::Tensor rel2_b,
    at::Tensor out_w,
    double eps,
    double scale) {

    const int B = (int)nodes.size(0);
    auto opts = nodes.options();
    auto q = at::empty({B, H, N, HD}, opts);
    auto k = at::empty({B, H, N, HD}, opts);
    auto v = at::empty({B, H, N, HD}, opts);
    auto ctx = at::empty({B, N, D}, opts);
    auto out = at::empty({B, N, D}, opts);

    at::cuda::CUDAGuard device_guard(nodes.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    ln_qkv_kernel<<<B * N, 128, 0, stream>>>(
        nodes.data_ptr<float>(), norm_w.data_ptr<float>(), norm_b.data_ptr<float>(),
        q_w.data_ptr<float>(), k_w.data_ptr<float>(), v_w.data_ptr<float>(),
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), B, (float)eps);

    dim3 grid(B, H, N);
    attn_kernel<<<grid, 32, 0, stream>>>(
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(),
        pair_feats.data_ptr<float>(), edge_mask.data_ptr<float>(),
        rel1_w.data_ptr<float>(), rel1_b.data_ptr<float>(),
        rel2_w.data_ptr<float>(), rel2_b.data_ptr<float>(),
        ctx.data_ptr<float>(), B, (float)scale);

    out_kernel<<<B * N, 128, 0, stream>>>(
        nodes.data_ptr<float>(), ctx.data_ptr<float>(), out_w.data_ptr<float>(), out.data_ptr<float>(), B);

    return out;
}
"""

_ext = load_inline(
    name="gf_vla_fused_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["gf_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, dim: int, pair_dim: int, num_heads: int):
        super().__init__()
        assert dim % num_heads == 0
        self.dim = dim
        self.pair_dim = pair_dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)
        self.norm = nn.LayerNorm(dim)
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.rel_mlp = nn.Sequential(
            nn.Linear(pair_dim, pair_dim),
            nn.SiLU(),
            nn.Linear(pair_dim, num_heads),
        )
        self.out_proj = nn.Linear(dim, dim, bias=False)

    def forward(self, nodes: torch.Tensor, pair_feats: torch.Tensor, edge_mask: torch.Tensor) -> torch.Tensor:
        return _ext.gf_forward(
            nodes,
            pair_feats,
            edge_mask,
            self.norm.weight,
            self.norm.bias,
            self.q_proj.weight,
            self.k_proj.weight,
            self.v_proj.weight,
            self.rel_mlp[0].weight,
            self.rel_mlp[0].bias,
            self.rel_mlp[2].weight,
            self.rel_mlp[2].bias,
            self.out_proj.weight,
            self.norm.eps,
            self.scale,
        )