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
from torch.utils.cpp_extension import load_inline

cpp_sources = r'''
#include <torch/extension.h>
torch::Tensor tri_forward_cuda(
    torch::Tensor pair_repr, torch::Tensor pair_mask, torch::Tensor a, torch::Tensor b, torch::Tensor out,
    torch::Tensor ln_in_w, torch::Tensor ln_in_b,
    torch::Tensor proj_a_w, torch::Tensor proj_a_b,
    torch::Tensor gate_a_w, torch::Tensor gate_a_b,
    torch::Tensor proj_b_w, torch::Tensor proj_b_b,
    torch::Tensor gate_b_w, torch::Tensor gate_b_b,
    torch::Tensor ln_out_w, torch::Tensor ln_out_b,
    torch::Tensor out_proj_w, torch::Tensor out_proj_b,
    torch::Tensor out_gate_w, torch::Tensor out_gate_b);
'''

cuda_sources = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#define C 64
#define H 32
#define N 32
#define BATCH 2
#define ROWS (BATCH*N*N)
__device__ __forceinline__ float warp_sum(float v){v+=__shfl_down_sync(0xffffffff,v,16);v+=__shfl_down_sync(0xffffffff,v,8);v+=__shfl_down_sync(0xffffffff,v,4);v+=__shfl_down_sync(0xffffffff,v,2);v+=__shfl_down_sync(0xffffffff,v,1);return v;}
__device__ __forceinline__ float sigm(float x){return 1.0f/(1.0f+expf(-x));}
__global__ void compute_ab_warp_kernel(const float* __restrict__ x_in,const float* __restrict__ mask,const float* __restrict__ ln_w,const float* __restrict__ ln_b,const float* __restrict__ pa_w,const float* __restrict__ pa_b,const float* __restrict__ ga_w,const float* __restrict__ ga_b,const float* __restrict__ pb_w,const float* __restrict__ pb_b,const float* __restrict__ gb_w,const float* __restrict__ gb_b,float* __restrict__ a,float* __restrict__ b){
    int lane=threadIdx.x&31; int wid=threadIdx.x>>5; int row=blockIdx.x*4+wid; if(row>=ROWS)return;
    __shared__ float svals[4][C]; const float* xin=x_in+row*C;
    float x0=xin[lane], x1=xin[lane+32]; float sum=warp_sum(x0+x1); float mean=__shfl_sync(0xffffffff,sum,0)*(1.0f/C);
    float d0=x0-mean,d1=x1-mean; float vsum=warp_sum(d0*d0+d1*d1); float inv=rsqrtf(__shfl_sync(0xffffffff,vsum,0)*(1.0f/C)+1.0e-5f);
    svals[wid][lane]=d0*inv*ln_w[lane]+ln_b[lane]; svals[wid][lane+32]=d1*inv*ln_w[lane+32]+ln_b[lane+32]; __syncwarp();
    int h=lane; float pa=pa_b[h],ga=ga_b[h],pb=pb_b[h],gb=gb_b[h]; const float*wpa=pa_w+h*C; const float*wga=ga_w+h*C; const float*wpb=pb_w+h*C; const float*wgb=gb_w+h*C;
    #pragma unroll
    for(int c=0;c<C;++c){float xv=svals[wid][c]; pa+=xv*wpa[c]; ga+=xv*wga[c]; pb+=xv*wpb[c]; gb+=xv*wgb[c];}
    float m=mask[row]; a[row*H+h]=pa*sigm(ga)*m; b[row*H+h]=pb*sigm(gb)*m;
}
__global__ void einsum_final_warp_kernel(const float* __restrict__ pair_repr,const float* __restrict__ a,const float* __restrict__ b,const float* __restrict__ ln_w,const float* __restrict__ ln_b,const float* __restrict__ op_w,const float* __restrict__ op_b,const float* __restrict__ og_w,const float* __restrict__ og_b,float* __restrict__ out){
    int lane=threadIdx.x&31; int wid=threadIdx.x>>5; int row=blockIdx.x*4+wid; if(row>=ROWS)return; int j=row&31; int i=(row>>5)&31; int bb=row>>10; int base=bb*N*N*H; int h=lane; float acc=0.0f;
    #pragma unroll
    for(int k=0;k<N;++k){acc += a[base+((k*N+i)*H+h)] * b[base+((k*N+j)*H+h)];}
    float sum=warp_sum(acc); float mean=__shfl_sync(0xffffffff,sum,0)*(1.0f/H); float d=acc-mean; float vsum=warp_sum(d*d); float inv=rsqrtf(__shfl_sync(0xffffffff,vsum,0)*(1.0f/H)+1.0e-5f);
    __shared__ float nval[4][H]; nval[wid][lane]=d*inv*ln_w[lane]+ln_b[lane]; __syncwarp(); const float* xin=pair_repr+row*C;
    #pragma unroll
    for(int rep=0;rep<2;++rep){int c=lane+rep*32; float v=op_b[c]; const float*wop=op_w+c*H;
        #pragma unroll
        for(int hh=0;hh<H;++hh) v += nval[wid][hh]*wop[hh];
        float g=og_b[c]; const float*wog=og_w+c*C;
        #pragma unroll
        for(int q=0;q<C;++q) g += xin[q]*wog[q];
        out[row*C+c]=sigm(g)*v;}
}
torch::Tensor tri_forward_cuda(torch::Tensor pair_repr,torch::Tensor pair_mask,torch::Tensor a,torch::Tensor b,torch::Tensor out,torch::Tensor ln_in_w,torch::Tensor ln_in_b,torch::Tensor proj_a_w,torch::Tensor proj_a_b,torch::Tensor gate_a_w,torch::Tensor gate_a_b,torch::Tensor proj_b_w,torch::Tensor proj_b_b,torch::Tensor gate_b_w,torch::Tensor gate_b_b,torch::Tensor ln_out_w,torch::Tensor ln_out_b,torch::Tensor out_proj_w,torch::Tensor out_proj_b,torch::Tensor out_gate_w,torch::Tensor out_gate_b){
    dim3 block(128); dim3 grid((ROWS+3)/4);
    compute_ab_warp_kernel<<<grid,block>>>(pair_repr.data_ptr<float>(),pair_mask.data_ptr<float>(),ln_in_w.data_ptr<float>(),ln_in_b.data_ptr<float>(),proj_a_w.data_ptr<float>(),proj_a_b.data_ptr<float>(),gate_a_w.data_ptr<float>(),gate_a_b.data_ptr<float>(),proj_b_w.data_ptr<float>(),proj_b_b.data_ptr<float>(),gate_b_w.data_ptr<float>(),gate_b_b.data_ptr<float>(),a.data_ptr<float>(),b.data_ptr<float>());
    einsum_final_warp_kernel<<<grid,block>>>(pair_repr.data_ptr<float>(),a.data_ptr<float>(),b.data_ptr<float>(),ln_out_w.data_ptr<float>(),ln_out_b.data_ptr<float>(),out_proj_w.data_ptr<float>(),out_proj_b.data_ptr<float>(),out_gate_w.data_ptr<float>(),out_gate_b.data_ptr<float>(),out.data_ptr<float>());
    return out;
}
'''
_ext=load_inline(name='tri_incoming_ext_final',cpp_sources=cpp_sources,cuda_sources=cuda_sources,functions=['tri_forward_cuda'],verbose=False,extra_cuda_cflags=['-O3','--use_fast_math'])
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.register_buffer('_buf_a', torch.empty((2,32,32,hidden_dim)), persistent=False); self.register_buffer('_buf_b', torch.empty((2,32,32,hidden_dim)), persistent=False); self.register_buffer('_buf_out', torch.empty((2,32,32,pair_dim)), persistent=False)
    def forward(self,pair_repr,pair_mask):
        if self._buf_a.device != pair_repr.device:
            self._buf_a = torch.empty((2,32,32,32), device=pair_repr.device, dtype=pair_repr.dtype); self._buf_b = torch.empty((2,32,32,32), device=pair_repr.device, dtype=pair_repr.dtype); self._buf_out = torch.empty((2,32,32,64), device=pair_repr.device, dtype=pair_repr.dtype)
        return _ext.tri_forward_cuda(pair_repr,pair_mask,self._buf_a,self._buf_b,self._buf_out,self.layer_norm_in.weight,self.layer_norm_in.bias,self.proj_a.weight,self.proj_a.bias,self.gate_a.weight,self.gate_a.bias,self.proj_b.weight,self.proj_b.bias,self.gate_b.weight,self.gate_b.bias,self.layer_norm_out.weight,self.layer_norm_out.bias,self.out_proj.weight,self.out_proj.bias,self.out_gate.weight,self.out_gate.bias)