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 torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

torch::Tensor fk_forward_cuda(torch::Tensor joint_angles, double link_scale);

torch::Tensor fk_forward(torch::Tensor joint_angles, double link_scale) {
    return fk_forward_cuda(joint_angles, link_scale);
}
'''

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

__global__ void fk7_warp_kernel(const float* __restrict__ a, float* __restrict__ out, float scale) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int lane = threadIdx.x & 31;
    int warp = tid >> 5;
    if (warp >= 32 || lane >= 7) return;

    const unsigned mask = 0x7fu;
    float h = a[warp * 7 + lane];
    float v;

    v = __shfl_up_sync(mask, h, 1); if (lane >= 1) h += v;
    v = __shfl_up_sync(mask, h, 2); if (lane >= 2) h += v;
    v = __shfl_up_sync(mask, h, 4); if (lane >= 4) h += v;

    float s, c;
    __sincosf(h, &s, &c);
    float x = scale * c;
    float y = scale * s;

    v = __shfl_up_sync(mask, x, 1); if (lane >= 1) x += v;
    v = __shfl_up_sync(mask, x, 2); if (lane >= 2) x += v;
    v = __shfl_up_sync(mask, x, 4); if (lane >= 4) x += v;

    v = __shfl_up_sync(mask, y, 1); if (lane >= 1) y += v;
    v = __shfl_up_sync(mask, y, 2); if (lane >= 2) y += v;
    v = __shfl_up_sync(mask, y, 4); if (lane >= 4) y += v;

    long long o = (long long)warp * 14 + lane * 2;
    out[o] = x;
    out[o + 1] = y;
}

__global__ void fk_kernel(const float* __restrict__ angles, float* __restrict__ out,
                          int B, int J, float link_scale) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;
    const float* row = angles + (long long)b * J;
    float heading = 0.0f;
    float x = 0.0f;
    float y = 0.0f;
    long long o = (long long)b * J * 2;
    for (int j = 0; j < J; ++j) {
        heading += row[j];
        float s, c;
        __sincosf(heading, &s, &c);
        x += link_scale * c;
        y += link_scale * s;
        out[o + j * 2] = x;
        out[o + j * 2 + 1] = y;
    }
}

torch::Tensor fk_forward_cuda(torch::Tensor joint_angles, double link_scale_d) {
    const int B = (int)joint_angles.size(0);
    const int J = (int)joint_angles.size(1);
    auto out = torch::empty({B, J, 2}, joint_angles.options());
    const float scale = (float)link_scale_d;

    if (B == 32 && J == 7) {
        fk7_warp_kernel<<<8, 128, 0, at::cuda::getCurrentCUDAStream()>>>(
            joint_angles.data_ptr<float>(), out.data_ptr<float>(), scale);
    } else {
        fk_kernel<<<(B + 127) / 128, 128, 0, at::cuda::getCurrentCUDAStream()>>>(
            joint_angles.data_ptr<float>(), out.data_ptr<float>(), B, J, scale);
    }
    return out;
}
'''

_ext = load_inline(
    name='fk_planar_ext_v3_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['fk_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, num_joints: int, link_scale: float = 1.0):
        super().__init__()
        self.num_joints = num_joints
        self.link_scale = float(link_scale)

    def forward(self, joint_angles: torch.Tensor) -> torch.Tensor:
        return _ext.fk_forward(joint_angles, self.link_scale)