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>

template <typename scalar_t>
__device__ __forceinline__ scalar_t silu(scalar_t x) {
    return x / (scalar_t(1.0) + exp(-x));
}

// Kernel for fused MLP + max pooling
// Each thread processes one (b, i) pair and iterates over j
template <typename scalar_t, int D_MAX, int P_MAX, int H_MAX>
__global__ void fused_mlp_maxpool_kernel_templated(
    const scalar_t* __restrict__ nodes,
    const scalar_t* __restrict__ pair_feats,
    const scalar_t* __restrict__ edge_mask,
    const scalar_t* __restrict__ w1,
    const scalar_t* __restrict__ b1,
    const scalar_t* __restrict__ w2,
    const scalar_t* __restrict__ b2,
    scalar_t* __restrict__ pooled,
    int B, int N, int D, int P, int H) {
    
    int b = blockIdx.y;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (b >= B || i >= N) return;
    
    // Load node i into local memory (registers if small enough)
    scalar_t node_i[D_MAX];
    for (int d = 0; d < D; d++) {
        node_i[d] = nodes[(b * N + i) * D + d];
    }
    
    // Negative infinity for masking
    scalar_t neg_inf = -1e30f;
    
    // Initialize pooled result
    scalar_t pooled_out[D_MAX];
    for (int d = 0; d < D; d++) {
        pooled_out[d] = neg_inf;
    }
    
    // Iterate over j
    for (int j = 0; j < N; j++) {
        // Check edge mask
        scalar_t mask = edge_mask[b * N * N + i * N + j];
        if (mask < scalar_t(0.5)) continue;
        
        // Load node j
        scalar_t node_j[D_MAX];
        for (int d = 0; d < D; d++) {
            node_j[d] = nodes[(b * N + j) * D + d];
        }
        
        // Load pair features for (i,j)
        scalar_t pf[P_MAX];
        for (int p = 0; p < P; p++) {
            pf[p] = pair_feats[((b * N + i) * N + j) * P + p];
        }
        
        // First linear: input = [node_i, node_j, pf], size = 2D + P
        // Output = hidden, size = H
        scalar_t hidden1[H_MAX];
        
        for (int h = 0; h < H; h++) {
            scalar_t sum = b1[h];
            // node_i contribution
            for (int d = 0; d < D; d++) {
                sum += node_i[d] * w1[d * H + h];
            }
            // node_j contribution
            for (int d = 0; d < D; d++) {
                sum += node_j[d] * w1[(D + d) * H + h];
            }
            // pair_feat contribution
            for (int p = 0; p < P; p++) {
                sum += pf[p] * w1[(2 * D + p) * H + h];
            }
            hidden1[h] = sum;
        }
        
        // SiLU activation
        for (int h = 0; h < H; h++) {
            hidden1[h] = silu(hidden1[h]);
        }
        
        // Second linear: input = hidden, size = H, output = D
        scalar_t output[D_MAX];
        for (int d = 0; d < D; d++) {
            scalar_t sum = b2[d];
            for (int h = 0; h < H; h++) {
                sum += hidden1[h] * w2[h * D + d];
            }
            output[d] = sum;
        }
        
        // Max pooling
        for (int d = 0; d < D; d++) {
            if (output[d] > pooled_out[d]) {
                pooled_out[d] = output[d];
            }
        }
    }
    
    // Store pooled result
    for (int d = 0; d < D; d++) {
        pooled[(b * N + i) * D + d] = pooled_out[d];
    }
}

// LayerNorm kernel
template <typename scalar_t>
__global__ void layernorm_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ output,
    int D, int total_elements) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_elements / D) return;
    
    int base = idx * D;
    
    // Compute mean
    scalar_t mean = scalar_t(0.0);
    for (int d = 0; d < D; d++) {
        mean += input[base + d];
    }
    mean /= scalar_t(D);
    
    // Compute variance
    scalar_t var = scalar_t(0.0);
    for (int d = 0; d < D; d++) {
        scalar_t diff = input[base + d] - mean;
        var += diff * diff;
    }
    var /= scalar_t(D);
    
    // Normalize with eps=1e-5
    scalar_t inv_std = rsqrt(var + scalar_t(1e-5));
    
    // Apply weight and bias
    for (int d = 0; d < D; d++) {
        output[base + d] = (input[base + d] - mean) * inv_std * weight[d] + bias[d];
    }
}

// Residual add kernel
template <typename scalar_t>
__global__ void residual_add_kernel(
    const scalar_t* __restrict__ nodes,
    const scalar_t* __restrict__ normalized,
    scalar_t* __restrict__ output,
    int total_elements) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_elements) return;
    
    output[idx] = nodes[idx] + normalized[idx];
}

torch::Tensor forward_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_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(1);
    
    auto options = nodes.options();
    auto pooled = torch::empty({B, N, D}, options);
    auto normalized = torch::empty({B, N, D}, options);
    auto output = torch::empty({B, N, D}, options);
    
    dim3 block(32);
    dim3 grid((N + block.x - 1) / block.x, B);
    
    AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), "fused_mlp_maxpool", ([&] {
        fused_mlp_maxpool_kernel_templated<scalar_t, 128, 64, 256><<<grid, block>>>(
            nodes.data_ptr<scalar_t>(),
            pair_feats.data_ptr<scalar_t>(),
            edge_mask.data_ptr<scalar_t>(),
            w1.data_ptr<scalar_t>(),
            b1.data_ptr<scalar_t>(),
            w2.data_ptr<scalar_t>(),
            b2.data_ptr<scalar_t>(),
            pooled.data_ptr<scalar_t>(),
            B, N, D, P, H);
    }));
    
    // LayerNorm
    int total = B * N;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), "layernorm", ([&] {
        layernorm_kernel<scalar_t><<<blocks, threads>>>(
            pooled.data_ptr<scalar_t>(),
            ln_weight.data_ptr<scalar_t>(),
            ln_bias.data_ptr<scalar_t>(),
            normalized.data_ptr<scalar_t>(),
            D, total * D);
    }));
    
    // Residual add
    int total_elements = B * N * D;
    blocks = (total_elements + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES(nodes.scalar_type(), "residual_add", ([&] {
        residual_add_kernel<scalar_t><<<blocks, threads>>>(
            nodes.data_ptr<scalar_t>(),
            normalized.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            total_elements);
    }));
    
    return output;
}
"""

cpp_source = "torch::Tensor forward_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_weight, torch::Tensor ln_bias);"

fused_module = load_inline(
    name='fused_scene_graph_v3',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['forward_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=True
)


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
        
        # Create PyTorch layers to hold weights
        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:
        # Get weights from PyTorch layers
        w1 = self.fuse[0].weight.t().contiguous()  # (in_features, out_features) -> transpose
        b1 = self.fuse[0].bias.contiguous()
        w2 = self.fuse[2].weight.t().contiguous()
        b2 = self.fuse[2].bias.contiguous()
        ln_weight = self.out_norm.weight.contiguous()
        ln_bias = self.out_norm.bias.contiguous()
        
        return fused_module.forward_cuda(
            nodes, pair_feats, edge_mask,
            w1, b1, w2, b2, ln_weight, ln_bias
        )


batch_size = 4
num_objects = 24
dim = 96
pair_dim = 16
hidden = 192


def get_inputs():
    nodes = torch.randn(batch_size, num_objects, dim)
    pair_feats = torch.randn(batch_size, num_objects, num_objects, pair_dim)
    edge_mask = torch.ones(batch_size, num_objects, num_objects)
    return [nodes, pair_feats, edge_mask]


def get_init_inputs():
    return [dim, pair_dim, hidden]