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 triangular_forward(
    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 <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math.h>

static __device__ __forceinline__ float sigmoidf_dev(float x) {
    return 1.0f / (1.0f + expf(-x));
}

// One block per (B,N,N) row. Computes input LayerNorm and all a/b projections.
__global__ void compute_ab_kernel(
    const float* __restrict__ x0, 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 rows, int C, int H)
{
    int row = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float red[128];
    __shared__ float xn[64];
    __shared__ float vals[128];

    float v = 0.0f;
    if (tid < C) v = x0[row * C + tid];
    red[tid] = (tid < C) ? v : 0.0f;
    __syncthreads();
    for (int s = 64; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid+s]; __syncthreads(); }
    float mean = red[0] * (1.0f / (float)C);

    float d = 0.0f;
    if (tid < C) { float t = v - mean; d = t * t; }
    red[tid] = (tid < C) ? d : 0.0f;
    __syncthreads();
    for (int s = 64; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid+s]; __syncthreads(); }
    float rstd = rsqrtf(red[0] * (1.0f / (float)C) + 1.0e-5f);

    if (tid < C) xn[tid] = (v - mean) * rstd * ln_w[tid] + ln_b[tid];
    __syncthreads();

    int h = tid & 31;
    int which = tid >> 5;
    if (h < H) {
        const float* w;
        const float* bias;
        if (which == 0) { w = pa_w; bias = pa_b; }
        else if (which == 1) { w = ga_w; bias = ga_b; }
        else if (which == 2) { w = pb_w; bias = pb_b; }
        else { w = gb_w; bias = gb_b; }
        float sum = bias[h];
        const float* wh = w + h * C;
        #pragma unroll
        for (int c = 0; c < 64; ++c) sum += wh[c] * xn[c];
        vals[tid] = sum;
    }
    __syncthreads();

    if (tid < H) {
        float m = mask[row];
        float av = vals[tid] * sigmoidf_dev(vals[32 + tid]) * m;
        float bv = vals[64 + tid] * sigmoidf_dev(vals[96 + tid]) * m;
        a[row * H + tid] = av;
        b[row * H + tid] = bv;
    }
}

// One thread per (b,i,j,h); sum over k.
__global__ void einsum_kernel(const float* __restrict__ a, const float* __restrict__ b,
                              float* __restrict__ out, int B, int N, int H) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N * N * H;
    if (idx >= total) return;
    int h = idx % H;
    int t = idx / H;
    int j = t % N;
    t /= N;
    int i = t % N;
    int bb = t / N;
    float sum = 0.0f;
    int base_a = ((bb * N + i) * N) * H + h;
    int base_b = ((bb * N + j) * N) * H + h;
    #pragma unroll
    for (int k = 0; k < 32; ++k) {
        sum += a[base_a + k * H] * b[base_b + k * H];
    }
    out[idx] = sum;
}

// One block per output row. LayerNorm over H, out projection, out gate, multiply.
__global__ void final_kernel(
    const float* __restrict__ mid, const float* __restrict__ pair_repr,
    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 rows, int C, int H)
{
    int row = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float red[128];
    __shared__ float yn[32];

    float v = 0.0f;
    if (tid < H) v = mid[row * H + tid];
    red[tid] = (tid < H) ? v : 0.0f;
    __syncthreads();
    for (int s = 64; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid+s]; __syncthreads(); }
    float mean = red[0] * (1.0f / (float)H);

    float d = 0.0f;
    if (tid < H) { float t = v - mean; d = t * t; }
    red[tid] = (tid < H) ? d : 0.0f;
    __syncthreads();
    for (int s = 64; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid+s]; __syncthreads(); }
    float rstd = rsqrtf(red[0] * (1.0f / (float)H) + 1.0e-5f);
    if (tid < H) yn[tid] = (v - mean) * rstd * ln_w[tid] + ln_b[tid];
    __syncthreads();

    if (tid < C) {
        int c = tid;
        float pv = op_b[c];
        const float* ow = op_w + c * H;
        #pragma unroll
        for (int h = 0; h < 32; ++h) pv += ow[h] * yn[h];

        float gv = og_b[c];
        const float* gw = og_w + c * C;
        const float* xr = pair_repr + row * C;
        #pragma unroll
        for (int k = 0; k < 64; ++k) gv += gw[k] * xr[k];
        out[row * C + c] = sigmoidf_dev(gv) * pv;
    }
}

torch::Tensor triangular_forward(
    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)
{
    const int B = pair_repr.size(0);
    const int N = pair_repr.size(1);
    const int C = pair_repr.size(3);
    const int H = proj_a_w.size(0);
    const int rows = B * N * N;
    auto optsH = pair_repr.options();
    auto a = torch::empty({B, N, N, H}, optsH);
    auto b = torch::empty({B, N, N, H}, optsH);
    auto mid = torch::empty({B, N, N, H}, optsH);
    auto out = torch::empty_like(pair_repr);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    compute_ab_kernel<<<rows, 128, 0, stream>>>(
        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>(), rows, C, H);

    int total = B * N * N * H;
    einsum_kernel<<<(total + 255) / 256, 256, 0, stream>>>(a.data_ptr<float>(), b.data_ptr<float>(), mid.data_ptr<float>(), B, N, H);

    final_kernel<<<rows, 128, 0, stream>>>(
        mid.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>(),
        out.data_ptr<float>(), rows, C, H);
    return out;
}
'''

_ext = load_inline(
    name='tri_mul_outgoing_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['triangular_forward'],
    extra_cflags=['-O3'],
    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)

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