import torch
import torch.nn as nn


class Model(nn.Module):
    """
    SE(3)-invariant linear layer.  Maps node features that are augmented
    with 3-D coordinate information through a linear transform that
    respects rotational and translational invariance.  The layer operates
    on (scalar_features, vector_features) pairs where scalar features
    transform trivially and vector features transform as 3-D vectors.
    Used in equivariant GNNs (EGNN, PaiNN, TFN).
    """

    def __init__(self, scalar_in, scalar_out, vector_in, vector_out):
        super().__init__()
        self.scalar_linear = nn.Linear(scalar_in + vector_in, scalar_out)
        self.vector_weight = nn.Parameter(torch.randn(vector_out, vector_in))
        nn.init.xavier_uniform_(self.vector_weight)
        self.vector_gate = nn.Linear(scalar_in + vector_in, vector_out)

    def forward(
        self, scalar_feat: torch.Tensor, vector_feat: torch.Tensor
    ) -> tuple:
        """
        Args:
            scalar_feat: (B, N, scalar_in) – invariant node features
            vector_feat: (B, N, vector_in, 3) – equivariant vector features
        Returns:
            (scalar_out, vector_out):
                scalar_out: (B, N, scalar_out)
                vector_out: (B, N, vector_out, 3)
        """
        # Invariant scalars from vector norms
        vec_norm = vector_feat.norm(dim=-1)  # (B, N, vector_in)
        combined = torch.cat([scalar_feat, vec_norm], dim=-1)

        s_out = self.scalar_linear(combined)

        # Equivariant vector transform: linear combination of input vectors
        v_out = torch.einsum("oi,bnid->bnod", self.vector_weight, vector_feat)
        # Gating by scalar-derived signal
        gate = torch.sigmoid(self.vector_gate(combined)).unsqueeze(-1)  # (B, N, vector_out, 1)
        v_out = gate * v_out

        return s_out, v_out


scalar_in = 64
scalar_out = 64
vector_in = 16
vector_out = 16
num_nodes = 128
batch_size = 8


def get_inputs():
    p = popcorn_pri
    bs = p.jitter_int(batch_size)
    nn = p.jitter_int(num_nodes)
    s = torch.randn(bs, nn, scalar_in)
    v = torch.randn(bs, nn, vector_in, 3)
    return [s, v]


def get_init_inputs():
    return [scalar_in, scalar_out, vector_in, vector_out]

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

# Compile CUDA kernel at module load time
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// Fused kernel: einsum + sigmoid gating
// v_out[b,n,o,d] = sigmoid(sum_k gate_w[o,k] * combined[b,n,k] + gate_b[o]) * sum_i weight[o,i] * v_in[b,n,i,d]
__global__ void vec_transform_kernel(
    const float* __restrict__ combined,
    const float* __restrict__ vector_weight,
    const float* __restrict__ gate_weight,
    const float* __restrict__ gate_bias,
    const float* __restrict__ vector_in,
    float* __restrict__ vector_out,
    int batch_size, int num_nodes, int vector_in_dim, int vector_out_dim, int combined_dim
) {
    // Each thread handles one (b, n, o, d) output element
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * num_nodes * vector_out_dim * 3;
    
    if (idx < total) {
        int d = idx % 3;
        int o = (idx / 3) % vector_out_dim;
        int n = (idx / (3 * vector_out_dim)) % num_nodes;
        int b = idx / (3 * vector_out_dim * num_nodes);
        
        int bn_idx = b * num_nodes + n;
        
        // Compute gate value: sigmoid(sum_k gate_w[o,k] * combined[b,n,k] + gate_b[o])
        float gate_val = gate_bias[o];
        for (int k = 0; k < combined_dim; k++) {
            gate_val += gate_weight[o * combined_dim + k] * combined[bn_idx * combined_dim + k];
        }
        gate_val = 1.0f / (1.0f + expf(-gate_val));  // sigmoid
        
        // Compute weighted vector sum: sum_i weight[o,i] * v_in[b,n,i,d]
        float v_val = 0.0f;
        for (int i = 0; i < vector_in_dim; i++) {
            int v_in_idx = ((b * num_nodes + n) * vector_in_dim + i) * 3 + d;
            v_val += vector_weight[o * vector_in_dim + i] * vector_in[v_in_idx];
        }
        
        // Apply gate
        vector_out[idx] = gate_val * v_val;
    }
}

// Vectorized version: process 4 elements per thread (o*d grouped)
__global__ void vec_transform_kernel_vec4(
    const float* __restrict__ combined,
    const float* __restrict__ vector_weight,
    const float* __restrict__ gate_weight,
    const float* __restrict__ gate_bias,
    const float* __restrict__ vector_in,
    float* __restrict__ vector_out,
    int batch_size, int num_nodes, int vector_in_dim, int vector_out_dim, int combined_dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int bn = batch_size * num_nodes;
    int total = bn * vector_out_dim;  // Each thread handles 3 elements (one o for all d)
    
    if (idx < total) {
        int o = idx % vector_out_dim;
        int bn_idx = idx / vector_out_dim;
        
        // Compute gate once for all 3 dimensions
        float gate_val = gate_bias[o];
        const float* combined_ptr = combined + bn_idx * combined_dim;
        const float* gate_w_ptr = gate_weight + o * combined_dim;
        
        #pragma unroll 8
        for (int k = 0; k < combined_dim; k++) {
            gate_val += gate_w_ptr[k] * combined_ptr[k];
        }
        gate_val = 1.0f / (1.0f + expf(-gate_val));
        
        // Compute weighted vector for each of 3 dimensions
        const float* weight_ptr = vector_weight + o * vector_in_dim;
        
        for (int d = 0; d < 3; d++) {
            float v_val = 0.0f;
            #pragma unroll 4
            for (int i = 0; i < vector_in_dim; i++) {
                int v_in_idx = (bn_idx * vector_in_dim + i) * 3 + d;
                v_val += weight_ptr[i] * vector_in[v_in_idx];
            }
            vector_out[(bn_idx * vector_out_dim + o) * 3 + d] = gate_val * v_val;
        }
    }
}

torch::Tensor vec_transform_cuda(
    torch::Tensor combined,
    torch::Tensor vector_weight,
    torch::Tensor gate_weight,
    torch::Tensor gate_bias,
    torch::Tensor vector_in
) {
    int batch_size = combined.size(0);
    int num_nodes = combined.size(1);
    int combined_dim = combined.size(2);
    int vector_in_dim = vector_in.size(2);
    int vector_out_dim = vector_weight.size(0);
    
    auto vector_out = torch::empty({batch_size, num_nodes, vector_out_dim, 3}, combined.options());
    
    int total = batch_size * num_nodes * vector_out_dim;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    vec_transform_kernel_vec4<<<blocks, threads>>>(
        combined.data_ptr<float>(),
        vector_weight.data_ptr<float>(),
        gate_weight.data_ptr<float>(),
        gate_bias.data_ptr<float>(),
        vector_in.data_ptr<float>(),
        vector_out.data_ptr<float>(),
        batch_size, num_nodes, vector_in_dim, vector_out_dim, combined_dim
    );
    
    return vector_out;
}
"""

cpp_source = """
torch::Tensor vec_transform_cuda(
    torch::Tensor combined,
    torch::Tensor vector_weight,
    torch::Tensor gate_weight,
    torch::Tensor gate_bias,
    torch::Tensor vector_in
);
"""

vec_module = load_inline(
    name="vec_transform_module",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["vec_transform_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"]
)


class ModelNew(nn.Module):
    """
    SE(3)-invariant linear layer with optimized CUDA kernel for vector transformation.
    """

    def __init__(self, scalar_in, scalar_out, vector_in, vector_out):
        super().__init__()
        self.scalar_linear = nn.Linear(scalar_in + vector_in, scalar_out)
        self.vector_weight = nn.Parameter(torch.empty(vector_out, vector_in))
        nn.init.xavier_uniform_(self.vector_weight)
        self.vector_gate = nn.Linear(scalar_in + vector_in, vector_out)
        
        self.vector_in = vector_in
        self.vector_out = vector_out

    def forward(
        self, scalar_feat: torch.Tensor, vector_feat: torch.Tensor
    ) -> tuple:
        """
        Args:
            scalar_feat: (B, N, scalar_in) - invariant node features
            vector_feat: (B, N, vector_in, 3) - equivariant vector features
        Returns:
            (scalar_out, vector_out):
                scalar_out: (B, N, scalar_out)
                vector_out: (B, N, vector_out, 3)
        """
        # Invariant scalars from vector norms
        vec_norm = vector_feat.norm(dim=-1)  # (B, N, vector_in)
        combined = torch.cat([scalar_feat, vec_norm], dim=-1)  # (B, N, scalar_in + vector_in)

        # Scalar output via optimized linear
        s_out = self.scalar_linear(combined)

        # Vector output via fused CUDA kernel (einsum + sigmoid gating)
        v_out = vec_module.vec_transform_cuda(
            combined,
            self.vector_weight,
            self.vector_gate.weight,
            self.vector_gate.bias,
            vector_feat
        )

        return s_out, v_out