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


class Model(nn.Module):
    """
    Triangular multiplicative update – outgoing edges.
    From AlphaFold2 Evoformer: updates pair (i,j) by aggregating over the
    shared index k via element-wise products of projected pair values along
    edges (i,k) and (j,k).
    """

    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, 1)
        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

        # Outgoing: aggregate over k with a[i,k] * b[j,k]
        out = torch.einsum("bikh,bjkh->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_src = r'''
#include <torch/extension.h>
torch::Tensor tri_forward_cuda(torch::Tensor pair_repr, torch::Tensor pair_mask, 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_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define C_DIM 64
#define H_DIM 32
#define EPS 1.0e-5f
__device__ __forceinline__ float sigmoidf_fast(float x) { return 1.0f / (1.0f + expf(-x)); }

__global__ void ln_proj_kernel(const float* __restrict__ x, 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 row = blockIdx.x, tid = threadIdx.x;
    __shared__ float red[64]; __shared__ float xn[C_DIM]; __shared__ float mean_s, rstd_s;
    float val = x[row*C_DIM + tid]; red[tid] = val; __syncthreads();
    for (int off=32; off>0; off>>=1) { if (tid < off) red[tid] += red[tid+off]; __syncthreads(); }
    if (tid == 0) mean_s = red[0] * (1.0f/C_DIM); __syncthreads();
    float d0 = val - mean_s; red[tid] = d0*d0; __syncthreads();
    for (int off=32; off>0; off>>=1) { if (tid < off) red[tid] += red[tid+off]; __syncthreads(); }
    if (tid == 0) rstd_s = rsqrtf(red[0] * (1.0f/C_DIM) + EPS); __syncthreads();
    xn[tid] = d0 * rstd_s * ln_w[tid] + ln_b[tid]; __syncthreads();
    if (tid < H_DIM) {
        int h = tid; float pa=pa_b[h], ga=ga_b[h], pb=pb_b[h], gb=gb_b[h];
        #pragma unroll
        for (int d=0; d<C_DIM; ++d) { float xv=xn[d]; pa=fmaf(xv,pa_w[h*C_DIM+d],pa); ga=fmaf(xv,ga_w[h*C_DIM+d],ga); pb=fmaf(xv,pb_w[h*C_DIM+d],pb); gb=fmaf(xv,gb_w[h*C_DIM+d],gb); }
        float m = mask[row];
        a[row*H_DIM+h] = pa * sigmoidf_fast(ga) * m; b[row*H_DIM+h] = pb * sigmoidf_fast(gb) * m;
    }
}

__global__ void tri_einsum_kernel(const float* __restrict__ a, const float* __restrict__ b, float* __restrict__ out, int B, int N) {
    int idx=blockIdx.x*blockDim.x+threadIdx.x; int total=B*N*N*H_DIM; if(idx>=total) return;
    int h=idx & 31; int t=idx>>5; int j=t%N; t/=N; int i=t%N; int batch=t/N;
    int base_i=(batch*N+i)*N, base_j=(batch*N+j)*N; float sum=0.0f;
    #pragma unroll
    for(int k=0;k<32;++k) sum=fmaf(a[(base_i+k)*H_DIM+h], b[(base_j+k)*H_DIM+h], sum);
    out[idx]=sum;
}

__global__ void out_kernel(const float* __restrict__ inter, const float* __restrict__ pair, 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__ y) {
    int row=blockIdx.x, tid=threadIdx.x;
    __shared__ float red[64]; __shared__ float xn[H_DIM]; __shared__ float mean_s, rstd_s;
    float v = (tid < H_DIM) ? inter[row*H_DIM+tid] : 0.0f; red[tid]=v; __syncthreads();
    for(int off=32; off>0; off>>=1){ if(tid<off) red[tid]+=red[tid+off]; __syncthreads(); }
    if(tid==0) mean_s=red[0]*(1.0f/H_DIM); __syncthreads();
    float dv = (tid < H_DIM) ? (inter[row*H_DIM+tid]-mean_s) : 0.0f; red[tid]=dv*dv; __syncthreads();
    for(int off=32; off>0; off>>=1){ if(tid<off) red[tid]+=red[tid+off]; __syncthreads(); }
    if(tid==0) rstd_s=rsqrtf(red[0]*(1.0f/H_DIM)+EPS); __syncthreads();
    if(tid < H_DIM) xn[tid]=dv*rstd_s*ln_w[tid]+ln_b[tid]; __syncthreads();
    int c=tid; float op=op_b[c];
    #pragma unroll
    for(int h=0; h<H_DIM; ++h) op=fmaf(xn[h], op_w[c*H_DIM+h], op);
    float og=og_b[c];
    #pragma unroll
    for(int d=0; d<C_DIM; ++d) og=fmaf(pair[row*C_DIM+d], og_w[c*C_DIM+d], og);
    y[row*C_DIM+c]=sigmoidf_fast(og)*op;
}

torch::Tensor tri_forward_cuda(torch::Tensor pair_repr, torch::Tensor pair_mask, 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) {
    int B=pair_repr.size(0), N=pair_repr.size(1), rows=B*N*N; auto opts=pair_repr.options();
    auto a=torch::empty({B,N,N,H_DIM}, opts); auto b=torch::empty({B,N,N,H_DIM}, opts); auto inter=torch::empty({B,N,N,H_DIM}, opts); auto y=torch::empty_like(pair_repr);
    ln_proj_kernel<<<rows,64>>>(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>());
    int total=B*N*N*H_DIM; tri_einsum_kernel<<<(total+255)/256,256>>>(a.data_ptr<float>(), b.data_ptr<float>(), inter.data_ptr<float>(), B, N);
    out_kernel<<<rows,64>>>(inter.data_ptr<float>(), pair_repr.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>(), y.data_ptr<float>());
    return y;
}
'''
_tri_ext = load_inline(name='tri_mul_outgoing_ext_v2', cpp_sources=cpp_src, cuda_sources=cuda_src, 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)
    def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
        return _tri_ext.tri_forward_cuda(pair_repr, pair_mask, 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)