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

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

__global__ void se3_linear_kernel(
    const float* __restrict__ scalar_feat,
    const float* __restrict__ vector_feat,
    const float* __restrict__ scalar_w,
    const float* __restrict__ scalar_b,
    const float* __restrict__ vector_w,
    const float* __restrict__ gate_w,
    const float* __restrict__ gate_b,
    float* __restrict__ s_out,
    float* __restrict__ v_out,
    int B, int N, int si, int so, int vi, int vo
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N;
    if (idx >= total) return;

    int b = idx / N;
    int n = idx % N;

    float norm[16];
    for(int i=0; i<vi; i++) {
        float sum = 0.0f;
        for(int d=0; d<3; d++) {
            float v = vector_feat[((b * N + n) * vi + i) * 3 + d];
            sum += v * v;
        }
        norm[i] = sqrtf(sum);
    }

    float combined[80];
    for(int i=0; i<si; i++) {
        combined[i] = scalar_feat[(b * N + n) * si + i];
    }
    for(int i=0; i<vi; i++) {
        combined[si + i] = norm[i];
    }

    int cd = si + vi;

    for(int o=0; o<so; o++) {
        float sum = scalar_b[o];
        for(int i=0; i<cd; i++) {
            sum += combined[i] * scalar_w[o * cd + i];
        }
        s_out[(b * N + n) * so + o] = sum;
    }

    float v_trans[48];
    for(int o=0; o<vo; o++) {
        for(int d=0; d<3; d++) {
            float sum = 0.0f;
            for(int i=0; i<vi; i++) {
                sum += vector_w[o * vi + i] * vector_feat[((b * N + n) * vi + i) * 3 + d];
            }
            v_trans[o * 3 + d] = sum;
        }
    }

    for(int o=0; o<vo; o++) {
        float sum = gate_b[o];
        for(int i=0; i<cd; i++) {
            sum += combined[i] * gate_w[o * cd + i];
        }
        float g = 1.0f / (1.0f + expf(-sum));
        for(int d=0; d<3; d++) {
            v_out[((b * N + n) * vo + o) * 3 + d] = g * v_trans[o * 3 + d];
        }
    }
}

std::vector<torch::Tensor> se3_linear_cuda(
    torch::Tensor scalar_feat,
    torch::Tensor vector_feat,
    torch::Tensor scalar_weight,
    torch::Tensor scalar_bias,
    torch::Tensor vector_weight,
    torch::Tensor vector_gate_weight,
    torch::Tensor vector_gate_bias
) {
    int B = scalar_feat.size(0);
    int N = scalar_feat.size(1);
    int si = scalar_feat.size(2);
    int vi = vector_feat.size(2);
    int vo = vector_weight.size(0);
    int so = scalar_weight.size(0);

    auto options = torch::TensorOptions().device(scalar_feat.device()).dtype(torch::kFloat32);
    auto s_out_t = torch::empty({B, N, so}, options);
    auto v_out_t = torch::empty({B, N, vo, 3}, options);

    int total = B * N;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;

    se3_linear_kernel<<<blocks, threads>>>(
        scalar_feat.data_ptr<float>(),
        vector_feat.data_ptr<float>(),
        scalar_weight.data_ptr<float>(),
        scalar_bias.data_ptr<float>(),
        vector_weight.data_ptr<float>(),
        vector_gate_weight.data_ptr<float>(),
        vector_gate_bias.data_ptr<float>(),
        s_out_t.data_ptr<float>(),
        v_out_t.data_ptr<float>(),
        B, N, si, so, vi, vo
    );

    return {s_out_t, v_out_t};
}
'''

cpp_source = '''
std::vector<torch::Tensor> se3_linear_cuda(torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_weight, torch::Tensor scalar_bias, torch::Tensor vector_weight, torch::Tensor vector_gate_weight, torch::Tensor vector_gate_bias);
'''

module = load_inline(
    name='se3_linear',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['se3_linear_cuda'],
    verbose=False,
    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)
        self._module = module

    def forward(self, scalar_feat: torch.Tensor, vector_feat: torch.Tensor):
        return self._module.se3_linear_cuda(
            scalar_feat.contiguous(),
            vector_feat.contiguous(),
            self.scalar_linear.weight,
            self.scalar_linear.bias,
            self.vector_weight,
            self.vector_gate.weight,
            self.vector_gate.bias
        )