import math

import torch
import torch.nn as nn

# 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).


class Model(nn.Module):
    """
    **Batched forward kinematics (FK)** for a **planar serial** revolute chain.

    Fixed link lengths ``L_j`` (buffer), joint angles ``θ_j`` (batch ``B``). For each
    link, the heading is ``Σ_{k≤j} θ_k``; joint positions are the cumulative sum of
    ``L_j [cos(heading), sin(heading)]``—fully **vectorized** over ``B`` and ``J``.

    Returns positions of the **end of each link** in the base frame ``(B, J, 2)``
    (last row is the end-effector in the plane). Typical fusion target for
    batched geometry in calibration, whole-body control, and VLA motion heads.
    """

    def __init__(self, num_joints: int, link_scale: float = 1.0):
        super().__init__()
        assert num_joints >= 1
        self.num_joints = num_joints
        lengths = torch.ones(num_joints, dtype=torch.float32) * float(link_scale)
        self.register_buffer("link_lengths", lengths)

    def forward(self, joint_angles: torch.Tensor) -> torch.Tensor:
        """
        Args:
            joint_angles: (B, J) radians, one revolute per link.
        Returns:
            positions: (B, J, 2) world-frame (x, y) at each link end; ``[..., -1, :]`` is EE.
        """
        B, J = joint_angles.shape
        assert J == self.num_joints

        heading = torch.cumsum(joint_angles, dim=-1)
        c = torch.cos(heading)
        s = torch.sin(heading)
        direction = torch.stack((c, s), dim=-1)
        displacements = direction * self.link_lengths.view(1, J, 1)
        return torch.cumsum(displacements, dim=1)


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

    def __init__(self, num_joints: int, link_scale: float = 1.0):
        super().__init__()
        self._impl = Model(num_joints, link_scale=link_scale)

    def forward(self, joint_angles: torch.Tensor) -> torch.Tensor:
        return self._impl(joint_angles)


batch_size = 32
num_joints = 7
link_scale = 0.15


def get_inputs():
    p = popcorn_pri
    return [torch.empty(p.jitter_int(batch_size), num_joints).uniform_(-math.pi, math.pi)]


def get_init_inputs():
    return [num_joints, link_scale]

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

cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// One thread block per batch element, each thread handles one joint
__global__ void forward_kinematics_kernel(
    const float* __restrict__ joint_angles,  // (B, J)
    const float* __restrict__ link_lengths,  // (J,)
    float* __restrict__ positions,            // (B, J, 2)
    int B, int J
) {
    int batch_idx = blockIdx.x;
    int joint_idx = threadIdx.x;
    
    if (batch_idx >= B || joint_idx >= J) return;
    
    // Load joint angles for this batch into shared memory
    extern __shared__ float shared_data[];
    float* angles_shared = shared_data;          // J elements
    float* heading_shared = &shared_data[J];     // J elements
    float* cos_shared = &shared_data[2*J];       // J elements
    float* sin_shared = &shared_data[3*J];       // J elements
    float* disp_x_shared = &shared_data[4*J];    // J elements
    float* disp_y_shared = &shared_data[5*J];    // J elements
    
    // Load angle for this joint
    angles_shared[joint_idx] = joint_angles[batch_idx * J + joint_idx];
    __syncthreads();
    
    // Compute cumulative sum (heading) using parallel scan
    // For small J, use simple approach
    float heading = 0.0f;
    for (int i = 0; i <= joint_idx; ++i) {
        heading += angles_shared[i];
    }
    heading_shared[joint_idx] = heading;
    __syncthreads();
    
    // Compute cos/sin of heading
    float c = cosf(heading);
    float s = sinf(heading);
    cos_shared[joint_idx] = c;
    sin_shared[joint_idx] = s;
    
    // Get link length and compute displacement
    float L = link_lengths[joint_idx];
    disp_x_shared[joint_idx] = c * L;
    disp_y_shared[joint_idx] = s * L;
    __syncthreads();
    
    // Compute cumulative sum of displacements
    float cum_x = 0.0f;
    float cum_y = 0.0f;
    for (int i = 0; i <= joint_idx; ++i) {
        cum_x += disp_x_shared[i];
        cum_y += disp_y_shared[i];
    }
    
    // Write output
    positions[batch_idx * J * 2 + joint_idx * 2] = cum_x;
    positions[batch_idx * J * 2 + joint_idx * 2 + 1] = cum_y;
}

torch::Tensor forward_kinematics_cuda(torch::Tensor joint_angles, torch::Tensor link_lengths) {
    TORCH_CHECK(joint_angles.dim() == 2, "joint_angles must be 2D");
    TORCH_CHECK(link_lengths.dim() == 1, "link_lengths must be 1D");
    
    int B = joint_angles.size(0);
    int J = joint_angles.size(1);
    
    TORCH_CHECK(link_lengths.size(0) == J, "link_lengths size mismatch");
    
    auto positions = torch::empty({B, J, 2}, joint_angles.options());
    
    // Shared memory: 6 * J floats (angles, heading, cos, sin, disp_x, disp_y)
    size_t shared_mem_size = 6 * J * sizeof(float);
    
    // One block per batch, J threads per block
    int threads_per_block = J;
    int blocks = B;
    
    forward_kinematics_kernel<<<blocks, threads_per_block, shared_mem_size>>>(
        joint_angles.data_ptr<float>(),
        link_lengths.data_ptr<float>(),
        positions.data_ptr<float>(),
        B, J
    );
    
    return positions;
}
"""

cpp_source = r"""
torch::Tensor forward_kinematics_cuda(torch::Tensor joint_angles, torch::Tensor link_lengths);
"""

# Load the CUDA extension
fk_extension = load_inline(
    name="forward_kinematics_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["forward_kinematics_cuda"],
    extra_cuda_cflags=["-O3"],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, num_joints: int, link_scale: float = 1.0):
        super().__init__()
        assert num_joints >= 1
        self.num_joints = num_joints
        lengths = torch.ones(num_joints, dtype=torch.float32) * float(link_scale)
        self.register_buffer("link_lengths", lengths)

    def forward(self, joint_angles: torch.Tensor) -> torch.Tensor:
        B, J = joint_angles.shape
        assert J == self.num_joints
        
        # Ensure inputs are contiguous and on GPU
        joint_angles = joint_angles.contiguous()
        link_lengths = self.link_lengths.contiguous()
        
        return fk_extension.forward_kinematics_cuda(joint_angles, link_lengths)