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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

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

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

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

    int b = idx / N;
    int i = idx % N;

    extern __shared__ float sdata[]; // for reduction, but per block we'll handle max

    // For simplicity, each thread computes for one (b,i), loops over j and D
    float max_val[96]; // assume D <=96
    for (int d=0; d<D; d++) max_val[d] = -1e30f;

    float neg = -1e30f; // approx -inf

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

        // compute fused for this i,j
        // first hidden
        for (int h=0; h<H; h++) {
            float sum = b1[h];
            // hi part: nodes[b,i,:] dot W1[h, :D]
            for (int d=0; d<D; d++) {
                sum += nodes[b*N*D + i*D + d] * W1[h * (2*D + P) + d];
            }
            // hj part
            for (int d=0; d<D; d++) {
                sum += nodes[b*N*D + j*D + d] * W1[h * (2*D + P) + D + d];
            }
            // pair
            for (int p=0; p<P; p++) {
                sum += pair_feats[b*N*N*P + i*N*P + j*P + p] * W1[h * (2*D + P) + 2*D + p];
            }
            float act = silu(sum);
            // now second linear for each d, accumulate max
            // actually need to do W2 for full vector
        }
        // optimize: better to compute full fused vector
    }

    // re-write properly for vector
    for (int j=0; j<N; j++) {
        if (edge_mask[b*N*N + i*N + j] < 0.5f) continue;
        for (int d=0; d<D; d++) {
            float sum = b2[d];
            // compute hidden activations on fly? inefficient but N small
            for (int h=0; h<H; h++) {
                float hsum = b1[h];
                for (int dd=0; dd<D; dd++) {
                    hsum += nodes[b*N*D + i*D + dd] * W1[h*(2*D+P) + dd];
                }
                for (int dd=0; dd<D; dd++) {
                    hsum += nodes[b*N*D + j*D + dd] * W1[h*(2*D+P) + D + dd];
                }
                for (int p=0; p<P; p++) {
                    hsum += pair_feats[b*N*N*P + i*N*P + j*P + p] * W1[h*(2*D+P) + 2*D + p];
                }
                sum += silu(hsum) * W2[d*H + h];
            }
            if (sum > max_val[d]) max_val[d] = sum;
        }
    }

    // now layernorm and add
    float mean = 0.0f;
    for (int d=0; d<D; d++) mean += max_val[d];
    mean /= D;
    float var = 0.0f;
    for (int d=0; d<D; d++) {
        float t = max_val[d] - mean;
        var += t*t;
    }
    var /= D;
    float inv_std = rsqrtf(var + 1e-5f);
    for (int d=0; d<D; d++) {
        float ln = (max_val[d] - mean) * inv_std * ln_weight[d] + ln_bias[d];
        output[b*N*D + i*D + d] = nodes[b*N*D + i*D + d] + ln;
    }
}

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 ln_weight, torch::Tensor ln_bias
) {
    int B = nodes.size(0);
    int N = nodes.size(1);
    int D = nodes.size(2);
    int P = pair_feats.size(3);
    int H = W1.size(0);

    auto output = torch::zeros_like(nodes);

    int threads = 256;
    int blocks = (B * N + threads - 1) / threads;

    fused_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_weight.data_ptr<float>(),
        ln_bias.data_ptr<float>(),
        output.data_ptr<float>(),
        B, N, D, P, H
    );

    return output;
}
'''

cpp_source = '''
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 ln_weight, torch::Tensor ln_bias
);
'''

module = load_inline(
    name="fused_scene_graph",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["fused_forward"],
    verbose=True,
    extra_cuda_cflags=["-O3"]
)

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
        return module.fused_forward(nodes, pair_feats, edge_mask, W1, b1, W2, b2, ln_w, ln_b)