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

cpp_source = """
void launch_compute_norms_and_concat(
    at::Tensor scalar_feat, at::Tensor vec_feat, at::Tensor combined,
    int64_t B, int64_t N, int64_t S, int64_t V);

void launch_vector_transform(
    at::Tensor vec_feat, at::Tensor weight, at::Tensor v_out,
    int64_t B, int64_t N, int64_t V_in, int64_t V_out);

void launch_apply_gate(
    at::Tensor v_out, at::Tensor gate,
    int64_t B, int64_t N, int64_t V_out);
"""

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// Fused kernel: compute vector norms and concatenate with scalar features
__global__ void compute_norms_and_concat_kernel(
    const float* __restrict__ scalar_feat,
    const float* __restrict__ vec_feat,
    float* __restrict__ combined,
    int B, int N, int S, int V) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N;
    if (idx >= total) return;
    
    int b = idx / N;
    int n = idx % N;
    int base_scalar = (b * N + n) * S;
    int base_vec = (b * N + n) * V;
    int base_combined = (b * N + n) * (S + V);
    
    // Copy scalar features
    for (int s = 0; s < S; s++) {
        combined[base_combined + s] = scalar_feat[base_scalar + s];
    }
    
    // Compute norms and store
    for (int v = 0; v < V; v++) {
        float x = vec_feat[(base_vec + v) * 3 + 0];
        float y = vec_feat[(base_vec + v) * 3 + 1];
        float z = vec_feat[(base_vec + v) * 3 + 2];
        float norm = sqrtf(x*x + y*y + z*z);
        combined[base_combined + S + v] = norm;
    }
}

// Vector transform kernel: v_out = W @ v_in for each position
__global__ void vector_transform_kernel(
    const float* __restrict__ vec_feat,
    const float* __restrict__ weight,
    float* __restrict__ v_out,
    int B, int N, int V_in, int V_out) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N * V_out;
    if (idx >= total) return;
    
    int vo = idx % V_out;
    int rem = idx / V_out;
    int n = rem % N;
    int b = rem / N;
    
    float sum_x = 0.0f, sum_y = 0.0f, sum_z = 0.0f;
    
    for (int vi = 0; vi < V_in; vi++) {
        float w = weight[vo * V_in + vi];
        int vec_idx = ((b * N + n) * V_in + vi) * 3;
        sum_x += w * vec_feat[vec_idx + 0];
        sum_y += w * vec_feat[vec_idx + 1];
        sum_z += w * vec_feat[vec_idx + 2];
    }
    
    int out_idx = ((b * N + n) * V_out + vo) * 3;
    v_out[out_idx + 0] = sum_x;
    v_out[out_idx + 1] = sum_y;
    v_out[out_idx + 2] = sum_z;
}

// Apply gating: v_out = gate * v_out
__global__ void apply_gate_kernel(
    float* __restrict__ v_out,
    const float* __restrict__ gate,
    int B, int N, int V_out) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N * V_out * 3;
    if (idx >= total) return;
    
    int c = idx % 3;
    int rem = idx / 3;
    int vo = rem % V_out;
    rem = rem / V_out;
    int n = rem % N;
    int b = rem / N;
    
    int v_out_idx = ((b * N + n) * V_out + vo) * 3 + c;
    int gate_idx = (b * N + n) * V_out + vo;
    v_out[v_out_idx] = v_out[v_out_idx] * gate[gate_idx];
}

// C++ wrapper implementations
void launch_compute_norms_and_concat(
    at::Tensor scalar_feat, at::Tensor vec_feat, at::Tensor combined,
    int64_t B, int64_t N, int64_t S, int64_t V) {
    int total = B * N;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    compute_norms_and_concat_kernel<<<blocks, threads>>>(
        scalar_feat.data_ptr<float>(),
        vec_feat.data_ptr<float>(),
        combined.data_ptr<float>(),
        B, N, S, V);
}

void launch_vector_transform(
    at::Tensor vec_feat, at::Tensor weight, at::Tensor v_out,
    int64_t B, int64_t N, int64_t V_in, int64_t V_out) {
    int total = B * N * V_out;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    vector_transform_kernel<<<blocks, threads>>>(
        vec_feat.data_ptr<float>(),
        weight.data_ptr<float>(),
        v_out.data_ptr<float>(),
        B, N, V_in, V_out);
}

void launch_apply_gate(
    at::Tensor v_out, at::Tensor gate,
    int64_t B, int64_t N, int64_t V_out) {
    int total = B * N * V_out * 3;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    apply_gate_kernel<<<blocks, threads>>>(
        v_out.data_ptr<float>(),
        gate.data_ptr<float>(),
        B, N, V_out);
}
"""

se3_ops = load_inline(
    name='se3_ops_v4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_compute_norms_and_concat', 'launch_vector_transform', 'launch_apply_gate'],
    extra_cuda_cflags=['-O3'],
)

class ModelNew(nn.Module):
    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:
        B, N, S = scalar_feat.shape
        _, _, V, _ = vector_feat.shape
        
        # Ensure contiguous
        scalar_feat = scalar_feat.contiguous()
        vector_feat = vector_feat.contiguous()
        
        # Allocate combined tensor
        combined = torch.empty(B, N, S + V, device=scalar_feat.device, dtype=scalar_feat.dtype)
        
        # Compute norms and concatenate
        se3_ops.launch_compute_norms_and_concat(
            scalar_feat, vector_feat, combined, B, N, S, V)
        
        # Scalar linear output
        combined_flat = combined.view(-1, S + V)
        s_out = self.scalar_linear(combined_flat).view(B, N, -1)
        
        # Vector transform
        v_out = torch.empty(B, N, self.vector_weight.shape[0], 3, device=vector_feat.device, dtype=vector_feat.dtype)
        se3_ops.launch_vector_transform(
            vector_feat, self.vector_weight, v_out, B, N, V, self.vector_weight.shape[0])
        
        # Gate computation
        gate = torch.sigmoid(self.vector_gate(combined_flat)).view(B, N, -1)
        
        # Apply gate
        se3_ops.launch_apply_gate(v_out, gate, B, N, self.vector_weight.shape[0])
        
        return s_out, v_out