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>

__global__ void se3_linear_kernel(
    const float* __restrict__ scalar_feat,
    const float* __restrict__ vector_feat,
    const float* __restrict__ scalar_weight,
    const float* __restrict__ scalar_bias,
    const float* __restrict__ vector_weight,
    const float* __restrict__ gate_weight,
    const float* __restrict__ gate_bias,
    float* __restrict__ scalar_out,
    float* __restrict__ vector_out,
    int B, int N, int scalar_in, int scalar_out_dim, int vector_in, int vector_out_dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N;
    if (idx >= total) return;

    extern __shared__ float smem[];
    float* vec_norm = smem;
    float* combined = vec_norm + vector_in;
    float* gate_local = combined + scalar_in + vector_in;

    // vec_norm
    #pragma unroll 4
    for (int i = 0; i < vector_in; ++i) {
        float sum = 0.0f;
        for (int d = 0; d < 3; ++d) {
            float v = vector_feat[((idx * vector_in) + i) * 3 + d];
            sum += v * v;
        }
        vec_norm[i] = sqrtf(sum);
    }

    // combined
    for (int i = 0; i < scalar_in; ++i) {
        combined[i] = scalar_feat[idx * scalar_in + i];
    }
    for (int i = 0; i < vector_in; ++i) {
        combined[scalar_in + i] = vec_norm[i];
    }

    // scalar linear
    for (int o = 0; o < scalar_out_dim; ++o) {
        float sum = scalar_bias[o];
        for (int i = 0; i < scalar_in + vector_in; ++i) {
            sum += combined[i] * scalar_weight[o * (scalar_in + vector_in) + i];
        }
        scalar_out[idx * scalar_out_dim + o] = sum;
    }

    // gate
    for (int o = 0; o < vector_out_dim; ++o) {
        float sum = gate_bias[o];
        for (int i = 0; i < scalar_in + vector_in; ++i) {
            sum += combined[i] * gate_weight[o * (scalar_in + vector_in) + i];
        }
        gate_local[o] = 1.0f / (1.0f + expf(-sum));
    }

    // vector transform + gate
    for (int o = 0; o < vector_out_dim; ++o) {
        float g = gate_local[o];
        for (int d = 0; d < 3; ++d) {
            float sum = 0.0f;
            for (int i = 0; i < vector_in; ++i) {
                sum += vector_weight[o * vector_in + i] * vector_feat[(idx * vector_in + i) * 3 + d];
            }
            vector_out[(idx * vector_out_dim + o) * 3 + d] = g * sum;
        }
    }
}

std::tuple<torch::Tensor, torch::Tensor> se3_forward_cuda(
    torch::Tensor scalar_feat,
    torch::Tensor vector_feat,
    torch::Tensor scalar_weight,
    torch::Tensor scalar_bias,
    torch::Tensor vector_weight,
    torch::Tensor gate_weight,
    torch::Tensor gate_bias
) {
    int B = scalar_feat.size(0);
    int N = scalar_feat.size(1);
    int sc_in = scalar_feat.size(2);
    int v_in = vector_feat.size(2);
    int sc_out = scalar_weight.size(0);
    int v_out = vector_weight.size(0);

    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(scalar_feat.device());
    auto s_out = torch::empty({B, N, sc_out}, options);
    auto v_out = torch::empty({B, N, v_out, 3}, options);

    int total = B * N;
    dim3 blocks((total + 255) / 256);
    dim3 threads(256);
    size_t smem = (v_in + sc_in + v_in + v_out) * sizeof(float);

    se3_linear_kernel<<<blocks, threads, smem>>>(
        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>(),
        gate_weight.data_ptr<float>(),
        gate_bias.data_ptr<float>(),
        s_out.data_ptr<float>(),
        v_out.data_ptr<float>(),
        B, N, sc_in, sc_out, v_in, v_out
    );

    return std::make_tuple(s_out, v_out);
}
'''

cpp_source = '''
#include <torch/extension.h>
std::tuple<torch::Tensor, torch::Tensor> se3_forward_cuda(
    torch::Tensor scalar_feat, torch::Tensor vector_feat,
    torch::Tensor scalar_weight, torch::Tensor scalar_bias,
    torch::Tensor vector_weight, torch::Tensor gate_weight, torch::Tensor gate_bias);
'''

mod = load_inline(
    name='se3_mod',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['se3_forward_cuda'],
    verbose=False
)

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, vector_feat):
        return mod.se3_forward_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
        )