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):
    """
    GF-VLA style **fused pairwise → node** scene-graph readout.

    For every ordered pair (i, j), fuse ``[node_i || node_j || pair_feat_ij]`` with
    an MLP, mask invalid graph edges, then **max-pool over outgoing neighbors j**
    to produce an updated per-object signal (object-centric scene graph refresh).

    Complements ``7_GFVLA_FusedSceneGraphUpdate`` (attention) with an MLP + reduce
    fusion pattern common for relation classifiers and affordance heads.
    """

    def __init__(self, dim: int, pair_dim: int, hidden: int):
        super().__init__()
        self.dim = dim
        self.pair_dim = pair_dim
        self.fuse = nn.Sequential(
            nn.Linear(2 * dim + pair_dim, hidden),
            nn.SiLU(),
            nn.Linear(hidden, dim),
        )
        self.out_norm = nn.LayerNorm(dim)

    def forward(
        self,
        nodes: torch.Tensor,
        pair_feats: torch.Tensor,
        edge_mask: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            nodes: (B, N, D)
            pair_feats: (B, N, N, P)
            edge_mask: (B, N, N) float {0,1}
        Returns:
            (B, N, D) residual ``nodes + LN( max_j fuse(...) )``
        """
        B, N, D = nodes.shape
        assert pair_feats.shape == (B, N, N, self.pair_dim)

        hi = nodes.unsqueeze(2).expand(-1, -1, N, -1)
        hj = nodes.unsqueeze(1).expand(-1, N, -1, -1)
        fused = self.fuse(torch.cat([hi, hj, pair_feats], dim=-1))

        neg = torch.finfo(fused.dtype).min
        fused = fused.masked_fill(edge_mask.unsqueeze(-1) < 0.5, neg)
        pooled, _ = fused.max(dim=2)
        return nodes + self.out_norm(pooled)


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

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

    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
hidden = 192


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, hidden]

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

_cpp_src = 'torch::Tensor fused_forward(torch::Tensor nodes, torch::Tensor pair_feats, torch::Tensor edge_mask, torch::Tensor w1, torch::Tensor b1, torch::Tensor w2, torch::Tensor b2, torch::Tensor gamma, torch::Tensor beta);'
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>

__global__ void pre_node_kernel(const float* __restrict__ nodes, const float* __restrict__ w1,
                                float* __restrict__ pre_i, float* __restrict__ pre_j,
                                int rows, int D, int P, int H) {
    int row = blockIdx.x;
    int h = threadIdx.x;
    if (h >= H) return;
    const float* n = nodes + row * D;
    const float* wi = w1 + h * (2 * D + P);
    float ai = 0.0f, aj = 0.0f;
    #pragma unroll 4
    for (int k = 0; k < D; ++k) {
        float x = n[k]; ai = fmaf(x, wi[k], ai); aj = fmaf(x, wi[D + k], aj);
    }
    pre_i[row * H + h] = ai;
    pre_j[row * H + h] = aj;
}

__global__ void hidden_pair_kernel(const float* __restrict__ pair_feats, const float* __restrict__ w1,
                                   const float* __restrict__ b1, const float* __restrict__ pre_i,
                                   const float* __restrict__ pre_j, float* __restrict__ hidden,
                                   int B, int N, int D, int P, int H) {
    int pair = blockIdx.x;
    int h = threadIdx.x;
    if (h >= H) return;
    int j = pair % N;
    int tmp = pair / N;
    int i = tmp % N;
    int b = tmp / N;
    int row_i = b * N + i;
    int row_j = b * N + j;
    const float* wi = w1 + h * (2 * D + P) + 2 * D;
    float acc = b1[h] + pre_i[row_i * H + h] + pre_j[row_j * H + h];
    const float* pf = pair_feats + pair * P;
    #pragma unroll
    for (int k = 0; k < 16; ++k) { if (k < P) acc = fmaf(pf[k], wi[k], acc); }
    float sig = 1.0f / (1.0f + expf(-acc));
    hidden[pair * H + h] = acc * sig;
}

__inline__ __device__ float warp_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(0xffffffff, v, offset);
    return v;
}
__inline__ __device__ float warp_max(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffff, v, offset));
    return v;
}

__global__ void second_pair_kernel(const float* __restrict__ hidden, const float* __restrict__ w2,
                                   const float* __restrict__ b2, float* __restrict__ pair_out,
                                   int total_pairs, int D, int H) {
    int pair = blockIdx.y;
    int warp = threadIdx.x >> 5;
    int lane = threadIdx.x & 31;
    int d = blockIdx.x * 8 + warp;
    if (d >= D) return;
    const float* hp = hidden + pair * H;
    const float* wd = w2 + d * H;
    float sum = 0.0f;
    for (int h = lane; h < H; h += 32) sum = fmaf(hp[h], wd[h], sum);
    sum = warp_sum(sum);
    if (lane == 0) pair_out[pair * D + d] = sum + b2[d];
}

__global__ void reduce_pool_kernel(const float* __restrict__ pair_out, const float* __restrict__ edge_mask,
                                   float* __restrict__ pooled, int B, int N, int D) {
    int row = blockIdx.y;
    int warp = threadIdx.x >> 5;
    int lane = threadIdx.x & 31;
    int d = blockIdx.x * 4 + warp;
    if (d >= D) return;
    float v = -FLT_MAX;
    if (lane < N && edge_mask[row * N + lane] >= 0.5f) {
        int pair = row * N + lane;
        v = pair_out[pair * D + d];
    }
    v = warp_max(v);
    if (lane == 0) pooled[row * D + d] = v;
}

__global__ void layernorm_kernel(const float* __restrict__ nodes, const float* __restrict__ pooled,
                                 const float* __restrict__ gamma, const float* __restrict__ beta,
                                 float* __restrict__ out, int rows, int D) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    float sum = 0.0f, sumsq = 0.0f;
    for (int d = tid; d < D; d += blockDim.x) { float x = pooled[row * D + d]; sum += x; sumsq = fmaf(x, x, sumsq); }
    __shared__ float smem[256];
    smem[tid] = sum; smem[tid + blockDim.x] = sumsq;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) { smem[tid] += smem[tid + s]; smem[tid + blockDim.x] += smem[tid + blockDim.x + s]; }
        __syncthreads();
    }
    float mean = smem[0] / (float)D;
    float var = smem[blockDim.x] / (float)D - mean * mean;
    float inv = rsqrtf(var + 1.0e-5f);
    for (int d = tid; d < D; d += blockDim.x) {
        float x = pooled[row * D + d]; out[row * D + d] = nodes[row * D + d] + (x - mean) * inv * gamma[d] + beta[d];
    }
}

torch::Tensor fused_forward(torch::Tensor nodes, torch::Tensor pair_feats, torch::Tensor edge_mask,
                            torch::Tensor w1, torch::Tensor b1, torch::Tensor w2, torch::Tensor b2,
                            torch::Tensor gamma, torch::Tensor beta) {
    int B = nodes.size(0), N = nodes.size(1), D = nodes.size(2), P = pair_feats.size(3), H = w1.size(0);
    auto opts = nodes.options();
    auto pre_i = torch::empty({B, N, H}, opts);
    auto pre_j = torch::empty({B, N, H}, opts);
    auto hidden = torch::empty({B, N, N, H}, opts);
    auto pair_out = torch::empty({B, N, N, D}, opts);
    auto pooled = torch::empty({B, N, D}, opts);
    auto out = torch::empty_like(nodes);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int total_pairs = B * N * N;
    pre_node_kernel<<<B*N, 256, 0, stream>>>(nodes.data_ptr<float>(), w1.data_ptr<float>(), pre_i.data_ptr<float>(), pre_j.data_ptr<float>(), B*N, D, P, H);
    hidden_pair_kernel<<<total_pairs, 256, 0, stream>>>(pair_feats.data_ptr<float>(), w1.data_ptr<float>(), b1.data_ptr<float>(), pre_i.data_ptr<float>(), pre_j.data_ptr<float>(), hidden.data_ptr<float>(), B, N, D, P, H);
    second_pair_kernel<<<dim3((D + 7) / 8, total_pairs), 256, 0, stream>>>(hidden.data_ptr<float>(), w2.data_ptr<float>(), b2.data_ptr<float>(), pair_out.data_ptr<float>(), total_pairs, D, H);
    reduce_pool_kernel<<<dim3((D + 3) / 4, B * N), 128, 0, stream>>>(pair_out.data_ptr<float>(), edge_mask.data_ptr<float>(), pooled.data_ptr<float>(), B, N, D);
    layernorm_kernel<<<B * N, 128, 0, stream>>>(nodes.data_ptr<float>(), pooled.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(), B * N, D);
    return out;
}
'''
_ext = load_inline(name='gf_vla_fused_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['fused_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)
class ModelNew(nn.Module):
    def __init__(self, dim: int, pair_dim: int, hidden: int):
        super().__init__(); self.dim = dim; self.pair_dim = pair_dim
        self.fuse = nn.Sequential(nn.Linear(2 * dim + pair_dim, hidden), nn.SiLU(), nn.Linear(hidden, dim)); self.out_norm = nn.LayerNorm(dim)
    def forward(self, nodes: torch.Tensor, pair_feats: torch.Tensor, edge_mask: torch.Tensor) -> torch.Tensor:
        return _ext.fused_forward(nodes, pair_feats, edge_mask, self.fuse[0].weight, self.fuse[0].bias, self.fuse[2].weight, self.fuse[2].bias, self.out_norm.weight, self.out_norm.bias)