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


class Model(nn.Module):
    """
    Invariant Point Attention (IPA) from AlphaFold2's structure module.
    Combines sequence-space attention with 3-D geometric queries/keys/values
    expressed in local residue frames.  The attention logits incorporate
    Euclidean distances in global frame, making it SE(3)-equivariant.
    """

    def __init__(self, node_dim, pair_dim, num_heads, num_query_points, num_value_points):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = node_dim // num_heads
        self.num_query_points = num_query_points
        self.num_value_points = num_value_points
        assert node_dim % num_heads == 0

        d = self.head_dim
        self.q_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.k_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.v_proj = nn.Linear(node_dim, num_heads * d, bias=False)

        self.q_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.k_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.v_pts = nn.Linear(node_dim, num_heads * num_value_points * 3, bias=False)

        self.pair_bias = nn.Linear(pair_dim, num_heads, bias=False)

        self.head_weights = nn.Parameter(torch.zeros(num_heads))

        out_features = num_heads * d + num_heads * num_value_points * 3
        self.out_proj = nn.Linear(out_features, node_dim)

        self.layer_norm = nn.LayerNorm(node_dim)

    def _apply_frame(self, pts, rotations, translations):
        """Transform local-frame points to global frame."""
        # pts: (B, N, P, 3), rotations: (B, N, 3, 3), translations: (B, N, 3)
        global_pts = torch.einsum("bnpc,bncd->bnpd", pts, rotations) + translations.unsqueeze(2)
        return global_pts

    def _apply_inv_frame(self, pts, rotations, translations):
        """Transform global-frame points to local frame."""
        centered = pts - translations.unsqueeze(2)
        local_pts = torch.einsum("bnpc,bndc->bnpd", centered, rotations)
        return local_pts

    def forward(
        self,
        node_repr: torch.Tensor,
        pair_repr: torch.Tensor,
        translations: torch.Tensor,
        rotations: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            node_repr:    (B, N, node_dim)
            pair_repr:    (B, N, N, pair_dim)
            translations: (B, N, 3)
            rotations:    (B, N, 3, 3)
        Returns:
            updated node representation (B, N, node_dim)
        """
        B, N, D = node_repr.shape
        h = self.num_heads
        d = self.head_dim
        Qp = self.num_query_points
        Vp = self.num_value_points

        x = self.layer_norm(node_repr)

        q = self.q_proj(x).view(B, N, h, d)
        k = self.k_proj(x).view(B, N, h, d)
        v = self.v_proj(x).view(B, N, h, d)

        # Points in local frame, then transform to global
        q_pts_local = self.q_pts(x).view(B, N, h * Qp, 3)
        k_pts_local = self.k_pts(x).view(B, N, h * Qp, 3)
        v_pts_local = self.v_pts(x).view(B, N, h * Vp, 3)

        q_pts = self._apply_frame(q_pts_local, rotations, translations).view(B, N, h, Qp, 3)
        k_pts = self._apply_frame(k_pts_local, rotations, translations).view(B, N, h, Qp, 3)
        v_pts = self._apply_frame(v_pts_local, rotations, translations).view(B, N, h, Vp, 3)

        # Scalar attention logits
        attn_scalar = torch.einsum("bihd,bjhd->bhij", q, k) / math.sqrt(d)

        # Point attention: negative squared distance between query and key points
        # q_pts: (B, N_i, h, Qp, 3), k_pts: (B, N_j, h, Qp, 3)
        q_expand = q_pts.unsqueeze(2)  # (B, N, 1, h, Qp, 3)
        k_expand = k_pts.unsqueeze(1)  # (B, 1, N, h, Qp, 3)
        pt_dist = ((q_expand - k_expand) ** 2).sum(dim=(-1, -2))  # (B, N, N, h)

        w = F.softplus(self.head_weights)
        pt_attn = -(w.view(1, 1, 1, h) * pt_dist) / 2.0
        pt_attn = pt_attn.permute(0, 3, 1, 2)  # (B, h, N, N)

        # Pair bias
        pair_bias = self.pair_bias(pair_repr).permute(0, 3, 1, 2)  # (B, h, N, N)

        attn = attn_scalar + pt_attn + pair_bias
        attn = F.softmax(attn, dim=-1)  # (B, h, N_i, N_j)

        # Aggregate scalar values
        out_scalar = torch.einsum("bhij,bjhd->bihd", attn, v).reshape(B, N, h * d)

        # Aggregate point values in global frame, then transform back to local
        out_pts = torch.einsum("bhij,bjhpc->bihpc", attn, v_pts)  # (B, N, h, Vp, 3)
        out_pts = out_pts.reshape(B, N, h * Vp, 3)
        out_pts_local = self._apply_inv_frame(out_pts, rotations, translations)
        out_pts_flat = out_pts_local.reshape(B, N, h * Vp * 3)

        out = torch.cat([out_scalar, out_pts_flat], dim=-1)
        return self.out_proj(out)


node_dim = 64
pair_dim = 32
num_heads = 4
num_query_points = 4
num_value_points = 8
seq_len = 32
batch_size = 2


def get_inputs():
    p = popcorn_pri
    node_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), node_dim)
    pair_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len), pair_dim)
    translations = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), 3)
    rotations = torch.zeros(p.jitter_int(batch_size), p.jitter_int(seq_len), 3, 3)
    for b in range(p.jitter_int(batch_size)):
        for n in range(p.jitter_int(seq_len)):
            rotations[b, n] = torch.eye(3)
    return [node_repr, pair_repr, translations, rotations]


def get_init_inputs():
    return [node_dim, pair_dim, num_heads, num_query_points, num_value_points]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor ipa_forward(
    torch::Tensor node_repr, torch::Tensor pair_repr,
    torch::Tensor translations, torch::Tensor rotations,
    torch::Tensor ln_w, torch::Tensor ln_b,
    torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
    torch::Tensor qpts_w, torch::Tensor kpts_w, torch::Tensor vpts_w,
    torch::Tensor pair_w, torch::Tensor head_w,
    torch::Tensor out_w, torch::Tensor out_b,
    int64_t num_heads, int64_t num_query_points, int64_t num_value_points);
'''

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

static __device__ __forceinline__ float softplus_dev(float x) {
    if (x > 20.0f) return x;
    if (x < -20.0f) return expf(x);
    return log1pf(expf(x));
}

__global__ void ln_proj_kernel(
    const float* __restrict__ node,
    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__ qpts_w,
    const float* __restrict__ kpts_w,
    const float* __restrict__ vpts_w,
    float* __restrict__ q,
    float* __restrict__ k,
    float* __restrict__ v,
    float* __restrict__ qpl,
    float* __restrict__ kpl,
    float* __restrict__ vpl,
    int BN, int D, int H, int QP, int VP)
{
    int row = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float red[128];
    __shared__ float xn[128];

    float s = 0.0f;
    for (int i = tid; i < D; i += blockDim.x) s += node[row * D + i];
    red[tid] = s;
    __syncthreads();
    for (int stride = blockDim.x/2; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float mean = red[0] / (float)D;

    s = 0.0f;
    for (int i = tid; i < D; i += blockDim.x) {
        float a = node[row * D + i] - mean;
        s += a * a;
    }
    red[tid] = s;
    __syncthreads();
    for (int stride = blockDim.x/2; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float inv_std = rsqrtf(red[0] / (float)D + 1.0e-5f);
    for (int i = tid; i < D; i += blockDim.x) {
        xn[i] = (node[row * D + i] - mean) * inv_std * ln_w[i] + ln_b[i];
    }
    __syncthreads();

    int dhead = D / H;
    int qkv = H * dhead;
    int qpo = H * QP * 3;
    int vpo = H * VP * 3;
    int total = 3*qkv + 2*qpo + vpo;

    for (int o = tid; o < total; o += blockDim.x) {
        const float* W;
        float* O;
        int oo, out_stride;
        if (o < qkv) { W = q_w; O = q; oo = o; out_stride = qkv; }
        else if (o < 2*qkv) { W = k_w; O = k; oo = o - qkv; out_stride = qkv; }
        else if (o < 3*qkv) { W = v_w; O = v; oo = o - 2*qkv; out_stride = qkv; }
        else if (o < 3*qkv + qpo) { W = qpts_w; O = qpl; oo = o - 3*qkv; out_stride = qpo; }
        else if (o < 3*qkv + 2*qpo) { W = kpts_w; O = kpl; oo = o - 3*qkv - qpo; out_stride = qpo; }
        else { W = vpts_w; O = vpl; oo = o - 3*qkv - 2*qpo; out_stride = vpo; }
        float acc = 0.0f;
        const float* wrow = W + oo * D;
        #pragma unroll
        for (int ii = 0; ii < 64; ++ii) {
            if (ii < D) acc += xn[ii] * wrow[ii];
        }
        O[row * out_stride + oo] = acc;
    }
}

__global__ void attn_kernel(
    const float* __restrict__ pair_repr,
    const float* __restrict__ translations,
    const float* __restrict__ rotations,
    const float* __restrict__ q,
    const float* __restrict__ k,
    const float* __restrict__ v,
    const float* __restrict__ qpl,
    const float* __restrict__ kpl,
    const float* __restrict__ vpl,
    const float* __restrict__ pair_w,
    const float* __restrict__ head_w,
    float* __restrict__ cat,
    int B, int N, int D, int PD, int H, int QP, int VP)
{
    int bid = blockIdx.x;
    int h = bid % H;
    int i = (bid / H) % N;
    int b = bid / (H * N);
    int tid = threadIdx.x;
    int dhead = D / H;
    int qkv_stride = H * dhead;
    int qp_stride = H * QP * 3;
    int vp_stride = H * VP * 3;
    int row_i = b * N + i;

    __shared__ float logits[64];
    __shared__ float attn[64];

    if (tid < N) {
        int j = tid;
        int row_j = b * N + j;
        float sc = 0.0f;
        for (int dd = 0; dd < dhead; ++dd) {
            sc += q[row_i * qkv_stride + h * dhead + dd] * k[row_j * qkv_stride + h * dhead + dd];
        }
        sc *= rsqrtf((float)dhead);

        float dist = 0.0f;
        for (int p = 0; p < QP; ++p) {
            float ql0 = qpl[row_i * qp_stride + (h*QP + p)*3 + 0];
            float ql1 = qpl[row_i * qp_stride + (h*QP + p)*3 + 1];
            float ql2 = qpl[row_i * qp_stride + (h*QP + p)*3 + 2];
            float kl0 = kpl[row_j * qp_stride + (h*QP + p)*3 + 0];
            float kl1 = kpl[row_j * qp_stride + (h*QP + p)*3 + 1];
            float kl2 = kpl[row_j * qp_stride + (h*QP + p)*3 + 2];
            for (int c = 0; c < 3; ++c) {
                float qg = ql0 * rotations[(row_i*3 + 0)*3 + c] + ql1 * rotations[(row_i*3 + 1)*3 + c] + ql2 * rotations[(row_i*3 + 2)*3 + c] + translations[row_i*3 + c];
                float kg = kl0 * rotations[(row_j*3 + 0)*3 + c] + kl1 * rotations[(row_j*3 + 1)*3 + c] + kl2 * rotations[(row_j*3 + 2)*3 + c] + translations[row_j*3 + c];
                float diff = qg - kg;
                dist += diff * diff;
            }
        }
        float pb = 0.0f;
        const float* pr = pair_repr + ((b * N + i) * N + j) * PD;
        const float* pw = pair_w + h * PD;
        for (int pp = 0; pp < PD; ++pp) pb += pr[pp] * pw[pp];
        float w = softplus_dev(head_w[h]);
        logits[j] = sc + pb - 0.5f * w * dist;
    }
    __syncthreads();

    if (tid == 0) {
        float m = logits[0];
        for (int j = 1; j < N; ++j) m = fmaxf(m, logits[j]);
        float s = 0.0f;
        for (int j = 0; j < N; ++j) {
            float e = expf(logits[j] - m);
            attn[j] = e;
            s += e;
        }
        float invs = 1.0f / s;
        for (int j = 0; j < N; ++j) attn[j] *= invs;
    }
    __syncthreads();

    int out_width = H * dhead + H * VP * 3;
    if (tid < dhead) {
        float acc = 0.0f;
        for (int j = 0; j < N; ++j) {
            int row_j = b * N + j;
            acc += attn[j] * v[row_j * qkv_stride + h * dhead + tid];
        }
        cat[row_i * out_width + h * dhead + tid] = acc;
    }
    if (tid < VP * 3) {
        int vp = tid / 3;
        int ld = tid - vp * 3;
        float glob[3];
        for (int c = 0; c < 3; ++c) {
            float acc = 0.0f;
            for (int j = 0; j < N; ++j) {
                int row_j = b * N + j;
                float vl0 = vpl[row_j * vp_stride + (h*VP + vp)*3 + 0];
                float vl1 = vpl[row_j * vp_stride + (h*VP + vp)*3 + 1];
                float vl2 = vpl[row_j * vp_stride + (h*VP + vp)*3 + 2];
                float vg = vl0 * rotations[(row_j*3 + 0)*3 + c] + vl1 * rotations[(row_j*3 + 1)*3 + c] + vl2 * rotations[(row_j*3 + 2)*3 + c] + translations[row_j*3 + c];
                acc += attn[j] * vg;
            }
            glob[c] = acc;
        }
        float centered0 = glob[0] - translations[row_i*3 + 0];
        float centered1 = glob[1] - translations[row_i*3 + 1];
        float centered2 = glob[2] - translations[row_i*3 + 2];
        float local = centered0 * rotations[(row_i*3 + ld)*3 + 0] + centered1 * rotations[(row_i*3 + ld)*3 + 1] + centered2 * rotations[(row_i*3 + ld)*3 + 2];
        cat[row_i * out_width + H*dhead + h * VP * 3 + tid] = local;
    }
}

__global__ void out_proj_kernel(
    const float* __restrict__ cat,
    const float* __restrict__ out_w,
    const float* __restrict__ out_b,
    float* __restrict__ out,
    int BN, int D, int INW)
{
    int row = blockIdx.x;
    int od = threadIdx.x;
    if (od < D) {
        float acc = out_b[od];
        const float* w = out_w + od * INW;
        const float* x = cat + row * INW;
        for (int f = 0; f < INW; ++f) acc += x[f] * w[f];
        out[row * D + od] = acc;
    }
}

torch::Tensor ipa_forward(
    torch::Tensor node_repr, torch::Tensor pair_repr,
    torch::Tensor translations, torch::Tensor rotations,
    torch::Tensor ln_w, torch::Tensor ln_b,
    torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
    torch::Tensor qpts_w, torch::Tensor kpts_w, torch::Tensor vpts_w,
    torch::Tensor pair_w, torch::Tensor head_w,
    torch::Tensor out_w, torch::Tensor out_b,
    int64_t num_heads, int64_t num_query_points, int64_t num_value_points)
{
    int B = (int)node_repr.size(0);
    int N = (int)node_repr.size(1);
    int D = (int)node_repr.size(2);
    int PD = (int)pair_repr.size(3);
    int H = (int)num_heads;
    int QP = (int)num_query_points;
    int VP = (int)num_value_points;
    int BN = B * N;
    int dhead = D / H;
    auto opts = node_repr.options();
    auto q = torch::empty({BN, H*dhead}, opts);
    auto k = torch::empty({BN, H*dhead}, opts);
    auto v = torch::empty({BN, H*dhead}, opts);
    auto qpl = torch::empty({BN, H*QP*3}, opts);
    auto kpl = torch::empty({BN, H*QP*3}, opts);
    auto vpl = torch::empty({BN, H*VP*3}, opts);
    auto cat = torch::empty({BN, H*dhead + H*VP*3}, opts);
    auto out = torch::empty({B, N, D}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    ln_proj_kernel<<<BN, 128, 0, stream>>>(
        node_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>(),
        qpts_w.data_ptr<float>(), kpts_w.data_ptr<float>(), vpts_w.data_ptr<float>(),
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(),
        qpl.data_ptr<float>(), kpl.data_ptr<float>(), vpl.data_ptr<float>(),
        BN, D, H, QP, VP);
    attn_kernel<<<B*N*H, 64, 0, stream>>>(
        pair_repr.data_ptr<float>(), translations.data_ptr<float>(), rotations.data_ptr<float>(),
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(),
        qpl.data_ptr<float>(), kpl.data_ptr<float>(), vpl.data_ptr<float>(),
        pair_w.data_ptr<float>(), head_w.data_ptr<float>(), cat.data_ptr<float>(),
        B, N, D, PD, H, QP, VP);
    int inw = H*dhead + H*VP*3;
    out_proj_kernel<<<BN, 64, 0, stream>>>(cat.data_ptr<float>(), out_w.data_ptr<float>(), out_b.data_ptr<float>(), out.data_ptr<float>(), BN, D, inw);
    return out;
}
'''

_ipa_ext = load_inline(
    name='ipa_fused_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['ipa_forward'],
    verbose=False,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, node_dim, pair_dim, num_heads, num_query_points, num_value_points):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = node_dim // num_heads
        self.num_query_points = num_query_points
        self.num_value_points = num_value_points
        assert node_dim % num_heads == 0
        d = self.head_dim
        self.q_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.k_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.v_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.q_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.k_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.v_pts = nn.Linear(node_dim, num_heads * num_value_points * 3, bias=False)
        self.pair_bias = nn.Linear(pair_dim, num_heads, bias=False)
        self.head_weights = nn.Parameter(torch.zeros(num_heads))
        out_features = num_heads * d + num_heads * num_value_points * 3
        self.out_proj = nn.Linear(out_features, node_dim)
        self.layer_norm = nn.LayerNorm(node_dim)

    def forward(self, node_repr, pair_repr, translations, rotations):
        return _ipa_ext.ipa_forward(
            node_repr, pair_repr, translations, rotations,
            self.layer_norm.weight, self.layer_norm.bias,
            self.q_proj.weight, self.k_proj.weight, self.v_proj.weight,
            self.q_pts.weight, self.k_pts.weight, self.v_pts.weight,
            self.pair_bias.weight, self.head_weights,
            self.out_proj.weight, self.out_proj.bias,
            self.num_heads, self.num_query_points, self.num_value_points)