import torch
import torch.nn as nn
import torch.nn.functional as F

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).


def _quaternion_to_rotation_matrix(q: torch.Tensor) -> torch.Tensor:
    """
    Unit quaternion ``q = (w, x, y, z)`` (scalar-first) to rotation matrix ``R``,
    shape ``(B, 3, 3)``, with ``p' = p @ R.T`` for row vectors ``p`` (``N×3``).
    """
    q = F.normalize(q, dim=-1, eps=1e-8)
    w, x, y, z = q.unbind(dim=-1)

    xx = x * x
    yy = y * y
    zz = z * z
    xy = x * y
    xz = x * z
    yz = y * z
    wx = w * x
    wy = w * y
    wz = w * z

    m00 = 1.0 - 2.0 * (yy + zz)
    m01 = 2.0 * (xy - wz)
    m02 = 2.0 * (xz + wy)
    m10 = 2.0 * (xy + wz)
    m11 = 1.0 - 2.0 * (xx + zz)
    m12 = 2.0 * (yz - wx)
    m20 = 2.0 * (xz - wy)
    m21 = 2.0 * (yz + wx)
    m22 = 1.0 - 2.0 * (xx + yy)

    row0 = torch.stack((m00, m01, m02), dim=-1)
    row1 = torch.stack((m10, m11, m12), dim=-1)
    row2 = torch.stack((m20, m21, m22), dim=-1)
    return torch.stack((row0, row1, row2), dim=-2)


class Model(nn.Module):
    """
    **Fused vectorized pose transform**: rotation from unit quaternion ``q`` and
    translation ``t`` applied to a batch of 3D points.

    One forward fuses **quaternion → ``3×3``** (Hamilton, scalar-first ``w,x,y,z``)
    with **batched matmul** and **broadcast translation**—typical for LiDAR / camera
    extrinsics, hand-eye calibration, and world ↔ body frame changes in VLA stacks.

    ``points' = points @ R.T + t`` with ``points`` row-wise ``(B, N, 3)``.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        points: torch.Tensor,
        quat: torch.Tensor,
        trans: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            points: (B, N, 3) points in the **source** frame.
            quat: (B, 4) unit quaternion ``(w, x, y, z)`` mapping source → target rotation.
            trans: (B, 3) translation in the **target** frame (same convention as ``R``).
        Returns:
            (B, N, 3) points in the **target** frame.
        """
        B, N, _ = points.shape
        assert quat.shape == (B, 4) and trans.shape == (B, 3)

        R = _quaternion_to_rotation_matrix(quat)
        return torch.matmul(points, R.transpose(-1, -2)) + trans.unsqueeze(1)


class ModelNew(nn.Module):
    """KernelBench candidate entry point; default delegates to ``Model``."""

    def __init__(self):
        super().__init__()
        self._impl = Model()

    def forward(
        self,
        points: torch.Tensor,
        quat: torch.Tensor,
        trans: torch.Tensor,
    ) -> torch.Tensor:
        return self._impl(points, quat, trans)


batch_size = 8
num_points = 512


def get_inputs():
    p = popcorn_pri
    points = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_points), 3)
    quat = torch.randn(p.jitter_int(batch_size), 4)
    trans = torch.randn(p.jitter_int(batch_size), 3)
    return [points, quat, trans]


def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor quat_transform_cuda(torch::Tensor points, torch::Tensor quat, torch::Tensor trans);

torch::Tensor quat_transform(torch::Tensor points, torch::Tensor quat, torch::Tensor trans) {
    return quat_transform_cuda(points, quat, trans);
}
"""

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

__global__ void quat_pose_kernel(
    const float* __restrict__ points,
    const float* __restrict__ quat,
    const float* __restrict__ trans,
    float* __restrict__ out,
    int B,
    int N)
{
    int b = blockIdx.x;
    if (b >= B) return;

    __shared__ float s[12]; // 9 rotation entries followed by translation

    if (threadIdx.x == 0) {
        const float* qptr = quat + b * 4;
        float w = qptr[0];
        float x = qptr[1];
        float y = qptr[2];
        float z = qptr[3];

        float n2 = w*w + x*x + y*y + z*z;
        float n = sqrtf(n2);
        float inv = (n > 1.0e-8f) ? (1.0f / n) : 1.0e8f;
        w *= inv; x *= inv; y *= inv; z *= inv;

        float xx = x * x;
        float yy = y * y;
        float zz = z * z;
        float xy = x * y;
        float xz = x * z;
        float yz = y * z;
        float wx = w * x;
        float wy = w * y;
        float wz = w * z;

        s[0] = 1.0f - 2.0f * (yy + zz);      // m00
        s[1] = 2.0f * (xy - wz);             // m01
        s[2] = 2.0f * (xz + wy);             // m02
        s[3] = 2.0f * (xy + wz);             // m10
        s[4] = 1.0f - 2.0f * (xx + zz);      // m11
        s[5] = 2.0f * (yz - wx);             // m12
        s[6] = 2.0f * (xz - wy);             // m20
        s[7] = 2.0f * (yz + wx);             // m21
        s[8] = 1.0f - 2.0f * (xx + yy);      // m22

        const float* tptr = trans + b * 3;
        s[9]  = tptr[0];
        s[10] = tptr[1];
        s[11] = tptr[2];
    }
    __syncthreads();

    const int base = b * N * 3;
    for (int n = threadIdx.x; n < N; n += blockDim.x) {
        const int idx = base + n * 3;
        float px = points[idx + 0];
        float py = points[idx + 1];
        float pz = points[idx + 2];

        // points @ R.T + trans, where R is row-major as built above.
        out[idx + 0] = px * s[0] + py * s[1] + pz * s[2] + s[9];
        out[idx + 1] = px * s[3] + py * s[4] + pz * s[5] + s[10];
        out[idx + 2] = px * s[6] + py * s[7] + pz * s[8] + s[11];
    }
}

torch::Tensor quat_transform_cuda(torch::Tensor points, torch::Tensor quat, torch::Tensor trans) {
    TORCH_CHECK(points.is_cuda(), "points must be a CUDA tensor");
    TORCH_CHECK(quat.is_cuda(), "quat must be a CUDA tensor");
    TORCH_CHECK(trans.is_cuda(), "trans must be a CUDA tensor");
    TORCH_CHECK(points.scalar_type() == at::kFloat, "points must be float32");
    TORCH_CHECK(quat.scalar_type() == at::kFloat, "quat must be float32");
    TORCH_CHECK(trans.scalar_type() == at::kFloat, "trans must be float32");
    TORCH_CHECK(points.dim() == 3 && points.size(2) == 3, "points must have shape (B,N,3)");
    TORCH_CHECK(quat.dim() == 2 && quat.size(1) == 4, "quat must have shape (B,4)");
    TORCH_CHECK(trans.dim() == 2 && trans.size(1) == 3, "trans must have shape (B,3)");
    TORCH_CHECK(points.size(0) == quat.size(0) && points.size(0) == trans.size(0), "batch dimensions must match");
    TORCH_CHECK(points.is_contiguous() && quat.is_contiguous() && trans.is_contiguous(), "inputs must be contiguous");

    auto out = torch::empty_like(points);
    const int B = (int)points.size(0);
    const int N = (int)points.size(1);

    if (B > 0 && N > 0) {
        const int threads = 256;
        dim3 grid(B);
        quat_pose_kernel<<<grid, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            points.data_ptr<float>(), quat.data_ptr<float>(), trans.data_ptr<float>(), out.data_ptr<float>(), B, N);
    }
    return out;
}
"""

_ext = load_inline(
    name="quat_pose_transform_inline_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["quat_transform"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, points: torch.Tensor, quat: torch.Tensor, trans: torch.Tensor) -> torch.Tensor:
        return _ext.quat_transform(points, quat, trans)