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_cuda(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>

// Specialized for the benchmark dimensions:
// scalar_in=64, scalar_out=64, vector_in=16, vector_out=16.
constexpr int SI = 64;
constexpr int SO = 64;
constexpr int VI = 16;
constexpr int VO = 16;
constexpr int COMB = 80;

__global__ __launch_bounds__(128)
void se3_fused_kernel(const float* __restrict__ scalar,
                      const float* __restrict__ vec,
                      const float* __restrict__ sw,
                      const float* __restrict__ sb,
                      const float* __restrict__ vw,
                      const float* __restrict__ gw,
                      const float* __restrict__ gb,
                      float* __restrict__ sout,
                      float* __restrict__ vout,
                      int rows) {
    int row = blockIdx.x;
    if (row >= rows) return;
    int tid = threadIdx.x;

    __shared__ float comb[COMB];
    __shared__ float gate[VO];
    __shared__ float vtmp[VO * 3];

    const float* srow = scalar + row * SI;
    const float* vrow = vec + row * VI * 3;

    // Build [scalar_features, vector_norms].
    if (tid < SI) {
        comb[tid] = __ldg(srow + tid);
    }
    if (tid < VI) {
        float x = __ldg(vrow + tid * 3 + 0);
        float y = __ldg(vrow + tid * 3 + 1);
        float z = __ldg(vrow + tid * 3 + 2);
        comb[SI + tid] = sqrtf(fmaf(x, x, fmaf(y, y, z * z)));
    }

    // Compute the vector linear transform while other threads prepare comb.
    // One thread per (vector_out, xyz) output.
    if (tid >= 80 && tid < 128) {
        int idx = tid - 80;      // 0..47
        int o = idx / 3;
        int d = idx - o * 3;
        float acc = 0.0f;
        #pragma unroll
        for (int i = 0; i < VI; ++i) {
            acc = fmaf(__ldg(vw + o * VI + i), __ldg(vrow + i * 3 + d), acc);
        }
        vtmp[idx] = acc;
    }

    __syncthreads();

    // Scalar linear layer: 64 outputs.
    if (tid < SO) {
        float acc = __ldg(sb + tid);
        const float* wrow = sw + tid * COMB;
        #pragma unroll
        for (int k = 0; k < COMB; ++k) {
            acc = fmaf(comb[k], __ldg(wrow + k), acc);
        }
        sout[row * SO + tid] = acc;
    }

    // Gate linear + sigmoid: 16 outputs, in parallel with scalar outputs.
    if (tid >= 64 && tid < 80) {
        int o = tid - 64;
        float acc = __ldg(gb + o);
        const float* wrow = gw + o * COMB;
        #pragma unroll
        for (int k = 0; k < COMB; ++k) {
            acc = fmaf(comb[k], __ldg(wrow + k), acc);
        }
        gate[o] = 1.0f / (1.0f + expf(-acc));
    }

    __syncthreads();

    // Apply gate and write vector output.
    if (tid < VO * 3) {
        int idx = tid;
        int o = idx / 3;
        vout[row * VO * 3 + idx] = gate[o] * vtmp[idx];
    }
}

std::vector<torch::Tensor> se3_forward_cuda(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) {
    const int64_t B = scalar_feat.size(0);
    const int64_t N = scalar_feat.size(1);
    const int rows = static_cast<int>(B * N);

    auto sout = torch::empty({B, N, SO}, scalar_feat.options());
    auto vout = torch::empty({B, N, VO, 3}, vector_feat.options());

    dim3 block(128);
    dim3 grid(rows);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    se3_fused_kernel<<<grid, block, 0, stream>>>(
        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>(),
        sout.data_ptr<float>(),
        vout.data_ptr<float>(),
        rows);

    return {sout, vout};
}
'''

_ext = load_inline(
    name='se3_invariant_fused_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['se3_forward_cuda'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
    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:
        s, v = _ext.se3_forward_cuda(
            scalar_feat,
            vector_feat,
            self.scalar_linear.weight,
            self.scalar_linear.bias,
            self.vector_weight,
            self.vector_gate.weight,
            self.vector_gate.bias,
        )
        return s, v