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

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

__global__ void compute_outer_mean(
    const float* __restrict__ x,
    const float* __restrict__ mask,
    const float* __restrict__ Wa,
    const float* __restrict__ ba,
    const float* __restrict__ Wb,
    const float* __restrict__ bb,
    const float* __restrict__ Wout,
    const float* __restrict__ bout,
    float* __restrict__ out,
    int B, int S, int N, int D, int H, int P
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N * N * P;
    if (idx >= total) return;

    int p = idx % P;
    int j = (idx / P) % N;
    int i = (idx / (P * N)) % N;
    int b = idx / (P * N * N);

    float accum[1024];
    for (int hh = 0; hh < H*H; hh++) accum[hh] = 0.0f;

    float denom = 0.0f;
    for (int s = 0; s < S; s++) {
        denom += mask[(b * S + s) * N + i];
    }
    denom = fmaxf(denom, 1.0f);

    for (int s=0; s<S; s++) {
        float m_i = mask[(b * S + s) * N + i];
        float m_j = mask[(b * S + s) * N + j];

        float a[32], bbv[32];
        for (int h=0; h<H; h++) {
            float va = ba[h];
            for (int d=0; d<D; d++) {
                va += Wa[h * D + d] * x[((b * S + s) * N + i) * D + d];
            }
            a[h] = va * m_i;
        }
        for (int h=0; h<H; h++) {
            float vb = bb[h];
            for (int d=0; d<D; d++) {
                vb += Wb[h * D + d] * x[((b * S + s) * N + j) * D + d];
            }
            bbv[h] = vb * m_j;
        }
        for (int h1=0; h1<H; h1++) {
            for (int h2=0; h2<H; h2++) {
                accum[h1*H + h2] += a[h1] * bbv[h2];
            }
        }
    }

    float res = bout[p];
    for (int hh=0; hh < H*H; hh++) {
        res += Wout[p * (H*H) + hh] * (accum[hh] / denom);
    }
    out[idx] = res;
}

torch::Tensor outer_product_mean_cuda(
    torch::Tensor msa_repr,
    torch::Tensor msa_mask,
    torch::Tensor ln_weight,
    torch::Tensor ln_bias,
    torch::Tensor proj_a_weight,
    torch::Tensor proj_a_bias,
    torch::Tensor proj_b_weight,
    torch::Tensor proj_b_bias,
    torch::Tensor out_proj_weight,
    torch::Tensor out_proj_bias
) {
    auto B = msa_repr.size(0);
    auto S = msa_repr.size(1);
    auto N = msa_repr.size(2);
    auto D = msa_repr.size(3);
    auto H = proj_a_weight.size(0);
    auto P = out_proj_bias.size(0);

    auto x = torch::layer_norm(msa_repr, {D}, ln_weight, ln_bias, 1e-5);

    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(msa_repr.device());
    auto output = torch::empty({B, N, N, P}, options);

    int total = B * N * N * P;
    dim3 block(256);
    dim3 grid((total + 255) / 256);

    compute_outer_mean<<<grid, block>>>(
        x.data_ptr<float>(),
        msa_mask.data_ptr<float>(),
        proj_a_weight.data_ptr<float>(),
        proj_a_bias.data_ptr<float>(),
        proj_b_weight.data_ptr<float>(),
        proj_b_bias.data_ptr<float>(),
        out_proj_weight.data_ptr<float>(),
        out_proj_bias.data_ptr<float>(),
        output.data_ptr<float>(),
        B, S, N, D, H, P
    );
    cudaDeviceSynchronize();
    return output;
}
'''

cpp_source = '''
torch::Tensor outer_product_mean_cuda(
    torch::Tensor msa_repr,
    torch::Tensor msa_mask,
    torch::Tensor ln_weight,
    torch::Tensor ln_bias,
    torch::Tensor proj_a_weight,
    torch::Tensor proj_a_bias,
    torch::Tensor proj_b_weight,
    torch::Tensor proj_b_bias,
    torch::Tensor out_proj_weight,
    torch::Tensor out_proj_bias
);
'''

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)
        self.ext = load_inline(
            name='outer_mean_ext',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['outer_product_mean_cuda'],
            verbose=False
        )

    def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:
        return self.ext.outer_product_mean_cuda(
            msa_repr, msa_mask,
            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
        )