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

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void quaternion_transform_kernel(
    const float* __restrict__ points,
    const float* __restrict__ quats,
    const float* __restrict__ trans,
    float* __restrict__ output,
    int B, int N
) {
    int b = blockIdx.x;
    int n_start = threadIdx.x;
    int stride = blockDim.x;

    // Load and normalize quaternion for this batch
    float q0 = quats[b * 4 + 0];
    float q1 = quats[b * 4 + 1];
    float q2 = quats[b * 4 + 2];
    float q3 = quats[b * 4 + 3];

    float norm = sqrtf(q0*q0 + q1*q1 + q2*q2 + q3*q3 + 1e-8f);
    float inv_norm = 1.0f / norm;
    float w = q0 * inv_norm;
    float x = q1 * inv_norm;
    float y = q2 * inv_norm;
    float z = q3 * inv_norm;

    // Compute rotation matrix elements (R)
    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;

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

    // Load translation
    float tx = trans[b * 3 + 0];
    float ty = trans[b * 3 + 1];
    float tz = trans[b * 3 + 2];

    // Each thread processes multiple points with stride
    for (int n = n_start; n < N; n += stride) {
        float px = points[b * N * 3 + n * 3 + 0];
        float py = points[b * N * 3 + n * 3 + 1];
        float pz = points[b * N * 3 + n * 3 + 2];

        // p' = p @ R.T  =>  p @ R^T means dot with columns of R, but since we have rows:
        // Actually: output = p @ R^T  =>  [px,py,pz] dot with each row of R? Wait:
        // R is row-major as constructed: rows are [m00,m01,m02], etc.
        // p @ R^T = [px*m00 + py*m10 + pz*m20, px*m01+py*m11+pz*m21, px*m02+py*m12+pz*m22]
        float ox = px * m00 + py * m10 + pz * m20;
        float oy = px * m01 + py * m11 + pz * m21;
        float oz = px * m02 + py * m12 + pz * m22;

        output[b * N * 3 + n * 3 + 0] = ox + tx;
        output[b * N * 3 + n * 3 + 1] = oy + ty;
        output[b * N * 3 + n * 3 + 2] = oz + tz;
    }
}

torch::Tensor quaternion_transform_cuda(
    torch::Tensor points,
    torch::Tensor quat,
    torch::Tensor trans
) {
    int B = points.size(0);
    int N = points.size(1);

    auto output = torch::empty({B, N, 3}, points.options());

    dim3 grid(B);
    dim3 block(256);

    quaternion_transform_kernel<<<grid, block>>>(
        points.data_ptr<float>(),
        quat.data_ptr<float>(),
        trans.data_ptr<float>(),
        output.data_ptr<float>(),
        B, N
    );

    return output;
}
'''

cpp_source = '''
torch::Tensor quaternion_transform_cuda(torch::Tensor points, torch::Tensor quat, torch::Tensor trans);
'''

module = load_inline(
    name='quaternion_transform',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['quaternion_transform_cuda'],
    verbose=False
)

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

    def forward(self, points, quat, trans):
        return module.quaternion_transform_cuda(points, quat, trans)