import torch
import torch.nn as nn
import math


class Model(nn.Module):
    """
    Outer-product mean: converts MSA representation to a pair
    representation update.  Used in AlphaFold2 to communicate information
    from the MSA stack to the pair stack.  For each pair (i,j), the
    outer product of projected MSA features at positions i and j is
    averaged over the MSA depth dimension.
    """

    def __init__(self, msa_dim, pair_dim, hidden_dim=32):
        super().__init__()
        self.layer_norm = nn.LayerNorm(msa_dim)
        self.proj_a = nn.Linear(msa_dim, hidden_dim)
        self.proj_b = nn.Linear(msa_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim * hidden_dim, pair_dim)

    def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            msa_repr: (B, S, N, msa_dim) – MSA representation
            msa_mask: (B, S, N)          – 1 for valid rows/positions
        Returns:
            pair update (B, N, N, pair_dim)
        """
        x = self.layer_norm(msa_repr)
        a = self.proj_a(x)  # (B, S, N, H)
        b = self.proj_b(x)

        mask = msa_mask.unsqueeze(-1)  # (B, S, N, 1)
        a = a * mask
        b = b * mask

        outer = torch.einsum("bsih,bsjk->bijhk", a, b)  # (B, N, N, H, H)
        denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1)
        outer = outer / denom  # mean over S

        B, N1, N2, H1, H2 = outer.shape
        outer = outer.reshape(B, N1, N2, H1 * H2)
        return self.out_proj(outer)


msa_dim = 64
pair_dim = 64
hidden_dim = 32
seq_len = 32
msa_depth = 8
batch_size = 2


def get_inputs():
    p = popcorn_pri
    msa_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(msa_depth), p.jitter_int(seq_len), msa_dim)
    msa_mask = torch.ones(p.jitter_int(batch_size), p.jitter_int(msa_depth), p.jitter_int(seq_len))
    return [msa_repr, msa_mask]


def get_init_inputs():
    return [msa_dim, 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 opm_forward(torch::Tensor msa, torch::Tensor mask,
                          torch::Tensor ln_w, torch::Tensor ln_b,
                          torch::Tensor wa, torch::Tensor ba,
                          torch::Tensor wb, torch::Tensor bb,
                          torch::Tensor wo, torch::Tensor bo);
'''

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

__global__ __launch_bounds__(64) void proj_kernel(const float* __restrict__ msa, const float* __restrict__ mask,
                            const float* __restrict__ ln_w, const float* __restrict__ ln_b,
                            const float* __restrict__ wa, const float* __restrict__ ba,
                            const float* __restrict__ wb, const float* __restrict__ bb,
                            float* __restrict__ A, float* __restrict__ Bp,
                            int Bsz) {
    int row = blockIdx.x;
    int n = row & 31;
    int tmp = row >> 5;
    int s = tmp & 7;
    int b = tmp >> 3;
    const float* x = msa + (((b*8 + s)*32 + n)*64);
    __shared__ float sh[64];
    int tid = threadIdx.x;
    float v = x[tid];
    sh[tid] = v;
    __syncthreads();
    if (tid < 32) sh[tid] += sh[tid+32]; __syncthreads();
    if (tid < 16) sh[tid] += sh[tid+16]; __syncthreads();
    if (tid < 8) sh[tid] += sh[tid+8]; __syncthreads();
    if (tid < 4) sh[tid] += sh[tid+4]; __syncthreads();
    if (tid < 2) sh[tid] += sh[tid+2]; __syncthreads();
    if (tid < 1) sh[tid] += sh[tid+1]; __syncthreads();
    float mean = sh[0] * (1.0f/64.0f);
    float d = v - mean;
    sh[tid] = d*d;
    __syncthreads();
    if (tid < 32) sh[tid] += sh[tid+32]; __syncthreads();
    if (tid < 16) sh[tid] += sh[tid+16]; __syncthreads();
    if (tid < 8) sh[tid] += sh[tid+8]; __syncthreads();
    if (tid < 4) sh[tid] += sh[tid+4]; __syncthreads();
    if (tid < 2) sh[tid] += sh[tid+2]; __syncthreads();
    if (tid < 1) sh[tid] += sh[tid+1]; __syncthreads();
    float rstd = rsqrtf(sh[0]*(1.0f/64.0f) + 1.0e-5f);
    float m = mask[(b*8 + s)*32 + n];
    int h = tid & 31;
    int is_b = tid >> 5;
    const float* w = (is_b ? wb : wa) + h*64;
    float acc = 0.0f;
    #pragma unroll
    for (int c=0;c<64;++c) {
        float xn = (x[c] - mean) * rstd;
        xn = xn * ln_w[c] + ln_b[c];
        acc += xn * w[c];
    }
    acc += is_b ? bb[h] : ba[h];
    acc *= m;
    float* out = (is_b ? Bp : A) + (((b*8 + s)*32 + n)*32 + h);
    *out = acc;
}

__global__ __launch_bounds__(256) void transb4_kernel(const float* __restrict__ Bp, const float* __restrict__ wo,
                              float* __restrict__ T, long long total4) {
    long long idx = (long long)blockIdx.x * 256 + threadIdx.x;
    if (idx >= total4) return;
    int p4 = (idx & 15) << 2;
    long long q = idx >> 4;
    int h = q & 31; q >>= 5;
    int j = q & 31; q >>= 5;
    int s = q & 7;  q >>= 3;
    int b = (int)q;
    const float* bv = Bp + (((b*8 + s)*32 + j)*32);
    const float* w0 = wo + (p4+0)*1024 + h*32;
    const float* w1 = wo + (p4+1)*1024 + h*32;
    const float* w2 = wo + (p4+2)*1024 + h*32;
    const float* w3 = wo + (p4+3)*1024 + h*32;
    float a0=0.f,a1=0.f,a2=0.f,a3=0.f;
    #pragma unroll
    for (int k=0;k<32;++k) {
        float bv0 = bv[k];
        a0 += bv0 * w0[k];
        a1 += bv0 * w1[k];
        a2 += bv0 * w2[k];
        a3 += bv0 * w3[k];
    }
    float* tp = T + ((((b*8 + s)*32 + j)*32 + h)*64 + p4);
    tp[0]=a0; tp[1]=a1; tp[2]=a2; tp[3]=a3;
}

__global__ __launch_bounds__(256) void output4_kernel(const float* __restrict__ A, const float* __restrict__ T,
                              const float* __restrict__ mask, const float* __restrict__ bo,
                              float* __restrict__ out, long long total4) {
    long long idx = (long long)blockIdx.x * 256 + threadIdx.x;
    if (idx >= total4) return;
    int p4 = (idx & 15) << 2;
    long long q = idx >> 4;
    int j = q & 31; q >>= 5;
    int i = q & 31; q >>= 5;
    int b = (int)q;
    float a0=0.f,a1=0.f,a2=0.f,a3=0.f;
    #pragma unroll
    for (int s=0;s<8;++s) {
        const float* av = A + (((b*8 + s)*32 + i)*32);
        const float* tv = T + ((((b*8 + s)*32 + j)*32)*64 + p4);
        #pragma unroll
        for (int h=0;h<32;++h) {
            float aa = av[h];
            const float* t = tv + h*64;
            a0 += aa * t[0];
            a1 += aa * t[1];
            a2 += aa * t[2];
            a3 += aa * t[3];
        }
    }
    float denom = 0.f;
    #pragma unroll
    for (int s=0;s<8;++s) denom += mask[(b*8 + s)*32 + i];
    if (denom < 1.f) denom = 1.f;
    float* o = out + ((b*32 + i)*32 + j)*64 + p4;
    float inv = 1.0f / denom;
    o[0] = a0*inv + bo[p4+0];
    o[1] = a1*inv + bo[p4+1];
    o[2] = a2*inv + bo[p4+2];
    o[3] = a3*inv + bo[p4+3];
}

torch::Tensor opm_forward(torch::Tensor msa, torch::Tensor mask,
                          torch::Tensor ln_w, torch::Tensor ln_b,
                          torch::Tensor wa, torch::Tensor ba,
                          torch::Tensor wb, torch::Tensor bb,
                          torch::Tensor wo, torch::Tensor bo) {
    int Bsz = msa.size(0);
    auto opts = msa.options();
    auto A = torch::empty({Bsz, 8, 32, 32}, opts);
    auto Bp = torch::empty({Bsz, 8, 32, 32}, opts);
    auto T = torch::empty({Bsz, 8, 32, 32, 64}, opts);
    auto out = torch::empty({Bsz, 32, 32, 64}, opts);
    proj_kernel<<<Bsz*8*32, 64>>>(msa.data_ptr<float>(), mask.data_ptr<float>(), ln_w.data_ptr<float>(), ln_b.data_ptr<float>(), wa.data_ptr<float>(), ba.data_ptr<float>(), wb.data_ptr<float>(), bb.data_ptr<float>(), A.data_ptr<float>(), Bp.data_ptr<float>(), Bsz);
    long long nt4 = (long long)Bsz*8*32*32*16;
    transb4_kernel<<<(nt4+255)/256, 256>>>(Bp.data_ptr<float>(), wo.data_ptr<float>(), T.data_ptr<float>(), nt4);
    long long no4 = (long long)Bsz*32*32*16;
    output4_kernel<<<(no4+255)/256, 256>>>(A.data_ptr<float>(), T.data_ptr<float>(), mask.data_ptr<float>(), bo.data_ptr<float>(), out.data_ptr<float>(), no4);
    return out;
}
'''
_ext = load_inline(name='opm_ext_v2', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['opm_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, msa_dim, pair_dim, hidden_dim=32):
        super().__init__()
        self.layer_norm = nn.LayerNorm(msa_dim)
        self.proj_a = nn.Linear(msa_dim, hidden_dim)
        self.proj_b = nn.Linear(msa_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim * hidden_dim, pair_dim)
    def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:
        return _ext.opm_forward(msa_repr.contiguous(), msa_mask.contiguous(), self.layer_norm.weight, self.layer_norm.bias, self.proj_a.weight, self.proj_a.bias, self.proj_b.weight, self.proj_b.bias, self.out_proj.weight, self.out_proj.bias)