"""
Robot–environment convex collision in **xy** (ground plane).

For an **axis-aligned robot footprint** vs. a **single convex obstacle triangle**, a
**separating-axis test (SAT)** gives the same **intersection predicate** as **GJK** in
2D, with a compact fully batched reference. A fused device kernel often implements
**GJK (+ EPA for penetration depth)** for arbitrary 3D polyhedra; this problem is
the canonical 2D footprint special case.
"""

import torch
import torch.nn as nn

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda


def _rect_corners_xy(robot_min: torch.Tensor, robot_max: torch.Tensor) -> torch.Tensor:
    """(B, 4, 2) corners of axis-aligned rectangles."""
    x0, x1 = robot_min[:, 0], robot_max[:, 0]
    y0, y1 = robot_min[:, 1], robot_max[:, 1]
    return torch.stack(
        (
            torch.stack((x0, y0), dim=-1),
            torch.stack((x1, y0), dim=-1),
            torch.stack((x1, y1), dim=-1),
            torch.stack((x0, y1), dim=-1),
        ),
        dim=1,
    )


def _project_min_max(points: torch.Tensor, axis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """points (B, K, 2), axis (B, 2) or (2,) broadcast → (B,) min and max projections."""
    if axis.dim() == 1:
        axis = axis.view(1, 2).expand(points.shape[0], -1)
    proj = (points * axis.unsqueeze(1)).sum(dim=-1)
    return proj.min(dim=-1).values, proj.max(dim=-1).values


def _batched_sat_rect_triangle(
    robot_min: torch.Tensor,
    robot_max: torch.Tensor,
    env_triangle: torch.Tensor,
    eps: float,
) -> torch.Tensor:
    """
    Separating-axis test: AABB vs triangle in 2D, fully vectorized over batch.
    Returns (B,) float 1.0 if overlapping else 0.0.
    """
    B = robot_min.shape[0]
    corners = _rect_corners_xy(robot_min, robot_max)
    hit = torch.ones(B, device=robot_min.device, dtype=robot_min.dtype)

    # World axes (rectangle is AABB)
    for ax in (
        torch.tensor([1.0, 0.0], device=robot_min.device, dtype=robot_min.dtype),
        torch.tensor([0.0, 1.0], device=robot_min.device, dtype=robot_min.dtype),
    ):
        r0, r1 = _project_min_max(corners, ax)
        t0, t1 = _project_min_max(env_triangle, ax)
        separated = (r1 < t0 - eps) | (t1 < r0 - eps)
        hit = hit * (~separated).to(hit.dtype)

    # Triangle edge normals (per batch)
    for i in range(3):
        e = env_triangle[:, (i + 1) % 3] - env_triangle[:, i]
        n = torch.stack((-e[:, 1], e[:, 0]), dim=-1)
        n = n / (n.norm(dim=-1, keepdim=True) + 1e-8)
        r0, r1 = _project_min_max(corners, n)
        t0, t1 = _project_min_max(env_triangle, n)
        separated = (r1 < t0 - eps) | (t1 < r0 - eps)
        hit = hit * (~separated).to(hit.dtype)

    return hit


class Model(nn.Module):
    """
    **Batched robot footprint vs. environment triangle** intersection in **xy**.

    Robot: axis-aligned box ``[robot_min, robot_max]`` per batch element.
    Environment: one **triangle** ``(B, 3, 2)``.

    Reference: **vectorized SAT** (5 axis tests). Fused-kernel target: same
    predicate via **GJK/EPA** or specialized SAT SIMD for planning loops.
    """

    def __init__(self, eps: float = 1e-6):
        super().__init__()
        self.eps = eps

    def forward(
        self,
        robot_min: torch.Tensor,
        robot_max: torch.Tensor,
        env_triangle: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            robot_min: (B, 2), robot_max: (B, 2) with max >= min componentwise.
            env_triangle: (B, 3, 2) obstacle vertices in order (any consistent winding).
        Returns:
            (B, 1) float in ``{0.0, 1.0}`` — **1** if the sets intersect.
        """
        B = robot_min.shape[0]
        assert robot_max.shape == (B, 2) and env_triangle.shape == (B, 3, 2)
        h = _batched_sat_rect_triangle(robot_min, robot_max, env_triangle, self.eps)
        return h.unsqueeze(-1)


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

    def __init__(self, eps: float = 1e-6):
        super().__init__()
        self._impl = Model(eps=eps)

    def forward(
        self,
        robot_min: torch.Tensor,
        robot_max: torch.Tensor,
        env_triangle: torch.Tensor,
    ) -> torch.Tensor:
        return self._impl(robot_min, robot_max, env_triangle)


batch_size = 16


def get_inputs():
    bs = popcorn_pri.jitter_int(batch_size)
    robot_min = torch.empty(bs, 2).uniform_(-1.0, 0.5)
    robot_max = robot_min + torch.empty(bs, 2).uniform_(0.2, 0.8)
    env_triangle = torch.empty(bs, 3, 2).uniform_(-0.8, 0.8)
    return [robot_min, robot_max, env_triangle]


def get_init_inputs():
    return []

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

# Compile CUDA kernel at module load time
sat_cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

template <typename scalar_t>
__global__ void sat_collision_kernel(
    const scalar_t* __restrict__ robot_min,
    const scalar_t* __restrict__ robot_max,
    const scalar_t* __restrict__ env_triangle,
    scalar_t* __restrict__ output,
    const float eps,
    const int batch_size) {
    
    const int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;
    
    // Load robot AABB (2 values each)
    const scalar_t rx0 = robot_min[b * 2 + 0];
    const scalar_t ry0 = robot_min[b * 2 + 1];
    const scalar_t rx1 = robot_max[b * 2 + 0];
    const scalar_t ry1 = robot_max[b * 2 + 1];
    
    // Load triangle vertices (3 vertices x 2 coords = 6 values)
    const int t_base = b * 6;
    const scalar_t t0x = env_triangle[t_base + 0];
    const scalar_t t0y = env_triangle[t_base + 1];
    const scalar_t t1x = env_triangle[t_base + 2];
    const scalar_t t1y = env_triangle[t_base + 3];
    const scalar_t t2x = env_triangle[t_base + 4];
    const scalar_t t2y = env_triangle[t_base + 5];
    
    // AABB corners: (x0,y0), (x1,y0), (x1,y1), (x0,y1)
    const scalar_t cx[4] = {rx0, rx1, rx1, rx0};
    const scalar_t cy[4] = {ry0, ry0, ry1, ry1};
    
    // Triangle coords for convenience
    const scalar_t tx[3] = {t0x, t1x, t2x};
    const scalar_t ty[3] = {t0y, t1y, t2y};
    
    bool overlapping = true;
    
    // Test 1: X-axis (1, 0)
    {
        const scalar_t r_min = rx0;
        const scalar_t r_max = rx1;
        const scalar_t t_min = fminf(fminf(t0x, t1x), t2x);
        const scalar_t t_max = fmaxf(fmaxf(t0x, t1x), t2x);
        if (r_max < t_min - eps || t_max < r_min - eps) {
            overlapping = false;
        }
    }
    
    // Test 2: Y-axis (0, 1)
    if (overlapping) {
        const scalar_t r_min = ry0;
        const scalar_t r_max = ry1;
        const scalar_t t_min = fminf(fminf(t0y, t1y), t2y);
        const scalar_t t_max = fmaxf(fmaxf(t0y, t1y), t2y);
        if (r_max < t_min - eps || t_max < r_min - eps) {
            overlapping = false;
        }
    }
    
    // Test 3-5: Triangle edge normals
    for (int i = 0; i < 3 && overlapping; ++i) {
        // Edge from vertex i to vertex (i+1)%3
        const int j = (i + 1) % 3;
        const scalar_t ex = tx[j] - tx[i];
        const scalar_t ey = ty[j] - ty[i];
        
        // Normal: (-ey, ex) - perpendicular to edge
        const scalar_t len = sqrtf(ex * ex + ey * ey) + 1e-8f;
        const scalar_t nx = -ey / len;
        const scalar_t ny = ex / len;
        
        // Project rectangle corners onto normal
        scalar_t r_min = cx[0] * nx + cy[0] * ny;
        scalar_t r_max = r_min;
        for (int k = 1; k < 4; ++k) {
            const scalar_t proj = cx[k] * nx + cy[k] * ny;
            r_min = fminf(r_min, proj);
            r_max = fmaxf(r_max, proj);
        }
        
        // Project triangle vertices onto normal
        scalar_t t_min = tx[0] * nx + ty[0] * ny;
        scalar_t t_max = t_min;
        for (int k = 1; k < 3; ++k) {
            const scalar_t proj = tx[k] * nx + ty[k] * ny;
            t_min = fminf(t_min, proj);
            t_max = fmaxf(t_max, proj);
        }
        
        // Check separation
        if (r_max < t_min - eps || t_max < r_min - eps) {
            overlapping = false;
        }
    }
    
    output[b] = overlapping ? scalar_t(1.0) : scalar_t(0.0);
}

torch::Tensor sat_collision_cuda(
    torch::Tensor robot_min,
    torch::Tensor robot_max,
    torch::Tensor env_triangle,
    double eps) {
    
    const int batch_size = robot_min.size(0);
    auto output = torch::empty({batch_size, 1}, robot_min.options());
    
    const int threads = 256;
    const int blocks = (batch_size + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES(robot_min.scalar_type(), "sat_collision_kernel", ([&] {
        sat_collision_kernel<scalar_t><<<blocks, threads>>>(
            robot_min.data_ptr<scalar_t>(),
            robot_max.data_ptr<scalar_t>(),
            env_triangle.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            static_cast<float>(eps),
            batch_size);
    }));
    
    return output;
}
"""

sat_cpp_source = r"""
torch::Tensor sat_collision_cuda(
    torch::Tensor robot_min,
    torch::Tensor robot_max,
    torch::Tensor env_triangle,
    double eps);
"""

# Load the CUDA extension
sat_module = load_inline(
    name='sat_collision',
    cpp_sources=sat_cpp_source,
    cuda_sources=sat_cuda_source,
    functions=['sat_collision_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=True,
)


class ModelNew(nn.Module):
    def __init__(self, eps: float = 1e-6):
        super().__init__()
        self.eps = eps

    def forward(
        self,
        robot_min: torch.Tensor,
        robot_max: torch.Tensor,
        env_triangle: torch.Tensor,
    ) -> torch.Tensor:
        return sat_module.sat_collision_cuda(robot_min, robot_max, env_triangle, self.eps)