import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    """
    Triangular multiplicative update – incoming edges.
    Mirror of the outgoing variant: updates pair (i,j) by aggregating over
    the shared index k via edges (k,i) and (k,j).
    """

    def __init__(self, pair_dim, hidden_dim):
        super().__init__()
        self.layer_norm_in = nn.LayerNorm(pair_dim)
        self.proj_a = nn.Linear(pair_dim, hidden_dim)
        self.gate_a = nn.Linear(pair_dim, hidden_dim)
        self.proj_b = nn.Linear(pair_dim, hidden_dim)
        self.gate_b = nn.Linear(pair_dim, hidden_dim)
        self.layer_norm_out = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, pair_dim)
        self.out_gate = nn.Linear(pair_dim, pair_dim)

    def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            pair_repr: (B, N, N, pair_dim)
            pair_mask: (B, N, N)
        Returns:
            updated pair representation (B, N, N, pair_dim)
        """
        mask = pair_mask.unsqueeze(-1)
        x = self.layer_norm_in(pair_repr)

        a = self.proj_a(x) * torch.sigmoid(self.gate_a(x)) * mask  # (B, N, N, H)
        b = self.proj_b(x) * torch.sigmoid(self.gate_b(x)) * mask

        # Incoming: aggregate over k with a[k,i] * b[k,j]
        out = torch.einsum("bkih,bkjh->bijh", a, b)

        out = self.layer_norm_out(out)
        out = self.out_proj(out)
        out = torch.sigmoid(self.out_gate(pair_repr)) * out
        return out


pair_dim = 64
hidden_dim = 32
seq_len = 32
batch_size = 2


def get_inputs():
    p = popcorn_pri
    pair_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len), pair_dim)
    pair_mask = torch.ones(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len))
    return [pair_repr, pair_mask]


def get_init_inputs():
    return [pair_dim, hidden_dim]

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
import math

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

constexpr int BLOCK_SIZE = 16;

// LayerNorm kernel
template <typename scalar_t>
__global__ void layer_norm_kernel(
    const scalar_t* __restrict__ input,
    scalar_t* __restrict__ output,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    int N, int dim, int stride) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N) return;
    
    const scalar_t* in_ptr = input + idx * stride;
    scalar_t* out_ptr = output + idx * stride;
    
    // Compute mean
    scalar_t mean = 0;
    for (int i = 0; i < dim; i++) {
        mean += in_ptr[i];
    }
    mean /= dim;
    
    // Compute variance
    scalar_t var = 0;
    for (int i = 0; i < dim; i++) {
        scalar_t diff = in_ptr[i] - mean;
        var += diff * diff;
    }
    var /= dim;
    
    scalar_t inv_std = rsqrt(var + scalar_t(1e-5));
    
    // Normalize and apply affine transform
    for (int i = 0; i < dim; i++) {
        out_ptr[i] = (in_ptr[i] - mean) * inv_std * weight[i] + bias[i];
    }
}

// Linear + sigmoid gating kernel
template <typename scalar_t>
__global__ void linear_gated_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ proj_weight,
    const scalar_t* __restrict__ proj_bias,
    const scalar_t* __restrict__ gate_weight,
    const scalar_t* __restrict__ gate_bias,
    scalar_t* __restrict__ output,
    const scalar_t* __restrict__ mask,
    int batch_size, int N, int in_dim, int out_dim) {
    
    int b = blockIdx.x;
    int ij = blockIdx.y * blockDim.x + threadIdx.x;
    int i = ij / N;
    int j = ij % N;
    
    if (i >= N || j >= N) return;
    
    int input_offset = ((b * N + i) * N + j) * in_dim;
    int output_offset = ((b * N + i) * N + j) * out_dim;
    
    scalar_t mask_val = mask ? mask[(b * N + i) * N + j] : scalar_t(1.0);
    
    for (int h = 0; h < out_dim; h++) {
        scalar_t proj_val = proj_bias[h];
        scalar_t gate_val = gate_bias[h];
        
        for (int d = 0; d < in_dim; d++) {
            scalar_t x = input[input_offset + d];
            proj_val += x * proj_weight[d * out_dim + h];
            gate_val += x * gate_weight[d * out_dim + h];
        }
        
        // Sigmoid
        gate_val = scalar_t(1.0) / (scalar_t(1.0) + exp(-gate_val));
        
        output[output_offset + h] = proj_val * gate_val * mask_val;
    }
}

// Einsum kernel: bkih,bkjh->bijh
template <typename scalar_t>
__global__ void einsum_incoming_kernel(
    const scalar_t* __restrict__ a,
    const scalar_t* __restrict__ b,
    scalar_t* __restrict__ out,
    int batch_size, int N, int H) {
    
    int b_idx = blockIdx.z;
    int h = blockIdx.y;
    int ij = blockIdx.x * blockDim.x + threadIdx.x;
    int i = ij / N;
    int j = ij % N;
    
    if (i >= N || j >= N) return;
    
    scalar_t sum = 0;
    for (int k = 0; k < N; k++) {
        int a_idx = ((b_idx * N + k) * N + i) * H + h;
        int b_idx_local = ((b_idx * N + k) * N + j) * H + h;
        sum += a[a_idx] * b[b_idx_local];
    }
    
    out[((b_idx * N + i) * N + j) * H + h] = sum;
}

// Output projection kernel
template <typename scalar_t>
__global__ void output_proj_kernel(
    const scalar_t* __restrict__ pair_repr,
    const scalar_t* __restrict__ norm_out,
    const scalar_t* __restrict__ proj_weight,
    const scalar_t* __restrict__ proj_bias,
    const scalar_t* __restrict__ gate_weight,
    const scalar_t* __restrict__ gate_bias,
    scalar_t* __restrict__ output,
    int batch_size, int N, int in_dim, int out_dim) {
    
    int b = blockIdx.x;
    int ij = blockIdx.y * blockDim.x + threadIdx.x;
    int i = ij / N;
    int j = ij % N;
    
    if (i >= N || j >= N) return;
    
    int input_offset = ((b * N + i) * N + j) * in_dim;
    int pair_offset = ((b * N + i) * N + j) * out_dim;
    int output_offset = pair_offset;
    
    // Compute projection
    for (int h = 0; h < out_dim; h++) {
        scalar_t proj_val = proj_bias[h];
        for (int d = 0; d < in_dim; d++) {
            proj_val += norm_out[input_offset + d] * proj_weight[d * out_dim + h];
        }
        
        // Compute gate
        scalar_t gate_val = gate_bias[h];
        for (int d = 0; d < out_dim; d++) {
            gate_val += pair_repr[pair_offset + d] * gate_weight[d * out_dim + h];
        }
        gate_val = scalar_t(1.0) / (scalar_t(1.0) + exp(-gate_val));
        
        output[output_offset + h] = gate_val * proj_val;
    }
}

// Host functions
torch::Tensor layer_norm_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias) {
    auto sizes = input.sizes();
    int total = 1;
    for (int i = 0; i < sizes.size() - 1; i++) total *= sizes[i];
    int dim = sizes[sizes.size() - 1];
    
    auto output = torch::empty_like(input);
    
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "layer_norm_cuda", ([&] {
        layer_norm_kernel<scalar_t><<<blocks, threads>>>(
            input.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            weight.data_ptr<scalar_t>(),
            bias.data_ptr<scalar_t>(),
            total, dim, dim);
    }));
    
    return output;
}

torch::Tensor linear_gated_cuda(
    torch::Tensor input,
    torch::Tensor proj_weight, torch::Tensor proj_bias,
    torch::Tensor gate_weight, torch::Tensor gate_bias,
    torch::Tensor mask, int out_dim) {
    
    auto sizes = input.sizes();
    int batch_size = sizes[0];
    int N = sizes[1];
    int in_dim = sizes[3];
    
    auto output = torch::zeros({batch_size, N, N, out_dim}, input.options());
    
    dim3 blocks(batch_size, (N * N + 255) / 256);
    const int threads = 256;
    
    bool has_mask = mask.defined() && mask.numel() > 0;
    
    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "linear_gated_cuda", ([&] {
        linear_gated_kernel<scalar_t><<<blocks, threads>>>(
            input.data_ptr<scalar_t>(),
            proj_weight.data_ptr<scalar_t>(),
            proj_bias.data_ptr<scalar_t>(),
            gate_weight.data_ptr<scalar_t>(),
            gate_bias.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            has_mask ? mask.data_ptr<scalar_t>() : nullptr,
            batch_size, N, in_dim, out_dim);
    }));
    
    return output;
}

torch::Tensor einsum_incoming_cuda(torch::Tensor a, torch::Tensor b) {
    auto sizes = a.sizes();
    int batch_size = sizes[0];
    int N = sizes[1];
    int H = sizes[3];
    
    auto output = torch::zeros({batch_size, N, N, H}, a.options());
    
    dim3 blocks((N * N + 255) / 256, H, batch_size);
    const int threads = 256;
    
    AT_DISPATCH_FLOATING_TYPES(a.scalar_type(), "einsum_incoming_cuda", ([&] {
        einsum_incoming_kernel<scalar_t><<<blocks, threads>>>(
            a.data_ptr<scalar_t>(),
            b.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            batch_size, N, H);
    }));
    
    return output;
}

torch::Tensor output_proj_cuda(
    torch::Tensor pair_repr,
    torch::Tensor norm_out,
    torch::Tensor proj_weight, torch::Tensor proj_bias,
    torch::Tensor gate_weight, torch::Tensor gate_bias,
    int out_dim) {
    
    auto sizes = norm_out.sizes();
    int batch_size = sizes[0];
    int N = sizes[1];
    int in_dim = sizes[3];
    
    auto output = torch::zeros({batch_size, N, N, out_dim}, norm_out.options());
    
    dim3 blocks(batch_size, (N * N + 255) / 256);
    const int threads = 256;
    
    AT_DISPATCH_FLOATING_TYPES(norm_out.scalar_type(), "output_proj_cuda", ([&] {
        output_proj_kernel<scalar_t><<<blocks, threads>>>(
            pair_repr.data_ptr<scalar_t>(),
            norm_out.data_ptr<scalar_t>(),
            proj_weight.data_ptr<scalar_t>(),
            proj_bias.data_ptr<scalar_t>(),
            gate_weight.data_ptr<scalar_t>(),
            gate_bias.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            batch_size, N, in_dim, out_dim);
    }));
    
    return output;
}
"""

cpp_source = """
torch::Tensor layer_norm_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias);
torch::Tensor linear_gated_cuda(torch::Tensor input, torch::Tensor proj_weight, torch::Tensor proj_bias, torch::Tensor gate_weight, torch::Tensor gate_bias, torch::Tensor mask, int out_dim);
torch::Tensor einsum_incoming_cuda(torch::Tensor a, torch::Tensor b);
torch::Tensor output_proj_cuda(torch::Tensor pair_repr, torch::Tensor norm_out, torch::Tensor proj_weight, torch::Tensor proj_bias, torch::Tensor gate_weight, torch::Tensor gate_bias, int out_dim);
"""

module = load_inline(
    name='tri_mul_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['layer_norm_cuda', 'linear_gated_cuda', 'einsum_incoming_cuda', 'output_proj_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self, pair_dim, hidden_dim):
        super().__init__()
        self.layer_norm_in = nn.LayerNorm(pair_dim)
        self.proj_a = nn.Linear(pair_dim, hidden_dim)
        self.gate_a = nn.Linear(pair_dim, hidden_dim)
        self.proj_b = nn.Linear(pair_dim, hidden_dim)
        self.gate_b = nn.Linear(pair_dim, hidden_dim)
        self.layer_norm_out = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, pair_dim)
        self.out_gate = nn.Linear(pair_dim, pair_dim)
        
        self.pair_dim = pair_dim
        self.hidden_dim = hidden_dim

    def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
        pair_repr = pair_repr.contiguous()
        pair_mask = pair_mask.contiguous()
        
        # LayerNorm input
        x = module.layer_norm_cuda(
            pair_repr,
            self.layer_norm_in.weight,
            self.layer_norm_in.bias
        )
        
        # Linear + gated for a
        a = module.linear_gated_cuda(
            x,
            self.proj_a.weight.transpose(0, 1).contiguous(),
            self.proj_a.bias,
            self.gate_a.weight.transpose(0, 1).contiguous(),
            self.gate_a.bias,
            pair_mask,
            self.hidden_dim
        )
        
        # Linear + gated for b
        b = module.linear_gated_cuda(
            x,
            self.proj_b.weight.transpose(0, 1).contiguous(),
            self.proj_b.bias,
            self.gate_b.weight.transpose(0, 1).contiguous(),
            self.gate_b.bias,
            pair_mask,
            self.hidden_dim
        )
        
        # Einsum incoming
        out = module.einsum_incoming_cuda(a, b)
        
        # LayerNorm output
        out = module.layer_norm_cuda(
            out,
            self.layer_norm_out.weight,
            self.layer_norm_out.bias
        )
        
        # Output projection with gate
        result = module.output_proj_cuda(
            pair_repr,
            out,
            self.out_proj.weight.transpose(0, 1).contiguous(),
            self.out_proj.bias,
            self.out_gate.weight.transpose(0, 1).contiguous(),
            self.out_gate.bias,
            self.pair_dim
        )
        
        return result