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_runtime.h>
#include <cmath>
#include <vector>

__global__ void compute_norm_kernel(const float* __restrict__ vector_feat, float* __restrict__ vec_norm, int total_bn, int vector_in) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = total_bn * vector_in;
    if (idx >= total) return;
    int bn = idx / vector_in;
    int vi = idx % vector_in;
    const float* vptr = vector_feat + bn * vector_in * 3 + vi * 3;
    float sum = vptr[0]*vptr[0] + vptr[1]*vptr[1] + vptr[2]*vptr[2];
    vec_norm[idx] = sqrtf(sum);
}

__global__ void matmul_bias_kernel(const float* __restrict__ input, const float* __restrict__ weight, const float* __restrict__ bias, float* __restrict__ output, int M, int K, int Out) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    int col = blockIdx.y * blockDim.y + threadIdx.y;
    if (row >= M || col >= Out) return;
    float sum = bias[col];
    for (int k = 0; k < K; ++k) {
        sum += input[row * K + k] * weight[col * K + k];
    }
    output[row * Out + col] = sum;
}

__global__ void sigmoid_kernel(float* data, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    float x = data[idx];
    data[idx] = 1.0f / (1.0f + expf(-x));
}

__global__ void vector_transform_kernel(const float* __restrict__ vector_feat, const float* __restrict__ v_weight, const float* __restrict__ gate, float* __restrict__ v_out, int total_bn, int vector_in, int vector_out) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = total_bn * vector_out * 3;
    if (idx >= total) return;
    int bn = idx / (vector_out * 3);
    int rem = idx % (vector_out * 3);
    int vo = rem / 3;
    int d = rem % 3;
    float g = gate[bn * vector_out + vo];
    float sum = 0.0f;
    const float* vw = v_weight + vo * vector_in;
    const float* vf_base = vector_feat + bn * vector_in * 3;
    for (int vi = 0; vi < vector_in; ++vi) {
        sum += vw[vi] * vf_base[vi * 3 + d];
    }
    v_out[idx] = g * sum;
}

std::vector<torch::Tensor> se3_forward(
    torch::Tensor scalar_feat,
    torch::Tensor vector_feat,
    torch::Tensor scalar_w,
    torch::Tensor scalar_b,
    torch::Tensor vec_w,
    torch::Tensor gate_w,
    torch::Tensor gate_b
) {
    auto opts = scalar_feat.options();
    int B = scalar_feat.size(0);
    int N = scalar_feat.size(1);
    int scalar_in = scalar_feat.size(2);
    int vector_in = vector_feat.size(2);
    int scalar_out = scalar_w.size(0);
    int vector_out = vec_w.size(0);
    int M = B * N;
    int K_comb = scalar_in + vector_in;

    auto vec_norm = torch::empty({B, N, vector_in}, opts);
    compute_norm_kernel<<<(M * vector_in + 255)/256, 256>>>(
        vector_feat.data_ptr<float>(), vec_norm.data_ptr<float>(), M, vector_in);

    auto combined = torch::cat({scalar_feat, vec_norm}, -1);

    auto s_out = torch::empty({B, N, scalar_out}, opts);
    dim3 block(16, 16);
    dim3 grid((M + 15)/16, (scalar_out + 15)/16);
    matmul_bias_kernel<<<grid, block>>>(
        combined.data_ptr<float>(), scalar_w.data_ptr<float>(), scalar_b.data_ptr<float>(),
        s_out.data_ptr<float>(), M, K_comb, scalar_out);

    auto gate_lin = torch::empty({M, vector_out}, opts);
    dim3 grid_g((M + 15)/16, (vector_out + 15)/16);
    matmul_bias_kernel<<<grid_g, block>>>(
        combined.data_ptr<float>(), gate_w.data_ptr<float>(), gate_b.data_ptr<float>(),
        gate_lin.data_ptr<float>(), M, K_comb, vector_out);

    sigmoid_kernel<<<(M * vector_out + 255)/256, 256>>>(gate_lin.data_ptr<float>(), M * vector_out);

    auto gate_sig = gate_lin.view({B, N, vector_out});

    auto v_out = torch::empty({B, N, vector_out, 3}, opts);
    vector_transform_kernel<<<(M * vector_out * 3 + 255)/256, 256>>>(
        vector_feat.data_ptr<float>(), vec_w.data_ptr<float>(), gate_sig.data_ptr<float>(),
        v_out.data_ptr<float>(), M, vector_in, vector_out);

    return {s_out, v_out};
}
'''

cpp_source = '''
std::vector<torch::Tensor> se3_forward(torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_w, torch::Tensor scalar_b, torch::Tensor vec_w, torch::Tensor gate_w, torch::Tensor gate_b);
'''

_module = load_inline(
    name='se3_mod',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['se3_forward'],
    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)
        self._mod = _module

    def forward(self, scalar_feat: torch.Tensor, vector_feat: torch.Tensor):
        outs = self._mod.se3_forward(
            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
        )
        return outs[0], outs[1]