"""
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
_cpp_src=r"""
#include <torch/extension.h>
torch::Tensor sat_forward_out(torch::Tensor robot_min, torch::Tensor robot_max, torch::Tensor env_triangle, torch::Tensor out, double eps);
"""
_cuda_src=r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
__device__ __forceinline__ void mm3(float a,float b,float c,float &mn,float &mx){mn=fminf(a,fminf(b,c)); mx=fmaxf(a,fmaxf(b,c));}
__device__ __forceinline__ bool sep(float nx,float ny,float xmin,float ymin,float xmax,float ymax,float x0,float y0,float x1,float y1,float x2,float y2,float eps){float cx=(xmin+xmax)*0.5f,cy=(ymin+ymax)*0.5f,ex=(xmax-xmin)*0.5f,ey=(ymax-ymin)*0.5f;float cd=cx*nx+cy*ny,r=fabsf(nx)*ex+fabsf(ny)*ey;float p0=x0*nx+y0*ny,p1=x1*nx+y1*ny,p2=x2*nx+y2*ny,mn,mx;mm3(p0,p1,p2,mn,mx);return (cd+r<mn-eps)||(mx<cd-r-eps);} 
__global__ __launch_bounds__(16,1) void sat_kernel(const float* __restrict__ rmin,const float* __restrict__ rmax,const float* __restrict__ tri,float* __restrict__ out,int B,float eps){int i=threadIdx.x+blockIdx.x*16;if(i>=B)return;float xmin=rmin[i*2],ymin=rmin[i*2+1],xmax=rmax[i*2],ymax=rmax[i*2+1];const float*t=tri+i*6;float x0=t[0],y0=t[1],x1=t[2],y1=t[3],x2=t[4],y2=t[5];float mn,mx;mm3(x0,x1,x2,mn,mx);if((xmax<mn-eps)||(mx<xmin-eps)){out[i]=0.f;return;}mm3(y0,y1,y2,mn,mx);if((ymax<mn-eps)||(mx<ymin-eps)){out[i]=0.f;return;}float ex=x1-x0,ey=y1-y0;float inv=1.f/(sqrtf(ex*ex+ey*ey)+1.0e-8f);if(sep(-ey*inv,ex*inv,xmin,ymin,xmax,ymax,x0,y0,x1,y1,x2,y2,eps)){out[i]=0.f;return;}ex=x2-x1;ey=y2-y1;inv=1.f/(sqrtf(ex*ex+ey*ey)+1.0e-8f);if(sep(-ey*inv,ex*inv,xmin,ymin,xmax,ymax,x0,y0,x1,y1,x2,y2,eps)){out[i]=0.f;return;}ex=x0-x2;ey=y0-y2;inv=1.f/(sqrtf(ex*ex+ey*ey)+1.0e-8f);if(sep(-ey*inv,ex*inv,xmin,ymin,xmax,ymax,x0,y0,x1,y1,x2,y2,eps)){out[i]=0.f;return;}out[i]=1.f;}
torch::Tensor sat_forward_out(torch::Tensor robot_min, torch::Tensor robot_max, torch::Tensor env_triangle, torch::Tensor out, double eps){int B=(int)robot_min.size(0);if(B>0){sat_kernel<<<(B+15)/16,16>>>(robot_min.data_ptr<float>(),robot_max.data_ptr<float>(),env_triangle.data_ptr<float>(),out.data_ptr<float>(),B,(float)eps);}return out;}
"""
_ext=load_inline(name="sat_rect_tri_ext_final",cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=["sat_forward_out"],verbose=False,extra_cuda_cflags=["-O3","--use_fast_math"])
class ModelNew(nn.Module):
    def __init__(self,eps:float=1e-6): super().__init__(); self.eps=float(eps); self._out=None
    def forward(self,robot_min,robot_max,env_triangle):
        B=robot_min.shape[0]
        if self._out is None or self._out.shape[0]!=B or self._out.device!=robot_min.device or self._out.dtype!=robot_min.dtype:
            self._out=torch.empty((B,1),device=robot_min.device,dtype=robot_min.dtype)
        return _ext.sat_forward_out(robot_min,robot_max,env_triangle,self._out,self.eps)