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 = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <cmath>

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

__device__ void atomicMaxFloat(float *addr, float val) {
    unsigned int *addr_as_ui = (unsigned int *)addr;
    unsigned int old = *addr_as_ui, assumed;
    do {
        assumed = old;
        if (__int_as_float(assumed) >= val) break;
        old = atomicCAS(addr_as_ui, assumed, __float_as_int(val));
    } while (assumed != old);
}

__global__ void fused_max_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,
    float *__restrict__ pooled,
    int B, int N, int D, int P, int H) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * N * N) return;
    int b = idx / (N * N);
    int tmp = idx % (N * N);
    int i = tmp / N;
    int j = tmp % N;
    if (edge_mask[b * N * N + i * N + j] < 0.5f) return;

    for (int dout = 0; dout < D; dout++) {
        float val = b2[dout];
        for (int h = 0; h < H; h++) {
            float sum = b1[h];
            int w1_base = h * (2 * D + P);
            const float *ni = nodes + (b * N + i) * D;
            const float *nj = nodes + (b * N + j) * D;
            const float *wp = w1 + w1_base;
            for (int d = 0; d < D; d++) sum += ni[d] * wp[d];
            for (int d = 0; d < D; d++) sum += nj[d] * wp[D + d];
            const float *pf = pair_feats + ((b * N + i) * N + j) * P;
            for (int p = 0; p < P; p++) sum += pf[p] * wp[2 * D + p];
            val += silu(sum) * w2[dout * H + h];
        }
        atomicMaxFloat(&pooled[(b * N + i) * D + dout], val);
    }
}

__global__ void ln_residual_kernel(
    const float *__restrict__ nodes,
    const float *__restrict__ pooled,
    const float *__restrict__ nw, const float *__restrict__ nb,
    float *__restrict__ out,
    int B, int N, int D, float eps) {
    extern __shared__ float sdata[];
    int row = blockIdx.x;
    if (row >= B * N) return;
    const float *p = pooled + row * D;
    const float *nd = nodes + row * D;
    float *o = out + row * D;

    // mean
    float sum = 0.0f;
    for (int d = threadIdx.x; d < D; d += blockDim.x) sum += p[d];
    sdata[threadIdx.x] = sum;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    float mean = sdata[0] / D;

    // var
    float varsum = 0.0f;
    for (int d = threadIdx.x; d < D; d += blockDim.x) {
        float diff = p[d] - mean;
        varsum += diff * diff;
    }
    sdata[threadIdx.x] = varsum;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    float rstd = rsqrtf(sdata[0] / D + eps);

    for (int d = threadIdx.x; d < D; d += blockDim.x) {
        float nval = (p[d] - mean) * rstd * nw[d] + nb[d];
        o[d] = nd[d] + nval;
    }
}

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 norm_w, torch::Tensor norm_b) {
    int B = nodes.size(0), N = nodes.size(1), D = nodes.size(2);
    int P = pair_feats.size(3), H = b1.size(0);
    auto pooled = torch::full({B, N, D}, -FLT_MAX, nodes.options());
    auto output = torch::empty_like(nodes);
    int threads = 256;
    fused_max_kernel<<<(B * N * N + threads - 1) / threads, 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>(),
        pooled.data_ptr<float>(), B, N, D, P, H);
    ln_residual_kernel<<<B * N, threads, threads * sizeof(float)>>>(
        nodes.data_ptr<float>(), pooled.data_ptr<float>(),
        norm_w.data_ptr<float>(), norm_b.data_ptr<float>(),
        output.data_ptr<float>(), B, N, D, 1e-5f);
    return output;
}
'''

cpp_source = '''
torch::Tensor fused_forward(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
                            torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor);
'''

mod = load_inline(name='fsg', cpp_sources=[cpp_source], cuda_sources=[cuda_source], functions=['fused_forward'], 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.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: torch.Tensor, pair_feats: torch.Tensor, edge_mask: torch.Tensor) -> torch.Tensor:
        w1, b1 = self.fuse[0].weight, self.fuse[0].bias
        w2, b2 = self.fuse[2].weight, self.fuse[2].bias
        nw, nb = self.out_norm.weight, self.out_norm.bias
        return mod.fused_forward(nodes, pair_feats, edge_mask, w1, b1, w2, b2, nw, nb)