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_repr, torch::Tensor msa_mask,
                          torch::Tensor ln_w, torch::Tensor ln_b,
                          torch::Tensor pa_w, torch::Tensor pa_b,
                          torch::Tensor pb_w, torch::Tensor pb_b,
                          torch::Tensor out_w, torch::Tensor out_b);
'''

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

#define C 64
#define H 32
#define P 64
#define S 8
#define N 32
#define BATCH 2

__inline__ __device__ float warp_reduce_sum(float v) {
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
    return v;
}

__inline__ __device__ float block_reduce_sum(float v) {
    __shared__ float shared[8];
    __shared__ float result;
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_reduce_sum(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    v = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) v = warp_reduce_sum(v);
    if (threadIdx.x == 0) result = v;
    __syncthreads();
    return result;
}

__global__ void project_kernel(const float* __restrict__ x,
                               const float* __restrict__ ln_w, const float* __restrict__ ln_b,
                               const float* __restrict__ pa_w, const float* __restrict__ pa_b,
                               const float* __restrict__ pb_w, const float* __restrict__ pb_b,
                               float* __restrict__ a, float* __restrict__ b) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const float* xr = x + row * C;

    float sum = 0.0f;
    for (int c = tid; c < C; c += blockDim.x) sum += xr[c];
    float mean = block_reduce_sum(sum) * (1.0f / C);

    float ss = 0.0f;
    for (int c = tid; c < C; c += blockDim.x) {
        float d = xr[c] - mean;
        ss += d * d;
    }
    float inv = rsqrtf(block_reduce_sum(ss) * (1.0f / C) + 1.0e-5f);

    __shared__ float norm[C];
    for (int c = tid; c < C; c += blockDim.x) {
        norm[c] = (xr[c] - mean) * inv * ln_w[c] + ln_b[c];
    }
    __syncthreads();

    if (tid < H) {
        int h = tid;
        float va = pa_b[h];
        float vb = pb_b[h];
        const float* wa = pa_w + h * C;
        const float* wb = pb_w + h * C;
        #pragma unroll
        for (int c = 0; c < C; ++c) {
            float nv = norm[c];
            va += nv * wa[c];
            vb += nv * wb[c];
        }
        a[row * H + h] = va;
        b[row * H + h] = vb;
    }
}

__global__ void outer_kernel(const float* __restrict__ a, const float* __restrict__ b,
                             float* __restrict__ outer) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = BATCH * N * N * H * H;
    if (idx >= total) return;
    int t = idx;
    int k = t & (H - 1); t >>= 5;
    int h = t & (H - 1); t >>= 5;
    int j = t & (N - 1); t >>= 5;
    int i = t & (N - 1); t >>= 5;
    int bb = t;
    float acc = 0.0f;
    #pragma unroll
    for (int s = 0; s < S; ++s) {
        int base = (bb * S + s) * N;
        acc += a[(base + i) * H + h] * b[(base + j) * H + k];
    }
    outer[idx] = acc * 0.125f;
}

__global__ void bias_kernel(float* __restrict__ out, const float* __restrict__ bias) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = BATCH * N * N * P;
    if (idx < total) out[idx] += bias[idx & (P - 1)];
}

torch::Tensor opm_forward(torch::Tensor msa_repr, torch::Tensor msa_mask,
                          torch::Tensor ln_w, torch::Tensor ln_b,
                          torch::Tensor pa_w, torch::Tensor pa_b,
                          torch::Tensor pb_w, torch::Tensor pb_b,
                          torch::Tensor out_w, torch::Tensor out_b) {
    auto opts = msa_repr.options();
    auto a = torch::empty({BATCH * S * N, H}, opts);
    auto bb = torch::empty({BATCH * S * N, H}, opts);
    auto outer = torch::empty({BATCH * N * N, H * H}, opts);
    auto out = torch::empty({BATCH, N, N, P}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    project_kernel<<<BATCH * S * N, 128, 0, stream>>>(
        msa_repr.data_ptr<float>(), ln_w.data_ptr<float>(), ln_b.data_ptr<float>(),
        pa_w.data_ptr<float>(), pa_b.data_ptr<float>(), pb_w.data_ptr<float>(), pb_b.data_ptr<float>(),
        a.data_ptr<float>(), bb.data_ptr<float>());

    outer_kernel<<<(BATCH * N * N * H * H + 255) / 256, 256, 0, stream>>>(
        a.data_ptr<float>(), bb.data_ptr<float>(), outer.data_ptr<float>());

    float alpha = 1.0f, beta = 0.0f;
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, stream);
    cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N,
                P, BATCH * N * N, H * H,
                &alpha,
                out_w.data_ptr<float>(), H * H,
                outer.data_ptr<float>(), H * H,
                &beta,
                out.data_ptr<float>(), P);
    bias_kernel<<<(BATCH * N * N * P + 255) / 256, 256, 0, stream>>>(out.data_ptr<float>(), out_b.data_ptr<float>());
    return out;
}
'''

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