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 = r"""
#include <torch/extension.h>

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_w,
                            torch::Tensor ln_b,
                            int64_t dim,
                            int64_t pair_dim,
                            int64_t hidden);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__global__ void fc1_silu_kernel(const float* __restrict__ nodes,
                                const float* __restrict__ pair_feats,
                                const float* __restrict__ w1,
                                const float* __restrict__ b1,
                                float* __restrict__ hidden_out,
                                int B, int N, int D, int P, int H) {
    int pair_idx = blockIdx.x;
    int j = pair_idx % N;
    int t = pair_idx / N;
    int i = t % N;
    int b = t / N;
    int IN = 2 * D + P;

    const float* ni = nodes + (b * N + i) * D;
    const float* nj = nodes + (b * N + j) * D;
    const float* pf = pair_feats + ((b * N + i) * N + j) * P;
    float* hout = hidden_out + ((b * N + i) * N + j) * H;

    for (int k = threadIdx.x; k < H; k += blockDim.x) {
        const float* wk = w1 + k * IN;
        float acc = b1[k];
        #pragma unroll 4
        for (int d = 0; d < D; ++d) acc = fmaf(ni[d], wk[d], acc);
        #pragma unroll 4
        for (int d = 0; d < D; ++d) acc = fmaf(nj[d], wk[D + d], acc);
        #pragma unroll 4
        for (int p = 0; p < P; ++p) acc = fmaf(pf[p], wk[2 * D + p], acc);
        float sig = 1.0f / (1.0f + expf(-acc));
        hout[k] = acc * sig;
    }
}

__global__ void fc2_max_kernel(const float* __restrict__ hidden,
                               const float* __restrict__ edge_mask,
                               const float* __restrict__ w2,
                               const float* __restrict__ b2,
                               float* __restrict__ pooled,
                               int B, int N, int D, int H) {
    extern __shared__ float sh[];
    int idx = blockIdx.x;
    int d = idx % D;
    int row = idx / D; // b*N+i
    int tid = threadIdx.x;
    int b = row / N;
    int i = row - b * N;

    float maxv = -FLT_MAX;
    const float* wd = w2 + d * H;
    for (int j = 0; j < N; ++j) {
        float m = edge_mask[(row * N) + j];
        float part = 0.0f;
        if (m >= 0.5f) {
            const float* hptr = hidden + (row * N + j) * H;
            for (int k = tid; k < H; k += blockDim.x) {
                part = fmaf(hptr[k], wd[k], part);
            }
        }
        sh[tid] = part;
        __syncthreads();
        for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
            if (tid < s) sh[tid] += sh[tid + s];
            __syncthreads();
        }
        if (tid == 0 && m >= 0.5f) {
            float val = sh[0] + b2[d];
            maxv = fmaxf(maxv, val);
        }
        __syncthreads();
    }
    if (tid == 0) pooled[row * D + d] = maxv;
}

__global__ void layernorm_residual_kernel(const float* __restrict__ nodes,
                                          const float* __restrict__ pooled,
                                          const float* __restrict__ ln_w,
                                          const float* __restrict__ ln_b,
                                          float* __restrict__ out,
                                          int rows, int D) {
    extern __shared__ float sh[];
    int row = blockIdx.x;
    int tid = threadIdx.x;

    float sum = 0.0f;
    for (int d = tid; d < D; d += blockDim.x) sum += pooled[row * D + d];
    sh[tid] = sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) sh[tid] += sh[tid + s];
        __syncthreads();
    }
    float mean = sh[0] / (float)D;

    float vsum = 0.0f;
    for (int d = tid; d < D; d += blockDim.x) {
        float x = pooled[row * D + d] - mean;
        vsum = fmaf(x, x, vsum);
    }
    sh[tid] = vsum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) sh[tid] += sh[tid + s];
        __syncthreads();
    }
    float invstd = rsqrtf(sh[0] / (float)D + 1.0e-5f);

    for (int d = tid; d < D; d += blockDim.x) {
        int off = row * D + d;
        float y = (pooled[off] - mean) * invstd;
        out[off] = nodes[off] + y * ln_w[d] + ln_b[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 ln_w,
                            torch::Tensor ln_b,
                            int64_t dim,
                            int64_t pair_dim,
                            int64_t hidden_size) {
    int B = (int)nodes.size(0);
    int N = (int)nodes.size(1);
    int D = (int)dim;
    int P = (int)pair_dim;
    int H = (int)hidden_size;

    auto opts = nodes.options();
    torch::Tensor hidden = torch::empty({B, N, N, H}, opts);
    torch::Tensor pooled = torch::empty({B, N, D}, opts);
    torch::Tensor out = torch::empty({B, N, D}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int threads1 = 256;
    fc1_silu_kernel<<<B * N * N, threads1, 0, stream>>>(
        nodes.data_ptr<float>(), pair_feats.data_ptr<float>(), w1.data_ptr<float>(), b1.data_ptr<float>(),
        hidden.data_ptr<float>(), B, N, D, P, H);

    int threads2 = 256;
    fc2_max_kernel<<<B * N * D, threads2, threads2 * sizeof(float), stream>>>(
        hidden.data_ptr<float>(), edge_mask.data_ptr<float>(), w2.data_ptr<float>(), b2.data_ptr<float>(),
        pooled.data_ptr<float>(), B, N, D, H);

    int threads3 = 128;
    layernorm_residual_kernel<<<B * N, threads3, threads3 * sizeof(float), stream>>>(
        nodes.data_ptr<float>(), pooled.data_ptr<float>(), ln_w.data_ptr<float>(), ln_b.data_ptr<float>(),
        out.data_ptr<float>(), B * N, D);
    return out;
}
"""

_ext = load_inline(
    name="gf_vla_fused_pairwise_node_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    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:
        return _ext.fused_forward(
            nodes.contiguous(),
            pair_feats.contiguous(),
            edge_mask.contiguous(),
            self.fuse[0].weight.contiguous(),
            self.fuse[0].bias.contiguous(),
            self.fuse[2].weight.contiguous(),
            self.fuse[2].bias.contiguous(),
            self.out_norm.weight.contiguous(),
            self.out_norm.bias.contiguous(),
            self.dim,
            self.pair_dim,
            self.hidden,
        )