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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cfloat>

__device__ __forceinline__ float silu(float x) {
    return x * (1.0f / (1.0f + __expf(-x)));
}

__global__ void scene_graph_kernel(
    const float* __restrict__ nodes,
    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,
    const float* __restrict__ ln_w,
    const float* __restrict__ ln_b,
    float* __restrict__ out,
    int B, int N, int D, int P, int H
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N * D;
    if (idx >= total) return;

    int d = idx % D;
    int i = (idx / D) % N;
    int b = idx / (D * N);

    float maxv = -FLT_MAX;

    int node_offset = b * N * D;
    int mask_offset = b * N * N;
    int pair_offset_base = b * N * N * P;
    int w1_row_stride = 2 * D + P;

    for (int j = 0; j < N; j++) {
        if (edge_mask[mask_offset + i * N + j] < 0.5f) continue;

        float hval = 0.0f;
        for (int hh = 0; hh < H; hh++) {
            float sum = b1[hh];
            const float* ni = nodes + node_offset + i * D;
            const float* w1_row = w1 + hh * w1_row_stride;
            for (int dd = 0; dd < D; dd++) {
                sum += ni[dd] * w1_row[dd];
            }
            const float* nj = nodes + node_offset + j * D;
            for (int dd = 0; dd < D; dd++) {
                sum += nj[dd] * w1_row[D + dd];
            }
            const float* pr = pair_feats + pair_offset_base + (i * N + j) * P;
            for (int pp = 0; pp < P; pp++) {
                sum += pr[pp] * w1_row[2 * D + pp];
            }
            float act = silu(sum);
            hval += act * w2[d * H + hh];
        }
        hval += b2[d];
        if (hval > maxv) maxv = hval;
    }

    if (maxv == -FLT_MAX) maxv = 0.0f;

    out[idx] = maxv;
}

torch::Tensor scene_graph_cuda(
    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 ln_w, torch::Tensor ln_b
) {
    auto B = nodes.size(0);
    auto N = nodes.size(1);
    auto D = nodes.size(2);
    auto P = pair_feats.size(3);
    auto H = w1.size(0);
    auto out = torch::empty({B, N, D}, nodes.options());
    int total = B * N * D;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    scene_graph_kernel<<<blocks, threads>>>(
        nodes.data_ptr<float>(),
        pair_feats.data_ptr<float>(),
        edge_mask.data_ptr<float>(),
        w1.data_ptr<float>(),
        b1.data_ptr<float>(),
        w2.data_ptr<float>(),
        b2.data_ptr<float>(),
        ln_w.data_ptr<float>(),
        ln_b.data_ptr<float>(),
        out.data_ptr<float>(),
        B, N, D, P, H
    );
    return out;
}
'''

cpp_source = r'''
torch::Tensor scene_graph_cuda(
    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 ln_w, torch::Tensor ln_b);
'''

mod = load_inline(
    name='scene_graph',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['scene_graph_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math']
)

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.hidden = hidden
        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, pair_feats, edge_mask):
        w1 = self.fuse[0].weight
        b1 = self.fuse[0].bias
        w2 = self.fuse[2].weight
        b2 = self.fuse[2].bias
        ln_w = self.out_norm.weight
        ln_b = self.out_norm.bias
        pooled = mod.scene_graph_cuda(nodes.contiguous(), pair_feats.contiguous(), edge_mask.contiguous(), w1, b1, w2, b2, ln_w, ln_b)
        pooled = self.out_norm(pooled)
        return nodes + pooled