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_src = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> se3_forward(
    torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_w,
    torch::Tensor scalar_b, torch::Tensor vector_w, torch::Tensor gate_w, torch::Tensor gate_b);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

__global__ void se3_forward_64_16_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__ scalar_out,
    float* __restrict__ vector_out)
{
    int row = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float comb[80];
    __shared__ float gate[16];
    __shared__ float vec[48];

    const float* srow = scalar_feat + ((size_t)row << 6);
    const float* vrow = vector_feat + (size_t)row * 48;

    if (tid < 64) comb[tid] = srow[tid];
    if (tid < 48) vec[tid] = vrow[tid];
    __syncthreads();

    if (tid < 16) {
        float x = vec[tid*3 + 0];
        float y = vec[tid*3 + 1];
        float z = vec[tid*3 + 2];
        comb[64 + tid] = sqrtf(x*x + y*y + z*z);
    }
    __syncthreads();

    if (tid < 80) {
        int o = tid;
        float acc;
        const float* w;
        if (o < 64) {
            w = scalar_w + (size_t)o * 80;
            acc = scalar_b[o];
        } else {
            int g = o - 64;
            w = gate_w + (size_t)g * 80;
            acc = gate_b[g];
        }
        #pragma unroll
        for (int i = 0; i < 80; ++i) {
            acc = fmaf(w[i], comb[i], acc);
        }
        if (o < 64) {
            scalar_out[((size_t)row << 6) + o] = acc;
        } else {
            int g = o - 64;
            gate[g] = 1.0f / (1.0f + expf(-acc));
        }
    }
    __syncthreads();

    if (tid < 48) {
        int o = tid / 3;
        int d = tid - o * 3;
        const float* w = vector_w + ((size_t)o << 4);
        float acc = 0.0f;
        #pragma unroll
        for (int i = 0; i < 16; ++i) {
            acc = fmaf(w[i], vec[i*3 + d], acc);
        }
        vector_out[(size_t)row * 48 + tid] = gate[o] * acc;
    }
}

std::vector<torch::Tensor> se3_forward(
    torch::Tensor scalar_feat, torch::Tensor vector_feat, torch::Tensor scalar_w,
    torch::Tensor scalar_b, torch::Tensor vector_w, torch::Tensor gate_w, torch::Tensor gate_b)
{
    int64_t B = scalar_feat.size(0);
    int64_t N = scalar_feat.size(1);
    auto s_out = torch::empty({B, N, scalar_w.size(0)}, scalar_feat.options());
    auto v_out = torch::empty({B, N, vector_w.size(0), 3}, vector_feat.options());
    int rows = (int)(B * N);
    se3_forward_64_16_kernel<<<rows, 128, 0, at::cuda::getCurrentCUDAStream()>>>(
        scalar_feat.data_ptr<float>(), vector_feat.data_ptr<float>(),
        scalar_w.data_ptr<float>(), scalar_b.data_ptr<float>(), vector_w.data_ptr<float>(),
        gate_w.data_ptr<float>(), gate_b.data_ptr<float>(), s_out.data_ptr<float>(), v_out.data_ptr<float>());
    return {s_out, v_out};
}
'''

_ext = load_inline(name='se3_fused_ext_v3', cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=['se3_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'],
                   extra_cflags=['-O3'], 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: torch.Tensor, vector_feat: torch.Tensor) -> tuple:
        return tuple(_ext.se3_forward(scalar_feat, vector_feat, self.scalar_linear.weight,
                                      self.scalar_linear.bias, self.vector_weight,
                                      self.vector_gate.weight, self.vector_gate.bias))