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


class Model(nn.Module):
    """
    Triangular multiplicative update – incoming edges.
    Mirror of the outgoing variant: updates pair (i,j) by aggregating over
    the shared index k via edges (k,i) and (k,j).
    """

    def __init__(self, pair_dim, hidden_dim):
        super().__init__()
        self.layer_norm_in = nn.LayerNorm(pair_dim)
        self.proj_a = nn.Linear(pair_dim, hidden_dim)
        self.gate_a = nn.Linear(pair_dim, hidden_dim)
        self.proj_b = nn.Linear(pair_dim, hidden_dim)
        self.gate_b = nn.Linear(pair_dim, hidden_dim)
        self.layer_norm_out = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, pair_dim)
        self.out_gate = nn.Linear(pair_dim, pair_dim)

    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)
        """
        mask = pair_mask.unsqueeze(-1)
        x = self.layer_norm_in(pair_repr)

        a = self.proj_a(x) * torch.sigmoid(self.gate_a(x)) * mask  # (B, N, N, H)
        b = self.proj_b(x) * torch.sigmoid(self.gate_b(x)) * mask

        # Incoming: aggregate over k with a[k,i] * b[k,j]
        out = torch.einsum("bkih,bkjh->bijh", a, b)

        out = self.layer_norm_out(out)
        out = self.out_proj(out)
        out = torch.sigmoid(self.out_gate(pair_repr)) * out
        return out


pair_dim = 64
hidden_dim = 32
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, hidden_dim]

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
import math

# CUDA source for the einsum kernel: "bkih,bkjh->bijh"
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>

// Threadblock size for the reduction kernel
#define TILE_SIZE 16

// Custom kernel for the einsum "bkih,bkjh->bijh"
// Computes: out[b,i,j,h] = sum_k (a[b,k,i,h] * b[b,k,j,h])
__global__ void einsum_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ out,
    const int B, const int N, const int H
) {
    // Each thread computes one output element (b, i, j, h)
    // Block dimensions: (TILE_SIZE, TILE_SIZE, 1) for (i, j) pairs
    // Grid: (B * H blocks in z, ceil(N/TILE_SIZE) in y, ceil(N/TILE_SIZE) in x)
    
    const int tid_x = threadIdx.x;
    const int tid_y = threadIdx.y;
    
    const int i = blockIdx.x * TILE_SIZE + tid_x;
    const int j = blockIdx.y * TILE_SIZE + tid_y;
    const int bh = blockIdx.z;  // Combined b*h index
    const int b = bh / H;
    const int h = bh % H;
    
    if (i >= N || j >= N) return;
    
    // Shared memory for tiling k dimension
    __shared__ float a_tile[TILE_SIZE][TILE_SIZE];
    __shared__ float b_tile[TILE_SIZE][TILE_SIZE];
    
    float sum = 0.0f;
    
    const int num_k_tiles = (N + TILE_SIZE - 1) / TILE_SIZE;
    
    for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) {
        // Load tiles into shared memory
        // a[b,k,i,h] with shape (B, N, N, H)
        // b[b,k,j,h] with shape (B, N, N, H)
        
        const int k_start = k_tile * TILE_SIZE;
        
        // Cooperative loading
        if (tid_y < TILE_SIZE) {
            const int k_load = k_start + tid_y;
            if (k_load < N) {
                // a[b, k, i, h]
                a_tile[tid_y][tid_x] = a[b * N * N * H + k_load * N * H + i * H + h];
            } else {
                a_tile[tid_y][tid_x] = 0.0f;
            }
        }
        
        if (tid_x < TILE_SIZE) {
            const int k_load = k_start + tid_x;
            if (k_load < N) {
                // b[b, k, j, h]
                b_tile[tid_x][tid_y] = b[b * N * N * H + k_load * N * H + j * H + h];
            } else {
                b_tile[tid_x][tid_y] = 0.0f;
            }
        }
        
        __syncthreads();
        
        // Compute partial sum for this tile
        #pragma unroll
        for (int k_local = 0; k_local < TILE_SIZE; ++k_local) {
            sum += a_tile[k_local][tid_x] * b_tile[k_local][tid_y];
        }
        
        __syncthreads();
    }
    
    // Write output: out[b, i, j, h] with shape (B, N, N, H)
    out[b * N * N * H + i * N * H + j * H + h] = sum;
}

// Alternative: Use cuBLAS batched GEMM for the einsum
// Reshape: a(B,N,N,H) -> (B,H,N,N) transposed to (B,H,N,N)
// The computation becomes batched matmul

torch::Tensor einsum_cuda(torch::Tensor a, torch::Tensor b) {
    // a: (B, N, N, H), b: (B, N, N, H)
    // output: (B, N, N, H)
    
    TORCH_CHECK(a.dim() == 4, "a must be 4D");
    TORCH_CHECK(b.dim() == 4, "b must be 4D");
    TORCH_CHECK(a.sizes() == b.sizes(), "a and b must have same shape");
    TORCH_CHECK(a.is_contiguous(), "a must be contiguous");
    TORCH_CHECK(b.is_contiguous(), "b must be contiguous");
    
    const int B = a.size(0);
    const int N = a.size(1);
    const int H = a.size(3);
    
    auto out = torch::zeros({B, N, N, H}, a.options());
    
    const dim3 blockSize(TILE_SIZE, TILE_SIZE, 1);
    const dim3 gridSize((N + TILE_SIZE - 1) / TILE_SIZE,
                        (N + TILE_SIZE - 1) / TILE_SIZE,
                        B * H);
    
    einsum_kernel<<<gridSize, blockSize>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        out.data_ptr<float>(),
        B, N, H
    );
    
    return out;
}

// Fused kernel: combines gating + einsum
__global__ void fused_gated_einsum_kernel(
    const float* __restrict__ proj_a,
    const float* __restrict__ gate_a,
    const float* __restrict__ proj_b,
    const float* __restrict__ gate_b,
    const float* __restrict__ mask,
    float* __restrict__ out,
    const int B, const int N, const int H
) {
    const int tid_x = threadIdx.x;
    const int tid_y = threadIdx.y;
    
    const int i = blockIdx.x * TILE_SIZE + tid_x;
    const int j = blockIdx.y * TILE_SIZE + tid_y;
    const int bh = blockIdx.z;
    const int b = bh / H;
    const int h = bh % H;
    
    if (i >= N || j >= N) return;
    
    __shared__ float a_tile[TILE_SIZE][TILE_SIZE];
    __shared__ float b_tile[TILE_SIZE][TILE_SIZE];
    
    float sum = 0.0f;
    
    const int num_k_tiles = (N + TILE_SIZE - 1) / TILE_SIZE;
    
    for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) {
        const int k_start = k_tile * TILE_SIZE;
        
        // Load and apply gating on-the-fly
        if (tid_y < TILE_SIZE) {
            const int k_load = k_start + tid_y;
            if (k_load < N) {
                // a_val = proj_a * sigmoid(gate_a) * mask
                int idx = b * N * N * H + k_load * N * H + i * H + h;
                float a_val = proj_a[idx] * (1.0f / (1.0f + expf(-gate_a[idx])));
                float mask_val = mask[b * N * N + k_load * N + i];
                a_tile[tid_y][tid_x] = a_val * mask_val;
            } else {
                a_tile[tid_y][tid_x] = 0.0f;
            }
        }
        
        if (tid_x < TILE_SIZE) {
            const int k_load = k_start + tid_x;
            if (k_load < N) {
                int idx = b * N * N * H + k_load * N * H + j * H + h;
                float b_val = proj_b[idx] * (1.0f / (1.0f + expf(-gate_b[idx])));
                float mask_val = mask[b * N * N + k_load * N + j];
                b_tile[tid_x][tid_y] = b_val * mask_val;
            } else {
                b_tile[tid_x][tid_y] = 0.0f;
            }
        }
        
        __syncthreads();
        
        #pragma unroll
        for (int k_local = 0; k_local < TILE_SIZE; ++k_local) {
            sum += a_tile[k_local][tid_x] * b_tile[k_local][tid_y];
        }
        
        __syncthreads();
    }
    
    out[b * N * N * H + i * N * H + j * H + h] = sum;
}

torch::Tensor fused_gated_einsum_cuda(
    torch::Tensor proj_a,
    torch::Tensor gate_a,
    torch::Tensor proj_b,
    torch::Tensor gate_b,
    torch::Tensor mask
) {
    const int B = proj_a.size(0);
    const int N = proj_a.size(1);
    const int H = proj_a.size(3);
    
    auto out = torch::zeros({B, N, N, H}, proj_a.options());
    
    const dim3 blockSize(TILE_SIZE, TILE_SIZE, 1);
    const dim3 gridSize((N + TILE_SIZE - 1) / TILE_SIZE,
                        (N + TILE_SIZE - 1) / TILE_SIZE,
                        B * H);
    
    fused_gated_einsum_kernel<<<gridSize, blockSize>>>(
        proj_a.data_ptr<float>(),
        gate_a.data_ptr<float>(),
        proj_b.data_ptr<float>(),
        gate_b.data_ptr<float>(),
        mask.data_ptr<float>(),
        out.data_ptr<float>(),
        B, N, H
    );
    
    return out;
}
"""

# Load the CUDA extension
einsum_ext = load_inline(
    name='einsum_ext',
    cpp_sources='',
    cuda_sources=cuda_source,
    functions=['einsum_cuda', 'fused_gated_einsum_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Triangular multiplicative update - incoming edges.
    Optimized with custom CUDA kernel for the einsum operation.
    """

    def __init__(self, pair_dim, hidden_dim):
        super().__init__()
        self.layer_norm_in = nn.LayerNorm(pair_dim)
        self.proj_a = nn.Linear(pair_dim, hidden_dim)
        self.gate_a = nn.Linear(pair_dim, hidden_dim)
        self.proj_b = nn.Linear(pair_dim, hidden_dim)
        self.gate_b = nn.Linear(pair_dim, hidden_dim)
        self.layer_norm_out = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, pair_dim)
        self.out_gate = nn.Linear(pair_dim, pair_dim)
        
        self.hidden_dim = hidden_dim
        self.pair_dim = pair_dim

    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)
        """
        # Ensure inputs are contiguous for the CUDA kernel
        pair_repr = pair_repr.contiguous()
        pair_mask = pair_mask.contiguous()
        
        mask = pair_mask.unsqueeze(-1)
        x = self.layer_norm_in(pair_repr)

        # Compute projections and gates
        proj_a_x = self.proj_a(x)
        gate_a_x = self.gate_a(x)
        proj_b_x = self.proj_b(x)
        gate_b_x = self.gate_b(x)
        
        # Use the fused CUDA kernel for the einsum with gating
        # This computes: a = proj_a * sigmoid(gate_a) * mask, b = proj_b * sigmoid(gate_b) * mask
        # Then: out = einsum("bkih,bkjh->bijh", a, b)
        out = einsum_ext.fused_gated_einsum_cuda(
            proj_a_x.contiguous(),
            gate_a_x.contiguous(),
            proj_b_x.contiguous(),
            gate_b_x.contiguous(),
            pair_mask.float().contiguous()  # mask without the extra dimension
        )

        out = self.layer_norm_out(out)
        out = self.out_proj(out)
        out = torch.sigmoid(self.out_gate(pair_repr)) * out
        return out