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>
torch::Tensor gf_vla_forward(torch::Tensor nodes, torch::Tensor pair_feats, torch::Tensor edge_mask,
                             torch::Tensor norm_w, torch::Tensor norm_b,
                             torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
                             torch::Tensor rel_w1, torch::Tensor rel_b1,
                             torch::Tensor rel_w2, torch::Tensor rel_b2,
                             torch::Tensor out_w);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.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_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

template<int BLOCK> __device__ float block_sum(float v){__shared__ float smem[BLOCK];int tid=threadIdx.x;smem[tid]=v;__syncthreads();for(int s=BLOCK/2;s>0;s>>=1){if(tid<s)smem[tid]+=smem[tid+s];__syncthreads();}return smem[0];}
template<int BLOCK> __device__ float block_max(float v){__shared__ float smem[BLOCK];int tid=threadIdx.x;smem[tid]=v;__syncthreads();for(int s=BLOCK/2;s>0;s>>=1){if(tid<s)smem[tid]=fmaxf(smem[tid],smem[tid+s]);__syncthreads();}return smem[0];}

__global__ void ln_qkv_kernel(const float* __restrict__ nodes,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ Wq,const float* __restrict__ Wk,const float* __restrict__ Wv,float* __restrict__ q,float* __restrict__ k,float* __restrict__ v,int B,int D,int H,int Hd){
    const int BLOCK=128; int tid=threadIdx.x; int r=blockIdx.x; const float* x=nodes+r*D;
    float val=0.f; for(int d=tid;d<D;d+=BLOCK) val+=x[d]; float mean=block_sum<BLOCK>(val)/(float)D;
    val=0.f; for(int d=tid;d<D;d+=BLOCK){float z=x[d]-mean; val+=z*z;} float inv=rsqrtf(block_sum<BLOCK>(val)/(float)D+1.0e-5f);
    __shared__ float hbuf[128]; for(int d=tid;d<D;d+=BLOCK) hbuf[d]=(x[d]-mean)*inv*gamma[d]+beta[d]; __syncthreads();
    int b=r/24; int n=r-b*24;
    for(int od=tid;od<D;od+=BLOCK){float sq=0.f,sk=0.f,sv=0.f; const float* wq=Wq+od*D; const float* wk=Wk+od*D; const float* wv=Wv+od*D;
        #pragma unroll
        for(int d=0;d<96;++d){float hv=hbuf[d]; sq+=hv*wq[d]; sk+=hv*wk[d]; sv+=hv*wv[d];}
        int hh=od/Hd; int dd=od-hh*Hd; int idx=((b*H+hh)*24+n)*Hd+dd; q[idx]=sq; k[idx]=sk; v[idx]=sv;}
}

__global__ void attn_rel_ctx_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__ W1,const float* __restrict__ b1,const float* __restrict__ W2,const float* __restrict__ b2,float* __restrict__ ctx,int B,int N,int H,int Hd,int P){
    const int BLOCK=64; int tid=threadIdx.x; int bid=blockIdx.x; int i=bid%N; int tmp=bid/N; int h=tmp%H; int b=tmp/H; __shared__ float logits[32]; float logit=-INFINITY; int j=tid;
    if(j<N){float dot=0.f; const float* qptr=q+((b*H+h)*N+i)*Hd; const float* kptr=k+((b*H+h)*N+j)*Hd;
        #pragma unroll
        for(int d=0;d<24;++d) dot+=qptr[d]*kptr[d];
        float rel=b2[h]; const float* pf=pair_feats+(((b*N+i)*N+j)*P);
        #pragma unroll
        for(int m=0;m<16;++m){float a=b1[m]; const float* w1=W1+m*P;
            #pragma unroll
            for(int p=0;p<16;++p) a+=pf[p]*w1[p];
            float silu=a/(1.0f+expf(-a)); rel+=silu*W2[h*P+m];}
        logit=dot*0.2041241452319315f+rel; if(edge_mask[(b*N+i)*N+j]<0.5f) logit=-INFINITY; logits[j]=logit;}
    float maxv=block_max<BLOCK>((tid<N)?logit:-INFINITY); float e=0.f; if(tid<N){e=expf(logits[tid]-maxv); logits[tid]=e;} float sumv=block_sum<BLOCK>((tid<N)?e:0.f);
    for(int d=tid;d<Hd;d+=BLOCK){float acc=0.f;
        #pragma unroll
        for(int jj=0;jj<24;++jj) acc+=(logits[jj]/sumv)*v[((b*H+h)*N+jj)*Hd+d];
        ctx[(b*N+i)*(H*Hd)+h*Hd+d]=acc;}
}

__global__ void out_res_kernel(const float* __restrict__ ctx,const float* __restrict__ Wout,const float* __restrict__ nodes,float* __restrict__ out,int D){
    const int BLOCK=128; int tid=threadIdx.x; int r=blockIdx.x; const float* c=ctx+r*D;
    for(int od=tid;od<D;od+=BLOCK){float acc=0.f; const float* w=Wout+od*D;
        #pragma unroll
        for(int d=0;d<96;++d) acc+=c[d]*w[d];
        out[r*D+od]=nodes[r*D+od]+acc;}
}

torch::Tensor gf_vla_forward(torch::Tensor nodes, torch::Tensor pair_feats, torch::Tensor edge_mask, torch::Tensor norm_w, torch::Tensor norm_b, torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w, torch::Tensor rel_w1, torch::Tensor rel_b1, torch::Tensor rel_w2, torch::Tensor rel_b2, torch::Tensor out_w){
    CHECK_INPUT(nodes); CHECK_INPUT(pair_feats); CHECK_INPUT(edge_mask); CHECK_INPUT(norm_w); CHECK_INPUT(norm_b); CHECK_INPUT(q_w); CHECK_INPUT(k_w); CHECK_INPUT(v_w); CHECK_INPUT(rel_w1); CHECK_INPUT(rel_b1); CHECK_INPUT(rel_w2); CHECK_INPUT(rel_b2); CHECK_INPUT(out_w);
    int B=nodes.size(0), N=nodes.size(1), D=nodes.size(2), P=pair_feats.size(3), H=rel_w2.size(0), Hd=D/H; TORCH_CHECK(D==96 && N==24 && P==16 && H==4 && Hd==24,"optimized kernel supports benchmark shape only"); auto opts=nodes.options();
    auto q=torch::empty({B,H,N,Hd},opts), k=torch::empty({B,H,N,Hd},opts), vv=torch::empty({B,H,N,Hd},opts), ctx=torch::empty({B,N,D},opts), out=torch::empty_like(nodes); int rows=B*N;
    ln_qkv_kernel<<<rows,128>>>(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>(),vv.data_ptr<float>(),B,D,H,Hd);
    attn_rel_ctx_kernel<<<B*H*N,64>>>(q.data_ptr<float>(),k.data_ptr<float>(),vv.data_ptr<float>(),pair_feats.data_ptr<float>(),edge_mask.data_ptr<float>(),rel_w1.data_ptr<float>(),rel_b1.data_ptr<float>(),rel_w2.data_ptr<float>(),rel_b2.data_ptr<float>(),ctx.data_ptr<float>(),B,N,H,Hd,P);
    out_res_kernel<<<rows,128>>>(ctx.data_ptr<float>(),out_w.data_ptr<float>(),nodes.data_ptr<float>(),out.data_ptr<float>(),D); return out;
}
'''
_ext = load_inline(name='gf_vla_ext_v2', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['gf_vla_forward'], verbose=False, extra_cuda_cflags=['-O3','--use_fast_math'])

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_vla_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)