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


class Model(nn.Module):
    """
    Triangular self-attention over pair representations, as used in
    AlphaFold2's Evoformer stack.  Starting-node variant: for each
    row i of the pair matrix, standard multi-head self-attention is
    performed across the column dimension j, with an additive pair
    bias.  This lets positions (i,j) and (i,k) exchange information
    through their shared starting node i.
    """

    def __init__(self, pair_dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = pair_dim // num_heads
        assert pair_dim % num_heads == 0

        self.layer_norm = nn.LayerNorm(pair_dim)
        self.query = nn.Linear(pair_dim, pair_dim, bias=False)
        self.key = nn.Linear(pair_dim, pair_dim, bias=False)
        self.value = nn.Linear(pair_dim, pair_dim, bias=False)
        self.gate = nn.Linear(pair_dim, pair_dim)
        self.bias_proj = nn.Linear(pair_dim, num_heads, bias=False)
        self.out_proj = nn.Linear(pair_dim, pair_dim)
        self.dropout = nn.Dropout(dropout)

    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)
        """
        B, N, _, D = pair_repr.shape
        h = self.num_heads
        d = self.head_dim

        x = self.layer_norm(pair_repr)
        g = torch.sigmoid(self.gate(x))

        # Pair bias from the pair representation itself
        bias = self.bias_proj(x).permute(0, 3, 1, 2)  # (B, h, N, N)

        # Reshape to (B*N, N, D) — attention across columns for each row
        x_flat = x.reshape(B * N, N, D)
        q = self.query(x_flat).view(B * N, N, h, d).transpose(1, 2)  # (B*N, h, N, d)
        k = self.key(x_flat).view(B * N, N, h, d).transpose(1, 2)
        v = self.value(x_flat).view(B * N, N, h, d).transpose(1, 2)

        attn = (q @ k.transpose(-1, -2)) / math.sqrt(d)  # (B*N, h, N, N)
        attn = attn.view(B, N, h, N, N)

        # Add pair bias: bias (B, h, N, N) → (B, 1, h, N, N) broadcast across rows
        attn = attn + bias.unsqueeze(1)

        # Apply mask: pair_mask (B, N, N) → mask over the key dimension
        mask = pair_mask.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, N, N) – last dim is keys
        attn = attn.masked_fill(mask == 0, float("-inf"))
        attn = attn.view(B * N, h, N, N)
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        out = (attn @ v).transpose(1, 2).reshape(B, N, N, D)
        out = g * out
        return self.out_proj(out)


pair_dim = 64
num_heads = 4
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, num_heads]

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(torch::Tensor pair_repr, torch::Tensor pair_mask,
                          torch::Tensor ln_w, torch::Tensor ln_b,
                          torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
                          torch::Tensor gate_w, torch::Tensor gate_b,
                          torch::Tensor bias_w,
                          torch::Tensor out_w, torch::Tensor out_b,
                          int64_t num_heads);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math_constants.h>

#ifndef MAX_D
#define MAX_D 128
#endif
#ifndef MAX_N
#define MAX_N 64
#endif

__device__ __forceinline__ float warp_safe_exp(float x) { return expf(x); }

__global__ void norm_proj_kernel(const float* __restrict__ inp,
                                 const float* __restrict__ ln_w,
                                 const float* __restrict__ ln_b,
                                 const float* __restrict__ q_w,
                                 const float* __restrict__ k_w,
                                 const float* __restrict__ v_w,
                                 const float* __restrict__ gate_w,
                                 const float* __restrict__ gate_b,
                                 const float* __restrict__ bias_w,
                                 float* __restrict__ q,
                                 float* __restrict__ k,
                                 float* __restrict__ v,
                                 float* __restrict__ gate,
                                 float* __restrict__ bias,
                                 int B, int N, int D, int H) {
    int token = blockIdx.x;
    int tid = threadIdx.x;
    int total_tokens = B * N * N;
    if (token >= total_tokens) return;

    __shared__ float red[256];
    __shared__ float x[MAX_D];

    const float* base = inp + (long long)token * D;
    float s = 0.f;
    for (int c = tid; c < D; c += blockDim.x) s += base[c];
    red[tid] = s;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float mean = red[0] / (float)D;

    s = 0.f;
    for (int c = tid; c < D; c += blockDim.x) {
        float d = base[c] - mean;
        s += d * d;
    }
    red[tid] = s;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float invstd = rsqrtf(red[0] / (float)D + 1.0e-5f);

    for (int c = tid; c < D; c += blockDim.x) {
        x[c] = (base[c] - mean) * invstd * ln_w[c] + ln_b[c];
    }
    __syncthreads();

    int total_out = 4 * D + H;
    for (int o = tid; o < total_out; o += blockDim.x) {
        const float* wptr;
        int row;
        int kind;
        if (o < D) { kind = 0; row = o; wptr = q_w + (long long)row * D; }
        else if (o < 2 * D) { kind = 1; row = o - D; wptr = k_w + (long long)row * D; }
        else if (o < 3 * D) { kind = 2; row = o - 2 * D; wptr = v_w + (long long)row * D; }
        else if (o < 4 * D) { kind = 3; row = o - 3 * D; wptr = gate_w + (long long)row * D; }
        else { kind = 4; row = o - 4 * D; wptr = bias_w + (long long)row * D; }
        float acc = 0.f;
        #pragma unroll
        for (int c = 0; c < MAX_D; ++c) {
            if (c < D) acc = fmaf(x[c], wptr[c], acc);
        }
        if (kind == 0) q[(long long)token * D + row] = acc;
        else if (kind == 1) k[(long long)token * D + row] = acc;
        else if (kind == 2) v[(long long)token * D + row] = acc;
        else if (kind == 3) {
            acc += gate_b[row];
            gate[(long long)token * D + row] = 1.f / (1.f + expf(-acc));
        } else {
            int j = token % N;
            int tmp = token / N;
            int i = tmp % N;
            int b = tmp / N;
            bias[((long long)b * H + row) * N * N + i * N + j] = acc;
        }
    }
}

__global__ void attention_kernel(const float* __restrict__ q,
                                 const float* __restrict__ k,
                                 const float* __restrict__ v,
                                 const float* __restrict__ gate,
                                 const float* __restrict__ bias,
                                 const float* __restrict__ mask,
                                 float* __restrict__ mid,
                                 int B, int N, int D, int H, int Hd) {
    int idx = blockIdx.x;
    int tid = threadIdx.x;
    int h = idx % H;
    int t = idx / H;
    int j = t % N;
    t /= N;
    int row = t % N;
    int b = t / N;

    __shared__ float scores[MAX_N];
    __shared__ float red[64];

    long long qtok = ((long long)b * N + row) * N + j;
    float scale = rsqrtf((float)Hd);

    for (int kk = tid; kk < N; kk += blockDim.x) {
        long long ktok = ((long long)b * N + row) * N + kk;
        float dot = 0.f;
        #pragma unroll
        for (int dd = 0; dd < 32; ++dd) {
            if (dd < Hd) dot = fmaf(q[qtok * D + h * Hd + dd], k[ktok * D + h * Hd + dd], dot);
        }
        float m = mask[((long long)b * N + j) * N + kk];
        float sc = dot * scale + bias[((long long)b * H + h) * N * N + j * N + kk];
        scores[kk] = (m == 0.f) ? -CUDART_INF_F : sc;
    }
    __syncthreads();

    float lmax = -CUDART_INF_F;
    for (int kk = tid; kk < N; kk += blockDim.x) lmax = fmaxf(lmax, scores[kk]);
    red[tid] = lmax;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] = fmaxf(red[tid], red[tid + stride]);
        __syncthreads();
    }
    float mx = red[0];

    float lsum = 0.f;
    for (int kk = tid; kk < N; kk += blockDim.x) {
        float e = expf(scores[kk] - mx);
        scores[kk] = e;
        lsum += e;
    }
    red[tid] = lsum;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float invsum = 1.f / red[0];
    for (int kk = tid; kk < N; kk += blockDim.x) scores[kk] *= invsum;
    __syncthreads();

    for (int dd = tid; dd < Hd; dd += blockDim.x) {
        float acc = 0.f;
        for (int kk = 0; kk < N; ++kk) {
            long long ktok = ((long long)b * N + row) * N + kk;
            acc = fmaf(scores[kk], v[ktok * D + h * Hd + dd], acc);
        }
        int dim = h * Hd + dd;
        mid[qtok * D + dim] = acc * gate[qtok * D + dim];
    }
}

__global__ void out_proj_kernel(const float* __restrict__ mid,
                                const float* __restrict__ out_w,
                                const float* __restrict__ out_b,
                                float* __restrict__ out,
                                int total_tokens, int D) {
    int token = blockIdx.x;
    int tid = threadIdx.x;
    if (token >= total_tokens) return;
    for (int oc = tid; oc < D; oc += blockDim.x) {
        const float* w = out_w + (long long)oc * D;
        const float* x = mid + (long long)token * D;
        float acc = out_b[oc];
        #pragma unroll
        for (int c = 0; c < MAX_D; ++c) {
            if (c < D) acc = fmaf(x[c], w[c], acc);
        }
        out[(long long)token * D + oc] = acc;
    }
}

torch::Tensor tri_forward(torch::Tensor pair_repr, torch::Tensor pair_mask,
                          torch::Tensor ln_w, torch::Tensor ln_b,
                          torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
                          torch::Tensor gate_w, torch::Tensor gate_b,
                          torch::Tensor bias_w,
                          torch::Tensor out_w, torch::Tensor out_b,
                          int64_t num_heads) {
    int B = (int)pair_repr.size(0);
    int N = (int)pair_repr.size(1);
    int D = (int)pair_repr.size(3);
    int H = (int)num_heads;
    int Hd = D / H;
    auto opts = pair_repr.options();
    auto q = torch::empty({B, N, N, D}, opts);
    auto k = torch::empty({B, N, N, D}, opts);
    auto v = torch::empty({B, N, N, D}, opts);
    auto gate = torch::empty({B, N, N, D}, opts);
    auto bias = torch::empty({B, H, N, N}, opts);
    auto mid = torch::empty({B, N, N, D}, opts);
    auto out = torch::empty({B, N, N, D}, opts);

    int total_tokens = B * N * N;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    norm_proj_kernel<<<total_tokens, 256, 0, stream>>>(
        pair_repr.data_ptr<float>(), ln_w.data_ptr<float>(), ln_b.data_ptr<float>(),
        q_w.data_ptr<float>(), k_w.data_ptr<float>(), v_w.data_ptr<float>(),
        gate_w.data_ptr<float>(), gate_b.data_ptr<float>(), bias_w.data_ptr<float>(),
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), gate.data_ptr<float>(), bias.data_ptr<float>(),
        B, N, D, H);

    attention_kernel<<<total_tokens * H, 64, 0, stream>>>(
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), gate.data_ptr<float>(),
        bias.data_ptr<float>(), pair_mask.data_ptr<float>(), mid.data_ptr<float>(),
        B, N, D, H, Hd);

    out_proj_kernel<<<total_tokens, 128, 0, stream>>>(
        mid.data_ptr<float>(), out_w.data_ptr<float>(), out_b.data_ptr<float>(), out.data_ptr<float>(),
        total_tokens, D);
    return out;
}
"""

_ext = load_inline(
    name="tri_start_attn_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["tri_forward"],
    extra_cuda_cflags=["-O3"],
    extra_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, pair_dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = pair_dim // num_heads
        assert pair_dim % num_heads == 0
        self.layer_norm = nn.LayerNorm(pair_dim)
        self.query = nn.Linear(pair_dim, pair_dim, bias=False)
        self.key = nn.Linear(pair_dim, pair_dim, bias=False)
        self.value = nn.Linear(pair_dim, pair_dim, bias=False)
        self.gate = nn.Linear(pair_dim, pair_dim)
        self.bias_proj = nn.Linear(pair_dim, num_heads, bias=False)
        self.out_proj = nn.Linear(pair_dim, pair_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
        return _ext.tri_forward(pair_repr, pair_mask,
                                self.layer_norm.weight, self.layer_norm.bias,
                                self.query.weight, self.key.weight, self.value.weight,
                                self.gate.weight, self.gate.bias,
                                self.bias_proj.weight,
                                self.out_proj.weight, self.out_proj.bias,
                                int(self.num_heads))