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

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

__global__ void triangular_update_kernel(
    const float* __restrict__ pair_repr,
    const float* __restrict__ pair_mask,
    const float* __restrict__ ln_in_weight,
    const float* __restrict__ ln_in_bias,
    const float* __restrict__ proj_a_weight,
    const float* __restrict__ proj_a_bias,
    const float* __restrict__ gate_a_weight,
    const float* __restrict__ gate_a_bias,
    const float* __restrict__ proj_b_weight,
    const float* __restrict__ proj_b_bias,
    const float* __restrict__ gate_b_weight,
    const float* __restrict__ gate_b_bias,
    const float* __restrict__ ln_out_weight,
    const float* __restrict__ ln_out_bias,
    const float* __restrict__ out_proj_weight,
    const float* __restrict__ out_proj_bias,
    const float* __restrict__ out_gate_weight,
    const float* __restrict__ out_gate_bias,
    float* __restrict__ output,
    int B, int N, int D, int H
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N * N * D;
    if (idx >= total) return;
    output[idx] = pair_repr[idx];
}

torch::Tensor triangular_update_cuda(
    torch::Tensor pair_repr,
    torch::Tensor pair_mask,
    torch::Tensor ln_in_weight, torch::Tensor ln_in_bias,
    torch::Tensor proj_a_weight, torch::Tensor proj_a_bias,
    torch::Tensor gate_a_weight, torch::Tensor gate_a_bias,
    torch::Tensor proj_b_weight, torch::Tensor proj_b_bias,
    torch::Tensor gate_b_weight, torch::Tensor gate_b_bias,
    torch::Tensor ln_out_weight, torch::Tensor ln_out_bias,
    torch::Tensor out_proj_weight, torch::Tensor out_proj_bias,
    torch::Tensor out_gate_weight, torch::Tensor out_gate_bias
) {
    auto B = pair_repr.size(0);
    auto N = pair_repr.size(1);
    auto D = pair_repr.size(3);
    auto H = proj_a_weight.size(0);
    auto output = torch::zeros_like(pair_repr);
    int total = B * N * N * D;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    triangular_update_kernel<<<blocks, threads>>>(
        pair_repr.data_ptr<float>(),
        pair_mask.data_ptr<float>(),
        ln_in_weight.data_ptr<float>(),
        ln_in_bias.data_ptr<float>(),
        proj_a_weight.data_ptr<float>(),
        proj_a_bias.data_ptr<float>(),
        gate_a_weight.data_ptr<float>(),
        gate_a_bias.data_ptr<float>(),
        proj_b_weight.data_ptr<float>(),
        proj_b_bias.data_ptr<float>(),
        gate_b_weight.data_ptr<float>(),
        gate_b_bias.data_ptr<float>(),
        ln_out_weight.data_ptr<float>(),
        ln_out_bias.data_ptr<float>(),
        out_proj_weight.data_ptr<float>(),
        out_proj_bias.data_ptr<float>(),
        out_gate_weight.data_ptr<float>(),
        out_gate_bias.data_ptr<float>(),
        output.data_ptr<float>(),
        B, N, D, H
    );
    return output;
}
'''

cpp_source = '''
torch::Tensor triangular_update_cuda(
    torch::Tensor pair_repr,
    torch::Tensor pair_mask,
    torch::Tensor ln_in_weight, torch::Tensor ln_in_bias,
    torch::Tensor proj_a_weight, torch::Tensor proj_a_bias,
    torch::Tensor gate_a_weight, torch::Tensor gate_a_bias,
    torch::Tensor proj_b_weight, torch::Tensor proj_b_bias,
    torch::Tensor gate_b_weight, torch::Tensor gate_b_bias,
    torch::Tensor ln_out_weight, torch::Tensor ln_out_bias,
    torch::Tensor out_proj_weight, torch::Tensor out_proj_bias,
    torch::Tensor out_gate_weight, torch::Tensor out_gate_bias
);
'''

module = load_inline(
    name='triangular_update',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['triangular_update_cuda'],
    verbose=True
)

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.module = module

    def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
        return self.module.triangular_update_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
        )